Create main.py

main.py

@@ -0,0 +1,138 @@
+import os
+import uvicorn
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import httpx
+from duckduckgo_search import DDGS
+
+app = FastAPI(title="Edyx Situation Aware AI Pipeline")
+
+# Allow requests from the Edyx gateway/frontend
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+class ChatRequest(BaseModel):
+    message: str
+    messages: list = []
+    
+GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
+
+# Fallback check - if we were actually deploying on HF with a local GGUF, 
+# we would load llama-cpp-python here. For this stage, we'll setup the Groq primary pipeline.
+
+async def evaluate_needs_search(query: str) -> bool:
+    """Uses a fast, small model to determine if the query requires real-time data."""
+    if not GROQ_API_KEY:
+        return False
+        
+    system_prompt = """You are a highly efficient classification router.
+Determine if the user's query requires up-to-date, real-time information or current events data from the internet to answer accurately.
+Respond ONLY with "YES" if it requires search, or "NO" if it can be answered with general knowledge up to 2023.
+DO NOT provide any other text."""
+
+    try:
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                "https://api.groq.com/openai/v1/chat/completions",
+                headers={"Authorization": f"Bearer {GROQ_API_KEY}"},
+                json={
+                    "model": "llama3-8b-8192", # Fast and cheap for routing
+                    "messages": [
+                        {"role": "system", "content": system_prompt},
+                        {"role": "user", "content": query}
+                    ],
+                    "temperature": 0.1,
+                    "max_tokens": 10
+                },
+                timeout=10.0
+            )
+            response.raise_for_status()
+            result = response.json()
+            answer = result['choices'][0]['message']['content'].strip().upper()
+            return "YES" in answer
+    except Exception as e:
+        print(f"Routing evaluation error: {e}")
+        return False # Default to no search on error to save latency
+
+def perform_search(query: str, max_results: int = 3) -> str:
+    """Performs a web search using DuckDuckGo."""
+    try:
+        with DDGS() as ddgs:
+            results = list(ddgs.text(query, max_results=max_results))
+            
+        if not results:
+            return "No recent information found."
+            
+        context = "Here is some current information I found on the web regarding the user's query:\n\n"
+        for i, r in enumerate(results):
+            context += f"Source {i+1} [{r.get('title', 'No Title')}]: {r.get('body', '')}\n"
+        return context
+    except Exception as e:
+        print(f"Search error: {e}")
+        return "Search failed or was blocked."
+
+@app.post("/chat/completions")
+async def situation_aware_chat(request: ChatRequest):
+    if not GROQ_API_KEY:
+        raise HTTPException(status_code=500, detail="GROQ_API_KEY is not set in the environment.")
+
+    # 1. Evaluate if search is needed
+    user_query = request.message
+    needs_search = await evaluate_needs_search(user_query)
+    
+    context_injection = ""
+    if needs_search:
+        print(f"Query '{user_query}' requires search. Fetching data...")
+        context_injection = perform_search(user_query)
+        print("Search complete.")
+
+    # 2. Prepare the final prompt
+    system_base = "You are 'Situation Aware AI', an advanced assistant integrated into the Edyx platform."
+    if context_injection:
+        system_base += "\n\nThe user has asked a question that requires current knowledge. You have been provided with real-time web search results below. Incorporate this information seamlessly into your answer to provide the most up-to-date and accurate response. Do not mention that you 'searched the web' unless asked, just present the facts.\n\n" + context_injection
+
+    # Construct message array preserving history
+    final_messages = [{"role": "system", "content": system_base}]
+    
+    # Add previous history (excluding the current message if it's already in the list)
+    for msg in request.messages:
+        final_messages.append({"role": msg.get("role", "user"), "content": msg.get("content", "")})
+        
+    # Ensure current query is at the end if not provided in history block
+    if not request.messages or request.messages[-1].get("content") != user_query:
+         final_messages.append({"role": "user", "content": user_query})
+
+    # 3. Call Primary LLM
+    try:
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                "https://api.groq.com/openai/v1/chat/completions",
+                headers={"Authorization": f"Bearer {GROQ_API_KEY}"},
+                json={
+                    "model": "llama3-70b-8192", 
+                    "messages": final_messages,
+                    "temperature": 0.5,
+                    "max_tokens": 4096
+                },
+                timeout=30.0
+            )
+            response.raise_for_status()
+            result = response.json()
+            return result
+    except Exception as e:
+        print(f"Primary LLM Error: {e}")
+        # Here we would fallback to `llama-cpp-python` local inference
+        raise HTTPException(status_code=503, detail="Primary AI service is currently unavailable.")
+
+@app.get("/health")
+def health_check():
+    return {"status": "ok", "service": "edyx-situation-aware-pipeline"}
+
+if __name__ == "__main__":
+    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)