Update app.py

app.py

@@ -1,12 +1,11 @@
+import os
+import gc
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from llama_cpp import Llama
-from huggingface_hub import hf_hub_download
-import gc
-import os
 
-app = FastAPI(title="Xtime AI API")
+app = FastAPI(title="My Local Brains API")
 
 app.add_middleware(
     CORSMiddleware,
@@ -16,85 +15,106 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-REPO_ID = "Rid3/xtime-v1beta-gguf-storage"
+# Папка, где будут лежать ваши .gguf файлы
+MODELS_DIR = "./models"
+os.makedirs(MODELS_DIR, exist_ok=True)
 
 current_llm = None
 current_model_name = ""
 
-MODELS = {
-    "medium": "xtime-v1beta-n-m_1p.gguf",
-    "small":  "xtime-v1beta-xp-r_2.gguf",
-    "large":  "xtime-v1beta-q4_K_M.gguf"
-}
+def get_local_models():
+    """Возвращает список всех .gguf файлов в папке models"""
+    return [f for f in os.listdir(MODELS_DIR) if f.endswith('.gguf')]
 
-def load_model(model_key: str):
+def load_model(model_filename: str):
     global current_llm, current_model_name
-    filename = MODELS.get(model_key)
-    if not filename:
-        raise HTTPException(status_code=404, detail="Model not found")
+    
+    model_path = os.path.join(MODELS_DIR, model_filename)
+    
+    if not os.path.exists(model_path):
+        raise HTTPException(status_code=404, detail=f"Модель {model_filename} не найдена в папке {MODELS_DIR}")
 
-    if current_model_name == model_key and current_llm is not None:
-        return
+    if current_model_name == model_filename and current_llm is not None:
+        return # Модель уже загружена
 
-    print(f"--- Loading model: {filename} ({model_key}) ---")
+    print(f"--- Загрузка мозга: {model_filename} ---")
 
+    # Освобождаем память от предыдущей модели
     if current_llm is not None:
         del current_llm
         gc.collect()
 
     try:
-        model_path = hf_hub_download(repo_id=REPO_ID, filename=filename)
-
         current_llm = Llama(
             model_path=model_path,
-            n_ctx=2048,
+            n_ctx=4096,              # Размер контекста (можно увеличить до 8192, если хватает памяти)
             n_threads=os.cpu_count() or 4,
-            n_gpu_layers=0,
-            verbose=False,
-            chat_format=None   # важно! для Phi-2 не используем llama-3
+            n_gpu_layers=-1,         # -1 означает выгрузку всех возможных слоев на видеокарту (GPU)
+            verbose=False
         )
-        current_model_name = model_key
-        print(f"✅ Модель {model_key} успешно загружена")
+        current_model_name = model_filename
+        print(f"✅ Мозг '{model_filename}' успешно подключен!")
     except Exception as e:
-        print(f"❌ Ошибка загрузки {model_key}: {e}")
+        print(f"❌ Ошибка загрузки {model_filename}: {e}")
         raise HTTPException(status_code=500, detail=str(e))
 
-
-@app.on_event("startup")
-async def startup_event():
-    load_model("medium")   # по умолчанию самая стабильная
-
-
 class ChatRequest(BaseModel):
     prompt: str
-    model_type: str = "medium"
-    api_key: str = ""
-
+    model: str = "" # Имя файла, например: "my_brain_v1.gguf". Если пусто, возьмет первую доступную.
+    system_prompt: str = "Ты полезный, умный ИИ-ассистент."
+    max_tokens: int = 512
+    temperature: float = 0.7
+
+@app.get("/models")
+async def list_models():
+    """Посмотреть все доступные модели"""
+    return {
+        "available_models": get_local_models(), 
+        "current_loaded_model": current_model_name
+    }
 
 @app.post("/chat")
 async def chat(request: ChatRequest):
-    if request.model_type != current_model_name:
-        load_model(request.model_type)
+    global current_model_name
+    
+    # Определяем, какую модель грузить
+    target_model = request.model
+    if not target_model:
+        models = get_local_models()
+        if not models:
+             raise HTTPException(status_code=404, detail="В папке ./models нет ни одного .gguf файла!")
+        target_model = models[0] # Берем первую попавшуюся
+
+    # Загружаем, если еще не загружена
+    if target_model != current_model_name:
+        load_model(target_model)
 
     try:
-        # Для Phi-2 лучше использовать обычный create_completion
-        prompt = f"User: {request.prompt}\nAssistant:"
+        # Универсальный шаблон промпта (System + User). 
+        # Если вы используете модели формата Llama-3 или ChatML, шаблон можно поменять.
+        prompt = f"System: {request.system_prompt}\nUser: {request.prompt}\nAssistant:"
         
         output = current_llm.create_completion(
             prompt=prompt,
-            max_tokens=512,
-            temperature=0.7,
-            stop=["User:", "<|endoftext|>"]
+            max_tokens=request.max_tokens,
+            temperature=request.temperature,
+            stop=["User:", "System:", "<|endoftext|>", "<|im_end|>"]
         )
         
         response_text = output["choices"][0]["text"].strip()
-        return {"response": response_text}
+        return {
+            "response": response_text,
+            "model_used": current_model_name
+        }
 
     except Exception as e:
         print(f"Ошибка при генерации: {e}")
         raise HTTPException(status_code=500, detail="Ошибка генерации ответа")
 
-
 @app.get("/")
 async def health():
-    return {"status": "online", "model": current_model_name}
+    return {
+        "status": "online", 
+        "active_brain": current_model_name,
+        "instruction": f"Положите ваши .gguf файлы в папку {os.path.abspath(MODELS_DIR)}"
+    }