Update main.py

main.py

@@ -1,8 +1,12 @@
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
-from typing import List, Optional, Literal, Dict, Any, Union
+from typing import List, Optional, Literal, Dict, Any, Union, Generator
 import uvicorn
+import time
+import asyncio
+import json
+import re
 from duckai import DuckAI
 
 app = FastAPI(title="DuckAI OpenAI Compatible API")
@@ -69,10 +73,64 @@ class DuckAIParser:
         """
         # Very simple estimation - about 4 characters per token on average
         return len(text) // 4
+    
+    @staticmethod
+    def stream_response(response_text: str, request_id: str, model: str) -> Generator[str, None, None]:
+        """
+        Stream the response with simulated typing effect
+        """
+        # Split by words to simulate streaming
+        words = re.findall(r'\S+\s*', response_text)
+        accumulated_text = ""
+        chunk_id = 0
+        
+        for word in words:
+            accumulated_text += word
+            
+            # Create the delta message structure (OpenAI compatible)
+            delta_data = {
+                "id": request_id,
+                "object": "chat.completion.chunk",
+                "created": int(time.time()),
+                "model": model,
+                "choices": [
+                    {
+                        "index": 0,
+                        "delta": {"content": word},
+                        "finish_reason": None
+                    }
+                ]
+            }
+            
+            yield f"data: {json.dumps(delta_data)}\n\n"
+            
+            # Sleep for 0.06 seconds between words
+            time.sleep(0.06)
+        
+        # Send the final message with finish_reason
+        end_data = {
+            "id": request_id,
+            "object": "chat.completion.chunk",
+            "created": int(time.time()),
+            "model": model,
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {},
+                    "finish_reason": "stop"
+                }
+            ]
+        }
+        
+        yield f"data: {json.dumps(end_data)}\n\n"
+        yield "data: [DONE]\n\n"
 
-@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
-async def create_chat_completion(request: ChatCompletionRequest):
+@app.post("/v1/chat/completions")
+async def create_chat_completion(request: ChatCompletionRequest, response: Response):
     try:
+        # Generate a request ID
+        request_id = f"chatcmpl-duck-{hash(str(request.messages)) % 10000}"
+        
         # Parse the conversation history into DuckAI format
         conversation_text = DuckAIParser.parse_conversation_history(request.messages)
         
@@ -80,17 +138,26 @@ async def create_chat_completion(request: ChatCompletionRequest):
         duck_ai = DuckAI()
         result = duck_ai.chat(conversation_text, model=request.model)
         
-        # Extract the assistant's response
+        # Clean up the response
         assistant_response = result.strip()
         
         # Estimate token usage
         prompt_tokens = DuckAIParser.estimate_tokens(conversation_text)
         completion_tokens = DuckAIParser.estimate_tokens(assistant_response)
         
-        # Create OpenAI-compatible response
-        response = ChatCompletionResponse(
-            id=f"chatcmpl-duck-{hash(conversation_text) % 10000}",
-            created=int(__import__('time').time()),
+        # Handle streaming if requested
+        if request.stream:
+            response.headers["Content-Type"] = "text/event-stream"
+            return DuckAIParser.stream_response(
+                assistant_response, 
+                request_id, 
+                request.model
+            )
+        
+        # Regular response (non-streaming)
+        return ChatCompletionResponse(
+            id=request_id,
+            created=int(time.time()),
             model=request.model,
             choices=[
                 ChatCompletionChoice(
@@ -106,8 +173,6 @@ async def create_chat_completion(request: ChatCompletionRequest):
                 total_tokens=prompt_tokens + completion_tokens
             )
         )
-        
-        return response
     
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -115,7 +180,7 @@ async def create_chat_completion(request: ChatCompletionRequest):
 @app.get("/v1/models")
 async def list_models():
     """Return a list of available models"""
-    current_time = int(__import__('time').time())
+    current_time = int(time.time())
     return {
         "object": "list",
         "data": [
@@ -152,10 +217,10 @@ async def list_models():
         ]
     }
 
-# Adding a simple root endpoint for health check
-@app.get("/")
-async def root():
-    return {"status": "ok", "message": "DuckAI OpenAI Compatible API is running"}
+# Health check endpoint
+@app.get("/health")
+async def health_check():
+    return {"status": "ok"}
 
 if __name__ == "__main__":
     uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)