from fastapi import FastAPI
from pydantic import BaseModel
from typing import List
from langchain_community.tools.tavily_search import TavilySearchResults
import os
from langgraph.prebuilt import create_react_agent
from langchain_openai import AzureChatOpenAI
from dotenv import load_dotenv
load_dotenv()
os.environ['AZURE_OPENAI_API_KEY'] = os.getenv('AZURE_OPENAI_API_KEY')
os.environ['OPENAI_API_VERSION'] = os.getenv('OPENAI_API_VERSION')
os.environ['AZURE_OPENAI_ENDPOINT'] = os.getenv('AZURE_OPENAI_ENDPOINT')
os.environ['TAVILY_API_KEY'] = os.getenv('TAVILY_API_KEY')

MODEL_NAME = [
   "gpt-4o"
]

tool = TavilySearchResults(max_results=5)

tools = [tool, ]

app = FastAPI()

class RequestState(BaseModel):
    model_name: str
    system_prompt: str
    messages: List[str]

@app.post("/chat")
def chat_endpoint(request: RequestState):
    """API Endpoint to chat with bot"""
    if request.model_name not in MODEL_NAME:
        return {"error": "Model not found"}
    
    llm = AzureChatOpenAI(
    azure_deployment="gpt-4o", 
    temperature=0,
    max_tokens=None,
    timeout=None,
    max_retries=2,
)
    agent = create_react_agent(llm, tools=tools, prompt=request.system_prompt)
    state = {"messages": request.messages}
    result = agent.invoke(state)
    return result

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)