Create agent.py

agent.py

@@ -0,0 +1,92 @@
+import re
+from typing import TypedDict, Annotated, List
+import operator
+
+from langchain.tools import tool
+from langchain_google_community.search import GoogleSearchAPIWrapper
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import BaseMessage, ToolMessage
+from langgraph.graph import StateGraph, END
+from langgraph.prebuilt import ToolNode
+
+# This class defines the structure of the agent's state
+class AgentState(TypedDict):
+    messages: Annotated[List[BaseMessage], operator.add]
+
+def create_agent_graph(vector_store, nvidia_api_key, google_api_key, google_cse_id):
+    """Creates and compiles the LangGraph agent."""
+
+    # 1. Define Tools
+    @tool
+    def paper_qa_tool(query: str) -> str:
+        """
+        Answers specific, detailed questions about scientific papers on graph theory,
+        sparsity, and the pebble game. Use this for questions that reference specific
+        paper details or concepts.
+        """
+        print("--- Calling Paper Q&A Tool ---")
+        retriever = vector_store.as_retriever(search_kwargs={'k': 3})
+        context_docs = retriever.get_relevant_documents(query)
+        
+        # Simple cleaning to remove potential gibberish from parsed PDFs
+        gibberish_pattern = re.compile(r'/DAN <[A-Fa-f0-9]+>')
+        cleaned_docs = [doc for doc in context_docs if not gibberish_pattern.search(doc.page_content)]
+        
+        if not cleaned_docs:
+            return "No relevant information found in the documents after cleaning."
+        
+        context_text = "\n\n".join([doc.page_content for doc in cleaned_docs])
+        return context_text
+
+    search_wrapper = GoogleSearchAPIWrapper(google_api_key=google_api_key, google_cse_id=google_cse_id)
+
+    @tool
+    def web_search_tool(query: str) -> str:
+        """
+        Provides up-to-date answers from the web for general knowledge, definitions,
+        or topics not covered in the local scientific papers. Also provides source links.
+        """
+        print("--- Calling Web Search Tool ---")
+        results = search_wrapper.results(query, num_results=3)
+        return "\n".join([f"Title: {res['title']}\nLink: {res['link']}\nSnippet: {res['snippet']}\n" for res in results])
+
+    tools = [paper_qa_tool, web_search_tool]
+    tool_node = ToolNode(tools)
+
+    # 2. Define the Model
+    # We use ChatOpenAI pointed at the NVIDIA endpoint
+    model = ChatOpenAI(
+        model="meta/llama3-70b-instruct",
+        openai_api_key=nvidia_api_key,
+        openai_api_base="https://integrate.api.nvidia.com/v1/ ",
+        temperature=0.2
+    ).bind_tools(tools)
+
+    # 3. Define Graph Nodes
+    def call_model(state):
+        """The primary node that calls the LLM."""
+        print("--- AGENT: Thinking... ---")
+        response = model.invoke(state["messages"])
+        return {"messages": [response]}
+
+    def should_continue(state):
+        """Router: decides whether to call a tool or end the conversation."""
+        last_message = state["messages"][-1]
+        if last_message.tool_calls:
+            return "continue"
+        return "end"
+
+    # 4. Build and Compile the Graph
+    workflow = StateGraph(AgentState)
+    workflow.add_node("agent", call_model)
+    workflow.add_node("action", tool_node)
+    
+    workflow.set_entry_point("agent")
+    workflow.add_conditional_edges(
+        "agent",
+        should_continue,
+        {"continue": "action", "end": END},
+    )
+    workflow.add_edge("action", "agent")
+
+    return workflow.compile()