Update agentic.py

agentic.py

@@ -11,6 +11,7 @@ from langchain_community.tools import WikipediaQueryRun
 from langchain_community.utilities import WikipediaAPIWrapper
 from langchain_aws import ChatBedrock
 from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_community.document_loaders import UnstructuredExcelLoader
 # from langchain_google_vertexai import ChatVertexAI
 
 # from langfuse.callback import CallbackHandler
@@ -51,25 +52,26 @@ class State(TypedDict):
     video_node_result: AnyMessage
     audio_node_result: AnyMessage
     code_node_result: AnyMessage
+    excel_node_result: AnyMessage
     next_node: str
 
 ########################
 
 ######## MODELS ########
-def get_general_model():
+def get_light_model():
 
     llm_provider = os.getenv("LLM_PROVIDER", "mistral")
 
     if llm_provider == "mistral":
-        general_model = ChatMistralAI(
-            model="mistral-large-2411",#"ministral-8b-latest",#"mistral-small-latest",#
+        light_model = ChatMistralAI(
+            model="mistral-small-latest",
             temperature=0,
             max_retries=2,
             api_key=os.getenv("MISTRAL_API_KEY")
         )
 
     if llm_provider == "aws":
-        general_model = ChatBedrock(
+        light_model = ChatBedrock(
             model_id="arn:aws:bedrock:us-east-1:416545197702:inference-profile/us.amazon.nova-lite-v1:0",
             # provider="amazon",
             temperature=0,
@@ -78,7 +80,29 @@ def get_general_model():
             aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY")
         )
     
-    return general_model
+    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
+
+def get_big_model():
+
+    big_model = ChatMistralAI(
+        model="mistral-medium-2505",#"mistral-large-2411"
+        temperature=0,
+        max_retries=2,
+        api_key=os.getenv("MISTRAL_API_KEY")
+    )
+
+    return big_model
 
 def get_vision_model():
 
@@ -87,7 +111,7 @@ def get_vision_model():
     if vlm_provider == "openai":
         print("Spawning Open AI VLM")
         vision_model = ChatOpenAI(
-            model="gpt-4o",#-mini",
+            model="gpt-4o",
             temperature=0,
             max_tokens=None,
             timeout=None,
@@ -257,17 +281,19 @@ def download_input_file(task_id: str) -> str:
 
 ######## LLM associations ########
 
-general_model = get_vision_model()
+medium_model = get_medium_model()
+light_model = get_light_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 ########
 
-general_model = get_general_model()
-
 search_tools = [
     web_search,
     wikipedia_search,
@@ -275,7 +301,7 @@ search_tools = [
 
 download_file_tool = [ download_input_file ]
 
-web_search_node_agent = general_model.bind_tools(search_tools, parallel_tool_calls=False)
+web_search_node_agent = big_model.bind_tools(search_tools, parallel_tool_calls=False)
 
 def thinking_node(state: State) -> dict:
     """
@@ -313,7 +339,7 @@ Now provide your response immediately without any preamble in text but not in ma
 
     sys_msg = SystemMessage(content=prompt)
 
-    thinking_node_response = [general_model.invoke([sys_msg] + [state["thinking_node_result"]])]
+    thinking_node_response = [big_model.invoke([sys_msg] + [state["thinking_node_result"]])]
 
     thinking_node_response[-1].pretty_print()
 
@@ -357,7 +383,7 @@ Now provide your response immediately without any preamble in text but not in ma
 
     sys_msg = SystemMessage(content=prompt)
 
-    code_node_response = [general_model.invoke([sys_msg])]
+    code_node_response = [big_model.invoke([sys_msg])]
 
     code_node_response[-1].pretty_print()
 
@@ -425,7 +451,6 @@ Now provide your response immediately without any preamble in text but not in ma
         "web_wiki_search_node_result": web_wiki_search_node_response,
     }
 
-
 def vision_node(state: State) -> dict:
     """
 Vision model that can analyze images and pictures and answer questions about them.
@@ -446,9 +471,9 @@ You do not handle videos
 You do not handle audio
 You do not handle code
 
-1. You need to fully understand the question
-2. You must think hard about what is relevant in the image to make the best answer to the question
-4. Report your thought process in detail, explaining your reasoning step-by-step.
+1. You need to fully understand the question.
+2. You must think hard about what is relevant in the image to make the best answer to the question.
+3. 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.
@@ -532,7 +557,7 @@ You do not handle images or pictures.
 You do not handle audio.
 You do not handle code.
 
-1. You need to fully understand the question
+1. You need to fully understand the question.
 2. Carefully observe the video, paying attention to relevant details, actions, and context.
 3. Focus on the user's question.
 4. If the question requires counting, identifying, or describing, be precise and clear in your response.
@@ -615,7 +640,7 @@ You do not handle images or pictures.
 You do not handle video.
 You do not handle code.
 
-1. You need to fully understand the question:
+1. You need to fully understand the question.
 2. Carefully listen to the audio, paying attention to relevant details, actions, and context.
 3. Focus on the user's question.
 4. If the question requires counting, identifying, or describing, be precise and clear in your response.
@@ -669,7 +694,69 @@ Now provide your response immediately without any preamble in text but not in ma
         "audio_node_result": audio_node_response
     }
 
+def excel_node(state: State):
+    """
+Excel handler model that can analyze excel files and answer questions about it.
+This node does not handle images or pictures.
+This node does not handle video.
+This node does not handle code.
+This node does not handle audio.
+
+Args:
+    state (State): with question key inside
+
+Returns:
+    dict: A dictionary containing the response from the excel handler node, with the key 'excel_node_result' holding the list of messages generated by the excel handler model.
+    """
+
+    loader = UnstructuredExcelLoader(state["downloaded_file"], mode="elements")
+    docs = loader.load()
+
+    prompt = f"""
+You are a powerful assistant which handles excel files.
+You do not handle images or pictures.
+You do not handle video.
+You do not handle code.
+You do not handle audio
+
+1. You need to fully understand the question.
+2. You must analyze the excel file to answer the question.
+3. If the question requires counting, identifying, or describing, be precise and clear in your response.
+4. Do not make up information that is not in the excel file.
+
+Here is the question {state['question']}
+Here is the excel file loaded in a Document object: {docs}. You will find htlm content of the file in the 'text_as_html' key.
+
+Now provide your response immediately without any preamble in text but not in markdown.
+    """
+
+    response = big_model.invoke(
+        input=prompt,
+        # config={
+        #     "callbacks": [langfuse_handler]
+        # }
+    )
+
+    response.pretty_print()
+
+    return {
+        "excel_node_result": response
+    }
+
 def format_answer_node(state: State):
+    """
+Format answer node that formats the answer of the last node.
+This node does not handle images or pictures.
+This node does not handle video.
+This node does not handle audio.
+This node does not handle code.
+
+Args:
+    state (State): with question key inside, and all other nodes results
+
+Returns:
+    dict: A dictionary containing the response from the format answer node, with the key 'format_answer_node_result' holding the list of messages generated by the format answer model.
+    """
     
     prompt = """
 You are the best assistant for final answer formating.
@@ -696,7 +783,6 @@ FINAL ANSWER: [YOUR FINAL ANSWER]
 - If asked for a comma-separated list:
     - Apply the above rules for numbers and strings to each element in the list.
     - And take care of having a space after each comma.
-    - If there is a list with single letter for each item, write them with the same case as the question.
 
 ## Constraints
 - You must not answer if the constraints above are not respected.
@@ -710,7 +796,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"]:
+    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"]:
         result = state.get(node_result, "")
         if result:
             # Ensure result is a string. If it's a message object, extract its content.
@@ -722,7 +808,7 @@ Now provide your response immediately without any preamble in text but not in ma
 
     sys_msg = SystemMessage(content=prompt)
 
-    response = [general_model.invoke([sys_msg] + state["messages"]+ question + nodes_response)]
+    response = [big_model.invoke([sys_msg] + state["messages"]+ question + nodes_response)]
 
     return {
         "messages": response,
@@ -748,6 +834,7 @@ You do not handle code
         - .wav or .mp3 is for audio
         - a youtube url is for video
         - a .jpg, .png, .jpeg is for image
+        - a .xlsx or .xls is for excel
 3. You must think hard about what is relevant in the question to make the best choice for the next node
 4. You must not answer the question by yourself
 5. Report your thought process in detail, explaining your reasoning step-by-step.
@@ -759,6 +846,7 @@ Here are the nodes you can choose:
 - video_node: {video_node.__doc__}
 - audio_node: {audio_node.__doc__}
 - code_node: {code_node.__doc__}
+- excel_node: {excel_node.__doc__}
 
 Here is the question : {state['question']}
 Here is the file : {state.get("input_file", "no file to handle")}
@@ -773,11 +861,11 @@ Now provide your response immediately respecting this format in lower case: 'nex
 
     sys_msg = SystemMessage(content=system_prompt)
 
-    entry_node_response = [general_model.invoke([sys_msg] + state["messages"])]
+    entry_node_response = [light_model.invoke([sys_msg] + state["messages"])]
 
     entry_node_response[-1].pretty_print()
 
-    regex_result = re.search(r'.*next\s*node:\s*(?P<next_node>[a-z_]*).*$', entry_node_response[-1].content, re.IGNORECASE)
+    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)
 
     next_node = "END"
     if regex_result:
@@ -805,6 +893,7 @@ def build_graph():
     builder.add_node("audio_node", audio_node)
     builder.add_node("code_node", code_node)
     builder.add_node("thinking_node", thinking_node)
+    builder.add_node("excel_node", excel_node)
     builder.add_node("format_answer_node", format_answer_node)
 
     builder.add_edge(START, "entry_node")
@@ -819,6 +908,7 @@ def build_graph():
             "video_node": "video_node",
             "audio_node": "audio_node",
             "code_node": "code_node",
+            "excel_node": "excel_node",
             "thinking_node": "thinking_node"
         }
     )
@@ -828,6 +918,7 @@ def build_graph():
     builder.add_edge("video_node", "format_answer_node")
     builder.add_edge("audio_node", "format_answer_node")
     builder.add_edge("code_node", "format_answer_node")    
+    builder.add_edge("excel_node", "format_answer_node")
     builder.add_edge("thinking_node", "format_answer_node")
     builder.add_edge("format_answer_node", END)
 
@@ -848,12 +939,11 @@ if __name__ == "__main__":
     with open("./responses.json", "r") as responses:
         json_responses = json.loads(responses.read())
 
-    # input =  {
-    #     "task_id": "cca530fc-4052-43b2-b130-b30968d8aa44",
-    #     "question": "Review the chess position provided in the image. It is black's turn. Provide the correct next move for black which guarantees a win. Please provide your response in algebraic notation.",
-    #     "Level": "1",
-    #     "file_name": "cca530fc-4052-43b2-b130-b30968d8aa44.png"
-    # }
+    # 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"
+    # }]
 
     with open("questions.json", "r") as questions:
         json_questions = json.loads(questions.read())
@@ -897,5 +987,3 @@ if __name__ == "__main__":
             print(f"Expected: {json_responses.get(task_id, '')}")
             print(f"Got: {answer}")
 
-        time.sleep(10)
-