kb_browser/script.tsx

import React, { useState, useEffect, useCallback } from "react";
import { createRoot } from "react-dom/client";

interface Document {
  id: number;
  title: string;
  content: string;
  snippet: string;
  source: string;
  source_type: string;
  url?: string;
  relevance_score: number;
  rank: number;
  metadata?: Record<string, any>;
}

interface SearchResponse {
  documents: Document[];
  search_time: number;
  query: string;
  total_count: number;
}

interface KnowledgeBrowserProps {
  query?: string;
  results?: Document[];
  search_type?: "semantic" | "keyword" | "hybrid";
  max_results?: number;
  elem_id?: string;
  elem_classes?: string[];
  visible?: boolean;
  onSubmit?: (query: string) => void;
  onSelect?: (document: Document) => void;
}

const KnowledgeBrowser: React.FC<KnowledgeBrowserProps> = ({
  query: initialQuery = "",
  results: initialResults = [],
  search_type: initialSearchType = "semantic",
  max_results: initialMaxResults = 10,
  elem_id,
  elem_classes = [],
  visible = true,
  onSubmit,
  onSelect,
}) => {
  const [query, setQuery] = useState(initialQuery);
  const [results, setResults] = useState<Document[]>(initialResults);
  const [searchType, setSearchType] = useState(initialSearchType);
  const [maxResults, setMaxResults] = useState(initialMaxResults);
  const [isLoading, setIsLoading] = useState(false);
  const [expandedResults, setExpandedResults] = useState<Set<number>>(new Set());
  const [searchTime, setSearchTime] = useState(0);

  const performSearch = useCallback(async (searchQuery: string) => {
    if (!searchQuery.trim()) return;

    setIsLoading(true);
    
    try {
      const response = await fetch("/api/search", {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
        },
        body: JSON.stringify({
          query: searchQuery,
          search_type: searchType,
          limit: maxResults,
          offset: 0,
        }),
      });

      if (!response.ok) {
        throw new Error(`Search failed: ${response.statusText}`);
      }

      const data: SearchResponse = await response.json();
      setResults(data.documents || []);
      setSearchTime(data.search_time || 0);
      
      // Trigger submit event
      if (onSubmit) {
        onSubmit(searchQuery);
      }
    } catch (error) {
      console.error("Search error:", error);
      setResults([]);
    } finally {
      setIsLoading(false);
    }
  }, [searchType, maxResults, onSubmit]);

  const handleSubmit = (e: React.FormEvent) => {
    e.preventDefault();
    performSearch(query);
  };

  const handleKeyDown = (e: React.KeyboardEvent) => {
    if (e.key === "Enter" && !e.shiftKey) {
      e.preventDefault();
      performSearch(query);
    }
  };

  const toggleExpanded = (documentId: number) => {
    const newExpanded = new Set(expandedResults);
    if (newExpanded.has(documentId)) {
      newExpanded.delete(documentId);
    } else {
      newExpanded.add(documentId);
    }
    setExpandedResults(newExpanded);
  };

  const handleSelect = (document: Document) => {
    if (onSelect) {
      onSelect(document);
    }
  };

  const getSourceTypeBadgeClass = (sourceType: string) => {
    const baseClass = "source-type-badge";
    switch (sourceType) {
      case "academic":
        return `${baseClass} source-type-academic`;
      case "web":
        return `${baseClass} source-type-web`;
      case "pdf":
        return `${baseClass} source-type-pdf`;
      case "code":
        return `${baseClass} source-type-code`;
      default:
        return `${baseClass} source-type-web`;
    }
  };

  const formatRelevanceScore = (score: number) => {
    return `${Math.round(score * 100)}%`;
  };

  const formatSearchTime = (time: number) => {
    return time < 1 ? `${Math.round(time * 1000)}ms` : `${time.toFixed(2)}s`;
  };

  if (!visible) {
    return null;
  }

  const containerClasses = ["kb-browser", ...elem_classes].join(" ");

  return (
    <div id={elem_id} className={containerClasses}>
      {/* Search Interface */}
      <div className="search-interface">
        <form onSubmit={handleSubmit}>
          <div className="search-input-container">
            <input
              type="text"
              className="search-input"
              placeholder="Enter your search query..."
              value={query}
              onChange={(e) => setQuery(e.target.value)}
              onKeyDown={handleKeyDown}
              disabled={isLoading}
            />
            <button
              type="submit"
              className="search-button"
              disabled={!query.trim() || isLoading}
            >
              {isLoading ? (
                <>
                  <span className="loading-spinner"></span>
                  <span style={{ marginLeft: "8px" }}>Searching...</span>
                </>
              ) : (
                "Search"
              )}
            </button>
          </div>
          
          <div className="search-options">
            <label>
              Search Type:
              <select
                className="search-type-select"
                value={searchType}
                onChange={(e) => setSearchType(e.target.value as any)}
                disabled={isLoading}
              >
                <option value="semantic">Semantic</option>
                <option value="keyword">Keyword</option>
                <option value="hybrid">Hybrid</option>
              </select>
            </label>
            
            <label>
              Max Results:
              <select
                className="search-type-select"
                value={maxResults}
                onChange={(e) => setMaxResults(parseInt(e.target.value))}
                disabled={isLoading}
              >
                <option value={5}>5</option>
                <option value={10}>10</option>
                <option value={20}>20</option>
              </select>
            </label>
          </div>
        </form>
      </div>

      {/* Results Container */}
      <div className="results-container">
        {results.length > 0 && (
          <div className="results-header">
            <div className="results-info">
              <strong>{results.length}</strong> results found in{" "}
              <strong>{formatSearchTime(searchTime)}</strong>
            </div>
          </div>
        )}

        {isLoading ? (
          <div className="empty-state">
            <div className="loading-spinner" style={{ marginBottom: "12px" }}></div>
            <p>Searching knowledge base...</p>
          </div>
        ) : results.length === 0 && query ? (
          <div className="empty-state">
            <p>No results found for "{query}"</p>
            <p style={{ fontSize: "14px", marginTop: "8px" }}>
              Try adjusting your search terms or search type.
            </p>
          </div>
        ) : results.length === 0 ? (
          <div className="empty-state">
            <p>Enter a search query to find relevant documents</p>
          </div>
        ) : (
          <div>
            {results.map((result) => (
              <div key={result.id} className="result-card">
                <div className="result-header">
                  <div className="result-meta">
                    <span className={getSourceTypeBadgeClass(result.source_type)}>
                      {result.source_type}
                    </span>
                    <span className="relevance-score">
                      {formatRelevanceScore(result.relevance_score)}
                    </span>
                  </div>
                  
                  <h3 className="result-title">{result.title}</h3>
                  <p className="result-source">{result.source}</p>
                  
                  <div className="result-snippet">{result.snippet}</div>
                  
                  <div className="result-actions">
                    <button
                      className="action-button"
                      onClick={() => toggleExpanded(result.id)}
                    >
                      {expandedResults.has(result.id) ? "Collapse" : "Expand"}
                    </button>
                    
                    {result.url && (
                      <button
                        className="action-button"
                        onClick={() => window.open(result.url, "_blank")}
                      >
                        View Source
                      </button>
                    )}
                    
                    <button
                      className="action-button primary"
                      onClick={() => handleSelect(result)}
                    >
                      Use This
                    </button>
                  </div>
                </div>
                
                {expandedResults.has(result.id) && (
                  <div className="expanded-content">
                    <h4>Full Content</h4>
                    <p style={{ lineHeight: "1.6", marginBottom: "16px" }}>
                      {result.content.length > 500
                        ? `${result.content.substring(0, 500)}...`
                        : result.content}
                    </p>
                    
                    {result.metadata && Object.keys(result.metadata).length > 0 && (
                      <>
                        <h4>Metadata</h4>
                        <div className="metadata-grid">
                          {Object.entries(result.metadata).map(([key, value]) => (
                            <div key={key} className="metadata-item">
                              <span className="metadata-label">
                                {key.charAt(0).toUpperCase() + key.slice(1)}:
                              </span>
                              <span className="metadata-value">
                                {Array.isArray(value) ? value.join(", ") : String(value)}
                              </span>
                            </div>
                          ))}
                        </div>
                      </>
                    )}
                  </div>
                )}
              </div>
            ))}
          </div>
        )}
      </div>
    </div>
  );
};

// Initialize the component when DOM is ready
const initializeComponent = () => {
  const container = document.getElementById("kb-browser-root");
  if (container) {
    const root = createRoot(container);
    root.render(
      <KnowledgeBrowser
        onSubmit={(query) => {
          console.log("Search submitted:", query);
          // Dispatch custom event for Gradio integration
          window.dispatchEvent(new CustomEvent("kb-browser-submit", { 
            detail: { query } 
          }));
        }}
        onSelect={(document) => {
          console.log("Document selected:", document);
          // Dispatch custom event for Gradio integration
          window.dispatchEvent(new CustomEvent("kb-browser-select", { 
            detail: { document } 
          }));
        }}
      />
    );
  }
};

// Initialize when DOM is loaded
if (document.readyState === "loading") {
  document.addEventListener("DOMContentLoaded", initializeComponent);
} else {
  initializeComponent();
}

export default KnowledgeBrowser;