Update app.py

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from typing import List, Optional
 from llama_cpp import Llama
@@ -10,6 +10,7 @@ app = FastAPI()
 
 llm = None
 
+# Models
 class Message(BaseModel):
     role: str
     content: str
@@ -20,6 +21,17 @@ class ChatRequest(BaseModel):
     temperature: Optional[float] = 0.7
     max_tokens: Optional[int] = 256
 
+class ModelInfo(BaseModel):
+    id: str
+    name: str
+    description: str
+
+# Load your models info here or dynamically from disk/config
+AVAILABLE_MODELS = [
+    ModelInfo(id="llama2", name="Llama 2", description="Meta Llama 2 model"),
+    # Add more models if you want
+]
+
 @app.on_event("startup")
 def load_model():
     global llm
@@ -32,13 +44,21 @@ def load_model():
         raise RuntimeError(f"Model not found at path: {model_path}")
     llm = Llama(model_path=model_path)
 
-@app.get("/")
-async def root():
-    return {"message": "API is running"}
+@app.get("/health")
+async def health_check():
+    return {"status": "ok"}
 
-@app.get("/api/tags")
-async def api_tags():
-    return []
+@app.get("/models")
+async def list_models():
+    # Return available models info
+    return [model.dict() for model in AVAILABLE_MODELS]
+
+@app.get("/models/{model_id}")
+async def get_model(model_id: str):
+    for model in AVAILABLE_MODELS:
+        if model.id == model_id:
+            return model.dict()
+    raise HTTPException(status_code=404, detail="Model not found")
 
 @app.post("/chat")
 async def chat(req: ChatRequest):
@@ -46,7 +66,11 @@ async def chat(req: ChatRequest):
     if llm is None:
         return {"error": "Model not initialized."}
 
-    # Build prompt from messages, Ollama uses system/user/assistant roles
+    # Validate model - simple check
+    if req.model not in [m.id for m in AVAILABLE_MODELS]:
+        raise HTTPException(status_code=400, detail="Unsupported model")
+
+    # Construct prompt from messages
     prompt = ""
     for m in req.messages:
         prompt += f"{m.role}: {m.content}\n"