Update agent.py

agent.py

@@ -1,258 +1,258 @@
-########## Imports ############
-from langchain_google_genai import ChatGoogleGenerativeAI
-import os
-from typing import TypedDict, List, Dict, Any, Optional
-from langgraph.graph import StateGraph, START, END
-from langchain_openai import ChatOpenAI
-from langchain_core.messages import HumanMessage
-
-from langchain_community.tools import WikipediaQueryRun
-from langchain_community.utilities import WikipediaAPIWrapper
-import string
-
-from langchain_experimental.tools import PythonREPLTool
-import ast, json
-
-from langchain_community.tools import DuckDuckGoSearchRun
-
-
-########## State ############
-class InfoState(TypedDict):
-    question: str
-    answer_type: Optional[str] # WebInfo  - WIKI - MATH
-    answer_code : Optional[str]
-    main_parts: Optional[List[str]]
-    tool_answer : Optional[list[str]]
-    final_answer : Optional[str]
-
-
-
-######### Nodes ############
-def get_wiki_relate(state: InfoState) -> InfoState:
-    """
-    Tool to Get the wikipedia info from keywords extracted from preprocessing at main_parts.
-    
-    Uses: Wikipedia API
-    Returns: tool_answer (summary)
-    """
-    print("Using Wikipedia...")
-    # Create the Wikipedia utility
-    wiki = WikipediaAPIWrapper(
-        lang="en",        # Wikipedia language
-        top_k_results=1,  # how many results to fetch
-        doc_content_chars_max=2000
-    )
-
-    # Make a tool from it
-    wiki_tool = WikipediaQueryRun(api_wrapper=wiki)
-
-    try:
-      wiki_answer = wiki_tool.run(" ".join(state["main_parts"]) + " full wikipedia article about this topic")
-      state['tool_answer'] = wiki_answer
-      return state
-    except Exception as e:
-      print("Rate limit Exception")
-      state['tool_answer'] = ""
-      return state
-
-
-def execute_code(state: InfoState) -> InfoState :
-    """Tool to calculate any math using python code or get current date time."""
-    print("Execut Code...")
-    python_tool = PythonREPLTool()
-    code = state["answer_code"]
-    state["tool_answer"]=python_tool.run(code)
-    return state
-
-def get_code(state:InfoState) -> InfoState:
-    """From prompt get the code to run."""
-    print("Getting Code (Gemini)...")
-    prompt = (
-    f"You are a strict code generator. "
-    f"Given the question: '{state['question']}', "
-    f"return ONLY valid Python code that computes the answer IF the question is about math, date, or time. "
-    f"Otherwise, return exactly: print('not valid')\n\n"
-    f"Rules:\n"
-    f"- Output ONLY the code or print('not valid')\n"
-    f"- No explanations, no markdown, no extra text\n"
-    f"- No quotes around the code\n"
-    f"- Use print() to show the result\n"
-    f"- Import modules only if needed (e.g. datetime, math)"
-)
-
-    # 2️⃣ Call Gemini
-    model = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
-    response = model.invoke([HumanMessage(content=prompt)]).content.strip()
-
-    state["answer_code"] = response
-    
-
-    return state
-
-
-
-def preprocess_text(state: dict) -> InfoState: 
-    
-    """
-    Preprocess text to get the keywords to help get results directly from wikipedia.
-    
-    Input: raw question
-    Output: main_parts (list of keywords)
-    """
-    print("Preprocess text (Gemini)...")
-    # 1️⃣ Prepare the prompt
-    prompt =  (
-    "We want to find the best-matching English Wikipedia pages for a factual question, "
-    "so we must extract only the essential topic names or entities that Wikipedia likely has pages for. "
-    "These should include the main subject (e.g., a person, event, place, or concept) and any directly relevant subtopic "
-    "(like 'Discography', 'Career', or 'History') if they help narrow the search.\n\n"
-
-    "Rules:\n"
-    "- Output 1 to 3 items maximum.\n"
-    "- Use English Wikipedia title format (capitalize each main word).\n"
-    "- Translate non-English names or terms to English.\n"
-    "- Exclude question words, pronouns, and filler terms.\n"
-    "- Fix spelling errors if necessary.\n"
-    "- Prefer specific Wikipedia topics over vague ones.\n\n"
-
-    "Example:\n"
-    "Q: 'Who built the Eiffel Tower?'\n"
-    "A: [\"Eiffel Tower\", \"Gustave Eiffel\"]\n\n"
-    f"Question: '{state['question']}'\n\n"
-    "Output ONLY a valid JSON list as described — no explanations, markdown, or extra formatting."
-    )
-
-
-    # 2️⃣ Call Gemini
-    model = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
-    response = model.invoke([HumanMessage(content=prompt)]).content.strip()
-
-    # 3️⃣ Try to safely parse
-    try:
-        # First, try JSON
-        state["main_parts"] = json.loads(response)
-    except json.JSONDecodeError:
-        try:
-            # If not JSON, try Python literal
-            state["main_parts"] = ast.literal_eval(response)
-        except Exception:
-            # If both fail, store fallback info
-            print("⚠️ Model returned invalid content:", response)
-            state["main_parts"] = []
-
-    return state
-
-
-
-
-def get_answer(state: InfoState) -> InfoState :
-    """
-    Final Node that returns the final answer organized.
-    
-    Combines: tool_answer → final_answer
-    """
-    print("Getting Answer (Gemini)...")
-
-    prompt = (
-        "You are a knowledgeable assistant that answers questions based on context and common factual knowledge.\n"
-        "Use the context first, but if it clearly lacks the needed details, you may rely on well-known public facts "
-        "(such as from Wikipedia) that logically complete the answer.\n\n"
-        f"Question: {state['question']}\n"
-        f"Context:\n{state.get('tool_answer')}\n\n"
-        "Instructions:\n"
-        "- Focus on producing one short factual answer.\n"
-        "- Do not include tool names, prefixes, or metadata.\n"
-        "- If the context contains partial hints, you can infer the answer from general knowledge of the same topic.\n"
-        "- If absolutely nothing is relevant, reply: I don't know.\n\n"
-        "Final Answer:"
-    )
-
-    model = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
-    state["final_answer"] = (model.invoke([HumanMessage(content=prompt)]).content)
-
-    return state
-
-def get_type(state: InfoState) -> InfoState:
-    """Choose which tool to use based on question type (WIKI, SEARCH, CODE)."""
-    print("Getting Type (Gemini)...")
-
-    prompt = "According to the Question " +state["question"] + " Select the best tool to answer WIKI if it's informatative or science question, WebInfo if it need up to data news,MATH if math or date or time,LLM if other or have link  You're very serious,just give one word from given"
-    model = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
-    state["answer_type"] = (model.invoke([HumanMessage(content=prompt)]).content)
-
-    return state
-
-
-
-
-def get_search_results(state: InfoState) -> InfoState:
-    """Tool to search web for results using DuckDuckGo."""
-    print("Searching...")
-
-    search = DuckDuckGoSearchRun()
-
-    try:
-      state['tool_answer'] = search.run(state["question"])   #" " .join(state["main_parts"]))
-      return state
-    except Exception:
-      state['tool_answer'] = ""
-      return state
-
-
-def route(state: InfoState):
-    print(state["answer_type"])
-    return state["answer_type"]
-
-
-
-################# Graph ################
-def get_graph():
-        graph = StateGraph(InfoState)
-
-        # Add nodes
-        #graph.add_node("get_wiki_relate", get_wiki_relate)
-        graph.add_node("preprocess_text", preprocess_text)
-        graph.add_node("get_answer", get_answer)
-        graph.add_node("get_type", get_type)
-        graph.add_node("get_search_results", get_search_results)
-        graph.add_node("execute_code", execute_code)
-        graph.add_node("get_code", get_code)
-
-        # Add edges
-        graph.add_edge(START, "preprocess_text")
-        graph.add_edge("preprocess_text", "get_type")
-
-
-        # Add conditional edges
-        graph.add_conditional_edges(
-            "get_type",
-            route,
-            {
-                "WebInfo": "get_search_results",
-                "WIKI": "get_search_results",#"get_wiki_relate",
-                "MATH": "get_code",
-                "LLM": "get_answer"
-            }
-        )
-
-        # Add final edges
-        graph.add_edge("get_search_results", "get_answer")  
-        #graph.add_edge("get_wiki_relate", "get_answer")
-
-        graph.add_edge("get_code", "execute_code")
-        graph.add_edge("execute_code", "get_answer")
-
-        graph.add_edge("get_answer", END)
-
-
-        # Compile the graph
-        compiled_graph = graph.compile()
-        return compiled_graph
-
-def ask(compiled_graph,question):
-        legitimate_result = compiled_graph.invoke({
-            "question": question,
-
-        })
-
-        return legitimate_result['final_answer'],legitimate_result
+########## Imports ############
+from langchain_google_genai import ChatGoogleGenerativeAI
+import os
+from typing import TypedDict, List, Dict, Any, Optional
+from langgraph.graph import StateGraph, START, END
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage
+
+from langchain_community.tools import WikipediaQueryRun
+from langchain_community.utilities import WikipediaAPIWrapper
+import string
+
+from langchain_experimental.tools import PythonREPLTool
+import ast, json
+
+from langchain_community.tools import DuckDuckGoSearchRun
+
+
+########## State ############
+class InfoState(TypedDict):
+    question: str
+    answer_type: Optional[str] # WebInfo  - WIKI - MATH
+    answer_code : Optional[str]
+    main_parts: Optional[List[str]]
+    tool_answer : Optional[list[str]]
+    final_answer : Optional[str]
+
+
+
+######### Nodes ############
+def get_wiki_relate(state: InfoState) -> InfoState:
+    """
+    Tool to Get the wikipedia info from keywords extracted from preprocessing at main_parts.
+    
+    Uses: Wikipedia API
+    Returns: tool_answer (summary)
+    """
+    print("Using Wikipedia...")
+    # Create the Wikipedia utility
+    wiki = WikipediaAPIWrapper(
+        lang="en",        # Wikipedia language
+        top_k_results=1,  # how many results to fetch
+        doc_content_chars_max=2000
+    )
+
+    # Make a tool from it
+    wiki_tool = WikipediaQueryRun(api_wrapper=wiki)
+
+    try:
+      wiki_answer = wiki_tool.run(" ".join(state["main_parts"]) + " full wikipedia article about this topic")
+      state['tool_answer'] = wiki_answer
+      return state
+    except Exception as e:
+      print("Rate limit Exception")
+      state['tool_answer'] = ""
+      return state
+
+
+def execute_code(state: InfoState) -> InfoState :
+    """Tool to calculate any math using python code or get current date time."""
+    print("Execut Code...")
+    python_tool = PythonREPLTool()
+    code = state["answer_code"]
+    state["tool_answer"]=python_tool.run(code)
+    return state
+
+def get_code(state:InfoState) -> InfoState:
+    """From prompt get the code to run."""
+    print("Getting Code (Gemini)...")
+    prompt = (
+    f"You are a strict code generator. "
+    f"Given the question: '{state['question']}', "
+    f"return ONLY valid Python code that computes the answer IF the question is about math, date, or time. "
+    f"Otherwise, return exactly: print('not valid')\n\n"
+    f"Rules:\n"
+    f"- Output ONLY the code or print('not valid')\n"
+    f"- No explanations, no markdown, no extra text\n"
+    f"- No quotes around the code\n"
+    f"- Use print() to show the result\n"
+    f"- Import modules only if needed (e.g. datetime, math)"
+)
+
+    # 2️⃣ Call Gemini
+    model = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
+    response = model.invoke([HumanMessage(content=prompt)]).content.strip()
+
+    state["answer_code"] = response
+    
+
+    return state
+
+
+
+def preprocess_text(state: dict) -> InfoState: 
+    
+    """
+    Preprocess text to get the keywords to help get results directly from wikipedia.
+    
+    Input: raw question
+    Output: main_parts (list of keywords)
+    """
+    print("Preprocess text (Gemini)...")
+    # 1️⃣ Prepare the prompt
+    prompt =  (
+    "We want to find the best-matching English Wikipedia pages for a factual question, "
+    "so we must extract only the essential topic names or entities that Wikipedia likely has pages for. "
+    "These should include the main subject (e.g., a person, event, place, or concept) and any directly relevant subtopic "
+    "(like 'Discography', 'Career', or 'History') if they help narrow the search.\n\n"
+
+    "Rules:\n"
+    "- Output 1 to 3 items maximum.\n"
+    "- Use English Wikipedia title format (capitalize each main word).\n"
+    "- Translate non-English names or terms to English.\n"
+    "- Exclude question words, pronouns, and filler terms.\n"
+    "- Fix spelling errors if necessary.\n"
+    "- Prefer specific Wikipedia topics over vague ones.\n\n"
+
+    "Example:\n"
+    "Q: 'Who built the Eiffel Tower?'\n"
+    "A: [\"Eiffel Tower\", \"Gustave Eiffel\"]\n\n"
+    f"Question: '{state['question']}'\n\n"
+    "Output ONLY a valid JSON list as described — no explanations, markdown, or extra formatting."
+    )
+
+
+    # 2️⃣ Call Gemini
+    model = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
+    response = model.invoke([HumanMessage(content=prompt)]).content.strip()
+
+    # 3️⃣ Try to safely parse
+    try:
+        # First, try JSON
+        state["main_parts"] = json.loads(response)
+    except json.JSONDecodeError:
+        try:
+            # If not JSON, try Python literal
+            state["main_parts"] = ast.literal_eval(response)
+        except Exception:
+            # If both fail, store fallback info
+            print("⚠️ Model returned invalid content:", response)
+            state["main_parts"] = []
+
+    return state
+
+
+
+
+def get_answer(state: InfoState) -> InfoState :
+    """
+    Final Node that returns the final answer organized.
+    
+    Combines: tool_answer → final_answer
+    """
+    print("Getting Answer (Gemini)...")
+
+    prompt = (
+        "You are a knowledgeable assistant that answers questions based on context and common factual knowledge.\n"
+        "Use the context first, but if it clearly lacks the needed details, you may rely on well-known public facts "
+        "(such as from Wikipedia) that logically complete the answer.\n\n"
+        f"Question: {state['question']}\n"
+        f"Context:\n{state.get('tool_answer')}\n\n"
+        "Instructions:\n"
+        "- Focus on producing one short factual answer.\n"
+        "- Do not include tool names, prefixes, or metadata.\n"
+        "- If the context contains partial hints, you can infer the answer from general knowledge of the same topic.\n"
+        "- If absolutely nothing is relevant, reply: I don't know.\n\n"
+        "Final Answer:"
+    )
+
+    model = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
+    state["final_answer"] = (model.invoke([HumanMessage(content=prompt)]).content)
+
+    return state
+
+def get_type(state: InfoState) -> InfoState:
+    """Choose which tool to use based on question type (WIKI, SEARCH, CODE)."""
+    print("Getting Type (Gemini)...")
+
+    prompt = "According to the Question " +state["question"] + " Select the best tool to answer WIKI if it's informatative or science question, WebInfo if it need up to data news,MATH if math or date or time,LLM if other or have link  You're very serious,just give one word from given"
+    model = ChatGoogleGenerativeAI(model="gemini-2.5-flash")
+    state["answer_type"] = (model.invoke([HumanMessage(content=prompt)]).content)
+
+    return state
+
+
+
+
+def get_search_results(state: InfoState) -> InfoState:
+    """Tool to search web for results using DuckDuckGo."""
+    print("Searching...")
+
+    search = DuckDuckGoSearchRun()
+
+    try:
+      state['tool_answer'] = search.run(state["question"])   #" " .join(state["main_parts"]))
+      return state
+    except Exception:
+      state['tool_answer'] = ""
+      return state
+
+
+def route(state: InfoState):
+    print(state["answer_type"])
+    return state["answer_type"]
+
+
+
+################# Graph ################
+def get_graph():
+        graph = StateGraph(InfoState)
+
+        # Add nodes
+        #graph.add_node("get_wiki_relate", get_wiki_relate)
+        graph.add_node("preprocess_text", preprocess_text)
+        graph.add_node("get_answer", get_answer)
+        graph.add_node("get_type", get_type)
+        graph.add_node("get_search_results", get_search_results)
+        graph.add_node("execute_code", execute_code)
+        graph.add_node("get_code", get_code)
+
+        # Add edges
+        graph.add_edge(START, "preprocess_text")
+        graph.add_edge("preprocess_text", "get_type")
+
+
+        # Add conditional edges
+        graph.add_conditional_edges(
+            "get_type",
+            route,
+            {
+                "WebInfo": "get_search_results",
+                "WIKI": "get_search_results",#"get_wiki_relate",
+                "MATH": "get_code",
+                "LLM": "get_answer"
+            }
+        )
+
+        # Add final edges
+        graph.add_edge("get_search_results", "get_answer")  
+        #graph.add_edge("get_wiki_relate", "get_answer")
+
+        graph.add_edge("get_code", "execute_code")
+        graph.add_edge("execute_code", "get_answer")
+
+        graph.add_edge("get_answer", END)
+
+
+        # Compile the graph
+        compiled_graph = graph.compile()
+        return compiled_graph
+
+def ask(compiled_graph,question):
+        legitimate_result = compiled_graph.invoke({
+            "question": question,
+
+        })
+
+        return legitimate_result['final_answer']