agent.py

from langchain_community.document_loaders import WikipediaLoader
from langchain_community.document_loaders import ArxivLoader
from langchain_core.tools import tool
from langgraph_supervisor import create_supervisor
from langchain.chat_models import init_chat_model
import os
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from langchain_experimental.agents import create_pandas_dataframe_agent
from openai import OpenAI
import re
from pydantic import BaseModel
from typing import List, Dict
import pandas as pd

from fonctions import clean_response, response_from_agent
class ToolInput(BaseModel):
    data: List[Dict]
    question: str

api_open_ai_agent_key=os.environ["OPENAI_API_KEY"]
client = OpenAI(api_key=api_open_ai_agent_key)
llm_4o = ChatOpenAI(
    model_name="gpt-4o",
    openai_api_key=api_open_ai_agent_key,  # ou variable d’environnement
)
llm_4_1 = ChatOpenAI(
    model_name='gpt-4.1',
    openai_api_key=api_open_ai_agent_key,  # ou variable d’environnement
)
llm_reasoning = ChatOpenAI(
    model_name = "o3-2025-04-16",
    openai_api_key=api_open_ai_agent_key,  # ou variable d’environnement
)



llm_reasoning_small = ChatOpenAI(
    model_name = "o4-mini-2025-04-16",
    openai_api_key=api_open_ai_agent_key,  # ou variable d’environnement
)

text_exemple = """
Set: {x, y, z}

Table:
   * | x | y | z
   --------------
   x | x | y | z
   y | y | x | x
   z | z | x | y

→ z * y = x, but y * z = x → equal  
→ y * z = x, but z * y = x → equal  
→ y * y = x, but y * y = x → equal  
→ All pairs commute → No counter-example

Answer: The operation is commutative.
"""
exemple_2 = """
Table:
   * | a | b
   ------------
   a | a | b
   b | a | a

→ a * b = b, but b * a = a ≠ b  
→ Counter-example: a, b  
Answer: a, b

"""

@tool
def wiki_search(query: str) -> dict[str, str]:
    """Search Wikipedia for a query and return maximum 2 results.

    Args:
        query: The search query."""
    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
            for doc in search_docs
        ])
    return {"wiki_results": formatted_search_docs}


@tool
def arvix_search(query: str) -> dict[str, str]:
    """Search Arxiv for a query and return maximum 3 result.

    Args:
        query: The search query."""
    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
            for doc in search_docs
        ])
    return {"arvix_results": formatted_search_docs}



prompt_web="""
        "You are a websearch agent.\n\n"
        "INSTRUCTIONS:\n"
        "- Assist ONLY with internet related tasks. DO NOT do any math\n"
        "- After you're done with your tasks, respond to the supervisor directly\n"
        "- Respond ONLY with the results of your work, do NOT include ANY other text."
        "- You can browse the web then give your results to your supervisor "
"""
@tool
def web_search_openai_tool(query: str) -> dict[str, str]:
    """Call web search and the output is a structured text answering the query)

    Args:
        query: The search query."""
    response = client.responses.create(
        model="gpt-4o",
        tools=[{"type": "web_search_preview"}],
        input=prompt_web + query
    )

    return {"web_results": response.output_text}
@tool
def add(a: float, b: float):
    """Add two numbers."""
    return a + b

@tool
def multiply(a: float, b: float):
    """Multiply two numbers."""
    return a * b

@tool
def divide(a: float, b: float):
    """Divide two numbers."""
    return a / b


@tool
def create_agent_and_answer(dict_data, question) -> str:
    """ From a dataframe, can anwser any question
    The input should be like :
    - dict_data= a dict
    - question= a str. Exemple : "Quel est l'âge moyen ?",

    """
    df = pd.DataFrame(dict_data)
    agent_excel = create_pandas_dataframe_agent(llm_4_1, df, verbose=True, allow_dangerous_code=True)
    text = agent_excel.run(question)

    return text


research_agent = create_react_agent(
    model=llm_4_1,
    tools=[wiki_search, arvix_search],
    prompt=(
        "You are a research agent.\n\n"
        "INSTRUCTIONS:\n"
        "- Assist ONLY with research-related tasks, DO NOT do any math\n"
        "- After you're done with your tasks, respond to the supervisor directly\n"
        "- Respond ONLY with the results of your work, do NOT include ANY other text."
        "- You only have access to arxiv or wikipedia, no other website"
    ),
    name="research_agent",
)

web_search_openai_agent = create_react_agent(
    model=llm_4_1,
    tools=[web_search_openai_tool],
    prompt=(
        "You are a websearch agent.\n\n"
        "INSTRUCTIONS:\n"
        "- Assist ONLY with internet related tasks. DO NOT do any math\n"
        "- After you're done with your tasks, respond to the supervisor directly\n"
        "- Respond ONLY with the results of your work, do NOT include ANY other text."
        "- You can browse the web then give your results to your supervisor "
    ),
    name="web_search_openai_agent",
)

math_agent = create_react_agent(
    model=llm_reasoning_small,

    tools=[add, multiply, divide],
    prompt=(
        "You are a mathematical reasoning agent.\n\n"
        "INSTRUCTIONS:\n"
        "- You are given mathematical structures such as sets, operations, and tables.\n"
        "- You must analyze them for properties like commutativity, associativity, identity, etc.\n"
        "- When given a Cayley table, check if x * y == y * x for all pairs to test commutativity.\n"
        "- Return a precise answer, including any counter-example pairs if they exist.\n"
        "- If a counter-example exists, return the set of involved elements in alphabetical order as a comma-separated list.\n"
        "- Do not rely on numeric calculation only — work symbolically when needed.\n"
        f"- Here are two examples : Example 1:{text_exemple}.\n"
        f"Example 2:{exemple_2}."),





    name="math_agent")


reflexion_agent = create_react_agent(
    model=llm_reasoning_small,

    tools=[],
    prompt=(
        "You are an intelligent agent\n\n"
        "INSTRUCTIONS:\n"
        "- Assist ONLY when you are called \n"
        "- You are the most intelligent agent"
        "- Your job is to solve hard problems when your supervisor ask you to do so"
        "- Give him a precise answer. "
    ),
    name="reflexion_agent",
)
agent_excel_new = create_react_agent(
    model=llm_4o,

    tools=[create_agent_and_answer],
    prompt=(
        "You are an agent psecialized with Excel files.\n\n"
        "INSTRUCTIONS:\n"
        "- Assist ONLY when an Excel file is mentionned \n"
        "- First, you will receive a dict that you can give to the associated file. If you don't have one, ask it to the supervisor.\n"
        "- Then you use the 'create_agent_and_answer' tool to answer the question.   \n"
        "- Once you have got a response from the 'create_agent_and_answer_tool', transmit it to your supervisor \n"
    ),
    name="agent_excel_new",
)
supervisor = create_supervisor(
    model=init_chat_model("openai:gpt-4.1", api_key = api_open_ai_agent_key),
    agents=[research_agent, web_search_openai_agent, agent_excel_new, reflexion_agent],
    prompt=(
        "You are a supervisor managing four agents:\n"
        "- research_agent: Specialised in ArXiv and Wikipedia. Assign research-related tasks to this agent.\n"
        "- reflexion: Called when the supervisor need a reflexion, not general knowledge. Can handle  math-related tasks such as solving equations, performing calculations or working on abstract maths subject such as matrix or demonstrating subjects.\n"
        "- web_search_openai_agent: Can browse the web to find up-to-date and relevant information. Assign web-related tasks to this agent.\n"
        "- agent_excel_new: Can understand tabular data. If an excel file is mentioned, call this agent. \n"
        "Assign work to one agent at a time. Do not call agents in parallel.\n"
        "The reflexion agent is your best weapon when the is a complex question. Call him only one time maximum by question.\n"
        " If there is an attached file, it will already loaded. Juste give the information to the agent. \n"
        "When a new question arises, if it is about an information that you can find on Wikipedia,  first consult the research_agent — it may provide useful information.\n"
        "If research_agent yields no results, then delegate the task to web_search_openai_agent.\n"
        "Each time you receive information from an agent, you have to analyze, process it then decide what to do (call an agent or give your final answer).\n"
        "As soon as you get a question, you have to analyze it and determine which agent is the most competent. Call at least one for each question.\n"
        "IMPORTANT : Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number NEITHER use units such as $, percent sign, or the currency unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string. If no punctuation is precised, don't add any. If you are asked for a price, don't precise the format, only the number. Respect the requested format"

    ),
    add_handoff_back_messages=True,
    output_mode="full_history",
).compile()