agents/researcher.py

from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langgraph.prebuilt import create_react_agent
from langchain.tools import tool
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain.chains.llm import LLMChain
from langchain_core.messages import SystemMessage
from langchain_core.output_parsers.string import StrOutputParser

from agents.prompts import (
    language_detection_prompt,
    research_agent_prompt,
    summarize_prompt,
    translate_prompt,
)


load_dotenv()

# Initialize model
model = init_chat_model("gemini-2.0-flash", model_provider="google_genai")


# Chain using your model
language_detection_chain = LLMChain(llm=model, prompt=language_detection_prompt, output_parser=StrOutputParser())

# Tool function
@tool
def detect_search_language(craft: str) -> str:
    """Uses an LLM to decide the best language to search for information about a given craft."""
    return language_detection_chain.run({"craft": craft})


# Tool 2: Search the web using Tavily
@tool
def web_search_in_language(query: str) -> str:
    """Search the internet for a given query (in any language) and return a few relevant results."""
    search_tool = TavilySearchResults(k=5)
    return search_tool.run(query)

# Tool 3: Translate text to English
translate_chain = LLMChain(llm=model, prompt=translate_prompt, output_parser=StrOutputParser())

@tool
def translate_to_english(text: str) -> str:
    """Translates a given text into English.

    Args:
        text (str): input text

    Returns:
        str: English translation
    """
    return translate_chain.run({"text": text})


# Tool 4: Summarize translated content
summarize_chain = LLMChain(llm=model, prompt=summarize_prompt, output_parser=StrOutputParser())

@tool
def summarize_craft_intro(text: str) -> str:
    """Summarizes a given text about a specific craft as a craft introduction for beginners.

    Args:
        text (str): text about a craft

    Returns:
        str: summary
    """
    return summarize_chain.run({"text": text})


# Define the agent
craft_research_agent = create_react_agent(
    model=model,
    tools=[
        detect_search_language,
        web_search_in_language,
        translate_to_english,
        summarize_craft_intro
    ],
   prompt = SystemMessage(content=research_agent_prompt.format()),
    name="craft_research_agent"
)



# Example usage
if __name__ == "__main__":
    response = craft_research_agent.invoke({"input": "I want to learn Bulgarian lacework"})
    print(response)