Update app.py

app.py

@@ -20,40 +20,241 @@ llm = ChatGroq(
     groq_api_key="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"  # Replace with your actual Groq API key
 )
 
-# Define the prompt template with elite cybersecurity expertise
-prompt_template = PromptTemplate(
-    input_variables=["query", "context"],
-    template="""
-    Context: You are an elite cybersecurity AI with comprehensive
-    mastery of all domains, including network security, cloud security, threat intelligence, cryptography, and incident response. Your expertise spans
-    enterprise-grade strategies, current threat landscapes (2023-2024), and actionable mitigation tactics. Prioritize concise, technical, and
-    ROI-driven insights.
+# Define all prompt templates
+prompt_templates = {
+    # Cybersecurity Threats and Mitigation
+    "common_threats": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Provide a comprehensive overview of the most common cybersecurity threats faced by organizations on a daily basis, including details on threat_1, threat_2, and threat_3. Also, provide effective measures to mitigate these risks and protect critical data and systems.
+        """
+    ),
+    "task_prioritization": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Provide a guide on how cybersecurity professionals prioritize their tasks and responsibilities, focusing on the most critical areas such as threat detection, response times, and resource allocation.
+        """
+    ),
+    "network_traffic_tools": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        List and describe the most effective tools and software used for monitoring network traffic. Include tools for real-time analysis, anomaly detection, and reporting.
+        """
+    ),
+    "vulnerability_assessments": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Provide best practices for conducting vulnerability assessments and penetration tests, including recommended frequencies and methodologies to ensure systems are adequately tested for vulnerabilities.
+        """
+    ),
+    "security_policies": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Explain the role cybersecurity professionals have in developing, updating, and enforcing security policies within an organization. Include considerations for evolving threats and compliance requirements.
+        """
+    ),
+    "staying_updated": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Describe the methods and tools cybersecurity professionals use to stay up-to-date on the latest cybersecurity threats, trends, and vulnerabilities, including ongoing education and industry resources.
+        """
+    ),
+    "immediate_incidents": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Identify and describe the types of cybersecurity incidents that require immediate attention, such as data breaches, malware attacks, and denial-of-service attacks. Provide guidance on how to respond to each incident type.
+        """
+    ),
+    "collaboration_it_teams": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Discuss how cybersecurity professionals work with IT teams to ensure system security, focusing on areas such as patch management, incident response, and ongoing risk management.
+        """
+    ),
+    "incident_investigation": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Outline the steps involved in investigating and resolving a security incident, including initial detection, containment, root cause analysis, and reporting.
+        """
+    ),
+    "securing_remote_workers": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Provide strategies for securing remote workers and their devices, including the use of VPNs, multi-factor authentication, and endpoint protection measures.
+        """
+    ),
+    "disaster_recovery": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Explain the responsibilities of cybersecurity professionals in ensuring that disaster recovery and business continuity plans are developed, tested, and maintained to address security challenges.
+        """
+    ),
+    "user_access_management": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Describe the best practices for managing user access and privileges, including role-based access control (RBAC), least privilege principles, and audit trails for sensitive systems.
+        """
+    ),
+    "cloud_security": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Provide a list of best practices for securing cloud-based infrastructure, including the use of strong authentication, data encryption, and continuous monitoring.
+        """
+    ),
+    "security_kpis": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Discuss the key performance indicators (KPIs) used by cybersecurity professionals to measure the effectiveness of security programs, such as incident response times, patching cycles, and vulnerability remediation rates.
+        """
+    ),
+    "employee_security_education": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Describe the methods used by cybersecurity professionals to educate employees on security best practices, including training programs, phishing simulations, and awareness campaigns.
+        """
+    ),
+    "common_challenges": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Identify and discuss the common challenges that cybersecurity professionals face, including resource limitations, evolving threats, and the complexities of compliance.
+        """
+    ),
+    "compliance_standards": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Provide an overview of how cybersecurity professionals ensure compliance with industry standards and regulations, such as GDPR, HIPAA, and PCI DSS, including regular audits and reporting.
+        """
+    ),
+    "encryption_role": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Explain the role of encryption in protecting sensitive data, focusing on encryption methods, data-at-rest vs. data-in-transit, and how encryption helps mitigate the risks of data breaches.
+        """
+    ),
+    "mobile_device_security": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Provide strategies for managing and securing mobile devices and applications, including mobile device management (MDM), app whitelisting, and secure communication methods.
+        """
+    ),
+    "security_audits": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Outline the steps involved in conducting security audits and risk assessments, including identifying potential threats, assessing vulnerabilities, and recommending mitigation strategies.
+        """
+    ),
+    "patch_management": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Describe the best practices for managing patch updates and ensuring software security, including patch management policies, vulnerability scanning, and prioritizing patches based on risk.
+        """
+    ),
+    "wireless_iot_security": PromptTemplate(
+        input_variables=["query", "context"],
+        template="""
+        Context: {context}
+        Query: {query}
+        Provide a comprehensive guide on securing wireless networks and IoT devices, including the use of encryption, network segmentation, and regular vulnerability assessments.
+        """
+    ),
+    # Add all other prompts here...
+}
 
-    Response Rules:
-    - Structure responses using the pyramid principle (key takeaway first).
-    - Maximum 500 words per response.
-    - Use technical terminology appropriately (e.g., OWASP Top 10, MITRE ATT&CK, NIST references).
-    - Include critical data points:
-      - CVE IDs for vulnerabilities.
-      - CVSS scores where applicable.
-      - Latest compliance standards (e.g., ISO 27001:2022, NIST CSF 2.0).
-    - Format complex concepts clearly:
-      → Security through obscurity
-      → Zero-trust architecture
+# Initialize chains for each prompt
+chains = {key: LLMChain(llm=llm, prompt=prompt) for key, prompt in prompt_templates.items()}
 
-    Source Integration:
-    - Cite only authoritative sources (e.g., CISA alerts, RFCs, vendor advisories).
-    - Include timestamps for exploit disclosures.
-    - Flag conflicting industry perspectives where relevant.
-
-    Context: {context}
-    Query: {query}
-
-    Provide a concise, actionable, and enterprise-focused response** based on your expertise and the provided context.
-    """
-)
-
-chain = LLMChain(llm=llm, prompt=prompt_template)
+# Classify user input to determine the appropriate prompt
+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"
+    # Add more conditions for other prompts...
+    else:
+        return "general"
 
 @app.post("/search")
 async def process_search(search_query: SearchQuery):
@@ -61,8 +262,14 @@ async def process_search(search_query: SearchQuery):
         # Set default context if not provided
         context = search_query.context or "You are a cybersecurity expert."
 
-        # Process the query using LangChain with context
-        response = chain.run(query=search_query.query, context=context)
+        # Classify the query
+        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)
+        else:
+            response = "I am not sure how to help with that. Please provide more details."
 
         return {
             "status": "success",