Update main.py

main.py

@@ -1,33 +1,39 @@
 from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, Field
 import httpx
 import os
-from typing import List, Dict, Any, Optional
+import json
+import time
+import uuid
+from typing import List, Dict, Any, Optional, AsyncGenerator
 
 # --- Configuration ---
-# Your actual Inference API key should be set as an environment variable
 INFERENCE_API_KEY = os.environ.get("INFERENCE_API_KEY", "inference-00050468cc1c4a20bd5ca0997c752329")
 INFERENCE_API_URL = "https://api.inference.net/v1/chat/completions"
 SEARCH_API_URL = "https://rkihacker-brave.hf.space/search"
+MODEL_NAME = "Binglity-Lite"
+BACKEND_MODEL = "meta-llama/llama-3.1-8b-instruct/fp-8"
 
-# --- System Prompt ---
+# --- A More Advanced System Prompt ---
 SYSTEM_PROMPT = """
-You are "Binglity-Lite", a large language model acting as a helpful AI assistant.
-Your primary function is to provide accurate, comprehensive, and helpful answers by synthesizing information from real-time web search results.
-When you are given a user's query and a set of search results, you must follow these rules:
-1.  Carefully analyze the user's query to understand their intent.
-2.  Review all the provided search results to gather relevant facts, data, and perspectives.
-3.  Construct a single, coherent, and well-written response that directly answers the user's query.
-4.  Base your answer **only** on the information found in the provided search results. Do not use any prior knowledge.
-5.  If the search results do not contain enough information to answer the question, state that you couldn't find a definitive answer based on the search.
-6.  Do not list the search results. Instead, integrate the information from them into your response.
+You are "Binglity-Lite", a state-of-the-art AI assistant. Your purpose is to provide accurate, unbiased, and comprehensive answers by synthesizing information from real-time web search results.
+
+**Your Instructions:**
+1.  **Analyze the User's Query**: Deeply understand the user's question, intent, and the specific information they are seeking.
+2.  **Synthesize, Don't List**: Do not simply list or summarize the search results. Your primary task is to integrate the information from the multiple sources provided into a single, cohesive, and well-structured response.
+3.  **Be Factual and Unbiased**: Base your entire response ONLY on the information contained within the provided search results. Do not introduce any external knowledge or personal opinions.
+4.  **Handle Contradictions**: If the search results present conflicting information, acknowledge the discrepancy and present the different viewpoints as found in the sources.
+5.  **Address Insufficient Information**: If the search results do not contain enough information to provide a complete answer, explicitly state that. Do not speculate or fill in the gaps.
+6.  **Maintain a Helpful Tone**: Your persona is knowledgeable, helpful, and neutral.
+7.  **Structure for Clarity**: Use clear language and logical formatting (like paragraphs or bullet points if appropriate) to make the information easy to understand.
 """
 
 # --- FastAPI App ---
 app = FastAPI(
     title="Binglity-Lite API",
-    description="A web search-powered chat completions API.",
-    version="1.0.0",
+    description="A web search-powered, streaming-capable chat completions API.",
+    version="1.1.0",
 )
 
 # --- Pydantic Models for OpenAI Compatibility ---
@@ -38,14 +44,12 @@ class ChatMessage(BaseModel):
 class ChatCompletionRequest(BaseModel):
     model: str
     messages: List[ChatMessage]
-    max_tokens: Optional[int] = 1024
+    max_tokens: Optional[int] = 2048
     temperature: Optional[float] = 0.7
+    stream: Optional[bool] = False
 
 # --- Web Search Function ---
 async def perform_web_search(query: str) -> List[Dict[str, Any]]:
-    """
-    Performs a web search using an external API.
-    """
     async with httpx.AsyncClient() as client:
         try:
             response = await client.get(
@@ -61,102 +65,127 @@ async def perform_web_search(query: str) -> List[Dict[str, Any]]:
             print(f"An unexpected error occurred during web search: {str(e)}")
             return []
 
-# --- Helper to format search results for the LLM ---
 def format_search_results_for_prompt(results: List[Dict[str, Any]]) -> str:
-    """
-    Formats the list of search result dictionaries into a string for the LLM prompt.
-    """
     if not results:
-        return "No search results found."
-
-    formatted = "Here are the web search results:\n\n"
+        return "No relevant search results were found. Please inform the user that you could not find information on their query."
+    
+    formatted = "### Web Search Results ###\n\n"
     for i, result in enumerate(results):
-        formatted += f"Result [{i+1}]:\n"
+        formatted += f"Source [{i+1}]:\n"
         formatted += f"Title: {result.get('title', 'N/A')}\n"
         formatted += f"URL: {result.get('url', 'N/A')}\n"
-        formatted += f"Description: {result.get('description', 'N/A')}\n\n"
+        formatted += f"Content: {result.get('description', 'N/A')}\n\n"
     return formatted
 
-# --- API Endpoint ---
-@app.post("/v1/chat/completions")
-async def chat_completions(request: ChatCompletionRequest):
+# --- Streaming Logic ---
+async def stream_response_generator(payload: Dict[str, Any]) -> AsyncGenerator[str, None]:
     """
-    Implements a chat completions endpoint compatible with OpenAI's API.
-    It performs a web search based on the user's last message.
+    Yields chunks from the inference API, formatted for OpenAI compatibility.
     """
-    if request.model != "Binglity-Lite":
-        raise HTTPException(
-            status_code=400,
-            detail=f"Model not supported. Please use 'Binglity-Lite'. You used '{request.model}'.",
-        )
-
-    if not request.messages:
-        raise HTTPException(status_code=400, detail="The 'messages' field is required.")
+    headers = {
+        "Authorization": f"Bearer {INFERENCE_API_KEY}",
+        "Content-Type": "application/json",
+        "Accept": "text/event-stream"
+    }
+    
+    # Create a unique ID for the response stream
+    response_id = f"chatcmpl-{uuid.uuid4()}"
+    created_time = int(time.time())
+
+    async with httpx.AsyncClient(timeout=300.0) as client:
+        async with client.stream("POST", INFERENCE_API_URL, json=payload, headers=headers) as response:
+            if response.status_code != 200:
+                error_content = await response.aread()
+                raise HTTPException(status_code=response.status_code, detail=f"Error from inference API: {error_content.decode()}")
+
+            # Stream the response line by line
+            async for line in response.aiter_lines():
+                if line.startswith("data:"):
+                    line_data = line[5:].strip()
+                    if line_data == "[DONE]":
+                        # Send the final data chunk and the done message
+                        yield f"data: {json.dumps({'id': response_id, 'model': MODEL_NAME, 'object': 'chat.completion.chunk', 'created': created_time, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]})}\n\n"
+                        yield "data: [DONE]\n\n"
+                        break
+                    
+                    try:
+                        chunk = json.loads(line_data)
+                        # Reformat the chunk to be OpenAI compliant
+                        formatted_chunk = {
+                            "id": response_id,
+                            "object": "chat.completion.chunk",
+                            "created": created_time,
+                            "model": MODEL_NAME,
+                            "choices": [{
+                                "index": 0,
+                                "delta": chunk["choices"][0].get("delta", {}),
+                                "finish_reason": chunk["choices"][0].get("finish_reason")
+                            }]
+                        }
+                        yield f"data: {json.dumps(formatted_chunk)}\n\n"
+                    except json.JSONDecodeError:
+                        print(f"Could not decode JSON from line: {line_data}")
+                        continue
 
-    # Extract the last user message as the query
-    user_query = request.messages[-1].content
-    user_role = request.messages[-1].role
+# --- API Endpoint ---
+@app.post("/v1/chat/completions")
+async def chat_completions(request: ChatCompletionRequest):
+    if request.model != MODEL_NAME:
+        raise HTTPException(status_code=400, detail=f"Model not supported. Please use '{MODEL_NAME}'.")
 
-    if user_role.lower() != 'user':
-        raise HTTPException(status_code=400, detail="The last message must be from the 'user'.")
+    user_query = request.messages[-1].content if request.messages else ""
+    if not user_query or request.messages[-1].role.lower() != 'user':
+        raise HTTPException(status_code=400, detail="The last message must be from the 'user' and contain content.")
 
     # 1. Perform Web Search
     search_results = await perform_web_search(user_query)
     formatted_results = format_search_results_for_prompt(search_results)
-
-    # 2. Construct the prompt for the external LLM
-    final_prompt = f"User Query: {user_query}\n\n{formatted_results}"
-
-    # 3. Call the external Inference API
-    headers = {
-        "Authorization": f"Bearer {INFERENCE_API_KEY}",
-        "Content-Type": "application/json",
-    }
-
-    # The payload for the external API uses our system prompt and the combined user query + search results
+    
+    # 2. Construct prompt for the backend model
+    final_user_prompt = f"User's question: \"{user_query}\"\n\nBased ONLY on the provided search results below, answer the user's question.\n\n{formatted_results}"
+    
+    # 3. Prepare payload for Inference API
     payload = {
-        "model": "meta-llama/llama-3.1-8b-instruct/fp-8", # The actual model used by the inference API
+        "model": BACKEND_MODEL,
         "messages": [
             {"role": "system", "content": SYSTEM_PROMPT},
-            {"role": "user", "content": final_prompt},
+            {"role": "user", "content": final_user_prompt},
         ],
         "max_tokens": request.max_tokens,
         "temperature": request.temperature,
+        "stream": request.stream,
     }
 
-    async with httpx.AsyncClient(timeout=60.0) as client:
-        try:
-            response = await client.post(INFERENCE_API_URL, json=payload, headers=headers)
-            response.raise_for_status()
-            model_response = response.json()
-
-            # 4. Format the response to be OpenAI API compliant
-            # This part depends on the exact structure of the inference API's response
-            # Assuming it's similar to OpenAI's, we extract the message content
-            generated_content = model_response["choices"][0]["message"]["content"]
-
-            api_response = {
-                "id": model_response.get("id", "chatcmpl-binglity-lite-123"),
-                "object": "chat.completion",
-                "created": model_response.get("created", 0),
-                "model": "Binglity-Lite",
-                "choices": [{
-                    "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "content": generated_content,
-                    },
-                    "finish_reason": "stop",
-                }],
-                "usage": model_response.get("usage", {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}),
-            }
-            return api_response
-
-        except httpx.HTTPStatusError as e:
-            raise HTTPException(status_code=e.response.status_code, detail=f"Error from inference API: {e.response.text}")
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}")
-
+    # 4. Handle streaming or single response
+    if request.stream:
+        return StreamingResponse(stream_response_generator(payload), media_type="text/event-stream")
+    else:
+        # Standard non-streaming request
+        headers = {"Authorization": f"Bearer {INFERENCE_API_KEY}"}
+        async with httpx.AsyncClient(timeout=120.0) as client:
+            try:
+                response = await client.post(INFERENCE_API_URL, json=payload, headers=headers)
+                response.raise_for_status()
+                model_response = response.json()
+
+                # Format response to be OpenAI API compliant
+                return {
+                    "id": model_response.get("id", f"chatcmpl-{uuid.uuid4()}"),
+                    "object": "chat.completion",
+                    "created": model_response.get("created", int(time.time())),
+                    "model": MODEL_NAME,
+                    "choices": [{
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": model_response["choices"][0]["message"]["content"],
+                        },
+                        "finish_reason": "stop",
+                    }],
+                    "usage": model_response.get("usage", {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}),
+                }
+            except httpx.HTTPStatusError as e:
+                raise HTTPException(status_code=e.response.status_code, detail=f"Error from inference API: {e.response.text}")
 
 @app.get("/")
 def read_root():