Update agentic.py

agentic.py

@@ -46,7 +46,7 @@ class State(TypedDict):
     input_file: str
     downloaded_file: Optional[str]
     task_id: str
-    web_wiki_search_node_result: AnyMessage
+    web_search_node_result: AnyMessage
     thinking_node_result: AnyMessage
     vision_node_result: AnyMessage
     video_node_result: AnyMessage
@@ -58,20 +58,20 @@ class State(TypedDict):
 ########################
 
 ######## MODELS ########
-def get_light_model():
+def get_general_model():
 
     llm_provider = os.getenv("LLM_PROVIDER", "mistral")
 
     if llm_provider == "mistral":
-        light_model = ChatMistralAI(
-            model="mistral-small-latest",
+        general_model = ChatMistralAI(
+            model="mistral-large-2411",#"ministral-8b-latest",#"mistral-small-latest",
             temperature=0,
             max_retries=2,
             api_key=os.getenv("MISTRAL_API_KEY")
         )
 
     if llm_provider == "aws":
-        light_model = ChatBedrock(
+        general_model = ChatBedrock(
             model_id="arn:aws:bedrock:us-east-1:416545197702:inference-profile/us.amazon.nova-lite-v1:0",
             # provider="amazon",
             temperature=0,
@@ -80,23 +80,12 @@ def get_light_model():
             aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY")
         )
     
-    return light_model
-
-def get_medium_model():
-
-    medium_model = ChatMistralAI(
-        model="ministral-8b-latest",#"mistral-large-2411",
-        temperature=0,
-        max_retries=2,
-        api_key=os.getenv("MISTRAL_API_KEY")
-    )
-
-    return medium_model
+    return general_model
 
 def get_big_model():
 
     big_model = ChatMistralAI(
-        model="mistral-medium-2505",#"mistral-large-2411"
+        model="mistral-medium-2505",
         temperature=0,
         max_retries=2,
         api_key=os.getenv("MISTRAL_API_KEY")
@@ -281,15 +270,13 @@ def download_input_file(task_id: str) -> str:
 
 ######## LLM associations ########
 
-medium_model = get_medium_model()
-light_model = get_light_model()
+general_model = get_general_model()
 big_model = get_big_model()
 
 vision_model = get_vision_model()
 video_handler_model = get_video_handler_model()
 audio_handler_model = get_audio_handler_model()
 
-
 ########################
 
 ######## Nodes Definition ########
@@ -329,7 +316,8 @@ You do not handle code
 
 1. You need to fully understand the question
 2. You must think hard about what is relevant in the question to make the best answer
-3. Report your thought process in detail, explaining your reasoning step-by-step.
+3. If there are calculations or maths, you need to verify twice before answering.
+4. Report your thought process in detail, explaining your reasoning step-by-step.
 
 Here is the question {state['question']}
 Now provide your response immediately without any preamble in text but not in markdown.
@@ -339,7 +327,7 @@ Now provide your response immediately without any preamble in text but not in ma
 
     sys_msg = SystemMessage(content=prompt)
 
-    thinking_node_response = [big_model.invoke([sys_msg] + [state["thinking_node_result"]])]
+    thinking_node_response = [general_model.invoke([sys_msg] + [state["thinking_node_result"]])]
 
     thinking_node_response[-1].pretty_print()
 
@@ -391,7 +379,7 @@ Now provide your response immediately without any preamble in text but not in ma
         "code_node_result": code_node_response,
     }
 
-def web_wiki_search_node(state: State) -> dict:
+def web_search_node(state: State) -> dict:
     """
 A powerful node to answer questions and make research on the web based on the question provided in the state.
 This node does not handle files
@@ -404,7 +392,7 @@ Args:
     state (State): A dictionary containing the current state of the agent, including the 'question' key which holds the question to be answered.
 
 Returns:
-    dict: A dictionary containing the response from the web search node, with the key 'web_wiki_search_node_result' holding the list of messages generated by the general model.
+    dict: A dictionary containing the response from the web search node, with the key 'web_search_node_result' holding the list of messages generated by the general model.
     """
 
     prompt = f"""
@@ -439,16 +427,16 @@ Here is the question {state['question']}
 Now provide your response immediately without any preamble in text but not in markdown.
     """
 
-    state["web_wiki_search_node_result"] = state.get("web_wiki_search_node_result", "")
+    state["web_search_node_result"] = state.get("web_search_node_result", "")
 
     sys_msg = SystemMessage(content=prompt)
 
-    web_wiki_search_node_response = [web_search_node_agent.invoke([sys_msg] + [state["web_wiki_search_node_result"]])]
+    web_search_node_response = [web_search_node_agent.invoke([sys_msg] + [state["web_search_node_result"]])]
 
-    web_wiki_search_node_response[-1].pretty_print()
+    web_search_node_response[-1].pretty_print()
 
     return {
-        "web_wiki_search_node_result": web_wiki_search_node_response,
+        "web_search_node_result": web_search_node_response,
     }
 
 def vision_node(state: State) -> dict:
@@ -796,7 +784,7 @@ Now provide your response immediately without any preamble in text but not in ma
 
     question = [HumanMessage(content=state["question"])]
 
-    for node_result in ["web_wiki_search_node_result", "vision_node_result", "video_node_result", "audio_node_result", "thinking_node_result", "code_node_result", "excel_node_result"]:
+    for node_result in ["web_search_node_result", "vision_node_result", "video_node_result", "audio_node_result", "thinking_node_result", "code_node_result", "excel_node_result"]:
         result = state.get(node_result, "")
         if result:
             # Ensure result is a string. If it's a message object, extract its content.
@@ -808,7 +796,7 @@ Now provide your response immediately without any preamble in text but not in ma
 
     sys_msg = SystemMessage(content=prompt)
 
-    response = [big_model.invoke([sys_msg] + state["messages"]+ question + nodes_response)]
+    response = [general_model.invoke([sys_msg] + state["messages"]+ question + nodes_response)]
 
     return {
         "messages": response,
@@ -841,7 +829,7 @@ You do not handle code
 
 Here are the nodes you can choose:
 - thinking_node: {thinking_node.__doc__}
-- web_wiki_search_node: {web_wiki_search_node.__doc__}
+- web_search_node: {web_search_node.__doc__}
 - vision_node: {vision_node.__doc__}
 - video_node: {video_node.__doc__}
 - audio_node: {audio_node.__doc__}
@@ -865,7 +853,7 @@ Now provide your response immediately respecting this format in lower case: 'nex
 
     entry_node_response[-1].pretty_print()
 
-    regex_result = re.search(r'.*next\s*node.*(?P<next_node>thinking_node|web_wiki_search_node|vision_node|video_node|audio_node|code_node|excel_node).*$', entry_node_response[-1].content, re.IGNORECASE)
+    regex_result = re.search(r'.*next\s*node.*(?P<next_node>thinking_node|web_search_node|vision_node|video_node|audio_node|code_node|excel_node).*$', entry_node_response[-1].content, re.IGNORECASE)
 
     next_node = "END"
     if regex_result:
@@ -887,7 +875,7 @@ Now provide your response immediately respecting this format in lower case: 'nex
 def build_graph():
     builder = StateGraph(State)
     builder.add_node("entry_node", entry_node)
-    builder.add_node("web_wiki_search_node", web_wiki_search_node)
+    builder.add_node("web_search_node", web_search_node)
     builder.add_node("vision_node", vision_node)
     builder.add_node("video_node", video_node)
     builder.add_node("audio_node", audio_node)
@@ -903,7 +891,7 @@ def build_graph():
         "entry_node",
         lambda state: state["next_node"],
         {
-            "web_wiki_search_node": "web_wiki_search_node",
+            "web_search_node": "web_search_node",
             "vision_node": "vision_node",
             "video_node": "video_node",
             "audio_node": "audio_node",
@@ -913,7 +901,7 @@ def build_graph():
         }
     )
     # After specialized node, go to END
-    builder.add_edge("web_wiki_search_node", "format_answer_node")
+    builder.add_edge("web_search_node", "format_answer_node")
     builder.add_edge("vision_node", "format_answer_node")
     builder.add_edge("video_node", "format_answer_node")
     builder.add_edge("audio_node", "format_answer_node")
@@ -939,14 +927,15 @@ if __name__ == "__main__":
     with open("./responses.json", "r") as responses:
         json_responses = json.loads(responses.read())
 
-    # json_questions =  [{
-    #     "question": "The attached Excel file contains the sales of menu items for a local fast-food chain. What were the total sales that the chain made from food (not including drinks)? Express your answer in USD with two decimal places.",
-    #     "file_name": "7bd855d8-463d-4ed5-93ca-5fe35145f733.xlsx", 
-    #     "task_id": "7bd855d8-463d-4ed5-93ca-5fe35145f733"
-    # }]
+    json_questions =  [{
+        "task_id": "840bfca7-4f7b-481a-8794-c560c340185d",
+        "question": "On June 6, 2023, an article by Carolyn Collins Petersen was published in Universe Today. This article mentions a team that produced a paper about their observations, linked at the bottom of the article. Find this paper. Under what NASA award number was the work performed by R. G. Arendt supported by?",
+        "Level": "1",
+        "file_name": ""
+    }]
 
-    with open("questions.json", "r") as questions:
-        json_questions = json.loads(questions.read())
+    # with open("questions.json", "r") as questions:
+    #     json_questions = json.loads(questions.read())
 
     for input in json_questions: