tinyflow-llm 0.1.0__py3-none-any.whl

tinyflow/providers/gemini_llm.py

@@ -0,0 +1,130 @@
+import os
+from typing import AsyncGenerator, List, Optional
+
+from google import genai
+from google.genai import types
+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
+
+
+class GeminiProvider(BaseLLM):
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model: Optional[str] = None,
+        base_url: Optional[str] = None,
+    ):
+        # Configuration precedence: Manual > Environment Variable > Error
+        self.api_key = api_key or os.getenv("GEMINI_API_KEY") or os.getenv("LLM_API_KEY")
+        self.model = (
+            model or os.getenv("GEMINI_MODEL") or os.getenv("LLM_MODEL") or "gemini-1.5-flash"
+        )
+        self.base_url = base_url  # Gemini supports custom endpoints via client options
+
+        if not self.api_key:
+            raise ValueError(
+                "Gemini API key required. Set GEMINI_API_KEY env var or pass api_key."
+            )
+
+        self.client = genai.Client(
+            api_key=self.api_key,
+        )
+
+    def _convert_messages(self, messages: List[Message]) -> List[types.Content]:
+        """
+        Gemini requires converting OpenAI-style messages to Google Content type
+        Note: This is a simplified conversion; complex Tool Results require more refined handling
+        """
+        gemini_contents = []
+        for m in messages:
+            role = "user" if m.role in ["user", "tool"] else "model"
+            if m.role == "system":
+                # Gemini V1 SDK suggests passing system prompt via config's system_instruction
+                # Temporarily doing simple merging here, or extracting separately during generation
+                continue
+
+            parts = [types.Part(text=m.content)] if m.content else []
+
+            # Simple conversion logic; production needs to handle function_call conversion
+            gemini_contents.append(types.Content(role=role, parts=parts))
+
+        return gemini_contents
+
+    def _get_system_prompt(self, messages: List[Message]) -> Optional[str]:
+        for m in messages:
+            if m.role == "system":
+                return m.content
+        return None
+
+    @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:
+        contents = self._convert_messages(messages)
+        system_instr = self._get_system_prompt(messages)
+
+        # Gemini supports passing functions directly, SDK automatically parses Schema
+        # But we already have Tool; for simplicity, momentarily disabling Tools for Gemini
+        if tools:
+            pass
+
+        response = await self.client.aio.models.generate_content(
+            model=self.model,
+            contents=contents,
+            config=types.GenerateContentConfig(system_instruction=system_instr),
+        )
+
+        return LLMResponse(content=response.text, raw_response=response)
+
+    def stream(
+        self, messages: List[Message], tools: Optional[List[Tool]] = None
+    ) -> AsyncGenerator[StreamChunk, None]:
+        contents = self._convert_messages(messages)
+        system_instr = self._get_system_prompt(messages)
+
+        async def _stream():
+            # Call streaming interface
+            response_stream = await self.client.aio.models.generate_content_stream(
+                model=self.model,
+                contents=contents,
+                config=types.GenerateContentConfig(system_instruction=system_instr),
+            )
+
+            async for chunk in response_stream:
+                if chunk.text:
+                    yield StreamChunk(content=chunk.text)
+
+        return _stream()
+
+    def stream_text(
+        self, messages: List[Message], tools: Optional[List[Tool]] = None
+    ) -> AsyncGenerator[StreamPart, None]:
+        """Stream response with rich part types."""
+        contents = self._convert_messages(messages)
+        system_instr = self._get_system_prompt(messages)
+
+        async def _stream_text():
+            response_stream = await self.client.aio.models.generate_content_stream(
+                model=self.model,
+                contents=contents,
+                config=types.GenerateContentConfig(system_instruction=system_instr),
+            )
+
+            async for chunk in response_stream:
+                if chunk.text:
+                    yield TextStreamDelta(type="text", text=chunk.text)
+
+        return _stream_text()

tinyflow/providers/openai_llm.py

@@ -0,0 +1,198 @@
+import logging
+import os
+from typing import AsyncGenerator, List, Optional
+
+from openai import AsyncOpenAI
+from tenacity import retry, stop_after_attempt, wait_exponential
+
+from tinyflow.core.tools import Tool
+from tinyflow.core.types import (
+    LLMResponse,
+    Message,
+    ReasoningStreamDelta,
+    StreamChunk,
+    StreamPart,
+    TextStreamDelta,
+    ToolCallComplete,
+    ToolCallStreamDelta,
+    ToolCallStreamStart,
+)
+from tinyflow.providers.base.llm import BaseLLM
+
+logger = logging.getLogger("tinyflow.providers.openai")
+
+
+class OpenAIProvider(BaseLLM):
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        model: Optional[str] = None,
+        base_url: Optional[str] = None,
+        timeout: float = 60.0,
+    ):
+        # Configuration precedence: Manual > Environment Variable > Error
+        self.api_key = api_key or os.getenv("OPENAI_API_KEY") or os.getenv("LLM_API_KEY")
+        self.model = model or os.getenv("OPENAI_MODEL") or os.getenv("LLM_MODEL") or "gpt-4o"
+        self.base_url = base_url or os.getenv("OPENAI_BASE_URL") or os.getenv("LLM_BASE_URL")
+        self.timeout = timeout
+
+        logging.info(
+            f"OpenAIProvider Initialized: model='{self.model}', base_url='{self.base_url}'"
+        )
+
+        if not self.api_key:
+            raise ValueError(
+                "OpenAI API key required. Set OPENAI_API_KEY env var or pass api_key."
+            )
+
+        self.client = AsyncOpenAI(
+            api_key=self.api_key,
+            base_url=self.base_url,
+            timeout=self.timeout,
+        )
+
+    def _convert_messages(self, messages: List[Message]):
+        openai_msgs = []
+        for m in messages:
+            msg = {"role": m.role, "content": m.content}
+            if m.tool_calls:
+                msg["tool_calls"] = m.tool_calls
+            if m.tool_call_id:
+                msg["tool_call_id"] = m.tool_call_id
+            if m.name:
+                msg["name"] = m.name
+
+            if m.reasoning_content:
+                msg["reasoning_content"] = m.reasoning_content
+
+            openai_msgs.append(msg)
+        return openai_msgs
+
+    @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"OpenAI Generate call with {len(messages)} messages")
+        openai_messages = self._convert_messages(messages)
+        openai_tools = [t.to_openai_schema(strict=True) for t in tools] if tools else None
+
+        response = await self.client.chat.completions.create(
+            model=self.model,
+            messages=openai_messages,
+            tools=openai_tools,  # pyright: ignore[reportArgumentType]
+        )
+        msg = response.choices[0].message
+        logger.debug(
+            f"OpenAI Generate response: {msg.content[:50] if msg.content else 'None'}..."
+        )
+
+        # Process tool_calls format
+        tool_calls_dict = None
+        if msg.tool_calls:
+            tool_calls_dict = [
+                {
+                    "id": tc.id,
+                    "function": {
+                        "name": tc.function.name,
+                        "arguments": tc.function.arguments,
+                    },
+                    "type": "function",
+                }
+                for tc in msg.tool_calls
+            ]
+
+        return LLMResponse(content=msg.content, tool_calls=tool_calls_dict, raw_response=response)
+
+    def stream(
+        self, messages: List[Message], tools: Optional[List[Tool]] = None
+    ) -> AsyncGenerator[StreamChunk, None]:
+        openai_messages = self._convert_messages(messages)
+
+        async def _stream():
+            stream = await self.client.chat.completions.create(
+                model=self.model, messages=openai_messages, stream=True
+            )
+
+            async for chunk in stream:
+                if not chunk.choices:
+                    continue
+                delta = chunk.choices[0].delta
+                if delta.content:
+                    yield StreamChunk(content=delta.content)
+
+        return _stream()
+
+    def stream_text(
+        self, messages: List[Message], tools: Optional[List[Tool]] = None
+    ) -> AsyncGenerator[StreamPart, None]:
+        openai_messages = self._convert_messages(messages)
+        openai_tools = [t.to_openai_schema(strict=True) for t in tools] if tools else None
+
+        async def _stream_text():
+            stream = await self.client.chat.completions.create(
+                model=self.model,
+                messages=openai_messages,
+                tools=openai_tools,  # type: ignore[reportCallIssue]
+                stream=True,
+            )
+
+            current_tool_call: Optional[dict] = None
+
+            async for chunk in stream:
+                if not chunk.choices:
+                    continue
+
+                delta = chunk.choices[0].delta
+
+                if delta.content:
+                    yield TextStreamDelta(type="text", text=delta.content)
+
+                # Support DeepSeek/o1 "reasoning_content"
+                if hasattr(delta, "reasoning_content") and delta.reasoning_content:
+                    yield ReasoningStreamDelta(type="reasoning", text=delta.reasoning_content)
+
+                if delta.tool_calls:
+                    for tc_delta in delta.tool_calls:
+                        if tc_delta.id:
+                            current_tool_call = {
+                                "id": tc_delta.id,
+                                "name": tc_delta.function.name,
+                                "args": tc_delta.function.arguments or "",
+                            }
+                            yield ToolCallStreamStart(
+                                type="tool-call-streaming-start",
+                                tool_call_id=current_tool_call["id"],
+                                tool_name=current_tool_call["name"],
+                            )
+                        else:
+                            if current_tool_call:
+                                current_tool_call["args"] += tc_delta.function.arguments or ""
+                                yield ToolCallStreamDelta(
+                                    type="tool-call-delta",
+                                    tool_call_id=current_tool_call["id"],
+                                    tool_name=current_tool_call["name"],
+                                    args_text_delta=tc_delta.function.arguments or "",
+                                )
+
+                if chunk.choices[0].finish_reason:
+                    if current_tool_call:
+                        import json
+
+                        try:
+                            args = json.loads(current_tool_call["args"])
+                        except json.JSONDecodeError:
+                            args = current_tool_call["args"]
+
+                        yield ToolCallComplete(
+                            type="tool-call",
+                            tool_call_id=current_tool_call["id"],
+                            tool_name=current_tool_call["name"],
+                            input=args,
+                        )
+                        current_tool_call = None
+
+        return _stream_text()

tinyflow/tools/builtin/__init__.py

@@ -0,0 +1,36 @@
+"""Built-in tools for TinyFlow Agent framework.
+
+This package provides ready-to-use tools for common Agent tasks:
+- Search: Tavily AI and DuckDuckGo
+- Web Reading: Jina AI Reader
+- Code Execution: Python Interpreter (restricted)
+"""
+
+from tinyflow.tools.builtin.code_execution import (
+    create_python_interpreter_tool,
+    python_interpreter_tool,
+)
+from tinyflow.tools.builtin.search import (
+    create_duckduckgo_search_tool,
+    create_tavily_search_tool,
+    duckduckgo_search_tool,
+    tavily_search_tool,
+)
+from tinyflow.tools.builtin.web_reader import (
+    create_jina_reader_tool,
+    jina_reader_tool,
+)
+
+__all__ = [
+    # Search tools
+    "create_tavily_search_tool",
+    "create_duckduckgo_search_tool",
+    "tavily_search_tool",
+    "duckduckgo_search_tool",
+    # Web reader tools
+    "create_jina_reader_tool",
+    "jina_reader_tool",
+    # Code execution
+    "create_python_interpreter_tool",
+    "python_interpreter_tool",
+]

tinyflow/tools/builtin/code_execution.py

@@ -0,0 +1,143 @@
+"""Python interpreter tool for code execution.
+
+⚠️ SECURITY WARNING: This tool executes arbitrary Python code.
+Use only in sandboxed environments. Never expose to untrusted users.
+"""
+
+import io
+import sys
+import traceback
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, Field
+
+from tinyflow.core.tools import Tool
+
+
+class PythonInterpreterInput(BaseModel):
+    """Input parameters for Python interpreter."""
+
+    code: str = Field(
+        description="The Python code to execute. "
+        "Use print() to output results. "
+        "Available libraries: os, sys, math, json, re, datetime, collections, itertools"
+    )
+    timeout: int = Field(default=30, description="Maximum execution time in seconds (1-60)")
+
+
+def create_python_interpreter_tool(
+    allowed_modules: Optional[list] = None, allow_file_access: bool = False
+) -> Tool:
+    """Create a Python interpreter tool instance.
+
+    ⚠️ SECURITY WARNING:
+    This tool executes arbitrary Python code. By default, it runs in a
+    restricted environment with limited builtins. However, it should only
+    be used in sandboxed environments.
+
+    Args:
+        allowed_modules: List of module names to allow importing.
+                        Defaults to safe standard library modules.
+        allow_file_access: Whether to allow file system access. Default False.
+
+    Returns:
+        Configured Tool instance for Python code execution.
+    """
+    # Default safe modules
+    if allowed_modules is None:
+        allowed_modules = [
+            "math",
+            "json",
+            "re",
+            "datetime",
+            "collections",
+            "itertools",
+            "random",
+            "statistics",
+            "typing",
+        ]
+
+    def execute_python_code(args: PythonInterpreterInput) -> str:
+        """Execute Python code in a restricted environment."""
+        max(1, min(args.timeout, 60))
+
+        # Create restricted globals
+        restricted_globals: Dict[str, Any] = {
+            "__builtins__": {
+                "print": print,
+                "len": len,
+                "range": range,
+                "enumerate": enumerate,
+                "zip": zip,
+                "map": map,
+                "filter": filter,
+                "sum": sum,
+                "min": min,
+                "max": max,
+                "abs": abs,
+                "round": round,
+                "pow": pow,
+                "divmod": divmod,
+                "bool": bool,
+                "int": int,
+                "float": float,
+                "str": str,
+                "list": list,
+                "dict": dict,
+                "tuple": tuple,
+                "set": set,
+                "frozenset": frozenset,
+                "type": type,
+                "isinstance": isinstance,
+                "hasattr": hasattr,
+                "getattr": getattr,
+                "Exception": Exception,
+            }
+        }
+
+        # Import allowed modules
+        for module_name in allowed_modules:
+            try:
+                module = __import__(module_name)
+                restricted_globals[module_name] = module
+            except ImportError:
+                pass
+
+        # Capture stdout
+        old_stdout = sys.stdout
+        sys.stdout = io.StringIO()
+
+        try:
+            # Execute code with timeout-like protection (via limited operations)
+            # Note: For true sandboxing, consider using subprocess or RestrictedPython
+            exec(args.code, restricted_globals)
+
+            # Get output
+            output = sys.stdout.getvalue()
+
+            if not output.strip():
+                return "Code executed successfully (no output)."
+
+            return f"Output:\n{output}"
+
+        except Exception as e:
+            error_msg = f"Error executing code: {type(e).__name__}: {str(e)}"
+            # Include traceback for debugging
+            tb = traceback.format_exc()
+            return f"{error_msg}\n\nTraceback:\n{tb}"
+        finally:
+            sys.stdout = old_stdout
+
+    return Tool(
+        name="python_interpreter",
+        description="Execute Python code in a restricted environment. "
+        "Available: math, json, re, datetime, collections, itertools. "
+        "Use print() to see output. "
+        "⚠️ Security: Code runs in restricted mode but use with caution.",
+        parameters=PythonInterpreterInput,
+        execute=execute_python_code,
+    )
+
+
+# Convenience instance with default restrictions
+python_interpreter_tool = create_python_interpreter_tool()

tinyflow/tools/builtin/search.py

@@ -0,0 +1,145 @@
+"""Built-in search tools using mainstream APIs.
+
+This module provides search capabilities through:
+- Tavily AI (high-quality AI search, requires API key)
+- DuckDuckGo (free, no API key required)
+"""
+
+import os
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from tinyflow.core.tools import Tool
+
+
+class TavilySearchInput(BaseModel):
+    """Input parameters for Tavily search."""
+
+    query: str = Field(description="The search query to execute")
+    max_results: int = Field(default=5, description="Maximum number of results to return (1-10)")
+    include_answer: bool = Field(
+        default=True, description="Include an AI-generated answer summary"
+    )
+
+
+class DuckDuckGoSearchInput(BaseModel):
+    """Input parameters for DuckDuckGo search."""
+
+    query: str = Field(description="The search query to execute")
+    max_results: int = Field(default=5, description="Maximum number of results to return")
+
+
+def _get_tavily_client(api_key: Optional[str] = None):
+    """Get or create Tavily client with the provided or environment API key."""
+    try:
+        from tavily import TavilyClient
+    except ImportError:
+        raise ImportError(
+            "tavily-python is required for TavilySearchTool. "
+            "Install with: pip install tavily-python"
+        )
+
+    # Configuration precedence: Manual > Environment Variable > Error
+    key = api_key or os.getenv("TAVILY_API_KEY")
+    if not key:
+        raise ValueError(
+            "Tavily API key required. Set TAVILY_API_KEY environment variable "
+            "or pass api_key parameter to the tool."
+        )
+
+    return TavilyClient(api_key=key)
+
+
+def create_tavily_search_tool(api_key: Optional[str] = None) -> Tool:
+    """Create a Tavily search tool instance.
+
+    Args:
+        api_key: Tavily API key. If not provided, uses TAVILY_API_KEY env var.
+
+    Returns:
+        Configured Tool instance for Tavily search.
+    """
+
+    def execute_tavily_search(args: TavilySearchInput) -> str:
+        """Execute Tavily search and format results."""
+        client = _get_tavily_client(api_key)
+
+        response = client.search(
+            query=args.query,
+            max_results=args.max_results,
+            include_answer=args.include_answer,
+        )
+
+        # Format results
+        results_text = f"Search results for: {args.query}\n\n"
+
+        if args.include_answer and "answer" in response:
+            results_text += f"AI Summary: {response['answer']}\n\n"
+
+        results = response.get("results", [])
+        for i, result in enumerate(results, 1):
+            title = result.get("title", "No title")
+            content = result.get("content", "No content")
+            url = result.get("url", "")
+            results_text += f"{i}. {title}\n{content}\nURL: {url}\n\n"
+
+        return results_text.strip()
+
+    return Tool(
+        name="tavily_search",
+        description="Search the web using Tavily AI for high-quality, relevant results. "
+        "Best for research, fact-checking, and finding current information.",
+        parameters=TavilySearchInput,
+        execute=execute_tavily_search,
+    )
+
+
+def create_duckduckgo_search_tool() -> Tool:
+    """Create a DuckDuckGo search tool instance.
+
+    Note: DuckDuckGo search does not require an API key.
+
+    Returns:
+        Configured Tool instance for DuckDuckGo search.
+    """
+
+    def execute_duckduckgo_search(args: DuckDuckGoSearchInput) -> str:
+        """Execute DuckDuckGo search and format results."""
+        try:
+            from duckduckgo_search import DDGS
+        except ImportError:
+            raise ImportError(
+                "duckduckgo-search is required for DuckDuckGoSearchTool. "
+                "Install with: pip install duckduckgo-search"
+            )
+
+        with DDGS() as ddgs:
+            results = list(ddgs.text(args.query, max_results=args.max_results))
+
+        # Format results
+        results_text = f"Search results for: {args.query}\n\n"
+
+        if not results:
+            return f"No results found for: {args.query}"
+
+        for i, result in enumerate(results, 1):
+            title = result.get("title", "No title")
+            body = result.get("body", "No description")
+            href = result.get("href", "")
+            results_text += f"{i}. {title}\n{body}\nURL: {href}\n\n"
+
+        return results_text.strip()
+
+    return Tool(
+        name="duckduckgo_search",
+        description="Search the web using DuckDuckGo. Free, no API key required. "
+        "Good for general web searches when Tavily is not available.",
+        parameters=DuckDuckGoSearchInput,
+        execute=execute_duckduckgo_search,
+    )
+
+
+# Convenience instances (will use env vars for API keys)
+tavily_search_tool = create_tavily_search_tool()
+duckduckgo_search_tool = create_duckduckgo_search_tool()