Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions api/agents/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,14 @@
from .relevancy_agent import RelevancyAgent
from .follow_up_agent import FollowUpAgent
from .response_formatter_agent import ResponseFormatterAgent
from .healer_agent import HealerAgent
from .utils import parse_response

__all__ = [
"AnalysisAgent",
"RelevancyAgent",
"FollowUpAgent",
"ResponseFormatterAgent",
"HealerAgent",
"parse_response"
]
21 changes: 18 additions & 3 deletions api/agents/analysis_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,18 +18,29 @@ def get_analysis( # pylint: disable=too-many-arguments, too-many-positional-arg
db_description: str,
instructions: str | None = None,
memory_context: str | None = None,
database_type: str | None = None,
) -> dict:
"""Get analysis of user query against database schema."""
formatted_schema = self._format_schema(combined_tables)
# Add system message with database type if not already present
if not self.messages or self.messages[0].get("role") != "system":
self.messages.insert(0, {
"role": "system",
"content": (
f"You are a SQL expert. TARGET DATABASE: "
f"{database_type.upper() if database_type else 'UNKNOWN'}"
)
})

prompt = self._build_prompt(
user_query, formatted_schema, db_description, instructions, memory_context
user_query, formatted_schema, db_description,
instructions, memory_context, database_type
)
self.messages.append({"role": "user", "content": prompt})
completion_result = completion(
model=Config.COMPLETION_MODEL,
messages=self.messages,
temperature=0,
top_p=1,
)

response = completion_result.choices[0].message.content
Expand Down Expand Up @@ -158,7 +169,8 @@ def _format_foreign_keys(self, foreign_keys: dict) -> str:

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
db_description: str, instructions, memory_context: str | None = None,
database_type: str | None = None,
) -> str:
"""
Build the prompt for Claude to analyze the query.
Expand All @@ -169,6 +181,7 @@ def _build_prompt( # pylint: disable=too-many-arguments, too-many-positional-a
db_description: Description of the database
instructions: Custom instructions for the query
memory_context: User and database memory context from previous interactions
database_type: Target database type (sqlite, postgresql, mysql, etc.)

Returns:
The formatted prompt for Claude
Expand Down Expand Up @@ -196,6 +209,8 @@ def _build_prompt( # pylint: disable=too-many-arguments, too-many-positional-a
prompt = f"""
You must strictly follow the instructions below. Deviations will result in a penalty to your confidence score.

TARGET DATABASE: {database_type.upper() if database_type else 'UNKNOWN'}

MANDATORY RULES:
- Always explain if you cannot fully follow the instructions.
- Always reduce the confidence score if instructions cannot be fully applied.
Expand Down
Loading