stratifyai 0.1.0__py3-none-any.whl

stratifyai/chat/stratifyai_anthropic.py

@@ -0,0 +1,194 @@
+"""Anthropic chat interface for StratifyAI.
+
+Provides convenient functions for Anthropic Claude chat completions.
+Model must be specified for each request.
+
+Environment Variable: ANTHROPIC_API_KEY
+
+Usage:
+    # Model is always required
+    from stratifyai.chat import anthropic
+    response = await anthropic.chat("Hello!", model="claude-sonnet-4-5")
+    
+    # Builder pattern (model required)
+    client = (
+        anthropic
+        .with_model("claude-opus-4-5")
+        .with_system("You are a helpful assistant")
+        .with_developer("Use markdown")
+    )
+    response = await client.chat("Hello!")
+"""
+
+import asyncio
+from typing import AsyncIterator, Optional, Union
+
+from stratifyai import LLMClient
+from stratifyai.models import ChatResponse, Message
+from stratifyai.chat.builder import ChatBuilder, create_module_builder
+
+# Default configuration (no default model - must be specified)
+DEFAULT_TEMPERATURE = 0.7
+DEFAULT_MAX_TOKENS = None
+
+# Module-level client (lazy initialization)
+_client: Optional[LLMClient] = None
+
+
+def _get_client() -> LLMClient:
+    """Get or create the module-level client."""
+    global _client
+    if _client is None:
+        _client = LLMClient(provider="anthropic")
+    return _client
+
+
+# Module-level builder for chaining
+_builder = create_module_builder(
+    provider="anthropic",
+    default_temperature=DEFAULT_TEMPERATURE,
+    default_max_tokens=DEFAULT_MAX_TOKENS,
+    client_factory=_get_client,
+)
+
+
+# Builder pattern methods (delegate to _builder)
+def with_model(model: str) -> ChatBuilder:
+    """Set the model to use. Returns a new ChatBuilder for chaining."""
+    return _builder.with_model(model)
+
+
+def with_system(prompt: str) -> ChatBuilder:
+    """Set the system prompt. Returns a new ChatBuilder for chaining."""
+    return _builder.with_system(prompt)
+
+
+def with_developer(instructions: str) -> ChatBuilder:
+    """Set developer instructions. Returns a new ChatBuilder for chaining."""
+    return _builder.with_developer(instructions)
+
+
+def with_temperature(temperature: float) -> ChatBuilder:
+    """Set the temperature. Returns a new ChatBuilder for chaining."""
+    return _builder.with_temperature(temperature)
+
+
+def with_max_tokens(max_tokens: int) -> ChatBuilder:
+    """Set max tokens. Returns a new ChatBuilder for chaining."""
+    return _builder.with_max_tokens(max_tokens)
+
+
+def with_options(**kwargs) -> ChatBuilder:
+    """Set additional options. Returns a new ChatBuilder for chaining."""
+    return _builder.with_options(**kwargs)
+
+
+async def chat(
+    prompt: Union[str, list[Message]],
+    *,
+    model: str,
+    system: Optional[str] = None,
+    temperature: float = DEFAULT_TEMPERATURE,
+    max_tokens: Optional[int] = DEFAULT_MAX_TOKENS,
+    stream: bool = False,
+    **kwargs,
+) -> Union[ChatResponse, AsyncIterator[ChatResponse]]:
+    """
+    Send an async chat completion request to Anthropic Claude.
+
+    Args:
+        prompt: User message string or list of Message objects.
+        model: Model name (required). E.g., "claude-sonnet-4-5", "claude-opus-4-5"
+        system: Optional system prompt (ignored if prompt is list of Messages).
+        temperature: Sampling temperature (0.0-1.0). Default: 0.7
+        max_tokens: Maximum tokens to generate. Default: None (model default)
+        stream: Whether to stream the response. Default: False
+        **kwargs: Additional parameters passed to the API.
+
+    Returns:
+        ChatResponse object, or AsyncIterator[ChatResponse] if streaming.
+
+    Example:
+        >>> from stratifyai.chat import anthropic
+        >>> response = await anthropic.chat("What is Python?", model="claude-sonnet-4-5")
+        >>> print(response.content)
+    """
+    client = _get_client()
+
+    # Build messages list
+    if isinstance(prompt, str):
+        messages = []
+        if system:
+            messages.append(Message(role="system", content=system))
+        messages.append(Message(role="user", content=prompt))
+    else:
+        messages = prompt
+
+    return await client.chat(
+        model=model,
+        messages=messages,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        stream=stream,
+        **kwargs,
+    )
+
+
+async def chat_stream(
+    prompt: Union[str, list[Message]],
+    *,
+    model: str,
+    system: Optional[str] = None,
+    temperature: float = DEFAULT_TEMPERATURE,
+    max_tokens: Optional[int] = DEFAULT_MAX_TOKENS,
+    **kwargs,
+) -> AsyncIterator[ChatResponse]:
+    """
+    Send an async streaming chat completion request to Anthropic Claude.
+
+    Args:
+        prompt: User message string or list of Message objects.
+        model: Model name (required). E.g., "claude-sonnet-4-5"
+        system: Optional system prompt (ignored if prompt is list of Messages).
+        temperature: Sampling temperature (0.0-1.0). Default: 0.7
+        max_tokens: Maximum tokens to generate. Default: None (model default)
+        **kwargs: Additional parameters passed to the API.
+
+    Yields:
+        ChatResponse chunks.
+
+    Example:
+        >>> from stratifyai.chat import anthropic
+        >>> async for chunk in anthropic.chat_stream("Tell me a story", model="claude-sonnet-4-5"):
+        ...     print(chunk.content, end="", flush=True)
+    """
+    return await chat(
+        prompt,
+        model=model,
+        system=system,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        stream=True,
+        **kwargs,
+    )
+
+
+def chat_sync(
+    prompt: Union[str, list[Message]],
+    *,
+    model: str,
+    system: Optional[str] = None,
+    temperature: float = DEFAULT_TEMPERATURE,
+    max_tokens: Optional[int] = DEFAULT_MAX_TOKENS,
+    **kwargs,
+) -> ChatResponse:
+    """Synchronous wrapper for chat(). Model is required."""
+    return asyncio.run(chat(
+        prompt,
+        model=model,
+        system=system,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        stream=False,
+        **kwargs,
+    ))

stratifyai/chat/stratifyai_bedrock.py

@@ -0,0 +1,200 @@
+"""AWS Bedrock chat interface for StratifyAI.
+
+Provides convenient functions for AWS Bedrock chat completions.
+Model must be specified for each request.
+
+Requires: AWS credentials (env vars, ~/.aws/credentials, or IAM role)
+
+Usage:
+    # Model is always required
+    from stratifyai.chat import bedrock
+    response = await bedrock.chat("Hello!", model="anthropic.claude-3-5-sonnet-20241022-v2:0")
+    
+    # Builder pattern (model required)
+    client = (
+        bedrock
+        .with_model("meta.llama3-3-70b-instruct-v1:0")
+        .with_system("You are a helpful assistant")
+        .with_developer("Use markdown")
+    )
+    response = await client.chat("Hello!")
+"""
+
+import asyncio
+from typing import AsyncIterator, Optional, Union
+
+from stratifyai import LLMClient
+from stratifyai.models import ChatResponse, Message
+from stratifyai.chat.builder import ChatBuilder, create_module_builder
+
+# Default configuration (no default model - must be specified)
+DEFAULT_TEMPERATURE = 0.7
+DEFAULT_MAX_TOKENS = None
+
+# Module-level client (lazy initialization)
+_client: Optional[LLMClient] = None
+
+
+def _get_client() -> LLMClient:
+    """Get or create the module-level client."""
+    global _client
+    if _client is None:
+        _client = LLMClient(provider="bedrock")
+    return _client
+
+
+# Module-level builder for chaining
+_builder = create_module_builder(
+    provider="bedrock",
+    default_temperature=DEFAULT_TEMPERATURE,
+    default_max_tokens=DEFAULT_MAX_TOKENS,
+    client_factory=_get_client,
+)
+
+
+# Builder pattern methods (delegate to _builder)
+def with_model(model: str) -> ChatBuilder:
+    """Set the model to use. Returns a new ChatBuilder for chaining."""
+    return _builder.with_model(model)
+
+
+def with_system(prompt: str) -> ChatBuilder:
+    """Set the system prompt. Returns a new ChatBuilder for chaining."""
+    return _builder.with_system(prompt)
+
+
+def with_developer(instructions: str) -> ChatBuilder:
+    """Set developer instructions. Returns a new ChatBuilder for chaining."""
+    return _builder.with_developer(instructions)
+
+
+def with_temperature(temperature: float) -> ChatBuilder:
+    """Set the temperature. Returns a new ChatBuilder for chaining."""
+    return _builder.with_temperature(temperature)
+
+
+def with_max_tokens(max_tokens: int) -> ChatBuilder:
+    """Set max tokens. Returns a new ChatBuilder for chaining."""
+    return _builder.with_max_tokens(max_tokens)
+
+
+def with_options(**kwargs) -> ChatBuilder:
+    """Set additional options. Returns a new ChatBuilder for chaining."""
+    return _builder.with_options(**kwargs)
+
+
+async def chat(
+    prompt: Union[str, list[Message]],
+    *,
+    model: str,
+    system: Optional[str] = None,
+    temperature: float = DEFAULT_TEMPERATURE,
+    max_tokens: Optional[int] = DEFAULT_MAX_TOKENS,
+    stream: bool = False,
+    **kwargs,
+) -> Union[ChatResponse, AsyncIterator[ChatResponse]]:
+    """
+    Send a chat completion request to AWS Bedrock.
+
+    Args:
+        prompt: User message string or list of Message objects.
+        model: Model name (required). E.g., "anthropic.claude-3-5-sonnet-20241022-v2:0"
+        system: Optional system prompt (ignored if prompt is list of Messages).
+        temperature: Sampling temperature (0.0-1.0). Default: 0.7
+        max_tokens: Maximum tokens to generate. Default: None (model default)
+        stream: Whether to stream the response. Default: False
+        **kwargs: Additional parameters passed to the API.
+
+    Returns:
+        ChatResponse object, or AsyncIterator[ChatResponse] if streaming.
+
+    Example:
+        >>> from stratifyai.chat import bedrock
+        >>> response = await bedrock.chat("What is Python?", model="anthropic.claude-3-5-sonnet-20241022-v2:0")
+        >>> print(response.content)
+
+        # Use a different model
+        >>> response = bedrock.chat(
+        ...     "Explain AI",
+        ...     model="meta.llama3-3-70b-instruct-v1:0"
+        ... )
+    """
+    client = _get_client()
+
+    # Build messages list
+    if isinstance(prompt, str):
+        messages = []
+        if system:
+            messages.append(Message(role="system", content=system))
+        messages.append(Message(role="user", content=prompt))
+    else:
+        messages = prompt
+
+    return await client.chat(
+        model=model,
+        messages=messages,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        stream=stream,
+        **kwargs,
+    )
+
+
+async def chat_stream(
+    prompt: Union[str, list[Message]],
+    *,
+    model: str,
+    system: Optional[str] = None,
+    temperature: float = DEFAULT_TEMPERATURE,
+    max_tokens: Optional[int] = DEFAULT_MAX_TOKENS,
+    **kwargs,
+) -> AsyncIterator[ChatResponse]:
+    """
+    Send a streaming chat completion request to AWS Bedrock.
+
+    Args:
+        prompt: User message string or list of Message objects.
+        model: Model name (required). E.g., "anthropic.claude-3-5-sonnet-20241022-v2:0"
+        system: Optional system prompt (ignored if prompt is list of Messages).
+        temperature: Sampling temperature (0.0-1.0). Default: 0.7
+        max_tokens: Maximum tokens to generate. Default: None (model default)
+        **kwargs: Additional parameters passed to the API.
+
+    Yields:
+        ChatResponse chunks.
+
+    Example:
+        >>> from stratifyai.chat import bedrock
+        >>> async for chunk in bedrock.chat_stream("Tell me a story", model="anthropic.claude-3-5-sonnet-20241022-v2:0"):
+        ...     print(chunk.content, end="", flush=True)
+    """
+    return await chat(
+        prompt,
+        model=model,
+        system=system,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        stream=True,
+        **kwargs,
+    )
+
+
+def chat_sync(
+    prompt,
+    *,
+    model: str,
+    system=None,
+    temperature=DEFAULT_TEMPERATURE,
+    max_tokens=DEFAULT_MAX_TOKENS,
+    **kwargs,
+):
+    """Synchronous wrapper for chat()."""
+    return asyncio.run(chat(
+        prompt,
+        model=model,
+        system=system,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        stream=False,
+        **kwargs,
+    ))

stratifyai/chat/stratifyai_deepseek.py

@@ -0,0 +1,194 @@
+"""DeepSeek chat interface for StratifyAI.
+
+Provides convenient functions for DeepSeek chat completions.
+Model must be specified for each request.
+
+Environment Variable: DEEPSEEK_API_KEY
+
+Usage:
+    # Model is always required
+    from stratifyai.chat import deepseek
+    response = await deepseek.chat("Hello!", model="deepseek-chat")
+    
+    # Builder pattern (model required)
+    client = (
+        deepseek
+        .with_model("deepseek-reasoner")
+        .with_system("You are a helpful assistant")
+        .with_developer("Use markdown")
+    )
+    response = await client.chat("Hello!")
+"""
+
+import asyncio
+from typing import AsyncIterator, Optional, Union
+
+from stratifyai import LLMClient
+from stratifyai.models import ChatResponse, Message
+from stratifyai.chat.builder import ChatBuilder, create_module_builder
+
+# Default configuration (no default model - must be specified)
+DEFAULT_TEMPERATURE = 0.7
+DEFAULT_MAX_TOKENS = None
+
+# Module-level client (lazy initialization)
+_client: Optional[LLMClient] = None
+
+
+def _get_client() -> LLMClient:
+    """Get or create the module-level client."""
+    global _client
+    if _client is None:
+        _client = LLMClient(provider="deepseek")
+    return _client
+
+
+# Module-level builder for chaining
+_builder = create_module_builder(
+    provider="deepseek",
+    default_temperature=DEFAULT_TEMPERATURE,
+    default_max_tokens=DEFAULT_MAX_TOKENS,
+    client_factory=_get_client,
+)
+
+
+# Builder pattern methods (delegate to _builder)
+def with_model(model: str) -> ChatBuilder:
+    """Set the model to use. Returns a new ChatBuilder for chaining."""
+    return _builder.with_model(model)
+
+
+def with_system(prompt: str) -> ChatBuilder:
+    """Set the system prompt. Returns a new ChatBuilder for chaining."""
+    return _builder.with_system(prompt)
+
+
+def with_developer(instructions: str) -> ChatBuilder:
+    """Set developer instructions. Returns a new ChatBuilder for chaining."""
+    return _builder.with_developer(instructions)
+
+
+def with_temperature(temperature: float) -> ChatBuilder:
+    """Set the temperature. Returns a new ChatBuilder for chaining."""
+    return _builder.with_temperature(temperature)
+
+
+def with_max_tokens(max_tokens: int) -> ChatBuilder:
+    """Set max tokens. Returns a new ChatBuilder for chaining."""
+    return _builder.with_max_tokens(max_tokens)
+
+
+def with_options(**kwargs) -> ChatBuilder:
+    """Set additional options. Returns a new ChatBuilder for chaining."""
+    return _builder.with_options(**kwargs)
+
+
+async def chat(
+    prompt: Union[str, list[Message]],
+    *,
+    model: str,
+    system: Optional[str] = None,
+    temperature: float = DEFAULT_TEMPERATURE,
+    max_tokens: Optional[int] = DEFAULT_MAX_TOKENS,
+    stream: bool = False,
+    **kwargs,
+) -> Union[ChatResponse, AsyncIterator[ChatResponse]]:
+    """
+    Send a chat completion request to DeepSeek.
+
+    Args:
+        prompt: User message string or list of Message objects.
+        model: Model name (required). E.g., "deepseek-chat", "deepseek-reasoner"
+        system: Optional system prompt (ignored if prompt is list of Messages).
+        temperature: Sampling temperature (0.0-2.0). Default: 0.7
+        max_tokens: Maximum tokens to generate. Default: None (model default)
+        stream: Whether to stream the response. Default: False
+        **kwargs: Additional parameters passed to the API.
+
+    Returns:
+        ChatResponse object, or AsyncIterator[ChatResponse] if streaming.
+
+    Example:
+        >>> from stratifyai.chat import deepseek
+        >>> response = await deepseek.chat("What is Python?", model="deepseek-chat")
+        >>> print(response.content)
+    """
+    client = _get_client()
+
+    # Build messages list
+    if isinstance(prompt, str):
+        messages = []
+        if system:
+            messages.append(Message(role="system", content=system))
+        messages.append(Message(role="user", content=prompt))
+    else:
+        messages = prompt
+
+    return await client.chat(
+        model=model,
+        messages=messages,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        stream=stream,
+        **kwargs,
+    )
+
+
+async def chat_stream(
+    prompt: Union[str, list[Message]],
+    *,
+    model: str,
+    system: Optional[str] = None,
+    temperature: float = DEFAULT_TEMPERATURE,
+    max_tokens: Optional[int] = DEFAULT_MAX_TOKENS,
+    **kwargs,
+) -> AsyncIterator[ChatResponse]:
+    """
+    Send a streaming chat completion request to DeepSeek.
+
+    Args:
+        prompt: User message string or list of Message objects.
+        model: Model name (required). E.g., "deepseek-chat"
+        system: Optional system prompt (ignored if prompt is list of Messages).
+        temperature: Sampling temperature (0.0-2.0). Default: 0.7
+        max_tokens: Maximum tokens to generate. Default: None (model default)
+        **kwargs: Additional parameters passed to the API.
+
+    Yields:
+        ChatResponse chunks.
+
+    Example:
+        >>> from stratifyai.chat import deepseek
+        >>> async for chunk in deepseek.chat_stream("Tell me a story", model="deepseek-chat"):
+        ...     print(chunk.content, end="", flush=True)
+    """
+    return await chat(
+        prompt,
+        model=model,
+        system=system,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        stream=True,
+        **kwargs,
+    )
+
+
+def chat_sync(
+    prompt,
+    *,
+    model: str,
+    system=None,
+    temperature=DEFAULT_TEMPERATURE,
+    max_tokens=DEFAULT_MAX_TOKENS,
+    **kwargs,
+):
+    """Synchronous wrapper for chat()."""
+    return asyncio.run(chat(
+        prompt,
+        model=model,
+        system=system,
+        temperature=temperature,
+        max_tokens=max_tokens,
+        stream=False,
+        **kwargs,
+    ))