main.py

from typing import TypedDict
from dotenv import load_dotenv
from langchain.agents import create_agent
from pydantic import BaseModel
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import PydanticOutputParser
from langchain.agents.middleware import ModelRequest, dynamic_prompt
from tools import wiki_tool, save_to_txt

load_dotenv()

class ResearchResponse(BaseModel):
    topic:str
    summary:str
    sources:list[str]
    tools_used:list[str]

class Context(TypedDict):
    user_role:str

@dynamic_prompt
def prompt(request: ModelRequest) -> str:
    user_role = request.runtime.context.get('user_role',"user")
    base_prompt = (
        "You are a research assistant that always gives detailed and reliable information. "
        "After generating the research result, always call the `save_to_txt` tool "
        "to save the output to a text file."
    )
    if user_role == "student":
        return f"{base_prompt} Explain concepts simply and avoid jargon."
    elif user_role == "researcher":
        return f"{base_prompt} Provide detailed technical responses."
    return base_prompt

llm = ChatOpenAI(model="gpt-4o-mini")

parser = PydanticOutputParser(pydantic_object=ResearchResponse)

tools = [wiki_tool, save_to_txt]

agent = create_agent(
    model="gpt-4o-mini",
    tools=tools,
    middleware=[prompt],
    context_schema=Context,
)

query = input("What can i help you research? ")
raw_response = agent.invoke(
    {"messages": [{"role": "user", "content": query}]},
    context={"user_role": "researcher"}  # örnek role
)

messages = raw_response.get("messages", [])
last_ai_message = None
for m in reversed(messages):
    if m.__class__.__name__ == "AIMessage" and m.content:
        last_ai_message = m.content
        break

if not last_ai_message:
    print("⚠️ No AI message found in response.")
else:
    print("\n🧠 --- Research Output ---\n")
    print(last_ai_message.replace(". ", ".\n"))
    print("\n--------------------------")