|
|
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)
|
|
|
|
