Update app.py

app.py

@@ -4,13 +4,15 @@ from src.utils.hf_packager import HFSpacePackager
 from agents.planner import plan_task
 from agents.executor import execute_step
 from agents.critic import review_result
-from memory import init_memory, add_to_memory, search_memory, get_memory
+import time, logging
+from fastapi import FastAPI, Request, HTTPException
+from pydantic import BaseModel, constr
+from memory import init_memory, add_to_memory, search_memory, get_memory, get_summary
 import uuid
 from datetime import datetime
 import os
 import psutil
 import json
-import time
 import subprocess
 from sentence_transformers import SentenceTransformer
 import numpy as np
@@ -18,12 +20,12 @@ import faiss
 from fastapi import FastAPI
 from ctransformers import AutoModelForCausalLM
 from src.core.cognitive_engine import CognitiveEngine
-from pydantic import BaseModel
-from memory import add_to_memory, get_memory
+
+logging.basicConfig(filename="app.log", level=logging.INFO)
 
 class GenerateRequest(BaseModel):
-    prompt: str
-    session_id: str
+    prompt: constr(min_length=5, max_length=1000)
+    session_id: constr(min_length=3, max_length=50)
 
 # Load LLM model
 llm_model = AutoModelForCausalLM.from_pretrained(
@@ -52,9 +54,21 @@ def status():
 
 @app.post("/generate")
 def generate_post(data: GenerateRequest):
-    response = llm_model(data.prompt)
-    add_to_memory(data.session_id, data.prompt, response)
-    return {"session_id": data.session_id, "response": response}
+    start = time.time()
+    try:
+        response = llm_model(data.prompt)
+        latency = round(time.time() - start, 2)
+
+        add_to_memory(data.session_id, data.prompt, response)
+        logging.info(f"[{data.session_id}] Prompt processed in {latency}s")
+
+        return {
+            "session_id": data.session_id,
+            "response": response,
+            "latency": latency
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/memory")
 def view_memory():
@@ -64,6 +78,18 @@ def view_memory():
 def session_memory(session_id: str):
     return {"session_id": session_id, "history": get_memory(session_id)}
 
+@app.get("/summary/{session_id}")
+def memory_summary(session_id: str):
+    return {
+        "session_id": session_id,
+        "summary": get_summary(session_id)
+    }
+    
+@app.middleware("http")
+async def auth_middleware(request: Request, call_next):
+    if request.url.path.startswith("/generate") and request.headers.get("X-API-Key") != "your-secret":
+        return JSONResponse(status_code=401, content={"detail": "Invalid API Key"})
+    return await call_next(request)
 
 # Test the model at startup (optional)
 if __name__ == "__main__":