From 38a378b774921acb2b841c4bddc449628d3c474c Mon Sep 17 00:00:00 2001 From: Vardhman Gupta Date: Fri, 2 Oct 2026 07:09:08 +0530 Subject: [PATCH 1/3] feat(chat): add multi-provider chat/LLM abstraction (#285) --- src/dynavec/chat/__init__.py | 45 +++++ src/dynavec/chat/anthropic.py | 143 ++++++++++++++ src/dynavec/chat/base.py | 89 +++++++++ src/dynavec/chat/bedrock.py | 176 +++++++++++++++++ src/dynavec/chat/openai.py | 137 ++++++++++++++ tests/test_chat_models.py | 345 ++++++++++++++++++++++++++++++++++ 6 files changed, 935 insertions(+) create mode 100644 src/dynavec/chat/__init__.py create mode 100644 src/dynavec/chat/anthropic.py create mode 100644 src/dynavec/chat/base.py create mode 100644 src/dynavec/chat/bedrock.py create mode 100644 src/dynavec/chat/openai.py create mode 100644 tests/test_chat_models.py diff --git a/src/dynavec/chat/__init__.py b/src/dynavec/chat/__init__.py new file mode 100644 index 0000000..e9736d6 --- /dev/null +++ b/src/dynavec/chat/__init__.py @@ -0,0 +1,45 @@ +"""Chat model abstraction and provider implementations.""" + +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(): + from dynavec.chat.openai import OpenAIChatModel + return OpenAIChatModel + +def _import_anthropic(): + from dynavec.chat.anthropic import AnthropicChatModel + return AnthropicChatModel + +def _import_bedrock(): + from dynavec.chat.bedrock import BedrockChatModel + return BedrockChatModel + +def __getattr__(name: str): + if name == "OpenAIChatModel": + return _import_openai() + elif name == "AnthropicChatModel": + return _import_anthropic() + elif name == "BedrockChatModel": + return _import_bedrock() + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + +__all__.extend(["OpenAIChatModel", "AnthropicChatModel", "BedrockChatModel"]) diff --git a/src/dynavec/chat/anthropic.py b/src/dynavec/chat/anthropic.py new file mode 100644 index 0000000..e4e34c2 --- /dev/null +++ b/src/dynavec/chat/anthropic.py @@ -0,0 +1,143 @@ +"""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 == "text_stream": + yield ChatChunk(content=event.text) + elif event.type == "input_json": + yield ChatChunk( + tool_calls=[ + ToolCall(id="", name="", arguments=event.partial_json) + ] + ) diff --git a/src/dynavec/chat/base.py b/src/dynavec/chat/base.py new file mode 100644 index 0000000..7a015f1 --- /dev/null +++ b/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(): + return list(self.stream(messages, tools, **kwargs)) + + chunks = await asyncio.to_thread(_sync_stream) + for chunk in chunks: + yield chunk diff --git a/src/dynavec/chat/bedrock.py b/src/dynavec/chat/bedrock.py new file mode 100644 index 0000000..710ecc8 --- /dev/null +++ b/src/dynavec/chat/bedrock.py @@ -0,0 +1,176 @@ +"""AWS Bedrock chat model.""" + +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 BedrockChatModel(ChatModel): + """Chat with AWS Bedrock via the unified Converse API. + + Parameters + ---------- + model: + Model id, e.g. ``"anthropic.claude-3-haiku-20240307-v1:0"``. + region_name: + Optional AWS region name. + """ + + def __init__( + self, + model: str = "anthropic.claude-3-haiku-20240307-v1:0", + region_name: str | None = None, + ) -> None: + try: + import boto3 + except ImportError as exc: + raise MissingDependencyError("BedrockChatModel", "boto3", "boto3") from exc + + self.model = model + self._client = boto3.client("bedrock-runtime", region_name=region_name) + + def _convert_messages(self, messages: list[Message]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + system_prompts: list[dict[str, Any]] = [] + bedrock_msgs: list[dict[str, Any]] = [] + + for msg in messages: + if msg.role == "system": + if msg.content: + system_prompts.append({"text": msg.content}) + continue + + content_blocks: list[dict[str, Any]] = [] + if msg.content: + content_blocks.append({"text": msg.content}) + + if msg.tool_calls: + for tc in msg.tool_calls: + content_blocks.append({ + "toolUse": { + "toolUseId": 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({ + "toolResult": { + "toolUseId": msg.tool_call_id, + "content": [{"text": msg.content or ""}], + "status": "success" + } + }) + bedrock_msgs.append({ + "role": "user", + "content": content_blocks + }) + continue + + bedrock_msgs.append({ + "role": "assistant" if msg.role == "assistant" else "user", + "content": content_blocks + }) + + return system_prompts, bedrock_msgs + + def _convert_tools(self, tools: list[Tool] | None) -> dict[str, Any] | None: + if not tools: + return None + return { + "tools": [ + { + "toolSpec": { + "name": t.name, + "description": t.description, + "inputSchema": { + "json": t.parameters + } + } + } + for t in tools + ] + } + + def invoke( + self, messages: list[Message], tools: list[Tool] | None = None, **kwargs: Any + ) -> ChatResult: + system, bedrock_msgs = self._convert_messages(messages) + bedrock_tools = self._convert_tools(tools) + + args: dict[str, Any] = { + "modelId": self.model, + "messages": bedrock_msgs, + **kwargs + } + if system: + args["system"] = system + if bedrock_tools: + args["toolConfig"] = bedrock_tools + + resp = self._client.converse(**args) + + output_msg = resp["output"]["message"] + role = output_msg["role"] + + text_content = "" + tool_calls = [] + + for block in output_msg.get("content", []): + if "text" in block: + text_content += block["text"] + elif "toolUse" in block: + tu = block["toolUse"] + tool_calls.append( + ToolCall( + id=tu["toolUseId"], + name=tu["name"], + arguments=json.dumps(tu["input"]) + ) + ) + + stop_reason = resp.get("stopReason") + out = Message(role=role, content=text_content if text_content else None, tool_calls=tool_calls) + return ChatResult(message=out, finish_reason=stop_reason) + + def stream( + self, messages: list[Message], tools: list[Tool] | None = None, **kwargs: Any + ) -> Iterator[ChatChunk]: + system, bedrock_msgs = self._convert_messages(messages) + bedrock_tools = self._convert_tools(tools) + + args: dict[str, Any] = { + "modelId": self.model, + "messages": bedrock_msgs, + **kwargs + } + if system: + args["system"] = system + if bedrock_tools: + args["toolConfig"] = bedrock_tools + + resp = self._client.converse_stream(**args) + stream = resp.get("stream") + if not stream: + return + + for event in stream: + if "contentBlockDelta" in event: + delta = event["contentBlockDelta"]["delta"] + if "text" in delta: + yield ChatChunk(content=delta["text"]) + elif "toolUse" in delta: + yield ChatChunk( + tool_calls=[ + ToolCall( + id="", + name="", + arguments=delta["toolUse"].get("input", "") + ) + ] + ) diff --git a/src/dynavec/chat/openai.py b/src/dynavec/chat/openai.py new file mode 100644 index 0000000..5490710 --- /dev/null +++ b/src/dynavec/chat/openai.py @@ -0,0 +1,137 @@ +"""OpenAI chat model (bring your own OPENAI_API_KEY).""" + +from __future__ import annotations + +from collections.abc import Iterator +from typing import Any + +from ..exceptions import MissingDependencyError +from .base import ChatChunk, ChatModel, ChatResult, Message, Tool, ToolCall + + +class OpenAIChatModel(ChatModel): + """Chat with OpenAI's models. + + Parameters + ---------- + model: + Model id, e.g. ``"gpt-4o"``. + api_key: + Optional; falls back to the ``OPENAI_API_KEY`` environment variable. + base_url: + Optional custom endpoint URL. + """ + + def __init__( + self, + model: str = "gpt-4o", + api_key: str | None = None, + base_url: str | None = None, + ) -> None: + try: + from openai import OpenAI + except ImportError as exc: + raise MissingDependencyError("OpenAIChatModel", "openai", "openai") from exc + + self.model = model + self._client = OpenAI(api_key=api_key, base_url=base_url) + + def _convert_messages(self, messages: list[Message]) -> list[dict[str, Any]]: + out = [] + for msg in messages: + d: dict[str, Any] = {"role": msg.role} + + if msg.content is not None: + d["content"] = msg.content + + if msg.tool_calls: + d["tool_calls"] = [ + { + "id": tc.id, + "type": "function", + "function": {"name": tc.name, "arguments": tc.arguments}, + } + for tc in msg.tool_calls + ] + + if msg.tool_call_id: + d["tool_call_id"] = msg.tool_call_id + d["name"] = "tool" # some models/APIs require a name + + out.append(d) + return out + + def _convert_tools(self, tools: list[Tool] | None) -> list[dict[str, Any]] | None: + if not tools: + return None + return [ + { + "type": "function", + "function": { + "name": t.name, + "description": t.description, + "parameters": t.parameters, + }, + } + for t in tools + ] + + def invoke( + self, messages: list[Message], tools: list[Tool] | None = None, **kwargs: Any + ) -> ChatResult: + """Invoke the chat model synchronously.""" + openai_msgs = self._convert_messages(messages) + openai_tools = self._convert_tools(tools) + + args: dict[str, Any] = {"model": self.model, "messages": openai_msgs, **kwargs} + if openai_tools: + args["tools"] = openai_tools + + resp = self._client.chat.completions.create(**args) + choice = resp.choices[0] + + out_msg = Message(role=choice.message.role or "assistant", content=choice.message.content) + if choice.message.tool_calls: + out_msg.tool_calls = [ + ToolCall( + id=tc.id, + name=tc.function.name, + arguments=tc.function.arguments, + ) + for tc in choice.message.tool_calls + ] + + return ChatResult(message=out_msg, finish_reason=choice.finish_reason) + + def stream( + self, messages: list[Message], tools: list[Tool] | None = None, **kwargs: Any + ) -> Iterator[ChatChunk]: + """Stream the chat model response.""" + openai_msgs = self._convert_messages(messages) + openai_tools = self._convert_tools(tools) + + args: dict[str, Any] = {"model": self.model, "messages": openai_msgs, "stream": True, **kwargs} + if openai_tools: + args["tools"] = openai_tools + + resp = self._client.chat.completions.create(**args) + + for chunk in resp: + if not chunk.choices: + continue + + choice = chunk.choices[0] + delta = choice.delta + + tool_calls = [] + if delta.tool_calls: + for tc in delta.tool_calls: + tool_calls.append( + ToolCall( + id=tc.id or "", + name=tc.function.name if (tc.function and tc.function.name) else "", + arguments=tc.function.arguments if (tc.function and tc.function.arguments) else "", + ) + ) + + yield ChatChunk(content=delta.content, tool_calls=tool_calls) diff --git a/tests/test_chat_models.py b/tests/test_chat_models.py new file mode 100644 index 0000000..d86443d --- /dev/null +++ b/tests/test_chat_models.py @@ -0,0 +1,345 @@ +"""Comprehensive tests for chat model providers.""" + +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + +import dynavec.chat as chat +from dynavec.chat.anthropic import AnthropicChatModel +from dynavec.chat.base import Message, Tool +from dynavec.chat.bedrock import BedrockChatModel +from dynavec.chat.openai import OpenAIChatModel +from dynavec.exceptions import MissingDependencyError + + +# --------------------------------------------------------------------------- +# OpenAI Fakes +# --------------------------------------------------------------------------- +class FakeOpenAIChatCompletions: + def __init__(self): + self.last_kwargs = None + + def create(self, **kwargs): + self.last_kwargs = kwargs + if kwargs.get("stream"): + return [ + SimpleNamespace( + choices=[ + SimpleNamespace( + delta=SimpleNamespace( + content="Streamed ", + tool_calls=[ + SimpleNamespace( + id="call_1", + function=SimpleNamespace( + name="search", arguments='{"q":"dynavec"}' + ), + ) + ], + ) + ) + ] + ) + ] + + if kwargs.get("tools"): + return SimpleNamespace( + choices=[ + SimpleNamespace( + message=SimpleNamespace( + role="assistant", + content=None, + tool_calls=[ + SimpleNamespace( + id="call_abc", + function=SimpleNamespace( + name="search", + arguments='{"query": "vector search"}', + ), + ) + ], + ), + finish_reason="tool_calls", + ) + ] + ) + + return SimpleNamespace( + choices=[ + SimpleNamespace( + message=SimpleNamespace( + role="assistant", + content="Hello from fake OpenAI", + tool_calls=[], + ), + finish_reason="stop", + ) + ] + ) + + +class FakeOpenAIClient: + def __init__(self): + self.chat = SimpleNamespace(completions=FakeOpenAIChatCompletions()) + + +# --------------------------------------------------------------------------- +# Anthropic Fakes +# --------------------------------------------------------------------------- +class FakeAnthropicMessages: + def __init__(self): + self.last_kwargs = None + + def create(self, **kwargs): + self.last_kwargs = kwargs + if kwargs.get("tools"): + return SimpleNamespace( + content=[ + SimpleNamespace( + type="tool_use", + id="toolu_123", + name="get_weather", + input={"location": "SF"}, + ) + ], + stop_reason="tool_use", + ) + return SimpleNamespace( + content=[ + SimpleNamespace( + type="text", + text="Hello from fake Anthropic", + ) + ], + stop_reason="end_turn", + ) + + def stream(self, **kwargs): + self.last_kwargs = kwargs + + class FakeStreamContext: + def __enter__(self): + return [ + SimpleNamespace(type="text_stream", text="Streamed "), + SimpleNamespace(type="input_json", partial_json='{"loc":"NY"}'), + ] + + def __exit__(self, exc_type, exc_val, exc_tb): + pass + + return FakeStreamContext() + + +class FakeAnthropicClient: + def __init__(self): + self.messages = FakeAnthropicMessages() + + +# --------------------------------------------------------------------------- +# Bedrock Fakes +# --------------------------------------------------------------------------- +class FakeBedrockClient: + def __init__(self): + self.last_kwargs = None + + def converse(self, **kwargs): + self.last_kwargs = kwargs + if "toolConfig" in kwargs: + return { + "output": { + "message": { + "role": "assistant", + "content": [ + { + "toolUse": { + "toolUseId": "tool_bedrock_1", + "name": "lookup", + "input": {"id": "123"}, + } + } + ], + } + }, + "stopReason": "tool_use", + } + return { + "output": { + "message": { + "role": "assistant", + "content": [{"text": "Hello from fake Bedrock"}], + } + }, + "stopReason": "end_turn", + } + + def converse_stream(self, **kwargs): + self.last_kwargs = kwargs + return { + "stream": [ + { + "contentBlockDelta": { + "delta": {"text": "Streamed "}, + } + }, + { + "contentBlockDelta": { + "delta": {"toolUse": {"input": '{"arg":"val"}'}}, + } + }, + ] + } + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- +def test_chat_module_exports(): + assert hasattr(chat, "ChatModel") + assert hasattr(chat, "Message") + assert hasattr(chat, "Tool") + assert hasattr(chat, "ToolCall") + assert hasattr(chat, "ChatResult") + assert hasattr(chat, "ChatChunk") + assert hasattr(chat, "OpenAIChatModel") + assert hasattr(chat, "AnthropicChatModel") + assert hasattr(chat, "BedrockChatModel") + + +def test_openai_chat_model(monkeypatch): + fake_client = FakeOpenAIClient() + + def fake_init(self, *args, **kwargs): + self.model = "gpt-4o" + self._client = fake_client + + monkeypatch.setattr("dynavec.chat.openai.OpenAIChatModel.__init__", fake_init) + + model = OpenAIChatModel() + + # Plain invoke + res = model.invoke([Message(role="user", content="Hi")]) + assert res.message.content == "Hello from fake OpenAI" + assert res.finish_reason == "stop" + + # Tool invoke + tool = Tool( + name="search", + description="Search vector store", + parameters={"type": "object", "properties": {"query": {"type": "string"}}}, + ) + res_tool = model.invoke([Message(role="user", content="Search query")], tools=[tool]) + assert res_tool.finish_reason == "tool_calls" + assert len(res_tool.message.tool_calls) == 1 + assert res_tool.message.tool_calls[0].name == "search" + assert json.loads(res_tool.message.tool_calls[0].arguments) == {"query": "vector search"} + + # Stream + chunks = list(model.stream([Message(role="user", content="Hi")])) + assert chunks[0].content == "Streamed " + assert len(chunks[0].tool_calls) == 1 + + +def test_anthropic_chat_model(monkeypatch): + fake_client = FakeAnthropicClient() + + def fake_init(self, *args, **kwargs): + self.model = "claude-3-5-sonnet-20240620" + self._client = fake_client + + monkeypatch.setattr("dynavec.chat.anthropic.AnthropicChatModel.__init__", fake_init) + + model = AnthropicChatModel() + + # Plain invoke with system prompt + res = model.invoke( + [ + Message(role="system", content="You are a helpful assistant."), + Message(role="user", content="Hi"), + ] + ) + assert res.message.content == "Hello from fake Anthropic" + assert res.finish_reason == "end_turn" + assert fake_client.messages.last_kwargs.get("system") == "You are a helpful assistant." + + # Tool invoke + tool = Tool( + name="get_weather", + description="Get current weather", + parameters={"type": "object", "properties": {"location": {"type": "string"}}}, + ) + res_tool = model.invoke([Message(role="user", content="What's weather?")], tools=[tool]) + assert res_tool.finish_reason == "tool_use" + assert len(res_tool.message.tool_calls) == 1 + assert res_tool.message.tool_calls[0].id == "toolu_123" + assert json.loads(res_tool.message.tool_calls[0].arguments) == {"location": "SF"} + + # Stream + chunks = list(model.stream([Message(role="user", content="Hi")])) + assert chunks[0].content == "Streamed " + assert chunks[1].tool_calls[0].arguments == '{"loc":"NY"}' + + +def test_bedrock_chat_model(monkeypatch): + fake_client = FakeBedrockClient() + + def fake_init(self, *args, **kwargs): + self.model = "anthropic.claude-3-haiku-20240307-v1:0" + self._client = fake_client + + monkeypatch.setattr("dynavec.chat.bedrock.BedrockChatModel.__init__", fake_init) + + model = BedrockChatModel() + + # Plain invoke + res = model.invoke([Message(role="user", content="Hi")]) + assert res.message.content == "Hello from fake Bedrock" + assert res.finish_reason == "end_turn" + + # Tool invoke + tool = Tool( + name="lookup", + description="Lookup record", + parameters={"type": "object", "properties": {"id": {"type": "string"}}}, + ) + res_tool = model.invoke([Message(role="user", content="Lookup")], tools=[tool]) + assert res_tool.finish_reason == "tool_use" + assert len(res_tool.message.tool_calls) == 1 + assert res_tool.message.tool_calls[0].name == "lookup" + + # Stream + chunks = list(model.stream([Message(role="user", content="Hi")])) + assert chunks[0].content == "Streamed " + assert chunks[1].tool_calls[0].arguments == '{"arg":"val"}' + + +@pytest.mark.asyncio +async def test_async_chat_model_delegation(monkeypatch): + def fake_init(self, *args, **kwargs): + self.model = "gpt-4o" + self._client = FakeOpenAIClient() + + monkeypatch.setattr("dynavec.chat.openai.OpenAIChatModel.__init__", fake_init) + + model = OpenAIChatModel() + res = await model.ainvoke([Message(role="user", content="Hi")]) + assert res.message.content == "Hello from fake OpenAI" + + streamed_chunks = [] + async for chunk in model.astream([Message(role="user", content="Hi")]): + streamed_chunks.append(chunk) + + assert len(streamed_chunks) == 1 + assert streamed_chunks[0].content == "Streamed " + + +def test_missing_dependency_guards(monkeypatch): + import sys + + # Simulate missing openai module + monkeypatch.setitem(sys.modules, "openai", None) + with pytest.raises(MissingDependencyError) as exc_info: + OpenAIChatModel() + assert "OpenAIChatModel" in str(exc_info.value) From c50cef33cb425e6af33b0e924734c10dbe9f344e Mon Sep 17 00:00:00 2001 From: Vardhman Gupta Date: Fri, 2 Oct 2026 07:15:22 +0530 Subject: [PATCH 2/3] fix(types): add return annotations and anthropic typecheck dependencies --- pyproject.toml | 5 +++++ src/dynavec/chat/__init__.py | 24 ++++++++++++++++++------ src/dynavec/chat/base.py | 2 +- 3 files changed, 24 insertions(+), 7 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 4f31016..cda2c69 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/dynavec/chat/__init__.py b/src/dynavec/chat/__init__.py index e9736d6..836f81e 100644 --- a/src/dynavec/chat/__init__.py +++ b/src/dynavec/chat/__init__.py @@ -1,5 +1,9 @@ """Chat model abstraction and provider implementations.""" +from __future__ import annotations + +from typing import Any + from dynavec.chat.base import ( ChatChunk, ChatModel, @@ -20,26 +24,34 @@ "ToolCall", ] + # Optional imports for providers -def _import_openai(): +def _import_openai() -> Any: from dynavec.chat.openai import OpenAIChatModel + return OpenAIChatModel -def _import_anthropic(): + +def _import_anthropic() -> Any: from dynavec.chat.anthropic import AnthropicChatModel + return AnthropicChatModel -def _import_bedrock(): + +def _import_bedrock() -> Any: from dynavec.chat.bedrock import BedrockChatModel + return BedrockChatModel -def __getattr__(name: str): + +def __getattr__(name: str) -> Any: if name == "OpenAIChatModel": return _import_openai() - elif name == "AnthropicChatModel": + if name == "AnthropicChatModel": return _import_anthropic() - elif name == "BedrockChatModel": + if name == "BedrockChatModel": return _import_bedrock() raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + __all__.extend(["OpenAIChatModel", "AnthropicChatModel", "BedrockChatModel"]) diff --git a/src/dynavec/chat/base.py b/src/dynavec/chat/base.py index 7a015f1..876c126 100644 --- a/src/dynavec/chat/base.py +++ b/src/dynavec/chat/base.py @@ -81,7 +81,7 @@ async def astream( **kwargs: Any ) -> AsyncIterator[ChatChunk]: """Stream the chat model response asynchronously. Default delegates to thread.""" - def _sync_stream(): + def _sync_stream() -> list[ChatChunk]: return list(self.stream(messages, tools, **kwargs)) chunks = await asyncio.to_thread(_sync_stream) From 04030ce014ef521602ead697a61e91ea02564237 Mon Sep 17 00:00:00 2001 From: Vardhman Gupta Date: Fri, 2 Oct 2026 07:20:54 +0530 Subject: [PATCH 3/3] fix(anthropic): use content_block_delta event types in streaming --- src/dynavec/chat/anthropic.py | 17 +++++++++++++---- tests/test_chat_models.py | 12 ++++++++++-- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/src/dynavec/chat/anthropic.py b/src/dynavec/chat/anthropic.py index e4e34c2..2eb4aa3 100644 --- a/src/dynavec/chat/anthropic.py +++ b/src/dynavec/chat/anthropic.py @@ -133,11 +133,20 @@ def stream( with self._client.messages.stream(**args) as stream: for event in stream: - if event.type == "text_stream": - yield ChatChunk(content=event.text) - elif event.type == "input_json": + 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=event.partial_json) + ToolCall( + id="", + name="", + arguments=delta.partial_json, + ) ] ) diff --git a/tests/test_chat_models.py b/tests/test_chat_models.py index d86443d..820c1a1 100644 --- a/tests/test_chat_models.py +++ b/tests/test_chat_models.py @@ -123,8 +123,16 @@ def stream(self, **kwargs): class FakeStreamContext: def __enter__(self): return [ - SimpleNamespace(type="text_stream", text="Streamed "), - SimpleNamespace(type="input_json", partial_json='{"loc":"NY"}'), + SimpleNamespace( + type="content_block_delta", + delta=SimpleNamespace(type="text_delta", text="Streamed "), + ), + SimpleNamespace( + type="content_block_delta", + delta=SimpleNamespace( + type="input_json_delta", partial_json='{"loc":"NY"}' + ), + ), ] def __exit__(self, exc_type, exc_val, exc_tb):