Update app.py

app.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 import os
 from dotenv import load_dotenv
 import requests
@@ -8,6 +8,7 @@ from typing import Dict, Any, List
 from pydantic import BaseModel
 import time
 import json
+import asyncio
 
 load_dotenv()
 
@@ -32,6 +33,7 @@ class ChatCompletionRequest(BaseModel):
     model: str
     messages: List[ChatMessage]
     temperature: float = 0.7
+    stream: bool = False
 
 def count_tokens(text: str) -> int:
     # Используем тот же алгоритм, что и в прямом API
@@ -47,6 +49,37 @@ def clean_assistant_response(text: str) -> str:
         text = text[:-3].strip()
     return text
 
+async def stream_response(response_text: str):
+    # Разбиваем текст на части для стриминга
+    words = response_text.split()
+    for i in range(0, len(words), 2):
+        chunk = " ".join(words[i:i+2]) + " "
+        yield f"data: {json.dumps({
+            'id': f'chatcmpl-{os.urandom(12).hex()}',
+            'object': 'chat.completion.chunk',
+            'created': int(time.time()),
+            'model': 'phi4-r1',
+            'choices': [{
+                'index': 0,
+                'delta': {'content': chunk},
+                'finish_reason': None
+            }]
+        }, ensure_ascii=False)}\n\n"
+        await asyncio.sleep(0.1)  # Небольшая задержка между чанками
+    
+    # Отправляем финальное сообщение
+    yield f"data: {json.dumps({
+        'id': f'chatcmpl-{os.urandom(12).hex()}',
+        'object': 'chat.completion.chunk',
+        'created': int(time.time()),
+        'model': 'phi4-r1',
+        'choices': [{
+            'index': 0,
+            'delta': {},
+            'finish_reason': 'stop'
+        }]
+    }, ensure_ascii=False)}\n\n"
+
 @app.get("/")
 async def root():
     response = JSONResponse({"status": "FastFlowWrapper is running"})
@@ -109,6 +142,17 @@ async def create_chat_completion(request: ChatCompletionRequest):
         flowise_response = response.json()
         assistant_response = clean_assistant_response(flowise_response.get("text", ""))
         
+        # Если запрошен стриминг
+        if request.stream:
+            return StreamingResponse(
+                stream_response(assistant_response),
+                media_type="text/event-stream"
+            )
+        
+        # Подсчитываем токены
+        prompt_tokens = count_tokens(last_message.content)
+        completion_tokens = count_tokens(assistant_response)
+        
         response = JSONResponse({
             "id": "chatcmpl-" + os.urandom(12).hex(),
             "object": "chat.completion",
@@ -126,8 +170,9 @@ async def create_chat_completion(request: ChatCompletionRequest):
                 }
             ],
             "usage": {
-                "completion_tokens": 0,  # Устанавливаем 0, так как это не важно
-                "total_tokens": 0  # Устанавливаем 0, так как это не важно
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": prompt_tokens + completion_tokens
             },
             "stats": {},
             "system_fingerprint": "phi4-r1"