Update app.py

app.py

@@ -4,9 +4,11 @@ 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
 
-app = FastAPI(title="My Local Brains API")
+app = FastAPI(title="Xtime GGUF Remote API")
 
+# Настройка CORS для удаленного подключения
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -15,106 +17,87 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Папка, где будут лежать ваши .gguf файлы
-MODELS_DIR = "./models"
-os.makedirs(MODELS_DIR, exist_ok=True)
-
+# Глобальные переменные для хранения текущей модели в памяти
 current_llm = None
-current_model_name = ""
+current_model_id = "" # format: repo_id/filename
 
-def get_local_models():
-    """Возвращает список всех .gguf файлов в папке models"""
-    return [f for f in os.listdir(MODELS_DIR) if f.endswith('.gguf')]
+class ChatRequest(BaseModel):
+    repo_id: str          # Ссылка на репозиторий (напр. "bartowski/Llama-3.2-3B-Instruct-GGUF")
+    filename: str         # Имя файла (напр. "Llama-3.2-3B-Instruct-Q4_K_M.gguf")
+    prompt: str           # Текст пользователя
+    system_prompt: str = "You are a helpful assistant."
+    max_tokens: int = 512
+    temperature: float = 0.7
 
-def load_model(model_filename: str):
-    global current_llm, current_model_name
+def load_model_if_new(repo_id: str, filename: str):
+    """Загружает модель, если она еще не в памяти или если пришла новая ссылка"""
+    global current_llm, current_model_id
     
-    model_path = os.path.join(MODELS_DIR, model_filename)
+    new_model_id = f"{repo_id}/{filename}"
     
-    if not os.path.exists(model_path):
-        raise HTTPException(status_code=404, detail=f"Модель {model_filename} не найдена в папке {MODELS_DIR}")
-
-    if current_model_name == model_filename and current_llm is not None:
-        return # Модель уже загружена
+    # Если модель уже загружена, просто выходим
+    if current_llm is not None and current_model_id == new_model_id:
+        return
 
-    print(f"--- Загрузка мозга: {model_filename} ---")
-
-    # Освобождаем память от предыдущей модели
+    print(f"--- Загрузка новой модели: {new_model_id} ---")
+    
+    # Очистка памяти перед загрузкой новой модели
     if current_llm is not None:
         del current_llm
         gc.collect()
 
     try:
+        # Скачивание файла с Hugging Face (использует кэш, если файл уже есть)
+        model_path = hf_hub_download(repo_id=repo_id, filename=filename)
+        
+        # Инициализация Llama
         current_llm = Llama(
             model_path=model_path,
-            n_ctx=4096,              # Размер контекста (можно увеличить до 8192, если хватает памяти)
+            n_ctx=2048,
             n_threads=os.cpu_count() or 4,
-            n_gpu_layers=-1,         # -1 означает выгрузку всех возможных слоев на видеокарту (GPU)
+            n_gpu_layers=0, # Установите > 0, если у вас есть GPU
             verbose=False
         )
-        current_model_name = model_filename
-        print(f"✅ Мозг '{model_filename}' успешно подключен!")
+        current_model_id = new_model_id
+        print(f"✅ Модель {filename} успешно загружена и готова")
     except Exception as e:
-        print(f"❌ Ошибка загрузки {model_filename}: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-class ChatRequest(BaseModel):
-    prompt: 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
-    }
+        print(f"❌ Ошибка при загрузке модели: {e}")
+        raise HTTPException(status_code=500, detail=f"Failed to load model: {str(e)}")
 
 @app.post("/chat")
 async def chat(request: ChatRequest):
-    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)
+    """Эндпоинт для чата, который сам переключает модели"""
+    # 1. Проверяем/загружаем модель
+    load_model_if_new(request.repo_id, request.filename)
 
     try:
-        # Универсальный шаблон промпта (System + User). 
-        # Если вы используете модели формата Llama-3 или ChatML, шаблон можно поменять.
-        prompt = f"System: {request.system_prompt}\nUser: {request.prompt}\nAssistant:"
+        # 2. Генерация ответа (базовый формат, подходящий для большинства моделей)
+        formatted_prompt = f"System: {request.system_prompt}\nUser: {request.prompt}\nAssistant:"
         
         output = current_llm.create_completion(
-            prompt=prompt,
+            prompt=formatted_prompt,
             max_tokens=request.max_tokens,
             temperature=request.temperature,
-            stop=["User:", "System:", "<|endoftext|>", "<|im_end|>"]
+            stop=["User:", "System:", "</s>", "<|endoftext|>"]
         )
         
-        response_text = output["choices"][0]["text"].strip()
         return {
-            "response": response_text,
-            "model_used": current_model_name
+            "response": output["choices"][0]["text"].strip(),
+            "model_id": current_model_id
         }
-
     except Exception as e:
-        print(f"Ошибка при генерации: {e}")
-        raise HTTPException(status_code=500, detail="Ошибка генерации ответа")
+        print(f"Ошибка генерации: {e}")
+        raise HTTPException(status_code=500, detail="Generation error")
 
-@app.get("/")
+@app.get("/health")
 async def health():
+    """Проверка состояния сервера"""
     return {
-        "status": "online", 
-        "active_brain": current_model_name,
-        "instruction": f"Положите ваши .gguf файлы в папку {os.path.abspath(MODELS_DIR)}"
-    }
+        "status": "online",
+        "current_model": current_model_id if current_model_id else "None"
+    }
+
+if __name__ == "__main__":
+    import uvicorn
+    # Запуск на порту 7860 (стандарт для HF Spaces)
+    uvicorn.run(app, host="0.0.0.0", port=7860)