Update agent.py

agent.py

@@ -110,90 +110,68 @@ def build_graph():
     Build and return a StateGraph using a Hugging Face chat LLM with tools.
     """
     try:
-        # For HuggingFace Spaces, use the Inference API approach
         hf_token = os.getenv("HUGGINGFACE_TOKEN") or os.getenv("HF_TOKEN") or os.getenv("HF_API_TOKEN")
         
         if hf_token:
             print("Using HuggingFace Inference API...")
             from langchain_huggingface import HuggingFaceEndpoint
             
-            # Use HF Inference API instead of local models
             llm = HuggingFaceEndpoint(
                 repo_id="microsoft/DialoGPT-medium",
                 huggingfacehub_api_token=hf_token,
                 model_kwargs={"temperature": 0.1, "max_new_tokens": 512}
             )
             
-            # Wrap for chat interface
             llm = ChatHuggingFace(llm=llm)
             print("✓ Successfully initialized HF Inference API")
             
         else:
-            print("No HF token found, creating mock LLM for demo...")
-            # Create a simple mock LLM for demonstration
+            print("No HF token found, creating mock LLM for demo…")
             class SimpleMockLLM:
                 def bind_tools(self, tools):
                     return self
                 
                 def invoke(self, messages):
                     from langchain_core.messages import AIMessage
-                    
-                    # Extract the last message
-                    if hasattr(messages, '__iter__') and messages:
-                        last_msg = messages[-1]
-                        if hasattr(last_msg, 'content'):
-                            content = last_msg.content.lower()
-                        else:
-                            content = str(last_msg).lower()
-                    else:
-                        content = str(messages).lower()
-                    
-                    # Simple rule-based responses
+                    last_msg = messages[-1] if messages else None
+                    content = getattr(last_msg, 'content', str(last_msg)).lower() if last_msg else ""
                     if any(word in content for word in ['math', 'calculate', 'add', 'multiply']):
                         return AIMessage(content="I can help with math! Try asking me to add, multiply, subtract, or divide numbers.")
                     elif any(word in content for word in ['search', 'find', 'look up']):
                         return AIMessage(content="I can search Wikipedia, Arxiv, or the web for information. What would you like me to search for?")
                     else:
-                        return AIMessage(content=f"Hello! I'm a demo assistant. I can help with math calculations and searches. You said: {content[:100]}...")
+                        return AIMessage(content=f"Hello! I'm a demo assistant. You said: {content[:100]}...")
             
             llm = SimpleMockLLM()
             print("✓ Created demo LLM")
         
     except Exception as e:
         print(f"Error initializing LLM: {e}")
-        # Fallback to basic mock
         class BasicMockLLM:
             def bind_tools(self, tools):
                 return self
             
             def invoke(self, messages):
                 from langchain_core.messages import AIMessage
-                return AIMessage(content="Demo mode: I'm a basic assistant. Please configure HuggingFace token for full functionality.")
+                return AIMessage(content="Demo mode: Please configure a token for full functionality.")
         
         llm = BasicMockLLM()
         print("✓ Using basic fallback LLM")
     
     llm_with_tools = llm.bind_tools(tools)
 
-    # retriever node: use vector store if available
     def retriever(state: MessagesState):
         if supabase:
             query = state["messages"][-1].content
             doc = vector_store.similarity_search(query, k=1)[0]
             content = doc.page_content
-            if "Final answer :" in content:
-                answer = content.split("Final answer :")[-1].strip()
-            else:
-                answer = content.strip()
+            answer = content.split("Final answer :")[-1].strip() if "Final answer :" in content else content.strip()
             return {"messages": [AIMessage(content=answer)]}
-        # no supabase: pass through
         return {"messages": state["messages"]}
 
-    # assistant node: invoke LLM and tools
     def assistant(state: MessagesState):
         return {"messages": [llm_with_tools.invoke(state["messages"])]}
 
-    # assemble graph
     g = StateGraph(MessagesState)
     g.add_node("retriever", retriever)
     g.add_node("assistant", assistant)
@@ -202,7 +180,6 @@ def build_graph():
     g.add_conditional_edges("assistant", tools_condition)
     g.add_node("tools", ToolNode(tools))
     g.add_edge("tools", "assistant")
-
     g.set_entry_point("retriever")
     g.set_finish_point("assistant")
-    return g.compile()
+    return g.compile()