Update app/app.py

app/app.py

@@ -1,45 +1,57 @@
-from fastapi import FastAPI, Request
-from pydantic import BaseModel
 import os
 import logging
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 from app.policy_vector_db import PolicyVectorDB, ensure_db_populated
-from llama_cpp import Llama
 
-app = FastAPI()
+# --- GGUF Model Config ---
+MODEL_PATH = "dop-phi-1.5-Q4_K_M.gguf"
+MODEL_URL = (
+    "https://huggingface.co/Kalpokoch/QuantizedFineTunedPhi1.5/resolve/main/dop-phi-1.5-Q4_K_M.gguf"
+)
+
+# Logging setup
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("app")
 
-# --- Load Vector DB ---
-VECTOR_DB_DIR = "/app/vector_database"
-CHUNKS_FILE = "processed_chunks.json"
-vector_db = PolicyVectorDB(persist_directory=VECTOR_DB_DIR)
-ensure_db_populated(vector_db, CHUNKS_FILE)
-
-# --- Load Quantized GGUF Model ---
-MODEL_PATH = "dop-phi-1.5-Q4_K_M.gguf"
+# --- Download model if not present ---
+def maybe_download_model():
+    if not os.path.exists(MODEL_PATH):
+        logger.info("Downloading GGUF model from Hugging Face...")
+        import requests
+
+        try:
+            with requests.get(MODEL_URL, stream=True) as r:
+                r.raise_for_status()
+                with open(MODEL_PATH, "wb") as f:
+                    for chunk in r.iter_content(chunk_size=8192):
+                        if chunk:
+                            f.write(chunk)
+            logger.info("Model downloaded successfully.")
+        except Exception as e:
+            logger.error(f"Failed to download GGUF model: {e}")
+            raise
+
+maybe_download_model()
+
+# Initialize FastAPI
+app = FastAPI()
 
-llm = Llama(
-    model_path=MODEL_PATH,
-    n_ctx=2048,
-    n_threads=os.cpu_count(),
-    use_mlock=False
-)
+# Initialize vector DB
+vector_db = PolicyVectorDB(persist_directory="vector_database", relevance_threshold=0.0)
+ensure_db_populated(vector_db, "processed_chunks.json")
 
+# Request schema
 class QueryRequest(BaseModel):
     query: str
+    top_k: int = 5
 
+# API endpoint
 @app.post("/query")
-async def query_policy(request: QueryRequest):
-    query_text = request.query
-    chunks = vector_db.search(query_text, top_k=5)
-
-    if not chunks:
-        return {"answer": "No relevant information found.", "chunks": []}
-
-    context = "\n\n".join([c['text'] for c in chunks])
-    prompt = f"Answer the question based on the context below:\n\nContext:\n{context}\n\nQuestion: {query_text}\nAnswer:"
-
-    response = llm(prompt, max_tokens=512, stop=["\n\n", "###"], echo=False)
-    answer = response["choices"][0]["text"].strip()
-
-    return {"answer": answer, "chunks": chunks}
+async def query_vector_db(req: QueryRequest):
+    try:
+        results = vector_db.search(query_text=req.query, top_k=req.top_k)
+        return {"results": results}
+    except Exception as e:
+        logger.error(f"Query error: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail="Internal server error")