alita-sdk 0.3.257__py3-none-any.whl → 0.3.562__py3-none-any.whl

alita_sdk/community/inventory/toolkit.py

@@ -0,0 +1,173 @@
+"""
+Inventory Retrieval Toolkit.
+
+Provides LangChain-compatible toolkit for querying pre-built knowledge graphs.
+This is a pure retrieval toolkit - for ingestion, use the IngestionPipeline.
+
+The toolkit can be added to any agent to provide knowledge graph context.
+"""
+
+from typing import Any, Optional, List, Dict, Literal, ClassVar
+
+from langchain_core.tools import BaseTool
+from langchain_core.tools import BaseToolkit
+from pydantic import BaseModel, Field, ConfigDict, create_model
+
+from .retrieval import InventoryRetrievalApiWrapper
+from ...tools.base.tool import BaseAction
+from ...tools.utils import clean_string, get_max_toolkit_length
+
+
+class InventoryRetrievalToolkit(BaseToolkit):
+    """
+    Toolkit for querying pre-built Knowledge Graphs.
+    
+    This toolkit provides retrieval-only access to a knowledge graph
+    that was built using the IngestionPipeline. It can be added to
+    any agent to provide codebase context.
+    
+    Available tools:
+    - search_graph: Search entities by name, type, or layer
+    - get_entity: Get detailed entity info with relations
+    - get_entity_content: Retrieve source code via citation
+    - impact_analysis: Analyze upstream/downstream dependencies
+    - get_related_entities: Get related entities
+    - get_stats: Graph statistics
+    - get_citations: Citation summaries
+    - list_entities_by_type: List entities of a type
+    - list_entities_by_layer: List entities in a layer
+    """
+    
+    tools: List[BaseTool] = []
+    toolkit_max_length: ClassVar[int] = 0
+    
+    # All available tools in this toolkit
+    AVAILABLE_TOOLS: ClassVar[List[str]] = [
+        "search_graph",
+        "get_entity",
+        "get_entity_content",
+        "impact_analysis",
+        "get_related_entities",
+        "get_stats",
+        "get_citations",
+        "list_entities_by_type",
+        "list_entities_by_layer",
+    ]
+    
+    @staticmethod
+    def toolkit_config_schema() -> BaseModel:
+        """Return the configuration schema for this toolkit."""
+        selected_tools = {
+            x['name']: x['args_schema'].model_json_schema() 
+            for x in InventoryRetrievalApiWrapper.model_construct().get_available_tools()
+        }
+        InventoryRetrievalToolkit.toolkit_max_length = get_max_toolkit_length(selected_tools)
+        
+        model = create_model(
+            'inventory',
+            graph_path=(str, Field(
+                description="Path to the knowledge graph JSON file (required). "
+                           "The graph should be built using IngestionPipeline."
+            )),
+            base_directory=(Optional[str], Field(
+                default=None,
+                description="Base directory for local content retrieval. "
+                           "If set, get_entity_content will read from local files."
+            )),
+            selected_tools=(List[Literal[tuple(selected_tools)]], Field(
+                default=[],
+                json_schema_extra={'args_schemas': selected_tools}
+            )),
+            __config__=ConfigDict(json_schema_extra={
+                'metadata': {
+                    "label": "Knowledge Graph Retrieval",
+                    "icon_url": "inventory-icon.svg",
+                    "max_length": InventoryRetrievalToolkit.toolkit_max_length,
+                    "categories": ["knowledge management"],
+                    "extra_categories": [
+                        "knowledge graph", 
+                        "code context", 
+                        "impact analysis", 
+                        "code search"
+                    ],
+                    "description": (
+                        "Query a pre-built knowledge graph for codebase context. "
+                        "Use IngestionPipeline to build the graph first."
+                    ),
+                }
+            })
+        )
+        return model
+    
+    @classmethod
+    def get_toolkit(
+        cls,
+        selected_tools: Optional[List[str]] = None,
+        toolkit_name: str = "inventory",
+        graph_path: Optional[str] = None,
+        base_directory: Optional[str] = None,
+        source_toolkits: Optional[Dict[str, Any]] = None,
+        **kwargs
+    ) -> "InventoryRetrievalToolkit":
+        """
+        Create and return an InventoryRetrievalToolkit instance.
+        
+        Args:
+            selected_tools: List of tool names to include (None = all tools)
+            toolkit_name: Name for this toolkit instance
+            graph_path: Path to the knowledge graph JSON file
+            base_directory: Base directory for local content retrieval
+            source_toolkits: Dict of source toolkits for remote content retrieval
+            **kwargs: Additional configuration
+            
+        Returns:
+            Configured InventoryRetrievalToolkit instance
+        """
+        # Validate selected tools
+        if selected_tools is None:
+            selected_tools = cls.AVAILABLE_TOOLS
+        else:
+            selected_tools = [t for t in selected_tools if t in cls.AVAILABLE_TOOLS]
+            if not selected_tools:
+                selected_tools = cls.AVAILABLE_TOOLS
+        
+        # Create API wrapper
+        api_wrapper = InventoryRetrievalApiWrapper(
+            graph_path=graph_path or kwargs.get('graph_path', ''),
+            base_directory=base_directory or kwargs.get('base_directory'),
+            source_toolkits=source_toolkits or kwargs.get('source_toolkits', {}),
+        )
+        
+        # Get available tools from wrapper
+        available_tools = api_wrapper.get_available_tools()
+        
+        # Build tool mapping
+        tool_map = {t['name']: t for t in available_tools}
+        
+        # Use clean toolkit name for context (max 1000 chars in description)
+        toolkit_context = f" [Toolkit: {clean_string(toolkit_name, 0)}]" if toolkit_name else ''
+        
+        tools = []
+        for tool_name in selected_tools:
+            if tool_name in tool_map:
+                tool_info = tool_map[tool_name]
+                # Add toolkit context to description with character limit
+                description = tool_info['description']
+                if toolkit_context and len(description + toolkit_context) <= 1000:
+                    description = description + toolkit_context
+                tools.append(BaseAction(
+                    api_wrapper=api_wrapper,
+                    name=tool_name,
+                    description=description,
+                    args_schema=tool_info['args_schema']
+                ))
+        
+        return cls(tools=tools)
+    
+    def get_tools(self) -> List[BaseTool]:
+        """Return list of tools in this toolkit."""
+        return self.tools
+
+
+# Keep backward compatibility alias
+InventoryToolkit = InventoryRetrievalToolkit

alita_sdk/community/inventory/toolkit_utils.py

@@ -0,0 +1,176 @@
+"""
+Toolkit configuration and instantiation utilities for inventory ingestion.
+
+This module provides functions to load toolkit configurations, instantiate source
+toolkits from various sources (filesystem, GitHub, ADO), and get LLM instances
+for entity extraction.
+"""
+
+import json
+import os
+import re
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from alita_sdk.alita_client import AlitaClient
+
+
+def load_toolkit_config(toolkit_path: str) -> Dict[str, Any]:
+    """
+    Load and parse a toolkit config JSON file.
+    
+    Supports environment variable substitution for values like ${GITHUB_PAT}.
+    
+    Args:
+        toolkit_path: Path to the toolkit configuration JSON file
+        
+    Returns:
+        Dictionary containing the parsed and environment-resolved configuration
+        
+    Example:
+        >>> config = load_toolkit_config("configs/github_toolkit.json")
+        >>> config['type']
+        'github'
+    """
+    with open(toolkit_path, 'r') as f:
+        config = json.load(f)
+    
+    # Recursively resolve environment variables
+    def resolve_env_vars(obj):
+        if isinstance(obj, str):
+            # Match ${VAR_NAME} pattern
+            pattern = r'\$\{([^}]+)\}'
+            matches = re.findall(pattern, obj)
+            for var_name in matches:
+                env_value = os.environ.get(var_name, '')
+                obj = obj.replace(f'${{{var_name}}}', env_value)
+            return obj
+        elif isinstance(obj, dict):
+            return {k: resolve_env_vars(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [resolve_env_vars(item) for item in obj]
+        return obj
+    
+    return resolve_env_vars(config)
+
+
+def get_llm_for_config(
+    client: AlitaClient,
+    model: Optional[str] = None,
+    temperature: float = 0.0
+):
+    """
+    Get LLM instance from Alita client for entity extraction.
+    
+    Args:
+        client: AlitaClient instance
+        model: Model name (defaults to 'gpt-4o-mini' if not specified)
+        temperature: Temperature for the model (default 0.0 for deterministic output)
+        
+    Returns:
+        LLM instance configured with the specified model and parameters
+        
+    Example:
+        >>> client = AlitaClient(...)
+        >>> llm = get_llm_for_config(client, model='gpt-4o', temperature=0.0)
+    """
+    model_name = model or 'gpt-4o-mini'
+    
+    return client.get_llm(
+        model_name=model_name,
+        model_config={
+            'temperature': temperature,
+            'max_tokens': 4096
+        }
+    )
+
+
+def get_source_toolkit(toolkit_config: Dict[str, Any]):
+    """
+    Instantiate a source toolkit from configuration.
+    
+    Supports filesystem, GitHub, and Azure DevOps (ADO) toolkit types. For SDK-based
+    toolkits (GitHub, ADO), automatically handles configuration mapping and toolkit
+    instantiation from the registry.
+    
+    Args:
+        toolkit_config: Toolkit configuration dictionary with 'type' key
+                       and type-specific parameters
+                       
+    Returns:
+        Instantiated toolkit object ready for ingestion
+        
+    Raises:
+        ValueError: If toolkit type is unsupported or configuration is invalid
+        
+    Example:
+        >>> # Filesystem toolkit
+        >>> config = {'type': 'filesystem', 'base_path': '/path/to/code'}
+        >>> toolkit = get_source_toolkit(config)
+        
+        >>> # GitHub toolkit
+        >>> config = {
+        ...     'type': 'github',
+        ...     'github_token': 'ghp_...',
+        ...     'github_repository': 'owner/repo',
+        ...     'github_branch': 'main'
+        ... }
+        >>> toolkit = get_source_toolkit(config)
+    """
+    from alita_sdk.community.inventory.filesystem_toolkit import FilesystemToolkit
+    from alita_sdk.community.toolkits import AVAILABLE_TOOLS
+    
+    toolkit_type = toolkit_config.get('type')
+    
+    if toolkit_type == 'filesystem':
+        base_path = toolkit_config.get('base_path')
+        if not base_path:
+            raise ValueError("Filesystem toolkit requires 'base_path' configuration")
+        return FilesystemToolkit(base_path=Path(base_path))
+    
+    # Handle SDK toolkits (GitHub, ADO)
+    if toolkit_type not in AVAILABLE_TOOLS:
+        raise ValueError(
+            f"Unknown toolkit type: {toolkit_type}. "
+            f"Available types: filesystem, {', '.join(AVAILABLE_TOOLS.keys())}"
+        )
+    
+    toolkit_class = AVAILABLE_TOOLS[toolkit_type]
+    
+    # Flatten nested config if needed
+    config_for_init = {}
+    for key, value in toolkit_config.items():
+        if key == 'type':
+            continue
+        if isinstance(value, dict):
+            # Flatten nested dicts
+            config_for_init.update(value)
+        else:
+            config_for_init[key] = value
+    
+    # Map field names for specific toolkit types
+    if toolkit_type == 'github':
+        field_mapping = {
+            'github_token': 'token',
+            'github_repository': 'repository',
+            'github_branch': 'branch'
+        }
+        config_for_init = {
+            field_mapping.get(k, k): v 
+            for k, v in config_for_init.items()
+        }
+    elif toolkit_type == 'ado':
+        field_mapping = {
+            'ado_token': 'token',
+            'ado_organization': 'organization',
+            'ado_project': 'project',
+            'ado_repository': 'repository',
+            'ado_branch': 'branch'
+        }
+        config_for_init = {
+            field_mapping.get(k, k): v 
+            for k, v in config_for_init.items()
+        }
+    
+    # Instantiate toolkit
+    return toolkit_class(**config_for_init)