from sentence_transformers import SentenceTransformer from langchain_groq import ChatGroq from langchain_core.prompts import ChatPromptTemplate import os from dotenv import load_dotenv from collections import deque # Load environment variables from .env file load_dotenv() class QueryHandler: def __init__(self, collection, llm_model_name='llama-3.3-70b-versatile', embed_model_name='sentence-transformers/all-MiniLM-L6-v2', max_history=5): self.collection = collection self.embed_model = SentenceTransformer(embed_model_name) # Load the Groq API key from environment variable self.groq_api_key = os.getenv('GROQ_API_KEY') if not self.groq_api_key: raise ValueError("GROQ_API_KEY environment variable not set.") # Initialize the LLM client with the API key self.llm = ChatGroq(groq_api_key=self.groq_api_key, model_name=llm_model_name) # Session history management self.history = deque(maxlen=max_history) # Store last 'max_history' interactions def extract_filters(self, user_prompt): """ Use the LLM to extract filter criteria from the user prompt. """ extraction_prompt = ChatPromptTemplate.from_template( """ Analyze the following query and extract filters for restaurant recommendations: - Cuisine type (if mentioned): Look for specific cuisine preferences. - Location: Identify the city or locality. - Price range: Look for a maximum or minimum price (e.g., $100). If any of these details are missing, return them as None. Query: {input} Provide the extracted filters in JSON format like: {{ "cuisine_type": "string or None", "location": "string or None", "price": "int or None" }} """ ) prompt = extraction_prompt.format(input=user_prompt) response = self.llm.invoke(prompt) try: filters = eval(response.content) # Convert string response to dictionary except Exception as e: filters = {"cuisine_type": None, "location": None, "price": None} return filters def query(self, user_prompt, n_results=5): """ Process the user query and retrieve filtered results from ChromaDB. """ # Generate embedding for the user prompt query_embedding = self.embed_model.encode(user_prompt).tolist() # Extract filters using the LLM filters = self.extract_filters(user_prompt) # Create the filter dictionary for ChromaDB chromadb_filter = {} if filters.get("cuisine_type"): chromadb_filter["cuisine_type"] = filters["cuisine_type"] if filters.get("location"): chromadb_filter["location"] = filters["location"] if filters.get("price"): chromadb_filter["price"] = {"$lte": filters["price"]} # Perform the query in ChromaDB with filters results = self.collection.query( query_embeddings=[query_embedding], n_results=n_results, where=chromadb_filter ) return results def get_recent_context(self): # Retrieve recent history for context, including LLM-generated responses context_parts = [] for idx, (prompt, results, llm_response) in enumerate(self.history): metadatas = results.get('metadatas', []) flat_metadatas = [item for sublist in metadatas for item in sublist] context = "\n".join( f"{idx + 1}. Restaurant: {metadata.get('restaurant_name', 'N/A')}, " f"Location: {metadata.get('location', 'N/A')}, " f"Locality: {metadata.get('locality', 'N/A')}, " f"City: {metadata.get('city', 'N/A')}, " f"Votes: {metadata.get('votes', 'N/A')}, " f"Cost: {metadata.get('cost', 'N/A')}, " f"Rating: {metadata.get('rating', 'N/A')}" for metadata in flat_metadatas ) # Include the LLM-generated response for each entry context_parts.append(f"Query: {prompt}\nResults:\n{context}\nResponse:\n{llm_response}") return "\n".join(context_parts) def generate_response(self, results, user_prompt): # Gather recent context from history recent_context = self.get_recent_context() # Current query's results metadatas = results.get('metadatas', []) flat_metadatas = [item for sublist in metadatas for item in sublist] # Construct recommendation text and collect image paths recommendation_text = "" image_paths = [] # To store image paths for rendering for metadata in flat_metadatas: recommendation_text += ( f"Restaurant: {metadata.get('restaurant_name', 'N/A')}, " f"Location: {metadata.get('location', 'N/A')}, " f"Locality: {metadata.get('locality', 'N/A')}, " f"City: {metadata.get('city', 'N/A')}, " f"Votes: {metadata.get('votes', 'N/A')}, " f"Cost: {metadata.get('cost', 'N/A')}, " f"Rating: {metadata.get('rating', 'N/A')}\n" ) if metadata.get('image_path'): # Append valid image paths image_paths.append(metadata['image_path']) # LLM prompt template with session context prompt_template = ChatPromptTemplate.from_template( """ Previous Conversations: {recent_context} Here is the data relevant to your query: - **Name**: The restaurant's name. - **Location, Locality, and City**: The address details of the restaurant. - **Cuisine**: Types of cuisines offered by the restaurant. - **Rating**: Average rating on a scale of 5. - **Votes**: Number of people who have rated the restaurant. - **Cost**: Average cost of dining. You are a restaurant recommender with knowledge of restaurants, cuisines, ratings, and costs across various cities in India. Respond to queries based on the provided details and recommend the most relevant options based on user preferences. If information is incomplete, provide the best suggestions and encourage follow-up questions. Use the recent context only if you require additional information from previous conversations. If the query is not related to the information available to you, make a generic response and let the user know you would be more than happy to help with restaurant recommendations. {context} Question: {input} """ ) # Generate the prompt by formatting the template prompt = prompt_template.format(recent_context=recent_context, context=recommendation_text, input=user_prompt) # Use the LLM to generate the output response = self.llm.invoke(prompt) # Pass the prompt string directly # Store the prompt, results, and LLM-generated response in the history self.history.append((user_prompt, results, response.content)) return response.content, image_paths # Return both recommendation text and image paths