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 crewai import Agent, Task, Crew
 import os
 
 # Initialize FastAPI app
@@ -19,25 +20,63 @@ llm = ChatGroq(
     groq_api_key="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"  # Replace with your actual Groq API key
 )
 
-prompt_template = PromptTemplate(
-    input_variables=["query"],
-    template="Please provide a detailed response to the following query: {query}"
+# Define the classifier agent
+classifier_agent = Agent(
+    role='Classifier',
+    goal='Understand the context of the user query and generate up to 5 suggestions.',
+    backstory='You are an AI that specializes in understanding user queries and providing relevant suggestions.',
+    llm=llm,
+    verbose=True
+)
+
+# Define the task for the classifier agent
+classifier_task = Task(
+    description='Analyze the user query and generate up to 5 suggestions based on the context.',
+    agent=classifier_agent,
+    expected_output='A list of up to 5 suggestions related to the user query.'
+)
+
+# Define the main agent for processing the query
+main_agent = Agent(
+    role='Main Agent',
+    goal='Provide a detailed response to the user query.',
+    backstory='You are an AI that specializes in providing detailed and accurate responses to user queries.',
+    llm=llm,
+    verbose=True
+)
+
+# Define the task for the main agent
+main_task = Task(
+    description='Provide a detailed response to the user query.',
+    agent=main_agent,
+    expected_output='A detailed and accurate response to the user query.'
+)
+
+# Create the crew
+crew = Crew(
+    agents=[classifier_agent, main_agent],
+    tasks=[classifier_task, main_task],
+    verbose=2
 )
-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)
+        # Process the query using CrewAI
+        result = crew.kickoff(inputs={'query': search_query.query})
+        
+        # Extract the response and suggestions from the result
+        response = result['outputs']['main_agent']
+        suggestions = result['outputs']['classifier_agent']
         
         return {
             "status": "success",
-            "response": response
+            "response": response,
+            "suggestions": suggestions
         }
     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 is running"}