Update app.py

app.py

@@ -8,9 +8,10 @@ import os
 # Initialize FastAPI app
 app = FastAPI()
 
-# Create a request model
+# Create a request model with context
 class SearchQuery(BaseModel):
     query: str
+    context: str = None  # Optional context field
 
 # Initialize LangChain with Groq
 llm = ChatGroq(
@@ -19,17 +20,25 @@ llm = ChatGroq(
     groq_api_key="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"  # Replace with your actual Groq API key
 )
 
+# Define the prompt template with cybersecurity expertise
 prompt_template = PromptTemplate(
-    input_variables=["query"],
-    template="Please provide a detailed response to the following query: {query}"
+    input_variables=["query", "context"],
+    template="""
+    Context: You are a cybersecurity expert with extensive experience in all sub-streams of the industry, including but not limited to network security, application security, cloud security, threat intelligence, penetration testing, and incident response. {context}
+    Query: {query}
+    Please provide a detailed and professional response to the query based on your expertise in cybersecurity and the provided context.
+    """
 )
 chain = LLMChain(llm=llm, prompt=prompt_template)
 
 @app.post("/search")
 async def process_search(search_query: SearchQuery):
     try:
-        # Process the query using LangChain
-        response = chain.run(query=search_query.query)
+        # 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)
         
         return {
             "status": "success",
@@ -40,4 +49,4 @@ async def process_search(search_query: SearchQuery):
 
 @app.get("/")
 async def root():
-    return {"message": "Search API is running"} 
+    return {"message": "Search API is running"}