kailash 0.1.4__py3-none-any.whl → 0.1.5__py3-none-any.whl

kailash/nodes/ai/agents.py

@@ -7,7 +7,64 @@ from kailash.nodes.base import Node, NodeParameter, register_node
 
 @register_node()
 class ChatAgent(Node):
-    """Chat-based AI agent node."""
+    """
+    Chat-based AI agent node for conversational interactions.
+
+    This node provides a conversational AI interface that maintains context across
+    multiple message exchanges. It supports various LLM configurations and can be
+    customized with system prompts to create specialized conversational agents for
+    different domains and use cases.
+
+    Design Philosophy:
+        The ChatAgent embodies the principle of contextual conversation, maintaining
+        the full dialogue history to provide coherent and relevant responses. It
+        abstracts away the complexities of LLM APIs while providing a consistent
+        interface for chat-based interactions across different providers.
+
+    Upstream Dependencies:
+        - User interfaces or APIs providing conversation messages
+        - Context injection systems adding relevant information
+        - Authentication systems for user-specific interactions
+        - Workflow orchestrators managing conversation flow
+
+    Downstream Consumers:
+        - Response formatting nodes processing agent outputs
+        - Logging systems recording conversations
+        - Analytics nodes analyzing interaction patterns
+        - UI components displaying chat responses
+
+    Configuration:
+        The agent can be configured with different models, temperature settings,
+        and token limits. System prompts allow specialization for specific domains
+        or behaviors without code changes.
+
+    Implementation Details:
+        - Maintains conversation history with role-based messages
+        - Prepends system prompt to establish agent behavior
+        - Currently uses mock responses for testing (production would use LLM APIs)
+        - Supports streaming responses (when integrated with real LLMs)
+        - Implements token counting for cost management
+        - Thread-safe for concurrent conversations
+
+    Error Handling:
+        - Validates message format and required fields
+        - Handles empty or malformed conversations gracefully
+        - Returns appropriate responses for API failures
+        - Implements retry logic for transient errors
+
+    Side Effects:
+        - May log conversations for debugging (configurable)
+        - Updates internal conversation state
+        - May trigger external LLM API calls (in production)
+
+    Examples:
+        >>> # Test parameter structure without constructor validation
+        >>> agent = ChatAgent.__new__(ChatAgent)
+        >>> params = agent.get_parameters()
+        >>> assert "messages" in params
+        >>> assert "model" in params
+        >>> assert "temperature" in params
+    """
 
     def get_parameters(self) -> Dict[str, NodeParameter]:
         return {
@@ -97,7 +154,64 @@ class ChatAgent(Node):
 
 @register_node()
 class RetrievalAgent(Node):
-    """Retrieval-augmented generation agent."""
+    """
+    Retrieval-augmented generation (RAG) agent for knowledge-based responses.
+
+    This node implements a RAG pipeline that retrieves relevant documents based on
+    a query and optionally generates answers using the retrieved context. It combines
+    information retrieval techniques with language generation to provide accurate,
+    grounded responses based on provided documents.
+
+    Design Philosophy:
+        The RetrievalAgent addresses the hallucination problem in LLMs by grounding
+        responses in retrieved documents. It implements a two-stage process: first
+        finding relevant information, then synthesizing an answer based only on that
+        information. This ensures factual accuracy and traceability.
+
+    Upstream Dependencies:
+        - Document ingestion nodes providing indexed content
+        - Query processing nodes enhancing user queries
+        - Embedding generation nodes (in production implementations)
+        - Vector databases or search indices
+
+    Downstream Consumers:
+        - Response formatting nodes presenting answers
+        - Citation generation nodes adding references
+        - Quality assessment nodes evaluating retrieval accuracy
+        - UI components displaying results with sources
+
+    Configuration:
+        The agent can be configured with retrieval parameters like top_k results
+        and similarity thresholds. Answer generation can be toggled based on use
+        case requirements.
+
+    Implementation Details:
+        - Currently uses keyword-based similarity (production would use embeddings)
+        - Supports various document formats (dict with content, or strings)
+        - Implements relevance scoring and ranking
+        - Filters results by similarity threshold
+        - Generates contextual answers from retrieved documents
+        - Maintains retrieval provenance for transparency
+
+    Error Handling:
+        - Handles empty document sets gracefully
+        - Validates query format and parameters
+        - Returns empty results for no matches
+        - Provides meaningful responses even with limited retrieval
+
+    Side Effects:
+        - No persistent side effects
+        - May trigger embedding generation (in production)
+        - May access external vector databases
+
+    Examples:
+        >>> # Test parameter structure without constructor validation
+        >>> agent = RetrievalAgent.__new__(RetrievalAgent)
+        >>> params = agent.get_parameters()
+        >>> assert "query" in params
+        >>> assert "documents" in params
+        >>> assert "top_k" in params
+    """
 
     def get_parameters(self) -> Dict[str, NodeParameter]:
         return {

kailash/nodes/ai/ai_providers.py

@@ -17,12 +17,69 @@ class BaseAIProvider(ABC):
 
     This abstract class defines the common interface and shared functionality
     for providers that may support LLM operations, embedding operations, or both.
+    It establishes a unified pattern for provider initialization, capability
+    detection, and error handling across different AI services.
 
     Design Philosophy:
-    - Single source of truth for provider availability
-    - Shared client management and initialization
-    - Common error handling patterns
-    - Flexible support for providers with different capabilities
+        The BaseAIProvider follows the principle of "capability-based architecture"
+        where providers declare their capabilities explicitly. This allows for
+        flexible provider implementations that may support chat, embeddings, or
+        both, while maintaining a consistent interface. The design promotes:
+        - Single source of truth for provider availability
+        - Shared client management and initialization
+        - Common error handling patterns
+        - Flexible support for providers with different capabilities
+
+    Upstream Dependencies:
+        - Configuration systems providing API keys and credentials
+        - Environment variable loaders for secure credential management
+        - Package managers ensuring required dependencies
+        - Network infrastructure for API access
+
+    Downstream Consumers:
+        - LLMAgentNode: Uses chat capabilities for conversational AI
+        - EmbeddingGeneratorNode: Uses embedding capabilities for vector generation
+        - Provider selection logic choosing appropriate implementations
+        - Error handling systems catching provider-specific exceptions
+
+    Configuration:
+        Each provider implementation handles its own configuration needs,
+        typically through environment variables or explicit parameters.
+        Common patterns include API keys, endpoints, and model selections.
+
+    Implementation Details:
+        - Lazy initialization of clients to avoid unnecessary connections
+        - Cached availability checks to reduce repeated validation
+        - Capability dictionary for runtime feature detection
+        - Abstract methods enforce implementation of core functionality
+        - Thread-safe design for concurrent usage
+
+    Error Handling:
+        - Provider availability checked before operations
+        - Graceful degradation when providers are unavailable
+        - Standardized error responses across different providers
+        - Detailed error messages for debugging
+
+    Side Effects:
+        - May establish network connections to AI services
+        - May consume API quotas during availability checks
+        - Caches client instances for performance
+
+    Examples:
+        >>> # Provider implementation
+        >>> class MyProvider(BaseAIProvider):
+        ...     def __init__(self):
+        ...         super().__init__()
+        ...         self._capabilities = {"chat": True, "embeddings": False}
+        ...
+        ...     def is_available(self) -> bool:
+        ...         # Check API key, dependencies, etc.
+        ...         return True
+        >>>
+        >>> provider = MyProvider()
+        >>> assert provider.supports_chat() == True
+        >>> assert provider.supports_embeddings() == False
+        >>> assert provider.is_available() == True
     """
 
     def __init__(self):
@@ -70,8 +127,78 @@ class LLMProvider(BaseAIProvider):
     """
     Abstract base class for providers that support LLM chat operations.
 
-    Providers that support chat operations should inherit from this class
-    and implement the chat() method.
+    This class extends BaseAIProvider to define the interface for language model
+    providers. It ensures consistent chat operation interfaces across different
+    LLM services while allowing provider-specific optimizations and features.
+
+    Design Philosophy:
+        LLMProvider standardizes the chat interface while preserving flexibility
+        for provider-specific features. It follows the OpenAI message format as
+        the de facto standard, enabling easy switching between providers. The
+        design supports both simple completions and advanced features like
+        streaming, function calling, and custom parameters.
+
+    Upstream Dependencies:
+        - BaseAIProvider: Inherits core provider functionality
+        - Message formatting systems preparing chat inputs
+        - Token counting utilities for cost management
+        - Rate limiting systems managing API quotas
+
+    Downstream Consumers:
+        - LLMAgentNode: Primary consumer for chat operations
+        - ChatAgent: Uses for conversational interactions
+        - A2AAgentNode: Leverages for agent communication
+        - Response processing nodes handling outputs
+
+    Configuration:
+        Provider-specific parameters are passed through kwargs, allowing:
+        - Model selection (model parameter)
+        - Temperature and sampling parameters
+        - Token limits and stop sequences
+        - Provider-specific features (tools, functions, etc.)
+
+    Implementation Details:
+        - Standardized message format: List[Dict[str, str]]
+        - Messages contain 'role' and 'content' fields minimum
+        - Supports system, user, and assistant roles
+        - Response format standardized across providers
+        - Streaming support through callbacks (implementation-specific)
+
+    Error Handling:
+        - Invalid message format validation
+        - API error standardization
+        - Rate limit handling with retry guidance
+        - Token limit exceeded handling
+        - Network error recovery strategies
+
+    Side Effects:
+        - API calls consume tokens/credits
+        - May log conversations for debugging
+        - Updates internal usage metrics
+        - May trigger rate limiting
+
+    Examples:
+        >>> class MyLLMProvider(LLMProvider):
+        ...     def is_available(self) -> bool:
+        ...         return True  # Check actual availability
+        ...
+        ...     def chat(self, messages, **kwargs):
+        ...         # Simulate LLM response
+        ...         return {
+        ...             "success": True,
+        ...             "content": "Response to: " + messages[-1]["content"],
+        ...             "model": kwargs.get("model", "default"),
+        ...             "usage": {"prompt_tokens": 10, "completion_tokens": 5}
+        ...         }
+        >>>
+        >>> provider = MyLLMProvider()
+        >>> messages = [
+        ...     {"role": "system", "content": "You are helpful."},
+        ...     {"role": "user", "content": "Hello!"}
+        ... ]
+        >>> response = provider.chat(messages, model="gpt-4")
+        >>> assert response["success"] == True
+        >>> assert "content" in response
     """
 
     def __init__(self):
@@ -97,8 +224,79 @@ class EmbeddingProvider(BaseAIProvider):
     """
     Abstract base class for providers that support embedding generation.
 
-    Providers that support embedding operations should inherit from this class
-    and implement the embed() and get_model_info() methods.
+    This class extends BaseAIProvider to define the interface for embedding
+    providers. It standardizes how text is converted to vector representations
+    across different embedding services while supporting provider-specific
+    optimizations and model configurations.
+
+    Design Philosophy:
+        EmbeddingProvider abstracts the complexity of different embedding models
+        and services behind a simple, consistent interface. It handles batching,
+        dimension management, and normalization while allowing providers to
+        optimize for their specific architectures. The design supports both
+        sentence and document embeddings with appropriate chunking strategies.
+
+    Upstream Dependencies:
+        - BaseAIProvider: Inherits core provider functionality
+        - Text preprocessing nodes preparing embedding inputs
+        - Chunking strategies for long documents
+        - Tokenization utilities for size management
+
+    Downstream Consumers:
+        - EmbeddingGeneratorNode: Primary consumer for vector generation
+        - Vector databases storing embeddings
+        - Similarity search implementations
+        - Clustering and classification systems
+
+    Configuration:
+        Provider-specific parameters include:
+        - Model selection for different embedding sizes/qualities
+        - Batch size limits for efficiency
+        - Normalization preferences
+        - Dimension specifications
+
+    Implementation Details:
+        - Batch processing for efficiency
+        - Automatic text truncation/chunking for model limits
+        - Vector normalization options
+        - Dimension validation and consistency
+        - Cache-friendly operations for repeated texts
+
+    Error Handling:
+        - Empty text handling
+        - Text length validation
+        - Batch size limit enforcement
+        - Model availability checking
+        - Dimension mismatch detection
+
+    Side Effects:
+        - API calls consume embedding quotas
+        - May cache embeddings for efficiency
+        - Updates usage metrics
+        - May trigger rate limiting
+
+    Examples:
+        >>> class MyEmbeddingProvider(EmbeddingProvider):
+        ...     def is_available(self) -> bool:
+        ...         return True
+        ...
+        ...     def embed(self, texts, **kwargs):
+        ...         # Simulate embedding generation
+        ...         return [[0.1, 0.2, 0.3] for _ in texts]
+        ...
+        ...     def get_model_info(self):
+        ...         return {
+        ...             "name": "my-embedding-model",
+        ...             "dimensions": 3,
+        ...             "max_tokens": 512
+        ...         }
+        >>>
+        >>> provider = MyEmbeddingProvider()
+        >>> embeddings = provider.embed(["Hello", "World"])
+        >>> assert len(embeddings) == 2
+        >>> assert len(embeddings[0]) == 3
+        >>> info = provider.get_model_info()
+        >>> assert info["dimensions"] == 3
     """
 
     def __init__(self):