| import os |
| from groq import Groq |
| from typing import List, Dict, Any |
|
|
| class QueryEngine: |
| def __init__(self, api_key: str = None): |
| """Initialize the query engine with Groq API""" |
| if not api_key: |
| api_key = os.getenv("GROQ_API_KEY") |
| if not api_key: |
| raise ValueError("Groq API key is required") |
| |
| self.client = Groq(api_key=api_key) |
| self.model = "llama3-70b-8192" |
| |
| def generate_response(self, query: str, context_docs: List[Dict[str, Any]] = None): |
| """Generate a response using the LLM""" |
| |
| context = "" |
| if context_docs: |
| for doc in context_docs.get('documents', [[]])[0]: |
| context += f"{doc}\n\n" |
| |
| |
| prompt = f"""You are an AI assistant that helps with document analysis and answering questions. |
| |
| Context information: |
| {context} |
| |
| User question: {query} |
| |
| Please provide a helpful, accurate, and concise answer based on the context information provided. If the context doesn't contain relevant information, say so instead of making up an answer.""" |
| |
| |
| response = self.client.chat.completions.create( |
| model=self.model, |
| messages=[ |
| {"role": "system", "content": "You are a helpful AI assistant."}, |
| {"role": "user", "content": prompt} |
| ], |
| temperature=0.2, |
| max_tokens=1024 |
| ) |
| |
| return response.choices[0].message.content |
|
|
| def generate_sql_query(self, question: str, table_info: str): |
| """Generate an SQL query for the given question and table information""" |
| prompt = f"""You are an SQL and data analysis expert. Generate an appropriate SQL query using SQLite syntax for the question provided, without any explanations or code comments. |
| |
| Table Information: |
| {table_info} |
| |
| User Question: {question} |
| |
| Generate only the SQL query, nothing else.""" |
| |
| |
| response = self.client.chat.completions.create( |
| model=self.model, |
| messages=[ |
| {"role": "system", "content": "You are an SQL expert."}, |
| {"role": "user", "content": prompt} |
| ], |
| temperature=0.1, |
| max_tokens=512 |
| ) |
| |
| return response.choices[0].message.content.strip() |
|
|
| def analyze_sql_results(self, question: str, sql_query: str, results: str): |
| """Analyze the results of an SQL query""" |
| prompt = f"""You are a data analyst. Analyze the following SQL query results and provide a clear, concise interpretation. |
| |
| User Question: {question} |
| SQL Query: {sql_query} |
| Query Results: |
| {results} |
| |
| Provide a clear analysis of these results that directly answers the user's question.""" |
| |
| |
| response = self.client.chat.completions.create( |
| model=self.model, |
| messages=[ |
| {"role": "system", "content": "You are a data analyst."}, |
| {"role": "user", "content": prompt} |
| ], |
| temperature=0.2, |
| max_tokens=1024 |
| ) |
| |
| return response.choices[0].message.content |
