Update agentic.py

agentic.py

@@ -18,7 +18,8 @@ from langfuse.callback import CallbackHandler
 import base64
 import json
 import time
-from pprint import pprint
+import requests
+
 
 # import boto3
 
@@ -34,30 +35,41 @@ load_dotenv()
 # Initialize Langfuse CallbackHandler for LangGraph/Langchain (tracing)
 langfuse_handler = CallbackHandler()
 
+######## STATE ########
 class State(TypedDict):
     """
     A class representing the state of the agent.
     """
-    answer: str
     question: str
     messages: Annotated[list[AnyMessage], add_messages]
     input_file: str
+    downloaded_file: Optional[str]
+    task_id: str
+    web_wiki_search_node_result: AnyMessage
+    thinking_node_result: AnyMessage
+    vision_node_result: AnyMessage
+    video_node_result: AnyMessage
+    audio_node_result: AnyMessage
+    code_node_result: AnyMessage
+    next: str
 
+########################
 
-def get_assistant_model():
+######## MODELS ########
+def get_general_model():
 
     llm_provider = os.getenv("LLM_PROVIDER", "mistral")
 
     if llm_provider == "mistral":
-        assistant_model = ChatMistralAI(
-            model="mistral-small-latest",#"ministral-8b-latest",#"mistral-large-2411",#
+        general_model = ChatMistralAI(
+            model="mistral-large-2411",#"ministral-8b-latest",#"mistral-small-latest",#"mistral-small-latest",#
             temperature=0,
             max_retries=2,
             api_key=os.getenv("MISTRAL_API_KEY")
         )
 
     if llm_provider == "aws":
-        assistant_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,
@@ -66,7 +78,7 @@ def get_assistant_model():
             aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY")
         )
     
-    return assistant_model
+    return general_model
 
 def get_vision_model():
 
@@ -107,6 +119,21 @@ def get_video_handler_model():
 
     return video_handler_model
 
+def get_audio_handler_model():
+    audio_handler_model = ChatOpenAI(
+        model="gpt-4o-audio-preview-2024-12-17",#,gpt-4o-mini-audio-preview-2024-12-17",#
+        temperature=0,
+        max_tokens=None,
+        timeout=None,
+        max_retries=2,
+        api_key=os.getenv("OPENAI_API_KEY"),
+    )
+
+    return audio_handler_model
+
+########################
+
+######## Functions ########
 
 def download_youtube_content(url: str, output_path: Optional[str] = None) -> None:
     """
@@ -173,41 +200,263 @@ def download_youtube_content(url: str, output_path: Optional[str] = None) -> Non
 
     return video_file_name
 
-def search_webpage(state: State, url: str)-> str:
-    """
-    Search a web page based on the current state.
-    """
-    # Simulate a web page search and return a result
-    return "search_webpage"
 
 web_search = DuckDuckGoSearchRun()
 wikipedia_search = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
 
-assistant_model = get_assistant_model()
+def download_input_file(task_id: str) -> str:
+    """
+    Download the file specified in state input_file key.
+    You only need the task_id to download the file.
+
+    Args:
+        task_id (str): The task_id of the file to download.
+
+    Returns:
+        str: The path to the downloaded file.
+    """
+
+    output_path = os.path.join(os.getcwd(), 'downloads')
+
+    api_url = os.getenv("DEFAULT_API_URL")
+
+    # Create output directory if it doesn't exist
+    os.makedirs(output_path, exist_ok=True)
+
+    # Construct the full URL
+    url = f"{api_url}/files/{task_id}"
+
+    try:
+        # Send a GET request to download the file
+        response = requests.get(url, stream=True)
+        response.raise_for_status()  # Raise an error for bad status codes
+
+        headers = dict(response.headers)
+        attachement = headers["content-disposition"]
+
+        regex_result = re.search(r'filename="(.*)"', attachement)
+        filename = regex_result.group(1)
+
+        # Define the output file path
+        output_file_path = os.path.join(output_path, filename)
+
+        # Write the file to the output path
+        with open(output_file_path, 'wb') as file:
+            for chunk in response.iter_content(chunk_size=8192):
+                file.write(chunk)
+
+        print(f"File downloaded successfully and saved to: {output_file_path}")
+    
+        return output_file_path
+
+    except requests.exceptions.RequestException as e:
+        print(f"An error occurred while downloading the file: {str(e)}")
+        return ""
+
+########################
+
+######## LLM associations ########
+
+general_model = get_general_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,
+]
+
+download_file_tool = [ download_input_file ]
+
+web_search_node_agent = general_model.bind_tools(search_tools, parallel_tool_calls=False)
 
-def vision_model_call(state: State) -> str:
+def thinking_node(state: State) -> dict:
     """
-        Vision model that can analyze images and picture and answer questions about them.
+A powerful node to answer general questions, reflection, maths, deduction, prediction.
+This node does not handle files
+This node does not handle images or pictures
+This node does not handle videos
+This node does not handle audio
+This node does not handle code
+
+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 'thinking_node_result' holding the list of messages generated by the general model.
+    """
+
+    prompt = f"""
+You are a powerful assistant that answers general questions, reflection, maths, deduction, prediction.
+You do not handle files
+You do not handle images or pictures
+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 question to make the best answer
+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.
+    """
+
+    state["thinking_node_result"] = state.get("thinking_node_result", "")
+
+    sys_msg = SystemMessage(content=prompt)
+
+    thinking_node_response = [general_model.invoke([sys_msg] + [state["thinking_node_result"]])]
+
+    thinking_node_response[-1].pretty_print()
+
+    return {
+        "thinking_node_result": thinking_node_response,
+    }
 
-        Args:
-            state (State) with input_file key and question key inside
+def code_node(state: State) -> dict:
+    """
+A powerful node to handle and understand code.
+This node does not handle images or pictures
+This node does not handle videos
+This node does not handle audio
+This node does not access the web
+
+Args:
+    state (State): A dictionary containing the current state of the agent, including the 'question' key which holds the question to be answered.
 
-        Returns:
-            answer(string)
+Returns:
+    dict: A dictionary containing the response from the web search node, with the key 'code_node_result' holding the list of messages generated by the general model.
     """
+
+    with open(state["downloaded_file"], "r") as code_file:
+        code = code_file.read()
+
     prompt = f"""
-    You are a powerful vision assistant, you can analyze images and answer question about the picture
-    The question is : {state["question"]}
+You are a powerful assistant that handle and understand code.
+You do not handle images or pictures
+You do not handle videos
+You do not handle audio
+
+1. You need to fully understand the question.
+2. You must think hard about the code and predict the result to answer the question.
+3. Report your thought process in detail, explaining your reasoning step-by-step.
+
+Here is the question : {state['question']}
+Here is the code : {code}
+
+Now provide your response immediately without any preamble.
+    """
+
+    sys_msg = SystemMessage(content=prompt)
+
+    code_node_response = [general_model.invoke([sys_msg])]
+
+    code_node_response[-1].pretty_print()
+
+    return {
+        "code_node_result": code_node_response,
+    }
+
+def web_wiki_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
+This node does not handle images or pictures
+This node does not handle videos
+This node does not handle audio
+This node does not handle code
+
+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.
+    """
+
+    prompt = f"""
+You are a powerful assistant that makes research on the web in order to give the best answer to the question.
+You do not handle files
+You do not handle images or pictures
+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 question to make the best search with write words
+3. You must use the best of the tools you have to answer the question precisly
+4. Report your thought process in detail, explaining your reasoning step-by-step.
+
+Here are the tools available:
+web_search:
+    {web_search.description}
+    Args:
+        {web_search.args_schema}
+    Returns:
+        {web_search.response_format}
+
+wikipedia_search:
+    {wikipedia_search.description}
+    Args:
+        {wikipedia_search.args_schema}
+    Returns:
+        {wikipedia_search.response_format}
+
+Here is the question {state['question']}
+Now provide your response immediately without any preamble.
+    """
+
+    state["web_wiki_search_node_result"] = state.get("web_wiki_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_wiki_search_node_response[-1].pretty_print()
+
+    return {
+        "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.
+This node does not handle videos.
+This node does not handle audio.
+This node does not handle code.
+
+Args:
+    state (State): A dictionary containing the current state of the agent, including the 'question' key which holds the question to be answered and the 'input_file' key which holds the path to the image file.
+Returns:
+    dict: A dictionary containing the response from the vision node, with the key 'vision_node_result' holding the list of messages generated by the vision model.
+    """
+
+    prompt = f"""
+You are a powerful vision assistant, you can analyze images and answer question about the picture
+You do not download files
+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.
+
+Here is the question {state['question']}
+Now provide your response immediately without any preamble.
     """
-    # You need to consider system prompt of the first assistant to answer.
-    # Here is the system prompt: {state}
-    # """
 
     image_base64 = ""
     try:
-        with open(state["input_file"], "rb") as image_file:
+        with open(state["downloaded_file"], "rb") as image_file:
             image_bytes = image_file.read()
 
             image_base64 = base64.b64encode(image_bytes).decode("utf-8")
@@ -244,46 +493,68 @@ def vision_model_call(state: State) -> str:
             }
         ]
 
-        response = vision_model.invoke(
+        vision_node_response = [vision_model.invoke(
             input=message,
-            config={
-                "callbacks": [langfuse_handler]
-            }
-        )
+            # config={
+            #     "callbacks": [langfuse_handler]
+            # }
+        )]
 
-        return response.content + "\n\n"
+        vision_node_response[-1].pretty_print()
+
+        return {
+            "vision_node_result": vision_node_response
+        }
 
     except Exception as e:
         # A butler should handle errors gracefully
         error_msg = f"Error extracting text: {str(e)}"
         print(error_msg)
-        return ""
+        return {}
     
-def video_handler_model_call(state: State, video_url: str) -> str:
+def video_node(state: State) -> str:
     """
-        Video handler model that can analyze videos and answer questions about them.
+Video handler model that can analyze videos and answer questions about them.
+This node does not handle images or pictures.
+This node does not handle audio.
+This node does not handle code.
 
-        Args:
-            state (State): with question key inside
-            video_url (str): URL of the YouTube video to analyze.
+Args:
+    state (State): A dictionary containing the current state of the agent, including the 'question' key which holds the question to be answered.
 
-        Returns:
-            answer(string)
+Returns:
+    dict: A dictionary containing the response from the video handler node, with the key 'video_node_result' holding the list of messages generated by the video handler model.
     """
 
     prompt = f"""
-    You are a highly capable video analysis assistant. Your task is to watch and analyze the provided video content and answer the user's question as accurately and concisely as possible.
-
-    Instructions:
-    - Carefully observe the video, paying attention to relevant details, actions, and context.
-    - Focus on the user's question: "{state['question']}"
-    - If the question requires counting, identifying, or describing, be precise and clear in your response.
-    - If you are unsure, state what you can infer from the video.
-    - Do not make up information that is not visible or inferable from the video.
-
+You are a highly capable video analysis assistant. Your task is to watch and analyze the provided video content and answer the user's question as accurately and concisely as possible.
+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
+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.
+5. If you are unsure, state what you can infer from the video.
+6. Do not make up information that is not visible or inferable from the video.
+
+Here is the question {state['question']}
+Now provide your response immediately without any preamble.
     """
 
-    downloaded_video = download_youtube_content(url=video_url)
+    if re.search(r'youtube\.com', state["question"]):
+        # More flexible regex pattern to match YouTube URLs
+        regex_result = re.search(r"(?P<youtube_url>https://(?:www\.)?youtube\.com/watch\?v=[a-zA-Z0-9_-]+)", state["question"])
+        if regex_result:
+            video_url = regex_result.group("youtube_url")
+            downloaded_video = download_youtube_content(url=video_url)
+        else:
+            # Fallback if regex doesn't match
+            print("Could not extract YouTube URL from question. Using question as fallback.")
+            downloaded_video = state["downloaded_file"]
+    else:
+        downloaded_video = state["downloaded_file"]
 
     print(f"Downloaded video: {downloaded_video}")
 
@@ -311,113 +582,283 @@ def video_handler_model_call(state: State, video_url: str) -> str:
         }
     ]
 
-    response = video_handler_model.invoke(
+    video_node_response = [video_handler_model.invoke(
         input=message,
-        config={
-            "callbacks": [langfuse_handler]
+        # config={
+        #     "callbacks": [langfuse_handler]
+        # }
+    )]
+
+    video_node_response[-1].pretty_print()
+
+    return {
+        "video_node_result": video_node_response
+    }
+
+def audio_node(state: State) -> str:
+    """
+Audio handler model that can analyze audio 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.
+
+Args:
+    state (State): with question key inside
+
+Returns:
+    dict: A dictionary containing the response from the video handler node, with the key 'audioo_node_result' holding the list of messages generated by the audio handler model.
+    """
+
+    prompt = f"""
+You are a highly capable audio analysis assistant. Your task is to listen to and analyze the provided audio content and answer the user's question as accurately and concisely as possible.
+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:
+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.
+5. If you are unsure, state what you can infer from the audio.
+6. Do not make up information that is not audible or inferable from the audio.
+
+Here is the question {state['question']}
+Now provide your response immediately without any preamble.
+    """
+
+    downloaded_audio = state["downloaded_file"]
+
+    print(f"Downloaded audio: {downloaded_audio}")
+
+    audio_format = re.search(r'\.(\w+)$', downloaded_audio).group(1)
+
+    with open(downloaded_audio, "rb") as audio_file:
+        encoded_audio = base64.b64encode(audio_file.read()).decode()
+
+    os.remove(downloaded_audio)
+
+    message = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "text",
+                    "text": prompt,
+                },
+                {
+                    "type": "input_audio",
+                    "input_audio": {
+                        "data": encoded_audio,
+                        "format": audio_format,
+                    }
+                },
+            ]
         }
-    )
-        
-    return response.content + "\n\n"
-    
-# Tools
-tools = [
-    web_search,
-    # search_webpage,
-    wikipedia_search,
-    vision_model_call,
-    video_handler_model_call
-]
+    ]
 
-assistant_with_tools = assistant_model.bind_tools(tools, parallel_tool_calls=False)
+    audio_node_response = [audio_handler_model.invoke(
+        input=message,
+        # config={
+        #     "callbacks": [langfuse_handler]
+        # }
+    )]
 
+    audio_node_response[-1].pretty_print()
 
-def assistant(state: State)-> str:
-    # System message
-    textual_description_of_tool=f"""
-    web_search:
-        {web_search.description}
-        Args:
-            {web_search.args_schema}
-        Returns:
-            {web_search.response_format}
-    
-    wikipedia_search:
-        {wikipedia_search.description}
-        Args:
-            {wikipedia_search.args_schema}
-        Returns:
-            {wikipedia_search.response_format}
-
-    vision_model_call:
-        {vision_model_call.__doc__}
-        Args:
-            {vision_model_call.__annotations__}
-        Returns:
-            {vision_model_call.__annotations__['return']}
-
-    video_handler_model_call:
-        {video_handler_model_call.__doc__}
-        Args:
-            {video_handler_model_call.__annotations__}
-        Returns:
-            {video_handler_model_call.__annotations__['return']}
-"""
+    return {
+        "audio_node_result": audio_node_response
+    }
+
+def format_answer_node(state: State):
     
-    with open("./prompt.txt", "r") as prompt_file:
-        system_prompt = prompt_file.read()
+    prompt = """
+You are the best assistant for final answer formating.
+
+1. You must not change the content of the response of the last node.
+2. You must fully understand the question
+3. You must return the answer by following hard the format and the constraints
+4. Report your thought process in detail, explaining your reasoning step-by-step.
+
+5. Conclude your answer with the following template:
+FINAL ANSWER: [YOUR FINAL ANSWER]
+
+## Response Format
+- If asked for a number:
+    - Provide the number without commas, dollar signs, percent signs, or any units (unless specified).
+- If asked for a string:
+    - Write the string without articles (a, an, the).
+    - Don't answer a full sentence when a short version is enough.
+    - Do not use abbreviations (e.g., for cities).
+    - Write digits in plain text (e.g., "one" instead of "1") unless specified otherwise.
+- If asked for a comma-separated list:
+    - Apply the above rules for numbers and strings to each element in the list.
+
+## Constraints
+- You must not answer if the constraints above are not respected.
+- Your final answer should be provided in the format: FINAL ANSWER: [YOUR FINAL ANSWER]
+- Your final answer should be a number, a string, or a comma-separated list of numbers and/or strings, following the specified formatting rules.
+
+Now provide your response immediately without any preamble.
+    """
+
+    nodes_response = [HumanMessage(content="Here are the results of the previous nodes")]
 
-    sys_msg = SystemMessage(content=system_prompt+textual_description_of_tool)
+    question = [HumanMessage(content=state["question"])]
 
-    response = [assistant_with_tools.invoke([sys_msg] + state["messages"])]
+    for node_result in ["web_wiki_search_node_result", "vision_node_result", "video_node_result", "audio_node_result", "thinking_node_result", "code_node_result"]:
+        result = state.get(node_result, "")
+        if result:
+            # Ensure result is a string. If it's a message object, extract its content.
+            if hasattr(result, "content"):
+                content = result.content
+            else:
+                content = str(result)
+            nodes_response.append(HumanMessage(content=content))
+
+    sys_msg = SystemMessage(content=prompt)
+
+    response = [general_model.invoke([sys_msg] + state["messages"]+ question + nodes_response)]
 
     return {
-        "system_prompt": system_prompt,
         "messages": response,
-        "question": state["question"],
-        "answer": state.get("answer", "")
     }
 
+########################
 
-def build_graph():
-    builder = StateGraph(State)
+######## Entry Node ########
+def entry_node(state: State)-> str:
+    # System message
+    
+    system_prompt = f"""
+You are a powerful assistant that handle the user message and manage other nodes in order to provide the best answer to the question.
+You do not handle images or pictures
+You do not handle videos
+You do not handle audio
+You do not handle code
+
+1. You need to fully understand the subject of the question
+2. You need to understand the subject of the question with the question itself and the file extension
+    For example of extensions:
+        - .py is for code
+        - .wav or .mp3 is for audio
+        - a youtube url is for video
+        - a .jpg, .png, .jpeg is for image
+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.
+
+Here are the nodes you can choose:
+- thinking_node: {thinking_node.__doc__}
+- web_wiki_search_node: {web_wiki_search_node.__doc__}
+- vision_node: {vision_node.__doc__}
+- video_node: {video_node.__doc__}
+- audio_node: {audio_node.__doc__}
+- code_node: {code_node.__doc__}
+
+Here is the question : {state['question']}
+Here is the file : {state.get("input_file", "no file to handle")}
+
+Now provide your response immediately respecting this format: 'next node: \'the node name you choose\''.
+    """
+
+    downloaded = ""
+    # If there's an input file, download it directly
+    if state.get("input_file", None):
+        downloaded = download_input_file(state.get("task_id"))
+
+    sys_msg = SystemMessage(content=system_prompt)
+
+    entry_node_response = [general_model.invoke([sys_msg] + state["messages"])]
+
+    entry_node_response[-1].pretty_print()
+
+    regex_result = re.search(r'.*next\snode:\s+(?P<next_node>.*)$', entry_node_response[-1].content)
+
+    next_node = "END"
+    if regex_result:
+        # Extract the node name and remove any quotes around it
+        next_node = regex_result.group("next_node").strip().strip('\'"')
+
+    return {
+        "next": next_node,
+        "downloaded_file": downloaded
+    }
 
-    # Nodes
-    builder.add_node("assistant", assistant)
-    # builder.add_node("reviewer", reviewer)
-    builder.add_node("tools", ToolNode(tools))
+########################
 
-    # Edges
-    builder.add_edge(START, "assistant")
+######## Build Graph ########
 
+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("vision_node", vision_node)
+    builder.add_node("video_node", video_node)
+    builder.add_node("audio_node", audio_node)
+    builder.add_node("code_node", code_node)
+    builder.add_node("thinking_node", thinking_node)
+    builder.add_node("format_answer_node", format_answer_node)
+
+    builder.add_edge(START, "entry_node")
+
+    # Conditional routing from entry_node to specialized nodes
     builder.add_conditional_edges(
-        "assistant",
-        tools_condition,
+        "entry_node",
+        lambda state: state["next"],
+        {
+            "web_wiki_search_node": "web_wiki_search_node",
+            "vision_node": "vision_node",
+            "video_node": "video_node",
+            "audio_node": "audio_node",
+            "code_node": "code_node",
+            "thinking_node": "thinking_node"
+        }
     )
-    builder.add_edge("tools", "assistant")
+    # After specialized node, go to END
+    builder.add_edge("web_wiki_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")
+    builder.add_edge("code_node", "format_answer_node")    
+    builder.add_edge("thinking_node", "format_answer_node")
+    builder.add_edge("format_answer_node", END)
+
 
     return builder.compile()
 
+########################
+
 if __name__ == "__main__":
 
     agent_graph = build_graph()
 
-    file_name = ""
+    with open("image.png", "wb") as png:
+        png.write(agent_graph.get_graph(xray=True).draw_mermaid_png())
+    
+    # print(vision_node.__doc__)
 
-    # question = "In the video https://www.youtube.com/watch?v=L1vXCYZAYYM, what is the highest number of bird species to be on camera simultaneously?"
-    question = "Who are the pitchers with the number before and after Taishō Tamai's number as of July 2023? Give them to me in the form Pitcher Before, Pitcher After, use their last names only, in Roman characters."
-    # file_name = question_json.get("file_name", "")
+    input = {
+        "task_id": "9d191bce-651d-4746-be2d-7ef8ecadb9c2",
+        "question": "Examine the video at https://www.youtube.com/watch?v=1htKBjuUWec. What does Teal'c say in response to the question \"Isn't that hot?\"", 
+        "file_name": ""
+    }
+
+    question = input.get("question", "No question found")
+    file_name = input.get("file_name", "")
+    task_id = input.get("task_id", "")
 
     print(f"QUESTION : {question}")
     print(f"FILE: {file_name}")
 
-    messages = [HumanMessage(content=f"Can you answer this question please ? {question}")]
+    user_prompt = [HumanMessage(content=f"Can you answer the question please ?")]
+
+    user_input = {"messages": user_prompt, "question": question, "input_file": file_name, "task_id": task_id}
 
     messages = agent_graph.invoke(
-        input={"messages": messages, "question": question, "input_file": file_name}, 
+        input=user_input,
         config={
             "recursion_limit": 10,
-            "callbacks": [langfuse_handler]
+            # "callbacks": [langfuse_handler]
         }
     )