import os 
from fastapi import FastAPI 
from pydantic import BaseModel 
from tavily import TavilyClient 
from datetime import date 
from langchain_ollama.llms import OllamaLLM 

app = FastAPI() 

class PromptRequest(BaseModel): 
    prompt: str 
    temperature: float = 0.5 
    
@app.get("/") 
def health(): 
    return {"ok": True} 
    
today_date = date.today() 

def search_tool(query:str): 
    api_key = os.environ.get("TAVILY_API_KEY") 
    client = TavilyClient(api_key) 
    response = client.search(query=query, include_answer="advanced", search_depth="advanced") 
    return response 
    
@app.post("/aya:8b") 
async def generate_response(request: PromptRequest): 
    llm = OllamaLLM( 
        model="aya:8b", 
        temperature=request.temperature, 
        base_url="http://localhost:11436" 
    ) 
    tool_prompt = f"System Role: You are an autonomous AI Agent with real-time internet access. Current Date: {today_date} TOOL_DEFINITION: - Name: Search_tool - Activation Command: Search [Your Query Here]" 
    response = llm.invoke(f'{tool_prompt}, User-Query:-{request.prompt}') 
    
    if "Search " in response: 
        new_query = response.removeprefix("Search ") 
        search_response = search_tool(query=new_query) 
        new_response = llm.invoke(f"Extra_information:- {search_response} User-Query:- {request.prompt}") 
        return {"response": new_response} 
    else: 
        return {"response": response}    

@app.post("/phi4:14b") 
async def qwen_generate_response(request:PromptRequest): 
    llm = OllamaLLM( 
        model="phi4:14b", 
        temperature=request.temperature, 
        base_url="http://localhost:11435" 
    ) 
    tool_prompt = f"System Role: You are an autonomous AI Agent with real-time internet access. Current Date: {today_date} TOOL_DEFINITION: - Name: Search_tool - Activation Command: Search [Your Query Here]" 
    response = llm.invoke(f'{tool_prompt}, User-Query:-{request.prompt}') 
    
    if "Search " in response: 
        new_query = response.removeprefix("Search ") 
        search_response = search_tool(query=new_query) 
        new_response = llm.invoke(f"Extra_information:- {search_response} User-Query:- {request.prompt}") 
        return {"response": new_response} 
    else: 
        return {"response": response}