Update main.py

main.py

@@ -1,25 +1,22 @@
 import httpx
+import json
+import time
+import uuid
 from fastapi import FastAPI, HTTPException
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, Field
-from typing import List
-
-app = FastAPI(
-    title="Perplexity-like API",
-    description="An API that uses web search to answer questions with citations.",
-    version="1.0.0"
-)
+from typing import List, Optional
 
 # --- API Configuration ---
+# It's recommended to use environment variables for sensitive data in production.
 TYPEGPT_API_URL = "https://api.typegpt.net/v1/chat/completions"
-TYPEGPT_API_KEY = "sk-oPdaZC7n1JlDq0sJ5NSSyHe7sYaeAXeEuj0wX4Lk8hlOGPF8"
+TYPEGPT_API_KEY = "sk-oPdaZC7n1JlDq0sJ5NSSyHe7sYaeAXeEuj0wX4Lk8hlOGPF8" # Replace with your actual key
 SEARCH_API_URL = "https://superapis-bing.hf.space/search"
 
-
 # --- System Prompt ---
 # This prompt guides the AI to behave like a factual research assistant.
 SYSTEM_PROMPT = """
 You are an expert AI research assistant. Your primary goal is to provide accurate, comprehensive, and helpful answers based ONLY on the provided search results.
-
 Instructions:
 1.  Carefully analyze the user's query and the provided search results.
 2.  Synthesize an answer directly from the information found in the search results.
@@ -30,93 +27,163 @@ Instructions:
 7.  Structure your response in a clear and easy-to-read format. Start with a direct answer, followed by a more detailed explanation.
 """
 
-# --- Pydantic Models for API Request/Response ---
+# --- Pydantic Models ---
 
+# For incoming requests
 class ChatMessage(BaseModel):
     role: str
     content: str
 
 class ChatCompletionRequest(BaseModel):
     messages: List[ChatMessage] = Field(..., example=[{"role": "user", "content": "What are the benefits of learning Python?"}])
-    model: str = "gpt-4.1-mini" # Model is fixed but included for compatibility
+    model: str = "perplexity-like" # Model name can be customized
+    stream: bool = Field(default=False, description="Enable streaming response")
+
+# For outgoing streaming responses (OpenAI compatible)
+class ChatDelta(BaseModel):
+    content: Optional[str] = None
+    role: Optional[str] = None
+
+class ChatCompletionStreamChoice(BaseModel):
+    delta: ChatDelta
+    index: int = 0
+    finish_reason: Optional[str] = None
+
+class ChatCompletionStreamResponse(BaseModel):
+    id: str = Field(default_factory=lambda: f"chatcmpl-{uuid.uuid4().hex}")
+    object: str = "chat.completion.chunk"
+    created: int = Field(default_factory=lambda: int(time.time()))
+    model: str = "perplexity-like"
+    choices: List[ChatCompletionStreamChoice]
+
+# --- FastAPI App Initialization ---
+app = FastAPI(
+    title="Perplexity-like API",
+    description="An API that uses web search to answer questions with citations, supporting streaming.",
+    version="2.0.0"
+)
 
-class Choice(BaseModel):
-    message: ChatMessage
 
-class ChatCompletionResponse(BaseModel):
-    choices: List[Choice]
-    search_results: List[dict]
+# --- Streaming Logic ---
+async def stream_llm_response(payload: dict):
+    """
+    An async generator that streams the response from the language model.
+    """
+    start_time = time.time()
+    try:
+        async with httpx.AsyncClient(timeout=60.0) as client:
+            headers = {
+                "Authorization": f"Bearer {TYPEGPT_API_KEY}",
+                "Content-Type": "application/json"
+            }
+            async with client.stream("POST", TYPEGPT_API_URL, headers=headers, json=payload) as response:
+                # Check for errors from the upstream API
+                if response.status_code != 200:
+                    error_content = await response.aread()
+                    raise HTTPException(
+                        status_code=response.status_code,
+                        detail=f"Error from language model API: {error_content.decode()}"
+                    )
+
+                # Process the stream line by line
+                async for line in response.aiter_lines():
+                    if line.startswith("data: "):
+                        data_str = line.removeprefix("data: ")
+                        if data_str.strip() == "[DONE]":
+                            break
+                        try:
+                            chunk = json.loads(data_str)
+                            delta_content = chunk["choices"][0]["delta"].get("content")
+                            if delta_content:
+                                # Create a streaming-compliant response chunk
+                                stream_choice = ChatCompletionStreamChoice(delta=ChatDelta(content=delta_content))
+                                stream_response = ChatCompletionStreamResponse(choices=[stream_choice])
+                                yield f"data: {stream_response.model_dump_json()}\n\n"
+                        except (json.JSONDecodeError, KeyError, IndexError):
+                            # Skip malformed lines
+                            continue
+    
+    except httpx.RequestError as e:
+        # Handle network-related errors during the streaming request
+        error_message = f"HTTP Request Error during streaming: {e}"
+        stream_choice = ChatCompletionStreamChoice(delta=ChatDelta(content=f"\n\nERROR: {error_message}"))
+        stream_response = ChatCompletionStreamResponse(choices=[stream_choice])
+        yield f"data: {stream_response.model_dump_json()}\n\n"
+    
+    except Exception as e:
+        # Handle other unexpected errors
+        error_message = f"An unexpected error occurred during streaming: {e}"
+        stream_choice = ChatCompletionStreamChoice(delta=ChatDelta(content=f"\n\nERROR: {error_message}"))
+        stream_response = ChatCompletionStreamResponse(choices=[stream_choice])
+        yield f"data: {stream_response.model_dump_json()}\n\n"
+
+    # Send the final chunk with finish_reason
+    finally:
+        finish_time = time.time()
+        print(f"Stream finished in {finish_time - start_time:.2f} seconds.")
+        final_choice = ChatCompletionStreamChoice(delta=ChatDelta(), finish_reason="stop")
+        final_response = ChatCompletionStreamResponse(choices=[final_choice])
+        yield f"data: {final_response.model_dump_json()}\n\n"
+        yield "data: [DONE]\n\n"
 
 
 # --- API Endpoint ---
-
-@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
+@app.post("/v1/chat/completions")
 async def chat_completions(request: ChatCompletionRequest):
     """
-    Takes a user's chat history, performs a web search based on the latest query,
-    and uses the TypeGPT model to generate a factual, cited response.
+    Takes a user's query, performs a web search, and streams a factual,
+    cited response from a language model.
     """
     if not request.messages or request.messages[-1].role != "user":
         raise HTTPException(status_code=400, detail="Invalid request. The last message must be from the 'user'.")
 
     user_query = request.messages[-1].content
 
-    async with httpx.AsyncClient(timeout=30.0) as client:
-        # 1. Perform a web search
-        try:
-            search_params = {"keywords": user_query}
+    # 1. Perform a web search
+    try:
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            search_params = {"keywords": user_query, "max_results": 7}
             search_response = await client.get(SEARCH_API_URL, params=search_params)
             search_response.raise_for_status()
             search_results = search_response.json()
-        except httpx.RequestError as e:
-            raise HTTPException(status_code=502, detail=f"Error calling the search API: {e}")
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=f"Failed to process search results: {e}")
-
-        # 2. Format search results into a context for the language model
-        context = ""
-        for i, result in enumerate(search_results[:7]):  # Use top 7 results for richer context
-            context += f"Source [{i+1}]:\nTitle: {result.get('title', 'N/A')}\nSnippet: {result.get('snippet', '')}\nURL: {result.get('url', 'N/A')}\n\n"
-
-        # 3. Construct the prompt for the language model
-        final_prompt = f"""
-        **Search Results:**
-        {context}
-
-        **User Query:** "{user_query}"
+    except httpx.RequestError as e:
+        raise HTTPException(status_code=502, detail=f"Error calling the search API: {e}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to process search results: {e}")
+
+    # 2. Format search results into a context for the language model
+    # Using the 'description' field as per the new OpenAPI spec
+    context = ""
+    for i, result in enumerate(search_results):
+        context += f"Source [{i+1}]:\nTitle: {result.get('title', 'N/A')}\nDescription: {result.get('description', '')}\nURL: {result.get('url', 'N/A')}\n\n"
+
+    # 3. Construct the prompt for the language model
+    final_prompt = f"""
+    **Search Results:**
+    {context}
+    **User Query:** "{user_query}"
+    Please provide a comprehensive answer based on the search results above, following all instructions.
+    """
 
-        Please provide a comprehensive answer based on the search results above, following all instructions.
-        """
+    # 4. Prepare the payload for the TypeGPT language model
+    llm_payload = {
+        "model": "gpt-4.1-mini",
+        "messages": [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": final_prompt}
+        ],
+        "stream": True  # Enable streaming from the backing LLM
+    }
 
-        # 4. Get the response from the TypeGPT language model
-        try:
-            headers = {
-                "Authorization": f"Bearer {TYPEGPT_API_KEY}",
-                "Content-Type": "application/json"
-            }
-            # The payload now includes the system prompt and the user prompt with context
-            payload = {
-                "model": "gpt-4.1-mini",
-                "messages": [
-                    {"role": "system", "content": SYSTEM_PROMPT},
-                    {"role": "user", "content": final_prompt}
-                ]
-            }
-            llm_response = await client.post(TYPEGPT_API_URL, headers=headers, json=payload)
-            llm_response.raise_for_status()
-            llm_data = llm_response.json()
-            answer_content = llm_data['choices'][0]['message']['content']
-        except httpx.RequestError as e:
-            raise HTTPException(status_code=502, detail=f"Error calling language model API: {e}")
-        except (KeyError, IndexError) as e:
-            raise HTTPException(status_code=500, detail=f"Invalid response structure from language model API: {e}")
-
-    # 5. Format the final response
-    response_message = ChatMessage(role="assistant", content=answer_content)
-    response_choice = Choice(message=response_message)
-    
-    return ChatCompletionResponse(choices=[response_choice], search_results=search_results)
+    # 5. Return the streaming response
+    return StreamingResponse(stream_llm_response(llm_payload), media_type="text/event-stream")
 
+# --- Main execution ---
 if __name__ == "__main__":
     import uvicorn
+    # To run this app:
+    # 1. Save the code as main.py
+    # 2. Install necessary packages: pip install fastapi "uvicorn[standard]" httpx
+    # 3. Run in your terminal: uvicorn main:app --reload
+    # 4. Access the interactive docs at http://127.0.0.1:8000/docs
     uvicorn.run(app, host="0.0.0.0", port=8000)