Update app.py

app.py

@@ -1,47 +1,61 @@
-import os
-import subprocess
-
-from fastapi import FastAPI, Request, HTTPException
-import uvicorn
-
-app = FastAPI()
-
-# Load your API key from environment (set this in HF Secrets for security)
-OLLAMA_API_KEY = os.environ.get("OLLAMA_API_KEY", "change_me")
-
-@app.post("/generate")
-async def generate(request: Request):
-    """Endpoint that generates text based on the prompt."""
-    # 1. Check API key
-    auth_header = request.headers.get("Authorization")
-    if not auth_header or not auth_header.startswith("Bearer "):
-        raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
-
-    token = auth_header.split(" ")[1]
-    if token != OLLAMA_API_KEY:
-        raise HTTPException(status_code=401, detail="Invalid API key")
-
-    # 2. Parse the request JSON
-    body = await request.json()
-    prompt_text = body.get("prompt", "")
-    if not prompt_text:
-        raise HTTPException(status_code=400, detail="No prompt provided")
-
-    # 3. Option A: Call Ollama via its CLI directly
-    #    (The model is served in the background via `ollama serve`)
-    process = subprocess.Popen(
-        ["ollama", "run", prompt_text],
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
-        text=True
-    )
-    output, error = process.communicate()
-
-    if process.returncode != 0 or error:
-        raise HTTPException(status_code=500, detail=f"Ollama error: {error.strip()}")
-
-    return {"response": output.strip()}
-
-
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+import os
+import subprocess
+
+from fastapi import FastAPI, Request, HTTPException
+import uvicorn
+
+app = FastAPI()
+
+# Load your API key from the environment (defaults to "change_me")
+OLLAMA_API_KEY = os.environ.get("OLLAMA_API_KEY", "change_me")
+
+@app.post("/generate")
+async def generate(request: Request):
+    """Endpoint that generates text based on the prompt."""
+    # 1. Check API key
+    auth_header = request.headers.get("Authorization")
+    if not auth_header or not auth_header.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
+
+    token = auth_header.split(" ")[1]
+    if token != OLLAMA_API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+
+    # 2. Parse the request JSON
+    body = await request.json()
+    prompt_text = body.get("prompt", "")
+    if not prompt_text:
+        raise HTTPException(status_code=400, detail="No prompt provided")
+
+    # 3. Call Ollama via CLI (example: using a HF model)
+    #    If you want a default model, you can use just "ollama run" with the prompt.
+    #    If you want a specific HF model, specify it here:
+    #        "ollama run hf.co/abanm/Dubs-Q8_0-GGUF prompt_text"
+    #    For now, let's illustrate passing two args: model + the user prompt
+    command = [
+        "ollama",
+        "run",
+        prompt_text  # or "hf.co/abanm/Dubs-Q8_0-GGUF", prompt_text
+    ]
+
+    process = subprocess.Popen(
+        command,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True
+    )
+    output, error = process.communicate()
+
+    if process.returncode != 0 or error:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Ollama error: {error.strip()}"
+        )
+
+    return {"response": output.strip()}
+
+
+if __name__ == "__main__":
+    # For local debugging (outside Docker)
+    uvicorn.run(app, host="0.0.0.0", port=7860)
+