Update agent.py

agent.py

@@ -53,8 +53,8 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from huggingface_hub import login
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
 from langchain_huggingface import HuggingFaceEndpoint
-from langchain.agents import initialize_agent
-from langchain.agents import AgentType
+#from langchain.agents import initialize_agent
+#from langchain.agents import AgentType
 from typing import Union, List
 from functools import reduce
 import operator
@@ -446,13 +446,15 @@ llm = HuggingFaceEndpoint(
     max_new_tokens=512
 )
 
+
+# No longer required as Langgraph is replacing Langchain
 # Initialize LangChain agent
-agent = initialize_agent(
-    tools=tools,
-    llm=llm,
-    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-    verbose=True
-)
+#agent = initialize_agent(
+#    tools=tools,
+#    llm=llm,
+#    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+#    verbose=True
+#)
 
 
 
@@ -963,8 +965,14 @@ model_config = {
     "huggingfacehub_api_token": os.getenv("HF_TOKEN")
 }
 
+
 def build_graph(provider, model_config):
-    # Step 1: Initialize the LLM
+    from langchain_core.messages import SystemMessage, HumanMessage
+    from langgraph.graph import StateGraph, ToolNode
+    from langchain_core.runnables import RunnableLambda
+    from some_module import vector_store  # Make sure this is defined/imported
+
+    # Step 1: Get LLM
     def get_llm(provider: str, config: dict):
         if provider == "huggingface":
             from langchain_huggingface import HuggingFaceEndpoint
@@ -979,53 +987,31 @@ def build_graph(provider, model_config):
             raise ValueError(f"Unsupported provider: {provider}")
     
     llm = get_llm(provider, model_config)
-    
-    # -------------------------------
-    # Step 6: Define LangChain Tools
-    # -------------------------------
-    calc_tool = calculator  # Math operations tool
-    web_tool = web_search    # Web search tool
-    wiki_tool = wiki_search  # Wikipedia search tool
-    arxiv_tool = arxiv_search # Arxiv search tool
-    youtube_tool = get_youtube_transcript  # YouTube transcript extraction
-    video_tool = extract_video_id  # Video ID extraction tool
-    analyze_tool = analyze_attachment  # File analysis tool
-    wikiq_tool = wikidata_query  # Wikidata query tool
-    
-    # -------------------------------
-    # Step 7: Create the Planner-Agent Logic
-    # -------------------------------
-    # Define tools list
+
+    # Step 2: Define tools
     tools = [
-        wiki_tool, 
-        calc_tool, 
-        web_tool, 
-        arxiv_tool,
-        youtube_tool, 
-        video_tool, 
-        analyze_tool, 
-        wikiq_tool
+        wiki_search, 
+        calculator, 
+        web_search, 
+        arxiv_search,
+        get_youtube_transcript, 
+        extract_video_id, 
+        analyze_attachment, 
+        wikidata_query
     ]
-    
-    # Step 8: Bind tools to the LLM
-    llm_with_tools = llm.bind_tools(tools)
-
-    # Return the LLM with tools bound
-    return llm_with_tools
-
 
+    # Step 3: Bind tools to LLM
+    llm_with_tools = llm.bind_tools(tools)
 
-    
-    # Initialize system message
+    # Step 4: Build stateful graph logic
     sys_msg = SystemMessage(content="You are a helpful assistant.")
-    
-    # Define the retriever function
-    def retriever(state: MessagesState):
+
+    def retriever(state: dict):
         user_query = state["messages"][0].content
         similar_docs = vector_store.similarity_search(user_query)
-    
+
         if not similar_docs:
-            wiki_result = wiki_tool.run(user_query)
+            wiki_result = wiki_search.run(user_query)
             return {
                 "messages": [
                     sys_msg,
@@ -1041,34 +1027,36 @@ def build_graph(provider, model_config):
                     HumanMessage(content=f"Reference:\n\n{similar_docs[0].page_content}")
                 ]
             }
-    
-    # Define the assistant function
-    def assistant(state: MessagesState):
+
+    def assistant(state: dict):
         return {"messages": [llm_with_tools.invoke(state["messages"])]}
-    
-    # Define condition for tools usage
-    def tools_condition(state: MessagesState) -> str:
+
+    def tools_condition(state: dict) -> str:
         if "use tool" in state["messages"][-1].content.lower():
             return "tools"
         else:
             return "END"
-    
-    # Initialize the StateGraph
-    builder = StateGraph(MessagesState)
-    
-    # Add nodes to the graph
+
+    # Step 5: Define LangGraph StateGraph
+    builder = StateGraph(dict)  # Using dict as state type here
+
     builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
     builder.add_node("tools", ToolNode(tools))
-    
-    # Set the entry point
+
     builder.set_entry_point("retriever")
-    
-    # Define edges
     builder.add_edge("retriever", "assistant")
     builder.add_conditional_edges("assistant", tools_condition)
     builder.add_edge("tools", "assistant")
 
-    # Compile graph
-    return builder.compile()
+    graph = builder.compile()
+    return graph
+
+
+# call build_graph AFTER it’s defined
+agent = build_graph(provider, model_config)
+
+# Now you can use the agent like this:
+result = agent.invoke({"messages": [HumanMessage(content=question)]})
+