ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/bots/db/abstract.py

@@ -0,0 +1,556 @@
+"""
+Database Agent Architecture for AI-Parrot.
+
+This module provides an abstract base for database introspection agents
+that can analyze database schemas and generate queries from natural language.
+"""
+
+from abc import abstractmethod
+from typing import Dict, Any, List, Optional, Union
+import uuid
+import json
+import asyncio
+from string import Template
+from sqlalchemy.ext.asyncio import AsyncEngine
+from ..abstract import AbstractBot
+from ...tools.manager import (
+    ToolManager,
+)
+from ...stores.abstract import AbstractStore
+from .prompts import DB_AGENT_PROMPT, BASIC_HUMAN_PROMPT
+from .tools import (
+    SchemaSearchTool,
+    QueryGenerationTool,
+    DatabaseSchema,
+    TableMetadata,
+    ExplainQueryTool,
+)
+from ...models import AIMessage
+from ...tools.databasequery import DatabaseQueryTool
+from .cache import SchemaCache
+from ...memory import ConversationTurn
+
+class AbstractDBAgent(AbstractBot):
+    """
+    Abstract base class for database introspection agents.
+
+    This agent analyzes database schemas, stores metadata in a knowledge base,
+    and generates queries from natural language descriptions.
+    """
+    system_prompt_template: str = DB_AGENT_PROMPT
+    human_prompt_template = BASIC_HUMAN_PROMPT
+    _default_temperature: float = 0.0
+    max_tokens: int = 8192
+
+    def __init__(
+        self,
+        name: str = "DatabaseAgent",
+        credentials: Union[str, Dict[str, Any]] = None,
+        schema_name: Optional[str] = None,
+        knowledge_store: AbstractStore = None,
+        auto_analyze_schema: bool = True,
+        cache_ttl: int = 3600,  # Cache TTL in seconds
+        **kwargs
+    ):
+        """
+        Initialize the database agent.
+
+        Args:
+            name: Agent name
+            credentials: Database connection credentials
+            schema_name: Target schema name (optional)
+            knowledge_store: Vector store for schema metadata
+            auto_analyze_schema: Whether to automatically analyze schema on init
+        """
+        kwargs.setdefault('temperature', self._default_temperature)
+        super().__init__(name=name, **kwargs)
+        self.role = kwargs.get(
+            'role', 'Database Analysis Assistant'
+        )
+        self.goal = kwargs.get(
+            'goal', 'Help users interact with databases using natural language'
+        )
+        self.capabilities = kwargs.get(
+            'capabilities',
+            'Database schema analysis, query generation, and data retrieval'
+        )
+        self.backstory = kwargs.get(
+            'backstory',
+            'Expert database assistant with deep knowledge of SQL and data analysis'
+        )
+
+        self.credentials = credentials
+        self.schema_name = schema_name
+        self.knowledge_store = knowledge_store
+        # Schema cache for metadata
+        self.cache = SchemaCache(ttl=cache_ttl)
+        self.auto_analyze_schema = auto_analyze_schema
+
+        # Initialize database-specific components
+        self.engine: Optional[AsyncEngine] = None
+        self.schema_metadata: Optional[DatabaseSchema] = None
+
+        # Initialize tool manager
+        self.tool_manager = ToolManager(
+            logger=self.logger,
+            debug=self._debug
+        )
+
+        # Add database-specific tools
+        self._setup_database_tools()
+        try:
+            self.loop = asyncio.get_running_loop()
+        except RuntimeError:
+            raise RuntimeError(
+                "An event loop is required to initialize the AbstractDBAgent."
+            )
+        asyncio.set_event_loop(self.loop)
+
+    @property
+    def tools(self):
+        """Get list of registered tools from manager."""
+        return self.tool_manager.get_all_tools() if self.tool_manager else []
+
+    async def initialize_schema(self):
+        """Initialize database connection and analyze schema."""
+        try:
+            await self.connect_database()
+            self.schema_metadata = await self.extract_schema_metadata()
+
+            if self.knowledge_store:
+                await self.store_schema_in_knowledge_base()
+
+        except Exception as e:
+            self.logger.error(
+                f"Failed to initialize schema: {e}"
+            )
+            raise
+
+    def _setup_database_tools(self):
+        """Setup database-specific tools."""
+        # Add schema search tool
+        schema_search_tool = SchemaSearchTool(agent=self)
+        self.tool_manager.register_tool(schema_search_tool)
+
+        # Add query generation tool
+        query_gen_tool = QueryGenerationTool(agent=self)
+        self.tool_manager.register_tool(query_gen_tool)
+
+        # Add database query tool
+        db_query_tool = DatabaseQueryTool(agent=self)
+        self.tool_manager.register_tool(db_query_tool)
+
+        # Add explain query tool
+        explain_query_tool = ExplainQueryTool(agent=self)
+        self.tool_manager.register_tool(explain_query_tool)
+
+    async def configure(self, app=None) -> None:
+        """Configure the database agent."""
+        await super().configure(app)
+        await self.initialize_schema()
+
+    def _define_prompt(self, config: Optional[dict] = None, **kwargs):
+        """
+        Define the System Prompt and replace variables.
+        """
+        # setup the prompt variables:
+        if config:
+            for key, val in config.items():
+                setattr(self, key, val)
+
+        tmpl = Template(self.system_prompt_template)
+        final_prompt = tmpl.safe_substitute(
+            name=self.name,
+            role=getattr(self, 'role', 'Database Analysis Assistant'),
+            goal=getattr(self, 'goal', 'Help users interact with databases using natural language'),
+            capabilities=getattr(self, 'capabilities', 'Database schema analysis, query generation, and data retrieval'),
+            backstory=getattr(self, 'backstory', 'Expert database assistant with deep knowledge of SQL and data analysis'),
+            rationale=getattr(self, 'rationale', 'Providing precise and helpful database interactions'),
+            **kwargs
+        )
+        self.system_prompt_template = final_prompt
+
+    async def create_system_prompt(
+        self,
+        user_context: str = "",
+        context: str = "",
+        vector_context: str = "",
+        pre_context: str = "",
+        conversation_context: str = "",
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> str:
+        """
+        Create the complete system prompt for the LLM with user context support.
+
+        Args:
+            user_context: User-specific context for the database interaction
+            vector_context: Vector store context
+            conversation_context: Previous conversation context
+            metadata: Additional metadata
+            **kwargs: Additional template variables
+        """
+        # Process contexts
+        context_parts = []
+
+        # Add user context if provided
+        if user_context:
+            context_parts.append(f"""
+Use the following information about user's data to guide your responses:
+**User Context:**
+{user_context}
+Based on the user context above, please tailor your response to their specific:
+- Role and responsibilities
+- Technical expertise level
+- Business objectives
+- Preferred communication style
+            """)
+        if context:
+            context_parts.append(f"**Additional Context:**\n{context}")
+
+        # Add vector context
+        if vector_context:
+            context_parts.append(f"**Database Context:**\n{vector_context}")
+
+        # Add metadata
+        if metadata:
+            metadata_text = "**Database Metadata:**\n"
+            for key, value in metadata.items():
+                if key == 'sources' and isinstance(value, list):
+                    metadata_text += f"- {key}: {', '.join(value[:3])}{'...' if len(value) > 3 else ''}\n"
+                else:
+                    metadata_text += f"- {key}: {value}\n"
+            context_parts.append(metadata_text)
+
+        # Format conversation context
+        chat_history_section = ""
+        if conversation_context:
+            chat_history_section = f"**Conversation History:**\n{conversation_context}"
+
+        # Database-specific context
+        db_context_parts = []
+        if self.schema_metadata:
+            db_info = f"**Database Information:**\n"
+            db_info += f"- Database: {self.schema_metadata.database_name}\n"
+            db_info += f"- Type: {self.schema_metadata.database_type}\n"
+            db_info += f"- Tables: {len(self.schema_metadata.tables)}\n"
+            db_info += f"- Views: {len(self.schema_metadata.views)}\n"
+            db_context_parts.append(db_info)
+
+        # Apply template substitution
+        tmpl = Template(self.system_prompt_template)
+        system_prompt = tmpl.safe_substitute(
+            user_context=user_context,
+            pre_context=pre_context,
+            context="\n\n".join(context_parts) if context_parts else "",
+            database_context="\n\n".join(db_context_parts) if db_context_parts else "",
+            chat_history=chat_history_section,
+            **kwargs
+        )
+        return system_prompt
+
+    async def cleanup(self) -> None:
+        """Cleanup resources including cache connections."""
+        await self.cache.close()
+        await super().cleanup() if hasattr(super(), 'cleanup') else None
+
+    @abstractmethod
+    async def connect_database(self) -> None:
+        """Connect to the database. Must be implemented by subclasses."""
+        pass
+
+    @abstractmethod
+    async def extract_schema_metadata(self) -> DatabaseSchema:
+        """
+        Extract complete schema metadata from the database.
+        Must be implemented by subclasses based on database type.
+        """
+        pass
+
+    @abstractmethod
+    async def generate_query(
+        self,
+        natural_language_query: str,
+        target_tables: Optional[List[str]] = None,
+        query_type: str = "SELECT"
+    ) -> Dict[str, Any]:
+        """
+        Generate database query from natural language.
+        Must be implemented by subclasses based on database type.
+        """
+        pass
+
+    @abstractmethod
+    async def execute_query(self, query: str) -> Dict[str, Any]:
+        """
+        Execute a query against the database.
+        Must be implemented by subclasses based on database type.
+        """
+        pass
+
+    @abstractmethod
+    async def explain_query(self, query: str) -> str:
+        """
+        Explain a database query (e.g. EXPLAIN ANALYZE).
+        Must be implemented by subclasses based on database type.
+        """
+        pass
+
+    async def search_schema(
+        self,
+        search_term: str,
+        search_type: str = "all",
+        limit: int = 10
+    ) -> List[Dict[str, Any]]:
+        """
+        Search the database schema for tables or columns matching the term.
+        """
+        if not self.schema_metadata:
+            # Try to load from cache or re-extract
+             return []
+
+        results = []
+        term = search_term.lower()
+
+        # Helper to check string match
+        def matches(text: Optional[str]) -> bool:
+            return bool(text and term in text.lower())
+
+        for table in self.schema_metadata.tables:
+            table_match = False
+            # Check table name and description
+            if search_type in ["all", "tables"]:
+                full_name = f"{table.schema}.{table.name}" if table.schema else table.name
+                if matches(table.name) or matches(table.description) or matches(full_name):
+                    results.append({
+                        "type": "table",
+                        "name": table.name,
+                        "schema": table.schema,
+                        "description": table.description
+                    })
+                    table_match = True
+
+            # Check columns
+            if search_type in ["all", "columns"]:
+                for col in table.columns:
+                    if matches(col.get("name")) or matches(col.get("description")):
+                        # Add table context if not already added/searching specifically for columns
+                        if not table_match:
+                             results.append({
+                                "type": "column",
+                                "table": table.name,
+                                "schema": table.schema,
+                                "name": col.get("name"),
+                                "description": col.get("description"),
+                                "metadata": f"Type: {col.get('type')}"
+                            })
+
+        # Sort by relevance (exact matches first? undefined for now) and limit
+        return results[:limit]
+
+    async def store_schema_in_knowledge_base(self) -> None:
+        """Store schema metadata in the knowledge base for retrieval."""
+        if not self.knowledge_store or not self.schema_metadata:
+            return
+
+        documents = []
+
+        # Store table metadata
+        for table in self.schema_metadata.tables:
+            table_doc = {
+                "content": self._format_table_for_storage(table),
+                "metadata": {
+                    "type": "table_schema",
+                    "database": self.schema_metadata.database_name,
+                    "schema": table.schema,
+                    "tablename": table.tablename,
+                    "database_type": self.schema_metadata.database_type
+                }
+            }
+            documents.append(table_doc)
+
+        # Store view metadata
+        for view in self.schema_metadata.views:
+            view_doc = {
+                "content": self._format_table_for_storage(view, is_view=True),
+                "metadata": {
+                    "type": "view_schema",
+                    "database": self.schema_metadata.database_name,
+                    "schema": view.schema,
+                    "view_name": view.tablename,
+                    "database_type": self.schema_metadata.database_type
+                }
+            }
+            documents.append(view_doc)
+
+        # Store in knowledge base
+        await self.knowledge_store.add_documents(documents)
+
+    def _format_table_for_storage(self, table: TableMetadata, is_view: bool = False) -> str:
+        """Format table metadata for storage in knowledge base."""
+        object_type = "VIEW" if is_view else "TABLE"
+
+        content = f"""
+{object_type}: {table.schema}.{table.tablename}
+Description: {table.description or 'No description available'}
+
+Columns:
+"""
+        for col in table.columns:
+            nullable = "NULL" if col.get('nullable', True) else "NOT NULL"
+            default = f" DEFAULT {col['default']}" if col.get('default') else ""
+            content += f"  - {col['name']}: {col['type']} {nullable}{default}\n"
+            if col.get('description'):
+                content += f"    Description: {col['description']}\n"
+
+        if table.primary_keys:
+            content += f"\nPrimary Keys: {', '.join(table.primary_keys)}\n"
+
+        if table.foreign_keys:
+            content += "\nForeign Keys:\n"
+            for fk in table.foreign_keys:
+                content += f"  - {fk['column']} -> {fk['referenced_table']}.{fk['referenced_column']}\n"
+
+        if table.indexes:
+            content += "\nIndexes:\n"
+            for idx in table.indexes:
+                content += f"  - {idx['name']}: {', '.join(idx['columns'])}\n"
+
+        if table.sample_data:
+            content += "\nSample Data:\n"
+            for i, row in enumerate(table.sample_data[:3]):  # Limit to 3 rows
+                content += f"  Row {i+1}: {json.dumps(row, default=str)}\n"
+
+        return content
+
+    async def ask(
+        self,
+        question: str = None,
+        user_context: Optional[str] = None,
+        context: Optional[str] = None,
+        session_id: Optional[str] = None,
+        user_id: Optional[str] = None,
+        use_conversation_history: bool = True,
+        **kwargs
+    ) -> AIMessage:
+        """
+        Database-specific ask method with user context support and agentic mode.
+
+        Args:
+            question: The user's question about the database
+            user_context: User-specific context for database interaction
+            session_id: Session identifier for conversation history
+            user_id: User identifier
+            use_conversation_history: Whether to use conversation history
+            **kwargs: Additional arguments for LLM
+
+        Returns:
+            AIMessage: The response from the LLM
+        """
+        # Backwards compatibility for prompt
+        if question is None:
+            question = kwargs.get('prompt')
+        
+        prompt = question # internal usage expects prompt variable name in logic below
+        # Force agentic mode for database operations
+        effective_mode = "agentic"
+
+        # Override temperature to ensure consistent database operations
+        kwargs['temperature'] = kwargs.get('temperature', self._default_temperature)
+
+        # Generate session ID if not provided
+        if not session_id:
+            session_id = f"db_session_{hash(prompt + str(user_id))}"
+
+        try:
+            # Get conversation history if enabled
+            conversation_history = None
+            conversation_context = ""
+
+            if use_conversation_history and self.conversation_memory:
+                conversation_history = await self.get_conversation_history(user_id, session_id)
+                if not conversation_history:
+                    conversation_history = await self.create_conversation_history(
+                        user_id, session_id
+                    )
+                conversation_context = self.build_conversation_context(conversation_history)
+
+            # Get vector context from knowledge store if available
+            vector_context = ""
+            vector_metadata = {}
+            if self.knowledge_store:
+                try:
+                    # Search for relevant schema information
+                    search_results = await self.knowledge_store.similarity_search(
+                        prompt, k=5
+                    )
+                    if search_results:
+                        vector_context = "\n\n".join([doc.page_content for doc in search_results])
+                        vector_metadata = {
+                            'sources': [doc.metadata.get('source', 'unknown') for doc in search_results],
+                            'tables_referenced': [doc.metadata.get('table_name') for doc in search_results if doc.metadata.get('table_name')]
+                        }
+                except Exception as e:
+                    self.logger.warning(f"Error retrieving vector context: {e}")
+
+            # Create system prompt with user context
+            system_prompt = await self.create_system_prompt(
+                user_context=user_context,
+                context=context or "",
+                vector_context=vector_context,
+                conversation_context=conversation_context,
+                metadata=vector_metadata,
+                **kwargs
+            )
+
+            # Log the operation
+            self.logger.info(
+                f"Database query in agentic mode: use_tools=True, "
+                f"effective_mode={effective_mode}, available_tools={len(self.tools)}"
+            )
+
+            # Make the LLM call with tools enabled
+            async with self._llm as client:
+                response = await client.ask(
+                    prompt=prompt,
+                    model=kwargs.get('model', self._llm_model),
+                    system_prompt=system_prompt,
+                    user_id=user_id,
+                    session_id=session_id,
+                    temperature=kwargs.get('temperature', self._default_temperature),
+                    tools=self.tools if self.tools else None,
+                    **{k: v for k, v in kwargs.items() if k not in ['model', 'temperature']}
+                )
+
+            # Update conversation memory
+            if use_conversation_history and conversation_history:
+                assistant_content = str(response.output) if response.output is not None else (response.response or "")
+
+                turn = ConversationTurn(
+                    turn_id=str(uuid.uuid4()),
+                    user_id=user_id,
+                    user_message=prompt,
+                    assistant_response=assistant_content,
+                    metadata={
+                        'user_context': user_context,
+                        'effective_mode': effective_mode,
+                        'tools_used': [tool.name for tool in response.tool_calls] if response.tool_calls else []
+                    }
+                )
+                chatbot_key = getattr(self, 'chatbot_id', None)
+                if chatbot_key is not None:
+                    chatbot_key = str(chatbot_key)
+                await self.conversation_memory.add_turn(
+                    user_id,
+                    session_id,
+                    turn,
+                    chatbot_id=chatbot_key
+                )
+
+            return response
+
+        except Exception as e:
+            self.logger.error(
+                f"Error in database ask method: {e}"
+            )
+            raise