Update app.py

app.py

@@ -1,10 +1,8 @@
-from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
-
-from io import BytesIOimport
-import base64
-
-import requests
+import os
+from dotenv import load_dotenv
+from langchain_google_genai import ChatGoogleGenerativeAI
 
+import operator
 from langgraph.graph import StateGraph, START, END
 from typing import Annotated, Literal
 from typing_extensions import TypedDict
@@ -19,28 +17,20 @@ from langchain_experimental.utilities import PythonREPL
 class SmallerQuestions(BaseModel):
     questions: list[str]
 
-class Question(BaseModel):
-    question: str
-    nature: str
-    context: str
-    middle_answer: Annotated[list, operator.add]
-
 class Nature(BaseModel):
     nature: Literal["research", "computation", "interpretation"]
 
-class SearchableQuestion(BaseModel):
-    question: str
-    key_words: list[str]
-    context: str
-
-class Code(BaseModel):
-    code: str
-
 class OverallState(TypedDict):
-    first_question: str
-    smaller_questions: list[str]
-    middle_answer: Annotated[list, operator.add]
-    answer: str
+    initial_question: str = ""
+    smaller_questions: list[str] = []
+    current_question: str = ""
+    context: str = ""
+    key_words: str = ""
+    wiki_results: str = ""
+    web_results: str = ""
+    nature: Literal["research", "computation", "interpretation"] = "interpretation"
+    middle_answers: Annotated[list, operator.add] = []
+    answer: str = ""
 
 #Nodes
 
@@ -48,83 +38,87 @@ def BreakQuestion(state: OverallState):
     prompt = """Task: 
                 Break the following question down into a maximum of 7 smaller questions. 
                 Your response must be a list in which each item corresponds to one smaller question.
-                Beware that the answer to each of the questions you list should lead to the next question without holdbacks.
+                Beware that the answer to each one of the questions you list should lead to the next question without holdbacks.
                 
                 Question:
-                """ + state["first_question"]
+                """ + state["initial_question"]
     response = llm.with_structured_output(SmallerQuestions).invoke(prompt)
-    return {smaller_questions: response.questions}
+    return {"smaller_questions": response.questions,
+           "current_question": response.questions[0]}
 
-def DecideNature(state: Question):
+def DecideNature(state: OverallState):
     prompt = """Task:
-                The following question is answearable via one of these three approaches: "computation", "reasearch" or "interpretation".
-                Computation is a question that requires a python interpreter; research is a question that requires a web search in sites like wikipedia; interpretation is a question that needs only reasoning to be answered.
+                The following question is answerable via one of the following three approaches: "computation", "research" or "interpretation".
+                Computation is a question that requires a Python interpreter; research is a question that requires a web search in sites like Wikipedia; interpretation is a question that needs only reasoning to be answered.
                 Decide which of those is the best approach and return only the string corresponding to the one you choose, and nothing else.
                 Question: 
-                """ + state["question"]
+                """ + state["current_question"]
     response = llm.with_structured_output(Nature).invoke(prompt)
-    return {nature: response.nature}
+    return {"nature": response.nature}
 
-def KeyWords(state: SearchableQuestion):
+def KeyWords(state: OverallState):
     prompt = """Task:
-                Gather the key words from the following question to perform a search.
+                Gather the keywords from the following question to perform a search. Your response must return only keywords and nothing else.
                 Question:
-                """ + state["question"]
-    response = llm.with_structured_output(SearchableQuestion).invoke(prompt)
-    return {key_words: response.key_words}
+                """ + state["current_question"]
+    response = llm.invoke(prompt)
+    return {"key_words": response.content}
     
-def WikiSearch(state: SearchableQuestion):
+def WikiSearch(state: OverallState):
     wikipedia_api = WikipediaAPIWrapper(top_k_results=1,doc_content_chars_max=1000)
     wikipedia = WikipediaQueryRun(api_wrapper=wikipedia_api)
-    result = wikipedia.run({query: state["key_words"]})
-    return {context: result}
+    result = wikipedia.run({"query": state["key_words"]})
+    return {"wiki_results": result}
 
-def WebSearch(state: SearchableQuestion):
+def WebSearch(state: OverallState):
     search_tool = DuckDuckGoSearchRun()
     result = search_tool.invoke(state["key_words"])
-    return {context: result}
+    return {"web_results": result}
 
-def ContextAnswer(state: SearchableQuestion):
-    prompt = """Task:
-                I will give you a context and a question, and I want you to answer the question based on the given context and nothing else.
-                Context:
-                """ + state["context"] + "Question: " + state["question"]
-    response = llm.invoke(prompt)
-    return {middle_answer: [response]}
+def ContextAnswer(state: OverallState):
+    context = "Wiki results: " + state.get("wiki_results", "No wiki results found") + "\n\n Web results: " + state.get("web_results", "No web results found")
+    return {"context": context}
 
-def Computation(state: Question):
+def Computation(state: OverallState):
     python_repl = PythonREPL()
     prompt = """Task:
                 The following question requires a computation to be answered. Create a code that solves the problem.
                 Your response must be only the code that solves the problem and nothing else.
                 Question:
-                """ + state["question"]
-    response = llm.with_structured_output(Code).invoke(prompt)
-    code_result = python_repl.run(response.code)
-    structured_answer = "The result of the following problem is " + code_result + "Problem: " + state["question"]
-    return {middle_answer: [structured_answer]}
+                """ + state["current_question"]
+    response = llm.invoke(prompt)
+    code_result = python_repl.run(response.content)
+    structured_answer = "The result of the following problem is " + code_result + "Problem: " + state["current_question"]
+    return {"middle_answers": [structured_answer]}
 
-def Reasoning(state: Question):
+def Reasoning(state: OverallState):
     prompt = """Task:
-                Answer the following question using your reasoning. Your response must be contain only the answer to the question and nothing else.
+                Answer the following question using your reasoning. Your response must contain only the answer to the question and nothing else.
                 Question:
-                """ + state["question"]
+                """ + state["current_question"]
     response = llm.invoke(prompt)
-    return {middle_answer: [response]}
+    return {"middle_answers": [response.content]}
+
+def AdvanceToNextQuestion(state: OverallState):
+    current_index = state["smaller_questions"].index(state["current_question"])
+    if current_index < len(state["smaller_questions"]) - 1:
+        return {"current_question": state["smaller_questions"][current_index + 1],
+               "context": ""}
+    return {}
 
 def FinalAnswer(state: OverallState):
-    middle_answers = "\n\n".join(state["middle_answer"])
+    middle_answers = "\n\n".join(state["middle_answers"])
     prompt = """Task:
                 You will receive a question and some information from which you must be able to solve the question completely.
                 Do your best to keep your response as close to the given information as possible. Your response should contain only the final answer to the question and nothing else.
                 Question:
-                """ + state["question"] + "Context: " + middle_answers
+                """ + state["initial_question"] + "Context: " + middle_answers
     response = llm.invoke(prompt)
-    return {answer: response}
+    return {"answer": response.content}
 
 #Edges
 
-def Nature(state: Question):
+def Nature(state: OverallState):
     if state["nature"] == "computation":
         return Send("computation", state)
     elif state["nature"] == "research":
@@ -132,9 +126,27 @@ def Nature(state: Question):
     else:
         return Send("interpretation", state)
 
+def NextQuestion(state: OverallState):
+    if state["current_question"] == state["smaller_questions"][-1]:
+        return "final_answer"
+    else:
+        return "define_nature"
+
+        
+load_dotenv()
+
+llm = ChatGoogleGenerativeAI(
+    model="gemini-1.5-pro",  # or "gemini-1.5-flash" for faster/cheaper
+    temperature=0,
+    max_tokens=1024,
+    timeout=30,
+    max_retries=2,
+    google_api_key=os.getenv("GOOGLE_API_KEY")
+)
 
 builder = StateGraph(OverallState)
 builder.add_node("break_question", BreakQuestion)
+builder.add_node("advance", AdvanceToNextQuestion)
 builder.add_node("define_nature", DecideNature)
 builder.add_node("key_words", KeyWords)
 builder.add_node("wiki_search", WikiSearch)
@@ -146,11 +158,13 @@ builder.add_node("final_answer", FinalAnswer)
 builder.add_edge(START, "break_question")
 builder.add_conditional_edges("define_nature", Nature)
 builder.add_edge("key_words", "wiki_search")
-builder.add_edge("wiki_search", "web_search")
+builder.add_edge("key_words", "web_search")
 builder.add_edge("web_search", "answer")
-builder.add_edge("answer", "final_answer")
-builder.add_edge("computation", "final_answer")
-builder.add_edge("interpretation", "final_answer")
+builder.add_edge("wiki_search", "answer")
+builder.add_edge("answer", "advance")
+builder.add_edge("computation", "advance")
+builder.add_edge("interpretation", "advance")
+builder.add_conditional_edges("advance", NextQuestion)
 builder.add_edge("final_answer", END)
 
 graph = builder.compile()