diff --git a/api/agents/analysis_agent.py b/api/agents/analysis_agent.py index f274e8a8..661c3025 100644 --- a/api/agents/analysis_agent.py +++ b/api/agents/analysis_agent.py @@ -11,7 +11,7 @@ class AnalysisAgent(BaseAgent): """Agent for analyzing user queries and generating database analysis.""" - def get_analysis( + def get_analysis( # pylint: disable=too-many-arguments, too-many-positional-arguments self, user_query: str, combined_tables: list, @@ -156,7 +156,7 @@ def _format_foreign_keys(self, foreign_keys: dict) -> str: return fk_str - def _build_prompt( + def _build_prompt( # pylint: disable=too-many-arguments, too-many-positional-arguments self, user_input: str, formatted_schema: str, db_description: str, instructions, memory_context: str | None = None ) -> str: @@ -292,5 +292,5 @@ def _build_prompt( 12. For personal queries, FIRST check memory context for user identification. If user identity is found in memory context (user name, previous personal queries, etc.), the query IS translatable. 13. CRITICAL PERSONALIZATION CHECK: If missing user identification/personalization is a significant or primary component of the query (e.g., "show my orders", "my account balance", "my recent purchases", "how many employees I have", "products I own") AND no user identification is available in memory context or schema, set "is_sql_translatable" to false. However, if memory context contains user identification (like user name or previous successful personal queries), then personal queries ARE translatable even if they are the primary component of the query. - Again: OUTPUT ONLY VALID JSON. No explanations outside the JSON block. """ + Again: OUTPUT ONLY VALID JSON. No explanations outside the JSON block. """ # pylint: disable=line-too-long return prompt diff --git a/api/agents/follow_up_agent.py b/api/agents/follow_up_agent.py index e5d596f4..3798f3b5 100644 --- a/api/agents/follow_up_agent.py +++ b/api/agents/follow_up_agent.py @@ -30,14 +30,13 @@ """ -class FollowUpAgent(BaseAgent): +class FollowUpAgent(BaseAgent): # pylint: disable=too-few-public-methods """Agent for generating helpful follow-up questions when queries fail or are off-topic.""" def generate_follow_up_question( - self, + self, user_question: str, - analysis_result: dict, - found_tables: list = None + analysis_result: dict ) -> str: """ Generate helpful follow-up questions based on failed SQL translation. @@ -51,13 +50,16 @@ def generate_follow_up_question( Returns: str: Conversational follow-up response """ - + # Extract key information from analysis result - is_translatable = analysis_result.get("is_sql_translatable", False) if analysis_result else False + is_translatable = ( + analysis_result.get("is_sql_translatable", False) + if analysis_result else False + ) missing_info = analysis_result.get("missing_information", []) if analysis_result else [] ambiguities = analysis_result.get("ambiguities", []) if analysis_result else [] - explanation = analysis_result.get("explanation", "No detailed explanation available") if analysis_result else "No analysis result available" - + explanation = (analysis_result.get("explanation", "No detailed explanation available") + if analysis_result else "No analysis result available") # Prepare the prompt prompt = FOLLOW_UP_GENERATION_PROMPT.format( QUESTION=user_question, @@ -66,17 +68,17 @@ def generate_follow_up_question( AMBIGUITIES=ambiguities, EXPLANATION=explanation ) - + try: completion_result = completion( model=Config.COMPLETION_MODEL, messages=[{"role": "user", "content": prompt}], temperature=0.9 ) - + response = completion_result.choices[0].message.content.strip() return response - - except Exception as e: + + except Exception: # pylint: disable=broad-exception-caught # Fallback response if LLM call fails - return "I'm having trouble generating a follow-up question right now. Could you try rephrasing your question or providing more specific details about what you're looking for?" + return "Sorry, I couldn't generate a follow-up. Could you clarify your question a bit?" diff --git a/api/agents/relevancy_agent.py b/api/agents/relevancy_agent.py index 78ef6aa4..9317ec60 100644 --- a/api/agents/relevancy_agent.py +++ b/api/agents/relevancy_agent.py @@ -70,15 +70,6 @@ class RelevancyAgent(BaseAgent): # pylint: disable=too-few-public-methods """Agent for determining relevancy of queries to database schema.""" - def __init__(self, queries_history: list[str], result_history: list[str]): - """Initialize the relevancy agent with query and result history.""" - if result_history is None: - self.messages = [] - else: - self.messages = [] - for query, result in zip(queries_history[:-1], result_history): - self.messages.append({"role": "user", "content": query}) - self.messages.append({"role": "assistant", "content": result}) async def get_answer(self, user_question: str, database_desc: dict) -> dict: """Get relevancy assessment for user question against database description.""" diff --git a/api/agents/taxonomy_agent.py b/api/agents/taxonomy_agent.py index be527964..eb8895ad 100644 --- a/api/agents/taxonomy_agent.py +++ b/api/agents/taxonomy_agent.py @@ -1,6 +1,7 @@ """Taxonomy agent for taxonomy classification of questions and SQL queries.""" from litellm import completion +from api.agents.utils import BaseAgent from api.config import Config @@ -35,12 +36,10 @@ """ -class TaxonomyAgent: +class TaxonomyAgent(BaseAgent): # pylint: disable=too-few-public-methods """Agent for taxonomy classification of questions and SQL queries.""" - def __init__(self): - """Initialize the taxonomy agent.""" def get_answer(self, question: str, sql: str) -> str: """Get taxonomy classification for a question and SQL pair.""" diff --git a/api/agents/utils.py b/api/agents/utils.py index ac6605c0..53e678a0 100644 --- a/api/agents/utils.py +++ b/api/agents/utils.py @@ -4,7 +4,7 @@ from typing import Any, Dict -class BaseAgent: +class BaseAgent: # pylint: disable=too-few-public-methods """Base class for agents.""" def __init__(self, queries_history: list, result_history: list):