updated code base

agent.py

@@ -1,19 +1,46 @@
 import os
+import shutil
 from langchain_groq import ChatGroq
 from langchain.prompts import PromptTemplate
 from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.prebuilt import ToolNode, tools_condition
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader
-from langchain_core.messages import HumanMessage
+from langchain_core.messages import  SystemMessage, HumanMessage
 from langchain.tools import tool
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.runnables import Runnable
 from dotenv import load_dotenv
+from langchain.vectorstores import Chroma
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.tools.retriever import create_retriever_tool
+from typing import TypedDict, Annotated, List
 
 # Load environment variables from .env
 load_dotenv()
 
+# Custom Agent Prompt Template
+Agent_prompt_template = '''You are a helpful assistant following the REACT methodology and tasked with answering questions using a set of tools. 
+Once a question is asked,you have to Report your thoughts, and finish your answer with the following template: 
+FINAL ANSWER: [YOUR FINAL ANSWER]. 
+
+### **Instructions:**  
+- YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. 
+- If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. 
+- If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
+- Provide the answer in clear and professional language.  
+- **Limit** the FINAL ANSWER to 1000 words and keep it to the point.
+- **Do not** include any additional information or explanations in your FINAL ANSWER.
+- **Do not** include any information that is not relevant to the question.
+- **Validate** your answer before providing it.  
+- **Print** all the steps you take to arrive at your answer.
+
+Your answer should only start with "FINAL ANSWER: ", then follows with the answer. '''
+
+sys_msg = SystemMessage(content=Agent_prompt_template)
+
+
 # Initialize LLM
 def initialize_llm():
     """Initializes the ChatGroq LLM."""
@@ -247,103 +274,145 @@ def standard_deviation(numbers: list) -> float:
     variance = sum((x - mean_value) ** 2 for x in numbers) / len(numbers)
     return variance ** 0.5
 
-# Build the LangGraph
-def build_graph():
+# --- Vector Store + Retriever ---
+# State schema
+class MessagesState(TypedDict):
+    messages: Annotated[List[HumanMessage], "Messages in the conversation"]
+
+# === VECTOR STORE SETUP ===
+PERSIST_DIR = "./chroma_store"
+
+def initialize_chroma_store():
+    # Optional: clear existing store if desired
+    if os.path.exists(PERSIST_DIR):
+        shutil.rmtree(PERSIST_DIR)
+        os.makedirs(PERSIST_DIR)
+
+    # Initialize embeddings
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+
+    # Load existing or empty vector store
+    vectorstore = Chroma(
+        embedding_function=embeddings,
+        persist_directory=PERSIST_DIR
+    )
+    return vectorstore
+
+vector_store = initialize_chroma_store()
+
+# Create retriever tool
+retriever_tool = create_retriever_tool(
+    retriever=vector_store.as_retriever(),
+    name="Question Search",
+    description="A tool to retrieve similar questions from a vector store."
+)
+
+
+
+@tool
+def weather_tool(location: str) -> str:
     """
-    Builds the LangGraph with the defined tools and assistant node.
+    Fetches the weather for a location.
+
+    Args:
+        location (str): The location to fetch weather for.
 
     Returns:
-        RunnableGraph: The compiled LangGraph.
+        str: The weather information.
+    """
+    return get_weather(location, search_tool)
+
+@tool
+def web_search(query: str) -> str:
+    """Search the web for a given query and return the summary.
+    Args:
+        query (str): The search query.
     """
-    llm = initialize_llm()
-    search_tool = initialize_search_tool()
-    recommendation_chain = initialize_recommendation_chain(llm)
-
-    @tool
-    def weather_tool(location: str) -> str:
-        """
-        Fetches the weather for a location.
-
-        Args:
-            location (str): The location to fetch weather for.
-
-        Returns:
-            str: The weather information.
-        """
-        return get_weather(location, search_tool)
-    
-    @tool
-    def web_search(query: str) -> str:
-        """Search the web for a given query and return the summary.
-        Args:
-            query (str): The search query.
-        """
-        
-        search_tool = TavilySearchResults()
-        result = search_tool.run(query)
-        return result[0]['content']
-    
-    @tool
-    def wiki_search(query : str) -> str:
-        """Search Wikipedia for a given query and return the summary.
-        Args:
-            query (str): The search query.
-        """
     
-        search_docs = WikipediaLoader(query=query, load_max_docs=1).load()
-        formatted_search_docs = "\n\n----\n\n".join(
-            [
-                f'<Document Source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}">\n{doc.page_content}\n</Document>'
-                for doc in search_docs
-            ]
-        )
-        return formatted_search_docs
+    search_tool = TavilySearchResults()
+    result = search_tool.run(query)
+    return result[0]['content']
+
+@tool
+def wiki_search(query : str) -> str:
+    """Search Wikipedia for a given query and return the summary.
+    Args:
+        query (str): The search query.
+    """
 
-    @tool
-    def recommendation_tool(weather_condition: str) -> str:
-        """
-        Provides recommendations based on weather conditions.
+    search_docs = WikipediaLoader(query=query, load_max_docs=1).load()
+    formatted_search_docs = "\n\n----\n\n".join(
+        [
+            f'<Document Source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}">\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ]
+    )
+    return formatted_search_docs
+
+# @tool
+# def recommendation_tool(weather_condition: str) -> str:
+#     """
+#     Provides recommendations based on weather conditions.
+
+#     Args:
+#         weather_condition (str): The weather condition.
+
+#     Returns:
+#         str: The recommendations.
+#     """
+#     return get_recommendation(weather_condition, recommendation_chain)
+
+tools = [weather_tool, wiki_search, web_search,
+            add, subtract, multiply, divide, square, cube, power, factorial, mean, standard_deviation]
 
-        Args:
-            weather_condition (str): The weather condition.
+llm = ChatGroq(
+    temperature=0,
+    model_name="qwen-qwq-32b",  # Updated to working model
+    groq_api_key=os.getenv("GROQ_API_KEY")
+)
 
-        Returns:
-            str: The recommendations.
-        """
-        return get_recommendation(weather_condition, recommendation_chain)
+llm_with_tools = llm.bind_tools(tools)
 
-    tools = [weather_tool, recommendation_tool, wiki_search, web_search,
-             add, subtract, multiply, divide, square, cube, power, factorial, mean, standard_deviation]
+# === NODES ===
 
-    llm_with_tools = llm.bind_tools(tools)
+def retriever(state: MessagesState):
+    """Retrieve similar context and inject"""
+    query = state["messages"][0].content
+    similar_docs = vector_store.similarity_search(query)
 
-    def assistant(state: MessagesState):
-        """
-        Assistant node in the LangGraph.
+    if similar_docs:
+        ref_msg = HumanMessage(
+            content=f"Here is a similar question and answer for reference:\n\n{similar_docs[0].page_content}"
+        )
+        return {"messages": [sys_msg] + state["messages"] + [ref_msg]}
+    else:
+        return {"messages": [sys_msg] + state["messages"]}
 
-        Args:
-            state (MessagesState): The current state of the conversation.
+def assistant(state: MessagesState):
+    """Invoke LLM with tools"""
+    return {"messages": [llm_with_tools.invoke(state["messages"])]}
 
-        Returns:
-            dict: The next state of the conversation.
-        """
-        print("Entering assistant node...")
-        response = llm_with_tools.invoke(state["messages"])
-        print(f"Assistant says: {response.content}")
-        return {"messages": [response]}
+# === GRAPH BUILD ===
 
+def build_graph():
     builder = StateGraph(MessagesState)
+    builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
-    builder.add_node("tools", ToolNode(tools))
-    builder.set_entry_point("assistant")
+    builder.add_node("tools", ToolNode([retriever_tool] + tools))
+
+    builder.set_entry_point("retriever")
+    builder.add_edge("retriever", "assistant")
     builder.add_conditional_edges("assistant", tools_condition)
     builder.add_edge("tools", "assistant")
+
     return builder.compile()
 
+# === TEST ===
 if __name__ == "__main__":
+    question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
     graph = build_graph()
-    question = "How many albums were pulished by Mercedes Sosa?"
     messages = [HumanMessage(content=question)]
     result = graph.invoke({"messages": messages})
+
     for msg in result["messages"]:
-        msg.pretty_print()
+        msg.pretty_print()