kailash 0.3.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

kailash/nodes/ai/agents.py

@@ -1,6 +1,6 @@
 """AI agent nodes for the Kailash SDK."""
 
-from typing import Any, Dict
+from typing import Any
 
 from kailash.nodes.base import Node, NodeParameter, register_node
 
@@ -66,7 +66,7 @@ class ChatAgent(Node):
         >>> assert "temperature" in params
     """
 
-    def get_parameters(self) -> Dict[str, NodeParameter]:
+    def get_parameters(self) -> dict[str, NodeParameter]:
         return {
             "messages": NodeParameter(
                 name="messages",
@@ -104,7 +104,7 @@ class ChatAgent(Node):
             ),
         }
 
-    def run(self, **kwargs) -> Dict[str, Any]:
+    def run(self, **kwargs) -> dict[str, Any]:
         messages = kwargs["messages"]
         model = kwargs.get("model", "default")
         temperature = kwargs.get("temperature", 0.7)
@@ -213,7 +213,7 @@ class RetrievalAgent(Node):
         >>> assert "top_k" in params
     """
 
-    def get_parameters(self) -> Dict[str, NodeParameter]:
+    def get_parameters(self) -> dict[str, NodeParameter]:
         return {
             "query": NodeParameter(
                 name="query", type=str, required=True, description="Query for retrieval"
@@ -247,7 +247,7 @@ class RetrievalAgent(Node):
             ),
         }
 
-    def run(self, **kwargs) -> Dict[str, Any]:
+    def run(self, **kwargs) -> dict[str, Any]:
         query = kwargs["query"]
         documents = kwargs["documents"]
         top_k = kwargs.get("top_k", 5)
@@ -301,7 +301,7 @@ class RetrievalAgent(Node):
 class FunctionCallingAgent(Node):
     """Agent that can call functions based on input."""
 
-    def get_parameters(self) -> Dict[str, NodeParameter]:
+    def get_parameters(self) -> dict[str, NodeParameter]:
         return {
             "query": NodeParameter(
                 name="query", type=str, required=True, description="User query"
@@ -328,7 +328,7 @@ class FunctionCallingAgent(Node):
             ),
         }
 
-    def run(self, **kwargs) -> Dict[str, Any]:
+    def run(self, **kwargs) -> dict[str, Any]:
         query = kwargs["query"]
         available_functions = kwargs["available_functions"]
         context = kwargs.get("context", {})
@@ -403,7 +403,7 @@ class FunctionCallingAgent(Node):
 class PlanningAgent(Node):
     """Agent that creates execution plans."""
 
-    def get_parameters(self) -> Dict[str, NodeParameter]:
+    def get_parameters(self) -> dict[str, NodeParameter]:
         return {
             "goal": NodeParameter(
                 name="goal", type=str, required=True, description="Goal to achieve"
@@ -430,7 +430,7 @@ class PlanningAgent(Node):
             ),
         }
 
-    def run(self, **kwargs) -> Dict[str, Any]:
+    def run(self, **kwargs) -> dict[str, Any]:
         goal = kwargs["goal"]
         available_tools = kwargs["available_tools"]
         constraints = kwargs.get("constraints", {})

kailash/nodes/ai/ai_providers.py

@@ -8,7 +8,7 @@ separation between LLM and embedding capabilities.
 
 import hashlib
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, Union
+from typing import Any
 
 
 class BaseAIProvider(ABC):
@@ -101,9 +101,8 @@ class BaseAIProvider(ABC):
         Returns:
             bool: True if the provider can be used, False otherwise
         """
-        pass
 
-    def get_capabilities(self) -> Dict[str, bool]:
+    def get_capabilities(self) -> dict[str, bool]:
         """
         Get the capabilities supported by this provider.
 
@@ -206,7 +205,7 @@ class LLMProvider(BaseAIProvider):
         self._capabilities["chat"] = True
 
     @abstractmethod
-    def chat(self, messages: List[Dict[str, str]], **kwargs) -> Dict[str, Any]:
+    def chat(self, messages: list[dict[str, str]], **kwargs) -> dict[str, Any]:
         """
         Generate a chat completion using the provider's LLM.
 
@@ -217,7 +216,6 @@ class LLMProvider(BaseAIProvider):
         Returns:
             Dict containing the standardized response
         """
-        pass
 
 
 class EmbeddingProvider(BaseAIProvider):
@@ -304,7 +302,7 @@ class EmbeddingProvider(BaseAIProvider):
         self._capabilities["embeddings"] = True
 
     @abstractmethod
-    def embed(self, texts: List[str], **kwargs) -> List[List[float]]:
+    def embed(self, texts: list[str], **kwargs) -> list[list[float]]:
         """
         Generate embeddings for a list of texts.
 
@@ -315,10 +313,9 @@ class EmbeddingProvider(BaseAIProvider):
         Returns:
             List of embedding vectors
         """
-        pass
 
     @abstractmethod
-    def get_model_info(self, model: str) -> Dict[str, Any]:
+    def get_model_info(self, model: str) -> dict[str, Any]:
         """
         Get information about a specific embedding model.
 
@@ -328,7 +325,6 @@ class EmbeddingProvider(BaseAIProvider):
         Returns:
             Dict containing model information
         """
-        pass
 
 
 class UnifiedAIProvider(LLMProvider, EmbeddingProvider):
@@ -395,7 +391,7 @@ class OllamaProvider(UnifiedAIProvider):
 
         return self._available
 
-    def chat(self, messages: List[Dict[str, str]], **kwargs) -> Dict[str, Any]:
+    def chat(self, messages: list[dict[str, str]], **kwargs) -> dict[str, Any]:
         """Generate a chat completion using Ollama.
 
         Args:
@@ -470,7 +466,7 @@ class OllamaProvider(UnifiedAIProvider):
         except Exception as e:
             raise RuntimeError(f"Ollama error: {str(e)}")
 
-    def embed(self, texts: List[str], **kwargs) -> List[List[float]]:
+    def embed(self, texts: list[str], **kwargs) -> list[list[float]]:
         """
         Generate embeddings using Ollama.
 
@@ -506,7 +502,7 @@ class OllamaProvider(UnifiedAIProvider):
         except Exception as e:
             raise RuntimeError(f"Ollama embedding error: {str(e)}")
 
-    def get_model_info(self, model: str) -> Dict[str, Any]:
+    def get_model_info(self, model: str) -> dict[str, Any]:
         """Get information about an Ollama embedding model."""
         if model in self._model_cache:
             return self._model_cache[model]
@@ -575,7 +571,7 @@ class OpenAIProvider(UnifiedAIProvider):
 
         return self._available
 
-    def chat(self, messages: List[Dict[str, str]], **kwargs) -> Dict[str, Any]:
+    def chat(self, messages: list[dict[str, str]], **kwargs) -> dict[str, Any]:
         """
         Generate a chat completion using OpenAI.
 
@@ -655,7 +651,7 @@ class OpenAIProvider(UnifiedAIProvider):
         except Exception as e:
             raise RuntimeError(f"OpenAI error: {str(e)}")
 
-    def embed(self, texts: List[str], **kwargs) -> List[List[float]]:
+    def embed(self, texts: list[str], **kwargs) -> list[list[float]]:
         """
         Generate embeddings using OpenAI.
 
@@ -699,7 +695,7 @@ class OpenAIProvider(UnifiedAIProvider):
         except Exception as e:
             raise RuntimeError(f"OpenAI embedding error: {str(e)}")
 
-    def get_model_info(self, model: str) -> Dict[str, Any]:
+    def get_model_info(self, model: str) -> dict[str, Any]:
         """Get information about an OpenAI embedding model."""
         models = {
             "text-embedding-3-large": {
@@ -775,7 +771,7 @@ class AnthropicProvider(LLMProvider):
 
         return self._available
 
-    def chat(self, messages: List[Dict[str, str]], **kwargs) -> Dict[str, Any]:
+    def chat(self, messages: list[dict[str, str]], **kwargs) -> dict[str, Any]:
         """Generate a chat completion using Anthropic."""
         try:
             import anthropic
@@ -869,7 +865,7 @@ class CohereProvider(EmbeddingProvider):
 
         return self._available
 
-    def embed(self, texts: List[str], **kwargs) -> List[List[float]]:
+    def embed(self, texts: list[str], **kwargs) -> list[list[float]]:
         """Generate embeddings using Cohere."""
         try:
             import cohere
@@ -899,7 +895,7 @@ class CohereProvider(EmbeddingProvider):
         except Exception as e:
             raise RuntimeError(f"Cohere embedding error: {str(e)}")
 
-    def get_model_info(self, model: str) -> Dict[str, Any]:
+    def get_model_info(self, model: str) -> dict[str, Any]:
         """Get information about a Cohere embedding model."""
         models = {
             "embed-english-v3.0": {
@@ -1050,7 +1046,7 @@ class HuggingFaceProvider(EmbeddingProvider):
 
         return self._available_api or self._available_local
 
-    def embed(self, texts: List[str], **kwargs) -> List[List[float]]:
+    def embed(self, texts: list[str], **kwargs) -> list[list[float]]:
         """Generate embeddings using HuggingFace."""
         model = kwargs.get("model", "sentence-transformers/all-MiniLM-L6-v2")
         use_api = kwargs.get("use_api", self._available_api)
@@ -1067,8 +1063,8 @@ class HuggingFaceProvider(EmbeddingProvider):
             )
 
     def _embed_api(
-        self, texts: List[str], model: str, normalize: bool
-    ) -> List[List[float]]:
+        self, texts: list[str], model: str, normalize: bool
+    ) -> list[list[float]]:
         """Generate embeddings using HuggingFace Inference API."""
         try:
             import os
@@ -1106,8 +1102,8 @@ class HuggingFaceProvider(EmbeddingProvider):
             raise RuntimeError(f"HuggingFace API error: {str(e)}")
 
     def _embed_local(
-        self, texts: List[str], model: str, device: str, normalize: bool
-    ) -> List[List[float]]:
+        self, texts: list[str], model: str, device: str, normalize: bool
+    ) -> list[list[float]]:
         """Generate embeddings using local HuggingFace model."""
         try:
             import torch
@@ -1165,7 +1161,7 @@ class HuggingFaceProvider(EmbeddingProvider):
         except Exception as e:
             raise RuntimeError(f"HuggingFace local error: {str(e)}")
 
-    def get_model_info(self, model: str) -> Dict[str, Any]:
+    def get_model_info(self, model: str) -> dict[str, Any]:
         """Get information about a HuggingFace embedding model."""
         models = {
             "sentence-transformers/all-MiniLM-L6-v2": {
@@ -1235,7 +1231,7 @@ class MockProvider(UnifiedAIProvider):
         """Mock provider is always available."""
         return True
 
-    def chat(self, messages: List[Dict[str, str]], **kwargs) -> Dict[str, Any]:
+    def chat(self, messages: list[dict[str, str]], **kwargs) -> dict[str, Any]:
         """Generate mock LLM response."""
         last_user_message = ""
         for msg in reversed(messages):
@@ -1272,7 +1268,7 @@ class MockProvider(UnifiedAIProvider):
             "metadata": {},
         }
 
-    def embed(self, texts: List[str], **kwargs) -> List[List[float]]:
+    def embed(self, texts: list[str], **kwargs) -> list[list[float]]:
         """Generate mock embeddings."""
         model = kwargs.get("model", "mock-embedding")
         dimensions = kwargs.get("dimensions", 1536)
@@ -1300,7 +1296,7 @@ class MockProvider(UnifiedAIProvider):
 
         return embeddings
 
-    def get_model_info(self, model: str) -> Dict[str, Any]:
+    def get_model_info(self, model: str) -> dict[str, Any]:
         """Get information about a mock embedding model."""
         models = {
             "mock-embedding-small": {"dimensions": 384, "max_tokens": 512},
@@ -1335,8 +1331,8 @@ PROVIDERS = {
 
 
 def get_provider(
-    provider_name: str, provider_type: Optional[str] = None
-) -> Union[BaseAIProvider, LLMProvider, EmbeddingProvider]:
+    provider_name: str, provider_type: str | None = None
+) -> BaseAIProvider | LLMProvider | EmbeddingProvider:
     """
     Get an AI provider instance by name.
 
@@ -1408,8 +1404,8 @@ def get_provider(
 
 
 def get_available_providers(
-    provider_type: Optional[str] = None,
-) -> Dict[str, Dict[str, Any]]:
+    provider_type: str | None = None,
+) -> dict[str, dict[str, Any]]:
     """
     Get information about all available providers.
 
@@ -1440,9 +1436,12 @@ def get_available_providers(
             capabilities = provider.get_capabilities()
 
             # Apply filter if specified
-            if provider_type == "chat" and not capabilities.get("chat"):
-                continue
-            elif provider_type == "embeddings" and not capabilities.get("embeddings"):
+            if (
+                provider_type == "chat"
+                and not capabilities.get("chat")
+                or provider_type == "embeddings"
+                and not capabilities.get("embeddings")
+            ):
                 continue
 
             results[name] = {

kailash/nodes/ai/embedding_generator.py

@@ -1,7 +1,7 @@
 """Embedding Generator node for vector embeddings with support for multiple providers."""
 
 import time
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 from kailash.nodes.base import Node, NodeParameter, register_node
 
@@ -107,7 +107,7 @@ class EmbeddingGeneratorNode(Node):
         )
     """
 
-    def get_parameters(self) -> Dict[str, NodeParameter]:
+    def get_parameters(self) -> dict[str, NodeParameter]:
         return {
             "operation": NodeParameter(
                 name="operation",
@@ -222,7 +222,7 @@ class EmbeddingGeneratorNode(Node):
             ),
         }
 
-    def run(self, **kwargs) -> Dict[str, Any]:
+    def run(self, **kwargs) -> dict[str, Any]:
         operation = kwargs["operation"]
         provider = kwargs.get("provider", "mock")
         model = kwargs.get("model", "default")
@@ -343,16 +343,16 @@ class EmbeddingGeneratorNode(Node):
 
     def _embed_single_text(
         self,
-        text: Optional[str],
+        text: str | None,
         provider: str,
         model: str,
         cache_enabled: bool,
         cache_ttl: int,
-        dimensions: Optional[int],
+        dimensions: int | None,
         normalize: bool,
         timeout: int,
         max_retries: int,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         """Generate embedding for a single text."""
         if not text:
             return {
@@ -418,18 +418,18 @@ class EmbeddingGeneratorNode(Node):
 
     def _embed_batch_texts(
         self,
-        texts: List[str],
+        texts: list[str],
         provider: str,
         model: str,
         batch_size: int,
         chunk_size: int,
         cache_enabled: bool,
         cache_ttl: int,
-        dimensions: Optional[int],
+        dimensions: int | None,
         normalize: bool,
         timeout: int,
         max_retries: int,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         """Generate embeddings for a batch of texts."""
         if not texts:
             return {
@@ -538,10 +538,10 @@ class EmbeddingGeneratorNode(Node):
 
     def _calculate_similarity(
         self,
-        embedding_1: Optional[List[float]],
-        embedding_2: Optional[List[float]],
+        embedding_1: list[float] | None,
+        embedding_2: list[float] | None,
         metric: str,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         """Calculate similarity between two embedding vectors."""
         if not embedding_1 or not embedding_2:
             return {
@@ -587,17 +587,17 @@ class EmbeddingGeneratorNode(Node):
 
     def _embed_mcp_resource(
         self,
-        resource_uri: Optional[str],
+        resource_uri: str | None,
         provider: str,
         model: str,
         chunk_size: int,
         cache_enabled: bool,
         cache_ttl: int,
-        dimensions: Optional[int],
+        dimensions: int | None,
         normalize: bool,
         timeout: int,
         max_retries: int,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         """Embed content from an MCP resource."""
         if not resource_uri:
             return {
@@ -632,7 +632,7 @@ class EmbeddingGeneratorNode(Node):
 
         return result
 
-    def _generate_mock_embedding(self, text: str, dimensions: int) -> List[float]:
+    def _generate_mock_embedding(self, text: str, dimensions: int) -> list[float]:
         """Generate a mock embedding vector based on text content."""
         import hashlib
         import random
@@ -653,10 +653,10 @@ class EmbeddingGeneratorNode(Node):
         text: str,
         provider: str,
         model: str,
-        dimensions: Optional[int],
+        dimensions: int | None,
         timeout: int,
         max_retries: int,
-    ) -> List[float]:
+    ) -> list[float]:
         """Generate embedding using external provider."""
         try:
             from .ai_providers import get_provider
@@ -702,10 +702,10 @@ class EmbeddingGeneratorNode(Node):
         text: str,
         provider: str,
         model: str,
-        dimensions: Optional[int],
+        dimensions: int | None,
         timeout: int,
         max_retries: int,
-    ) -> List[float]:
+    ) -> list[float]:
         """Fallback implementation for backward compatibility."""
         # Handle Ollama provider
         if provider == "ollama":
@@ -739,7 +739,7 @@ class EmbeddingGeneratorNode(Node):
             f"{provider}:{model}:{text}", actual_dimensions
         )
 
-    def _chunk_text(self, text: str, chunk_size: int) -> List[str]:
+    def _chunk_text(self, text: str, chunk_size: int) -> list[str]:
         """Split text into chunks based on token limit."""
         # Simple word-based chunking (real implementation would use proper tokenization)
         words = text.split()
@@ -751,7 +751,7 @@ class EmbeddingGeneratorNode(Node):
 
         return chunks or [""]
 
-    def _average_embeddings(self, embeddings: List[List[float]]) -> List[float]:
+    def _average_embeddings(self, embeddings: list[list[float]]) -> list[float]:
         """Average multiple embedding vectors."""
         if not embeddings:
             return []
@@ -765,16 +765,16 @@ class EmbeddingGeneratorNode(Node):
 
         return averaged
 
-    def _normalize_vector(self, vector: List[float]) -> List[float]:
+    def _normalize_vector(self, vector: list[float]) -> list[float]:
         """Normalize vector to unit length."""
         magnitude = sum(x * x for x in vector) ** 0.5
         if magnitude == 0:
             return vector
         return [x / magnitude for x in vector]
 
-    def _cosine_similarity(self, vec1: List[float], vec2: List[float]) -> float:
+    def _cosine_similarity(self, vec1: list[float], vec2: list[float]) -> float:
         """Calculate cosine similarity between two vectors."""
-        dot_product = sum(a * b for a, b in zip(vec1, vec2))
+        dot_product = sum(a * b for a, b in zip(vec1, vec2, strict=False))
         magnitude1 = sum(a * a for a in vec1) ** 0.5
         magnitude2 = sum(b * b for b in vec2) ** 0.5
 
@@ -783,13 +783,13 @@ class EmbeddingGeneratorNode(Node):
 
         return dot_product / (magnitude1 * magnitude2)
 
-    def _euclidean_distance(self, vec1: List[float], vec2: List[float]) -> float:
+    def _euclidean_distance(self, vec1: list[float], vec2: list[float]) -> float:
         """Calculate Euclidean distance between two vectors."""
-        return sum((a - b) ** 2 for a, b in zip(vec1, vec2)) ** 0.5
+        return sum((a - b) ** 2 for a, b in zip(vec1, vec2, strict=False)) ** 0.5
 
-    def _dot_product(self, vec1: List[float], vec2: List[float]) -> float:
+    def _dot_product(self, vec1: list[float], vec2: list[float]) -> float:
         """Calculate dot product of two vectors."""
-        return sum(a * b for a, b in zip(vec1, vec2))
+        return sum(a * b for a, b in zip(vec1, vec2, strict=False))
 
     def _interpret_similarity(self, score: float, metric: str) -> str:
         """Provide human-readable interpretation of similarity score."""
@@ -823,15 +823,14 @@ class EmbeddingGeneratorNode(Node):
         content = f"{provider}:{model}:{text}"
         return f"emb_{hashlib.md5(content.encode()).hexdigest()[:16]}"
 
-    def _get_cached_embedding(self, cache_key: str) -> Optional[Dict[str, Any]]:
+    def _get_cached_embedding(self, cache_key: str) -> dict[str, Any] | None:
         """Retrieve embedding from cache (mock implementation)."""
         # Mock cache lookup - in real implementation, use Redis or similar
         return None
 
-    def _cache_embedding(self, cache_key: str, vector: List[float], ttl: int) -> None:
+    def _cache_embedding(self, cache_key: str, vector: list[float], ttl: int) -> None:
         """Store embedding in cache (mock implementation)."""
         # Mock cache storage - in real implementation, use Redis or similar
-        pass
 
     def _estimate_embedding_cost(self, tokens: int, provider: str, model: str) -> float:
         """Estimate embedding cost based on tokens and provider pricing."""