rakam-systems-core 0.1.1rc7__py3-none-any.whl

rakam_systems_core/ai_core/interfaces/llm_gateway.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Iterator, Optional, Type, TypeVar
+from pydantic import BaseModel
+from ..base import BaseComponent
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class LLMRequest(BaseModel):
+    """Standardized LLM request structure."""
+    system_prompt: Optional[str] = None
+    user_prompt: str
+    temperature: Optional[float] = None
+    max_tokens: Optional[int] = None
+    response_format: Optional[str] = None  # "text" or "json"
+    json_schema: Optional[Type[BaseModel]] = None
+    extra_params: Dict[str, Any] = {}
+    
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class LLMResponse(BaseModel):
+    """Standardized LLM response structure."""
+    content: str
+    parsed_content: Optional[Any] = None
+    usage: Optional[Dict[str, Any]] = None
+    model: Optional[str] = None
+    finish_reason: Optional[str] = None
+    metadata: Dict[str, Any] = {}
+    
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class LLMGateway(BaseComponent, ABC):
+    """Abstract base class for LLM gateway implementations.
+    
+    This gateway provides a standardized interface for interacting with various
+    LLM providers (OpenAI, Mistral, etc.) with support for:
+    - Text generation
+    - Structured output generation
+    - Streaming responses
+    - Token counting
+    """
+    
+    def __init__(
+        self,
+        name: str = "llm_gateway",
+        config: Optional[Dict[str, Any]] = None,
+        provider: Optional[str] = None,
+        model: Optional[str] = None,
+        default_temperature: float = 0.7,
+        api_key: Optional[str] = None,
+    ):
+        super().__init__(name, config)
+        self.provider = provider
+        self.model = model
+        self.default_temperature = default_temperature
+        self.api_key = api_key
+    
+    @abstractmethod
+    def generate(
+        self,
+        request: LLMRequest,
+    ) -> LLMResponse:
+        """Generate a response from the LLM.
+        
+        Args:
+            request: Standardized LLM request
+            
+        Returns:
+            Standardized LLM response
+        """
+        raise NotImplementedError
+    
+    @abstractmethod
+    def generate_structured(
+        self,
+        request: LLMRequest,
+        schema: Type[T],
+    ) -> T:
+        """Generate structured output conforming to a Pydantic schema.
+        
+        Args:
+            request: Standardized LLM request
+            schema: Pydantic model class to parse response into
+            
+        Returns:
+            Instance of the schema class
+        """
+        raise NotImplementedError
+    
+    def stream(
+        self,
+        request: LLMRequest,
+    ) -> Iterator[str]:
+        """Stream token/segment responses.
+        
+        Args:
+            request: Standardized LLM request
+            
+        Yields:
+            String chunks from the LLM
+        """
+        # Default implementation yields full response
+        response = self.generate(request)
+        yield response.content
+    
+    @abstractmethod
+    def count_tokens(
+        self,
+        text: str,
+        model: Optional[str] = None,
+    ) -> int:
+        """Count tokens in text.
+        
+        Args:
+            text: Text to count tokens for
+            model: Model name to determine encoding
+            
+        Returns:
+            Number of tokens
+        """
+        raise NotImplementedError
+    
+    # Legacy methods for backward compatibility
+    def run(self, prompt: str, **kwargs: Any) -> str:
+        """Legacy synchronous text completion."""
+        request = LLMRequest(
+            user_prompt=prompt,
+            system_prompt=kwargs.get("system_prompt"),
+            temperature=kwargs.get("temperature"),
+            max_tokens=kwargs.get("max_tokens"),
+            extra_params=kwargs,
+        )
+        response = self.generate(request)
+        return response.content

rakam_systems_core/ai_core/interfaces/loader.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+from ..base import BaseComponent
+
+if TYPE_CHECKING:
+    from ..vs_core import Node, VSFile
+
+
+class Loader(BaseComponent, ABC):
+    """
+    Abstract base class for document loaders.
+    
+    This class provides a common interface for loading documents into different formats.
+    Subclasses must implement the load_as_text, load_as_chunks, load_as_nodes, and load_as_vsfile methods.
+    """
+    
+    @abstractmethod
+    def run(self, source: str) -> List[str]:
+        """Load raw documents from a source (path, URL, id, etc.)."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def load_as_text(self, source: Union[str, Path]) -> str:
+        """
+        Load document and return as a single text string.
+        
+        Args:
+            source: Path to document file
+            
+        Returns:
+            Full text content of the document as a single string
+        """
+        raise NotImplementedError
+    
+    @abstractmethod
+    def load_as_chunks(self, source: Union[str, Path]) -> List[str]:
+        """
+        Load document and return as a list of text chunks.
+        
+        Args:
+            source: Path to document file
+            
+        Returns:
+            List of text chunks extracted from the document
+        """
+        raise NotImplementedError
+    
+    @abstractmethod
+    def load_as_nodes(
+        self,
+        source: Union[str, Path],
+        source_id: Optional[str] = None,
+        custom_metadata: Optional[Dict[str, Any]] = None
+    ) -> List["Node"]:
+        """
+        Load document and return as Node objects with metadata.
+        
+        Args:
+            source: Path to document file
+            source_id: Optional source identifier (defaults to file path)
+            custom_metadata: Optional custom metadata to attach to nodes
+            
+        Returns:
+            List of Node objects with text chunks and metadata
+        """
+        raise NotImplementedError
+    
+    @abstractmethod
+    def load_as_vsfile(
+        self,
+        file_path: Union[str, Path],
+        custom_metadata: Optional[Dict[str, Any]] = None
+    ) -> "VSFile":
+        """
+        Load document and return as VSFile object.
+        
+        Args:
+            file_path: Path to document file
+            custom_metadata: Optional custom metadata
+            
+        Returns:
+            VSFile object with nodes
+        """
+        raise NotImplementedError

rakam_systems_core/ai_core/interfaces/reranker.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List
+from ..base import BaseComponent
+
+class Reranker(BaseComponent, ABC):
+    @abstractmethod
+    def run(self, documents: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Reorder documents by relevance and return a new list."""
+        raise NotImplementedError

rakam_systems_core/ai_core/interfaces/retriever.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List
+from ..base import BaseComponent
+
+class Retriever(BaseComponent, ABC):
+    @abstractmethod
+    def run(self, query: str) -> List[Dict[str, Any]]:
+        """Return a list of candidate hits with metadata.
+        Output schema is intentionally loose to stay dependency‑free."""
+        raise NotImplementedError

rakam_systems_core/ai_core/interfaces/tool.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Dict, Optional, Awaitable
+import inspect
+from ..base import BaseComponent
+
+
+class ToolComponent(BaseComponent, ABC):
+    """
+    Represents a callable external or internal tool, compatible with Pydantic AI.
+    
+    This is the base class for all tools in the system. Tools can be functions
+    or callable objects that can be invoked by agents.
+    
+    Attributes:
+        name: Unique name for the tool
+        description: Human-readable description
+        function: The callable function (defaults to self.run)
+        json_schema: JSON schema for tool parameters
+        takes_ctx: Whether the tool takes context as first argument
+        is_async: Whether the function is async
+    """
+    
+    def __init__(
+        self,
+        name: str,
+        config: Optional[Dict[str, Any]] = None,
+        description: Optional[str] = None,
+        json_schema: Optional[Dict[str, Any]] = None,
+        takes_ctx: bool = False,
+    ) -> None:
+        """
+        Initialize a ToolComponent.
+        
+        Args:
+            name: Unique name for the tool
+            config: Optional configuration dictionary
+            description: Human-readable description of what the tool does
+            json_schema: JSON schema defining the tool's parameters
+            takes_ctx: Whether the tool takes context as first argument
+        """
+        super().__init__(name, config)
+        self.description = description or f"Tool: {name}"
+        self.json_schema = json_schema or self._generate_default_schema()
+        self.takes_ctx = takes_ctx
+        
+        # Set function to the run method for Pydantic AI compatibility
+        # Check if run method is async
+        self.function = self.run
+        self.is_async = inspect.iscoroutinefunction(self.run)
+    
+    def _generate_default_schema(self) -> Dict[str, Any]:
+        """
+        Generate a default JSON schema for the tool.
+        Subclasses can override this to provide custom schemas.
+        """
+        return {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "Input query or parameter for the tool"
+                }
+            },
+            "required": ["query"],
+            "additionalProperties": False,
+        }
+    
+    @abstractmethod
+    def run(self, *args: Any, **kwargs: Any) -> Any:
+        """
+        Execute the primary operation for the tool.
+        This method should be overridden by subclasses.
+        """
+        raise NotImplementedError
+    
+    @classmethod
+    def from_function(
+        cls,
+        function: Callable[..., Any],
+        name: str,
+        description: str,
+        json_schema: Dict[str, Any],
+        takes_ctx: bool = False,
+    ) -> "ToolComponent":
+        """
+        Create a ToolComponent from a function (Pydantic AI compatible).
+        
+        This factory method allows creating tool instances from standalone functions.
+        
+        Args:
+            function: The callable function
+            name: Unique name for the tool
+            description: Human-readable description
+            json_schema: JSON schema for parameters
+            takes_ctx: Whether the tool takes context as first argument
+            
+        Returns:
+            FunctionToolComponent instance wrapping the function
+        """
+        return FunctionToolComponent(
+            function=function,
+            name=name,
+            description=description,
+            json_schema=json_schema,
+            takes_ctx=takes_ctx,
+        )
+    
+    async def acall(self, *args: Any, **kwargs: Any) -> Any:
+        """
+        Async call for the tool.
+        Automatically handles both sync and async run methods.
+        """
+        if self.is_async:
+            return await self.run(*args, **kwargs)
+        else:
+            return self.run(*args, **kwargs)
+
+
+class FunctionToolComponent(ToolComponent):
+    """
+    A ToolComponent that wraps a standalone function.
+    
+    This is used internally by the from_function factory method to wrap
+    plain functions as ToolComponent instances.
+    """
+    
+    def __init__(
+        self,
+        function: Callable[..., Any],
+        name: str,
+        description: str,
+        json_schema: Dict[str, Any],
+        takes_ctx: bool = False,
+        config: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        # Store the function before calling super().__init__
+        self._wrapped_function = function
+        
+        # Initialize the parent with the description and schema
+        super().__init__(
+            name=name,
+            config=config,
+            description=description,
+            json_schema=json_schema,
+            takes_ctx=takes_ctx,
+        )
+        
+        # Override function and is_async based on the wrapped function
+        self.function = function
+        self.is_async = inspect.iscoroutinefunction(function)
+    
+    def run(self, *args: Any, **kwargs: Any) -> Any:
+        """Execute the wrapped function."""
+        return self._wrapped_function(*args, **kwargs)
+    
+    async def acall(self, *args: Any, **kwargs: Any) -> Any:
+        """Async call for the wrapped function."""
+        if self.is_async:
+            return await self._wrapped_function(*args, **kwargs)
+        else:
+            return self._wrapped_function(*args, **kwargs)

rakam_systems_core/ai_core/interfaces/tool_invoker.py

@@ -0,0 +1,260 @@
+"""
+Tool Invoker for uniform tool invocation across different modes.
+Supports both direct tool calls and MCP-based tool calls.
+"""
+from __future__ import annotations
+from typing import Any, Dict, Optional, Union
+import asyncio
+from .tool_registry import ToolRegistry, ToolMetadata, ToolMode
+from .tool import ToolComponent
+
+
+class ToolInvocationError(Exception):
+    """Base exception for tool invocation errors."""
+    pass
+
+
+class ToolNotFoundError(ToolInvocationError):
+    """Raised when a tool is not found in the registry."""
+    pass
+
+
+class MCPServerNotFoundError(ToolInvocationError):
+    """Raised when an MCP server is not available."""
+    pass
+
+
+class ToolInvoker:
+    """
+    Uniform interface for invoking tools regardless of their execution mode.
+    
+    Supports:
+    - Direct tool invocation (synchronous and asynchronous)
+    - MCP-based tool invocation via registered servers
+    - Automatic mode selection based on tool registration
+    - Error handling and validation
+    
+    Example:
+        >>> registry = ToolRegistry()
+        >>> invoker = ToolInvoker(registry)
+        >>> 
+        >>> # Invoke a direct tool
+        >>> result = await invoker.ainvoke("calculate", x=10, y=20)
+        >>> 
+        >>> # Invoke an MCP tool (with MCP server registered)
+        >>> invoker.register_mcp_server("search_server", mcp_server_instance)
+        >>> result = await invoker.ainvoke("web_search", query="Python")
+    """
+    
+    def __init__(self, registry: ToolRegistry):
+        """
+        Initialize the ToolInvoker.
+        
+        Args:
+            registry: ToolRegistry instance containing registered tools
+        """
+        self.registry = registry
+        self._mcp_servers: Dict[str, Any] = {}
+    
+    def register_mcp_server(self, server_name: str, server_instance: Any) -> None:
+        """
+        Register an MCP server for tool invocation.
+        
+        Args:
+            server_name: Name of the MCP server
+            server_instance: The MCP server instance
+        """
+        self._mcp_servers[server_name] = server_instance
+    
+    def unregister_mcp_server(self, server_name: str) -> bool:
+        """
+        Unregister an MCP server.
+        
+        Returns:
+            True if server was unregistered, False if not found
+        """
+        if server_name in self._mcp_servers:
+            del self._mcp_servers[server_name]
+            return True
+        return False
+    
+    def invoke(self, tool_name: str, **kwargs: Any) -> Any:
+        """
+        Synchronously invoke a tool by name.
+        
+        Args:
+            tool_name: Name of the tool to invoke
+            **kwargs: Arguments to pass to the tool
+        
+        Returns:
+            Result from the tool execution
+        
+        Raises:
+            ToolNotFoundError: If tool is not registered
+            ToolInvocationError: If tool invocation fails
+        """
+        metadata = self.registry.get_tool(tool_name)
+        if metadata is None:
+            raise ToolNotFoundError(f"Tool '{tool_name}' not found in registry")
+        
+        if metadata.mode == ToolMode.DIRECT:
+            return self._invoke_direct(metadata, **kwargs)
+        elif metadata.mode == ToolMode.MCP:
+            return self._invoke_mcp(metadata, **kwargs)
+        else:
+            raise ToolInvocationError(f"Unknown tool mode: {metadata.mode}")
+    
+    async def ainvoke(self, tool_name: str, **kwargs: Any) -> Any:
+        """
+        Asynchronously invoke a tool by name.
+        
+        Args:
+            tool_name: Name of the tool to invoke
+            **kwargs: Arguments to pass to the tool
+        
+        Returns:
+            Result from the tool execution
+        
+        Raises:
+            ToolNotFoundError: If tool is not registered
+            ToolInvocationError: If tool invocation fails
+        """
+        metadata = self.registry.get_tool(tool_name)
+        if metadata is None:
+            raise ToolNotFoundError(f"Tool '{tool_name}' not found in registry")
+        
+        if metadata.mode == ToolMode.DIRECT:
+            return await self._ainvoke_direct(metadata, **kwargs)
+        elif metadata.mode == ToolMode.MCP:
+            return await self._ainvoke_mcp(metadata, **kwargs)
+        else:
+            raise ToolInvocationError(f"Unknown tool mode: {metadata.mode}")
+    
+    def _invoke_direct(self, metadata: ToolMetadata, **kwargs: Any) -> Any:
+        """Invoke a direct tool synchronously."""
+        tool_instance = metadata.tool_instance
+        
+        if isinstance(tool_instance, ToolComponent):
+            # ToolComponent - call the function directly (it's now compatible with kwargs)
+            return tool_instance.function(**kwargs)
+        else:
+            raise ToolInvocationError(
+                f"Invalid tool instance type for '{metadata.name}': {type(tool_instance)}"
+            )
+    
+    async def _ainvoke_direct(self, metadata: ToolMetadata, **kwargs: Any) -> Any:
+        """Invoke a direct tool asynchronously."""
+        tool_instance = metadata.tool_instance
+        
+        if isinstance(tool_instance, ToolComponent):
+            # ToolComponent - use acall method for async support
+            return await tool_instance.acall(**kwargs)
+        else:
+            raise ToolInvocationError(
+                f"Invalid tool instance type for '{metadata.name}': {type(tool_instance)}"
+            )
+    
+    def _invoke_mcp(self, metadata: ToolMetadata, **kwargs: Any) -> Any:
+        """Invoke a tool via MCP server synchronously."""
+        if metadata.mcp_server not in self._mcp_servers:
+            raise MCPServerNotFoundError(
+                f"MCP server '{metadata.mcp_server}' not registered"
+            )
+        
+        server = self._mcp_servers[metadata.mcp_server]
+        tool_name = metadata.mcp_tool_name or metadata.name
+        
+        # Construct MCP message
+        message = {
+            'action': 'invoke_tool',
+            'tool_name': tool_name,
+            'arguments': kwargs,
+        }
+        
+        try:
+            # Send message to the tool component (receiver is the tool name)
+            result = server.send_message(
+                sender='tool_invoker',
+                receiver=tool_name,
+                message=message
+            )
+            return result
+        except Exception as e:
+            raise ToolInvocationError(
+                f"Failed to invoke MCP tool '{metadata.name}': {str(e)}"
+            ) from e
+    
+    async def _ainvoke_mcp(self, metadata: ToolMetadata, **kwargs: Any) -> Any:
+        """Invoke a tool via MCP server asynchronously."""
+        if metadata.mcp_server not in self._mcp_servers:
+            raise MCPServerNotFoundError(
+                f"MCP server '{metadata.mcp_server}' not registered"
+            )
+        
+        server = self._mcp_servers[metadata.mcp_server]
+        tool_name = metadata.mcp_tool_name or metadata.name
+        
+        # Construct MCP message
+        message = {
+            'action': 'invoke_tool',
+            'tool_name': tool_name,
+            'arguments': kwargs,
+        }
+        
+        try:
+            # Check if server has async send_message
+            if hasattr(server, 'asend_message'):
+                result = await server.asend_message(
+                    sender='tool_invoker',
+                    receiver=tool_name,
+                    message=message
+                )
+            else:
+                # Fall back to sync in thread pool
+                result = await asyncio.get_event_loop().run_in_executor(
+                    None,
+                    server.send_message,
+                    'tool_invoker',
+                    tool_name,
+                    message
+                )
+            return result
+        except Exception as e:
+            raise ToolInvocationError(
+                f"Failed to invoke MCP tool '{metadata.name}': {str(e)}"
+            ) from e
+    
+    def get_tool_info(self, tool_name: str) -> Optional[Dict[str, Any]]:
+        """
+        Get information about a registered tool.
+        
+        Args:
+            tool_name: Name of the tool
+        
+        Returns:
+            Dictionary with tool information or None if not found
+        """
+        metadata = self.registry.get_tool(tool_name)
+        if metadata is None:
+            return None
+        
+        return metadata.to_dict()
+    
+    def list_available_tools(self) -> Dict[str, Dict[str, Any]]:
+        """
+        List all available tools with their information.
+        
+        Returns:
+            Dictionary mapping tool names to their metadata
+        """
+        return {
+            metadata.name: metadata.to_dict()
+            for metadata in self.registry.get_all_tools()
+        }
+    
+    def __repr__(self) -> str:
+        return (
+            f"ToolInvoker(tools={len(self.registry)}, "
+            f"mcp_servers={len(self._mcp_servers)})"
+        )
+