Merge pull request #27 from ProspektStudio/david

Touch-ups

Author: dtom90

api/models.py

@@ -12,7 +12,7 @@
 
 BASE_PROMPT = """You are an expert in the topic of satellites and space stations.
 You are here to answer any questions you have regarding the topic.
-Give a short one-paragraph answer with the most relevant information."""
+Give a short one-paragraph answer with the most recent information."""
 
 def generate_prompt(group: str, name: str):
   prompt = f"Give me information about {name}"

api/rag_agent.py

@@ -14,10 +14,18 @@
 
 logger = uvicorn_logging.getLogger("uvicorn")
 
+# Model and file configuration
 EMBEDDINGS_MODEL = "text-embedding-004"
 WEB_PAGES_FILE = "webpages.txt"
 
-SYSTEM_PROMPT_FOR_COMBINED_TOOL = """You have access to the following tools:
+# RAG configuration
+CHUNK_SIZE = 1000
+CHUNK_OVERLAP = 200
+SIMILARITY_SEARCH_K = 10
+LLM_TEMPERATURE = 0.0
+
+SYSTEM_PROMPT_FOR_COMBINED_TOOL = """
+You have access to the following tools:
 1.  **`combined_search`**: Use this tool to provide a comprehensive response that combines both document-specific information and general knowledge. This tool will:
     * First search the document collection for relevant information
     * Then supplement the response with additional general knowledge
@@ -48,18 +56,33 @@ def __init__(
 
         # Load and process documents if provided
         web_pages = load_webpages(WEB_PAGES_FILE)
+        logger.info(f"Using {len(web_pages)} webpages from {WEB_PAGES_FILE}")
+        
         loader = WebBaseLoader(web_paths=web_pages)
         docs = loader.load()
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-        all_splits = text_splitter.split_documents(docs)
+        logger.info(f"Loaded {len(docs)} documents")
+        
+        # Preprocess documents
+        cleaned_docs = []
+        for doc in docs:
+            # Clean the content
+            cleaned_content = self._preprocess_content(doc.page_content)
+            if cleaned_content:  # Only keep documents with content after cleaning
+                doc.page_content = cleaned_content
+                cleaned_docs.append(doc)
+        logger.info(f"Cleaned {len(cleaned_docs)} documents")
+
+        # Split and store cleaned documents
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)
+        all_splits = text_splitter.split_documents(cleaned_docs)
         _ = self.vector_store.add_documents(documents=all_splits)
-        logger.info(f"Loaded {len(all_splits)} document chunks into vector store")
+        logger.info(f"Split into {len(all_splits)} chunks and loaded into vector store")
 
         # 2. Set up language model
         self.llm = init_chat_model(
           llm_model,
           model_provider=llm_model_provider,
-          temperature=0.0
+          temperature=LLM_TEMPERATURE
         )
 
         # 3. Set up memory
@@ -76,6 +99,24 @@ def __init__(
 
         logger.info(f"RAG Agent initialized with system prompt: \n\n{system_prompt}")
 
+    def _preprocess_content(self, content: str) -> str:
+        """
+        Simple preprocessing: replace separators with spaces and clean whitespace.
+        """
+        if not content:
+            return ""
+            
+        # Replace common separators with spaces
+        separators = ['\n', '\r', '\t', '|', '•', '→', '←', '↑', '↓', '↔', '↕', '↖', '↗', '↘', '↙']
+        for sep in separators:
+            content = content.replace(sep, ' ')
+            
+        # Clean up whitespace: replace multiple spaces with single space
+        import re
+        content = re.sub(r'\s+', ' ', content)
+        
+        return content.strip()
+
     def generate_combined_tool(self):
         @tool(response_format="content_and_artifact")
         def combined_search(query: str) -> tuple[str, list]:
@@ -84,10 +125,9 @@ def combined_search(query: str) -> tuple[str, list]:
             and general knowledge. This tool will first search the document collection and then
             supplement with general knowledge if needed.
             """
-            logger.info(f"Performing combined search for query: {query}")
 
             # First, try to retrieve from documents
-            retrieved_docs = self.vector_store.similarity_search(query, k=3)
+            retrieved_docs = self.vector_store.similarity_search(query, k=SIMILARITY_SEARCH_K)
             doc_content = ""
 
             if retrieved_docs:
@@ -98,12 +138,13 @@ def combined_search(query: str) -> tuple[str, list]:
                 )
 
             # Then, get general knowledge response
-            message = HumanMessage(content=f"User query: {query}\nPlease provide additional general knowledge that complements the following document information:\n{doc_content}")
+            message = HumanMessage(content=query)
             try:
                 general_knowledge = self.llm.invoke([message]).content
 
                 # Combine both responses
                 combined_response = f"Document Information:\n{doc_content}\n\nAdditional General Knowledge:\n{general_knowledge}"
+                logger.info(f"Completed combined search for query")
                 return combined_response, retrieved_docs
             except Exception as e:
                 logger.error(f"Error in combined search: {e}")
@@ -115,8 +156,7 @@ async def ask(self, prompt: str):
         """
         Sends a prompt to the agent and streams the final LLM response.
         """
-        print(f"\nUser Prompt: {prompt}")
-        print("Agent Response (streaming):")
+        logger.info(f"\nUser Prompt: {prompt}")
 
         # Track if we're in a tool call
         in_tool_call = False
@@ -141,6 +181,5 @@ async def ask(self, prompt: str):
                 if isinstance(chunk, AIMessageChunk):
                     # Only yield content if we're not in a tool call
                     if chunk.content:
-                        print(chunk.content, end="", flush=True)
                         final_response_buffer.append(chunk.content)
                         yield chunk.content

api/webpages.txt

@@ -4,16 +4,15 @@ https://www.spacex.com/vehicles/dragon/
 https://en.wikipedia.org/wiki/SpaceX_Dragon
 https://en.wikipedia.org/wiki/SpaceX_Dragon_2
 https://en.wikipedia.org/wiki/SpaceX_Crew-10
-https://en.wikipedia.org/wiki/Zarya_(ISS_module
+https://en.wikipedia.org/wiki/Zarya_(ISS_module)
 https://en.wikipedia.org/wiki/Tianhe_core_module
-https://en.wikipedia.org/wiki/Nauka_(ISS_module
+https://en.wikipedia.org/wiki/Nauka_(ISS_module)
 https://isstracker.pl/en/satellites/46260
 https://en.wikipedia.org/wiki/Wentian_module
 https://en.wikipedia.org/wiki/Mengtian_module
 https://en.wikipedia.org/wiki/Tianzhou_8
 https://en.wikipedia.org/wiki/Progress_MS-29
 https://nextspaceflight.com/launches/details/7637
-https://en.wikipedia.org/wiki/SpaceX_Crew-10
 https://nextspaceflight.com/launches/details/7624
 https://en.wikipedia.org/wiki/Shenzhou_20
 https://www.globalstar.com/

components/AiInfo.tsx

@@ -27,91 +27,33 @@ const agentInfo: Record<string, { name: string; description: string }> = {
   },
 }
 
-const Tooltip = ({ agent, text }: { agent: Agent, text: string }) => {
-  const [showTooltip, setShowTooltip] = useState(false);
-  const [tooltipPosition, setTooltipPosition] = useState({ x: 0, y: 0 });
-  const [mounted, setMounted] = useState(false);
-
-  useEffect(() => {
-    setMounted(true);
-    return () => setMounted(false);
-  }, []);
-
-  const handleMouseEnter = (e: React.MouseEvent) => {
-    const rect = e.currentTarget.getBoundingClientRect();
-    setTooltipPosition({
-      x: rect.left + rect.width / 2,
-      y: rect.top
-    });
-    setShowTooltip(true);
-  };
-
-  const handleMouseLeave = () => {
-    setShowTooltip(false);
-  };
-
-  return (
-    <>
-      <span 
-        style={{ 
-          cursor: 'help',
-          position: 'relative',
-          display: 'inline-block'
-        }} 
-        onMouseEnter={handleMouseEnter}
-        onMouseLeave={handleMouseLeave}
-      >
-        <svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round">
-          <circle cx="12" cy="12" r="10"></circle>
-          <path d="M12 16v-4"></path>
-          <path d="M12 8h.01"></path>
-        </svg>
-      </span>
-      {mounted && showTooltip && createPortal(
-        <div
-          style={{
-            position: 'fixed',
-            left: tooltipPosition.x,
-            top: tooltipPosition.y,
-            transform: `translate(${agent === Agent.LLM ? '-30%' : '-80%'}, -100%)`,
-            padding: '8px',
-            background: 'rgba(0, 0, 0, 0.8)',
-            color: 'white',
-            borderRadius: '4px',
-            fontSize: '12px',
-            zIndex: 1000,
-            pointerEvents: 'none',
-            width: '300px',
-            marginTop: '-8px'
-          }}
-        >
-          {text}
-        </div>,
-        document.body
-      )}
-    </>
-  );
-};
-
 const AiInfo: React.FC<AiInfoProps> = ({ selectedSatellite }) => {
 
   if (!selectedSatellite) return null;
 
-  const [agent, setAgent] = useState<keyof typeof agentInfo>('llm');
-  const [isHovered, setIsHovered] = useState(false);
+  const [agent, setAgent] = useState<keyof typeof agentInfo | null>(null);
   const [isLoading, setIsLoading] = useState<boolean>(false);
   const [llmSatelliteInfo, setLlmSatelliteInfo] = useState<string>('');
   const [ragSatelliteInfo, setRagSatelliteInfo] = useState<string>('');
-
-  const satelliteInfo = agent === 'llm' ? llmSatelliteInfo : ragSatelliteInfo;
-  const setSatelliteInfo = agent === 'llm' ? setLlmSatelliteInfo : setRagSatelliteInfo;
+  const satelliteInfo = agent === null ? '' : (agent === Agent.LLM ? llmSatelliteInfo : ragSatelliteInfo);
 
   useEffect(() => {
     setLlmSatelliteInfo('');
     setRagSatelliteInfo('');
   }, [selectedSatellite]);
 
-  const fetchSatelliteInfoInChuncks = async (): Promise<void> => {
+  const onAgentSelect = async (agent: keyof typeof agentInfo) => {
+    
+    // Check if we already have info for this agent
+    const currentInfo = agent === Agent.LLM ? llmSatelliteInfo : ragSatelliteInfo;
+    if (currentInfo) {
+      setAgent(agent);
+      return;
+    }
+
+    setAgent(agent);
+    const setSatelliteInfo = agent === Agent.LLM ? setLlmSatelliteInfo : setRagSatelliteInfo;
+
     setIsLoading(true);
     if (!selectedSatellite) return;
     try {
@@ -129,43 +71,17 @@ const AiInfo: React.FC<AiInfoProps> = ({ selectedSatellite }) => {
   return (
     <>
       <div className="toggle-buttons">
-        {Object.entries(agentInfo).map(([key, value]) => (
+        {Object.keys(agentInfo).map((key: keyof typeof agentInfo) => (
           <button
             key={key}
             className={`toggle-button w-full flex items-center justify-center gap-2 ${agent === key ? 'active' : ''}`}
-            onClick={() => setAgent(key)}
+            onClick={() => onAgentSelect(key)}
           >
             <span className="pt-1">{key.toUpperCase()}</span>
             <Tooltip agent={key as Agent} text={agentInfo[key].description} />
           </button>
         ))}
       </div>
-      {!satelliteInfo && !isLoading && (
-        <p className="mt-4">
-          <button 
-            onClick={() => fetchSatelliteInfoInChuncks()}
-            onMouseEnter={() => setIsHovered(true)}
-            onMouseLeave={() => setIsHovered(false)}
-            style={{
-              width: '100%',
-              marginBottom: '16px',
-              padding: '8px 16px',
-              background: isHovered ? 'rgba(59, 130, 246, 0.9)' : 'rgba(59, 130, 246, 0.8)',
-              border: '1px solid rgba(147, 197, 253, 0.3)',
-              borderRadius: '4px',
-              color: 'white',
-              fontSize: '14px',
-              fontFamily: 'Shentox',
-              fontWeight: 500,
-              cursor: 'pointer',
-              transition: 'all 0.2s ease-in-out',
-              transform: isHovered ? 'translateY(-1px)' : 'none'
-            }}
-          >
-            Get AI-Powered Satellite Info
-          </button>
-        </p>
-      )}
 
       <div className="satellite-info">
         {isLoading ? (
@@ -241,4 +157,70 @@ const AiInfo: React.FC<AiInfoProps> = ({ selectedSatellite }) => {
   );
 };
 
+const Tooltip = ({ agent, text }: { agent: Agent, text: string }) => {
+  const [showTooltip, setShowTooltip] = useState(false);
+  const [tooltipPosition, setTooltipPosition] = useState({ x: 0, y: 0 });
+  const [mounted, setMounted] = useState(false);
+
+  useEffect(() => {
+    setMounted(true);
+    return () => setMounted(false);
+  }, []);
+
+  const handleMouseEnter = (e: React.MouseEvent) => {
+    const rect = e.currentTarget.getBoundingClientRect();
+    setTooltipPosition({
+      x: rect.left + rect.width / 2,
+      y: rect.top
+    });
+    setShowTooltip(true);
+  };
+
+  const handleMouseLeave = () => {
+    setShowTooltip(false);
+  };
+
+  return (
+    <>
+      <span 
+        style={{ 
+          cursor: 'help',
+          position: 'relative',
+          display: 'inline-block'
+        }} 
+        onMouseEnter={handleMouseEnter}
+        onMouseLeave={handleMouseLeave}
+      >
+        <svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2" strokeLinecap="round" strokeLinejoin="round">
+          <circle cx="12" cy="12" r="10"></circle>
+          <path d="M12 16v-4"></path>
+          <path d="M12 8h.01"></path>
+        </svg>
+      </span>
+      {mounted && showTooltip && createPortal(
+        <div
+          style={{
+            position: 'fixed',
+            left: tooltipPosition.x,
+            top: tooltipPosition.y,
+            transform: `translate(${agent === Agent.LLM ? '-30%' : '-80%'}, -100%)`,
+            padding: '8px',
+            background: 'rgba(0, 0, 0, 0.8)',
+            color: 'white',
+            borderRadius: '4px',
+            fontSize: '12px',
+            zIndex: 1000,
+            pointerEvents: 'none',
+            width: '300px',
+            marginTop: '-8px'
+          }}
+        >
+          {text}
+        </div>,
+        document.body
+      )}
+    </>
+  );
+};
+
 export default AiInfo;

components/CustomDropdown.tsx

@@ -130,7 +130,8 @@ const CustomDropdown: React.FC<CustomDropdownProps> = ({
         }
 
         .dropdown-item.selected {
-          background: transparent;
+          background: rgba(128, 128, 128, 0.3);
+          cursor: default;
         }
 
         /* Scrollbar styling */