nia-mcp-server 1.0.18__py3-none-any.whl → 1.0.20__py3-none-any.whl

nia_mcp_server/__init__.py

@@ -2,4 +2,4 @@
 NIA MCP Server - Proxy server for NIA Knowledge Agent
 """
 
-__version__ = "1.0.18"
+__version__ = "1.0.20"

nia_mcp_server/api_client.py

@@ -28,7 +28,7 @@ class NIAApiClient:
         self.client = httpx.AsyncClient(
             headers={
                 "Authorization": f"Bearer {api_key}",
-                "User-Agent": "nia-mcp-server/1.0.0",
+                "User-Agent": "nia-mcp-server/1.0.20",
                 "Content-Type": "application/json"
             },
             timeout=720.0  # 12 minute timeout for deep research operations
@@ -552,15 +552,23 @@ class NIAApiClient:
                 for repo in repositories:
                     repo_list.append({"repository": repo})
             
-            # Build data source list  
+            # Build data source list
             source_list = []
             if data_sources:
-                for source_id in data_sources:
-                    # Handle both list of IDs and list of dicts
-                    if isinstance(source_id, dict):
-                        source_list.append(source_id)
+                for source in data_sources:
+                    # Handle flexible identifier formats:
+                    # 1. String directly (display_name, URL, or source_id) - NEW
+                    # 2. Dict with "source_id" (backwards compatible)
+                    # 3. Dict with "identifier" (new format)
+                    if isinstance(source, str):
+                        # Pass string directly - backend will resolve it
+                        source_list.append(source)
+                    elif isinstance(source, dict):
+                        # Keep dict format as-is (backwards compatible)
+                        source_list.append(source)
                     else:
-                        source_list.append({"source_id": source_id})
+                        # Convert other types to string
+                        source_list.append(str(source))
             
             # Validate at least one source
             if not repo_list and not source_list:
@@ -758,4 +766,125 @@ class NIAApiClient:
         except httpx.HTTPStatusError as e:
             raise self._handle_api_error(e)
         except Exception as e:
-            raise APIError(f"Failed to check local filesystem status: {str(e)}")
+            raise APIError(f"Failed to check local filesystem status: {str(e)}")
+
+    # ========================================================================
+    # CHROMA PACKAGE SEARCH METHODS
+    # ========================================================================
+
+    async def package_search_grep(
+        self,
+        registry: str,
+        package_name: str,
+        pattern: str,
+        version: Optional[str] = None,
+        language: Optional[str] = None,
+        filename_sha256: Optional[str] = None,
+        a: Optional[int] = None,
+        b: Optional[int] = None,
+        c: Optional[int] = None,
+        head_limit: Optional[int] = None,
+        output_mode: str = "content"
+    ) -> Dict[str, Any]:
+        """Execute grep search on package source code via Chroma."""
+        try:
+            payload = {
+                "registry": registry,
+                "package_name": package_name,
+                "pattern": pattern,
+                "version": version,
+                "language": language,
+                "filename_sha256": filename_sha256,
+                "a": a,
+                "b": b,
+                "c": c,
+                "head_limit": head_limit,
+                "output_mode": output_mode
+            }
+
+            # Remove None values
+            payload = {k: v for k, v in payload.items() if v is not None}
+
+            response = await self.client.post(
+                f"{self.base_url}/v2/package-search/grep",
+                json=payload
+            )
+            response.raise_for_status()
+            return response.json()
+
+        except httpx.HTTPStatusError as e:
+            raise self._handle_api_error(e)
+        except Exception as e:
+            raise APIError(f"Failed to search package with grep: {str(e)}")
+
+    async def package_search_hybrid(
+        self,
+        registry: str,
+        package_name: str,
+        semantic_queries: List[str],
+        version: Optional[str] = None,
+        filename_sha256: Optional[str] = None,
+        pattern: Optional[str] = None,
+        language: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """Execute hybrid semantic search on package source code via Chroma."""
+        try:
+            payload = {
+                "registry": registry,
+                "package_name": package_name,
+                "semantic_queries": semantic_queries,
+                "version": version,
+                "filename_sha256": filename_sha256,
+                "pattern": pattern,
+                "language": language
+            }
+
+            # Remove None values
+            payload = {k: v for k, v in payload.items() if v is not None}
+
+            response = await self.client.post(
+                f"{self.base_url}/v2/package-search/hybrid",
+                json=payload
+            )
+            response.raise_for_status()
+            return response.json()
+
+        except httpx.HTTPStatusError as e:
+            raise self._handle_api_error(e)
+        except Exception as e:
+            raise APIError(f"Failed to search package with hybrid search: {str(e)}")
+
+    async def package_search_read_file(
+        self,
+        registry: str,
+        package_name: str,
+        filename_sha256: str,
+        start_line: int,
+        end_line: int,
+        version: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """Read specific lines from a package file via Chroma."""
+        try:
+            payload = {
+                "registry": registry,
+                "package_name": package_name,
+                "filename_sha256": filename_sha256,
+                "start_line": start_line,
+                "end_line": end_line,
+                "version": version
+            }
+
+            # Remove None values
+            payload = {k: v for k, v in payload.items() if v is not None}
+
+            response = await self.client.post(
+                f"{self.base_url}/v2/package-search/read-file",
+                json=payload
+            )
+            response.raise_for_status()
+            return response.json()
+
+        except httpx.HTTPStatusError as e:
+            raise self._handle_api_error(e)
+        except Exception as e:
+            raise APIError(f"Failed to read package file: {str(e)}")