nia-mcp-server 1.0.23__py3-none-any.whl → 1.0.42__py3-none-any.whl

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.23",
+                "User-Agent": "nia-mcp-server/1.0.27",
                 "Content-Type": "application/json"
             },
             timeout=720.0  # 12 minute timeout for deep research operations
@@ -46,18 +46,36 @@ class NIAApiClient:
     
     def _handle_api_error(self, e: httpx.HTTPStatusError) -> APIError:
         """Convert HTTP errors to more specific API errors."""
-        error_detail = e.response.text
+        error_detail = ""
+        response_text = None
+
+        # Safely access response body; streaming responses may not be readable yet.
         try:
-            error_json = e.response.json()
-            error_detail = error_json.get("detail", error_detail)
-        except (json.JSONDecodeError, ValueError):
-            # Failed to parse JSON response, keep original error_detail
-            pass
-        
+            response_text = e.response.text
+            error_detail = response_text
+        except RuntimeError:
+            logger.warning(
+                "Unable to read streaming response body while handling HTTP error."
+            )
+            # Fall back to reason phrase if available.
+            error_detail = e.response.reason_phrase or ""
+
+        if response_text:
+            try:
+                error_json = e.response.json()
+                error_detail = error_json.get("detail", error_detail)
+                logger.debug(f"Parsed error JSON: {error_json}")
+            except (json.JSONDecodeError, ValueError) as parse_error:
+                # Failed to parse JSON response, keep original error_detail
+                logger.warning(f"Failed to parse error response as JSON: {parse_error}")
+
         status_code = e.response.status_code
-        
+
+        if not error_detail:
+            error_detail = f"HTTP {status_code} Error"
+
         # Log the full error for debugging
-        logger.error(f"API error - Status: {status_code}, Response: {error_detail}")
+        logger.error(f"API error - Status: {status_code}, Detail: {error_detail}")
         
         # Handle specific error cases
         if status_code == 401:
@@ -87,6 +105,9 @@ class NIAApiClient:
                 )
         elif status_code == 429:
             return APIError(f"Rate limit exceeded: {error_detail}", status_code, error_detail)
+        elif status_code == 400:
+            # Bad Request - return the full error detail from backend
+            return APIError(error_detail, status_code, error_detail)
         elif status_code == 404:
             return APIError(f"Resource not found: {error_detail}", status_code, error_detail)
         elif status_code == 500:
@@ -432,7 +453,104 @@ class NIAApiClient:
         except Exception as e:
             logger.error(f"Failed to rename repository: {e}")
             raise APIError(f"Failed to rename repository: {str(e)}")
-    
+            
+    async def get_github_tree(
+        self,
+        owner_repo: str,
+        branch: Optional[str] = None,
+        include_paths: Optional[List[str]] = None,
+        exclude_paths: Optional[List[str]] = None,
+        file_extensions: Optional[List[str]] = None,
+        exclude_extensions: Optional[List[str]] = None,
+        show_full_paths: bool = False
+    ) -> Dict[str, Any]:
+        """Get file tree directly from GitHub API (no FalkorDB dependency).
+
+        Args:
+            owner_repo: Repository in owner/repo format or repository ID
+            branch: Optional branch name (defaults to repository's default branch)
+            include_paths: Only include files in these paths (e.g., ["src/", "lib/"])
+            exclude_paths: Exclude files in these paths (e.g., ["node_modules/", "dist/"])
+            file_extensions: Only include these file extensions (e.g., [".py", ".js"])
+            exclude_extensions: Exclude these file extensions (e.g., [".md", ".lock"])
+            show_full_paths: Show full file paths instead of hierarchical tree
+
+        Returns:
+            GitHub tree structure with files, directories, and stats
+        """
+        try:
+            # Check if this looks like owner/repo format (contains /)
+            if '/' in owner_repo:
+                # First, get the repository ID
+                status = await self.get_repository_status(owner_repo)
+                if not status:
+                    raise APIError(f"Repository {owner_repo} not found", 404)
+
+                # Extract the repository ID from status
+                repo_id = status.get("repository_id") or status.get("id")
+                if not repo_id:
+                    # Try to get it from list as fallback
+                    repos = await self.list_repositories()
+                    for repo in repos:
+                        if repo.get("repository") == owner_repo:
+                            repo_id = repo.get("repository_id") or repo.get("id")
+                            break
+
+                if not repo_id:
+                    raise APIError(f"No repository ID found for {owner_repo}", 404)
+
+                # Get tree using the ID
+                params = {}
+                if branch:
+                    params["branch"] = branch
+                if include_paths:
+                    params["include_paths"] = ",".join(include_paths)
+                if exclude_paths:
+                    params["exclude_paths"] = ",".join(exclude_paths)
+                if file_extensions:
+                    params["file_extensions"] = ",".join(file_extensions)
+                if exclude_extensions:
+                    params["exclude_extensions"] = ",".join(exclude_extensions)
+                if show_full_paths:
+                    params["show_full_paths"] = "true"
+
+                response = await self.client.get(
+                    f"{self.base_url}/v2/repositories/{repo_id}/github-tree",
+                    params=params
+                )
+                response.raise_for_status()
+                return response.json()
+            else:
+                # Assume it's already a repository ID
+                params = {}
+                if branch:
+                    params["branch"] = branch
+                if include_paths:
+                    params["include_paths"] = ",".join(include_paths)
+                if exclude_paths:
+                    params["exclude_paths"] = ",".join(exclude_paths)
+                if file_extensions:
+                    params["file_extensions"] = ",".join(file_extensions)
+                if exclude_extensions:
+                    params["exclude_extensions"] = ",".join(exclude_extensions)
+                if show_full_paths:
+                    params["show_full_paths"] = "true"
+
+                response = await self.client.get(
+                    f"{self.base_url}/v2/repositories/{owner_repo}/github-tree",
+                    params=params
+                )
+                response.raise_for_status()
+                return response.json()
+
+        except httpx.HTTPStatusError as e:
+            raise self._handle_api_error(e)
+        except APIError:
+            raise
+        except Exception as e:
+            logger.error(f"Failed to get GitHub tree: {e}")
+            raise APIError(f"Failed to get GitHub tree: {str(e)}")
+
     # Data Source methods
     
     async def create_data_source(
@@ -449,15 +567,16 @@ class NIAApiClient:
     ) -> Dict[str, Any]:
         """Create a new documentation/web data source."""
         try:
+            effective_max_age = 3600 if max_age is None else max_age
+
             payload = {
                 "url": url,
                 "url_patterns": url_patterns or [],
-                "exclude_patterns": exclude_patterns or []
+                "exclude_patterns": exclude_patterns or [],
+                "max_age": effective_max_age
             }
             
             # Add optional parameters
-            if max_age is not None:
-                payload["max_age"] = max_age
             # Don't hardcode formats - let backend defaults apply
             # This allows screenshots to be captured by default
             if only_main_content is not None:
@@ -535,6 +654,213 @@ class NIAApiClient:
             logger.error(f"Failed to rename data source: {e}")
             raise APIError(f"Failed to rename data source: {str(e)}")
     
+    # Documentation Virtual Filesystem Methods
+    
+    async def get_doc_tree(self, source_id: str) -> Dict[str, Any]:
+        """Get filesystem tree structure of indexed documentation."""
+        try:
+            response = await self.client.get(f"{self.base_url}/v2/data-sources/{source_id}/tree")
+            response.raise_for_status()
+            return response.json()
+            
+        except httpx.HTTPStatusError as e:
+            raise self._handle_api_error(e)
+        except Exception as e:
+            logger.error(f"Failed to get documentation tree: {e}")
+            raise APIError(f"Failed to get documentation tree: {str(e)}")
+    
+    async def get_doc_ls(self, source_id: str, path: str = "/") -> Dict[str, Any]:
+        """List contents of a virtual directory in the documentation."""
+        try:
+            response = await self.client.get(
+                f"{self.base_url}/v2/data-sources/{source_id}/ls",
+                params={"path": path}
+            )
+            response.raise_for_status()
+            return response.json()
+            
+        except httpx.HTTPStatusError as e:
+            raise self._handle_api_error(e)
+        except Exception as e:
+            logger.error(f"Failed to list documentation directory: {e}")
+            raise APIError(f"Failed to list documentation directory: {str(e)}")
+    
+    async def get_doc_read(
+        self, 
+        source_id: str, 
+        path: str,
+        line_start: Optional[int] = None,
+        line_end: Optional[int] = None,
+        max_length: Optional[int] = None
+    ) -> Dict[str, Any]:
+        """Read content of a documentation page by virtual filesystem path.
+        
+        Args:
+            source_id: Data source ID
+            path: Virtual path to the page
+            line_start: Start line (1-based, inclusive)
+            line_end: End line (1-based, inclusive)
+            max_length: Max characters to return
+        """
+        try:
+            params = {"path": path}
+            if line_start is not None:
+                params["line_start"] = line_start
+            if line_end is not None:
+                params["line_end"] = line_end
+            if max_length is not None:
+                params["max_length"] = max_length
+            
+            response = await self.client.get(
+                f"{self.base_url}/v2/data-sources/{source_id}/read",
+                params=params
+            )
+            response.raise_for_status()
+            return response.json()
+            
+        except httpx.HTTPStatusError as e:
+            raise self._handle_api_error(e)
+        except Exception as e:
+            logger.error(f"Failed to read documentation file: {e}")
+            raise APIError(f"Failed to read documentation file: {str(e)}")
+    
+    async def post_doc_grep(
+        self, 
+        source_id: str, 
+        pattern: str, 
+        path: str = "/",
+        context_lines: Optional[int] = None,
+        A: Optional[int] = None,
+        B: Optional[int] = None,
+        case_sensitive: bool = False,
+        whole_word: bool = False,
+        fixed_string: bool = False,
+        max_matches_per_file: int = 10,
+        max_total_matches: int = 100,
+        output_mode: str = "content",
+        highlight: bool = False
+    ) -> Dict[str, Any]:
+        """Search documentation content with regex pattern.
+        
+        Args:
+            source_id: Data source ID
+            pattern: Regex pattern to search for
+            path: Limit search to this path prefix
+            context_lines: Lines before AND after (shorthand for A/B)
+            A: Lines after each match (like grep -A)
+            B: Lines before each match (like grep -B)
+            case_sensitive: Case-sensitive matching
+            whole_word: Match whole words only
+            fixed_string: Treat pattern as literal string
+            max_matches_per_file: Max matches per file
+            max_total_matches: Max total matches
+            output_mode: Output format ('content', 'files_with_matches', 'count')
+            highlight: Add >>markers<< around matched text
+        """
+        try:
+            body = {
+                    "pattern": pattern,
+                    "path": path,
+                "case_sensitive": case_sensitive,
+                "whole_word": whole_word,
+                "fixed_string": fixed_string,
+                "max_matches_per_file": max_matches_per_file,
+                "max_total_matches": max_total_matches,
+                "output_mode": output_mode,
+                "highlight": highlight
+            }
+            
+            # Only include optional context parameters if provided
+            if context_lines is not None:
+                body["context_lines"] = context_lines
+            if A is not None:
+                body["A"] = A
+            if B is not None:
+                body["B"] = B
+            
+            response = await self.client.post(
+                f"{self.base_url}/v2/data-sources/{source_id}/grep",
+                json=body
+            )
+            response.raise_for_status()
+            return response.json()
+            
+        except httpx.HTTPStatusError as e:
+            raise self._handle_api_error(e)
+        except Exception as e:
+            logger.error(f"Failed to grep documentation: {e}")
+            raise APIError(f"Failed to grep documentation: {str(e)}")
+    
+    async def post_code_grep(
+        self, 
+        repository: str, 
+        pattern: str, 
+        path: str = "",
+        context_lines: Optional[int] = None,
+        A: Optional[int] = None,
+        B: Optional[int] = None,
+        case_sensitive: bool = False,
+        whole_word: bool = False,
+        fixed_string: bool = False,
+        max_matches_per_file: int = 10,
+        max_total_matches: int = 100,
+        output_mode: str = "content",
+        highlight: bool = False,
+        exhaustive: bool = False
+    ) -> Dict[str, Any]:
+        """Search repository code with regex pattern.
+        
+        Args:
+            repository: Repository identifier (owner/repo format)
+            pattern: Regex pattern to search for
+            path: Limit search to this file path prefix
+            context_lines: Lines before AND after (shorthand for A/B)
+            A: Lines after each match (like grep -A)
+            B: Lines before each match (like grep -B)
+            case_sensitive: Case-sensitive matching
+            whole_word: Match whole words only
+            fixed_string: Treat pattern as literal string
+            max_matches_per_file: Max matches per file
+            max_total_matches: Max total matches
+            output_mode: Output format ('content', 'files_with_matches', 'count')
+            highlight: Add >>markers<< around matched text
+            exhaustive: When True, searches ALL chunks instead of BM25 top-k
+        """
+        try:
+            body = {
+                "pattern": pattern,
+                "path": path,
+                "case_sensitive": case_sensitive,
+                "whole_word": whole_word,
+                "fixed_string": fixed_string,
+                "max_matches_per_file": max_matches_per_file,
+                "max_total_matches": max_total_matches,
+                "output_mode": output_mode,
+                "highlight": highlight,
+                "exhaustive": exhaustive
+            }
+            
+            # Only include optional context parameters if provided
+            if context_lines is not None:
+                body["context_lines"] = context_lines
+            if A is not None:
+                body["A"] = A
+            if B is not None:
+                body["B"] = B
+            
+            response = await self.client.post(
+                f"{self.base_url}/v2/repositories/{quote(repository, safe='')}/grep",
+                json=body
+            )
+            response.raise_for_status()
+            return response.json()
+            
+        except httpx.HTTPStatusError as e:
+            raise self._handle_api_error(e)
+        except Exception as e:
+            logger.error(f"Failed to grep repository code: {e}")
+            raise APIError(f"Failed to grep repository code: {str(e)}")
+    
     async def query_unified(
         self,
         messages: List[Dict[str, str]],
@@ -569,11 +895,10 @@ class NIAApiClient:
                     else:
                         # Convert other types to string
                         source_list.append(str(source))
-            
-            # Validate at least one source
-            if not repo_list and not source_list:
-                raise Exception("No repositories or data sources specified")
-            
+
+            # NOTE: Don't validate here - let backend handle auto-hint generation
+            # The backend will generate hints if both lists are empty
+
             payload = {
                 "messages": messages,
                 "repositories": repo_list,
@@ -656,22 +981,54 @@ class NIAApiClient:
             payload = {
                 "query": query,
             }
-            
+
             if output_format:
                 payload["output_format"] = output_format
-            
+
             response = await self.client.post(
                 f"{self.base_url}/v2/deep-research",
                 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"Deep research failed: {str(e)}")
-    
+
+    async def oracle_research(
+        self,
+        query: str,
+        repositories: Optional[List[str]] = None,
+        data_sources: Optional[List[str]] = None,
+        output_format: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """Call the in-house Oracle research agent."""
+        try:
+            payload: Dict[str, Any] = {
+                "query": query,
+            }
+
+            if repositories:
+                payload["repositories"] = repositories
+            if data_sources:
+                payload["data_sources"] = data_sources
+            if output_format:
+                payload["output_format"] = output_format
+
+            response = await self.client.post(
+                f"{self.base_url}/v2/oracle",
+                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"Oracle research failed: {str(e)}")
+
     async def get_source_content(
         self,
         source_type: str,
@@ -723,76 +1080,6 @@ class NIAApiClient:
             raise self._handle_api_error(e)
         except Exception as e:
             raise APIError(f"Failed to submit bug report: {str(e)}")
-    
-    async def index_local_filesystem(
-        self,
-        directory_path: str,
-        inclusion_patterns: List[str] = None,
-        exclusion_patterns: List[str] = None,
-        max_file_size_mb: int = 50
-    ) -> Dict[str, Any]:
-        """Index a local filesystem directory."""
-        try:
-            payload = {
-                "directory_path": directory_path,
-                "inclusion_patterns": inclusion_patterns or [],
-                "exclusion_patterns": exclusion_patterns or [],
-                "max_file_size_mb": max_file_size_mb
-            }
-            
-            response = await self.client.post(
-                f"{self.base_url}/v2/local-filesystem",
-                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 index local filesystem: {str(e)}")
-    
-    async def scan_local_filesystem(
-        self,
-        directory_path: str,
-        inclusion_patterns: List[str] = None,
-        exclusion_patterns: List[str] = None,
-        max_file_size_mb: int = 50
-    ) -> Dict[str, Any]:
-        """Scan a local filesystem directory to preview what would be indexed."""
-        try:
-            payload = {
-                "directory_path": directory_path,
-                "inclusion_patterns": inclusion_patterns or [],
-                "exclusion_patterns": exclusion_patterns or [],
-                "max_file_size_mb": max_file_size_mb
-            }
-            
-            response = await self.client.post(
-                f"{self.base_url}/v2/local-filesystem/scan",
-                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 scan local filesystem: {str(e)}")
-    
-    async def check_local_filesystem_status(self, source_id: str) -> Dict[str, Any]:
-        """Check the indexing status of a local filesystem source."""
-        try:
-            response = await self.client.get(
-                f"{self.base_url}/v2/local-filesystem/{source_id}"
-            )
-            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 check local filesystem status: {str(e)}")
 
     # ========================================================================
     # CHROMA PACKAGE SEARCH METHODS
@@ -926,9 +1213,15 @@ class NIAApiClient:
         content: str,
         agent_source: str,
         tags: List[str] = None,
-        metadata: Dict[str, Any] = None
+        metadata: Dict[str, Any] = None,
+        nia_references: Optional[Dict[str, Any]] = None,
+        edited_files: Optional[List[Dict[str, Any]]] = None,
+        workspace_metadata: Optional[Dict[str, Any]] = None,
+        file_metadata: Optional[Dict[str, Any]] = None,
+        workspace_override: Optional[str] = None,
+        cwd: Optional[str] = None
     ) -> Dict[str, Any]:
-        """Save a conversation context for cross-agent sharing."""
+        """Save a conversation context for cross-agent sharing with workspace awareness."""
         try:
             payload = {
                 "title": title,
@@ -939,6 +1232,22 @@ class NIAApiClient:
                 "metadata": metadata or {}
             }
 
+            # Add new structured fields if provided
+            if nia_references is not None:
+                payload["nia_references"] = nia_references
+            if edited_files is not None:
+                payload["edited_files"] = edited_files
+
+            # Add workspace-aware fields
+            if workspace_metadata is not None:
+                payload["workspace_metadata"] = workspace_metadata
+            if file_metadata is not None:
+                payload["file_metadata"] = file_metadata
+            if workspace_override is not None:
+                payload["workspace_override"] = workspace_override
+            if cwd is not None:
+                payload["cwd"] = cwd
+
             response = await self.client.post(
                 f"{self.base_url}/v2/contexts",
                 json=payload
@@ -956,9 +1265,14 @@ class NIAApiClient:
         limit: int = 20,
         offset: int = 0,
         tags: Optional[str] = None,
-        agent_source: Optional[str] = None
+        agent_source: Optional[str] = None,
+        scope: Optional[str] = None,
+        workspace: Optional[str] = None,
+        directory: Optional[str] = None,
+        file_overlap: Optional[str] = None,
+        cwd: Optional[str] = None
     ) -> Dict[str, Any]:
-        """List user's conversation contexts with pagination and filtering."""
+        """List user's conversation contexts with pagination, filtering, and workspace awareness."""
         try:
             params = {
                 "limit": limit,
@@ -970,6 +1284,18 @@ class NIAApiClient:
             if agent_source:
                 params["agent_source"] = agent_source
 
+            # Add workspace-aware filters
+            if scope:
+                params["scope"] = scope
+            if workspace:
+                params["workspace"] = workspace
+            if directory:
+                params["directory"] = directory
+            if file_overlap:
+                params["file_overlap"] = file_overlap
+            if cwd:
+                params["cwd"] = cwd
+
             response = await self.client.get(
                 f"{self.base_url}/v2/contexts",
                 params=params
@@ -1054,7 +1380,7 @@ class NIAApiClient:
         tags: Optional[str] = None,
         agent_source: Optional[str] = None
     ) -> Dict[str, Any]:
-        """Search conversation contexts by content, title, or summary."""
+        """Search conversation contexts by content, title, or summary (keyword search)."""
         try:
             params = {
                 "q": query,
@@ -1076,4 +1402,90 @@ class NIAApiClient:
         except httpx.HTTPStatusError as e:
             raise self._handle_api_error(e)
         except Exception as e:
-            raise APIError(f"Failed to search contexts: {str(e)}")
+            raise APIError(f"Failed to search contexts: {str(e)}")
+
+    async def search_contexts_semantic(
+        self,
+        query: str,
+        limit: int = 20,
+        organization_id: Optional[str] = None,
+        cwd: Optional[str] = None,
+        include_highlights: bool = True,
+        workspace_filter: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """Search conversation contexts using semantic search (vector + BM25 hybrid)."""
+        try:
+            params = {
+                "q": query,
+                "limit": limit,
+                "include_highlights": include_highlights
+            }
+
+            if organization_id:
+                params["organization_id"] = organization_id
+            if cwd:
+                params["cwd"] = cwd
+            if workspace_filter:
+                params["workspace_filter"] = workspace_filter
+
+            response = await self.client.get(
+                f"{self.base_url}/v2/contexts/semantic-search",
+                params=params
+            )
+            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 contexts semantically: {str(e)}")
+
+    # =========================================================================
+    # Universal Search
+    # =========================================================================
+
+    async def universal_search(
+        self,
+        query: str,
+        top_k: int = 20,
+        include_repos: bool = True,
+        include_docs: bool = True,
+        alpha: float = 0.7,
+        compress_output: bool = False
+    ) -> Dict[str, Any]:
+        """
+        Search across ALL indexed public sources using TurboPuffer hybrid search.
+        
+        Args:
+            query: Natural language search query
+            top_k: Total number of results to return (default: 20, max: 100)
+            include_repos: Include repository sources (default: True)
+            include_docs: Include documentation sources (default: True)
+            alpha: Weight for vector search vs BM25 (default: 0.7 = 70% vector)
+            compress_output: Use AI (Gemini Flash) to compress results into concise answer
+            
+        Returns:
+            Dict with results, sources_searched, query_time_ms, optional errors, and optional answer
+        """
+        try:
+            payload = {
+                "query": query,
+                "top_k": top_k,
+                "include_repos": include_repos,
+                "include_docs": include_docs,
+                "alpha": alpha,
+                "compress_output": compress_output
+            }
+            
+            response = await self.client.post(
+                f"{self.base_url}/v2/universal-search",
+                json=payload
+            )
+            response.raise_for_status()
+            return response.json()
+            
+        except httpx.HTTPStatusError as e:
+            raise self._handle_api_error(e)
+        except Exception as e:
+            logger.error(f"Universal search failed: {e}")
+            raise APIError(f"Failed to perform universal search: {str(e)}")