Update app.py

app.py

@@ -3,6 +3,7 @@ from pydantic import BaseModel
 from langchain_groq import ChatGroq
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
+from transformers import pipeline
 import os
 
 # Initialize FastAPI app
@@ -22,7 +23,6 @@ llm = ChatGroq(
 
 # Define all prompt templates
 prompt_templates = {
-    # Cybersecurity Threats and Mitigation
     "common_threats": PromptTemplate(
         input_variables=["query", "context"],
         template="""
@@ -199,7 +199,6 @@ prompt_templates = {
         Provide a comprehensive guide on securing wireless networks and IoT devices, including the use of encryption, network segmentation, and regular vulnerability assessments.
         """
     ),
-    # General Prompt
     "general": PromptTemplate(
         input_variables=["query", "context"],
         template="""
@@ -213,56 +212,35 @@ prompt_templates = {
 # Initialize chains for each prompt
 chains = {key: LLMChain(llm=llm, prompt=prompt) for key, prompt in prompt_templates.items()}
 
-# Classify user input to determine the appropriate prompt
+# Initialize the zero-shot classifier
+classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
+
+# Define the possible question types (labels) based on your prompt templates
+question_types = list(prompt_templates.keys())
+
+# Classifier function using the model
 def classify_query(query: str) -> str:
-    # Simple keyword-based classification
-    if "common threats" in query.lower():
-        return "common_threats"
-    elif "task prioritization" in query.lower():
-        return "task_prioritization"
-    elif "network traffic tools" in query.lower():
-        return "network_traffic_tools"
-    elif "vulnerability assessments" in query.lower():
-        return "vulnerability_assessments"
-    elif "security policies" in query.lower():
-        return "security_policies"
-    elif "staying updated" in query.lower():
-        return "staying_updated"
-    elif "immediate incidents" in query.lower():
-        return "immediate_incidents"
-    elif "collaboration with IT teams" in query.lower():
-        return "collaboration_it_teams"
-    elif "incident investigation" in query.lower():
-        return "incident_investigation"
-    elif "securing remote workers" in query.lower():
-        return "securing_remote_workers"
-    elif "disaster recovery" in query.lower():
-        return "disaster_recovery"
-    elif "user access management" in query.lower():
-        return "user_access_management"
-    elif "cloud security" in query.lower():
-        return "cloud_security"
-    elif "security KPIs" in query.lower():
-        return "security_kpis"
-    elif "employee security education" in query.lower():
-        return "employee_security_education"
-    elif "common challenges" in query.lower():
-        return "common_challenges"
-    elif "compliance standards" in query.lower():
-        return "compliance_standards"
-    elif "encryption role" in query.lower():
-        return "encryption_role"
-    elif "mobile device security" in query.lower():
-        return "mobile_device_security"
-    elif "security audits" in query.lower():
-        return "security_audits"
-    elif "patch management" in query.lower():
-        return "patch_management"
-    elif "wireless and IoT security" in query.lower():
-        return "wireless_iot_security"
-    # Default to the general prompt
-    else:
-        return "general"
+    """
+    Classify the query using a zero-shot classification model.
+    Returns the most likely question type from the prompt templates.
+    """
+    try:
+        # Perform zero-shot classification
+        result = classifier(query, candidate_labels=question_types)
+        
+        # Get the label with the highest score
+        predicted_type = result["labels"][0]
+        confidence = result["scores"][0]
+        
+        # If confidence is too low (e.g., < 0.5), fallback to 'general'
+        if confidence < 0.5:
+            print(f"Low confidence ({confidence}) for query '{query}', falling back to 'general'")
+            return "general"
+        
+        return predicted_type
+    except Exception as e:
+        print(f"Error in classification: {e}")
+        return "general"  # Fallback to general in case of errors
 
 @app.post("/search")
 async def process_search(search_query: SearchQuery):
@@ -270,22 +248,36 @@ async def process_search(search_query: SearchQuery):
         # Set default context if not provided
         context = search_query.context or "You are a cybersecurity expert."
 
-        # Classify the query
+        # Classify the query using the model
         query_type = classify_query(search_query.query)
 
         # Process the query using the appropriate chain
         if query_type in chains:
-            response = chains[query_type].run(query=search_query.query, context=context)
+            raw_response = chains[query_type].run(query=search_query.query, context=context)
         else:
-            response = chains["general"].run(query=search_query.query, context=context)
+            raw_response = chains["general"].run(query=search_query.query, context=context)
+
+        # Structure the response according to the desired format
+        structured_response = {
+            "Clearly articulate your task and desired outcome": f"The task is to address the query: '{search_query.query}'. The desired outcome is a detailed, actionable response.",
+            "Offer relevant background information to guide the AI’s understanding": f"The query was processed with the context: '{context}', guiding the response to align with cybersecurity expertise.",
+            "Use Clear Language: Avoid ambiguity and complex wording": raw_response.strip(),  # The raw response from Grok, cleaned up
+            "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."
+        }
 
         return {
             "status": "success",
-            "response": response
+            "response": structured_response,
+            "classified_type": query_type  # Optional: return the classified type for debugging
         }
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 @app.get("/")
 async def root():
-    return {"message": "Search API is running"}
+    return {"message": "Search API with structured response is running"}
+
+# Run the app (optional, for local testing)
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)