agent.py

"""
agent.py — Wires the LLM, tools, and prompt together. No logic, no prompts here.
All reasoning is handled by the LLM. All computation is handled by tools.
All prompt structure lives in prompt.py.
"""
import os
from dotenv import load_dotenv
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.prebuilt import create_react_agent
from tools import (search_academic_source, save_papers, get_paper_batch,
                   save_output, read_output, cluster_and_visualize,
                   get_pajais_taxonomy, read_pdf_text, enrich_doi,
                   read_word_text, import_from_scratch)
from prompt import build_prompt

load_dotenv()

TOOLS = [search_academic_source, save_papers, get_paper_batch, save_output,
         read_output, cluster_and_visualize, get_pajais_taxonomy,
         read_pdf_text, enrich_doi, read_word_text, import_from_scratch]

def _llm():
    provider = os.getenv("LLM_PROVIDER", "groq").lower()
    if provider == "mistral":
        from langchain_mistralai import ChatMistralAI
        return ChatMistralAI(model=os.getenv("MISTRAL_BIG", "mistral-small-latest"), temperature=0)
    elif provider == "huggingface":
        from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
        repo_id = os.getenv("HF_MODEL_ID", "mistralai/Mistral-Small-24B-Instruct-v0.1")
        token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")
        llm = HuggingFaceEndpoint(repo_id=repo_id, huggingfacehub_api_token=token, temperature=0.01)
        return ChatHuggingFace(llm=llm)
    from langchain_groq import ChatGroq
    return ChatGroq(model=os.getenv("GROQ_BIG", "llama-3.3-70b-versatile"), temperature=0)

supervisor = create_react_agent(
    model=_llm(),
    tools=TOOLS,
    prompt=build_prompt(TOOLS),     # prompt is built from the live tool list
    checkpointer=InMemorySaver()
)

def chat_with_agent(message: str, thread_id: str) -> str:
    config = {"configurable": {"thread_id": thread_id}}
    res = supervisor.invoke({"messages": [{"role": "user", "content": message}]}, config)
    return res["messages"][-1].content