updating frontend

main.py

@@ -1,106 +1,326 @@
-import uuid
-import json
-import asyncio
-import time
+# import uuid
+# import json
+# import asyncio
+# import time
+# import os
+# import sys
+# from contextlib import asynccontextmanager
+# from loguru import logger
+# from fastapi import FastAPI, HTTPException, status, Response
+# from fastapi.middleware.cors import CORSMiddleware
+# from fastapi.staticfiles import StaticFiles
+# from fastapi.responses import HTMLResponse
+
+# # Import your existing schemas (Ensure schemas.py is in the same folder)
+# from schemas import ChatRequest, ChatResponse
+
+# # -------------------------------------------------
+# # 1. Loguru Configuration
+# # -------------------------------------------------
+# logger.remove()
+# logger.add(sys.stdout, format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level}</level> | <cyan>{extra[session_id]}</cyan> - {message}")
+# logger = logger.bind(session_id="SYSTEM")
+
+# # -------------------------------------------------
+# # 2. AI Logic (Replacing the MQTT Worker)
+# # -------------------------------------------------
+# # We define a direct function instead of publishing to MQTT
+# async def get_ai_response(question: str):
+#     """
+#     Replace this with your actual agent logic (e.g., LangChain or Groq).
+#     This simulates what your 'worker' used to do.
+#     """
+#     # Simulate processing time
+#     await asyncio.sleep(1) 
+#     return {
+#         "answer": f"I am your SmartCoffee assistant. You asked: {question}",
+#         "sources": ["knowledge_base_v1"],
+#         "timestamp": time.time()
+#     }
+
+# # -------------------------------------------------
+# # 3. App Lifespan
+# # -------------------------------------------------
+# @asynccontextmanager
+# async def lifespan(app: FastAPI):
+#     logger.info("Starting AI Agent on Hugging Face...")
+#     yield
+#     logger.info("Shutting down...")
+
+# # -------------------------------------------------
+# # 4. App Init
+# # -------------------------------------------------
+# app = FastAPI(title="SmartCoffee AI 2026", lifespan=lifespan)
+
+# # Allow CORS for local testing, though HF uses same-origin
+# app.add_middleware(
+#     CORSMiddleware,
+#     allow_origins=["*"],
+#     allow_methods=["*"],
+#     allow_headers=["*"],
+# )
+
+# # --- CRITICAL: Mount Static Files ---
+# # This serves your index.html, CSS, and JS
+# app.mount("/static", StaticFiles(directory="static"), name="static")
+
+# # -------------------------------------------------
+# # 5. Routes
+# # -------------------------------------------------
+
+# @app.get("/", response_class=HTMLResponse)
+# async def serve_frontend():
+#     """Serves the main chat interface"""
+#     try:
+#         with open("static/index.html", "r", encoding="utf-8") as f:
+#             return HTMLResponse(content=f.read())
+#     except FileNotFoundError:
+#         return HTMLResponse(content="<h1>index.html not found in /static</h1>", status_code=404)
+
+# @app.post("/api/v1/chat", response_model=ChatResponse)
+# async def chat(request: ChatRequest):
+#     if request.session_id == "default":
+#         request.session_id = f"hf_{uuid.uuid4().hex[:12]}"
+
+#     request_logger = logger.bind(session_id=request.session_id)
+#     request_logger.info(f"Processing request: {request.question}")
+
+#     try:
+#         # Instead of MQTT publish, call logic directly
+#         response = await get_ai_response(request.question)
+        
+#         request_logger.success("Response generated.")
+#         return ChatResponse(
+#             question=request.question,
+#             answer=response["answer"],
+#             sources=response.get("sources", []),
+#             session_id=request.session_id,
+#             timestamp=response.get("timestamp", time.time()),
+#         )
+#     except Exception as e:
+#         request_logger.error(f"Error: {str(e)}")
+#         raise HTTPException(status_code=500, detail="Internal AI Error")
+
+# @app.get("/health")
+# async def health():
+#     return {"status": "healthy", "platform": "Hugging Face"}
+
+
+
+
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.responses import HTMLResponse, RedirectResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel, Field, field_validator, validator
 import os
-import sys
+import re
+import time
+import uuid
 from contextlib import asynccontextmanager
-from loguru import logger
-from fastapi import FastAPI, HTTPException, status, Response
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.staticfiles import StaticFiles
-from fastapi.responses import HTMLResponse
-
-# Import your existing schemas (Ensure schemas.py is in the same folder)
-from schemas import ChatRequest, ChatResponse
-
-# -------------------------------------------------
-# 1. Loguru Configuration
-# -------------------------------------------------
-logger.remove()
-logger.add(sys.stdout, format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level}</level> | <cyan>{extra[session_id]}</cyan> - {message}")
-logger = logger.bind(session_id="SYSTEM")
-
-# -------------------------------------------------
-# 2. AI Logic (Replacing the MQTT Worker)
-# -------------------------------------------------
-# We define a direct function instead of publishing to MQTT
-async def get_ai_response(question: str):
-    """
-    Replace this with your actual agent logic (e.g., LangChain or Groq).
-    This simulates what your 'worker' used to do.
-    """
-    # Simulate processing time
-    await asyncio.sleep(1) 
-    return {
-        "answer": f"I am your SmartCoffee assistant. You asked: {question}",
-        "sources": ["knowledge_base_v1"],
-        "timestamp": time.time()
-    }
+import logging
+
+# Logging setup
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Space-specific: Use mounted dataset path
+KB_PATH = "/data/knowledge_base"
+
+# Groq client setup
+from groq import Groq
+client = Groq(api_key=os.getenv("GROQ_API_KEY"))
 
-# -------------------------------------------------
-# 3. App Lifespan
-# -------------------------------------------------
+# Space hardware: CPU-basic, limit memory
+MAX_SESSIONS = 50  # Lower for free tier
+
+# Lifespan for startup/shutdown
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    logger.info("Starting AI Agent on Hugging Face...")
+    logger.info("🚀 Starting up agent...")
+    # Load knowledge base here
+    await load_knowledge_base()
     yield
-    logger.info("Shutting down...")
-
-# -------------------------------------------------
-# 4. App Init
-# -------------------------------------------------
-app = FastAPI(title="SmartCoffee AI 2026", lifespan=lifespan)
-
-# Allow CORS for local testing, though HF uses same-origin
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
+    logger.info("🔌 Shutting down agent...")
+
+app = FastAPI(
+    title="SmartCoffee AI Agent",
+    description="AI Support Agent - Hugging Face Spaces Edition",
+    version="1.0.0",
+    lifespan=lifespan
 )
 
-# --- CRITICAL: Mount Static Files ---
-# This serves your index.html, CSS, and JS
-app.mount("/static", StaticFiles(directory="static"), name="static")
+# Mount static files (CSS/JS)
+app.mount("/static", StaticFiles(directory="."), name="static")
 
-# -------------------------------------------------
-# 5. Routes
-# -------------------------------------------------
+# Pydantic models
+class ChatRequest(BaseModel):
+    question: str = Field(..., min_length=3, max_length=300)
+    session_id: str = Field(default="default", pattern=r"^[a-zA-Z0-9_-]+$")
+    
+    question: str
 
-@app.get("/", response_class=HTMLResponse)
-async def serve_frontend():
-    """Serves the main chat interface"""
+    @field_validator('question')
+    @classmethod
+    def sanitize_input(cls, v: str) -> str:
+        # Standardize whitespace and strip
+        v = re.sub(r'\s+', ' ', v).strip()
+        
+        # Security check for prompt injection keywords
+        forbidden_keywords = ['ignore', 'system', 'admin', 'prompt']
+        if any(word in v.lower() for word in forbidden_keywords):
+            raise ValueError("Invalid input pattern")
+            
+        return v
+
+# In-memory session store (no Redis in free tier)
+sessions = {}
+
+async def load_knowledge_base():
+    """Load knowledge base from HF dataset at startup"""
+    from datasets import load_dataset
+    
+    logger.info("📚 Loading knowledge base...")
     try:
-        with open("static/index.html", "r", encoding="utf-8") as f:
-            return HTMLResponse(content=f.read())
-    except FileNotFoundError:
-        return HTMLResponse(content="<h1>index.html not found in /static</h1>", status_code=404)
+        dataset = load_dataset("YOUR_USERNAME/smartcoffee-kb", split="train")
+        # Process into text chunks
+        global knowledge_docs
+        knowledge_docs = [doc["text"] for doc in dataset]
+        logger.info(f"✅ Loaded {len(knowledge_docs)} documents")
+    except Exception as e:
+        logger.error(f"❌ Failed to load KB: {e}")
+        knowledge_docs = []
 
-@app.post("/api/v1/chat", response_model=ChatResponse)
-async def chat(request: ChatRequest):
-    if request.session_id == "default":
-        request.session_id = f"hf_{uuid.uuid4().hex[:12]}"
+# RAG function
+def rag_query(question: str) -> str:
+    from langchain_huggingface import HuggingFaceEmbeddings
+    from sklearn.metrics.pairwise import cosine_similarity
+    import numpy as np
+    
+    if not knowledge_docs:
+        return "Knowledge base not loaded."
+    
+    # Simple TF-IDF search (memory-efficient)
+    from sklearn.feature_extraction.text import TfidfVectorizer
+    
+    vectorizer = TfidfVectorizer(max_features=1000, stop_words='english')
+    doc_vectors = vectorizer.fit_transform(knowledge_docs)
+    question_vec = vectorizer.transform([question])
+    
+    # Get top 2 most similar docs
+    similarities = cosine_similarity(question_vec, doc_vectors).flatten()
+    top_indices = np.argsort(similarities)[-2:]
+    
+    context = "\n\n".join([knowledge_docs[i] for i in top_indices])
+    return context
 
-    request_logger = logger.bind(session_id=request.session_id)
-    request_logger.info(f"Processing request: {request.question}")
+# LLM call
+def generate_response(question: str, context: str, session_id: str) -> dict:
+    start_time = time.time()
+    
+    prompt = f"""You are SmartCoffee Support AI. Use ONLY this context:
 
+Context:
+{context}
+
+Question: {question}
+
+Answer concisely in 2-3 sentences. If unsure, say "I need to check with my team."
+
+Answer:"""
+    
     try:
-        # Instead of MQTT publish, call logic directly
-        response = await get_ai_response(request.question)
-        
-        request_logger.success("Response generated.")
-        return ChatResponse(
-            question=request.question,
-            answer=response["answer"],
-            sources=response.get("sources", []),
-            session_id=request.session_id,
-            timestamp=response.get("timestamp", time.time()),
+        response = client.chat.completions.create(
+            model="llama3-8b-8192",
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=200,
+            temperature=0.1
         )
+        
+        latency = time.time() - start_time
+        
+        return {
+            "answer": response.choices[0].message.content,
+            "latency": latency,
+            "tokens_in": response.usage.prompt_tokens,
+            "tokens_out": response.usage.completion_tokens,
+            "model": "groq-llama3-8b",
+            "sources": [f"doc_{i}" for i in range(2)]
+        }
+        
     except Exception as e:
-        request_logger.error(f"Error: {str(e)}")
-        raise HTTPException(status_code=500, detail="Internal AI Error")
+        logger.error(f"LLM error: {e}")
+        return {
+            "answer": "Sorry, I'm having trouble processing your request.",
+            "latency": time.time() - start_time,
+            "error": str(e)
+        }
+
+# Routes
+@app.get("/", response_class=HTMLResponse)
+async def serve_frontend():
+    """Serve the combined frontend"""
+    with open("index.html", "r", encoding="utf-8") as f:
+        return HTMLResponse(content=f.read())
+
+@app.post("/api/v1/chat")
+async def chat(request: ChatRequest):
+    try:
+        # Get session memory
+        session = sessions.get(request.session_id, {
+            "history": [],
+            "created_at": time.time()
+        })
+        
+        # Clean up old sessions
+        if len(sessions) > MAX_SESSIONS:
+            oldest = min(sessions, key=lambda k: sessions[k]["created_at"])
+            del sessions[oldest]
+        
+        # Add user message to history
+        session["history"].append({"role": "user", "content": request.question})
+        
+        # RAG query
+        context = rag_query(request.question)
+        
+        # Generate response
+        result = generate_response(request.question, context, request.session_id)
+        
+        # Add bot message to history
+        session["history"].append({"role": "bot", "content": result["answer"]})
+        sessions[request.session_id] = session
+        
+        return {
+            "question": request.question,
+            "answer": result["answer"],
+            "sources": result.get("sources", []),
+            "session_id": request.session_id,
+            "latency_ms": int(result["latency"] * 1000)
+        }
+        
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    except Exception as e:
+        logger.error(f"Unexpected error: {e}")
+        raise HTTPException(status_code=500, detail="Failed to process request")
 
 @app.get("/health")
 async def health():
-    return {"status": "healthy", "platform": "Hugging Face"}
+    return {
+        "status": "operational",
+        "sessions_active": len(sessions),
+        "kb_loaded": len(knowledge_docs) if 'knowledge_docs' in globals() else 0
+    }
+
+@app.get("/api/v1/metrics")
+async def metrics():
+    """Simple metrics endpoint"""
+    return {
+        "total_requests": sum(len(s.get("history", [])) for s in sessions.values()) // 2,
+        "active_sessions": len(sessions),
+        "uptime_seconds": int(time.time() - app.state.startup_time)
+    }
+
+
+@app.get("/")
+async def root():
+    return {"message": "Agent is running", "uptime": time.time() - app.state.startup_time}

requirements.txt

@@ -30,3 +30,4 @@ prometheus-client==0.23.1
 
 #loguru
 loguru==0.7.3
+datasets