"""
Shared data models and types for the MCP system
"""
from dataclasses import dataclass
from typing import Dict, List, Optional, Any, Union
from enum import Enum
class QueryStrategy(Enum):
"""Strategy for answering a question"""
SQL_ONLY = "sql_only"
SEMANTIC_ONLY = "semantic_only"
HYBRID = "hybrid"
SCHEMA_ONLY = "schema_only"
# Alias for backwards compatibility
SearchStrategy = QueryStrategy
@dataclass
class QuestionClassification:
"""Classification of a user question"""
strategy: QueryStrategy
confidence: float
reasoning: str
needs_sql: bool = False
needs_semantic: bool = False
needs_schema: bool = False
@dataclass
class SQLQuery:
"""Represents a SQL query with metadata"""
sql: str
description: str
confidence: float
safety_check: bool = False
estimated_rows: Optional[int] = None
@dataclass
class QueryResult:
"""Result of executing a query"""
success: bool
data: Optional[List[Dict[str, Any]]] = None
error: Optional[str] = None
rows_affected: int = 0
execution_time: float = 0.0
query: Optional[str] = None
@dataclass
class SemanticSearchResult:
"""Result from semantic/vector search"""
content: str
score: float
source: str
metadata: Dict[str, Any]
@dataclass
class SchemaInfo:
"""Database schema information"""
tables: List[Dict[str, Any]]
relationships: List[Dict[str, Any]]
indexes: List[Dict[str, Any]]
views: List[Dict[str, Any]]
functions: List[Dict[str, Any]]
summary: str
@dataclass
class SearchRequest:
"""Request for smart search"""
question: str
filters: Optional[Dict[str, Any]] = None
limit: int = 10
include_sql: bool = True
include_semantic: bool = True
include_schema: bool = True
@dataclass
class SearchResponse:
"""Response from smart search"""
success: bool
answer_markdown: str
sources_used: Dict[str, int]
strategy: QueryStrategy
sql_queries: List[SQLQuery]
semantic_results: List[SemanticSearchResult]
processing_time: float
error: Optional[str] = None
@dataclass
class DatabaseConnection:
"""Database connection configuration"""
host: str
port: int
database: str
username: str
password: str
schema: Optional[str] = None
def get_connection_string(self) -> str:
"""Get SQLAlchemy connection string"""
return f"postgresql://{self.username}:{self.password}@{self.host}:{self.port}/{self.database}"
@dataclass
class EmbeddingConfig:
"""Configuration for embeddings"""
model: str
dimension: int
provider: str # "openai", "huggingface", etc.
api_key: Optional[str] = None
