signalpilot-ai-internal 0.10.0__py3-none-any.whl → 0.11.24__py3-none-any.whl

signalpilot_ai_internal/oauth_token_store.py

@@ -0,0 +1,141 @@
+"""
+OAuth Token Store - Secure storage for OAuth tokens in .env format
+Stores tokens at <cache_dir>/connect/.env
+(e.g., ~/Library/Caches/SignalPilotAI/connect/.env on macOS)
+"""
+
+import logging
+from typing import Any, Dict, Optional
+
+from .signalpilot_home import get_signalpilot_home
+
+logger = logging.getLogger(__name__)
+
+
+class OAuthTokenStore:
+    """
+    Secure storage for OAuth tokens using .env format.
+    Tokens are stored with server-id prefixes for namespacing.
+    """
+
+    _instance = None
+
+    def __init__(self):
+        self._home_manager = get_signalpilot_home()
+
+    @classmethod
+    def get_instance(cls) -> 'OAuthTokenStore':
+        """Get singleton instance."""
+        if cls._instance is None:
+            cls._instance = OAuthTokenStore()
+        return cls._instance
+
+    def store_tokens(self, integration_id: str, mcp_server_id: str, env_vars: Dict[str, str]):
+        """
+        Store OAuth tokens for an integration.
+
+        Args:
+            integration_id: The integration ID (e.g., 'notion', 'slack')
+            mcp_server_id: The MCP server ID associated with this integration
+            env_vars: Environment variables containing tokens
+        """
+        # Store the registry entry (mapping server_id -> integration_id)
+        self._home_manager.set_oauth_registry_entry(mcp_server_id, integration_id)
+
+        # Store the actual tokens
+        self._home_manager.set_oauth_tokens(mcp_server_id, env_vars)
+
+        logger.info(f"[OAuthTokenStore] Stored tokens for {mcp_server_id}")
+
+    def get_tokens(self, mcp_server_id: str) -> Optional[Dict[str, str]]:
+        """
+        Get OAuth tokens for an MCP server.
+
+        Args:
+            mcp_server_id: The MCP server ID
+
+        Returns:
+            Environment variables dict or None if not found
+        """
+        return self._home_manager.get_oauth_tokens(mcp_server_id)
+
+    def get_integration_id(self, mcp_server_id: str) -> Optional[str]:
+        """
+        Get the integration ID for an MCP server.
+
+        Args:
+            mcp_server_id: The MCP server ID
+
+        Returns:
+            Integration ID or None if not found
+        """
+        registry = self._home_manager.get_oauth_registry()
+        return registry.get(mcp_server_id)
+
+    def is_oauth_server(self, mcp_server_id: str) -> bool:
+        """
+        Check if an MCP server is an OAuth integration.
+
+        Args:
+            mcp_server_id: The MCP server ID
+
+        Returns:
+            True if this server has stored OAuth tokens
+        """
+        registry = self._home_manager.get_oauth_registry()
+        return mcp_server_id in registry
+
+    def remove_tokens(self, mcp_server_id: str) -> bool:
+        """
+        Remove OAuth tokens for an MCP server.
+
+        Args:
+            mcp_server_id: The MCP server ID
+
+        Returns:
+            True if tokens were removed, False if not found
+        """
+        # Remove from registry
+        self._home_manager.remove_oauth_registry_entry(mcp_server_id)
+
+        # Remove the tokens
+        result = self._home_manager.remove_oauth_tokens(mcp_server_id)
+
+        if result:
+            logger.info(f"[OAuthTokenStore] Removed tokens for {mcp_server_id}")
+
+        return result
+
+    def update_tokens(self, mcp_server_id: str, env_vars: Dict[str, str]) -> bool:
+        """
+        Update OAuth tokens for an existing MCP server.
+
+        Args:
+            mcp_server_id: The MCP server ID
+            env_vars: New environment variables containing tokens
+
+        Returns:
+            True if tokens were updated, False if server not found
+        """
+        if not self.is_oauth_server(mcp_server_id):
+            logger.warning(f"[OAuthTokenStore] Server {mcp_server_id} not found for update")
+            return False
+
+        result = self._home_manager.set_oauth_tokens(mcp_server_id, env_vars)
+        if result:
+            logger.info(f"[OAuthTokenStore] Updated tokens for {mcp_server_id}")
+        return result
+
+    def get_all_oauth_servers(self) -> Dict[str, str]:
+        """
+        Get mapping of all OAuth MCP server IDs to their integration IDs.
+
+        Returns:
+            Dict mapping mcp_server_id -> integration_id
+        """
+        return self._home_manager.get_oauth_registry()
+
+
+def get_oauth_token_store() -> OAuthTokenStore:
+    """Get the singleton instance of the OAuth token store."""
+    return OAuthTokenStore.get_instance()

signalpilot_ai_internal/schema_search_config.yml

@@ -1,32 +1,38 @@
 logging:
-  level: 'WARNING'
+  level: "WARNING"
 
 embedding:
-  location: 'memory'
-  model: 'multi-qa-MiniLM-L6-cos-v1'
-  metric: 'cosine'
+  location: "memory" # Options: "memory", "vectordb" (coming soon)
+  model: "multi-qa-MiniLM-L6-cos-v1"
+  metric: "cosine" # Options: "cosine", "euclidean", "manhattan", "dot"
   batch_size: 32
   show_progress: false
-  cache_dir: '/tmp/.schema_search_cache'
+  cache_dir: "/tmp/.schema_search_cache"
 
 chunking:
-  strategy: 'raw'
+  strategy: "raw" # Options: "raw", "llm"
   max_tokens: 256
   overlap_tokens: 50
-  model: 'gpt-4o-mini'
+  model: "gpt-4o-mini"
 
 search:
-  strategy: 'bm25'
+  # Search strategy: "semantic" (embeddings), "bm25" (BM25 lexical), "fuzzy" (fuzzy string matching), "hybrid" (semantic + bm25)
+  strategy: "bm25"
   initial_top_k: 20
   rerank_top_k: 5
-  semantic_weight: 0.67
-  hops: 1
+  semantic_weight: 0.67 # For hybrid search (bm25_weight = 1 - semantic_weight)
+  hops: 1 # Number of foreign key hops for graph expansion (0-2 recommended)
 
 reranker:
-  model: null
+  # CrossEncoder model for reranking. Set to null to disable reranking
+  model: null # "Alibaba-NLP/gte-reranker-modernbert-base"
 
 schema:
   include_columns: true
   include_indices: true
   include_foreign_keys: true
   include_constraints: true
+
+output:
+  format: "markdown" # Options: "json", "markdown"
+  limit: 5 # Default number of results to return

signalpilot_ai_internal/schema_search_service.py

@@ -21,10 +21,64 @@ class SchemaSearchHandler(APIHandler):
         for key, value in os.environ.items():
             if key.endswith("_CONNECTION_JSON") and isinstance(value, str) and value.strip().startswith("{"):
                 config = json.loads(value)
+                
+                # Special handling for Databricks
+                if config.get("type") == "databricks":
+                    return self._build_databricks_url(config)
+                
                 url = config.get("connectionUrl")
                 if url:
                     return url
         return os.environ.get("DB_URL")
+    
+    def _build_databricks_url(self, config: dict) -> Optional[str]:
+        """Build Databricks URL in the format: databricks://token:{token}@{host}?http_path={http_path}&catalog={catalog}"""
+        import re
+        
+        # Extract host from connectionUrl
+        connection_url = config.get('connectionUrl', '')
+        if not connection_url:
+            return None
+        
+        url_match = re.match(r'https?://([^/]+)', connection_url)
+        if not url_match:
+            return None
+        
+        host = url_match.group(1)
+        
+        # Get access token based on auth type
+        auth_type = config.get('authType', 'pat')
+        if auth_type == 'pat':
+            token = config.get('accessToken', '')
+        else:
+            # For service principal, we would need to get OAuth token
+            # For now, return None to fallback to other methods
+            return None
+        
+        if not token:
+            return None
+        
+        # Get HTTP path
+        http_path = config.get('warehouseHttpPath') or config.get('httpPath', '')
+        if not http_path:
+            warehouse_id = config.get('warehouseId')
+            if warehouse_id:
+                http_path = f"/sql/1.0/warehouses/{warehouse_id}"
+            else:
+                return None
+        
+        # Get catalog (optional)
+        catalog = config.get('catalog', '')
+        schema = config.get('schema', '')
+
+        # Build the URL
+        query_parts = [f"http_path={http_path}"]
+        if catalog:
+            query_parts.append(f"catalog={catalog}")
+        if schema:
+            query_parts.append(f"schema={schema}")
+
+        return f"databricks://token:{token}@{host}?{'&'.join(query_parts)}"
 
     @tornado.web.authenticated
     async def post(self):
@@ -60,18 +114,25 @@ class SchemaSearchHandler(APIHandler):
 
         if db_url_lower.startswith("snowflake://"):
             self._ensure_snowflake_dependencies()
+        elif db_url_lower.startswith("databricks://"):
+            self._ensure_databricks_dependencies()
         elif db_url_lower.startswith("postgresql") or db_url_lower.startswith("postgres") or db_url_lower.startswith("mysql+pymysql"):
             pass
         else:
             self.set_status(400)
-            self.finish(json.dumps({"error": "Schema search currently supports PostgreSQL, MySQL, or Snowflake connections"}))
+            self.finish(json.dumps({"error": "Schema search currently supports PostgreSQL, MySQL, Snowflake, or Databricks connections"}))
             return
 
         engine = None
         try:
             engine = create_engine(db_url)
-            schema_search = SchemaSearch(engine=engine, config_path=str(self.CONFIG_PATH))
-            schema_search.index()
+            schema_search = SchemaSearch(
+                engine=engine,
+                config_path=str(self.CONFIG_PATH),
+                llm_api_key=os.environ.get("SCHEMA_SEARCH_LLM_API_KEY"),
+                llm_base_url=os.environ.get("SCHEMA_SEARCH_LLM_BASE_URL")
+            )
+            schema_search.index(force=False)
 
             limit = body.get("limit")
             if limit is not None:
@@ -84,7 +145,7 @@ class SchemaSearchHandler(APIHandler):
                 result = schema_search.search(query, limit=limit)
                 query_results.append({
                     "query": query,
-                    "results": result
+                    "results": self._coerce_schema_search_result(result)
                 })
 
             self.finish(json.dumps({"results": query_results}))
@@ -101,9 +162,29 @@ class SchemaSearchHandler(APIHandler):
     def _install_package(self, package: str) -> None:
         subprocess.check_call([sys.executable, "-m", "pip", "install", package])
 
+    def _coerce_schema_search_result(self, result):
+        if hasattr(result, "to_dict"):
+            try:
+                return result.to_dict()
+            except Exception:
+                pass
+        if isinstance(result, str):
+            try:
+                return json.loads(result)
+            except json.JSONDecodeError:
+                return result
+        return result
+
     def _ensure_snowflake_dependencies(self) -> None:
         try:
             import snowflake.sqlalchemy  # type: ignore  # noqa: F401
         except ImportError:
             self._install_package("snowflake-sqlalchemy")
             import snowflake.sqlalchemy  # type: ignore  # noqa: F401
+    
+    def _ensure_databricks_dependencies(self) -> None:
+        try:
+            import databricks.sqlalchemy  # type: ignore  # noqa: F401
+        except ImportError:
+            self._install_package("databricks-sqlalchemy")
+            import databricks.sqlalchemy  # type: ignore  # noqa: F401