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

parrot/tools/database/abstract.py

@@ -0,0 +1,286 @@
+from typing import Any, Dict, List, Optional
+from abc import ABC, abstractmethod
+from pydantic import Field
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
+from sqlalchemy.orm import sessionmaker
+from navconfig.logging import logging
+from ..abstract import (
+    AbstractTool,
+    ToolResult,
+    AbstractToolArgsSchema
+)
+from ...stores.abstract import AbstractStore
+from ...conf import asyncpg_sqlalchemy_url
+from .models import TableMetadata
+from .cache import SchemaMetadataCache
+
+
+class SchemaSearchArgs(AbstractToolArgsSchema):
+    """Arguments for schema search tool."""
+    search_term: str = Field(
+        description="Term to search for in table names, column names, or descriptions"
+    )
+    schema_name: Optional[str] = Field(
+        default=None,
+        description="Schema name to search in"
+    )
+    table_name: Optional[str] = Field(
+        default=None,
+        description="Table name to search in"
+    )
+    search_type: str = Field(
+        default="all",
+        description="Type of search: 'tables', 'columns', 'descriptions', or 'all'"
+    )
+    limit: int = Field(
+        default=5,
+        description="Maximum number of results to return"
+    )
+
+
+class AbstractSchemaManagerTool(AbstractTool, ABC):
+    """
+    Abstract base for database-specific schema management tools.
+
+    Handles all schema-related operations:
+    - Schema analysis and metadata extraction
+    - Schema search and discovery
+    - Metadata caching and retrieval
+    """
+
+    name = "SchemaManagerTool"
+    description = "Comprehensive schema management for database operations"
+    args_schema = SchemaSearchArgs
+
+    def __init__(
+        self,
+        allowed_schemas: List[str],
+        engine: AsyncEngine = None,
+        metadata_cache: SchemaMetadataCache = None,
+        vector_store: Optional[AbstractStore] = None,
+        dsn: Optional[str] = None,
+        database_type: str = "postgresql",
+        session_maker: Optional[sessionmaker] = None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.dsn = dsn or asyncpg_sqlalchemy_url
+        self.database_type = database_type
+        self.allowed_schemas = allowed_schemas
+        # Database components
+        self.engine: Optional[AsyncEngine] = engine or self._get_engine(
+            dsn=self.dsn,
+            search_path=",".join(allowed_schemas)
+        )
+        # Schema-aware components
+        self.metadata_cache = metadata_cache or SchemaMetadataCache(
+            vector_store=vector_store,  # Optional - can be None
+            lru_maxsize=500,  # Large cache for many tables
+            lru_ttl=1800     # 30 minutes
+        )
+        # Vector Store:
+        self.knowledge_store = vector_store
+
+        if session_maker:
+            self.session_maker = session_maker
+        else:
+            self.session_maker = sessionmaker(
+                self.engine,
+                class_=AsyncSession,
+                expire_on_commit=False
+            )
+
+        self.logger = logging.getLogger(f"{self.__class__.__name__}")
+        self.logger.debug(
+            f"Initialized with {len(allowed_schemas)} schemas: {allowed_schemas}"
+        )
+
+    def _get_engine(self, dsn: str, search_path: str) -> AsyncEngine:
+        """Create and return an AsyncEngine for the given DSN."""
+        from sqlalchemy.ext.asyncio import create_async_engine
+        return create_async_engine(
+            dsn,
+            echo=False,
+            pool_pre_ping=True,
+            pool_recycle=3600,
+            # Multi-schema search path
+            connect_args={
+                "server_settings": {
+                    "search_path": search_path
+                }
+            }
+        )
+
+    async def _execute(
+        self,
+        search_term: str,
+        search_type: str = "all",
+        limit: int = 10
+    ) -> ToolResult:
+        """Execute schema search with the provided parameters."""
+        try:
+            raw_results = await self.search_schema(search_term, search_type, limit)
+
+            formatted_results = []
+            for table in raw_results:
+                formatted_result = await self._format_table_result(table, search_term, search_type)
+                if formatted_result:
+                    formatted_results.append(formatted_result)
+
+            return ToolResult(
+                status="success",
+                result=formatted_results,
+                metadata={
+                    "search_term": search_term,
+                    "search_type": search_type,
+                    "results_count": len(formatted_results),
+                    "searched_schemas": self.allowed_schemas
+                }
+            )
+        except Exception as e:
+            self.logger.error(f"Schema search failed: {e}")
+            return ToolResult(
+                status="error",
+                result=None,
+                error=str(e),
+                metadata={"search_term": search_term}
+            )
+
+    async def analyze_all_schemas(self) -> Dict[str, int]:
+        """
+        Analyze all allowed schemas and populate metadata cache.
+        Returns dict of schema_name -> table_count.
+        """
+        self.logger.info(f"Analyzing schemas: {self.allowed_schemas}")
+
+        results = {}
+        total_tables = 0
+
+        for schema_name in self.allowed_schemas:
+            try:
+                table_count = await self.analyze_schema(schema_name)
+                results[schema_name] = table_count
+                total_tables += table_count
+                self.logger.info(f"Schema '{schema_name}': {table_count} tables/views analyzed")
+            except Exception as e:
+                self.logger.warning(f"Failed to analyze schema '{schema_name}': {e}")
+                results[schema_name] = 0
+                continue
+
+        self.logger.info(
+            f"Analysis completed. Total: {total_tables} tables across {len(self.allowed_schemas)} schemas"
+        )
+        return results
+
+    @abstractmethod
+    async def analyze_schema(self, schema_name: str) -> int:
+        """
+        Analyze individual schema and return table count.
+        Must be implemented by database-specific subclasses.
+        """
+        pass
+
+    @abstractmethod
+    async def analyze_table(
+        self,
+        session: AsyncSession,
+        schema_name: str,
+        table_name: str,
+        table_type: str,
+        comment: Optional[str]
+    ) -> TableMetadata:
+        """
+        Analyze individual table metadata.
+        Must be implemented by database-specific subclasses.
+        """
+        pass
+
+    async def search_schema(
+        self,
+        search_term: str,
+        search_type: str = "all",
+        limit: int = 10
+    ) -> List[TableMetadata]:
+        """Search database schema - returns raw TableMetadata for agent use."""
+        self.logger.debug(
+            f"🔍 SCHEMA SEARCH: '{search_term}' (type: {search_type}, limit: {limit})")
+
+        tables = await self.metadata_cache.search_similar_tables(
+            schema_names=self.allowed_schemas,
+            query=search_term,
+            limit=limit
+        )
+
+        self.logger.info(f"✅ SEARCH COMPLETE: Found {len(tables)} results")
+        return tables
+
+    async def _format_table_result(
+        self,
+        table: TableMetadata,
+        search_term: str,
+        search_type: str
+    ) -> Optional[Dict[str, Any]]:
+        """Format a table metadata object into a search result."""
+
+        # Always return since cache already did filtering
+        return {
+            "type": "table",
+            "schema": table.schema,
+            "tablename": table.tablename,
+            "full_name": table.full_name,
+            "table_type": table.table_type,
+            "description": table.comment,
+            "columns": [
+                {
+                    "name": col.get('name'),
+                    "type": col.get('type'),
+                    "nullable": col.get('nullable', True),
+                    "description": col.get('description')
+                }
+                for col in table.columns
+            ],
+            "row_count": table.row_count,
+            "sample_data": table.sample_data[:3] if table.sample_data else [],
+            "search_term": search_term,
+            "search_type": search_type
+        }
+
+    async def get_table_details(
+        self,
+        schema: str,
+        tablename: str
+    ) -> Optional[TableMetadata]:
+        """Get detailed metadata for a specific table."""
+        if schema not in self.allowed_schemas:
+            self.logger.warning(
+                f"Schema '{schema}' not in allowed schemas: {self.allowed_schemas}"
+            )
+            return None
+
+        try:
+            return await self.metadata_cache.get_table_metadata(schema, tablename)
+        except Exception as e:
+            self.logger.error(f"Failed to get table details for {schema}.{tablename}: {e}")
+            return None
+
+    async def get_schema_overview(self, schema_name: str) -> Optional[Dict[str, Any]]:
+        """Get overview of a specific schema."""
+        if schema_name not in self.allowed_schemas:
+            return None
+
+        if schema_meta := self.metadata_cache.get_schema_overview(schema_name):
+            return {
+                "schema": schema_meta.schema,
+                "database_name": schema_meta.database_name,
+                "table_count": schema_meta.table_count,
+                "view_count": schema_meta.view_count,
+                "total_rows": schema_meta.total_rows,
+                "last_analyzed": schema_meta.last_analyzed.isoformat() if schema_meta.last_analyzed else None,
+                "tables": list(schema_meta.tables.keys()),
+                "views": list(schema_meta.views.keys())
+            }
+        return None
+
+    def get_allowed_schemas(self) -> List[str]:
+        """Get the list of schemas this tool can search."""
+        return self.allowed_schemas.copy()

parrot/tools/database/bq.py

@@ -0,0 +1,115 @@
+from typing import Optional, List
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+from datetime import datetime
+from .abstract import AbstractSchemaManagerTool
+from .models import TableMetadata
+
+class BQSchemaSearchTool(AbstractSchemaManagerTool):
+    """BigQuery-specific schema manager tool."""
+
+    name = "BQSchemaSearchTool"
+    description = "Schema management for BigQuery databases"
+
+    async def analyze_schema(self, schema_name: str) -> int:
+        """Analyze individual BigQuery schema (dataset) and return table count."""
+        async with self.session_maker() as session:
+            try:
+                # Query INFORMATION_SCHEMA.TABLES for the dataset
+                query_str = f"SELECT table_name, table_type FROM `{schema_name}.INFORMATION_SCHEMA.TABLES` WHERE table_type IN ('BASE TABLE', 'VIEW')"
+
+                result = await session.execute(text(query_str))
+                tables_data = result.fetchall()
+
+                for table_row in tables_data:
+                    table_name = table_row.table_name
+                    table_type = table_row.table_type
+                    comment = None
+
+                    try:
+                        table_metadata = await self.analyze_table(
+                            session, schema_name, table_name, table_type, comment
+                        )
+                        await self.metadata_cache.store_table_metadata(table_metadata)
+                    except Exception as e:
+                        self.logger.warning(
+                            f"Failed to analyze table {schema_name}.{table_name}: {e}"
+                        )
+
+                return len(tables_data)
+            except Exception as e:
+                self.logger.error(f"Error accessing schema {schema_name}: {e}")
+                return 0
+
+    async def analyze_table(
+        self,
+        session: AsyncSession,
+        schema_name: str,
+        table_name: str,
+        table_type: str,
+        comment: Optional[str]
+    ) -> TableMetadata:
+        """Analyze individual BigQuery table metadata."""
+
+        # Get columns
+        columns_query = f"""
+            SELECT
+                column_name,
+                data_type,
+                is_nullable,
+                NULL as column_default,
+                NULL as character_maximum_length,
+                NULL as comment
+            FROM `{schema_name}.INFORMATION_SCHEMA.COLUMNS`
+            WHERE table_name = :table_name
+            ORDER BY ordinal_position
+        """
+
+        result = await session.execute(
+            text(columns_query),
+            {"table_name": table_name}
+        )
+
+        columns = []
+        for col_row in result.fetchall():
+            columns.append({
+                "name": col_row.column_name,
+                "type": col_row.data_type,
+                "nullable": col_row.is_nullable == "YES",
+                "default": col_row.column_default,
+                "max_length": col_row.character_maximum_length,
+                "comment": col_row.comment
+            })
+
+        # Primary keys are generally not as critical/available in standard BQ metadata queries
+        # in the same way as Postgres, so leaving empty for now.
+        primary_keys = []
+
+        row_count = None
+
+        sample_data = []
+        if table_type == 'BASE TABLE':
+            try:
+                sample_query = f'SELECT * FROM `{schema_name}.{table_name}` LIMIT 3'
+                sample_result = await session.execute(text(sample_query))
+                rows = sample_result.fetchall()
+                if rows:
+                    column_names = list(sample_result.keys())
+                    sample_data = [dict(zip(column_names, row)) for row in rows]
+            except Exception as e:
+                pass
+
+        return TableMetadata(
+            schema=schema_name,
+            tablename=table_name,
+            table_type=table_type,
+            full_name=f'`{schema_name}.{table_name}`',
+            comment=comment,
+            columns=columns,
+            primary_keys=primary_keys,
+            foreign_keys=[],
+            indexes=[],
+            row_count=row_count,
+            sample_data=sample_data,
+            last_accessed=datetime.now()
+        )

parrot/tools/database/cache.py

@@ -0,0 +1,284 @@
+# ============================================================================
+# SCHEMA-AWARE METADATA CACHE
+# ============================================================================
+from typing import Dict, List, Optional
+import re
+from cachetools import TTLCache
+from navconfig.logging import logging
+from .models import SchemaMetadata, TableMetadata
+from ...stores.abstract import AbstractStore
+
+class SchemaMetadataCache:
+    """Two-tier caching: LRU (hot data) + Optional Vector Store (cold/searchable data)."""
+
+    def __init__(
+        self,
+        vector_store: Optional[AbstractStore] = None,
+        lru_maxsize: int = 500,  # Increased for large schema count
+        lru_ttl: int = 1800      # 30 minutes
+    ):
+        # Tier 1: LRU Cache for frequently accessed metadata
+        self.hot_cache = TTLCache(maxsize=lru_maxsize, ttl=lru_ttl)
+
+        # Tier 2: Optional Vector Store for similarity search and persistence
+        self.vector_store = vector_store
+        self.vector_enabled = vector_store is not None
+
+        # Schema-level caches
+        self.schema_cache: Dict[str, SchemaMetadata] = {}
+        self.table_access_stats: Dict[str, int] = {}
+
+        self.logger = logging.getLogger("Parrot.SchemaMetadataCache")
+
+        if not self.vector_enabled:
+            print("Vector store not provided - using LRU cache only")
+
+    def _table_cache_key(self, schema_name: str, table_name: str) -> str:
+        """Generate cache key for table metadata."""
+        return f"table:{schema_name}:{table_name}"
+
+    def _schema_cache_key(self, schema_name: str) -> str:
+        """Generate cache key for schema metadata."""
+        return f"schema:{schema_name}"
+
+    async def get_table_metadata(
+        self,
+        schema_name: str,
+        table_name: str
+    ) -> Optional[TableMetadata]:
+        """Get table metadata with access tracking."""
+        cache_key = self._table_cache_key(schema_name, table_name)
+
+        # Check hot cache first
+        if cache_key in self.hot_cache:
+            metadata = self.hot_cache[cache_key]
+            self._track_access(cache_key)
+            return metadata
+
+        # Check schema cache
+        if schema_name in self.schema_cache:
+            schema_meta = self.schema_cache[schema_name]
+            all_objects = schema_meta.get_all_objects()
+            if table_name in all_objects:
+                metadata = all_objects[table_name]
+                # Promote to hot cache
+                self.hot_cache[cache_key] = metadata
+                self._track_access(cache_key)
+                return metadata
+
+        # Check vector store only if enabled
+        if self.vector_enabled:
+            search_results = await self._search_vector_store(schema_name, table_name)
+            if search_results:
+                # Store in hot cache
+                self.hot_cache[cache_key] = search_results
+                self._track_access(cache_key)
+                return search_results
+
+        return None
+
+    async def store_table_metadata(self, metadata: TableMetadata):
+        """Store table metadata in available cache tiers."""
+        cache_key = self._table_cache_key(metadata.schema, metadata.tablename)
+
+        # Store in hot cache
+        self.hot_cache[cache_key] = metadata
+
+        # Update schema cache
+        if metadata.schema not in self.schema_cache:
+            self.schema_cache[metadata.schema] = SchemaMetadata(
+                schema=metadata.schema,
+                database_name="navigator",  # Could be dynamic
+                table_count=0,
+                view_count=0
+            )
+
+        schema_meta = self.schema_cache[metadata.schema]
+        if metadata.table_type == 'BASE TABLE':
+            schema_meta.tables[metadata.tablename] = metadata
+        else:
+            schema_meta.views[metadata.tablename] = metadata
+
+        # Store in vector store only if enabled
+        if self.vector_enabled:
+            await self._store_in_vector_store(metadata)
+
+    async def search_similar_tables(
+        self,
+        schema_names: List[str],
+        query: str,
+        limit: int = 5
+    ) -> List[TableMetadata]:
+        """Search for similar tables within allowed schemas."""
+        if not self.vector_enabled:
+            # Fallback: search in LRU cache and schema cache
+            return self._search_cache_only(schema_names, query, limit)
+
+        # Search with multi-schema filter
+        search_query = f"schemas:{','.join(schema_names)} {query}"
+        try:
+            results = await self.vector_store.similarity_search(
+                search_query,
+                k=limit,
+                filter={"schema_name": {"$in": schema_names}}  # Multi-schema filter
+            )
+
+            # Convert results back to TableMetadata
+            return await self._convert_vector_results(results)
+        except Exception:
+            # Fallback to cache-only search
+            return self._search_cache_only(schema_names, query, limit)
+
+    def _search_cache_only(
+        self,
+        schema_names: List[str],
+        query: str,
+        limit: int
+    ) -> List[TableMetadata]:
+        """Fallback search using only cache when vector store unavailable."""
+        results = []
+        query_lower = query.lower()
+        keywords = self._extract_search_keywords(query_lower)
+
+        self.logger.notice(
+            f"🔍 SEARCH: Extracted keywords from '{query}': {keywords}"
+        )
+
+        # Search through schema caches
+        for schema_name in schema_names:
+            if schema_name in self.schema_cache:
+                schema_meta = self.schema_cache[schema_name]
+                all_objects = schema_meta.get_all_objects()
+
+                self.logger.debug(
+                    f"🔍 SEARCHING SCHEMA '{schema_name}': {len(all_objects)} tables available"
+                )
+
+                for table_name, table_meta in all_objects.items():
+                    score = self._calculate_relevance_score(table_name, table_meta, keywords)
+
+                    if score > 0:
+                        self.logger.debug(
+                            f"🔍 MATCH: {table_name} scored {score}"
+                        )
+                        # IMPORTANT: Return the actual TableMetadata object, not a tuple
+                        results.append(table_meta)  # FIX: was (table_meta_copy, score)
+
+                        if len(results) >= limit:
+                            break
+
+                if len(results) >= limit:
+                    break
+
+        # FIX: Sort by a different method since we're not storing scores anymore
+        # For now, just return in order found (cache already does some relevance filtering)
+        final_results = results[:limit]
+
+        self.logger.info(f"🔍 SEARCH: Found {len(final_results)} results")
+
+        # DEBUG: Log what we're returning
+        for i, table in enumerate(final_results):
+            self.logger.debug(f"  Result {i+1}: {table.schema}.{table.tablename}")
+
+        return final_results
+
+    def _extract_search_keywords(self, query: str) -> List[str]:
+        """Extract meaningful keywords from a natural language query."""
+        # Convert to lowercase and remove common stop words
+        stop_words = {
+            'get', 'show', 'find', 'list', 'select', 'by', 'from', 'the', 'a', 'an',
+            'and', 'or', 'of', 'to', 'in', 'on', 'at', 'for', 'with', 'top', 'all'
+        }
+
+        # Split on non-alphanumeric characters and filter
+        words = re.findall(r'\b[a-zA-Z]+\b', query.lower())
+        return [word for word in words if word not in stop_words and len(word) > 2]
+
+    def _calculate_relevance_score(
+        self,
+        table_name: str,
+        table_meta: TableMetadata,
+        keywords: List[str]
+    ) -> float:
+        """Calculate relevance score for a table based on keywords."""
+        score = 0.0
+
+        table_name_lower = table_name.lower()
+        column_names = [col['name'].lower() for col in table_meta.columns]
+
+        for keyword in keywords:
+            keyword_lower = keyword.lower()
+
+            if keyword_lower == table_name_lower:
+                score += 10.0
+                self.logger.debug(f"Exact table match: '{keyword}' == '{table_name}'")
+
+            elif keyword_lower in table_name_lower:
+                score += 5.0
+                self.logger.debug(f"Partial table match: '{keyword}' in '{table_name}'")
+
+            elif keyword_lower in column_names:
+                score += 8.0
+                self.logger.debug(f"Column match: '{keyword}' found in columns")
+
+            elif any(keyword_lower in col_name for col_name in column_names):
+                score += 3.0
+                self.logger.debug(f"Partial column match: '{keyword}' partially matches column")
+
+            elif table_meta.comment and keyword_lower in table_meta.comment.lower():
+                score += 2.0
+                self.logger.debug(f"Comment match: '{keyword}' in table comment")
+
+        return score
+
+    def get_schema_overview(self, schema_name: str) -> Optional[SchemaMetadata]:
+        """Get complete schema overview."""
+        return self.schema_cache.get(schema_name)
+
+    def get_hot_tables(self, schema_names: List[str], limit: int = 10) -> List[tuple[str, str, int]]:
+        """Get most frequently accessed tables across allowed schemas."""
+        schema_access = []
+
+        for schema_name in schema_names:
+            schema_prefix = f"table:{schema_name}:"
+            for key, count in self.table_access_stats.items():
+                if key.startswith(schema_prefix):
+                    table_name = key.replace(schema_prefix, "")
+                    schema_access.append((schema_name, table_name, count))
+
+        return sorted(schema_access, key=lambda x: x[2], reverse=True)[:limit]
+
+    def _track_access(self, cache_key: str):
+        """Track table access for hot table identification."""
+        self.table_access_stats[cache_key] = self.table_access_stats.get(cache_key, 0) + 1
+
+    async def _search_vector_store(self, schema_name: str, table_name: str) -> Optional[TableMetadata]:
+        """Search vector store for specific table."""
+        # Implementation depends on your vector store
+        if not self.vector_enabled:
+            return None
+        return None
+
+    async def _store_in_vector_store(self, metadata: TableMetadata):
+        """Store metadata in vector store."""
+        if not self.vector_enabled:
+            return
+        try:
+            document = {
+                "content": metadata.to_yaml_context(),
+                "metadata": {
+                    "type": "table_metadata",
+                    "schema": metadata.schema,
+                    "tablename": metadata.tablename,
+                    "table_type": metadata.table_type,
+                    "full_name": metadata.full_name
+                }
+            }
+            await self.vector_store.add_documents([document])
+        except Exception:
+            pass  # Silent failure for vector store issues
+
+    async def _convert_vector_results(self, results) -> List[TableMetadata]:
+        """Convert vector store results to TableMetadata objects."""
+        # Implementation depends on your vector store format
+        return []