Update main.py

main.py

@@ -1,6 +1,6 @@
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 import httpx
 import os
 import json
@@ -13,23 +13,27 @@ from typing import List, Dict, Any, Optional, AsyncGenerator
 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"
-NEWS_API_URL = "https://rkihacker-brave.hf.space/news" # Added News API URL
+NEWS_API_URL = "https://rkihacker-brave.hf.space/news"
+IMAGE_API_URL = "https://rkihacker-brave.hf.space/images" # Added Image API URL
 MODEL_NAME = "Binglity-Lite"
 BACKEND_MODEL = "meta-llama/llama-3.1-8b-instruct/fp-8"
 
 # --- Final Advanced System Prompt ---
 SYSTEM_PROMPT = """
-You are "Binglity-Lite", a highly advanced AI search assistant. Your purpose is to provide users with accurate, comprehensive, and trustworthy answers by synthesizing information from a given set of web and news search results.
+You are "Binglity-Lite", a highly advanced AI search assistant. Your purpose is to provide users with accurate, comprehensive, and trustworthy answers by synthesizing information from a given set of web, news, and image search results.
+
 **Core Directives:**
-1.  **Answer Directly**: Immediately address the user's question. **Do not** use introductory phrases like "Based on the search results..." or "Here is the information I found...". Your tone should be confident, objective, and encyclopedic.
+1.  **Answer Directly**: Immediately address the user's question. **Do not** use introductory phrases like "Based on the search results...". Your tone should be confident, objective, and encyclopedic.
 2.  **Synthesize, Don't Summarize**: Your primary task is to weave information from multiple sources into a single, cohesive, and well-structured answer. Do not simply describe what each source says one by one.
 3.  **Cite with Inline Markdown Links**: This is your most important instruction. When you present a fact or a piece of information from a source, you **must** cite it immediately using an inline Markdown link.
     *   **Format**: The format must be `[phrase or sentence containing the fact](URL)`. The URL must come from the `URL:` field of the provided source.
     *   **Example**: If a source with URL `https://example.com/science` says "The Earth is the third planet from the Sun", your output should be: "The Earth is the [third planet from the Sun](https://example.com/science)."
     *   **Rule**: Every piece of information in your answer must be attributable to a source via these inline links.
-4.  **Be Fact-Based**: Your entire response must be based **exclusively** on the information provided in the web and news search results. Do not use any outside knowledge.
-5.  **Filter for Relevance**: If a search result is not relevant to the user's query, ignore it completely. Do not mention it in your response.
-6.  **Handle Ambiguity**: If the search results are contradictory or insufficient to answer the question fully, state this clearly in your response, citing the conflicting sources.
+4.  **Be Fact-Based**: Your entire response must be based **exclusively** on the information provided in the search results. Do not use any outside knowledge.
+5.  **Interpret Image Results**: For image search results, use the title and context to describe the image if it's relevant to the user's query. Cite the source page URL.
+6.  **Filter for Relevance**: If a search result is not relevant to the user's query, ignore it completely. Do not mention it in your response.
+7.  **Handle Ambiguity**: If the search results are contradictory or insufficient to answer the question fully, state this clearly in your response, citing the conflicting sources.
+
 **Final Output Structure:**
 Your final response MUST be structured in two parts:
 1.  **The Synthesized Answer**: A well-written response that directly answers the user's query, with facts and statements properly cited using inline Markdown links as described above.
@@ -41,8 +45,8 @@ Your final response MUST be structured in two parts:
 # --- FastAPI App ---
 app = FastAPI(
     title="Binglity-Lite API",
-    description="A web search-powered, streaming-capable chat completions API.",
-    version="1.3.0", # Version updated
+    description="A web, news, and image search-powered, streaming-capable chat completions API.",
+    version="1.4.0",
 )
 
 # --- Pydantic Models for OpenAI Compatibility ---
@@ -57,49 +61,22 @@ class ChatCompletionRequest(BaseModel):
     temperature: Optional[float] = 0.7
     stream: Optional[bool] = False
 
-# --- Web Search Functions ---
-async def perform_web_search(query: str) -> List[Dict[str, Any]]:
-    async with httpx.AsyncClient() as client:
-        try:
-            response = await client.get(
-                SEARCH_API_URL,
-                params={"query": query, "max_results": 10}
-            )
-            response.raise_for_status()
-            results = response.json()
-            # Add source type to each result
-            for result in results:
-                result['source_type'] = 'Web'
-            return results
-        except httpx.HTTPStatusError as e:
-            print(f"Error from search API: {e.response.text}")
-            return []
-        except Exception as e:
-            print(f"An unexpected error occurred during web search: {str(e)}")
-            return []
-
-async def perform_news_search(query: str) -> List[Dict[str, Any]]:
-    """Performs a search against the news API."""
-    async with httpx.AsyncClient() as client:
-        try:
-            # Parameters can be adjusted as needed, e.g., region
-            response = await client.get(
-                NEWS_API_URL,
-                params={"query": query, "max_results": 10, "region": "en-US"}
-            )
-            response.raise_for_status()
-            results = response.json()
-            # Add source type to each result
-            for result in results:
-                result['source_type'] = 'News'
-            return results
-        except httpx.HTTPStatusError as e:
-            print(f"Error from news API: {e.response.text}")
-            return []
-        except Exception as e:
-            print(f"An unexpected error occurred during news search: {str(e)}")
-            return []
-
+# --- Search Functions ---
+async def perform_search(client: httpx.AsyncClient, url: str, query: str, source_type: str) -> List[Dict[str, Any]]:
+    """Generic function to perform a search against a given API."""
+    try:
+        response = await client.get(url, params={"query": query, "max_results": 10})
+        response.raise_for_status()
+        results = response.json()
+        for result in results:
+            result['source_type'] = source_type
+        return results
+    except httpx.HTTPStatusError as e:
+        print(f"Error from {source_type} API: {e.response.text}")
+        return []
+    except Exception as e:
+        print(f"An unexpected error occurred during {source_type} search: {str(e)}")
+        return []
 
 def format_search_results_for_prompt(results: List[Dict[str, Any]]) -> str:
     """Formats combined search results for the language model prompt."""
@@ -108,15 +85,22 @@ def format_search_results_for_prompt(results: List[Dict[str, Any]]) -> str:
     
     formatted = "### Search Results ###\n\n"
     for i, result in enumerate(results):
-        source_type = result.get('source_type', 'Search') # Default in case it's missing
+        source_type = result.get('source_type', 'Search')
         formatted += f"Source [{i+1}] ({source_type}):\n"
         formatted += f"Title: {result.get('title', 'N/A')}\n"
         formatted += f"URL: {result.get('url', 'N/A')}\n"
-        formatted += f"Content: {result.get('description', 'N/A')}\n\n"
+        
+        if source_type == 'Image':
+            formatted += f"Content: [Image Result] A picture titled '{result.get('title', 'N/A')}'\n"
+            formatted += f"Image URL: {result.get('image', 'N/A')}\n\n"
+        else:
+            formatted += f"Content: {result.get('description', 'N/A')}\n\n"
+            
     return formatted
 
 # --- Streaming Logic ---
 async def stream_response_generator(payload: Dict[str, Any]) -> AsyncGenerator[str, None]:
+    """Generates server-sent events for streaming responses."""
     headers = {
         "Authorization": f"Bearer {INFERENCE_API_KEY}",
         "Content-Type": "application/json",
@@ -133,7 +117,7 @@ async def stream_response_generator(payload: Dict[str, Any]) -> AsyncGenerator[s
 
             async for line in response.aiter_lines():
                 if line.startswith("data:"):
-                    line_data = line[5:].strip()
+                    line_data = line[len("data:"):].strip()
                     if line_data == "[DONE]":
                         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"
@@ -141,12 +125,18 @@ async def stream_response_generator(payload: Dict[str, Any]) -> AsyncGenerator[s
                     
                     try:
                         chunk = json.loads(line_data)
-                        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:
+                        # **ERROR FIX**: Check if 'choices' exists and is not empty before accessing
+                        if chunk.get("choices") and len(chunk["choices"]) > 0:
+                            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, IndexError):
                         continue
 
 # --- API Endpoint ---
@@ -159,24 +149,28 @@ async def chat_completions(request: ChatCompletionRequest):
     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.")
 
-    # Perform web and news searches concurrently
-    web_results, news_results = await asyncio.gather(
-        perform_web_search(user_query),
-        perform_news_search(user_query)
-    )
+    # Perform all searches concurrently
+    async with httpx.AsyncClient() as client:
+        search_tasks = [
+            perform_search(client, SEARCH_API_URL, user_query, "Web"),
+            perform_search(client, NEWS_API_URL, user_query, "News"),
+            perform_search(client, IMAGE_API_URL, user_query, "Image"),
+        ]
+        all_results = await asyncio.gather(*search_tasks)
 
     # Combine results and remove duplicates by URL
     combined_results = []
     seen_urls = set()
-    for result in web_results + news_results:
-        url = result.get('url')
-        if url and url not in seen_urls:
-            combined_results.append(result)
-            seen_urls.add(url)
+    for result_list in all_results:
+        for result in result_list:
+            url = result.get('url')
+            if url and url not in seen_urls:
+                combined_results.append(result)
+                seen_urls.add(url)
     
     formatted_results = format_search_results_for_prompt(combined_results)
     
-    final_user_prompt = f"User's question: \"{user_query}\"\n\nUse the web and news search results below to answer the user's question. Follow all rules in your system prompt exactly.\n\n{formatted_results}"
+    final_user_prompt = f"User's question: \"{user_query}\"\n\nUse the web, news, and image search results below to answer the user's question. Follow all rules in your system prompt exactly.\n\n{formatted_results}"
     
     payload = {
         "model": BACKEND_MODEL,
@@ -184,7 +178,9 @@ async def chat_completions(request: ChatCompletionRequest):
             {"role": "system", "content": SYSTEM_PROMPT},
             {"role": "user", "content": final_user_prompt},
         ],
-        "max_tokens": request.max_tokens, "temperature": request.temperature, "stream": request.stream,
+        "max_tokens": request.max_tokens,
+        "temperature": request.temperature,
+        "stream": request.stream,
     }
 
     if request.stream:
@@ -196,6 +192,11 @@ async def chat_completions(request: ChatCompletionRequest):
                 response = await client.post(INFERENCE_API_URL, json=payload, headers=headers)
                 response.raise_for_status()
                 model_response = response.json()
+                
+                # Ensure the response structure is valid before returning
+                if not model_response.get("choices") or len(model_response["choices"]) == 0:
+                     raise HTTPException(status_code=500, detail="Invalid response from inference API: 'choices' field is missing or empty.")
+
                 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",}],