main.py

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("/gemma4:e2b")
async def generate_response(request: PromptRequest):
    llm = OllamaLLM(
        model="gemma4:e2b",
        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("/qwen3.5:2b")
async def qwen_generate_response(request:PromptRequest):
    llm = OllamaLLM(
        model="qwen3.5:2b",
        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}