Update main.py

main.py

@@ -1,14 +1,13 @@
 import os
 import httpx
 import json
-from fastapi import FastAPI, Request, HTTPException
+from fastapi import FastAPI, Request
 from fastapi.responses import StreamingResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from qdrant_client import QdrantClient
 
 app = FastAPI()
 
-# 1. Настройка CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -19,113 +18,150 @@ app.add_middleware(
 
 COLLECTION_NAME = "equipment_registry"
 
-# Функция для эмбеддингов
-async def get_openrouter_embedding(text: str, api_key: str):
+
+async def get_embedding(text: str, api_key: str) -> list[float]:
+    """Получение эмбеддинга через OpenRouter"""
     async with httpx.AsyncClient() as client:
-        payload = {
-            "model": "openai/text-embedding-3-small",
-            "input": text,
-            "encoding_format": "float"
-        }
-        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
-        response = await client.post("https://openrouter.ai/api/v1/embeddings", headers=headers, json=payload, timeout=30.0)
+        response = await client.post(
+            "https://openrouter.ai/api/v1/embeddings",
+            headers={
+                "Authorization": f"Bearer {api_key}",
+                "Content-Type": "application/json"
+            },
+            json={
+                "model": "openai/text-embedding-3-small",
+                "input": text,
+                "encoding_format": "float"
+            },
+            timeout=30.0
+        )
         if response.status_code != 200:
             raise Exception(f"Embedding Error: {response.text}")
         return response.json()["data"][0]["embedding"]
 
+
+def search_qdrant(query_vector: list[float], url: str, api_key: str) -> str:
+    """Поиск контекста в Qdrant"""
+    client = QdrantClient(url=url, api_key=api_key)
+    results = client.search(
+        collection_name=COLLECTION_NAME,
+        query_vector=query_vector,
+        limit=5
+    )
+    contexts = [
+        res.payload.get("search_text", "")
+        for res in results
+        if res.payload
+    ]
+    return "\n\n".join(contexts)
+
+
 @app.post("/api/chat")
 async def chat(request: Request):
-    # ПРАВИЛЬНЫЕ ИМЕНА ИЗ ТВОИХ SECRETS
+    # Переменные окружения
     OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")
     QDRANT_URL = os.getenv("QDRANT_URL")
     QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
 
-    if not OPENROUTER_API_KEY or not QDRANT_URL or not QDRANT_API_KEY:
-        return JSONResponse(content={"error": "Missing environment variables on HF Space"}, status_code=500)
+    if not all([OPENROUTER_API_KEY, QDRANT_URL, QDRANT_API_KEY]):
+        return JSONResponse(
+            content={"error": "Missing environment variables"},
+            status_code=500
+        )
 
     try:
         body = await request.json()
         messages = body.get("messages", [])
+        
         if not messages:
             return JSONResponse(content={"error": "No messages"}, status_code=400)
 
         user_query = messages[-1].get("content", "")
 
-        # 1. Поиск в Qdrant
-        vector = await get_openrouter_embedding(user_query, OPENROUTER_API_KEY)
+        # 1. Получаем эмбеддинг запроса
+        query_vector = await get_embedding(user_query, OPENROUTER_API_KEY)
 
-        q_client = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
-        search_results = q_client.search(
-            collection_name=COLLECTION_NAME, 
-            query_vector=vector, 
-            limit=5
-        )
+        # 2. Ищем контекст в Qdrant
+        context = search_qdrant(query_vector, QDRANT_URL, QDRANT_API_KEY)
 
-        context = "\n\n".join([res.payload.get("search_text", "") for res in search_results if res.payload])
+        # 3. Системный промпт с контекстом
+        system_prompt = f"""Ты — инженерный ИИ-ассистент по реестру научного оборудования.
+Используй предоставленный контекст для точного ответа.
+Если данных нет в контексте, вежливо скажи об этом.
 
-        system_prompt = (
-            "Ты — инженерный ИИ-ассистент по реестру научного оборудования. "
-            "Используй предоставленный контекст для точного ответа. "
-            "Если данных нет, вежливо скажи об этом.\n\n"
-            f"КОНТЕКСТ:\n{context}"
-        )
+КОНТЕКСТ:
+{context}"""
 
-        # 2. Стриминг ответа (DeepSeek-V3 через OpenRouter)
-        async def generate():
-            payload = {
-                "model": "deepseek/deepseek-chat",
-                "messages": [{"role": "system", "content": system_prompt}] + messages,
-                "stream": True
-            }
-            headers = {
-                "Authorization": f"Bearer {OPENROUTER_API_KEY}", 
-                "Content-Type": "application/json",
-                "HTTP-Referer": "https://huggingface.co/", 
-                "X-OpenRouter-Title": "Inventory Assistant"
-            }
+        # 4. Формируем сообщения для LLM
+        llm_messages = [{"role": "system", "content": system_prompt}] + messages
 
+        # 5. Стриминг ответа
+        async def generate():
             async with httpx.AsyncClient(timeout=120.0) as client:
-                async with client.stream("POST", "https://openrouter.ai/api/v1/chat/completions", headers=headers, json=payload) as response:
+                async with client.stream(
+                    "POST",
+                    "https://openrouter.ai/api/v1/chat/completions",
+                    headers={
+                        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
+                        "Content-Type": "application/json",
+                        "HTTP-Referer": "https://huggingface.co/",
+                        "X-Title": "Equipment Assistant"
+                    },
+                    json={
+                        "model": "deepseek/deepseek-chat",
+                        "messages": llm_messages,
+                        "stream": True
+                    }
+                ) as response:
                     if response.status_code != 200:
-                        error_text = await response.aread()
-                        yield f"OpenRouter Error: {error_text.decode()}"
+                        error = await response.aread()
+                        yield f"Error: {error.decode()}"
                         return
 
                     async for line in response.aiter_lines():
                         if not line or not line.startswith("data: "):
                             continue
-                        
-                        data_str = line[6:].strip()
-                        if data_str == "[DONE]":
+
+                        data = line[6:].strip()
+                        if data == "[DONE]":
                             break
-                            
+
                         try:
-                            chunk = json.loads(data_str)
-                            if not chunk.get('choices'):
-                                continue
-                                
-                            content = chunk['choices'][0].get('delta', {}).get('content', '')
+                            chunk = json.loads(data)
+                            content = (
+                                chunk.get("choices", [{}])[0]
+                                .get("delta", {})
+                                .get("content", "")
+                            )
                             if content:
-                                # Отдаем чистый текст для useChat
                                 yield content
-                        except:
+                        except json.JSONDecodeError:
                             continue
 
         return StreamingResponse(
-            generate(), 
-            media_type="text/plain",
+            generate(),
+            media_type="text/event-stream",
             headers={
                 "Cache-Control": "no-cache",
-                "X-Accel-Buffering": "no" # Чтобы HF не тормозил поток
+                "Connection": "keep-alive",
+                "X-Accel-Buffering": "no"
             }
         )
 
     except Exception as e:
-        print(f"Chat Error: {str(e)}")
-        return JSONResponse(content={"detail": str(e)}, status_code=500)
+        print(f"Chat Error: {e}")
+        return JSONResponse(content={"error": str(e)}, status_code=500)
+
 
 @app.get("/")
 async def health():
-    # Проверка, видит ли сервер ключи
-    check = "OK" if os.getenv("OPENROUTER_API_KEY") else "MISSING"
-    return {"status": "running", "api_key": check}
+    """Проверка здоровья сервиса"""
+    has_keys = all([
+        os.getenv("OPENROUTER_API_KEY"),
+        os.getenv("QDRANT_URL"),
+        os.getenv("QDRANT_API_KEY")
+    ])
+    return {
+        "status": "running",
+        "keys_configured": has_keys
+    }