from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from langchain_groq import ChatGroq
from crewai import Agent, Task, Crew
import os

# Initialize FastAPI app
app = FastAPI()

# Create a request model
class SearchQuery(BaseModel):
    query: str

# Initialize LangChain with Groq
llm = ChatGroq(
    temperature=0.7,
    model_name="mixtral-8x7b-32768",
    groq_api_key="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"  # Replace with your actual Groq API key
)

# 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
)

@app.post("/search")
async def process_search(search_query: SearchQuery):
    try:
        # 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,
            "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"}