File size: 2,242 Bytes
3348cda | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | 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} |