Update app.py

app.py

@@ -6,6 +6,8 @@ from langchain.prompts import PromptTemplate
 from supabase import create_client, Client
 from datetime import datetime
 from typing import List, Dict
+import json
+import uuid
 
 # Initialize FastAPI app
 app = FastAPI()
@@ -15,17 +17,22 @@ SUPABASE_URL = "https://ykkbxlbonywjmvbyfvwo.supabase.co"
 SUPABASE_KEY = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6Inlra2J4bGJvbnl3am12YnlmdndvIiwicm9sZSI6ImFub24iLCJpYXQiOjE3Mzk5NTA2NjIsImV4cCI6MjA1NTUyNjY2Mn0.2BZul_igHKmZtQGhbwV3PvRsCikxviL8ogTKPD3XhuU"
 supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)
 
-# Create a request model with context and user_id
+# Create a request model with context, user_id, and optional thread_id
 class SearchQuery(BaseModel):
     query: str
     context: str = None  # Optional context field
-    user_id: str  # Required to identify the user for storing history
+    user_id: str  # UUID string to identify the user for storing history
+    thread_id: str = None  # Optional thread_id to append to an existing thread
 
 # Create a response model for history
 class ConversationHistory(BaseModel):
+    id: str  # UUID as string
+    user_id: str  # UUID as string
     query: str
-    response: Dict
+    response: str  # Response stored as TEXT in the DB
     timestamp: str
+    thread_id: str  # UUID as string
+    title: str = None  # Optional title
 
 # Initialize LangChain with Groq
 llm = ChatGroq(
@@ -34,7 +41,7 @@ llm = ChatGroq(
     groq_api_key="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"  # Replace with your actual Groq API key
 )
 
-# Define prompt templates (keeping all for future flexibility, but defaulting to "general")
+# Define prompt templates
 prompt_templates = {
     "common_threats": PromptTemplate(
         input_variables=["query", "context"],
@@ -52,46 +59,97 @@ prompt_templates = {
         Please provide a detailed and professional response to the query based on your expertise in cybersecurity and the provided context.
         """
     ),
-    # You can keep other templates here if you want to manually select them later
 }
 
 # Initialize chains for each prompt
 chains = {key: LLMChain(llm=llm, prompt=prompt) for key, prompt in prompt_templates.items()}
 
-# Helper function to get conversation history for a user
-def get_conversation_history(user_id: str) -> List[Dict]:
+# Helper function to get conversation history for a user (by user_id and optionally thread_id)
+def get_conversation_history(user_id: str, thread_id: str = None) -> List[Dict]:
     try:
-        response = supabase.table("conversation_history").select("*").eq("user_id", user_id).order("timestamp", desc=True).execute()
+        query = supabase.table("conversation_history").select("*").eq("user_id", user_id)
+        if thread_id:
+            query = query.eq("thread_id", thread_id)
+        response = query.order("timestamp", desc=True).execute()
         return response.data
     except Exception as e:
         print(f"Error retrieving history: {e}")
         return []
 
+# Helper function to get all threads for a user (distinct thread_ids with their titles)
+def get_user_threads(user_id: str) -> List[Dict]:
+    try:
+        # Select distinct threads with their titles
+        response = supabase.table("conversation_history")\
+            .select("thread_id, title")\
+            .eq("user_id", user_id)\
+            .order("timestamp", desc=True)\
+            .execute()
+        # Remove duplicates while preserving order
+        seen = set()
+        threads = []
+        for item in response.data:
+            if item["thread_id"] not in seen:
+                seen.add(item["thread_id"])
+                threads.append({
+                    "thread_id": item["thread_id"],
+                    "title": item["title"] or f"Thread {len(threads) + 1}"
+                })
+        return threads
+    except Exception as e:
+        print(f"Error retrieving threads: {e}")
+        return []
+
 # Helper function to save conversation to Supabase
-def save_conversation(user_id: str, query: str, response: Dict):
+def save_conversation(user_id: str, query: str, response: Dict, thread_id: str = None, title: str = None):
     try:
+        # If no thread_id is provided, generate a new one
+        if not thread_id:
+            thread_id = str(uuid.uuid4())
+        
+        # If no title is provided, generate a default one based on the query
+        if not title:
+            title = query[:50] + "..." if len(query) > 50 else query
+
         conversation = {
             "user_id": user_id,
             "query": query,
-            "response": response,
-            "timestamp": datetime.utcnow().isoformat()
+            "response": json.dumps(response),  # Convert response Dict to string
+            "timestamp": datetime.utcnow().isoformat(),
+            "thread_id": thread_id,
+            "title": title
         }
         supabase.table("conversation_history").insert(conversation).execute()
+        return thread_id
     except Exception as e:
         print(f"Error saving conversation: {e}")
+        raise
 
 @app.post("/search")
 async def process_search(search_query: SearchQuery):
     try:
+        # Validate user_id as UUID
+        try:
+            uuid.UUID(search_query.user_id)
+        except ValueError:
+            raise HTTPException(status_code=400, detail="Invalid user_id: must be a valid UUID")
+
+        # Validate thread_id as UUID if provided
+        if search_query.thread_id:
+            try:
+                uuid.UUID(search_query.thread_id)
+            except ValueError:
+                raise HTTPException(status_code=400, detail="Invalid thread_id: must be a valid UUID")
+
         # Set default context if not provided
         base_context = search_query.context or "You are a cybersecurity expert."
 
-        # Retrieve previous conversation history for context
-        history = get_conversation_history(search_query.user_id)
-        history_context = "\n".join([f"Previous Query: {item['query']}\nPrevious Response: {item['response']['Use Clear Language: Avoid ambiguity and complex wording']}" for item in history])
+        # Retrieve previous conversation history for context (within the same thread if thread_id is provided)
+        history = get_conversation_history(search_query.user_id, search_query.thread_id)
+        history_context = "\n".join([f"Previous Query: {item['query']}\nPrevious Response: {json.loads(item['response'])['Use Clear Language: Avoid ambiguity and complex wording']}" for item in history])
         full_context = f"{base_context}\n{history_context}" if history_context else base_context
 
-        # Default to the "general" prompt template (no classification)
+        # Default to the "general" prompt template
         query_type = "general"
 
         # Process the query using the general chain
@@ -105,21 +163,41 @@ async def process_search(search_query: SearchQuery):
             "Experiment with different prompt structures and learn from the results": f"This response uses the '{query_type}' template. Try rephrasing the query for alternative perspectives or more specificity."
         }
 
-        # Save the conversation to Supabase
-        save_conversation(search_query.user_id, search_query.query, structured_response)
+        # Save the conversation to Supabase (append to existing thread or create new)
+        thread_id = save_conversation(
+            user_id=search_query.user_id,
+            query=search_query.query,
+            response=structured_response,
+            thread_id=search_query.thread_id
+        )
 
         return {
             "status": "success",
             "response": structured_response,
+            "thread_id": thread_id,
             "classified_type": query_type
         }
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/history/{user_id}")
-async def get_history(user_id: str):
+async def get_history(user_id: str, thread_id: str = None):
     try:
-        history = get_conversation_history(user_id)
+        # Validate user_id as UUID
+        try:
+            uuid.UUID(user_id)
+        except ValueError:
+            raise HTTPException(status_code=400, detail="Invalid user_id: must be a valid UUID")
+
+        # Validate thread_id as UUID if provided
+        if thread_id:
+            try:
+                uuid.UUID(thread_id)
+            except ValueError:
+                raise HTTPException(status_code=400, detail="Invalid thread_id: must be a valid UUID")
+
+        # Get conversation history for the user (optionally filtered by thread_id)
+        history = get_conversation_history(user_id, thread_id)
         return {
             "status": "success",
             "history": history
@@ -127,9 +205,27 @@ async def get_history(user_id: str):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+@app.get("/threads/{user_id}")
+async def get_threads(user_id: str):
+    try:
+        # Validate user_id as UUID
+        try:
+            uuid.UUID(user_id)
+        except ValueError:
+            raise HTTPException(status_code=400, detail="Invalid user_id: must be a valid UUID")
+
+        # Get all threads for the user
+        threads = get_user_threads(user_id)
+        return {
+            "status": "success",
+            "threads": threads
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
 @app.get("/")
 async def root():
-    return {"message": "Search API with structured response and history is running"}
+    return {"message": "Search API with structured response, history, and threads is running"}
 
 if __name__ == "__main__":
     import uvicorn