Update main.py

main.py

@@ -1,45 +1,87 @@
-import os
+import requests
 from fastapi import FastAPI
 from pydantic import BaseModel
-from tavily import TavilyClient
 from datetime import date
-from langchain_ollama.llms import OllamaLLM
+from ollama import Client
 
 app = FastAPI()
+today_date = date.today()
 
 class PromptRequest(BaseModel):
     prompt: str
     temperature: float = 0.5
 
+TOOL_PROMPT = f"""You are an autonomous AI Agent with real-time internet access.
+Current Date: {today_date}
+
+You have access to one tool:
+- Search [query] : searches the internet and returns real-time results
+
+RULES:
+- If the user query needs current/real-time/recent information, respond with ONLY: Search [your search query]
+- If you can answer from your own knowledge, answer directly
+- Do NOT mix search command with other text
+- After getting search results, answer the user query using those results
+
+"""
+
+def search_tool(query: str) -> str:
+    try:
+        r = requests.get(
+            "https://s.jina.ai/",
+            params={"q": query},
+            headers={"Accept": "application/json", "X-Retain-Images": "none"},
+            timeout=10
+        )
+        data = r.json().get("data", [])[:5]
+        results = []
+        for item in data:
+            title = item.get("title", "")
+            desc = item.get("description", "") or item.get("content", "")[:300]
+            url = item.get("url", "")
+            results.append(f"- {title}\n  {desc}\n  Source: {url}")
+        return "\n\n".join(results) if results else "No results found."
+    except Exception as e:
+        return f"Search failed: {str(e)}"
+
+def run_llm(messages: list, temperature: float) -> str:
+    client = Client(host="http://localhost:11434")
+    response = client.chat(
+        model="qwen2.5-coder:7b",
+        messages=messages,
+        options={"temperature": temperature}
+    )
+    return response["message"]["content"].strip()
+
 @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("/qwen2.5-coder:7b")
 async def generate_response(request: PromptRequest):
-    llm = OllamaLLM(
-        model="qwen2.5-coder:7b",
-        temperature=request.temperature,
-        base_url="http://localhost:11434"
-    )
-    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}')
-    print(response)
-
-    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}")
-        print(new_response)
-        return {"response": new_response}
-    else:
-        return {"response": response}
+    # Step 1: ask LLM if it needs to search
+    messages = [
+        {"role": "system", "content": TOOL_PROMPT},
+        {"role": "user", "content": request.prompt}
+    ]
+    response = run_llm(messages, request.temperature)
+    print(f"[LLM RAW]: {response}")
+
+    # Step 2: if LLM wants to search
+    if response.strip().startswith("Search "):
+        query = response.strip().removeprefix("Search ").strip("[]").strip()
+        print(f"[SEARCH QUERY]: {query}")
+        search_results = search_tool(query)
+        print(f"[SEARCH RESULTS]: {search_results[:200]}...")
+
+        # Step 3: feed results back to LLM
+        messages.append({"role": "assistant", "content": response})
+        messages.append({
+            "role": "user",
+            "content": f"Search Results:\n{search_results}\n\nNow answer the original question: {request.prompt}"
+        })
+        final_response = run_llm(messages, request.temperature)
+        print(f"[FINAL]: {final_response}")
+        return {"response": final_response, "searched": True, "query": query}
 
+    return {"response": response, "searched": False}