Create agentic.py

agentic.py

@@ -0,0 +1,510 @@
+from langgraph.graph import StateGraph, START, END
+from typing_extensions import TypedDict, Annotated, Literal, Optional
+from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage
+from langgraph.graph.message import add_messages
+from langchain_mistralai import ChatMistralAI
+from langchain_openai import ChatOpenAI
+from langgraph.prebuilt import ToolNode, tools_condition
+from langchain_core.runnables.graph import MermaidDrawMethod
+from langchain_community.tools import DuckDuckGoSearchRun
+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_google_vertexai import ChatVertexAI
+
+from langfuse.callback import CallbackHandler
+
+import base64
+
+# import boto3
+
+from yt_dlp import YoutubeDL
+import os
+# from urllib.parse import urlparse, parse_qs
+import re
+from dotenv import load_dotenv
+
+# Load env vars from .env file
+load_dotenv()
+
+# Initialize Langfuse CallbackHandler for LangGraph/Langchain (tracing)
+langfuse_handler = CallbackHandler()
+
+class State(TypedDict):
+    """
+    A class representing the state of the agent.
+    """
+    answer: str
+    question: str
+    messages: Annotated[list[AnyMessage], add_messages]
+    input_file: str
+
+
+def get_assistant_model():
+
+    llm_provider = os.getenv("LLM_PROVIDER", "mistral")
+
+    if llm_provider == "mistral":
+        assistant_model = ChatMistralAI(
+            model="mistral-small-latest",#"ministral-8b-latest",#
+            temperature=0,
+            max_retries=2,
+            api_key=os.getenv("MISTRAL_API_KEY")
+        )
+
+    if llm_provider == "aws":
+        assistant_model = ChatBedrock(
+            model_id="arn:aws:bedrock:us-east-1:416545197702:inference-profile/us.amazon.nova-lite-v1:0",
+            # provider="amazon",
+            temperature=0,
+            region_name="eu-west-3",
+            aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID"),
+            aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY")
+        )
+    
+    return assistant_model
+
+def get_vision_model():
+
+    vlm_provider = os.getenv("VLM_PROVIDER", "mistral")
+
+    if vlm_provider == "openai":
+        print("Spawning Open AI VLM")
+        vision_model = ChatOpenAI(
+            model="gpt-4o",#-mini",
+            temperature=0,
+            max_tokens=None,
+            timeout=None,
+            max_retries=2,
+            api_key=os.getenv("OPENAI_API_KEY"),
+        )
+    
+    if vlm_provider == "mistral":
+        print("Spawning Mistral VLM")
+        vision_model = ChatMistralAI(
+            model="pixtral-12b-2409",#"mistral-small-latest","pixtral-large-latest",#
+            temperature=0,
+            max_retries=2,
+            api_key=os.getenv("MISTRAL_API_KEY")
+        )
+    
+    return vision_model
+
+def get_video_handler_model():
+
+    video_handler_model = ChatGoogleGenerativeAI(
+        model="gemini-2.0-flash",
+        temperature=0,
+        max_tokens=None,
+        timeout=None,
+        max_retries=2,
+        # other params...
+    )
+
+    return video_handler_model
+
+# reviewer_model = ChatMistralAI(
+#     model="mistral-small-latest",
+#     temperature=0,
+#     max_retries=2,
+#     api_key=os.getenv("MISTRAL_API_KEY"),
+# )
+
+def download_youtube_content(url: str, output_path: Optional[str] = None) -> None:
+    """
+    Download YouTube content (single video or playlist) in MP4 format only.
+    
+    Args:
+        url (str): URL of the YouTube video or playlist
+        output_path (str, optional): Directory to save the downloads. Defaults to './downloads'
+    """
+    # Set default output path if none provided
+    if output_path is None:
+        output_path = os.path.join(os.getcwd(), 'downloads')
+    
+    # Create output directory if it doesn't exist
+    os.makedirs(output_path, exist_ok=True)
+    
+    # Configure yt-dlp options for MP4 only
+    ydl_opts = {
+        'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best',
+        'merge_output_format': 'mp4',
+        'ignoreerrors': True,
+        'no_warnings': False,
+        'extract_flat': False,
+        # Disable all additional downloads
+        'writesubtitles': False,
+        'writethumbnail': False,
+        'writeautomaticsub': False,
+        'postprocessors': [{
+            'key': 'FFmpegVideoConvertor',
+            'preferedformat': 'mp4',
+        }],
+        # Clean up options
+        'keepvideo': False,
+        'clean_infojson': True
+    }
+
+
+    ydl_opts['outtmpl'] = os.path.join(output_path, '%(title)s.%(ext)s')
+    print("Detected single video URL. Downloading video...")
+
+    try:
+        with YoutubeDL(ydl_opts) as ydl:
+            # Download content
+            ydl.download([url])
+            print(f"\nDownload completed successfully! Files saved to: {output_path}")
+
+    except Exception as e:
+        print(f"An error occurred: {str(e)}")
+
+    result = os.listdir(output_path)
+
+    video_file_names = [x for x in result if re.match(r".*\.mp4$", x)]
+
+    if len(video_file_names) == 1:
+        video_file_name = video_file_names.pop()
+        video_file_name = f"{output_path}/{video_file_name}"
+    else:
+        video_file_name = None
+
+    for other_files in result:
+        if f"{output_path}/{other_files}" != video_file_name:
+            print(f"Removing file: {other_files}")
+            os.remove(os.path.join(output_path, other_files))
+
+    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()
+vision_model = get_vision_model()
+video_handler_model = get_video_handler_model()
+
+def vision_model_call(state: State) -> str:
+    """
+        Vision model that can analyze images and picture and answer questions about them.
+
+        Args:
+            state (State) with input_file key and question key inside
+
+        Returns:
+            answer(string)
+    """
+    prompt = f"""
+    You are a powerful vision assistant, you can analyze images and answer question about the picture
+    The question is : {state["question"]}
+    """
+    # 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:
+            image_bytes = image_file.read()
+
+            image_base64 = base64.b64encode(image_bytes).decode("utf-8")
+
+        mistral_image_handling = {
+            "type": "image_url",
+            "image_url": f"data:image/png;base64,{image_base64}",
+        }
+
+        openai_image_handling = {
+            "type": "image",
+            "source_type": "base64",
+            "mime_type": "image/png",  # or image/png, etc.
+            "data": image_base64,
+        }
+
+        vision_provider = "mistral"
+
+        if vision_provider == "openai":
+            image_handling = openai_image_handling
+        else:
+            image_handling = mistral_image_handling
+
+        message = [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": prompt,
+                    },
+                    image_handling
+                ]
+            }
+        ]
+
+        response = vision_model.invoke(
+            input=message,
+            config={
+                "callbacks": [langfuse_handler]
+            }
+        )
+
+        return response.content + "\n\n"
+
+    except Exception as e:
+        # A butler should handle errors gracefully
+        error_msg = f"Error extracting text: {str(e)}"
+        print(error_msg)
+        return ""
+    
+def video_handler_mode_call(state: State, video_url: str) -> str:
+    """
+        Video handler model that can analyze videos and answer questions about them.
+
+        Args:
+            state (State): with question key inside
+            video_url (str): URL of the YouTube video to analyze.
+
+        Returns:
+            answer(string)
+    """
+
+    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.
+
+    """
+
+    downloaded_video = download_youtube_content(url=video_url)
+
+    print(f"Downloaded video: {downloaded_video}")
+
+    video_mime_type = "video/mp4"
+
+    with open(downloaded_video, "rb") as video_file:
+        encoded_video = base64.b64encode(video_file.read()).decode("utf-8")
+
+    os.remove(downloaded_video)
+
+    message = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "text",
+                    "text": prompt,
+                },
+                {
+                    "type": "media",
+                    "data": encoded_video,  # Use base64 string directly
+                    "mime_type": video_mime_type,
+                },
+            ]
+        }
+    ]
+
+    response = video_handler_model.invoke(
+        input=message,
+        config={
+            "callbacks": [langfuse_handler]
+        }
+    )
+        
+    return response.content + "\n\n"
+
+# def router(state: State)-> Literal["OK", "retry", "failed"]:
+#     """Determine the next step based on FINAL ANSWER"""
+
+#     print(f"MESSAGES : {state['messages'][-1].content}")
+#     print(f"STATE : {state['system_prompt']}")
+#     if True:
+#         return "OK"
+#     else:
+#         if state.status == "ERROR":
+#             if state.error_count >= 3:
+#                 return "failed"
+#         return "retry"
+
+# def reviewer_validation(state: State) -> Literal["OK", "NOK"]:
+#     print(f"MESSAGES : {state['messages'][-1].content}")
+#     print(f"STATE : {state}")
+#     return "OK"
+#     if True:
+#         return "OK"
+#     else:
+#         return "NOK"
+    
+# Tools
+tools = [
+    web_search,
+    # search_webpage,
+    wikipedia_search,
+    vision_model_call,
+    video_handler_mode_call
+]
+
+assistant_with_tools = assistant_model.bind_tools(tools, parallel_tool_calls=False)
+# reviewer_with_tools = reviewer_model.bind_tools(tools, parallel_tool_calls=False)
+
+
+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_mode_call:
+        {video_handler_mode_call.__doc__}
+        Args:
+            {video_handler_mode_call.__annotations__}
+        Returns:
+            {video_handler_mode_call.__annotations__['return']}
+"""
+    
+    with open("./prompt.txt", "r") as prompt_file:
+        system_prompt = prompt_file.read()
+
+    file_prompt = f"If the question is about an image or a picture, and if you need to analyze a file, an image and so on... you must check the path {state['input_file']}"
+
+    video_handler_prompt = "If the question is about a video and if you need to analyze a video, you must use the video_handler_mode_call tool."
+
+    sys_msg = SystemMessage(content=system_prompt+file_prompt+video_handler_prompt+textual_description_of_tool)
+
+    response = [assistant_with_tools.invoke([sys_msg] + state["messages"])]
+
+    return {
+        "system_prompt": system_prompt,
+        "messages": response,
+        "question": state["question"],
+        "answer": state.get("answer", "")
+    }
+
+# def reviewer(state: State)-> str:
+
+#     with open("./prompt.txt", "r") as prompt_file:
+#         system_prompt = prompt_file.read()
+
+#     sys_msg = SystemMessage(content=f"""
+#     You are a powerful AI assistant reviewer.
+#     You must review the answer of the previous assistant.
+#     If you need you can correct the answer.
+#     You need to make sure the answer is formatted correctly.
+#     The response should not be sent if FINAL ANSWER: is not registered and does not respect the constraints of the initial prompt.
+                            
+#     Here is the prompt of the previous assistant :
+#     {system_prompt}
+
+#     Here is the question: {state['question']}
+#     """)
+
+#     response = [reviewer_with_tools.invoke([sys_msg] + state["messages"])]
+
+#     return {
+#         "messages": response,
+#     }
+
+def build_graph():
+    builder = StateGraph(State)
+
+    # Nodes
+    builder.add_node("assistant", assistant)
+    # builder.add_node("reviewer", reviewer)
+    builder.add_node("tools", ToolNode(tools))
+
+    # Edges
+    builder.add_edge(START, "assistant")
+
+    builder.add_conditional_edges(
+        "assistant",
+        tools_condition,
+    )
+    builder.add_edge("tools", "assistant")
+    # builder.add_edge("assistant", "reviewer")
+
+    # builder.add_conditional_edges(
+    #     "assistant",
+    #     router,
+    #     {
+    #         "OK": "reviewer",
+    #         "retry": "assistant",
+    #         "failed": END
+    #     }
+    # )
+    # builder.add_conditional_edges(
+    #     "reviewer",
+    #     reviewer_validation,
+    #     {
+    #         "OK": END,
+    #         "NOK": "assistant"
+    #     }
+    # )
+
+    return builder.compile()
+
+if __name__ == "__main__":
+
+    agent_graph = build_graph()
+
+    file_name = ""
+
+    # agent_graph.get_graph(xray=True).draw_mermaid_png(output_file_path="./graph.png", draw_method=MermaidDrawMethod.PYPPETEER)
+    question = "Where were the Vietnamese specimens described by Kuznetzov in Nedoshivina's 2010 paper eventually deposited? Just give me the city name without abbreviations."
+    # question = "Given this table defining * on the set S = {a, b, c, d, e}\n\n|*|a|b|c|d|e|\n|---|---|---|---|---|---|\n|a|a|b|c|b|d|\n|b|b|c|a|e|c|\n|c|c|a|b|b|a|\n|d|b|e|b|e|d|\n|e|d|b|a|d|c|\n\nprovide the subset of S involved in any possible counter-examples that prove * is not commutative. Provide your answer as a comma separated list of the elements in the set in alphabetical order."
+    # question = ".rewsna eht sa \"tfel\" drow eht fo etisoppo eht etirw ,ecnetnes siht dnatsrednu uoy fI"
+    # 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?"
+    # 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 = "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."
+    # question = "Can you describe the image ?"
+    # file_name = "images/cca530fc-4052-43b2-b130-b30968d8aa44.png"
+
+    messages = [HumanMessage(content=f"Can you answer this question please ? {question}")]
+
+    messages = agent_graph.invoke(
+        input={"messages": messages, "question": question, "input_file": file_name}, 
+        config={
+            "recursion_limit": 10,
+            "callbacks": [langfuse_handler]
+        }
+    )
+
+    for m in messages['messages']:
+        m.pretty_print()
+    
+    try:
+        regex_result = re.search(r"FINAL ANSWER:\s*(?P<answer>.*)$", messages['messages'][-1].content)
+        answer = regex_result.group("answer")
+    except:
+        regex_result = re.search(r"\s*(?P<answer>.*)$", messages['messages'][-1].content)
+        answer = regex_result.group("answer")
+
+    print(f"ANSWER : {answer}")
+    # from pprint import pprint