tinyflow-llm 0.1.0__py3-none-any.whl

tinyflow/embeddings/factory.py

@@ -0,0 +1,140 @@
+"""Embedding Factory
+
+Select different Embedding implementations based on configuration
+
+Supports:
+- openai: OpenAI API Embedding
+- local / sentence-transformers: Local Sentence Transformers model
+
+Configuration precedence: Parameters > settings > Environment variables > Default values
+"""
+
+from typing import Optional
+
+from tinyflow.config.helpers import filter_none_kwargs, get_config_value
+from tinyflow.config.settings import settings
+from tinyflow.embeddings.base import BaseEmbedding
+from tinyflow.embeddings.openai_embedding import OpenAIEmbedding
+
+
+class EmbeddingFactory:
+    """Embedding Factory class.
+
+    Configuration Precedence:
+        1. Explicit parameters (highest)
+        2. Environment variables (EMBEDDING_PROVIDER)
+        3. Settings (from .env or defaults)
+        4. Provider defaults (lowest)
+    """
+
+    @staticmethod
+    def create(
+        provider: Optional[str] = None,
+        **kwargs,
+    ) -> BaseEmbedding:
+        """Create Embedding provider instance.
+
+        Args:
+            provider: Provider name ('openai', 'local').
+                     Falls back to EMBEDDING_PROVIDER env var -> settings -> default.
+            **kwargs: Provider-specific configuration parameters.
+                     - OpenAI: api_key, model, base_url
+                     - Local: model_path, device
+
+        Returns:
+            Configured embedding provider instance.
+
+        Raises:
+            ValueError: If provider is not specified or supported.
+        """
+        provider = get_config_value(
+            provider,
+            env_key="EMBEDDING_PROVIDER",
+            settings_field="EMBEDDING_PROVIDER",
+            default_value="openai",
+        )
+
+        if not provider:
+            raise ValueError("Embedding Provider must be specified")
+
+        provider = provider.lower()
+
+        if provider == "openai":
+            openai_kwargs = {
+                "api_key": get_config_value(
+                    kwargs.get("api_key"), "EMBEDDING_API_KEY", "EMBEDDING_API_KEY"
+                ),
+                "base_url": get_config_value(
+                    kwargs.get("base_url"), "EMBEDDING_BASE_URL", "EMBEDDING_BASE_URL"
+                ),
+                "model": get_config_value(
+                    kwargs.get("model"), "EMBEDDING_MODEL", "EMBEDDING_MODEL"
+                ),
+            }
+            # Use LLM_API_KEY as fallback if EMBEDDING_API_KEY is not set
+            if not openai_kwargs["api_key"]:
+                openai_kwargs["api_key"] = get_config_value(
+                    None, "LLM_API_KEY", "LLM_API_KEY"
+                )
+
+            openai_kwargs = filter_none_kwargs(**openai_kwargs)
+            return OpenAIEmbedding(**openai_kwargs)
+
+        elif provider in ["local", "sentence-transformers"]:
+            # Lazy import to avoid hard dependency on torch/sentence_transformers
+            from tinyflow.embeddings.local_embedding import SentenceTransformerEmbedding
+
+            local_kwargs = {
+                "model_path": get_config_value(
+                    kwargs.get("model_path"),
+                    "EMBEDDING_MODEL_PATH",
+                    "EMBEDDING_MODEL_PATH",
+                ),
+                "device": get_config_value(
+                    kwargs.get("device"), "EMBEDDING_MODEL_DEVICE", "EMBEDDING_MODEL_DEVICE"
+                ),
+            }
+            local_kwargs = filter_none_kwargs(**local_kwargs)
+            return SentenceTransformerEmbedding(**local_kwargs)
+
+        else:
+            raise ValueError(
+                f"Unsupported embedding provider: {provider}. Supported: openai, local"
+            )
+
+        if not provider:
+            raise ValueError("Embedding Provider must be specified")
+
+        provider = provider.lower()
+
+        if provider == "openai":
+            openai_kwargs = {
+                "api_key": get_config_value(
+                    kwargs.get("api_key"), "EMBEDDING_API_KEY", "EMBEDDING_API_KEY"
+                ),
+                "base_url": get_config_value(
+                    kwargs.get("base_url"), "EMBEDDING_BASE_URL", "EMBEDDING_BASE_URL"
+                ),
+                "model": get_config_value(
+                    kwargs.get("model"), "EMBEDDING_MODEL", "EMBEDDING_MODEL"
+                ),
+            }
+            openai_kwargs = filter_none_kwargs(**openai_kwargs)
+            return OpenAIEmbedding(**openai_kwargs)
+
+        elif provider in ("local", "sentence-transformers"):
+            local_kwargs = {
+                "model_path": get_config_value(
+                    kwargs.get("model_path"), "LOCAL_MODEL", "EMBEDDING_MODEL_PATH"
+                ),
+                "device": get_config_value(
+                    kwargs.get("device"), "DEVICE", "EMBEDDING_MODEL_DEVICE"
+                ),
+            }
+            local_kwargs = filter_none_kwargs(**local_kwargs)
+            return SentenceTransformerEmbedding(**local_kwargs)
+        else:
+            raise ValueError(
+                f"Unsupported embedding provider: {provider}. "
+                "Supported providers: openai, local, sentence-transformers"
+            )

tinyflow/embeddings/local_embedding.py

@@ -0,0 +1,65 @@
+"""Sentence Transformers Local Embedding Implementation
+
+Use open source models for local embedding
+
+Note: This module requires torch and sentence_transformers.
+If installation fails, use tinyflow.embeddings.openai_embedding.OpenAIEmbedding instead.
+"""
+
+from typing import Any, List, Optional
+
+from sentence_transformers import SentenceTransformer
+
+from tinyflow.config.settings import settings
+from tinyflow.embeddings.base import BaseEmbedding
+
+
+class SentenceTransformerEmbedding(BaseEmbedding):
+    """Local Sentence Transformers Embedding Implementation
+
+    Supports:
+    - all-MiniLM-L6-v2 (384 dimensions)
+    - all-mpnet-base-v2 (768 dimensions)
+    - paraphrase-multilingual-MiniLM-L12-v2 (768 dimensions)
+
+    Configuration precedence: Parameters > settings > Environment variables > Default values
+    """
+
+    def __init__(
+        self,
+        model_path: Optional[str] = None,
+        device: Optional[str] = None,
+    ):
+        """Initialize Sentence Transformer Embedding
+
+        Args:
+            model_path: Model path or name (defaults to settings or environment variables)
+            device: Computation device (auto/cpu/cuda/mps, defaults to settings or environment variables)
+
+        Raises:
+            ImportError: If torch/sentence_transformers is not installed
+        """
+        model_path = model_path or settings.EMBEDDING_MODEL_PATH
+        device = device or settings.EMBEDDING_MODEL_DEVICE
+
+        self.model = SentenceTransformer(model_path)
+        self.device = device
+
+    def embed(self, texts: List[str]) -> List[List[float]]:
+        """Convert text list to vectors"""
+        self.model.to(self.device)
+
+        embeddings: Any = self.model.encode(
+            texts,
+            convert_to_numpy=True,
+            show_progress_bar=False,
+        )
+
+        return list(embeddings)
+
+    def get_dimension(self) -> int:
+        """Get vector dimension"""
+        dim = self.model.get_sentence_embedding_dimension()
+        if dim is None:
+            raise ValueError("Unable to determine embedding dimension")
+        return dim

tinyflow/embeddings/openai_embedding.py

@@ -0,0 +1,70 @@
+"""OpenAI Embedding Implementation
+
+Use OpenAI API for text embedding
+"""
+
+from typing import List, Optional
+
+from openai import AsyncOpenAI
+
+from tinyflow.config.settings import settings
+from tinyflow.embeddings.base import BaseEmbedding
+
+
+class OpenAIEmbedding(BaseEmbedding):
+    """OpenAI Embedding Implementation
+
+    Supports:
+    - text-embedding-3-small
+    - text-embedding-3-large
+    - text-embedding-ada-002
+
+    Configuration precedence: Parameters > settings > Environment variables > Default values
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        model: Optional[str] = None,
+    ):
+        """Initialize OpenAI Embedding
+
+        Args:
+            api_key: OpenAI API key (defaults to settings or environment variables)
+            base_url: API base URL (defaults to settings or environment variables)
+            model: Model name (defaults to settings or environment variables)
+        """
+        api_key = api_key or settings.EMBEDDING_API_KEY or settings.LLM_API_KEY or ""
+        base_url = base_url or settings.EMBEDDING_BASE_URL or settings.LLM_BASE_URL
+        model = model or settings.EMBEDDING_MODEL
+
+        self.client = AsyncOpenAI(api_key=api_key, base_url=base_url)
+        self.model = model
+
+    async def embed(self, texts: List[str]) -> List[List[float]]:
+        """Convert text list to vectors"""
+        # Batch processing, max 2048 texts per batch
+        embeddings = []
+        batch_size = 2048
+
+        for i in range(0, len(texts), batch_size):
+            batch = texts[i : i + batch_size]
+            response = await self.client.embeddings.create(
+                input=batch,
+                model=self.model,
+            )
+
+            batch_embeddings = [item.embedding for item in response.data]
+            embeddings.extend(batch_embeddings)
+
+        return embeddings
+
+    def get_dimension(self) -> int:
+        """Get vector dimension"""
+        model_dimensions = {
+            "text-embedding-3-small": 1536,
+            "text-embedding-3-large": 3072,
+            "text-embedding-ada-002": 1536,
+        }
+        return model_dimensions.get(self.model, 1536)

tinyflow/memory/__init__.py

@@ -0,0 +1,5 @@
+from .base import BaseMemory
+from .simple import SimpleMemory
+from .vector import VectorMemory
+
+__all__ = ["BaseMemory", "SimpleMemory", "VectorMemory"]

tinyflow/memory/base.py

@@ -0,0 +1,16 @@
+import abc
+from typing import List
+
+
+class BaseMemory(abc.ABC):
+    @abc.abstractmethod
+    async def add(self, text: str):
+        pass
+
+    @abc.abstractmethod
+    async def search(self, query: str, limit: int = 3) -> List[str]:
+        pass
+
+    @abc.abstractmethod
+    async def clear(self):
+        pass

tinyflow/memory/simple.py

@@ -0,0 +1,34 @@
+from typing import List
+
+from .base import BaseMemory
+
+
+class SimpleMemory(BaseMemory):
+    def __init__(self):
+        self.memories: List[str] = []
+
+    async def add(self, text: str):
+        if text not in self.memories:
+            self.memories.append(text)
+
+    async def search(self, query: str, limit: int = 3) -> List[str]:
+        if not self.memories:
+            return []
+
+        if len(self.memories) <= limit:
+            return self.memories
+
+        query_words = set(query.lower().split())
+
+        scored_memories = []
+        for i, text in enumerate(self.memories):
+            text_words = set(text.lower().split())
+            score = len(query_words.intersection(text_words))
+            score += i / 1000000
+            scored_memories.append((score, text))
+
+        scored_memories.sort(key=lambda x: x[0], reverse=True)
+        return [text for score, text in scored_memories[:limit]]
+
+    async def clear(self):
+        self.memories = []

tinyflow/memory/vector.py

@@ -0,0 +1,26 @@
+import logging
+from typing import List
+
+from tinyflow.vector.base import BaseVectorDB
+
+from .base import BaseMemory
+
+logger = logging.getLogger("tinyflow.memory.vector")
+
+
+class VectorMemory(BaseMemory):
+    def __init__(self, vector_db: BaseVectorDB):
+        self.db = vector_db
+
+    async def add(self, text: str):
+        logger.debug(f"Adding to vector memory: {text[:50]}...")
+        await self.db.add([text])
+
+    async def search(self, query: str, limit: int = 3) -> List[str]:
+        logger.debug(f"Searching vector memory for: {query}")
+        results = await self.db.search(query, limit=limit)
+        return [r["text"] for r in results]
+
+    async def clear(self):
+        logger.info("Clearing vector memory")
+        await self.db.clear()

tinyflow/providers/anthropic_llm.py

@@ -0,0 +1,132 @@
+import logging
+import os
+from typing import AsyncGenerator, List, Optional
+
+import anthropic
+from tenacity import retry, stop_after_attempt, wait_exponential
+
+from tinyflow.core.tools import Tool
+from tinyflow.core.types import (
+    LLMResponse,
+    Message,
+    StreamChunk,
+    StreamPart,
+    TextStreamDelta,
+)
+from tinyflow.providers.base.llm import BaseLLM
+
+logger = logging.getLogger("tinyflow.providers.anthropic")
+
+
+class AnthropicProvider(BaseLLM):
+    def __init__(
+        self,
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        model: Optional[str] = None,
+        timeout: float = 60.0,
+    ):
+        # Configuration precedence: Manual > Environment Variable > Error
+        self.api_key = api_key or os.getenv("ANTHROPIC_API_KEY") or os.getenv("LLM_API_KEY")
+        self.model = (
+            model
+            or os.getenv("ANTHROPIC_MODEL")
+            or os.getenv("LLM_MODEL")
+            or "claude-3-opus-20240229"
+        )
+        self.timeout = timeout
+
+        if not self.api_key:
+            raise ValueError(
+                "Anthropic API key required. Set ANTHROPIC_API_KEY env var or pass api_key."
+            )
+
+        self.client = anthropic.AsyncAnthropic(
+            base_url=base_url,
+            api_key=self.api_key,
+            timeout=self.timeout,
+        )
+
+    @retry(
+        wait=wait_exponential(multiplier=1, min=2, max=10),
+        stop=stop_after_attempt(3),
+        reraise=True,
+    )
+    async def generate(
+        self, messages: List[Message], tools: Optional[List[Tool]] = None
+    ) -> LLMResponse:
+        logger.debug(f"Anthropic Generate call with {len(messages)} messages")
+        system_prompt = ""
+        claude_messages = []
+
+        for m in messages:
+            if m.role == "system":
+                system_prompt += (m.content or "") + "\n"
+            else:
+                claude_messages.append({"role": m.role, "content": m.content})
+
+        response = await self.client.messages.create(
+            model=self.model,
+            max_tokens=1024,
+            system=system_prompt,
+            messages=claude_messages,
+        )
+
+        content_text = ""
+        for block in response.content:
+            if isinstance(block, anthropic.types.TextBlock):
+                content_text += block.text
+
+        logger.debug(f"Anthropic Generate response: {content_text[:50]}...")
+
+        return LLMResponse(content=content_text, raw_response=response)
+
+    def stream(
+        self, messages: List[Message], tools: Optional[List[Tool]] = None
+    ) -> AsyncGenerator[StreamChunk, None]:
+        system_prompt = ""
+        claude_messages = []
+
+        for m in messages:
+            if m.role == "system":
+                system_prompt += (m.content or "") + "\n"
+            else:
+                claude_messages.append({"role": m.role, "content": m.content})
+
+        async def _stream():
+            async with self.client.messages.stream(
+                model=self.model,
+                max_tokens=1024,
+                system=system_prompt,
+                messages=claude_messages,
+            ) as stream:
+                async for text in stream.text_stream:
+                    if text:
+                        yield StreamChunk(content=text)
+
+        return _stream()
+
+    def stream_text(
+        self, messages: List[Message], tools: Optional[List[Tool]] = None
+    ) -> AsyncGenerator[StreamPart, None]:
+        system_prompt = ""
+        claude_messages = []
+
+        for m in messages:
+            if m.role == "system":
+                system_prompt += (m.content or "") + "\n"
+            else:
+                claude_messages.append({"role": m.role, "content": m.content})
+
+        async def _stream_text():
+            async with self.client.messages.stream(
+                model=self.model,
+                max_tokens=1024,
+                system=system_prompt,
+                messages=claude_messages,
+            ) as stream:
+                async for text in stream.text_stream:
+                    if text:
+                        yield TextStreamDelta(type="text", text=text)
+
+        return _stream_text()

tinyflow/providers/base/factory.py

@@ -0,0 +1,81 @@
+import logging
+from typing import Optional
+
+from tinyflow.config.helpers import filter_none_kwargs, get_config_value
+from tinyflow.core.exceptions import ConfigurationError
+from tinyflow.providers.anthropic_llm import AnthropicProvider
+from tinyflow.providers.base.llm import BaseLLM
+from tinyflow.providers.gemini_llm import GeminiProvider
+from tinyflow.providers.openai_llm import OpenAIProvider
+
+logger = logging.getLogger("tinyflow.providers.factory")
+
+
+class LLMFactory:
+    """Factory for creating LLM provider instances.
+
+    Supports both configuration-driven (env vars) and explicit instantiation.
+    Configuration Precedence:
+        1. Explicit parameters (highest)
+        2. Environment variables (LLM_*)
+        3. Settings (from .env or defaults)
+        4. Provider defaults (lowest)
+    """
+
+    @staticmethod
+    def create(
+        provider: Optional[str] = None,
+        **kwargs,
+    ) -> BaseLLM:
+        """Create an LLM provider instance.
+
+        Args:
+            provider: Provider name ('openai', 'anthropic', 'gemini').
+                     Falls back to LLM_PROVIDER env var -> settings -> default.
+            **kwargs: Additional provider-specific arguments (api_key, model, base_url, etc.)
+                     Common args (model, api_key, base_url) are resolved against
+                     LLM_* env vars and settings if not provided.
+
+        Returns:
+            Configured provider instance.
+
+        Raises:
+            ValueError: If provider is not specified and not found in configuration.
+        """
+        # Determine provider
+        provider = get_config_value(
+            provider,
+            env_key="LLM_PROVIDER",
+            settings_field="LLM_PROVIDER",
+            default_value="openai",
+        )
+
+        if not provider:
+            raise ConfigurationError("LLM Provider must be specified")
+
+        provider = provider.lower()
+        logger.info(f"Creating LLM provider: {provider}")
+
+        common_kwargs = {
+            "base_url": get_config_value(kwargs.get("base_url"), "LLM_BASE_URL", "LLM_BASE_URL"),
+            "api_key": get_config_value(kwargs.get("api_key"), "LLM_API_KEY", "LLM_API_KEY"),
+            "model": get_config_value(kwargs.get("model"), "LLM_MODEL", "LLM_MODEL"),
+        }
+
+        provider_kwargs = {**kwargs, **common_kwargs}
+        provider_kwargs = filter_none_kwargs(**provider_kwargs)
+
+        if provider == "openai":
+            return OpenAIProvider(**provider_kwargs)
+        elif provider == "deepseek":
+            if "base_url" not in provider_kwargs:
+                provider_kwargs["base_url"] = "https://api.deepseek.com"
+            return OpenAIProvider(**provider_kwargs)
+        elif provider == "anthropic":
+            return AnthropicProvider(**provider_kwargs)
+        elif provider == "gemini":
+            return GeminiProvider(**provider_kwargs)
+        else:
+            raise ConfigurationError(
+                f"Unsupported provider: {provider}. Supported: openai, anthropic, gemini"
+            )

tinyflow/providers/base/llm.py

@@ -0,0 +1,37 @@
+from abc import ABC, abstractmethod
+from typing import AsyncGenerator, List, Optional
+
+from tinyflow.core.tools import Tool
+from tinyflow.core.types import LLMResponse, Message, StreamChunk, StreamPart
+
+
+class BaseLLM(ABC):
+    @abstractmethod
+    async def generate(
+        self, messages: List[Message], tools: Optional[List[Tool]] = None
+    ) -> LLMResponse:
+        """Generate a complete response from the LLM."""
+        pass
+
+    @abstractmethod
+    def stream(
+        self, messages: List[Message], tools: Optional[List[Tool]] = None
+    ) -> AsyncGenerator[StreamChunk, None]:
+        """Stream response chunks (legacy method - use stream_text for rich types)."""
+        pass
+
+    @abstractmethod
+    def stream_text(
+        self, messages: List[Message], tools: Optional[List[Tool]] = None
+    ) -> AsyncGenerator[StreamPart, None]:
+        """Stream response with rich part types (text, reasoning, tool calls).
+
+        Yields StreamPart objects that can be:
+        - TextStreamDelta: Text content delta
+        - ReasoningStreamDelta: Reasoning/thinking trace
+        - ToolCallStreamStart: Tool call streaming begins
+        - ToolCallStreamDelta: Tool call argument delta
+        - ToolCallComplete: Complete tool call
+        - ToolResultStream: Tool execution result
+        """
+        pass