feat(chat): add multi-provider chat/LLM abstraction (#285)

pyproject.toml

@@ -34,6 +34,7 @@ dependencies = [
 [project.optional-dependencies]
 # --- Pluggable, bring-your-own-key embedders (users pick one or more) ---
 openai = ["openai>=1.40"]
+anthropic = ["anthropic>=0.40"]
 gemini = ["google-generativeai>=0.8"]
 cohere = ["cohere>=5.0"]
 voyage = ["voyageai>=0.3.2"]
@@ -64,9 +65,11 @@ typecheck = [
     "haystack-ai>=2.0",
     "dspy>=3.3; python_version >= '3.10'",
     "semantic-kernel>=1.44; python_version >= '3.10'",
+    "anthropic>=0.40",
 ]
 all = [
     "openai>=1.40",
+    "anthropic>=0.40",
     "google-generativeai>=0.8",
     "cohere>=5.0",
     "voyageai>=0.3.2",
@@ -141,6 +144,8 @@ show_error_codes = true
 # dynavec code that uses them.
 [[tool.mypy.overrides]]
 module = [
+    "anthropic",
+    "anthropic.*",
     "bs4",
     "bs4.*",
     "boto3",

src/dynavec/chat/__init__.py

@@ -0,0 +1,57 @@
+"""Chat model abstraction and provider implementations."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from dynavec.chat.base import (
+    ChatChunk,
+    ChatModel,
+    ChatResult,
+    Message,
+    Role,
+    Tool,
+    ToolCall,
+)
+
+__all__ = [
+    "ChatChunk",
+    "ChatModel",
+    "ChatResult",
+    "Message",
+    "Role",
+    "Tool",
+    "ToolCall",
+]
+
+
+# Optional imports for providers
+def _import_openai() -> Any:
+    from dynavec.chat.openai import OpenAIChatModel
+
+    return OpenAIChatModel
+
+
+def _import_anthropic() -> Any:
+    from dynavec.chat.anthropic import AnthropicChatModel
+
+    return AnthropicChatModel
+
+
+def _import_bedrock() -> Any:
+    from dynavec.chat.bedrock import BedrockChatModel
+
+    return BedrockChatModel
+
+
+def __getattr__(name: str) -> Any:
+    if name == "OpenAIChatModel":
+        return _import_openai()
+    if name == "AnthropicChatModel":
+        return _import_anthropic()
+    if name == "BedrockChatModel":
+        return _import_bedrock()
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__.extend(["OpenAIChatModel", "AnthropicChatModel", "BedrockChatModel"])

src/dynavec/chat/anthropic.py

@@ -0,0 +1,152 @@
+"""Anthropic chat model (bring your own ANTHROPIC_API_KEY)."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterator
+from typing import Any
+
+from ..exceptions import MissingDependencyError
+from .base import ChatChunk, ChatModel, ChatResult, Message, Tool, ToolCall
+
+
+class AnthropicChatModel(ChatModel):
+    """Chat with Anthropic's Claude models.
+
+    Parameters
+    ----------
+    model:
+        Model id, e.g. ``"claude-3-5-sonnet-20240620"``.
+    api_key:
+        Optional; falls back to the ``ANTHROPIC_API_KEY`` environment variable.
+    """
+
+    def __init__(
+        self,
+        model: str = "claude-3-5-sonnet-20240620",
+        api_key: str | None = None,
+    ) -> None:
+        try:
+            from anthropic import Anthropic
+        except ImportError as exc:
+            raise MissingDependencyError("AnthropicChatModel", "anthropic", "anthropic") from exc
+
+        self.model = model
+        self._client = Anthropic(api_key=api_key)
+
+    def _convert_messages(self, messages: list[Message]) -> tuple[str, list[dict[str, Any]]]:
+        system_prompt = ""
+        anthropic_msgs = []
+
+        for msg in messages:
+            if msg.role == "system":
+                system_prompt += (msg.content or "") + "\n"
+                continue
+
+            content_blocks: list[dict[str, Any]] = []
+            if msg.content:
+                content_blocks.append({"type": "text", "text": msg.content})
+
+            if msg.tool_calls:
+                for tc in msg.tool_calls:
+                    content_blocks.append({
+                        "type": "tool_use",
+                        "id": tc.id,
+                        "name": tc.name,
+                        "input": json.loads(tc.arguments) if tc.arguments else {}
+                    })
+
+            if msg.role == "tool" and msg.tool_call_id:
+                content_blocks.append({
+                    "type": "tool_result",
+                    "tool_use_id": msg.tool_call_id,
+                    "content": msg.content or ""
+                })
+                anthropic_msgs.append({
+                    "role": "user",
+                    "content": content_blocks
+                })
+                continue
+
+            anthropic_msgs.append({
+                "role": "assistant" if msg.role == "assistant" else "user",
+                "content": content_blocks
+            })
+
+        return system_prompt.strip(), anthropic_msgs
+
+    def _convert_tools(self, tools: list[Tool] | None) -> list[dict[str, Any]] | None:
+        if not tools:
+            return None
+        return [
+            {
+                "name": t.name,
+                "description": t.description,
+                "input_schema": t.parameters,
+            }
+            for t in tools
+        ]
+
+    def invoke(
+        self, messages: list[Message], tools: list[Tool] | None = None, **kwargs: Any
+    ) -> ChatResult:
+        system, anthropic_msgs = self._convert_messages(messages)
+        anthropic_tools = self._convert_tools(tools)
+
+        args: dict[str, Any] = {"model": self.model, "messages": anthropic_msgs, "max_tokens": 1024, **kwargs}
+        if system:
+            args["system"] = system
+        if anthropic_tools:
+            args["tools"] = anthropic_tools
+
+        resp = self._client.messages.create(**args)
+
+        text_content = ""
+        tool_calls = []
+
+        for block in resp.content:
+            if block.type == "text":
+                text_content += block.text
+            elif block.type == "tool_use":
+                tool_calls.append(
+                    ToolCall(
+                        id=block.id,
+                        name=block.name,
+                        arguments=json.dumps(block.input)
+                    )
+                )
+
+        out_msg = Message(role="assistant", content=text_content if text_content else None, tool_calls=tool_calls)
+        return ChatResult(message=out_msg, finish_reason=resp.stop_reason)
+
+    def stream(
+        self, messages: list[Message], tools: list[Tool] | None = None, **kwargs: Any
+    ) -> Iterator[ChatChunk]:
+        system, anthropic_msgs = self._convert_messages(messages)
+        anthropic_tools = self._convert_tools(tools)
+
+        args: dict[str, Any] = {"model": self.model, "messages": anthropic_msgs, "max_tokens": 1024, **kwargs}
+        if system:
+            args["system"] = system
+        if anthropic_tools:
+            args["tools"] = anthropic_tools
+
+        with self._client.messages.stream(**args) as stream:
+            for event in stream:
+                if event.type != "content_block_delta":
+                    continue
+
+                delta = event.delta
+
+                if delta.type == "text_delta":
+                    yield ChatChunk(content=delta.text)
+                elif delta.type == "input_json_delta":
+                    yield ChatChunk(
+                        tool_calls=[
+                            ToolCall(
+                                id="",
+                                name="",
+                                arguments=delta.partial_json,
+                            )
+                        ]
+                    )

src/dynavec/chat/base.py

@@ -0,0 +1,89 @@
+"""Chat model abstraction for dynavec."""
+
+from __future__ import annotations
+
+import asyncio
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator, Iterator
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+Role = Literal["system", "user", "assistant", "tool"]
+
+@dataclass
+class ToolCall:
+    """A tool invocation requested by the model."""
+    id: str
+    name: str
+    arguments: str  # JSON string
+
+@dataclass
+class Message:
+    """A single message in a chat conversation."""
+    role: Role
+    content: str | None = None
+    tool_calls: list[ToolCall] = field(default_factory=list)
+    tool_call_id: str | None = None
+
+@dataclass
+class Tool:
+    """A tool/function specification provided to the model."""
+    name: str
+    description: str
+    parameters: dict[str, Any]  # JSON schema dict
+
+@dataclass
+class ChatResult:
+    """The final result of a chat invocation."""
+    message: Message
+    finish_reason: str | None = None
+
+@dataclass
+class ChatChunk:
+    """A streamed chunk from a chat invocation."""
+    content: str | None = None
+    tool_calls: list[ToolCall] = field(default_factory=list)
+
+class ChatModel(ABC):
+    """Base class for all chat/LLM backends."""
+
+    @abstractmethod
+    def invoke(
+        self, 
+        messages: list[Message], 
+        tools: list[Tool] | None = None, 
+        **kwargs: Any
+    ) -> ChatResult:
+        """Invoke the chat model synchronously."""
+
+    @abstractmethod
+    def stream(
+        self, 
+        messages: list[Message], 
+        tools: list[Tool] | None = None, 
+        **kwargs: Any
+    ) -> Iterator[ChatChunk]:
+        """Stream the chat model response."""
+
+    async def ainvoke(
+        self, 
+        messages: list[Message], 
+        tools: list[Tool] | None = None, 
+        **kwargs: Any
+    ) -> ChatResult:
+        """Invoke the chat model asynchronously. Default delegates to thread."""
+        return await asyncio.to_thread(self.invoke, messages, tools, **kwargs)
+
+    async def astream(
+        self, 
+        messages: list[Message], 
+        tools: list[Tool] | None = None, 
+        **kwargs: Any
+    ) -> AsyncIterator[ChatChunk]:
+        """Stream the chat model response asynchronously. Default delegates to thread."""
+        def _sync_stream() -> list[ChatChunk]:
+            return list(self.stream(messages, tools, **kwargs))
+
+        chunks = await asyncio.to_thread(_sync_stream)
+        for chunk in chunks:
+            yield chunk