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

signalpilot_ai_internal/schema_search_service.py

@@ -32,7 +32,7 @@ class SchemaSearchHandler(APIHandler):
         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}:443/{catalog}?http_path={http_path}"""
+        """Build Databricks URL in the format: databricks://token:{token}@{host}?http_path={http_path}&catalog={catalog}"""
         import re
         
         # Extract host from connectionUrl
@@ -69,14 +69,16 @@ class SchemaSearchHandler(APIHandler):
         
         # Get catalog (optional)
         catalog = config.get('catalog', '')
-        
+        schema = config.get('schema', '')
+
         # Build the URL
+        query_parts = [f"http_path={http_path}"]
         if catalog:
-            db_url = f"databricks://token:{token}@{host}:443/{catalog}?http_path={http_path}"
-        else:
-            db_url = f"databricks://token:{token}@{host}:443?http_path={http_path}"
-        
-        return db_url
+            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):
@@ -124,8 +126,13 @@ class SchemaSearchHandler(APIHandler):
         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:
@@ -138,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}))
@@ -155,6 +162,19 @@ 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