Update main.py

main.py

@@ -1,90 +1,129 @@
+import os
 import json
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, FileResponse
+from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 from llama_cpp import Llama
 from huggingface_hub import hf_hub_download
+from tavily import TavilyClient
 
-app = FastAPI(title="Fast Gemma API")
+app = FastAPI(title="Qwen Turbo Search API")
 
 # --- КОНФИГУРАЦИЯ ---
+TAVILY_API_KEY = os.getenv("TAVILY_API_KEY") 
+
 REPO_ID = "bartowski/Qwen2.5-1.5B-Instruct-GGUF"
-# Используем Q4_K_M - это золотая середина скорости и ума
 FILENAME = "Qwen2.5-1.5B-Instruct-Q6_K.gguf"
 
 llm = None
+tavily_client = None
 
+# --- ИНИЦИАЛИЗАЦИЯ ---
 @app.on_event("startup")
 def startup_event():
-    global llm
+    global llm, tavily_client
+    
+    if TAVILY_API_KEY:
+        tavily_client = TavilyClient(api_key=TAVILY_API_KEY)
+        print("✅ Tavily Search подключен")
+    else:
+        print("⚠️ Нет TAVILY_API_KEY. Поиск работать не будет.")
+
     print("🚀 Загрузка модели...")
     try:
-        model_path = hf_hub_download(
-            repo_id=REPO_ID,
-            filename=FILENAME,
-            cache_dir="./models"
-        )
-        
-        # --- ОПТИМИЗАЦИЯ ЗАГРУЗКИ ---
+        model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME, cache_dir="./models")
         llm = Llama(
             model_path=model_path,
-            n_ctx=4096,         # Увеличим контекст
-            n_threads=2,        # Лимит Hugging Face
-            n_batch=1024,       # <--- УСКОРЕНИЕ: Читаем промпт большими кусками
-            verbose=False       # Меньше мусора в логах
+            n_ctx=8192,
+            n_threads=2,
+            n_batch=1024,
+            verbose=False
         )
-        print("✅ Модель готова и ускорена!")
+        print("✅ Модель готова!")
     except Exception as e:
         print(f"❌ Ошибка: {e}")
 
-# --- МОДЕЛИ ДАННЫХ ---
+# --- ПОДКЛЮЧАЕМ ИНТЕРФЕЙС ---
+# Создай папку static рядом с main.py!
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+@app.get("/")
+def read_root():
+    # Отдаем наш HTML файл при входе на главную
+    return FileResponse('static/index.html')
+
+# --- ЛОГИКА ПОИСКА ---
+def perform_search(query: str):
+    if not tavily_client: return "Нет ключа Tavily.", []
+    print(f"🔎 Ищу: {query}")
+    try:
+        res = tavily_client.search(query=query, search_depth="advanced", max_results=5)
+        text = ""
+        sources = []
+        for i, r in enumerate(res['results']):
+            idx = i + 1
+            text += f"ИСТОЧНИК [{idx}]: {r['title']}\nТЕКСТ: {r['content']}\n\n"
+            sources.append({"id": idx, "title": r['title'], "url": r['url']})
+        return text, sources
+    except Exception as e:
+        print(f"Err: {e}")
+        return "Ошибка поиска.", []
+
+# --- API ---
 class Message(BaseModel):
     role: str
     content: str
 
 class ChatRequest(BaseModel):
     messages: list[Message]
-    temperature: float = 0.7
-    max_tokens: int = 500
-    stream: bool = False  # <--- Добавили флаг стриминга
+    temperature: float = 0.6
+    max_tokens: int = 2048
+    stream: bool = True
+    use_search: bool = False
 
-# --- ЭНДПОИНТ ---
 @app.post("/v1/chat/completions")
-def chat_completions(request: ChatRequest):
-    if not llm:
-        raise HTTPException(status_code=503, detail="Модель грузится...")
-
-    # Подготовка сообщений
-    messages_payload = [{"role": m.role, "content": m.content} for m in request.messages]
-
-    # === РЕЖИМ 1: STREAMING (МГНОВЕННЫЙ ОТВЕТ) ===
-    if request.stream:
-        def iter_response():
-            # Запрашиваем поток у модели
-            stream_gen = llm.create_chat_completion(
-                messages=messages_payload,
-                temperature=request.temperature,
-                max_tokens=request.max_tokens,
-                stream=True  # Включаем стрим в движке
-            )
-            # Читаем генератор по кусочкам
-            for chunk in stream_gen:
-                # Формируем формат Server-Sent Events (как у OpenAI)
-                yield f"data: {json.dumps(chunk)}\n\n"
-            yield "data: [DONE]\n\n"
-
-        return StreamingResponse(iter_response(), media_type="text/event-stream")
-
-    # === РЕЖИМ 2: ОБЫЧНЫЙ (ЖДЕМ ВЕСЬ ТЕКСТ) ===
+def chat_completions(req: ChatRequest):
+    if not llm: raise HTTPException(503, "Loading...")
+    
+    msgs = [{"role": m.role, "content": m.content} for m in req.messages]
+    
+    # Поиск
+    if req.use_search:
+        query = msgs[-1]['content']
+        context, sources = perform_search(query)
+        
+        sys_prompt = (
+            "Ты умный помощник. Отвечай на вопрос, используя ТОЛЬКО эти данные из интернета.\n"
+            "Обязательно указывай источники [1], [2].\n"
+            f"=== ДАННЫЕ ===\n{context}"
+        )
+        # Добавляем источники в конец последнего сообщения (для UI)
+        sources_md = "\n\n**Источники:**\n" + "\n".join([f"{s['id']}. [{s['title']}]({s['url']})" for s in sources])
+        
+        # Инъекция системного промпта
+        msgs.insert(0, {"role": "system", "content": sys_prompt})
     else:
-        response = llm.create_chat_completion(
-            messages=messages_payload,
-            temperature=request.temperature,
-            max_tokens=request.max_tokens,
-            stream=False
+        sources_md = ""
+
+    # Генерация
+    def iter_response():
+        stream = llm.create_chat_completion(
+            messages=msgs,
+            temperature=req.temperature,
+            max_tokens=req.max_tokens,
+            stream=True
         )
-        return response
+        for chunk in stream:
+            yield f"data: {json.dumps(chunk)}\n\n"
+        
+        # Если были источники, отправим их отдельным чанком в конце
+        if sources_md:
+            final_chunk = {
+                "choices": [{"delta": {"content": sources_md}, "finish_reason": None}]
+            }
+            yield f"data: {json.dumps(final_chunk)}\n\n"
+            
+        yield "data: [DONE]\n\n"
 
-@app.get("/")
-def home():
-    return {"status": "running", "optimization": "enabled"}
+    return StreamingResponse(iter_response(), media_type="text/event-stream")