tinyflow-llm 0.1.0__py3-none-any.whl

tinyflow/core/exceptions.py

@@ -0,0 +1,63 @@
+"""TinyFlow Exceptions
+
+This module defines the public exception hierarchy for the TinyFlow framework.
+Users can catch these exceptions to handle errors in a structured way.
+"""
+
+from typing import Any, Optional
+
+
+class TinyFlowError(Exception):
+    """Base class for all TinyFlow exceptions."""
+
+    def __init__(self, message: str):
+        super().__init__(message)
+        self.message = message
+
+
+class ConfigurationError(TinyFlowError):
+    """Raised when configuration is missing or invalid."""
+
+    pass
+
+
+class ProviderError(TinyFlowError):
+    """Raised when an LLM provider fails (API errors, rate limits, etc.)."""
+
+    def __init__(
+        self,
+        message: str,
+        provider: str,
+        status_code: Optional[int] = None,
+        raw_response: Optional[Any] = None,
+    ):
+        super().__init__(message)
+        self.provider = provider
+        self.status_code = status_code
+        self.raw_response = raw_response
+
+
+class ToolError(TinyFlowError):
+    """Raised when a tool execution fails."""
+
+    def __init__(self, message: str, tool_name: str):
+        super().__init__(message)
+        self.tool_name = tool_name
+
+
+class WorkflowError(TinyFlowError):
+    """Raised when workflow execution fails (dependencies, routing, etc.)."""
+
+    pass
+
+
+class VectorError(TinyFlowError):
+    """Raised when vector database operations fail."""
+
+    pass
+
+
+class MemoryError(TinyFlowError):
+    """Raised when memory operations fail."""
+
+    pass

tinyflow/core/logger.py

@@ -0,0 +1,13 @@
+import logging
+
+
+def setup_logger():
+    """Setup the main logger for the library with a NullHandler to prevent 'No handler found' warnings."""
+    logger = logging.getLogger("tinyflow")
+    # Only add NullHandler if no handlers are present to avoid duplication if called multiple times
+    if not logger.handlers:
+        logger.addHandler(logging.NullHandler())
+    return logger
+
+
+logger = setup_logger()

tinyflow/core/message.py

@@ -0,0 +1,6 @@
+from tinyflow.core.types import Message
+
+
+def convert_to_model_messages(dict_messages: list[dict]) -> list[Message]:
+    """Convert a list of dictionaries to a list of Message objects."""
+    return [Message(**msg) for msg in dict_messages]

tinyflow/core/protocol.py

@@ -0,0 +1,81 @@
+import json
+from typing import AsyncGenerator, Set
+
+from tinyflow.core.types import (
+    DataUIPart,
+    ReasoningUIPart,
+    TextUIPart,
+    ToolUIPart,
+    UIMessage,
+)
+
+
+class JSONStreamAdapter:
+    """
+    Adapts a UIMessage stream to a structured JSON delta stream.
+    Each yield is a valid JSON string representing a specific part update.
+    """
+
+    def __init__(self):
+        self._last_lengths = {}  # Track lengths per part type/id
+        self._emitted_states: Set[str] = set()
+
+    async def transform(self, message_stream: AsyncGenerator[UIMessage, None]) -> AsyncGenerator[str, None]:
+        async for message in message_stream:
+            if message.role != "assistant":
+                continue
+
+            for part in message.parts:
+                # Use a combination of type and optional ID for tracking
+                part_id = getattr(part, "tool_call_id", None) or getattr(part, "id", "default")
+                track_key = f"{part.type}:{part_id}"
+                
+                # 1. Handle Text & Reasoning (Incremental Deltas)
+                if isinstance(part, (TextUIPart, ReasoningUIPart)):
+                    current_text = part.text
+                    last_len = self._last_lengths.get(track_key, 0)
+                    
+                    if len(current_text) > last_len:
+                        delta = current_text[last_len:]
+                        yield json.dumps({
+                            "type": f"{part.type}-delta",
+                            "delta": delta,
+                            "index": message.parts.index(part) # Help frontend locate the part
+                        }) + "\n"
+                        self._last_lengths[track_key] = len(current_text)
+
+                # 2. Handle Tool Calls (State-based Events)
+                elif isinstance(part, ToolUIPart):
+                    state_key = f"tool:{part.tool_call_id}:{part.state}"
+                    if state_key not in self._emitted_states:
+                        payload = {
+                            "type": "tool-call" if "input" in part.state.value else "tool-result",
+                            "status": part.state.value,
+                            "toolCallId": part.tool_call_id,
+                            "toolName": part.tool_name,
+                        }
+                        if part.input is not None:
+                            payload["input"] = part.input
+                        if part.output is not None:
+                            payload["output"] = part.output
+                        if part.error_text:
+                            payload["error"] = part.error_text
+                            
+                        yield json.dumps(payload) + "\n"
+                        self._emitted_states.add(state_key)
+
+                # 3. Handle Custom Data
+                elif isinstance(part, DataUIPart):
+                    if track_key not in self._emitted_states:
+                        yield json.dumps({
+                            "type": "data",
+                            "id": part.id,
+                            "data": part.data
+                        }) + "\n"
+                        self._emitted_states.add(track_key)
+
+async def to_json_stream(message_stream: AsyncGenerator[UIMessage, None]) -> AsyncGenerator[str, None]:
+    """Helper function for the most common DX: agent.stream() -> to_json_stream()"""
+    adapter = JSONStreamAdapter()
+    async for chunk in adapter.transform(message_stream):
+        yield chunk

tinyflow/core/tools.py

@@ -0,0 +1,200 @@
+import inspect
+import logging
+from typing import Any, Callable, Dict, Optional, Type
+
+from docstring_parser import parse
+from pydantic import BaseModel, Field, create_model
+
+from tinyflow.core.exceptions import ToolError
+
+logger = logging.getLogger("tinyflow.core.tools")
+
+
+class Tool:
+    """
+    Explicit declarative tool definition (Vercel AI SDK style).
+
+    This is the unified tool class of the framework, providing clear control:
+    - name: Tool name
+    - description: Detailed description (LLM will decide when to call based on this)
+    - parameters: Pydantic model defining parameter structure
+    - execute: Execution function, receiving the validated Pydantic model instance
+
+    Example:
+        ```python
+        class WeatherInput(BaseModel):
+            city: str = Field(description="City name")
+            unit: str = Field(default="celsius", description="Temperature unit")
+
+        weather_tool = Tool(
+            name="get_weather",
+            description="Get weather information for a specific city.",
+            parameters=WeatherInput,
+            execute=lambda args: f"{args.city} is sunny"
+        )
+        ```
+    """
+
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        parameters: Type[BaseModel],
+        execute: Callable[[Any], Any],
+    ):
+        self.name = name
+        self.description = description
+        self.parameters = parameters
+        self.execute_func = execute
+        self.args_schema = parameters
+        self.func = execute
+        logger.debug(f"Tool initialized: {self.name}")
+
+    async def execute(self, **kwargs) -> Any:
+        """Execute the tool with automatic parameter validation."""
+        try:
+            validated_args = self.parameters.model_validate(kwargs)
+
+            if inspect.iscoroutinefunction(self.execute_func):
+                result = await self.execute_func(validated_args)
+            else:
+                result = self.execute_func(validated_args)
+
+            # If result is an Agent or WorkflowExecutor, run it
+            # This enables "Agent as a Tool" pattern
+            if hasattr(result, "run"):
+                if inspect.iscoroutinefunction(result.run):
+                    # For Agents/Workflows, we might want to return the result directly
+                    # or handle the run execution here.
+                    # Simpler approach: execute .run() and return output
+                    # For WorkflowExecutor, run() returns dict of results
+                    # For Agent, run() returns str
+                    output = await result.run(**kwargs)  # Pass args if compatible?
+                    # Actually, usually the tool function prepares the agent/workflow
+                    # and might run it with specific args derived from validated_args.
+                    # If the tool function *returns* the runner instance, we run it.
+                    # But usually the tool body does the running.
+                    # Let's assume the tool body returns the FINAL result (str/dict).
+                    # So no change needed here if tool body handles execution.
+
+                    # BUT, if we want to support returning an AsyncGenerator (Streaming Workflow),
+                    # we must NOT convert to str immediately.
+                    pass
+
+            return result
+        except Exception as e:
+            if hasattr(e, "errors"):
+                logger.warning(f"Validation Error for tool '{self.name}': {str(e)}")
+                # Raise ToolError for validation failure?
+                # Or keep string return for LLM loop?
+                # The execute method is typically called by Agent which wants a string.
+                # However, for library users executing tools directly, an exception is better.
+                # Let's wrap it in ToolError but format it friendly.
+                raise ToolError(f"Validation Error: {str(e)}", tool_name=self.name) from e
+
+            logger.error(f"Error executing tool '{self.name}': {str(e)}")
+            raise ToolError(f"Execution Error: {str(e)}", tool_name=self.name) from e
+
+    def to_openai_schema(self, strict: bool = False) -> Dict[str, Any]:
+        """
+        Convert to OpenAI tool format.
+
+        Extract field descriptions from Pydantic model Field descriptions to ensure
+        LLM accurately understands the purpose of each parameter.
+        """
+        schema = self.parameters.model_json_schema()
+
+        parameters = {
+            "type": "object",
+            "properties": schema.get("properties", {}),
+            "required": schema.get("required", []),
+        }
+
+        if strict:
+            parameters["additionalProperties"] = False
+            parameters["required"] = list(schema.get("properties", {}).keys())
+
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": parameters,
+                "strict": strict,
+            },
+        }
+
+
+def tool(
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+):
+    """
+    Tool decorator supporting automatic tool definition derivation from functions.
+
+    Supports two modes:
+    1. Decorating normal functions: Automatically parse description and parameter descriptions from docstring.
+    2. Decorating class methods: Use with Tool class.
+
+    Example:
+        ```python
+        @tool()
+        async def get_weather(city: str):
+            '''Get weather.
+
+            Args:
+                city: City name
+            '''
+            return "Sunny"
+
+        @tool(name="custom_name", description="Custom description")
+        def my_tool(input: str):
+            return input
+        ```
+    """
+
+    def decorator(func: Callable) -> Tool:
+        tool_name = name or func.__name__
+
+        tool_description = description
+        if tool_description is None:
+            doc = parse(func.__doc__ or "")
+            tool_description = doc.short_description or f"Tool: {tool_name}"
+
+        # Create Pydantic model from function signature
+        args_schema = _get_args_schema(func)
+
+        # Create execute function that unwraps Pydantic model to kwargs
+        def execute_wrapper(args):
+            kwargs = args.model_dump()
+            if inspect.iscoroutinefunction(func):
+                # Return the coroutine - it will be awaited by Tool.execute
+                return func(**kwargs)
+            else:
+                return func(**kwargs)
+
+        return Tool(
+            name=tool_name,
+            description=tool_description,
+            parameters=args_schema,
+            execute=execute_wrapper,
+        )
+
+    return decorator
+
+
+def _get_args_schema(func: Callable) -> Type[BaseModel]:
+    """Helper function to create Pydantic model from function signature."""
+    sig = inspect.signature(func)
+    doc = parse(func.__doc__ or "")
+
+    param_descriptions = {p.arg_name: p.description for p in doc.params}
+
+    fields = {}
+    for param_name, param in sig.parameters.items():
+        annotation = param.annotation if param.annotation != inspect.Parameter.empty else str
+        default = param.default if param.default != inspect.Parameter.empty else ...
+        description = param_descriptions.get(param_name, "")
+        fields[param_name] = (annotation, Field(default=default, description=description))
+
+    return create_model(f"{func.__name__}Schema", **fields)

tinyflow/core/types.py

@@ -0,0 +1,252 @@
+from enum import Enum
+from typing import Any, Dict, Generic, List, Literal, Optional, TypeVar, Union
+
+from pydantic import BaseModel, Field
+
+# =============================================================================
+# Legacy Types (Backward Compatibility)
+# =============================================================================
+
+
+class Message(BaseModel):
+    """Base message type for conversation history."""
+
+    role: Literal["system", "user", "assistant", "tool"]
+    content: Optional[str] = ""
+    name: Optional[str] = None
+    tool_calls: Optional[List[Dict[str, Any]]] = None
+    tool_call_id: Optional[str] = None
+    reasoning_content: Optional[str] = None
+
+
+class StreamChunk(BaseModel):
+    """Legacy stream chunk - deprecated in favor of StreamPart types."""
+
+    content: str
+    finish_reason: Optional[str] = None
+
+
+class LLMResponse(BaseModel):
+    """Response from LLM provider."""
+
+    content: Optional[str]
+    tool_calls: Optional[List[Dict[str, Any]]] = None
+    raw_response: Optional[object] = None
+    usage: Optional[dict] = None
+
+
+# =============================================================================
+# UI Message Part State Enums
+# =============================================================================
+
+
+class TextPartState(str, Enum):
+    """States for text content parts."""
+
+    STREAMING = "streaming"
+    DONE = "done"
+
+
+class ToolPartState(str, Enum):
+    """States for tool execution lifecycle.
+
+    States:
+        INPUT_STREAMING: Tool arguments being generated by LLM
+        INPUT_AVAILABLE: Arguments complete, ready for execution
+        OUTPUT_AVAILABLE: Tool executed, result available
+        OUTPUT_ERROR: Tool execution failed
+        APPROVAL_REQUESTED: Waiting for user approval
+    """
+
+    INPUT_STREAMING = "input-streaming"
+    INPUT_AVAILABLE = "input-available"
+    OUTPUT_AVAILABLE = "output-available"
+    OUTPUT_ERROR = "output-error"
+    APPROVAL_REQUESTED = "approval-requested"
+
+
+# =============================================================================
+# UI Message Part Types
+# =============================================================================
+
+
+class TextUIPart(BaseModel):
+    """Text content with streaming state."""
+
+    type: Literal["text"] = "text"
+    text: str
+    state: Optional[TextPartState] = None
+
+
+class ReasoningUIPart(BaseModel):
+    """Agent thinking/reasoning trace."""
+
+    type: Literal["reasoning"] = "reasoning"
+    text: str
+    state: Optional[TextPartState] = None
+    provider_metadata: Optional[Dict[str, Any]] = None
+
+
+class ToolUIPart(BaseModel):
+    """Tool invocation and result representation.
+
+    This part represents the complete lifecycle of a tool call:
+    1. INPUT_STREAMING: Tool arguments being generated by LLM
+    2. INPUT_AVAILABLE: Arguments complete, ready for execution
+    3. OUTPUT_AVAILABLE: Tool executed, result available
+    4. OUTPUT_ERROR: Tool execution failed
+    5. APPROVAL_REQUESTED: Waiting for user to approve execution
+    """
+
+    type: str = Field(..., description="Format: tool-{tool_name}")
+    tool_call_id: str
+    tool_name: str
+    state: ToolPartState
+
+    # Input data (when state is INPUT_* or OUTPUT_*)
+    input: Optional[Dict[str, Any]] = None
+
+    # Output data (when state is OUTPUT_AVAILABLE)
+    output: Optional[Any] = None
+
+    # Error text (when state is OUTPUT_ERROR)
+    error_text: Optional[str] = None
+
+    # Approval data (when state is APPROVAL_REQUESTED)
+    approval_id: Optional[str] = None
+
+    # Execution flag
+    provider_executed: Optional[bool] = None
+
+
+class FileUIPart(BaseModel):
+    """File or image attachment."""
+
+    type: Literal["file"] = "file"
+    media_type: str = Field(..., description="IANA media type (e.g., image/png)")
+    filename: Optional[str] = None
+    url: str = Field(..., description="Can be URL or data URL (base64)")
+
+
+class DataUIPart(BaseModel):
+    """Custom data part for structured data.
+
+    This allows passing arbitrary structured data to the frontend
+    for custom UI components.
+    """
+
+    type: str = Field(..., description="Format: data-{data_name}")
+    id: Optional[str] = None
+    data: Dict[str, Any]
+
+
+class StepStartUIPart(BaseModel):
+    """Marker for step boundary in multi-step tool execution."""
+
+    type: Literal["step-start"] = "step-start"
+
+
+# Union type for all UI message parts
+UIMessagePart = Union[
+    TextUIPart,
+    ReasoningUIPart,
+    ToolUIPart,
+    FileUIPart,
+    DataUIPart,
+    StepStartUIPart,
+]
+
+
+# =============================================================================
+# UI Message Type
+# =============================================================================
+
+METADATA = TypeVar("METADATA", bound=BaseModel)
+DATA_PARTS = TypeVar("DATA_PARTS", bound=BaseModel)
+TOOLS = TypeVar("TOOLS", bound=BaseModel)
+
+
+class UIMessage(BaseModel, Generic[METADATA, DATA_PARTS, TOOLS]):
+    """UI Message - Complete message representation for frontend rendering.
+
+    This is the single source of truth for application state, containing
+    full message history including metadata, data parts, and contextual
+    information for UI rendering.
+
+    Generic Parameters:
+        METADATA: Custom metadata type for additional message information
+        DATA_PARTS: Custom data part types for structured components
+        TOOLS: Tool definitions for type-safe tool interactions
+    """
+
+    id: str
+    role: Literal["system", "user", "assistant"]
+    metadata: Optional[METADATA] = None
+    parts: List[UIMessagePart]
+
+
+# =============================================================================
+# Stream Part Types (for Provider-Level Streaming)
+# =============================================================================
+
+
+class TextStreamDelta(BaseModel):
+    """Text delta in streaming response."""
+
+    type: Literal["text"] = "text"
+    text: str
+
+
+class ReasoningStreamDelta(BaseModel):
+    """Reasoning delta in streaming response."""
+
+    type: Literal["reasoning"] = "reasoning"
+    text: str
+
+
+class ToolCallStreamStart(BaseModel):
+    """Tool call streaming starts."""
+
+    type: Literal["tool-call-streaming-start"] = "tool-call-streaming-start"
+    tool_call_id: str
+    tool_name: str
+
+
+class ToolCallStreamDelta(BaseModel):
+    """Tool call argument delta (streaming)."""
+
+    type: Literal["tool-call-delta"] = "tool-call-delta"
+    tool_call_id: str
+    tool_name: str
+    args_text_delta: str
+
+
+class ToolCallComplete(BaseModel):
+    """Complete tool call (when streaming is off or complete)."""
+
+    type: Literal["tool-call"] = "tool-call"
+    tool_call_id: str
+    tool_name: str
+    input: Dict[str, Any]
+
+
+class ToolResultStream(BaseModel):
+    """Tool execution result."""
+
+    type: Literal["tool-result"] = "tool-result"
+    tool_call_id: str
+    tool_name: str
+    input: Dict[str, Any]
+    output: Any
+    is_error: bool = False
+
+
+# Union type for all stream parts
+StreamPart = Union[
+    TextStreamDelta,
+    ReasoningStreamDelta,
+    ToolCallStreamStart,
+    ToolCallStreamDelta,
+    ToolCallComplete,
+    ToolResultStream,
+]

tinyflow/embeddings/__init__.py

@@ -0,0 +1,4 @@
+"""Embedding package
+
+Contains different embedding model implementations
+"""

tinyflow/embeddings/base.py

@@ -0,0 +1,32 @@
+"""Embedding Abstraction Layer
+
+Supports multiple embedding model providers
+"""
+
+from abc import ABC, abstractmethod
+from typing import List
+
+
+class BaseEmbedding(ABC):
+    """Embedding abstract base class"""
+
+    @abstractmethod
+    async def embed(self, texts: List[str]) -> List[List[float]]:
+        """Convert text list to vectors
+
+        Args:
+            texts: List of text strings
+
+        Returns:
+            List of vectors (one for each text)
+        """
+        pass
+
+    @abstractmethod
+    def get_dimension(self) -> int:
+        """Get vector dimension
+
+        Returns:
+            Vector dimension
+        """
+        pass