From 10e56b62b34b2282ea27e37be225f908d3572ce2 Mon Sep 17 00:00:00 2001 From: Uzmaa7 Date: Mon, 5 Oct 2026 15:41:29 +0530 Subject: [PATCH 1/4] feat(orchestration): implement human-in-the-loop interrupts and resume workflow (#296) --- src/dynavec/agents/base.py | 68 ++++++++++++-- src/dynavec/agents/react.py | 134 ++++++++++++++++++++++---- src/dynavec/exceptions.py | 28 ++++++ tests/test_agents.py | 28 +++--- tests/test_hitl_interrupts.py | 172 ++++++++++++++++++++++++++++++++++ 5 files changed, 391 insertions(+), 39 deletions(-) create mode 100644 tests/test_hitl_interrupts.py diff --git a/src/dynavec/agents/base.py b/src/dynavec/agents/base.py index 9fea9b2..102fb9c 100644 --- a/src/dynavec/agents/base.py +++ b/src/dynavec/agents/base.py @@ -6,9 +6,12 @@ import json from collections.abc import Callable from dataclasses import dataclass, field -from typing import Any + from ..chat.base import Tool, ToolCall +import functools +from typing import Any, overload + @dataclass @@ -28,10 +31,10 @@ class AgentResult: output: str steps: list[AgentStep] = field(default_factory=list) finished: bool = True - termination_reason: str = "completed" # "completed", "max_steps_reached", "error" + termination_reason: str = "completed" # "completed", "max_steps_reached", "error", "interrupted" total_steps: int = 0 tool_calls_count: int = 0 - + interrupt_payload: dict[str, Any] | None = None @dataclass class PlanStep: @@ -116,7 +119,8 @@ def __init__( parameters: dict[str, Any] | None = None, ) -> None: self.fn = fn - self.name = name or fn.__name__ + # Extract name directly or fall back to function attribute + self.name = name or getattr(fn, "name", None) or getattr(fn, "__name__", None) or "unnamed_tool" self.description = description or (fn.__doc__ or f"Execute {self.name}").strip() self.parameters = ( parameters if parameters is not None else _generate_json_schema(fn) @@ -132,6 +136,8 @@ def to_chat_tool(self) -> Tool: def execute(self, arguments: dict[str, Any] | str | None = None) -> str: """Execute the wrapped function and return a string observation.""" + from ..exceptions import NodeInterrupt + parsed_args: dict[str, Any] = {} if isinstance(arguments, str): if arguments.strip(): @@ -167,6 +173,10 @@ def execute(self, arguments: dict[str, Any] | str | None = None) -> str: if isinstance(result, str): return result return json.dumps(result, ensure_ascii=False) + + except NodeInterrupt: + raise + except Exception as exc: # noqa: BLE001 return f"Error executing tool {self.name!r}: {exc}" @@ -174,19 +184,63 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any: return self.fn(*args, **kwargs) +@overload +def tool(fn: Callable[..., Any], /) -> AgentTool: ... + +@overload def tool( + *, name: str | None = None, description: str | None = None, parameters: dict[str, Any] | None = None, -) -> Callable[[Callable[..., Any]], AgentTool]: - """Decorator to convert a standard Python function into an AgentTool.""" +) -> Callable[[Callable[..., Any]], AgentTool]: ... + +def tool( + fn_or_name: Callable[..., Any] | str | None = None, + description: str | None = None, + parameters: dict[str, Any] | None = None, +) -> AgentTool | Callable[[Callable[..., Any]], AgentTool]: + """Decorator to convert a standard Python function into an AgentTool. + + Supports both `@tool` and `@tool(name="...", description="...")`. + """ + # Case 1: Used as `@tool` without parentheses + if callable(fn_or_name): + fn = fn_or_name + tool_name = getattr(fn, "__name__", "unnamed_tool") + return AgentTool( + fn=fn, + name=tool_name, + description=description, + parameters=parameters, + ) + + # Case 2: Used as `@tool(...)` with optional keyword arguments + custom_name = fn_or_name # In this branch, fn_or_name is a string or None def decorator(fn: Callable[..., Any]) -> AgentTool: + if isinstance(fn, AgentTool): + if custom_name: + fn.name = custom_name + if description: + fn.description = description + if parameters: + fn.parameters = parameters + return fn + + tool_name = custom_name or getattr(fn, "__name__", "unnamed_tool") return AgentTool( fn=fn, - name=name, + name=tool_name, description=description, parameters=parameters, ) return decorator + + +def interrupt(thread_id: str, node_id: str, payload: dict[str, Any] | None = None) -> None: + """Pause current run and wait for human input/approval.""" + from ..exceptions import NodeInterrupt + + raise NodeInterrupt(thread_id=thread_id, node_id=node_id, payload=payload) diff --git a/src/dynavec/agents/react.py b/src/dynavec/agents/react.py index 67b83ee..346afdf 100644 --- a/src/dynavec/agents/react.py +++ b/src/dynavec/agents/react.py @@ -9,6 +9,9 @@ from ..chat.base import ChatModel, Message, Tool from .base import AgentResult, AgentStep, AgentTool +from ..exceptions import NodeInterrupt +import json + DEFAULT_REACT_SYSTEM_PROMPT = """You are a helpful and precise reasoning agent. You solve tasks step-by-step using a ReAct (Reason + Act) approach. When you need to look up information or perform calculations, invoke the appropriate tools. @@ -36,21 +39,54 @@ def __init__( if tools: for t in tools: - agent_tool = t if isinstance(t, AgentTool) else AgentTool(t) + if isinstance(t, AgentTool): + agent_tool = t + elif callable(t): + agent_tool = AgentTool(t) + else: + continue self._tools_map[agent_tool.name] = agent_tool self._chat_tools.append(agent_tool.to_chat_tool()) - def run(self, goal: str, **kwargs: Any) -> AgentResult: - """Execute the ReAct loop synchronously until goal completion or max_steps.""" + def run(self, goal: str, thread_id: str = "default", **kwargs: Any) -> AgentResult: + """Execute the ReAct loop synchronously until goal completion, max_steps, or interrupt.""" messages: list[Message] = [] if self.system_prompt: messages.append(Message(role="system", content=self.system_prompt)) messages.append(Message(role="user", content=goal)) + return self._execute_loop(messages=messages, thread_id=thread_id, start_step=1, **kwargs) + + def resume( + self, + messages: list[Message], + human_decision: Any, + thread_id: str = "default", + start_step: int = 1, + **kwargs: Any, + ) -> AgentResult: + """Resume an interrupted run synchronously by injecting the human decision into context.""" + decision_str = ( + json.dumps(human_decision) + if isinstance(human_decision, dict) + else str(human_decision) + ) + messages.append(Message(role="user", content=f"Human Decision / Input: {decision_str}")) + return self._execute_loop( + messages=messages, thread_id=thread_id, start_step=start_step, **kwargs + ) + + def _execute_loop( + self, + messages: list[Message], + thread_id: str, + start_step: int, + **kwargs: Any, + ) -> AgentResult: steps: list[AgentStep] = [] total_tool_calls = 0 - for step_idx in range(1, self.max_steps + 1): + for step_idx in range(start_step, self.max_steps + 1): chat_res = self.model.invoke( messages, tools=self._chat_tools if self._chat_tools else None, @@ -60,7 +96,6 @@ def run(self, goal: str, **kwargs: Any) -> AgentResult: messages.append(msg) if not msg.tool_calls: - # Final response reached without further tool calls step = AgentStep( step_number=step_idx, thought=msg.content, @@ -77,15 +112,32 @@ def run(self, goal: str, **kwargs: Any) -> AgentResult: tool_calls_count=total_tool_calls, ) - # Execute tool calls step_observations: list[str] = [] for tc in msg.tool_calls: total_tool_calls += 1 tool_instance = self._tools_map.get(tc.name) - if tool_instance is not None: - obs = tool_instance.execute(tc.arguments) - else: - obs = f"Error: Tool {tc.name!r} is not registered in available tools." + + try: + if tool_instance is not None: + obs = tool_instance.execute(tc.arguments) + else: + obs = f"Error: Tool {tc.name!r} is not registered in available tools." + except NodeInterrupt as exc: + return AgentResult( + output=f"Execution interrupted at node {exc.node_id!r}", + steps=steps, + finished=False, + termination_reason="interrupted", + total_steps=step_idx, + tool_calls_count=total_tool_calls, + interrupt_payload={ + "thread_id": exc.thread_id, + "node_id": exc.node_id, + "payload": exc.payload, + "messages": messages, + "start_step": step_idx, + }, + ) step_observations.append(obs) messages.append( @@ -104,7 +156,6 @@ def run(self, goal: str, **kwargs: Any) -> AgentResult: ) steps.append(step) - # Reached max steps without completing last_output = steps[-1].thought or ( steps[-1].observations[-1] if steps[-1].observations else "" ) @@ -117,17 +168,46 @@ def run(self, goal: str, **kwargs: Any) -> AgentResult: tool_calls_count=total_tool_calls, ) - async def arun(self, goal: str, **kwargs: Any) -> AgentResult: - """Execute the ReAct loop asynchronously.""" + + async def arun(self, goal: str, thread_id: str = "default", **kwargs: Any) -> AgentResult: + """Execute the ReAct loop asynchronously until goal completion, max_steps, or interrupt.""" messages: list[Message] = [] if self.system_prompt: messages.append(Message(role="system", content=self.system_prompt)) messages.append(Message(role="user", content=goal)) + return await self._aexecute_loop(messages=messages, thread_id=thread_id, start_step=1, **kwargs) + + async def aresume( + self, + messages: list[Message], + human_decision: Any, + thread_id: str = "default", + start_step: int = 1, + **kwargs: Any, + ) -> AgentResult: + """Resume an interrupted run asynchronously by injecting human decision into context.""" + decision_str = ( + json.dumps(human_decision) + if isinstance(human_decision, dict) + else str(human_decision) + ) + messages.append(Message(role="user", content=f"Human Decision / Input: {decision_str}")) + return await self._aexecute_loop( + messages=messages, thread_id=thread_id, start_step=start_step, **kwargs + ) + + async def _aexecute_loop( + self, + messages: list[Message], + thread_id: str, + start_step: int, + **kwargs: Any, + ) -> AgentResult: steps: list[AgentStep] = [] total_tool_calls = 0 - for step_idx in range(1, self.max_steps + 1): + for step_idx in range(start_step, self.max_steps + 1): chat_res = await self.model.ainvoke( messages, tools=self._chat_tools if self._chat_tools else None, @@ -157,10 +237,28 @@ async def arun(self, goal: str, **kwargs: Any) -> AgentResult: for tc in msg.tool_calls: total_tool_calls += 1 tool_instance = self._tools_map.get(tc.name) - if tool_instance is not None: - obs = await asyncio.to_thread(tool_instance.execute, tc.arguments) - else: - obs = f"Error: Tool {tc.name!r} is not registered in available tools." + + try: + if tool_instance is not None: + obs = await asyncio.to_thread(tool_instance.execute, tc.arguments) + else: + obs = f"Error: Tool {tc.name!r} is not registered in available tools." + except NodeInterrupt as exc: + return AgentResult( + output=f"Execution interrupted at node {exc.node_id!r}", + steps=steps, + finished=False, + termination_reason="interrupted", + total_steps=step_idx, + tool_calls_count=total_tool_calls, + interrupt_payload={ + "thread_id": exc.thread_id, + "node_id": exc.node_id, + "payload": exc.payload, + "messages": messages, + "start_step": step_idx, + }, + ) step_observations.append(obs) messages.append( diff --git a/src/dynavec/exceptions.py b/src/dynavec/exceptions.py index 6db7f6a..869cec7 100644 --- a/src/dynavec/exceptions.py +++ b/src/dynavec/exceptions.py @@ -68,3 +68,31 @@ def __init__(self, feature: str, package: str, extra: str) -> None: self.feature = feature self.package = package self.extra = extra + + + + +class NodeInterrupt(DynavecError): + """Raised when an agent execution node pauses waiting for human approval or input.""" + + def __init__(self, thread_id: str, node_id: str, payload: dict | None = None) -> None: + super().__init__( + f"Execution interrupted at node {node_id!r} for thread {thread_id!r}. " + "Awaiting human input." + ) + self.thread_id = thread_id + self.node_id = node_id + self.payload = payload or {} + + +class ApprovalTimeoutError(DynavecError): + """Raised when a pending human approval expires before a decision is submitted.""" + + def __init__(self, thread_id: str, node_id: str, timeout_seconds: int) -> None: + super().__init__( + f"Approval request at node {node_id!r} for thread {thread_id!r} " + f"timed out after {timeout_seconds} seconds." + ) + self.thread_id = thread_id + self.node_id = node_id + self.timeout_seconds = timeout_seconds diff --git a/tests/test_agents.py b/tests/test_agents.py index 4ee91cb..c560bcc 100644 --- a/tests/test_agents.py +++ b/tests/test_agents.py @@ -74,26 +74,26 @@ def stream( def test_tool_decorator_and_schema_generation() -> None: """Verify @tool extracts docstrings, names, parameter types, and executes correctly.""" - @tool(name="custom_add", description="Add two integers.") - def add(a: int, b: int = 0) -> int: - """Add two numbers.""" + @tool + def custom_add(a: int, b: int = 0) -> int: + """Add two integers.""" return a + b - assert isinstance(add, AgentTool) - assert add.name == "custom_add" - assert add.description == "Add two integers." - assert add.parameters["type"] == "object" - assert add.parameters["properties"]["a"]["type"] == "integer" - assert add.parameters["properties"]["b"]["type"] == "integer" - assert "a" in add.parameters.get("required", []) - assert "b" not in add.parameters.get("required", []) + assert isinstance(custom_add, AgentTool) + assert custom_add.name == "custom_add" + assert custom_add.description == "Add two integers." + assert custom_add.parameters["type"] == "object" + assert custom_add.parameters["properties"]["a"]["type"] == "integer" + assert custom_add.parameters["properties"]["b"]["type"] == "integer" + assert "a" in custom_add.parameters.get("required", []) + assert "b" not in custom_add.parameters.get("required", []) # Test execution with dict - assert add.execute({"a": 5, "b": 10}) == "15" + assert custom_add.execute({"a": 5, "b": 10}) == "15" # Test execution with JSON string - assert add.execute('{"a": 20, "b": 22}') == "42" + assert custom_add.execute('{"a": 20, "b": 22}') == "42" # Test direct call - assert add(3, 4) == 7 + assert custom_add(3, 4) == 7 def test_react_agent_two_tool_task() -> None: diff --git a/tests/test_hitl_interrupts.py b/tests/test_hitl_interrupts.py new file mode 100644 index 0000000..fc6e432 --- /dev/null +++ b/tests/test_hitl_interrupts.py @@ -0,0 +1,172 @@ +"""Tests for Human-in-the-loop / interrupts functionality in dynavec agents.""" + +import asyncio +from collections.abc import AsyncIterator, Iterator +import pytest + +from dynavec.agents.base import interrupt, tool +from dynavec.agents.react import ReActAgent +from dynavec.chat.base import ChatModel, ChatResult, Message, ToolCall + + +class MockChatModel(ChatModel): + """Mock model that returns predefined tool calls or responses sequentially.""" + + def __init__(self, responses: list[ChatResult]) -> None: + super().__init__() + self.model_name = "mock-model" + self.responses = responses + self.call_count = 0 + + def invoke(self, messages: list[Message], tools=None, **kwargs) -> ChatResult: + res = self.responses[self.call_count] + self.call_count += 1 + return res + + async def ainvoke(self, messages: list[Message], tools=None, **kwargs) -> ChatResult: + return self.invoke(messages, tools=tools, **kwargs) + + def stream(self, messages: list[Message], tools=None, **kwargs) -> Iterator[ChatResult]: + res = self.invoke(messages, tools=tools, **kwargs) + yield res + + async def astream(self, messages: list[Message], tools=None, **kwargs) -> AsyncIterator[ChatResult]: + res = await self.ainvoke(messages, tools=tools, **kwargs) + yield res + + +@tool +def sensitive_action_tool(action: str) -> str: + """A tool that requires human approval before proceeding.""" + interrupt( + thread_id="thread-123", + node_id="approval_node", + payload={"action": action, "requires_approval": True}, + ) + return f"Executed {action}" + + +def test_interrupt_raised_and_caught(): + """Test that NodeInterrupt pauses execution and populates interrupt_payload.""" + mock_responses = [ + ChatResult( + message=Message( + role="assistant", + content="Requesting sensitive action.", + tool_calls=[ + ToolCall( + id="tc_1", + name="sensitive_action_tool", + arguments={"action": "delete_database"}, + ) + ], + ) + ) + ] + model = MockChatModel(mock_responses) + agent = ReActAgent(model=model, tools=[sensitive_action_tool]) + + result = agent.run("Delete the production database", thread_id="thread-123") + + assert result.finished is False + assert result.termination_reason == "interrupted" + assert result.interrupt_payload is not None + assert result.interrupt_payload["thread_id"] == "thread-123" + assert result.interrupt_payload["node_id"] == "approval_node" + assert result.interrupt_payload["payload"]["action"] == "delete_database" + + +def test_resume_execution(): + """Test resuming an interrupted agent with human decision.""" + initial_responses = [ + ChatResult( + message=Message( + role="assistant", + content="Requesting sensitive action.", + tool_calls=[ + ToolCall( + id="tc_1", + name="sensitive_action_tool", + arguments={"action": "delete_database"}, + ) + ], + ) + ) + ] + model = MockChatModel(initial_responses) + agent = ReActAgent(model=model, tools=[sensitive_action_tool]) + + result = agent.run("Delete database", thread_id="thread-123") + assert result.termination_reason == "interrupted" + + # Simulate resume after human approves + resume_responses = [ + ChatResult( + message=Message( + role="assistant", + content="Action was approved by human. Proceeding completed successfully.", + ) + ) + ] + resume_model = MockChatModel(resume_responses) + resume_agent = ReActAgent(model=resume_model, tools=[sensitive_action_tool]) + + saved_messages = result.interrupt_payload["messages"] + start_step = result.interrupt_payload["start_step"] + + final_result = resume_agent.resume( + messages=saved_messages, + human_decision={"approved": True}, + thread_id="thread-123", + start_step=start_step, + ) + + assert final_result.finished is True + assert final_result.termination_reason == "completed" + + +def test_async_interrupt_and_resume(): + """Test async arun and aresume functionality using asyncio.run.""" + async def _run_async(): + initial_responses = [ + ChatResult( + message=Message( + role="assistant", + content="Requesting sensitive action.", + tool_calls=[ + ToolCall( + id="tc_1", + name="sensitive_action_tool", + arguments={"action": "deploy_code"}, + ) + ], + ) + ) + ] + model = MockChatModel(initial_responses) + agent = ReActAgent(model=model, tools=[sensitive_action_tool]) + + result = await agent.arun("Deploy code to production", thread_id="thread-456") + assert result.termination_reason == "interrupted" + + resume_responses = [ + ChatResult( + message=Message( + role="assistant", + content="Deployment finished after human confirmation.", + ) + ) + ] + resume_agent = ReActAgent(model=MockChatModel(resume_responses), tools=[sensitive_action_tool]) + + final_result = await resume_agent.aresume( + messages=result.interrupt_payload["messages"], + human_decision="Approved by DevOps lead", + thread_id="thread-456", + start_step=result.interrupt_payload["start_step"], + ) + + assert final_result.finished is True + assert final_result.termination_reason == "completed" + + asyncio.run(_run_async()) \ No newline at end of file From c74380f72bd1961699a56d69e7d1c4ed60d74c26 Mon Sep 17 00:00:00 2001 From: Uzmaa7 Date: Mon, 5 Oct 2026 15:45:35 +0530 Subject: [PATCH 2/4] style: fix ruff linting and import formatting --- README.md | 44 +- benchmarks/benchmark_opq.py | 37 +- benchmarks/benchmark_quantization.py | 45 +- benchmarks/cost_model.py | 21 +- benchmarks/datasets.py | 6 +- benchmarks/report.py | 137 ++++-- benchmarks/run_benchmark.py | 28 +- docs/DEPLOY.md | 7 +- examples/crewai_tool.py | 11 +- examples/dashboard_demo.py | 30 +- examples/dspy_retriever.py | 4 +- examples/haystack_integration.py | 4 +- examples/hot_tier.py | 21 +- examples/ingest_office_docs.py | 3 + examples/langchain_rag.py | 4 +- examples/openai_1536_retrieval.py | 14 +- examples/pdf_ingestion.py | 8 +- examples/pinecone_vs_dynavec.py | 18 +- examples/quickstart.py | 25 +- examples/verify_provisioning.py | 8 +- notebooks/rrf_weight_fitting.ipynb | 94 +++- src/dynavec/agents/base.py | 28 +- src/dynavec/agents/react.py | 18 +- src/dynavec/async_client.py | 51 +-- src/dynavec/bm25.py | 42 +- src/dynavec/chat/anthropic.py | 83 ++-- src/dynavec/chat/base.py | 34 +- src/dynavec/chat/bedrock.py | 85 ++-- src/dynavec/chat/openai.py | 33 +- src/dynavec/cli.py | 524 +++++++++++----------- src/dynavec/client.py | 1 - src/dynavec/client_common.py | 48 +-- src/dynavec/config.py | 9 +- src/dynavec/embeddings/cache.py | 4 +- src/dynavec/embeddings/cached.py | 5 +- src/dynavec/embeddings/huggingface.py | 16 +- src/dynavec/embeddings/openai.py | 11 +- src/dynavec/eval/chart.py | 2 - src/dynavec/eval/trend.py | 8 +- src/dynavec/exceptions.py | 2 - src/dynavec/fusion.py | 11 +- src/dynavec/graph.py | 8 +- src/dynavec/integrations/dspy.py | 12 +- src/dynavec/integrations/haystack.py | 16 +- src/dynavec/integrations/langchain.py | 21 +- src/dynavec/integrations/llamaindex.py | 3 +- src/dynavec/mcp/__init__.py | 1 - src/dynavec/mcp/server.py | 8 +- src/dynavec/models.py | 4 +- src/dynavec/namespace.py | 4 +- src/dynavec/quantization.py | 48 +-- src/dynavec/retrievers.py | 8 +- src/dynavec/stores/async_dynamodb.py | 22 +- src/dynavec/stores/async_s3vectors.py | 45 +- src/dynavec/stores/s3vectors.py | 13 +- src/dynavec/utils.py | 2 +- tests/integration/test_live_aws.py | 6 +- tests/test_agents.py | 4 +- tests/test_async_client.py | 63 +-- tests/test_async_dynamodb.py | 23 +- tests/test_async_s3vectors.py | 51 +-- tests/test_bedrock_multimodal.py | 1 - tests/test_bm25_retriever.py | 25 +- tests/test_cache.py | 9 +- tests/test_chat_models.py | 4 +- tests/test_client_inmemory.py | 160 ++++--- tests/test_crewai_tool.py | 2 +- tests/test_describe.py | 2 +- tests/test_dynamodb.py | 4 +- tests/test_embedding_cache.py | 58 ++- tests/test_eval_trend.py | 11 +- tests/test_export_import.py | 13 +- tests/test_fusion.py | 9 + tests/test_graph_export.py | 39 +- tests/test_graph_properties.py | 1 - tests/test_haystack.py | 110 ++--- tests/test_hitl_interrupts.py | 12 +- tests/test_hot_tier.py | 37 +- tests/test_huggingface_embedder.py | 3 +- tests/test_list_vectors.py | 8 +- tests/test_markdown_source.py | 2 +- tests/test_mcp_server.py | 10 +- tests/test_metadata.py | 14 +- tests/test_metric_properties.py | 8 +- tests/test_namespace.py | 22 +- tests/test_openai_embeddings.py | 8 +- tests/test_provisioning.py | 16 +- tests/test_quantization.py | 1 + tests/test_retrievers.py | 4 +- tests/test_s3vectors.py | 112 +++-- tests/test_telemetry.py | 2 +- tests/test_utils_transforms.py | 1 + tools/build_docs.py | 574 +++++++++++++++++-------- typings/dspy/retrievers/retrieve.pyi | 4 +- 94 files changed, 1739 insertions(+), 1493 deletions(-) diff --git a/README.md b/README.md index 474d0b4..5b39e57 100644 --- a/README.md +++ b/README.md @@ -131,35 +131,41 @@ Replace `REGION` and `ACCOUNT_ID`. `dynamodb:Scan` is only needed for the GraphR ```python from dynavec import Dynavec, DynavecConfig, Document -from dynavec.embeddings import OpenAIEmbedder # or Gemini / Bedrock / SentenceTransformer +from dynavec.embeddings import OpenAIEmbedder # or Gemini / Bedrock / SentenceTransformer cfg = DynavecConfig( - vector_bucket="my-vectors", # S3 vector bucket - index="docs", # vector index - table="dynavec_docs", # DynamoDB table + vector_bucket="my-vectors", # S3 vector bucket + index="docs", # vector index + table="dynavec_docs", # DynamoDB table dimension=1536, distance_metric="cosine", region="us-east-1", - auto_provision=True, # create bucket + index + table if missing + auto_provision=True, # create bucket + index + table if missing ) db = Dynavec(cfg, embedder=OpenAIEmbedder(model="text-embedding-3-small")) db.upsert( [ - Document(id="a", text="The mitochondria is the powerhouse of the cell.", - metadata={"topic": "biology", "year": 2021}), - Document(id="b", text="Rockets reach orbit at roughly 28,000 km/h.", - metadata={"topic": "space", "year": 2023}), + Document( + id="a", + text="The mitochondria is the powerhouse of the cell.", + metadata={"topic": "biology", "year": 2021}, + ), + Document( + id="b", + text="Rockets reach orbit at roughly 28,000 km/h.", + metadata={"topic": "space", "year": 2023}, + ), ], - auto_metadata=True, # also attach hash/timestamp/word counts + auto_metadata=True, # also attach hash/timestamp/word counts ) hits = db.search( "how do cells make energy?", top_k=3, - filter={"topic": "biology"}, # S3 Vectors metadata pre-filter - rerank="mmr", # diversity-aware reranking + filter={"topic": "biology"}, # S3 Vectors metadata pre-filter + rerank="mmr", # diversity-aware reranking ) for h in hits: print(h.score, h.id, h.text) @@ -190,10 +196,12 @@ Expire ephemeral documents, session memories, or cache entries automatically via ```python # Pass ttl_seconds on Document or upsert() -db.upsert([ - Document(id="session-1", text="Temporary session context", ttl_seconds=3600), - Document(id="cached-doc", text="Ephemeral cache data", ttl_seconds=86400), -]) +db.upsert( + [ + Document(id="session-1", text="Temporary session context", ttl_seconds=3600), + Document(id="cached-doc", text="Ephemeral cache data", ttl_seconds=86400), + ] +) # Or set a default TTL across an entire batch ns = db.namespace("sessions") @@ -259,11 +267,13 @@ from dynavec.integrations.tools import make_retriever_fn retrieve = make_retriever_fn(db, namespace="kb", top_k=2) + @tool def search_knowledge_base(query: str) -> str: """Search the dynavec knowledge base for relevant passages.""" return retrieve(query) + agent = Agent(tools=[search_knowledge_base]) agent("Where is dynavec vector data stored?") ``` @@ -426,7 +436,7 @@ from dynavec.dashboard import serve rec = TelemetryRecorder() db = Dynavec(cfg, embedder=emb, cache=SemanticCache(), telemetry=rec) # ... your app runs searches; the recorder fills automatically ... -serve(rec, port=8779) # JSON API at http://127.0.0.1:8779 +serve(rec, port=8779) # JSON API at http://127.0.0.1:8779 ``` **2. Run the dashboard** (points at that API; falls back to sample data if unset): diff --git a/benchmarks/benchmark_opq.py b/benchmarks/benchmark_opq.py index 573da69..8617e9f 100644 --- a/benchmarks/benchmark_opq.py +++ b/benchmarks/benchmark_opq.py @@ -1,4 +1,3 @@ - """Benchmark PQ vs OPQ + PQ.""" from __future__ import annotations @@ -63,9 +62,7 @@ def recall_at_k( truth_set = set(truth[:k]) prediction_set = set(prediction[:k]) - recalls.append( - len(truth_set & prediction_set) / k - ) + recalls.append(len(truth_set & prediction_set) / k) return float(np.mean(recalls)) @@ -115,9 +112,7 @@ def benchmark_quantizer( k=10, ) - reconstruction_error = quantizer.reconstruction_error( - database - ) + reconstruction_error = quantizer.reconstruction_error(database) print(f"Fit time : {fit_time:.4f}s") print(f"Encode time : {encode_time:.4f}s") @@ -147,13 +142,9 @@ def main() -> None: m = 8 nbits = 8 - database = rng.normal( - size=(num_database, dimension) - ).astype(np.float32) + database = rng.normal(size=(num_database, dimension)).astype(np.float32) - queries = rng.normal( - size=(num_queries, dimension) - ).astype(np.float32) + queries = rng.normal(size=(num_queries, dimension)).astype(np.float32) print("=== PQ vs OPQ + PQ Benchmark ===") print(f"Database : {num_database} vectors") @@ -223,25 +214,13 @@ def main() -> None: print("\n=== Comparison ===") - recall_improvement = ( - opq_results["recall_10"] - - pq_results["recall_10"] - ) + recall_improvement = opq_results["recall_10"] - pq_results["recall_10"] - error_improvement = ( - pq_results["reconstruction_error"] - - opq_results["reconstruction_error"] - ) + error_improvement = pq_results["reconstruction_error"] - opq_results["reconstruction_error"] - print( - f"Recall@10 improvement : " - f"{recall_improvement:+.4f}" - ) + print(f"Recall@10 improvement : {recall_improvement:+.4f}") - print( - f"Reconstruction improvement: " - f"{error_improvement:+.6f}" - ) + print(f"Reconstruction improvement: {error_improvement:+.6f}") if __name__ == "__main__": diff --git a/benchmarks/benchmark_quantization.py b/benchmarks/benchmark_quantization.py index f3340ff..684d9cd 100644 --- a/benchmarks/benchmark_quantization.py +++ b/benchmarks/benchmark_quantization.py @@ -14,9 +14,7 @@ def generate_dataset(n_vectors=10_000, dim=32, seed=0): centers = rng.normal(size=(8, dim)).astype(np.float32) assign = rng.integers(0, 8, size=n_vectors) - x = centers[assign] + 0.05 * rng.normal( - size=(n_vectors, dim) - ).astype(np.float32) + x = centers[assign] + 0.05 * rng.normal(size=(n_vectors, dim)).astype(np.float32) return x.astype(np.float32) @@ -94,15 +92,11 @@ def benchmark(): scalar_codes, ) - scalar_reconstruction_error = scalar.reconstruction_error( - vectors - ) + scalar_reconstruction_error = scalar.reconstruction_error(vectors) scalar_bytes_per_vector = scalar.code_size_bytes scalar_total_bytes = scalar_codes.nbytes - scalar_compression = ( - raw_bytes_per_vector / scalar_bytes_per_vector - ) + scalar_compression = raw_bytes_per_vector / scalar_bytes_per_vector # --------------------------------------------------------- # Product Quantizer @@ -130,15 +124,11 @@ def benchmark(): pq_codes, ) - pq_reconstruction_error = pq.reconstruction_error( - vectors - ) + pq_reconstruction_error = pq.reconstruction_error(vectors) pq_bytes_per_vector = pq.code_size_bytes pq_total_bytes = pq_codes.nbytes - pq_compression = ( - raw_bytes_per_vector / pq_bytes_per_vector - ) + pq_compression = raw_bytes_per_vector / pq_bytes_per_vector # --------------------------------------------------------- # Results @@ -148,10 +138,7 @@ def benchmark(): print(f" Fit time: {scalar_fit_time:.4f} sec") print(f" Encode time: {scalar_encode_time:.4f} sec") print(f" Decode time: {scalar_decode_time:.4f} sec") - print( - f" Reconstruction error: " - f"{scalar_reconstruction_error:.8f}" - ) + print(f" Reconstruction error: {scalar_reconstruction_error:.8f}") print(f" Bytes/vector: {scalar_bytes_per_vector}") print(f" Total encoded size: {scalar_total_bytes / 1024:.2f} KB") print(f" Compression: {scalar_compression:.2f}x") @@ -161,32 +148,20 @@ def benchmark(): print(f" Fit time: {pq_fit_time:.4f} sec") print(f" Encode time: {pq_encode_time:.4f} sec") print(f" Decode time: {pq_decode_time:.4f} sec") - print( - f" Reconstruction error: " - f"{pq_reconstruction_error:.8f}" - ) + print(f" Reconstruction error: {pq_reconstruction_error:.8f}") print(f" Bytes/vector: {pq_bytes_per_vector}") print(f" Total encoded size: {pq_total_bytes / 1024:.2f} KB") print(f" Compression: {pq_compression:.2f}x") print() print("Comparison") - print( - f" Scalar/PQ compression: " - f"{scalar_compression / pq_compression:.2f}x" - ) + print(f" Scalar/PQ compression: {scalar_compression / pq_compression:.2f}x") if scalar_encode_time > 0: - print( - f" PQ encode / Scalar encode: " - f"{pq_encode_time / scalar_encode_time:.2f}x" - ) + print(f" PQ encode / Scalar encode: {pq_encode_time / scalar_encode_time:.2f}x") if scalar_decode_time > 0: - print( - f" PQ decode / Scalar decode: " - f"{pq_decode_time / scalar_decode_time:.2f}x" - ) + print(f" PQ decode / Scalar decode: {pq_decode_time / scalar_decode_time:.2f}x") print() print("=" * 80) diff --git a/benchmarks/cost_model.py b/benchmarks/cost_model.py index c839c01..ff758dd 100644 --- a/benchmarks/cost_model.py +++ b/benchmarks/cost_model.py @@ -19,12 +19,12 @@ @dataclass class Workload: - vectors: int # number of stored vectors - dim: int # embedding dimension + vectors: int # number of stored vectors + dim: int # embedding dimension queries_per_month: int writes_per_month: int = 0 - metadata_bytes: int = 500 # avg full metadata + text per doc (DynamoDB) - filterable_bytes: int = 120 # avg filterable metadata per vector (S3 Vectors) + metadata_bytes: int = 500 # avg full metadata + text per doc (DynamoDB) + filterable_bytes: int = 120 # avg filterable metadata per vector (S3 Vectors) @property def vector_bytes(self) -> int: @@ -147,12 +147,15 @@ def main() -> None: p.add_argument("--wpm", type=int, default=100_000, help="writes per month") args = p.parse_args() - w = Workload(vectors=args.vectors, dim=args.dim, - queries_per_month=args.qpm, writes_per_month=args.wpm) + w = Workload( + vectors=args.vectors, dim=args.dim, queries_per_month=args.qpm, writes_per_month=args.wpm + ) - print(f"\nWorkload: {w.vectors:,} vectors × {w.dim}d " - f"({w.raw_vector_gb:.1f} GB raw), {w.queries_per_month:,} queries/mo, " - f"{w.writes_per_month:,} writes/mo\n") + print( + f"\nWorkload: {w.vectors:,} vectors × {w.dim}d " + f"({w.raw_vector_gb:.1f} GB raw), {w.queries_per_month:,} queries/mo, " + f"{w.writes_per_month:,} writes/mo\n" + ) print("dynavec breakdown:") for k, v in dynavec_cost(w).items(): print(f" {k:<22} ${v:>10,.2f}") diff --git a/benchmarks/datasets.py b/benchmarks/datasets.py index 9ab3386..c88384c 100644 --- a/benchmarks/datasets.py +++ b/benchmarks/datasets.py @@ -9,8 +9,8 @@ @dataclass class Dataset: - vectors: np.ndarray # (N, dim) float32 - queries: np.ndarray # (Q, dim) float32 + vectors: np.ndarray # (N, dim) float32 + queries: np.ndarray # (Q, dim) float32 ids: list[str] ground_truth: np.ndarray # (Q, k) int indices into vectors (exact NN) @@ -55,7 +55,7 @@ def _l2_normalize(x: np.ndarray) -> np.ndarray: def _exact_topk(vectors: np.ndarray, queries: np.ndarray, k: int) -> np.ndarray: # cosine similarity == dot product on L2-normalized vectors - sims = queries @ vectors.T # (Q, N) + sims = queries @ vectors.T # (Q, N) return np.argpartition(-sims, kth=k - 1, axis=1)[:, :k] diff --git a/benchmarks/report.py b/benchmarks/report.py index 7c4f068..3201563 100644 --- a/benchmarks/report.py +++ b/benchmarks/report.py @@ -35,12 +35,12 @@ # Representative quality/latency profile (EDIT with measured values). # recall@10, warm p50/p95/p99 ms for a ~1M x 768 workload. _PROFILE = { - "dynavec": {"recall": 0.90, "p50": 45, "p95": 120, "p99": 180}, - "pinecone_serverless": {"recall": 0.95, "p50": 30, "p95": 70, "p99": 110}, - "opensearch_serverless": {"recall": 0.97, "p50": 15, "p95": 40, "p99": 70}, - "qdrant_cloud": {"recall": 0.98, "p50": 8, "p95": 20, "p99": 35}, - "weaviate_cloud": {"recall": 0.97, "p50": 10, "p95": 25, "p99": 45}, - "milvus_zilliz": {"recall": 0.98, "p50": 7, "p95": 18, "p99": 30}, + "dynavec": {"recall": 0.90, "p50": 45, "p95": 120, "p99": 180}, + "pinecone_serverless": {"recall": 0.95, "p50": 30, "p95": 70, "p99": 110}, + "opensearch_serverless": {"recall": 0.97, "p50": 15, "p95": 40, "p99": 70}, + "qdrant_cloud": {"recall": 0.98, "p50": 8, "p95": 20, "p99": 35}, + "weaviate_cloud": {"recall": 0.97, "p50": 10, "p95": 25, "p99": 45}, + "milvus_zilliz": {"recall": 0.98, "p50": 7, "p95": 18, "p99": 30}, } _LABELS = { @@ -73,16 +73,32 @@ def markdown_table(inp: ReportInputs) -> str: ("Latency p50 (ms)", {p: _PROFILE[p]["p50"] for p in products}, "min", "{:.0f}"), ("Latency p95 (ms)", {p: _PROFILE[p]["p95"] for p in products}, "min", "{:.0f}"), ("Cost ($/mo)", {p: costs[p] for p in products}, "min", "${:,.0f}"), - ("Serverless (scale-to-zero)", { - "dynavec": "Yes", "pinecone_serverless": "Yes", - "opensearch_serverless": "No (OCU floor)", "qdrant_cloud": "No (nodes)", - "weaviate_cloud": "No (nodes)", "milvus_zilliz": "No (CU)", - }, None, "{}"), - ("Data in your AWS account", { - "dynavec": "Yes", "pinecone_serverless": "No", - "opensearch_serverless": "Yes", "qdrant_cloud": "Self-host only", - "weaviate_cloud": "Self-host only", "milvus_zilliz": "Self-host only", - }, None, "{}"), + ( + "Serverless (scale-to-zero)", + { + "dynavec": "Yes", + "pinecone_serverless": "Yes", + "opensearch_serverless": "No (OCU floor)", + "qdrant_cloud": "No (nodes)", + "weaviate_cloud": "No (nodes)", + "milvus_zilliz": "No (CU)", + }, + None, + "{}", + ), + ( + "Data in your AWS account", + { + "dynavec": "Yes", + "pinecone_serverless": "No", + "opensearch_serverless": "Yes", + "qdrant_cloud": "Self-host only", + "weaviate_cloud": "Self-host only", + "milvus_zilliz": "Self-host only", + }, + None, + "{}", + ), ] header = "| Metric | " + " | ".join(_LABELS[p] for p in products) + " |" @@ -148,21 +164,30 @@ def cost_scale_chart(inp: ReportInputs, out_path: str) -> str: fig, ax = plt.subplots(figsize=(11, 6)) for p in products: ax.plot( - scales, series[p], marker="o", - color=_COLORS[p], linewidth=2.5 if p == "dynavec" else 1.6, + scales, + series[p], + marker="o", + color=_COLORS[p], + linewidth=2.5 if p == "dynavec" else 1.6, label=_LABELS[p], ) ax.set_xscale("log") ax.set_yscale("log") ax.set_xlabel("Number of vectors (log scale)") ax.set_ylabel("Estimated cost ($/month, log scale)") - ax.set_title(f"Monthly cost by scale — {inp.dim}d, {inp.qpm:,} queries/mo", - fontsize=14, weight="bold") + ax.set_title( + f"Monthly cost by scale — {inp.dim}d, {inp.qpm:,} queries/mo", fontsize=14, weight="bold" + ) ax.grid(True, which="both", alpha=0.25) ax.legend() - fig.text(0.5, -0.02, - "APPROX public list prices (see cost_model.py). Order-of-magnitude only.", - ha="center", fontsize=8, style="italic") + fig.text( + 0.5, + -0.02, + "APPROX public list prices (see cost_model.py). Order-of-magnitude only.", + ha="center", + fontsize=8, + style="italic", + ) fig.tight_layout() fig.savefig(out_path, dpi=140, bbox_inches="tight") plt.close(fig) @@ -196,9 +221,14 @@ def cost_matrix_chart(out_path: str, qpm: int = 1_000_000, wpm: int = 100_000) - for ax, dim in zip(axes, DIMENSIONS): for p in products: ys = [compare(Workload(n, dim, qpm, wpm))[p] for n in SCALES] - ax.plot(SCALES, ys, marker="o", - color=_COLORS[p], linewidth=2.4 if p == "dynavec" else 1.4, - label=_LABELS[p]) + ax.plot( + SCALES, + ys, + marker="o", + color=_COLORS[p], + linewidth=2.4 if p == "dynavec" else 1.4, + label=_LABELS[p], + ) ax.set_xscale("log") ax.set_yscale("log") ax.set_title(f"{dim}-dim embeddings", weight="bold") @@ -212,8 +242,11 @@ def cost_matrix_chart(out_path: str, qpm: int = 1_000_000, wpm: int = 100_000) - handles, labels = axes[0].get_legend_handles_labels() legend_ax.legend(handles, labels, loc="center", fontsize=12, title="Product") - fig.suptitle(f"Estimated monthly cost by scale & dimension · {qpm:,} queries/mo", - fontsize=16, weight="bold") + fig.suptitle( + f"Estimated monthly cost by scale & dimension · {qpm:,} queries/mo", + fontsize=16, + weight="bold", + ) fig.supxlabel("Number of vectors (log scale)") fig.supylabel("Cost $/month (log scale)") fig.tight_layout() @@ -228,19 +261,31 @@ def cost_by_dimension_chart(out_path: str, scale: int = 100_000_000, qpm: int = fig, ax = plt.subplots(figsize=(11, 6)) for p in products: ys = [compare(Workload(scale, d, qpm, 100_000))[p] for d in DIMENSIONS] - ax.plot(DIMENSIONS, ys, marker="s", - color=_COLORS[p], linewidth=2.4 if p == "dynavec" else 1.4, - label=_LABELS[p]) + ax.plot( + DIMENSIONS, + ys, + marker="s", + color=_COLORS[p], + linewidth=2.4 if p == "dynavec" else 1.4, + label=_LABELS[p], + ) ax.set_yscale("log") ax.set_xticks(DIMENSIONS) ax.set_xlabel("Embedding dimension") ax.set_ylabel("Cost $/month (log scale)") - ax.set_title(f"Cost by embedding dimension @ {_fmt_count(scale)} vectors", - fontsize=14, weight="bold") + ax.set_title( + f"Cost by embedding dimension @ {_fmt_count(scale)} vectors", fontsize=14, weight="bold" + ) ax.grid(True, which="both", alpha=0.25) ax.legend() - fig.text(0.5, -0.02, "APPROX list prices — see cost_model.py.", ha="center", - fontsize=8, style="italic") + fig.text( + 0.5, + -0.02, + "APPROX list prices — see cost_model.py.", + ha="center", + fontsize=8, + style="italic", + ) fig.tight_layout() fig.savefig(out_path, dpi=140, bbox_inches="tight") plt.close(fig) @@ -303,8 +348,10 @@ def generate(out_dir: str = "benchmarks/out", inp: ReportInputs | None = None) - table_md = markdown_table(inp) table_path = os.path.join(out_dir, "comparison.md") with open(table_path, "w") as f: - f.write(f"# dynavec benchmark comparison\n\n" - f"_Workload: {inp.vectors:,} × {inp.dim}d, {inp.qpm:,} queries/mo._\n\n") + f.write( + f"# dynavec benchmark comparison\n\n" + f"_Workload: {inp.vectors:,} × {inp.dim}d, {inp.qpm:,} queries/mo._\n\n" + ) f.write(table_md + "\n") bars = bar_chart(inp, os.path.join(out_dir, "quality_latency.png")) cost = cost_scale_chart(inp, os.path.join(out_dir, "cost_by_scale.png")) @@ -316,15 +363,21 @@ def generate(out_dir: str = "benchmarks/out", inp: ReportInputs | None = None) - scaling_path = os.path.join(out_dir, "scaling.md") with open(scaling_path, "w") as f: f.write("# dynavec scaling: dimensions × vector count (100K → 1B)\n") - f.write("\n_Cost from the cost model (APPROX public list prices). " - "Recall/latency unaffected by this sweep._\n") + f.write( + "\n_Cost from the cost model (APPROX public list prices). " + "Recall/latency unaffected by this sweep._\n" + ) f.write(scaling_tables_md(qpm=inp.qpm)) f.write("\n") return { - "table_md": table_path, "bar_chart": bars, "cost_chart": cost, - "cost_matrix": matrix, "cost_by_dimension": by_dim, - "storage_footprint": footprint, "scaling_md": scaling_path, + "table_md": table_path, + "bar_chart": bars, + "cost_chart": cost, + "cost_matrix": matrix, + "cost_by_dimension": by_dim, + "storage_footprint": footprint, + "scaling_md": scaling_path, } diff --git a/benchmarks/run_benchmark.py b/benchmarks/run_benchmark.py index e6e05e5..5035673 100644 --- a/benchmarks/run_benchmark.py +++ b/benchmarks/run_benchmark.py @@ -76,10 +76,7 @@ def __init__(self, args): def build(self, vectors, ids): self._id_to_index = {vid: i for i, vid in enumerate(ids)} - docs = [ - self._Document(id=vid, vector=vectors[i].tolist()) - for i, vid in enumerate(ids) - ] + docs = [self._Document(id=vid, vector=vectors[i].tolist()) for i, vid in enumerate(ids)] # upsert in batches so we don't build one giant list in memory B = 500 for s in range(0, len(docs), B): @@ -91,8 +88,10 @@ def query(self, q, k): def run(args) -> None: - print(f"Generating synthetic dataset: {args.n:,} × {args.dim}d, " - f"{args.queries} queries, k={args.k}") + print( + f"Generating synthetic dataset: {args.n:,} × {args.dim}d, " + f"{args.queries} queries, k={args.k}" + ) ds = make_synthetic(n=args.n, dim=args.dim, n_queries=args.queries, k=args.k) if args.backend == "local": @@ -122,18 +121,23 @@ def run(args) -> None: print(f"backend : {args.backend}") print(f"dataset : {args.n:,} vectors × {args.dim}d") print(f"recall@{args.k:<8}: {recall:.4f}") - print(f"latency (ms) : p50={pct['p50']:.2f} p95={pct['p95']:.2f} " - f"p99={pct['p99']:.2f} mean={pct['mean']:.2f}") + print( + f"latency (ms) : p50={pct['p50']:.2f} p95={pct['p95']:.2f} " + f"p99={pct['p99']:.2f} mean={pct['mean']:.2f}" + ) print(f"build time : {build_s:.2f}s") if args.backend == "local": - print("\nNote: local backend is exact NN (recall≈1.0). Use --backend " - "dynavec for real S3 Vectors recall & end-to-end latency.") + print( + "\nNote: local backend is exact NN (recall≈1.0). Use --backend " + "dynavec for real S3 Vectors recall & end-to-end latency." + ) print("=========================================\n") def main() -> None: - p = argparse.ArgumentParser(description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter) + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) p.add_argument("--backend", choices=["local", "dynavec"], default="local") p.add_argument("--n", type=int, default=10_000, help="number of vectors") p.add_argument("--dim", type=int, default=384) diff --git a/docs/DEPLOY.md b/docs/DEPLOY.md index 5a69433..da01ec9 100644 --- a/docs/DEPLOY.md +++ b/docs/DEPLOY.md @@ -90,10 +90,12 @@ python examples/openai_1536_retrieval.py Or load it in Python without exporting: ```python -from dotenv import load_dotenv # pip install python-dotenv +from dotenv import load_dotenv # pip install python-dotenv + load_dotenv() from dynavec import Dynavec, DynavecConfig -db = Dynavec(DynavecConfig(..., region="us-east-1")) # boto3 picks up the env creds + +db = Dynavec(DynavecConfig(..., region="us-east-1")) # boto3 picks up the env creds ``` Or pass keys explicitly (e.g. multi-account) instead of the env chain: @@ -101,6 +103,7 @@ Or pass keys explicitly (e.g. multi-account) instead of the env chain: ```python import os from dynavec import Dynavec, DynavecConfig, AWSCredentials + creds = AWSCredentials( access_key_id=os.environ["AWS_ACCESS_KEY_ID"], secret_access_key=os.environ["AWS_SECRET_ACCESS_KEY"], diff --git a/examples/crewai_tool.py b/examples/crewai_tool.py index 2ce0793..9a4b65c 100644 --- a/examples/crewai_tool.py +++ b/examples/crewai_tool.py @@ -1,7 +1,7 @@ """Wrap dynavec as a CrewAI tool so an agent can retrieve from it. - pip install "dynavec[sentence-transformers,crewai]" - python examples/crewai_tool.py +pip install "dynavec[sentence-transformers,crewai]" +python examples/crewai_tool.py """ from dynavec import Document, Dynavec, DynavecConfig @@ -24,7 +24,10 @@ db.upsert( [ Document(id="cw-1", text="CrewAI orchestrates multiple agents that collaborate on a task."), - Document(id="cw-2", text="A CrewAI tool is any callable an agent can invoke to fetch information."), + Document( + id="cw-2", + text="A CrewAI tool is any callable an agent can invoke to fetch information.", + ), Document(id="cw-3", text="dynavec keeps all vectors inside your own AWS account."), ], namespace="kb", @@ -35,4 +38,4 @@ # The tool behaves like any other CrewAI tool: give it to an Agent, or call it # directly the way CrewAI itself does, via .run(...). result = search_tool.run("what is a CrewAI tool?") -print(result) \ No newline at end of file +print(result) diff --git a/examples/dashboard_demo.py b/examples/dashboard_demo.py index d3e4426..cf25c0a 100644 --- a/examples/dashboard_demo.py +++ b/examples/dashboard_demo.py @@ -65,7 +65,7 @@ def put_vectors(self, vectors): def _matches(self, meta, flt): if not flt: return True - for clause in (flt["$and"] if "$and" in flt else [flt]): + for clause in flt["$and"] if "$and" in flt else [flt]: for k, v in clause.items(): if meta.get(k) != v: return False @@ -80,13 +80,19 @@ def query(self, query_vector, top_k, filter=None, return_metadata=True, return_d scored.sort(key=lambda x: x[1]) return [{"key": k, "distance": d, "metadata": m} for k, d, m in scored[:top_k]] - def query_pages(self, query_vector, top_k, filter=None, return_metadata=True, return_distance=True): + def query_pages( + self, query_vector, top_k, filter=None, return_metadata=True, return_distance=True + ): hits = self.query(query_vector, top_k, filter) for i in range(0, len(hits), 2): yield hits[i : i + 2] def get_vectors(self, keys, return_metadata=False): - return {k: {"vector": self._store[k][0], "metadata": self._store[k][1]} for k in keys if k in self._store} + return { + k: {"vector": self._store[k][0], "metadata": self._store[k][1]} + for k in keys + if k in self._store + } def delete_vectors(self, keys): for k in keys: @@ -112,9 +118,14 @@ def delete_many(self, ns, ids): TOPICS = ["food", "space", "ai", "aws", "bio"] QUERIES = [ - "apple pie recipe", "rocket to mars", "serverless vectors on aws", - "how do cells work", "transformer attention", "cheap vector database", - "dynamodb latency", "embedding models", + "apple pie recipe", + "rocket to mars", + "serverless vectors on aws", + "how do cells work", + "transformer attention", + "cheap vector database", + "dynamodb latency", + "embedding models", ] NAMESPACES = ["kb", "tenant-a", "tenant-b"] @@ -145,8 +156,11 @@ def main() -> None: for ns in NAMESPACES: db.upsert( [ - Document(id=f"{ns}-{i}", text=f"{random.choice(TOPICS)} document {i} about vectors", - metadata={"topic": random.choice(TOPICS)}) + Document( + id=f"{ns}-{i}", + text=f"{random.choice(TOPICS)} document {i} about vectors", + metadata={"topic": random.choice(TOPICS)}, + ) for i in range(40) ], namespace=ns, diff --git a/examples/dspy_retriever.py b/examples/dspy_retriever.py index 547ad82..4e4e220 100644 --- a/examples/dspy_retriever.py +++ b/examples/dspy_retriever.py @@ -1,7 +1,7 @@ """Use dynavec as a DSPy retrieval module. - pip install "dynavec[sentence-transformers,dspy]" - python examples/dspy_retriever.py +pip install "dynavec[sentence-transformers,dspy]" +python examples/dspy_retriever.py """ import dspy diff --git a/examples/haystack_integration.py b/examples/haystack_integration.py index b24710a..7316763 100644 --- a/examples/haystack_integration.py +++ b/examples/haystack_integration.py @@ -1,7 +1,7 @@ """Use dynavec as a Haystack DocumentStore and retriever. - pip install "dynavec[sentence-transformers,haystack]" - python examples/haystack_integration.py +pip install "dynavec[sentence-transformers,haystack]" +python examples/haystack_integration.py """ from haystack import Document diff --git a/examples/hot_tier.py b/examples/hot_tier.py index 78b2cf5..a1de8c2 100644 --- a/examples/hot_tier.py +++ b/examples/hot_tier.py @@ -22,7 +22,7 @@ region="us-east-1", auto_provision=True, # --- the hot tier --- - hot_tier=True, # keep a hot working set in RAM + hot_tier=True, # keep a hot working set in RAM hot_tier_max_vectors=200_000, # global RAM safety cap ) @@ -30,12 +30,19 @@ db.upsert( [ - Document(id="a", text="Mitochondria are the powerhouse of the cell.", - metadata={"topic": "biology"}), - Document(id="b", text="Rockets reach orbit at roughly 28,000 km/h.", - metadata={"topic": "space"}), - Document(id="c", text="Photosynthesis converts light into chemical energy.", - metadata={"topic": "biology"}), + Document( + id="a", + text="Mitochondria are the powerhouse of the cell.", + metadata={"topic": "biology"}, + ), + Document( + id="b", text="Rockets reach orbit at roughly 28,000 km/h.", metadata={"topic": "space"} + ), + Document( + id="c", + text="Photosynthesis converts light into chemical energy.", + metadata={"topic": "biology"}, + ), ] ) diff --git a/examples/ingest_office_docs.py b/examples/ingest_office_docs.py index ab457af..15cca4a 100644 --- a/examples/ingest_office_docs.py +++ b/examples/ingest_office_docs.py @@ -1,9 +1,11 @@ """Example: Ingesting Word (.docx), PowerPoint (.pptx), and Excel (.xlsx) files into dynavec.""" from pathlib import Path + from dynavec import Dynavec, DynavecConfig from dynavec.ingest import DocxSource, PptxSource, XlsxSource, ingest + def main(): cfg = DynavecConfig( vector_bucket="my-vectors", @@ -32,5 +34,6 @@ def main(): xlsx_count = ingest(db, XlsxSource(xlsx_file), namespace="spreadsheets") print(f"Ingested {xlsx_count} chunks from {xlsx_file}") + if __name__ == "__main__": main() diff --git a/examples/langchain_rag.py b/examples/langchain_rag.py index eeec57b..4e00bcf 100644 --- a/examples/langchain_rag.py +++ b/examples/langchain_rag.py @@ -1,7 +1,7 @@ """Use dynavec as a LangChain VectorStore / retriever. - pip install "dynavec[sentence-transformers,langchain]" - python examples/langchain_rag.py +pip install "dynavec[sentence-transformers,langchain]" +python examples/langchain_rag.py """ from dynavec import Dynavec, DynavecConfig diff --git a/examples/openai_1536_retrieval.py b/examples/openai_1536_retrieval.py index 4f6a68d..9f20e26 100644 --- a/examples/openai_1536_retrieval.py +++ b/examples/openai_1536_retrieval.py @@ -27,13 +27,21 @@ CORPUS = [ ("d1", "The mitochondria is the powerhouse of the cell.", {"topic": "biology"}), - ("d2", "Photosynthesis converts sunlight into chemical energy in plants.", {"topic": "biology"}), + ( + "d2", + "Photosynthesis converts sunlight into chemical energy in plants.", + {"topic": "biology"}, + ), ("d3", "Rockets must reach ~28,000 km/h to achieve low Earth orbit.", {"topic": "space"}), ("d4", "A black hole's event horizon is the point of no return.", {"topic": "space"}), ("d5", "Transformers use self-attention to model long-range dependencies.", {"topic": "ai"}), ("d6", "Vector databases power retrieval-augmented generation for LLMs.", {"topic": "ai"}), ("d7", "DynamoDB offers single-digit-millisecond key-value lookups.", {"topic": "aws"}), - ("d8", "Amazon S3 Vectors provides serverless approximate nearest-neighbor search.", {"topic": "aws"}), + ( + "d8", + "Amazon S3 Vectors provides serverless approximate nearest-neighbor search.", + {"topic": "aws"}, + ), ] QUERIES = [ @@ -54,7 +62,7 @@ def main() -> None: dimension=1536, distance_metric="cosine", region=REGION, - auto_provision=True, # creates bucket + index + table if missing + auto_provision=True, # creates bucket + index + table if missing ) print(f"Connecting to AWS ({REGION}) and provisioning resources ...") diff --git a/examples/pdf_ingestion.py b/examples/pdf_ingestion.py index 157a2b4..13323eb 100644 --- a/examples/pdf_ingestion.py +++ b/examples/pdf_ingestion.py @@ -1,7 +1,7 @@ """Ingest a PDF into dynavec (requires AWS credentials + S3 Vectors). - pip install "dynavec[ingest,sentence-transformers]" - python examples/pdf_ingestion.py path/to/document.pdf +pip install "dynavec[ingest,sentence-transformers]" +python examples/pdf_ingestion.py path/to/document.pdf """ import sys @@ -11,9 +11,7 @@ from dynavec.ingest import PDFSource, ingest if len(sys.argv) != 2: - raise SystemExit( - "usage: python examples/pdf_ingestion.py path/to/document.pdf" - ) + raise SystemExit("usage: python examples/pdf_ingestion.py path/to/document.pdf") # Local embedding model — no external API key required. embedder = SentenceTransformerEmbedder(model="all-MiniLM-L6-v2") diff --git a/examples/pinecone_vs_dynavec.py b/examples/pinecone_vs_dynavec.py index 8cab8b9..80df247 100644 --- a/examples/pinecone_vs_dynavec.py +++ b/examples/pinecone_vs_dynavec.py @@ -64,8 +64,12 @@ def main() -> None: print("Indexing in dynavec (S3 Vectors + DynamoDB) ...") db = Dynavec( DynavecConfig( - vector_bucket="dynavec-vs-pinecone", index="cmp", table="dynavec_cmp", - dimension=DIM, region=REGION, auto_provision=True, + vector_bucket="dynavec-vs-pinecone", + index="cmp", + table="dynavec_cmp", + dimension=DIM, + region=REGION, + auto_provision=True, ) ) db.upsert([Document(id=i, text=t, vector=v) for (i, t), v in zip(CORPUS, doc_vecs)]) @@ -79,7 +83,9 @@ def main() -> None: existing = [i["name"] for i in pc.list_indexes()] if idx_name not in existing: pc.create_index( - name=idx_name, dimension=DIM, metric="cosine", + name=idx_name, + dimension=DIM, + metric="cosine", spec=ServerlessSpec(cloud="aws", region=REGION), ) pindex = pc.Index(idx_name) @@ -103,8 +109,10 @@ def main() -> None: print(f"{q[:38]:<40}{_overlap(dv, pr_ids):>10.2f}{dv_ms:>13.0f}{pr_ms:>14.0f}") db.close() - print("\nNote: 'agree@k' is the fraction of dynavec's top-k also in Pinecone's " - "top-k on identical embeddings — a practical apples-to-apples retrieval check.") + print( + "\nNote: 'agree@k' is the fraction of dynavec's top-k also in Pinecone's " + "top-k on identical embeddings — a practical apples-to-apples retrieval check." + ) if __name__ == "__main__": diff --git a/examples/quickstart.py b/examples/quickstart.py index 2507d7e..57d58d1 100644 --- a/examples/quickstart.py +++ b/examples/quickstart.py @@ -1,7 +1,7 @@ """Minimal end-to-end dynavec example (requires AWS credentials + S3 Vectors). - pip install "dynavec[sentence-transformers]" - python examples/quickstart.py +pip install "dynavec[sentence-transformers]" +python examples/quickstart.py """ from dynavec import Document, Dynavec, DynavecConfig @@ -17,19 +17,26 @@ dimension=embedder.dimension, distance_metric="cosine", region="us-east-1", - auto_provision=True, # creates bucket + index + table if missing + auto_provision=True, # creates bucket + index + table if missing ) db = Dynavec(cfg, embedder=embedder) db.upsert( [ - Document(id="1", text="The mitochondria is the powerhouse of the cell.", - metadata={"topic": "biology"}), - Document(id="2", text="Rockets reach orbit at roughly 28,000 km/h.", - metadata={"topic": "space"}), - Document(id="3", text="Photosynthesis converts sunlight into chemical energy.", - metadata={"topic": "biology"}), + Document( + id="1", + text="The mitochondria is the powerhouse of the cell.", + metadata={"topic": "biology"}, + ), + Document( + id="2", text="Rockets reach orbit at roughly 28,000 km/h.", metadata={"topic": "space"} + ), + Document( + id="3", + text="Photosynthesis converts sunlight into chemical energy.", + metadata={"topic": "biology"}, + ), ], auto_metadata=True, ) diff --git a/examples/verify_provisioning.py b/examples/verify_provisioning.py index 91c30cc..366d1dc 100644 --- a/examples/verify_provisioning.py +++ b/examples/verify_provisioning.py @@ -34,8 +34,12 @@ def _ok(label, detail=""): def verify() -> None: cfg = DynavecConfig( - vector_bucket=BUCKET, index=INDEX, table=TABLE, - dimension=DIM, distance_metric="cosine", region=REGION, + vector_bucket=BUCKET, + index=INDEX, + table=TABLE, + dimension=DIM, + distance_metric="cosine", + region=REGION, auto_provision=True, ) print(f"Provisioning index {INDEX!r} in {REGION} ...") diff --git a/notebooks/rrf_weight_fitting.ipynb b/notebooks/rrf_weight_fitting.ipynb index 48ed87e..cfec4ec 100644 --- a/notebooks/rrf_weight_fitting.ipynb +++ b/notebooks/rrf_weight_fitting.ipynb @@ -2,6 +2,7 @@ "cells": [ { "cell_type": "markdown", + "id": "7fb27b941602401d91542211134fc71a", "metadata": {}, "source": [ "# Fitting RRF / Blend Weights from Labeled Data\n", @@ -15,6 +16,7 @@ }, { "cell_type": "markdown", + "id": "acae54e37e7d407bbb7b55eff062a284", "metadata": {}, "source": [ "## 1. Imports" @@ -23,16 +25,18 @@ { "cell_type": "code", "execution_count": null, + "id": "9a63283cbaf04dbcab1f6479b197f3a8", "metadata": {}, "outputs": [], "source": [ - "from dynavec import RRFWeightFitter, FitResult\n", + "from dynavec import FitResult, RRFWeightFitter\n", "from dynavec.eval import LabeledQuery\n", "from dynavec.models import SearchResult" ] }, { "cell_type": "markdown", + "id": "8dd0d8092fe74a7c96281538738b07e2", "metadata": {}, "source": [ "## 2. Create synthetic retriever results\n", @@ -48,12 +52,46 @@ { "cell_type": "code", "execution_count": null, + "id": "72eea5119410473aa328ad9291626812", "metadata": {}, "outputs": [], - "source": "def r(doc_id, score):\n return SearchResult(id=doc_id, score=score)\n\n# Each list = results for one query from that retriever\n# Shape: [retriever][query_index] -> list[SearchResult]\n\ndense_results = [\n # query 0: \"FD interest rate\" — dense puts doc_fd first ✓\n [r(\"doc_fd\", 0.95), r(\"doc_savings\", 0.60), r(\"doc_noise\", 0.20)],\n # query 1: \"how to open savings account\" — dense puts doc_savings first ✓\n [r(\"doc_savings\", 0.90), r(\"doc_loan\", 0.50), r(\"doc_noise\", 0.15)],\n # query 2: \"home loan EMI calculator\" — dense puts doc_loan first ✓\n [r(\"doc_loan\", 0.88), r(\"doc_fd\", 0.40), r(\"doc_noise\", 0.10)],\n # query 3: \"UPI transaction limit\" — both retrievers agree ✓\n [r(\"doc_upi\", 0.92), r(\"doc_fd\", 0.45), r(\"doc_savings\", 0.20)],\n]\n\nsparse_results = [\n # query 0: sparse buries doc_fd at rank 2 — noisy keyword match ✗\n [r(\"doc_savings\", 0.80), r(\"doc_noise\", 0.50), r(\"doc_fd\", 0.20)],\n # query 1: sparse buries doc_savings at rank 2 ✗\n [r(\"doc_loan\", 0.75), r(\"doc_noise\", 0.45), r(\"doc_savings\", 0.15)],\n # query 2: sparse buries doc_loan at rank 2 ✗\n [r(\"doc_fd\", 0.70), r(\"doc_noise\", 0.40), r(\"doc_loan\", 0.10)],\n # query 3: sparse also gets this right ✓\n [r(\"doc_upi\", 0.85), r(\"doc_fd\", 0.40), r(\"doc_savings\", 0.20)],\n]\n\nprint(f\"Retrievers: 2\")\nprint(f\"Queries: {len(dense_results)}\")" + "source": [ + "def r(doc_id, score):\n", + " return SearchResult(id=doc_id, score=score)\n", + "\n", + "\n", + "# Each list = results for one query from that retriever\n", + "# Shape: [retriever][query_index] -> list[SearchResult]\n", + "\n", + "dense_results = [\n", + " # query 0: \"FD interest rate\" — dense puts doc_fd first ✓\n", + " [r(\"doc_fd\", 0.95), r(\"doc_savings\", 0.60), r(\"doc_noise\", 0.20)],\n", + " # query 1: \"how to open savings account\" — dense puts doc_savings first ✓\n", + " [r(\"doc_savings\", 0.90), r(\"doc_loan\", 0.50), r(\"doc_noise\", 0.15)],\n", + " # query 2: \"home loan EMI calculator\" — dense puts doc_loan first ✓\n", + " [r(\"doc_loan\", 0.88), r(\"doc_fd\", 0.40), r(\"doc_noise\", 0.10)],\n", + " # query 3: \"UPI transaction limit\" — both retrievers agree ✓\n", + " [r(\"doc_upi\", 0.92), r(\"doc_fd\", 0.45), r(\"doc_savings\", 0.20)],\n", + "]\n", + "\n", + "sparse_results = [\n", + " # query 0: sparse buries doc_fd at rank 2 — noisy keyword match ✗\n", + " [r(\"doc_savings\", 0.80), r(\"doc_noise\", 0.50), r(\"doc_fd\", 0.20)],\n", + " # query 1: sparse buries doc_savings at rank 2 ✗\n", + " [r(\"doc_loan\", 0.75), r(\"doc_noise\", 0.45), r(\"doc_savings\", 0.15)],\n", + " # query 2: sparse buries doc_loan at rank 2 ✗\n", + " [r(\"doc_fd\", 0.70), r(\"doc_noise\", 0.40), r(\"doc_loan\", 0.10)],\n", + " # query 3: sparse also gets this right ✓\n", + " [r(\"doc_upi\", 0.85), r(\"doc_fd\", 0.40), r(\"doc_savings\", 0.20)],\n", + "]\n", + "\n", + "print(\"Retrievers: 2\")\n", + "print(f\"Queries: {len(dense_results)}\")" + ] }, { "cell_type": "markdown", + "id": "8edb47106e1a46a883d545849b8ab81b", "metadata": {}, "source": [ "## 3. Define labeled (ground-truth) queries\n", @@ -65,14 +103,15 @@ { "cell_type": "code", "execution_count": null, + "id": "10185d26023b46108eb7d9f57d49d2b3", "metadata": {}, "outputs": [], "source": [ "labeled_queries = [\n", - " LabeledQuery(query=\"FD interest rate\", relevant_ids={\"doc_fd\"}),\n", + " LabeledQuery(query=\"FD interest rate\", relevant_ids={\"doc_fd\"}),\n", " LabeledQuery(query=\"how to open savings account\", relevant_ids={\"doc_savings\"}),\n", - " LabeledQuery(query=\"home loan EMI calculator\", relevant_ids={\"doc_loan\"}),\n", - " LabeledQuery(query=\"UPI transaction limit\", relevant_ids={\"doc_upi\"}),\n", + " LabeledQuery(query=\"home loan EMI calculator\", relevant_ids={\"doc_loan\"}),\n", + " LabeledQuery(query=\"UPI transaction limit\", relevant_ids={\"doc_upi\"}),\n", "]\n", "\n", "print(\"Labeled queries:\")\n", @@ -82,6 +121,7 @@ }, { "cell_type": "markdown", + "id": "8763a12b2bbd4a93a75aff182afb95dc", "metadata": {}, "source": [ "## 4. Build the fitter and check baseline score\n", @@ -93,6 +133,7 @@ { "cell_type": "code", "execution_count": null, + "id": "7623eae2785240b9bd12b16a66d81610", "metadata": {}, "outputs": [], "source": [ @@ -100,8 +141,8 @@ "fitter = RRFWeightFitter(\n", " labeled_queries=labeled_queries,\n", " result_lists=[dense_results, sparse_results],\n", - " k=60, # RRF smoothing constant (standard default)\n", - " eval_k=3, # evaluate nDCG@3\n", + " k=60, # RRF smoothing constant (standard default)\n", + " eval_k=3, # evaluate nDCG@3\n", ")\n", "\n", "baseline = fitter.score([1.0, 1.0])\n", @@ -115,6 +156,7 @@ }, { "cell_type": "markdown", + "id": "7cdc8c89c7104fffa095e18ddfef8986", "metadata": {}, "source": [ "## 5. Fit weights with grid search\n", @@ -126,6 +168,7 @@ { "cell_type": "code", "execution_count": null, + "id": "b118ea5561624da68c537baed56e602f", "metadata": {}, "outputs": [], "source": [ @@ -140,6 +183,7 @@ }, { "cell_type": "markdown", + "id": "938c804e27f84196a10c8828c723f798", "metadata": {}, "source": [ "## 6. Compare fitting methods\n", @@ -153,6 +197,7 @@ { "cell_type": "code", "execution_count": null, + "id": "504fb2a444614c0babb325280ed9130a", "metadata": {}, "outputs": [], "source": [ @@ -161,20 +206,23 @@ "print(f\"{'Method':<12} {'nDCG@3':>8} {'Evals':>8} Weights\")\n", "print(f\"{'─' * 50}\")\n", "for res in [result_grid, result_random]:\n", - " w_str = ', '.join(f'{w:.3f}' for w in res.weights)\n", + " w_str = \", \".join(f\"{w:.3f}\" for w in res.weights)\n", " print(f\"{res.method:<12} {res.score:>8.4f} {res.n_evaluations:>8} [{w_str}]\")\n", "\n", "# Try bayesian if scipy is available\n", "try:\n", " result_bayesian = fitter.fit(method=\"bayesian\", n_points=100)\n", - " w_str = ', '.join(f'{w:.3f}' for w in result_bayesian.weights)\n", - " print(f\"{result_bayesian.method:<12} {result_bayesian.score:>8.4f} {result_bayesian.n_evaluations:>8} [{w_str}]\")\n", + " w_str = \", \".join(f\"{w:.3f}\" for w in result_bayesian.weights)\n", + " print(\n", + " f\"{result_bayesian.method:<12} {result_bayesian.score:>8.4f} {result_bayesian.n_evaluations:>8} [{w_str}]\"\n", + " )\n", "except Exception as e:\n", " print(f\"\\nNote: bayesian method unavailable ({e})\")" ] }, { "cell_type": "markdown", + "id": "59bbdb311c014d738909a11f9e486628", "metadata": {}, "source": [ "## 7. Sweep weights to visualize the score landscape\n", @@ -185,6 +233,7 @@ { "cell_type": "code", "execution_count": null, + "id": "b43b363d81ae4b689946ece5c682cd59", "metadata": {}, "outputs": [], "source": [ @@ -198,8 +247,15 @@ "\n", " plt.figure(figsize=(8, 4))\n", " plt.plot(w1_values, scores, color=\"steelblue\", linewidth=2)\n", - " plt.axvline(result_grid.weights[0], color=\"tomato\", linestyle=\"--\", label=f\"best w_dense={result_grid.weights[0]:.2f}\")\n", - " plt.axhline(baseline, color=\"gray\", linestyle=\":\", label=f\"equal weights baseline ({baseline:.3f})\")\n", + " plt.axvline(\n", + " result_grid.weights[0],\n", + " color=\"tomato\",\n", + " linestyle=\"--\",\n", + " label=f\"best w_dense={result_grid.weights[0]:.2f}\",\n", + " )\n", + " plt.axhline(\n", + " baseline, color=\"gray\", linestyle=\":\", label=f\"equal weights baseline ({baseline:.3f})\"\n", + " )\n", " plt.xlabel(\"w_dense (w_sparse = 1 - w_dense)\")\n", " plt.ylabel(\"mean nDCG@3\")\n", " plt.title(\"nDCG@3 vs. RRF weight on dense retriever\")\n", @@ -211,11 +267,12 @@ " print(f\"{'w_dense':>8} {'w_sparse':>8} {'nDCG@3':>8}\")\n", " for w1, s in zip(w1_values[::10], scores[::10]):\n", " marker = \" ← best\" if abs(w1 - result_grid.weights[0]) < 0.01 else \"\"\n", - " print(f\"{w1:8.2f} {1-w1:8.2f} {s:8.4f}{marker}\")" + " print(f\"{w1:8.2f} {1 - w1:8.2f} {s:8.4f}{marker}\")" ] }, { "cell_type": "markdown", + "id": "8a65eabff63a45729fe45fb5ade58bdc", "metadata": {}, "source": [ "## 8. Three-retriever example with random search\n", @@ -228,6 +285,7 @@ { "cell_type": "code", "execution_count": null, + "id": "c3933fab20d04ec698c2621248eb3be0", "metadata": {}, "outputs": [], "source": [ @@ -252,12 +310,13 @@ "print(f\"{'Method':<12} {'nDCG@3':>8} {'Evals':>8} Weights\")\n", "print(f\"{'─' * 60}\")\n", "for res in [result3_grid, result3_random]:\n", - " w_str = ', '.join(f'{w:.3f}' for w in res.weights)\n", + " w_str = \", \".join(f\"{w:.3f}\" for w in res.weights)\n", " print(f\"{res.method:<12} {res.score:>8.4f} {res.n_evaluations:>8} [{w_str}]\")" ] }, { "cell_type": "markdown", + "id": "4dd4641cc4064e0191573fe9c69df29b", "metadata": {}, "source": [ "## 9. Save and load fitted weights\n", @@ -269,6 +328,7 @@ { "cell_type": "code", "execution_count": null, + "id": "8309879909854d7188b41380fd92a7c3", "metadata": {}, "outputs": [], "source": [ @@ -293,6 +353,7 @@ }, { "cell_type": "markdown", + "id": "3ed186c9a28b402fb0bc4494df01f08d", "metadata": {}, "source": [ "## 10. Use the fitted weights in production\n", @@ -303,18 +364,19 @@ { "cell_type": "code", "execution_count": null, + "id": "cb1e1581032b452c9409d6c6813c49d1", "metadata": {}, "outputs": [], "source": [ "from dynavec import reciprocal_rank_fusion\n", "\n", "# In production, call each retriever, then fuse with learned weights\n", - "new_dense_results = [r(\"doc_fd\", 0.91), r(\"doc_savings\", 0.55)]\n", + "new_dense_results = [r(\"doc_fd\", 0.91), r(\"doc_savings\", 0.55)]\n", "new_sparse_results = [r(\"doc_savings\", 0.70), r(\"doc_fd\", 0.35)]\n", "\n", "fused = reciprocal_rank_fusion(\n", " [new_dense_results, new_sparse_results],\n", - " weights=result_grid.weights, # ← the weights RRFWeightFitter learned\n", + " weights=result_grid.weights, # ← the weights RRFWeightFitter learned\n", ")\n", "\n", "print(\"Fused ranking with fitted weights:\")\n", diff --git a/src/dynavec/agents/base.py b/src/dynavec/agents/base.py index 102fb9c..52c5f6c 100644 --- a/src/dynavec/agents/base.py +++ b/src/dynavec/agents/base.py @@ -6,12 +6,9 @@ import json from collections.abc import Callable from dataclasses import dataclass, field - - -from ..chat.base import Tool, ToolCall -import functools from typing import Any, overload +from ..chat.base import Tool, ToolCall @dataclass @@ -31,11 +28,14 @@ class AgentResult: output: str steps: list[AgentStep] = field(default_factory=list) finished: bool = True - termination_reason: str = "completed" # "completed", "max_steps_reached", "error", "interrupted" + termination_reason: str = ( + "completed" # "completed", "max_steps_reached", "error", "interrupted" + ) total_steps: int = 0 tool_calls_count: int = 0 interrupt_payload: dict[str, Any] | None = None + @dataclass class PlanStep: """A single step in a decomposed task plan.""" @@ -120,11 +120,11 @@ def __init__( ) -> None: self.fn = fn # Extract name directly or fall back to function attribute - self.name = name or getattr(fn, "name", None) or getattr(fn, "__name__", None) or "unnamed_tool" - self.description = description or (fn.__doc__ or f"Execute {self.name}").strip() - self.parameters = ( - parameters if parameters is not None else _generate_json_schema(fn) + self.name = ( + name or getattr(fn, "name", None) or getattr(fn, "__name__", None) or "unnamed_tool" ) + self.description = description or (fn.__doc__ or f"Execute {self.name}").strip() + self.parameters = parameters if parameters is not None else _generate_json_schema(fn) def to_chat_tool(self) -> Tool: """Convert to a dynavec.chat.Tool schema.""" @@ -161,9 +161,7 @@ def execute(self, arguments: dict[str, Any] | str | None = None) -> str: result = self.fn() elif len(params) == 1 and list(params.keys())[0] not in parsed_args: # If single param expected and keys don't match, pass the first val or raw dict - first_val = ( - next(iter(parsed_args.values())) if parsed_args else arguments - ) + first_val = next(iter(parsed_args.values())) if parsed_args else arguments result = self.fn(first_val) else: # Filter only valid keyword arguments @@ -173,7 +171,7 @@ def execute(self, arguments: dict[str, Any] | str | None = None) -> str: if isinstance(result, str): return result return json.dumps(result, ensure_ascii=False) - + except NodeInterrupt: raise @@ -187,6 +185,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any: @overload def tool(fn: Callable[..., Any], /) -> AgentTool: ... + @overload def tool( *, @@ -195,13 +194,14 @@ def tool( parameters: dict[str, Any] | None = None, ) -> Callable[[Callable[..., Any]], AgentTool]: ... + def tool( fn_or_name: Callable[..., Any] | str | None = None, description: str | None = None, parameters: dict[str, Any] | None = None, ) -> AgentTool | Callable[[Callable[..., Any]], AgentTool]: """Decorator to convert a standard Python function into an AgentTool. - + Supports both `@tool` and `@tool(name="...", description="...")`. """ # Case 1: Used as `@tool` without parentheses diff --git a/src/dynavec/agents/react.py b/src/dynavec/agents/react.py index 346afdf..ce06fe3 100644 --- a/src/dynavec/agents/react.py +++ b/src/dynavec/agents/react.py @@ -3,14 +3,13 @@ from __future__ import annotations import asyncio +import json from collections.abc import Callable, Sequence from typing import Any from ..chat.base import ChatModel, Message, Tool -from .base import AgentResult, AgentStep, AgentTool - from ..exceptions import NodeInterrupt -import json +from .base import AgentResult, AgentStep, AgentTool DEFAULT_REACT_SYSTEM_PROMPT = """You are a helpful and precise reasoning agent. You solve tasks step-by-step using a ReAct (Reason + Act) approach. @@ -67,9 +66,7 @@ def resume( ) -> AgentResult: """Resume an interrupted run synchronously by injecting the human decision into context.""" decision_str = ( - json.dumps(human_decision) - if isinstance(human_decision, dict) - else str(human_decision) + json.dumps(human_decision) if isinstance(human_decision, dict) else str(human_decision) ) messages.append(Message(role="user", content=f"Human Decision / Input: {decision_str}")) return self._execute_loop( @@ -168,7 +165,6 @@ def _execute_loop( tool_calls_count=total_tool_calls, ) - async def arun(self, goal: str, thread_id: str = "default", **kwargs: Any) -> AgentResult: """Execute the ReAct loop asynchronously until goal completion, max_steps, or interrupt.""" messages: list[Message] = [] @@ -176,7 +172,9 @@ async def arun(self, goal: str, thread_id: str = "default", **kwargs: Any) -> Ag messages.append(Message(role="system", content=self.system_prompt)) messages.append(Message(role="user", content=goal)) - return await self._aexecute_loop(messages=messages, thread_id=thread_id, start_step=1, **kwargs) + return await self._aexecute_loop( + messages=messages, thread_id=thread_id, start_step=1, **kwargs + ) async def aresume( self, @@ -188,9 +186,7 @@ async def aresume( ) -> AgentResult: """Resume an interrupted run asynchronously by injecting human decision into context.""" decision_str = ( - json.dumps(human_decision) - if isinstance(human_decision, dict) - else str(human_decision) + json.dumps(human_decision) if isinstance(human_decision, dict) else str(human_decision) ) messages.append(Message(role="user", content=f"Human Decision / Input: {decision_str}")) return await self._aexecute_loop( diff --git a/src/dynavec/async_client.py b/src/dynavec/async_client.py index f901d45..83d277d 100644 --- a/src/dynavec/async_client.py +++ b/src/dynavec/async_client.py @@ -72,18 +72,11 @@ def __init__( self._default_transform = as_pipeline(transform) self._cache = cache self._telemetry = telemetry - self._hot = ( - HotTier(config) - if config.hot_tier - else None - ) + self._hot = HotTier(config) if config.hot_tier else None self._exit_stack: AsyncExitStack | None = None - if ( - embedder is not None - and embedder.dimension != config.dimension - ): + if embedder is not None and embedder.dimension != config.dimension: raise ConfigurationError( f"Embedder dimension ({embedder.dimension}) " f"!= index dimension ({config.dimension}). " @@ -146,10 +139,7 @@ def _require_open(self) -> None: ) def _invalidate_cache(self, namespace: str) -> None: - if ( - self._cache is not None - and self.config.cache_invalidate_on_write - ): + if self._cache is not None and self.config.cache_invalidate_on_write: self._cache.invalidate(namespace) async def _resolve_query_vector( @@ -160,15 +150,12 @@ async def _resolve_query_vector( if vector is not None: if len(vector) != self.config.dimension: raise DimensionMismatchError( - f"Query vector dimension {len(vector)} " - f"!= {self.config.dimension}." + f"Query vector dimension {len(vector)} != {self.config.dimension}." ) return vector if query is None: - raise ValueError( - "Provide either 'query' text or a 'vector'." - ) + raise ValueError("Provide either 'query' text or a 'vector'.") if self.embedder is None: raise ConfigurationError( @@ -270,11 +257,7 @@ async def asearch_stream( yielded = 0 - effective_page_size = ( - page_size - if page_size is not None - else self.config.top_k_page_size - ) + effective_page_size = page_size if page_size is not None else self.config.top_k_page_size async for page in self._vectors.query_pages( query_vector=query_vector, @@ -336,10 +319,7 @@ async def _prepare( list[str], list[HotPayload], ]: - pipeline = ( - as_pipeline(transform) - or self._default_transform - ) + pipeline = as_pipeline(transform) or self._default_transform apply_transform_pipeline( docs, @@ -359,9 +339,7 @@ async def _prepare( texts = embedding_texts(to_embed) - vectors = await self.embedder.aembed_documents( - texts - ) + vectors = await self.embedder.aembed_documents(texts) assign_embeddings( docs, @@ -379,10 +357,7 @@ async def _prepare( async def aupsert( self, - documents: Sequence[ - Document | dict[str, Any] - ] - | None = None, + documents: Sequence[Document | dict[str, Any]] | None = None, *, namespace: str = "default", auto_metadata: bool = False, @@ -391,9 +366,7 @@ async def aupsert( ) -> UpsertResult: """Insert or overwrite documents asynchronously.""" if ttl_seconds is not None and ttl_seconds <= 0: - raise ValueError( - f"ttl_seconds must be positive, got {ttl_seconds}." - ) + raise ValueError(f"ttl_seconds must be positive, got {ttl_seconds}.") if not documents: return UpsertResult( @@ -404,9 +377,7 @@ async def aupsert( self._require_open() docs = [ - document - if isinstance(document, Document) - else Document(**document) + document if isinstance(document, Document) else Document(**document) for document in documents ] diff --git a/src/dynavec/bm25.py b/src/dynavec/bm25.py index ae0d4e8..1e7b373 100644 --- a/src/dynavec/bm25.py +++ b/src/dynavec/bm25.py @@ -19,11 +19,43 @@ _TOKEN_RE = re.compile(r"[a-zA-Z0-9]+(?:[-_.:][a-zA-Z0-9]+)*") _SUBTOKEN_RE = re.compile(r"[a-zA-Z0-9]+") -DEFAULT_STOPWORDS: frozenset[str] = frozenset({ - "a", "an", "and", "are", "as", "at", "be", "but", "by", "for", "if", "in", - "into", "is", "it", "no", "not", "of", "on", "or", "such", "that", "the", - "their", "then", "there", "these", "they", "this", "to", "was", "will", "with" -}) +DEFAULT_STOPWORDS: frozenset[str] = frozenset( + { + "a", + "an", + "and", + "are", + "as", + "at", + "be", + "but", + "by", + "for", + "if", + "in", + "into", + "is", + "it", + "no", + "not", + "of", + "on", + "or", + "such", + "that", + "the", + "their", + "then", + "there", + "these", + "they", + "this", + "to", + "was", + "will", + "with", + } +) def default_tokenize( diff --git a/src/dynavec/chat/anthropic.py b/src/dynavec/chat/anthropic.py index 2eb4aa3..616aee1 100644 --- a/src/dynavec/chat/anthropic.py +++ b/src/dynavec/chat/anthropic.py @@ -37,42 +37,45 @@ def __init__( 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 {} - }) - + 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 - }) + 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 - }) - + 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: @@ -92,31 +95,34 @@ def invoke( ) -> 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} + + 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) - ) + 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) + 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( @@ -124,8 +130,13 @@ def stream( ) -> 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} + + args: dict[str, Any] = { + "model": self.model, + "messages": anthropic_msgs, + "max_tokens": 1024, + **kwargs, + } if system: args["system"] = system if anthropic_tools: diff --git a/src/dynavec/chat/base.py b/src/dynavec/chat/base.py index 876c126..57ce5ee 100644 --- a/src/dynavec/chat/base.py +++ b/src/dynavec/chat/base.py @@ -10,80 +10,80 @@ 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 + 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 + 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 + 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 + 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 diff --git a/src/dynavec/chat/bedrock.py b/src/dynavec/chat/bedrock.py index 710ecc8..9f4787b 100644 --- a/src/dynavec/chat/bedrock.py +++ b/src/dynavec/chat/bedrock.py @@ -34,7 +34,9 @@ def __init__( 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]]]: + 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]] = [] @@ -50,32 +52,35 @@ def _convert_messages(self, messages: list[Message]) -> tuple[list[dict[str, Any 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 {} + 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" + content_blocks.append( + { + "toolResult": { + "toolUseId": msg.tool_call_id, + "content": [{"text": msg.content or ""}], + "status": "success", + } } - }) - bedrock_msgs.append({ - "role": "user", - "content": content_blocks - }) + ) + bedrock_msgs.append({"role": "user", "content": content_blocks}) continue - bedrock_msgs.append({ - "role": "assistant" if msg.role == "assistant" else "user", - "content": content_blocks - }) + bedrock_msgs.append( + { + "role": "assistant" if msg.role == "assistant" else "user", + "content": content_blocks, + } + ) return system_prompts, bedrock_msgs @@ -88,9 +93,7 @@ def _convert_tools(self, tools: list[Tool] | None) -> dict[str, Any] | None: "toolSpec": { "name": t.name, "description": t.description, - "inputSchema": { - "json": t.parameters - } + "inputSchema": {"json": t.parameters}, } } for t in tools @@ -103,39 +106,33 @@ def invoke( system, bedrock_msgs = self._convert_messages(messages) bedrock_tools = self._convert_tools(tools) - args: dict[str, Any] = { - "modelId": self.model, - "messages": bedrock_msgs, - **kwargs - } + 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"]) - ) + 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) + 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( @@ -144,11 +141,7 @@ def stream( system, bedrock_msgs = self._convert_messages(messages) bedrock_tools = self._convert_tools(tools) - args: dict[str, Any] = { - "modelId": self.model, - "messages": bedrock_msgs, - **kwargs - } + args: dict[str, Any] = {"modelId": self.model, "messages": bedrock_msgs, **kwargs} if system: args["system"] = system if bedrock_tools: @@ -167,10 +160,6 @@ def stream( elif "toolUse" in delta: yield ChatChunk( tool_calls=[ - ToolCall( - id="", - name="", - arguments=delta["toolUse"].get("input", "") - ) + ToolCall(id="", name="", arguments=delta["toolUse"].get("input", "")) ] ) diff --git a/src/dynavec/chat/openai.py b/src/dynavec/chat/openai.py index 5490710..1ff57a8 100644 --- a/src/dynavec/chat/openai.py +++ b/src/dynavec/chat/openai.py @@ -40,10 +40,10 @@ 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"] = [ { @@ -53,11 +53,11 @@ def _convert_messages(self, messages: list[Message]) -> list[dict[str, Any]]: } 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 @@ -82,14 +82,14 @@ def invoke( """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 = [ @@ -109,20 +109,25 @@ def stream( """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} + + 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: @@ -130,8 +135,10 @@ def stream( 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 "", + arguments=tc.function.arguments + if (tc.function and tc.function.arguments) + else "", ) ) - + yield ChatChunk(content=delta.content, tool_calls=tool_calls) diff --git a/src/dynavec/cli.py b/src/dynavec/cli.py index 817b6d6..8ef7f33 100644 --- a/src/dynavec/cli.py +++ b/src/dynavec/cli.py @@ -13,297 +13,301 @@ def _parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(prog="dynavec") - subparsers = parser.add_subparsers(dest="command") - subparsers.add_parser("version", help="show the installed dynavec version") - doctor = subparsers.add_parser("doctor", help="check AWS credentials and resource access") - doctor.add_argument("--bucket", help="S3 Vectors bucket name") - doctor.add_argument("--index", help="S3 Vectors index name") - doctor.add_argument("--table", help="DynamoDB table name") - doctor.add_argument("--region", help="AWS region") - doctor.add_argument("--profile", help="AWS profile name") - - export_cmd = subparsers.add_parser( - "export", help="dump vectors and documents for a namespace to JSONL" - ) - export_cmd.add_argument( - "--namespace", - "-n", - default="default", - help="namespace to export (default: 'default')", - ) - export_cmd.add_argument( - "--output", - "-o", - help="output JSONL file path (default: stdout)", - ) - export_cmd.add_argument("--bucket", help="S3 Vectors bucket name (or DYNAVEC_VECTOR_BUCKET)") - export_cmd.add_argument("--index", help="S3 Vectors index name (or DYNAVEC_INDEX)") - export_cmd.add_argument("--table", help="DynamoDB table name (or DYNAVEC_TABLE)") - export_cmd.add_argument("--dimension", type=int, help="vector dimension (default: 1536)") - export_cmd.add_argument("--region", help="AWS region") - export_cmd.add_argument("--profile", help="AWS profile name") - - import_cmd = subparsers.add_parser( - "import", help="restore vectors and documents for a namespace from JSONL" - ) - import_cmd.add_argument( - "--namespace", - "-n", - default="default", - help="target namespace (default: 'default')", - ) - import_cmd.add_argument( - "--input", - "-i", - help="input JSONL file path (default: stdin)", - ) - import_cmd.add_argument( - "--batch-size", - type=int, - default=100, - help="batch size for upserting records (default: 100)", - ) - import_cmd.add_argument("--bucket", help="S3 Vectors bucket name (or DYNAVEC_VECTOR_BUCKET)") - import_cmd.add_argument("--index", help="S3 Vectors index name (or DYNAVEC_INDEX)") - import_cmd.add_argument("--table", help="DynamoDB table name (or DYNAVEC_TABLE)") - import_cmd.add_argument( - "--dimension", type=int, help="vector dimension (inferred from input if omitted)" - ) - import_cmd.add_argument("--region", help="AWS region") - import_cmd.add_argument("--profile", help="AWS profile name") - - mcp = subparsers.add_parser("mcp", help="run the FastMCP server for AI clients") - mcp.add_argument( - "--transport", - choices=["stdio", "sse"], - default="stdio", - help="Transport mode (default: stdio)", - ) - mcp.add_argument( - "--port", - type=int, - default=8000, - help="Port for SSE transport (default: 8000)", - ) - return parser + parser = argparse.ArgumentParser(prog="dynavec") + subparsers = parser.add_subparsers(dest="command") + subparsers.add_parser("version", help="show the installed dynavec version") + doctor = subparsers.add_parser("doctor", help="check AWS credentials and resource access") + doctor.add_argument("--bucket", help="S3 Vectors bucket name") + doctor.add_argument("--index", help="S3 Vectors index name") + doctor.add_argument("--table", help="DynamoDB table name") + doctor.add_argument("--region", help="AWS region") + doctor.add_argument("--profile", help="AWS profile name") + + export_cmd = subparsers.add_parser( + "export", help="dump vectors and documents for a namespace to JSONL" + ) + export_cmd.add_argument( + "--namespace", + "-n", + default="default", + help="namespace to export (default: 'default')", + ) + export_cmd.add_argument( + "--output", + "-o", + help="output JSONL file path (default: stdout)", + ) + export_cmd.add_argument("--bucket", help="S3 Vectors bucket name (or DYNAVEC_VECTOR_BUCKET)") + export_cmd.add_argument("--index", help="S3 Vectors index name (or DYNAVEC_INDEX)") + export_cmd.add_argument("--table", help="DynamoDB table name (or DYNAVEC_TABLE)") + export_cmd.add_argument("--dimension", type=int, help="vector dimension (default: 1536)") + export_cmd.add_argument("--region", help="AWS region") + export_cmd.add_argument("--profile", help="AWS profile name") + + import_cmd = subparsers.add_parser( + "import", help="restore vectors and documents for a namespace from JSONL" + ) + import_cmd.add_argument( + "--namespace", + "-n", + default="default", + help="target namespace (default: 'default')", + ) + import_cmd.add_argument( + "--input", + "-i", + help="input JSONL file path (default: stdin)", + ) + import_cmd.add_argument( + "--batch-size", + type=int, + default=100, + help="batch size for upserting records (default: 100)", + ) + import_cmd.add_argument("--bucket", help="S3 Vectors bucket name (or DYNAVEC_VECTOR_BUCKET)") + import_cmd.add_argument("--index", help="S3 Vectors index name (or DYNAVEC_INDEX)") + import_cmd.add_argument("--table", help="DynamoDB table name (or DYNAVEC_TABLE)") + import_cmd.add_argument( + "--dimension", type=int, help="vector dimension (inferred from input if omitted)" + ) + import_cmd.add_argument("--region", help="AWS region") + import_cmd.add_argument("--profile", help="AWS profile name") + + mcp = subparsers.add_parser("mcp", help="run the FastMCP server for AI clients") + mcp.add_argument( + "--transport", + choices=["stdio", "sse"], + default="stdio", + help="Transport mode (default: stdio)", + ) + mcp.add_argument( + "--port", + type=int, + default=8000, + help="Port for SSE transport (default: 8000)", + ) + return parser def _session(profile: str | None, region: str | None) -> Any: - import boto3 + import boto3 - kwargs: dict[str, str] = {} - if profile: - kwargs["profile_name"] = profile - if region: - kwargs["region_name"] = region - return boto3.Session(**kwargs) + kwargs: dict[str, str] = {} + if profile: + kwargs["profile_name"] = profile + if region: + kwargs["region_name"] = region + return boto3.Session(**kwargs) def _resolve_resources(args: argparse.Namespace) -> tuple[str, str, str]: - import os - - bucket = args.bucket or os.environ.get("DYNAVEC_VECTOR_BUCKET") or os.environ.get("DYNAVEC_BUCKET") - index = args.index or os.environ.get("DYNAVEC_INDEX") - table = args.table or os.environ.get("DYNAVEC_TABLE") - - missing = [] - if not bucket: - missing.append("--bucket") - if not index: - missing.append("--index") - if not table: - missing.append("--table") - - if missing: - raise ValueError( - f"Missing required resource configuration: {', '.join(missing)} " - "(provide via flags or DYNAVEC_* environment variables)." - ) - return str(bucket), str(index), str(table) + import os + + bucket = ( + args.bucket or os.environ.get("DYNAVEC_VECTOR_BUCKET") or os.environ.get("DYNAVEC_BUCKET") + ) + index = args.index or os.environ.get("DYNAVEC_INDEX") + table = args.table or os.environ.get("DYNAVEC_TABLE") + + missing = [] + if not bucket: + missing.append("--bucket") + if not index: + missing.append("--index") + if not table: + missing.append("--table") + + if missing: + raise ValueError( + f"Missing required resource configuration: {', '.join(missing)} " + "(provide via flags or DYNAVEC_* environment variables)." + ) + return str(bucket), str(index), str(table) def _export(args: argparse.Namespace) -> int: - try: - bucket, index, table = _resolve_resources(args) - except ValueError as exc: - print(f"[FAIL] {exc}", file=sys.stderr) - return 1 - - session = _session(args.profile, args.region) - dimension = args.dimension or 1536 - config = DynavecConfig( - vector_bucket=bucket, - index=index, - table=table, - dimension=dimension, - region=args.region, - ) - db = Dynavec(config, boto_session=session) - - try: - if args.output and args.output != "-": - count = db.export_namespace(args.output, namespace=args.namespace) - print(f"Exported {count} documents from namespace '{args.namespace}' to {args.output}") - else: - count = db.export_namespace(sys.stdout, namespace=args.namespace) - print(f"Exported {count} documents from namespace '{args.namespace}'", file=sys.stderr) - return 0 - except Exception as exc: - print(f"[FAIL] Export failed: {exc}", file=sys.stderr) - return 1 + try: + bucket, index, table = _resolve_resources(args) + except ValueError as exc: + print(f"[FAIL] {exc}", file=sys.stderr) + return 1 + + session = _session(args.profile, args.region) + dimension = args.dimension or 1536 + config = DynavecConfig( + vector_bucket=bucket, + index=index, + table=table, + dimension=dimension, + region=args.region, + ) + db = Dynavec(config, boto_session=session) + + try: + if args.output and args.output != "-": + count = db.export_namespace(args.output, namespace=args.namespace) + print(f"Exported {count} documents from namespace '{args.namespace}' to {args.output}") + else: + count = db.export_namespace(sys.stdout, namespace=args.namespace) + print(f"Exported {count} documents from namespace '{args.namespace}'", file=sys.stderr) + return 0 + except Exception as exc: + print(f"[FAIL] Export failed: {exc}", file=sys.stderr) + return 1 def _import(args: argparse.Namespace) -> int: - import json - - try: - bucket, index, table = _resolve_resources(args) - except ValueError as exc: - print(f"[FAIL] {exc}", file=sys.stderr) - return 1 - - session = _session(args.profile, args.region) - dimension = args.dimension - - if dimension is None and args.input and args.input != "-": - try: - with open(args.input, encoding="utf-8") as f: - for line in f: - line = line.strip() - if line: - first_item = json.loads(line) - if "vector" in first_item and isinstance(first_item["vector"], list): - dimension = len(first_item["vector"]) - break - except Exception: - pass - - if dimension is None: - try: - client = session.client("s3vectors", region_name=args.region) - idx = client.get_index(vectorBucketName=bucket, indexName=index) - dimension = idx.get("index", {}).get("dimension") - except Exception: - pass - - if dimension is None: - dimension = 1536 - - config = DynavecConfig( - vector_bucket=bucket, - index=index, - table=table, - dimension=dimension, - region=args.region, - ) - db = Dynavec(config, boto_session=session) - - try: - if args.input and args.input != "-": - count = db.import_namespace( - args.input, namespace=args.namespace, batch_size=args.batch_size - ) - print(f"Imported {count} documents into namespace '{args.namespace}' from {args.input}") - else: - count = db.import_namespace( - sys.stdin, namespace=args.namespace, batch_size=args.batch_size - ) - print(f"Imported {count} documents into namespace '{args.namespace}'") - return 0 - except Exception as exc: - print(f"[FAIL] Import failed: {exc}", file=sys.stderr) - return 1 + import json + + try: + bucket, index, table = _resolve_resources(args) + except ValueError as exc: + print(f"[FAIL] {exc}", file=sys.stderr) + return 1 + + session = _session(args.profile, args.region) + dimension = args.dimension + + if dimension is None and args.input and args.input != "-": + try: + with open(args.input, encoding="utf-8") as f: + for line in f: + line = line.strip() + if line: + first_item = json.loads(line) + if "vector" in first_item and isinstance(first_item["vector"], list): + dimension = len(first_item["vector"]) + break + except Exception: + pass + + if dimension is None: + try: + client = session.client("s3vectors", region_name=args.region) + idx = client.get_index(vectorBucketName=bucket, indexName=index) + dimension = idx.get("index", {}).get("dimension") + except Exception: + pass + + if dimension is None: + dimension = 1536 + + config = DynavecConfig( + vector_bucket=bucket, + index=index, + table=table, + dimension=dimension, + region=args.region, + ) + db = Dynavec(config, boto_session=session) + + try: + if args.input and args.input != "-": + count = db.import_namespace( + args.input, namespace=args.namespace, batch_size=args.batch_size + ) + print(f"Imported {count} documents into namespace '{args.namespace}' from {args.input}") + else: + count = db.import_namespace( + sys.stdin, namespace=args.namespace, batch_size=args.batch_size + ) + print(f"Imported {count} documents into namespace '{args.namespace}'") + return 0 + except Exception as exc: + print(f"[FAIL] Import failed: {exc}", file=sys.stderr) + return 1 def _check(label: str, callback: Callable[[], str]) -> bool: - try: - detail = callback() - except Exception as exc: # noqa: BLE001 - print(f"[FAIL] {label}\n {exc}") - return False - print(f"[PASS] {label}") - if detail: - print(f" {detail}") - return True + try: + detail = callback() + except Exception as exc: # noqa: BLE001 + print(f"[FAIL] {label}\n {exc}") + return False + print(f"[PASS] {label}") + if detail: + print(f" {detail}") + return True def _doctor(args: argparse.Namespace) -> int: - print("Dynavec doctor\n") - session = None - checks_passed = True - - def get_session() -> Any: - nonlocal session - if session is None: - session = _session(args.profile, args.region) - return session - - checks_passed &= _check( - "AWS credentials / STS identity", - lambda: _identity(get_session()), - ) - - if args.bucket or args.index: - if not args.bucket or not args.index: - print("[FAIL] S3 Vectors configuration\n --bucket and --index must be provided together") - checks_passed = False - else: - checks_passed &= _check( - f"S3 Vectors index: {args.index}", - lambda: _check_s3vectors(get_session(), args.bucket, args.index, args.region), - ) - - if args.table: - checks_passed &= _check( - f"DynamoDB table: {args.table}", - lambda: _check_dynamodb(get_session(), args.table, args.region), - ) - - print("\nDoctor checks passed." if checks_passed else "\nDoctor checks failed.") - return 0 if checks_passed else 1 + print("Dynavec doctor\n") + session = None + checks_passed = True + + def get_session() -> Any: + nonlocal session + if session is None: + session = _session(args.profile, args.region) + return session + + checks_passed &= _check( + "AWS credentials / STS identity", + lambda: _identity(get_session()), + ) + + if args.bucket or args.index: + if not args.bucket or not args.index: + print( + "[FAIL] S3 Vectors configuration\n --bucket and --index must be provided together" + ) + checks_passed = False + else: + checks_passed &= _check( + f"S3 Vectors index: {args.index}", + lambda: _check_s3vectors(get_session(), args.bucket, args.index, args.region), + ) + + if args.table: + checks_passed &= _check( + f"DynamoDB table: {args.table}", + lambda: _check_dynamodb(get_session(), args.table, args.region), + ) + + print("\nDoctor checks passed." if checks_passed else "\nDoctor checks failed.") + return 0 if checks_passed else 1 def _identity(session: Any) -> str: - identity = session.client("sts").get_caller_identity() - return f"Account: {identity.get('Account', 'unknown')}" + identity = session.client("sts").get_caller_identity() + return f"Account: {identity.get('Account', 'unknown')}" def _check_s3vectors(session: Any, bucket: str, index: str, region: str | None) -> str: - client = session.client("s3vectors", region_name=region) - client.get_index(vectorBucketName=bucket, indexName=index) - return f"Bucket: {bucket}" + client = session.client("s3vectors", region_name=region) + client.get_index(vectorBucketName=bucket, indexName=index) + return f"Bucket: {bucket}" def _check_dynamodb(session: Any, table: str, region: str | None) -> str: - session.client("dynamodb", region_name=region).describe_table(TableName=table) - return "Accessible" + session.client("dynamodb", region_name=region).describe_table(TableName=table) + return "Accessible" def main(argv: list[str] | None = None) -> int: - args = _parser().parse_args(argv) - if args.command == "version": - print(__version__) - return 0 - if args.command == "doctor": - return _doctor(args) - if args.command == "export": - return _export(args) - if args.command == "import": - return _import(args) - if args.command == "mcp": - from .mcp.server import create_mcp_server - - server = create_mcp_server() - if args.transport == "sse": - server.settings.port = args.port - server.run(transport="sse") - else: - server.run(transport="stdio") - return 0 - _parser().print_help() - return 0 + args = _parser().parse_args(argv) + if args.command == "version": + print(__version__) + return 0 + if args.command == "doctor": + return _doctor(args) + if args.command == "export": + return _export(args) + if args.command == "import": + return _import(args) + if args.command == "mcp": + from .mcp.server import create_mcp_server + + server = create_mcp_server() + if args.transport == "sse": + server.settings.port = args.port + server.run(transport="sse") + else: + server.run(transport="stdio") + return 0 + _parser().print_help() + return 0 if __name__ == "__main__": - sys.exit(main()) + sys.exit(main()) diff --git a/src/dynavec/client.py b/src/dynavec/client.py index 737cd4a..0639202 100644 --- a/src/dynavec/client.py +++ b/src/dynavec/client.py @@ -262,7 +262,6 @@ def _prepare( default_ttl_seconds=default_ttl_seconds, ) - def _write( self, namespace: str, diff --git a/src/dynavec/client_common.py b/src/dynavec/client_common.py index 878510c..7485c21 100644 --- a/src/dynavec/client_common.py +++ b/src/dynavec/client_common.py @@ -21,11 +21,7 @@ def s3_key(namespace: str, doc_id: str) -> str: - return ( - f"{encode_key_component(namespace)}" - f"{KEY_SEPARATOR}" - f"{encode_key_component(doc_id)}" - ) + return f"{encode_key_component(namespace)}{KEY_SEPARATOR}{encode_key_component(doc_id)}" def split_key(key: str) -> tuple[str, str]: @@ -63,11 +59,7 @@ def apply_transform_pipeline( def documents_to_embed( docs: list[Document], ) -> list[EmbeddingTarget]: - return [ - (index, doc.text) - for index, doc in enumerate(docs) - if doc.vector is None - ] + return [(index, doc.text) for index, doc in enumerate(docs) if doc.vector is None] def embedding_texts( @@ -76,15 +68,9 @@ def embedding_texts( optional_texts = [text for _, text in targets] if any(text is None for text in optional_texts): - raise ConfigurationError( - "A document has neither text nor vector." - ) + raise ConfigurationError("A document has neither text nor vector.") - return [ - text - for text in optional_texts - if text is not None - ] + return [text for text in optional_texts if text is not None] def assign_embeddings( @@ -94,8 +80,7 @@ def assign_embeddings( ) -> None: if len(vectors) != len(targets): raise ConfigurationError( - f"Embedder returned {len(vectors)} vectors for " - f"{len(targets)} documents." + f"Embedder returned {len(vectors)} vectors for {len(targets)} documents." ) for (index, _), vector in zip(targets, vectors): @@ -105,8 +90,7 @@ def assign_embeddings( def single_embedding(vectors: list[list[float]], doc_id: str) -> list[float]: if len(vectors) != 1: raise ConfigurationError( - f"Embedder returned {len(vectors)} vectors for 1 document " - f"({doc_id!r})." + f"Embedder returned {len(vectors)} vectors for 1 document ({doc_id!r})." ) return vectors[0] @@ -133,10 +117,7 @@ def build_write_payloads( vector = doc.vector if vector is None: - raise ConfigurationError( - f"Embedder did not return a vector for document " - f"{doc.id!r}." - ) + raise ConfigurationError(f"Embedder did not return a vector for document {doc.id!r}.") if len(vector) != config.dimension: raise DimensionMismatchError( @@ -158,23 +139,14 @@ def build_write_payloads( doc.text, ) - doc_ttl_seconds = ( - doc.ttl_seconds - if doc.ttl_seconds is not None - else default_ttl_seconds - ) + doc_ttl_seconds = doc.ttl_seconds if doc.ttl_seconds is not None else default_ttl_seconds if doc_ttl_seconds is not None and doc_ttl_seconds <= 0: raise ValueError( - f"Document {doc.id!r} ttl_seconds must be positive, " - f"got {doc_ttl_seconds}." + f"Document {doc.id!r} ttl_seconds must be positive, got {doc_ttl_seconds}." ) - ttl_timestamp = ( - int(time.time() + doc_ttl_seconds) - if doc_ttl_seconds is not None - else None - ) + ttl_timestamp = int(time.time() + doc_ttl_seconds) if doc_ttl_seconds is not None else None if ttl_timestamp is not None: ddb_metadata["_ttl"] = ttl_timestamp diff --git a/src/dynavec/config.py b/src/dynavec/config.py index bcd1fbd..57a6fef 100644 --- a/src/dynavec/config.py +++ b/src/dynavec/config.py @@ -158,15 +158,16 @@ def __post_init__(self) -> None: valid_log_levels = {"DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"} if self.log_level.upper() not in valid_log_levels: raise ValueError(f"log_level must be one of {valid_log_levels}") - + if self.dimension > 4096: logger = logging.getLogger(__name__) - logger.warning("Amazon S3 Vectors currently supports a maximum embedding dimension of 4096. You have configured a dimension of %d. This may result in an API error during provisioning or writing data.", self.dimension) + logger.warning( + "Amazon S3 Vectors currently supports a maximum embedding dimension of 4096. You have configured a dimension of %d. This may result in an API error during provisioning or writing data.", + self.dimension, + ) if self.put_rps is not None and self.put_rps <= 0: raise ValueError("put_rps must be positive") if self.query_rps is not None and self.query_rps <= 0: raise ValueError("query_rps must be positive") - - diff --git a/src/dynavec/embeddings/cache.py b/src/dynavec/embeddings/cache.py index f65d364..05c59db 100644 --- a/src/dynavec/embeddings/cache.py +++ b/src/dynavec/embeddings/cache.py @@ -93,7 +93,7 @@ def __init__(self, max_size: int | None = 10_000) -> None: def get(self, key: str) -> Vector | None: if key not in self._store: return None - self._store.move_to_end(key) # mark as recently used + self._store.move_to_end(key) # mark as recently used return self._store[key] def set(self, key: str, vector: Vector) -> None: @@ -101,7 +101,7 @@ def set(self, key: str, vector: Vector) -> None: self._store.move_to_end(key) else: if self.max_size is not None and len(self._store) >= self.max_size: - self._store.popitem(last=False) # evict LRU + self._store.popitem(last=False) # evict LRU self._store[key] = vector def __len__(self) -> int: diff --git a/src/dynavec/embeddings/cached.py b/src/dynavec/embeddings/cached.py index 2357c66..cd2f2e3 100644 --- a/src/dynavec/embeddings/cached.py +++ b/src/dynavec/embeddings/cached.py @@ -133,7 +133,4 @@ def inner(self) -> Embedder: return self._inner def __repr__(self) -> str: - return ( - f"CachedEmbedder(inner={self._inner!r}, " - f"model_key={self._model_key!r})" - ) + return f"CachedEmbedder(inner={self._inner!r}, model_key={self._model_key!r})" diff --git a/src/dynavec/embeddings/huggingface.py b/src/dynavec/embeddings/huggingface.py index 207cddf..749385d 100644 --- a/src/dynavec/embeddings/huggingface.py +++ b/src/dynavec/embeddings/huggingface.py @@ -30,9 +30,7 @@ def __init__( ) -> None: """Initialize HFInferenceEmbedder.""" self.model_name = model_name - self.api_key = ( - api_key or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN") - ) + self.api_key = api_key or os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN") self.batch_size = batch_size self.dimension = dimension or _MODEL_DIMS.get(model_name, 384) @@ -41,7 +39,9 @@ def __init__( "Hugging Face API token is required. Pass api_key or set HF_TOKEN environment variable." ) - self.api_url = f"https://api-inference.huggingface.co/pipeline/feature-extraction/{self.model_name}" + self.api_url = ( + f"https://api-inference.huggingface.co/pipeline/feature-extraction/{self.model_name}" + ) self.headers = {"Authorization": f"Bearer {self.api_key}"} def embed_documents(self, texts: Sequence[str]) -> list[Vector]: @@ -62,13 +62,9 @@ def embed_documents(self, texts: Sequence[str]) -> list[Vector]: if response.status_code != 200: if response.status_code in (401, 403): - raise ValueError( - "Invalid Hugging Face API token or unauthorized request." - ) + raise ValueError("Invalid Hugging Face API token or unauthorized request.") elif response.status_code == 404: - raise ValueError( - f"Model '{self.model_name}' not found on Hugging Face Hub." - ) + raise ValueError(f"Model '{self.model_name}' not found on Hugging Face Hub.") else: raise RuntimeError( f"Hugging Face API request failed ({response.status_code}): {response.text}" diff --git a/src/dynavec/embeddings/openai.py b/src/dynavec/embeddings/openai.py index 05ad713..8360bd8 100644 --- a/src/dynavec/embeddings/openai.py +++ b/src/dynavec/embeddings/openai.py @@ -18,13 +18,13 @@ "text-embedding-ada-002": 1536, } + def _is_openai_retryable(exc: Exception) -> bool: """Return True for OpenAI rate-limit and server errors.""" status_code = getattr(exc, "status_code", None) - return status_code == 429 or ( - isinstance(status_code, int) and 500 <= status_code < 600 - ) + return status_code == 429 or (isinstance(status_code, int) and 500 <= status_code < 600) + def _openai_retry_after(exc: Exception) -> float | None: """Return the server-provided Retry-After delay, if available.""" @@ -88,7 +88,6 @@ def embed_documents(self, texts: list[str]) -> list[Vector]: out.extend(d.embedding for d in resp.data) return out - @retry( retry_on=_is_openai_retryable, retry_delay=_openai_retry_after, @@ -98,6 +97,4 @@ def _create_embeddings( ) -> CreateEmbeddingResponse: if dimensions is None: return self._client.embeddings.create(model=model, input=input) - return self._client.embeddings.create( - model=model, input=input, dimensions=dimensions - ) + return self._client.embeddings.create(model=model, input=input, dimensions=dimensions) diff --git a/src/dynavec/eval/chart.py b/src/dynavec/eval/chart.py index 32b6140..8772436 100644 --- a/src/dynavec/eval/chart.py +++ b/src/dynavec/eval/chart.py @@ -229,5 +229,3 @@ def plot_retrieval_metrics( plt.close(fig) return output_path - - diff --git a/src/dynavec/eval/trend.py b/src/dynavec/eval/trend.py index 5d91ab0..5ad92c8 100644 --- a/src/dynavec/eval/trend.py +++ b/src/dynavec/eval/trend.py @@ -33,9 +33,9 @@ class EvalRun: """A recorded eval run with metadata tags and metric snapshots.""" run_id: str - ts: float # epoch seconds when the run was recorded - tags: dict[str, str] # e.g. {"dataset": "squad", "model": "ada-002"} - metrics: dict[str, Any] # flat metric snapshot (recall@k, ndcg@k, faithfulness, …) + ts: float # epoch seconds when the run was recorded + tags: dict[str, str] # e.g. {"dataset": "squad", "model": "ada-002"} + metrics: dict[str, Any] # flat metric snapshot (recall@k, ndcg@k, faithfulness, …) def to_dict(self) -> dict[str, Any]: return { @@ -121,7 +121,7 @@ def record( ) self._runs.append(run) if len(self._runs) > self._max_runs: - self._runs = self._runs[-self._max_runs:] + self._runs = self._runs[-self._max_runs :] return run # ------------------------------------------------------------------ filter diff --git a/src/dynavec/exceptions.py b/src/dynavec/exceptions.py index 869cec7..35beffc 100644 --- a/src/dynavec/exceptions.py +++ b/src/dynavec/exceptions.py @@ -70,8 +70,6 @@ def __init__(self, feature: str, package: str, extra: str) -> None: self.extra = extra - - class NodeInterrupt(DynavecError): """Raised when an agent execution node pauses waiting for human approval or input.""" diff --git a/src/dynavec/fusion.py b/src/dynavec/fusion.py index 5e73b99..b17d63b 100644 --- a/src/dynavec/fusion.py +++ b/src/dynavec/fusion.py @@ -35,6 +35,7 @@ # FitResult # --------------------------------------------------------------------------- + @dataclass class FitResult: """Result of an ``RRFWeightFitter.fit()`` call. @@ -85,6 +86,7 @@ def load(cls, path: str | Path) -> FitResult: # RRFWeightFitter # --------------------------------------------------------------------------- + class RRFWeightFitter: """Learn per-retriever RRF weights from labeled evaluation queries. @@ -125,10 +127,7 @@ def score(self, weights: list[float]) -> float: for query_idx, labeled_query in enumerate(self.labeled_queries): # result_lists[retriever_idx][query_idx] → list[SearchResult] - per_retriever = [ - self.result_lists[r][query_idx] - for r in range(len(self.result_lists)) - ] + per_retriever = [self.result_lists[r][query_idx] for r in range(len(self.result_lists))] fused = reciprocal_rank_fusion(per_retriever, k=self.k, weights=weights) ranked_ids = [r.id for r in fused] @@ -241,7 +240,9 @@ def _neg_ndcg(x: np.ndarray) -> float: best_w = _softmax(result.x) best_score = -float(result.fun) - return FitResult(weights=best_w, score=best_score, method="bayesian", n_evaluations=eval_count) + return FitResult( + weights=best_w, score=best_score, method="bayesian", n_evaluations=eval_count + ) # -- public API ---------------------------------------------------------- diff --git a/src/dynavec/graph.py b/src/dynavec/graph.py index 37d5e86..8f27d63 100644 --- a/src/dynavec/graph.py +++ b/src/dynavec/graph.py @@ -50,6 +50,7 @@ def _edges_changed_or_retryable(exc: Exception) -> bool: code = getattr(exc, "response", {}).get("Error", {}).get("Code", "") return code == "ConditionalCheckFailedException" or is_retryable(exc) + EXPORT_FORMATS = ("mermaid", "dot") # A bare Mermaid flowchart node id: letters/digits/underscore, not keyword-shaped. @@ -100,9 +101,7 @@ def _mermaid_aliases(nodes: list[str]) -> dict[str, str]: reserved keyword) gets an `nN` alias carrying the real id as its label. Alias numbering follows sorted id order, so it is stable across exports. """ - bare = { - n for n in nodes if _MERMAID_SAFE_ID.match(n) and n.lower() not in _MERMAID_RESERVED - } + bare = {n for n in nodes if _MERMAID_SAFE_ID.match(n) and n.lower() not in _MERMAID_RESERVED} aliases: dict[str, str] = {} counter = 0 for node in nodes: @@ -308,8 +307,7 @@ def _drop_edges(self, ns: str, src: str, dst: str, relation: str | None = None) edge for edge in edges if not ( - edge.get("target") == dst - and (relation is None or edge.get("relation") == relation) + edge.get("target") == dst and (relation is None or edge.get("relation") == relation) ) ] if len(kept) == len(edges): diff --git a/src/dynavec/integrations/dspy.py b/src/dynavec/integrations/dspy.py index 821ed83..041d819 100644 --- a/src/dynavec/integrations/dspy.py +++ b/src/dynavec/integrations/dspy.py @@ -1,13 +1,13 @@ """DSPy retrieval module backed by dynavec. - import dspy - from dynavec.integrations.dspy import DynavecRM +import dspy +from dynavec.integrations.dspy import DynavecRM - rm = DynavecRM(dynavec_client, namespace="kb", k=4) - dspy.configure(rm=rm) +rm = DynavecRM(dynavec_client, namespace="kb", k=4) +dspy.configure(rm=rm) - retriever = dspy.Retrieve(k=4) - passages = retriever("what is retrieval-augmented generation?").passages +retriever = dspy.Retrieve(k=4) +passages = retriever("what is retrieval-augmented generation?").passages """ from __future__ import annotations diff --git a/src/dynavec/integrations/haystack.py b/src/dynavec/integrations/haystack.py index 247ec4e..7846e07 100644 --- a/src/dynavec/integrations/haystack.py +++ b/src/dynavec/integrations/haystack.py @@ -1,4 +1,5 @@ """Haystack integration for Dynavec.""" + from __future__ import annotations from dataclasses import asdict @@ -38,8 +39,11 @@ def _convert_filter( return {field: filters["value"]} return filters + + class DynavecDocumentStore: """Haystack DocumentStore backed by Dynavec.""" + def __init__(self, client: Dynavec, namespace: str = "default") -> None: self.client = client self.namespace = namespace @@ -53,6 +57,7 @@ def count_documents(self) -> int: hydrate=False, ) ) + def write_documents( self, documents: list[Document], @@ -77,9 +82,7 @@ def write_documents( elif policy == DuplicatePolicy.FAIL: duplicates = [document.id for document in documents if document.id in existing] if duplicates: - raise DuplicateDocumentError( - f"Documents with IDs already exist: {duplicates}" - ) + raise DuplicateDocumentError(f"Documents with IDs already exist: {duplicates}") dynavec_documents = [ DynavecDocument( @@ -111,8 +114,7 @@ def filter_documents( namespace=self.namespace, hydrate=True, ): - if filters and any(result.metadata.get(key) != value for key, value in filters.items() - ): + if filters and any(result.metadata.get(key) != value for key, value in filters.items()): continue documents.append( @@ -134,8 +136,7 @@ def to_dict(self) -> dict[str, Any]: "config": asdict(self.client.config), "namespace": self.namespace, }, - } - + } @classmethod def from_dict(cls, data: dict[str, Any]) -> DynavecDocumentStore: @@ -151,7 +152,6 @@ def from_dict(cls, data: dict[str, Any]) -> DynavecDocumentStore: ) - class DynavecRetriever: """Haystack retriever backed by Dynavec.""" diff --git a/src/dynavec/integrations/langchain.py b/src/dynavec/integrations/langchain.py index bafaf2e..96f7cbc 100644 --- a/src/dynavec/integrations/langchain.py +++ b/src/dynavec/integrations/langchain.py @@ -70,10 +70,7 @@ def add_texts( texts = list(texts) ids = ids or [str(uuid.uuid4()) for _ in texts] metadatas = metadatas or [{} for _ in texts] - docs = [ - DVDocument(id=i, text=t, metadata=m) - for i, t, m in zip(ids, texts, metadatas) - ] + docs = [DVDocument(id=i, text=t, metadata=m) for i, t, m in zip(ids, texts, metadatas)] self._client.upsert(docs, namespace=self._namespace) return ids @@ -86,11 +83,11 @@ def delete(self, ids: list[str] | None = None, **kwargs: Any) -> bool | None: def similarity_search( self, query: str, k: int = 4, filter: dict[str, Any] | None = None, **kwargs: Any ) -> list[LCDocument]: - results = self._client.search( - query, top_k=k, namespace=self._namespace, filter=filter - ) + results = self._client.search(query, top_k=k, namespace=self._namespace, filter=filter) return [ - LCDocument(page_content=r.text or "", metadata={**r.metadata, "id": r.id, "score": r.score}) + LCDocument( + page_content=r.text or "", metadata={**r.metadata, "id": r.id, "score": r.score} + ) for r in results ] @@ -119,9 +116,7 @@ async def asimilarity_search( def similarity_search_with_score( self, query: str, k: int = 4, filter: dict[str, Any] | None = None, **kwargs: Any ) -> list[tuple[LCDocument, float]]: - results = self._client.search( - query, top_k=k, namespace=self._namespace, filter=filter - ) + results = self._client.search(query, top_k=k, namespace=self._namespace, filter=filter) return [ ( LCDocument(page_content=r.text or "", metadata={**r.metadata, "id": r.id}), @@ -173,7 +168,9 @@ def max_marginal_relevance_search( mmr_lambda=lambda_mult, ) return [ - LCDocument(page_content=r.text or "", metadata={**r.metadata, "id": r.id, "score": r.score}) + LCDocument( + page_content=r.text or "", metadata={**r.metadata, "id": r.id, "score": r.score} + ) for r in results ] diff --git a/src/dynavec/integrations/llamaindex.py b/src/dynavec/integrations/llamaindex.py index 8443ff6..55f9d14 100644 --- a/src/dynavec/integrations/llamaindex.py +++ b/src/dynavec/integrations/llamaindex.py @@ -86,8 +86,7 @@ def _translate_filters(metadata_filters: MetadataFilters) -> dict[str, Any]: return {"$or": translated} raise ValueError( - "Unsupported LlamaIndex metadata filter condition for S3 Vectors: " - f"{condition.value}" + f"Unsupported LlamaIndex metadata filter condition for S3 Vectors: {condition.value}" ) diff --git a/src/dynavec/mcp/__init__.py b/src/dynavec/mcp/__init__.py index 0df5c80..fe39777 100644 --- a/src/dynavec/mcp/__init__.py +++ b/src/dynavec/mcp/__init__.py @@ -9,4 +9,3 @@ from .server import client_from_env, create_mcp_server, main __all__ = ["client_from_env", "create_mcp_server", "main"] - diff --git a/src/dynavec/mcp/server.py b/src/dynavec/mcp/server.py index af2d9df..e7eedd6 100644 --- a/src/dynavec/mcp/server.py +++ b/src/dynavec/mcp/server.py @@ -37,7 +37,9 @@ def _resolve_embedder(env: Mapping[str, str]) -> Embedder | None: if embedder_type == "openai": from ..embeddings.openai import OpenAIEmbedder - return OpenAIEmbedder(model=model or "text-embedding-3-small", api_key=env.get("OPENAI_API_KEY")) + return OpenAIEmbedder( + model=model or "text-embedding-3-small", api_key=env.get("OPENAI_API_KEY") + ) if embedder_type == "gemini": from ..embeddings.gemini import GeminiEmbedder @@ -221,7 +223,9 @@ def dynavec_graph_search( def main(argv: list[str] | None = None) -> int: """Run the FastMCP server with stdio or SSE transport.""" - parser = argparse.ArgumentParser(prog="dynavec mcp", description="Run the dynavec FastMCP server.") + parser = argparse.ArgumentParser( + prog="dynavec mcp", description="Run the dynavec FastMCP server." + ) parser.add_argument( "--transport", choices=["stdio", "sse"], diff --git a/src/dynavec/models.py b/src/dynavec/models.py index 3d34bb0..db29545 100644 --- a/src/dynavec/models.py +++ b/src/dynavec/models.py @@ -25,9 +25,7 @@ class Document: def __post_init__(self) -> None: if self.text is None and self.vector is None: - raise ValueError( - f"Document {self.id!r} must have either 'text' or 'vector' set." - ) + raise ValueError(f"Document {self.id!r} must have either 'text' or 'vector' set.") if self.ttl_seconds is not None and self.ttl_seconds <= 0: raise ValueError( f"Document {self.id!r} ttl_seconds must be positive, got {self.ttl_seconds}." diff --git a/src/dynavec/namespace.py b/src/dynavec/namespace.py index 45d6c6e..7fa91ac 100644 --- a/src/dynavec/namespace.py +++ b/src/dynavec/namespace.py @@ -38,9 +38,7 @@ def upsert( ttl_seconds: int | None = None, **kw: Any, ) -> UpsertResult: - return self._db.upsert( - documents, namespace=self._ns, ttl_seconds=ttl_seconds, **kw - ) + return self._db.upsert(documents, namespace=self._ns, ttl_seconds=ttl_seconds, **kw) def update( self, diff --git a/src/dynavec/quantization.py b/src/dynavec/quantization.py index 2ba26a1..87105f6 100644 --- a/src/dynavec/quantization.py +++ b/src/dynavec/quantization.py @@ -264,9 +264,7 @@ def encode(self, vectors: np.ndarray) -> np.ndarray: if vectors.ndim != 2: raise ValueError("vectors must be a 2D array") - codes = np.round( - (vectors - mins) / scales - 128 - ) + codes = np.round((vectors - mins) / scales - 128) return np.asarray(np.clip(codes, -128, 127), dtype=np.int8) @@ -276,8 +274,7 @@ def decode(self, codes: np.ndarray) -> np.ndarray: codes = np.asarray(codes, dtype=np.int8) return np.asarray( - (codes.astype(np.float32) + 128) * scales - + mins, + (codes.astype(np.float32) + 128) * scales + mins, dtype=np.float32, ) @@ -315,15 +312,11 @@ def fit(self, vectors: np.ndarray) -> OPQRotation: raise ValueError("vectors must be a 2D array") if x.shape[1] != self.dim: - raise ValueError( - f"Expected dimension={self.dim}, got {x.shape[1]}" - ) + raise ValueError(f"Expected dimension={self.dim}, got {x.shape[1]}") rng = np.random.default_rng(self.seed) - random_matrix = rng.normal( - size=(self.dim, self.dim) - ).astype(np.float32) + random_matrix = rng.normal(size=(self.dim, self.dim)).astype(np.float32) q, _ = np.linalg.qr(random_matrix) @@ -349,12 +342,9 @@ def inverse_transform(self, vectors: np.ndarray) -> np.ndarray: def _fitted_rotation(self) -> np.ndarray: if self._rotation is None: - raise RuntimeError( - "OPQRotation must be .fit() before use" - ) + raise RuntimeError("OPQRotation must be .fit() before use") return self._rotation - def _update_rotation( self, original: np.ndarray, @@ -402,9 +392,7 @@ def fit(self, vectors: np.ndarray) -> OptimizedProductQuantizer: self._dim = x.shape[1] if self._dim % self.m != 0: - raise ValueError( - f"m={self.m} must divide dimension={self._dim}" - ) + raise ValueError(f"m={self.m} must divide dimension={self._dim}") self._opq = OPQRotation( dim=self._dim, @@ -428,11 +416,7 @@ def fit(self, vectors: np.ndarray) -> OptimizedProductQuantizer: codes = self._pq.encode(rotated) reconstructed = self._pq.decode(codes) - error = float( - ((rotated - reconstructed) ** 2) - .sum(axis=1) - .mean() - ) + error = float(((rotated - reconstructed) ** 2).sum(axis=1).mean()) self._training_errors.append(error) self._opq._update_rotation( @@ -464,9 +448,7 @@ def asymmetric_distances( ) -> np.ndarray: pq, opq = self._fitted_components() - rotated_query = opq.transform( - np.asarray(query, dtype=np.float32) - ) + rotated_query = opq.transform(np.asarray(query, dtype=np.float32)) return pq.asymmetric_distances( rotated_query, @@ -476,15 +458,9 @@ def asymmetric_distances( def reconstruction_error(self, vectors: np.ndarray) -> float: x = np.asarray(vectors, dtype=np.float32) - reconstructed = self.decode( - self.encode(x) - ) + reconstructed = self.decode(self.encode(x)) - return float( - ((x - reconstructed) ** 2) - .sum(axis=1) - .mean() - ) + return float(((x - reconstructed) ** 2).sum(axis=1).mean()) @property def training_errors(self) -> list[float]: @@ -493,7 +469,5 @@ def training_errors(self) -> list[float]: def _fitted_components(self) -> tuple[ProductQuantizer, OPQRotation]: if self._pq is None or self._opq is None: - raise RuntimeError( - "OptimizedProductQuantizer must be .fit() before use" - ) + raise RuntimeError("OptimizedProductQuantizer must be .fit() before use") return self._pq, self._opq diff --git a/src/dynavec/retrievers.py b/src/dynavec/retrievers.py index 79005cb..6cf6209 100644 --- a/src/dynavec/retrievers.py +++ b/src/dynavec/retrievers.py @@ -600,7 +600,9 @@ def __init__( if hasattr(weights, "weights"): weights = weights.weights if len(weights) != 2: - raise ValueError("weights must contain exactly 2 elements [dense_weight, sparse_weight]") + raise ValueError( + "weights must contain exactly 2 elements [dense_weight, sparse_weight]" + ) dense_weight, sparse_weight = float(weights[0]), float(weights[1]) if dense_weight <= 0 or sparse_weight <= 0: @@ -668,7 +670,9 @@ def search( if hasattr(weights, "weights"): weights = weights.weights if len(weights) != 2: - raise ValueError("weights must contain exactly 2 elements [dense_weight, sparse_weight]") + raise ValueError( + "weights must contain exactly 2 elements [dense_weight, sparse_weight]" + ) w_dense, w_sparse = float(weights[0]), float(weights[1]) # Run dense vector search and sparse BM25 search in parallel diff --git a/src/dynavec/stores/async_dynamodb.py b/src/dynavec/stores/async_dynamodb.py index 496cdcb..4469c39 100644 --- a/src/dynavec/stores/async_dynamodb.py +++ b/src/dynavec/stores/async_dynamodb.py @@ -124,8 +124,7 @@ async def put_many( self, namespace: str, items: Sequence[ - tuple[str, str | None, Metadata] - | tuple[str, str | None, Metadata, int | None] + tuple[str, str | None, Metadata] | tuple[str, str | None, Metadata, int | None] ], ) -> None: """Upsert document tuples using DynamoDB's async batch writer.""" @@ -195,10 +194,7 @@ async def get_many( ddb = self._require_ddb() - keys = [ - {"pk": self._pk(namespace, doc_id)} - for doc_id in ids - ] + keys = [{"pk": self._pk(namespace, doc_id)} for doc_id in ids] out: dict[str, dict[str, Any]] = {} @@ -222,9 +218,7 @@ async def get_many( ): entry: dict[str, Any] = { "text": _read_text(item), - "metadata": _from_dynamo( - item.get("metadata", {}) - ), + "metadata": _from_dynamo(item.get("metadata", {})), } ttl_attr = self._config.dynamodb_ttl_attribute @@ -233,15 +227,9 @@ async def get_many( out[item["id"]] = entry - unprocessed = ( - response.get("UnprocessedKeys") or {} - ) + unprocessed = response.get("UnprocessedKeys") or {} - request = ( - unprocessed - if unprocessed - else None - ) + request = unprocessed if unprocessed else None log_store_event( self._logger, diff --git a/src/dynavec/stores/async_s3vectors.py b/src/dynavec/stores/async_s3vectors.py index d963ded..5744d5f 100644 --- a/src/dynavec/stores/async_s3vectors.py +++ b/src/dynavec/stores/async_s3vectors.py @@ -23,15 +23,9 @@ class AsyncS3VectorsStore: def __init__(self, config: DynavecConfig, boto_session: Any) -> None: self._config = config self._session = boto_session - self._put_limiter = ( - TokenBucket(config.put_rps) - if config.put_rps is not None - else None - ) + self._put_limiter = TokenBucket(config.put_rps) if config.put_rps is not None else None self._query_limiter = ( - TokenBucket(config.query_rps) - if config.query_rps is not None - else None + TokenBucket(config.query_rps) if config.query_rps is not None else None ) self._client_context: Any | None = None self._client: Any | None = None @@ -117,10 +111,7 @@ async def put_vectors( raise ValueError("max_workers must be greater than 0.") t0 = time.perf_counter() - chunks = [ - vectors[i : i + _PUT_LIMIT] - for i in range(0, len(vectors), _PUT_LIMIT) - ] + chunks = [vectors[i : i + _PUT_LIMIT] for i in range(0, len(vectors), _PUT_LIMIT)] semaphore = asyncio.Semaphore(min(max_workers, len(chunks))) @@ -139,10 +130,7 @@ async def upload_chunk( async with semaphore: await self._put_batch(payload) - tasks = [ - asyncio.create_task(upload_chunk(chunk)) - for chunk in chunks - ] + tasks = [asyncio.create_task(upload_chunk(chunk)) for chunk in chunks] try: await asyncio.gather(*tasks) @@ -164,7 +152,6 @@ async def upload_chunk( duration_ms=round((time.perf_counter() - t0) * 1000, 2), ) - def _query_kwargs( self, query_vector: list[float], @@ -213,8 +200,7 @@ async def query( if top_k > _MAX_TOP_K: raise ValueError( - f"top_k ({top_k}) exceeds Amazon S3 Vectors " - f"maximum limit of {_MAX_TOP_K}." + f"top_k ({top_k}) exceeds Amazon S3 Vectors maximum limit of {_MAX_TOP_K}." ) results: list[dict[str, Any]] = [] @@ -261,15 +247,10 @@ async def query_pages( if top_k > _MAX_TOP_K: raise ValueError( - f"top_k ({top_k}) exceeds Amazon S3 Vectors " - f"maximum limit of {_MAX_TOP_K}." + f"top_k ({top_k}) exceeds Amazon S3 Vectors maximum limit of {_MAX_TOP_K}." ) - effective_page_size = ( - page_size - if page_size is not None - else self._config.top_k_page_size - ) + effective_page_size = page_size if page_size is not None else self._config.top_k_page_size if effective_page_size is not None and effective_page_size <= 0: raise ValueError("page_size must be a positive integer.") @@ -320,10 +301,7 @@ async def query_pages( else: buffer.extend(vectors) - while ( - len(buffer) >= effective_page_size - and yielded < top_k - ): + while len(buffer) >= effective_page_size and yielded < top_k: chunk = buffer[:effective_page_size] buffer = buffer[effective_page_size:] @@ -341,15 +319,10 @@ async def query_pages( if page.get("NextToken") is None: break - if ( - effective_page_size is not None - and buffer - and yielded < top_k - ): + if effective_page_size is not None and buffer and yielded < top_k: remaining = top_k - yielded yield buffer[:remaining] - @async_retry() async def get_vectors( self, diff --git a/src/dynavec/stores/s3vectors.py b/src/dynavec/stores/s3vectors.py index 60d1053..9e418a6 100644 --- a/src/dynavec/stores/s3vectors.py +++ b/src/dynavec/stores/s3vectors.py @@ -45,9 +45,7 @@ def __init__(self, config: DynavecConfig, boto_session: Any | None = None) -> No session = boto_session or boto3.Session() self._config = config - self._put_limiter = ( - TokenBucket(config.put_rps) if config.put_rps is not None else None - ) + self._put_limiter = TokenBucket(config.put_rps) if config.put_rps is not None else None self._query_limiter = ( TokenBucket(config.query_rps) if config.query_rps is not None else None ) @@ -75,7 +73,6 @@ def _put_batch(self, payload: list[dict[str, Any]]) -> None: vectors=payload, ) - def put_vectors( self, vectors: list[tuple[str, list[float], Metadata]], @@ -212,9 +209,7 @@ def query_pages( f"top_k ({top_k}) exceeds Amazon S3 Vectors maximum limit of {_MAX_TOP_K}." ) - effective_page_size = ( - page_size if page_size is not None else self._config.top_k_page_size - ) + effective_page_size = page_size if page_size is not None else self._config.top_k_page_size if effective_page_size is not None and effective_page_size <= 0: raise ValueError("page_size must be a positive integer.") @@ -308,9 +303,7 @@ def list_pages( leaves the page size to Amazon S3 Vectors. """ if page_size is not None and not _LIST_PAGE_MIN <= page_size <= _LIST_PAGE_MAX: - raise ValueError( - f"page_size must be between {_LIST_PAGE_MIN} and {_LIST_PAGE_MAX}." - ) + raise ValueError(f"page_size must be between {_LIST_PAGE_MIN} and {_LIST_PAGE_MAX}.") kwargs: dict[str, Any] = { "vectorBucketName": self._config.vector_bucket, diff --git a/src/dynavec/utils.py b/src/dynavec/utils.py index 70e8715..9d5199f 100644 --- a/src/dynavec/utils.py +++ b/src/dynavec/utils.py @@ -28,6 +28,7 @@ def decode_key_component(value: str) -> str: """Restore a component produced by :func:`encode_key_component`.""" return value.replace("%23", KEY_SEPARATOR).replace("%25", "%") + # botocore error codes that are safe to retry (throttling / transient). _RETRYABLE_CODES = frozenset( { @@ -181,7 +182,6 @@ def __init__( self.last_time = time.monotonic() self._lock = threading.Lock() - def _acquire_wait_time(self) -> float | None: with self._lock: now = time.monotonic() diff --git a/tests/integration/test_live_aws.py b/tests/integration/test_live_aws.py index bee083c..2d0f79c 100644 --- a/tests/integration/test_live_aws.py +++ b/tests/integration/test_live_aws.py @@ -97,7 +97,7 @@ def test_provision_upsert_search_roundtrip(live_db): hits = _query_with_retry(live_db, vecs[0].tolist(), top_k=5) assert hits, "no results returned from live S3 Vectors index" assert hits[0].id == "v0" - assert hits[0].text == "document number 0" # hydrated from DynamoDB + assert hits[0].text == "document number 0" # hydrated from DynamoDB # metadata pre-filter must scope results even_hits = _query_with_retry(live_db, vecs[0].tolist(), top_k=10, filter={"even": True}) @@ -110,7 +110,9 @@ def test_update_and_delete(live_db): rng = np.random.default_rng(1) v = _unit(rng, 1, DIM)[0].tolist() - live_db.upsert([Document(id="u1", vector=v, text="original", metadata={"tag": "a"})], namespace="it") + live_db.upsert( + [Document(id="u1", vector=v, text="original", metadata={"tag": "a"})], namespace="it" + ) live_db.update("u1", namespace="it", metadata={"tag": "b"}, merge_metadata=False) got = live_db.get(["u1"], namespace="it")[0] diff --git a/tests/test_agents.py b/tests/test_agents.py index c560bcc..5ff3ec4 100644 --- a/tests/test_agents.py +++ b/tests/test_agents.py @@ -43,9 +43,7 @@ def invoke( if self.call_count >= len(self.responses): # Default response if script exhausted return ChatResult( - message=Message( - role="assistant", content="Scripted responses exhausted." - ), + message=Message(role="assistant", content="Scripted responses exhausted."), finish_reason="stop", ) diff --git a/tests/test_async_client.py b/tests/test_async_client.py index c625910..3700406 100644 --- a/tests/test_async_client.py +++ b/tests/test_async_client.py @@ -38,9 +38,7 @@ def embed_query( self, text: str, ) -> list[float]: - raise AssertionError( - "AsyncDynavec must not call sync embed_query()." - ) + raise AssertionError("AsyncDynavec must not call sync embed_query().") async def aembed_query( self, @@ -53,9 +51,7 @@ def embed_documents( self, texts: list[str], ) -> list[list[float]]: - raise AssertionError( - "AsyncDynavec must not call sync embed_documents()." - ) + raise AssertionError("AsyncDynavec must not call sync embed_documents().") async def aembed_documents( self, @@ -63,19 +59,14 @@ async def aembed_documents( ) -> list[list[float]]: self.async_calls.append(texts) - return [ - [1.0, 2.0, 3.0, 4.0] - for _ in texts - ] + return [[1.0, 2.0, 3.0, 4.0] for _ in texts] class FakeVectorStore: def __init__(self) -> None: self.entered = False self.exited = False - self.put_calls: list[ - tuple[list[tuple[str, list[float], dict[str, Any]]], int] - ] = [] + self.put_calls: list[tuple[list[tuple[str, list[float], dict[str, Any]]], int]] = [] self.query_results: list[dict[str, Any]] = [] self.query_calls: list[dict[str, Any]] = [] self.query_page_results: list[list[dict[str, Any]]] = [] @@ -115,14 +106,10 @@ async def __aexit__( async def put_vectors( self, - vectors: list[ - tuple[str, list[float], dict[str, Any]] - ], + vectors: list[tuple[str, list[float], dict[str, Any]]], max_workers: int = 8, ) -> None: - self.put_calls.append( - (vectors, max_workers) - ) + self.put_calls.append((vectors, max_workers)) async def query_pages( self, @@ -179,11 +166,7 @@ async def get_many( ) ) - return { - doc_id: self.documents[doc_id] - for doc_id in ids - if doc_id in self.documents - } + return {doc_id: self.documents[doc_id] for doc_id in ids if doc_id in self.documents} async def __aenter__( self, @@ -210,9 +193,7 @@ async def put_many( ] ], ) -> None: - self.put_calls.append( - (namespace, items) - ) + self.put_calls.append((namespace, items)) def _install_fake_stores( @@ -299,9 +280,7 @@ async def test_aupsert_requires_open_client() -> None: async def test_aupsert_uses_async_embedding_and_writes_both_stores() -> None: embedder = AsyncOnlyEmbedder() - client, vectors, documents = _client( - embedder=embedder - ) + client, vectors, documents = _client(embedder=embedder) async with client: result = await client.aupsert( @@ -320,9 +299,7 @@ async def test_aupsert_uses_async_embedding_and_writes_both_stores() -> None: assert result.count == 1 assert result.ids == ["a"] - assert embedder.async_calls == [ - ["hello"] - ] + assert embedder.async_calls == [["hello"]] assert len(vectors.put_calls) == 1 assert len(documents.put_calls) == 1 @@ -337,9 +314,7 @@ async def test_aupsert_uses_async_embedding_and_writes_both_stores() -> None: 4.0, ] - namespace, document_payload = ( - documents.put_calls[0] - ) + namespace, document_payload = documents.put_calls[0] assert namespace == "tenant" assert document_payload[0][0] == "a" @@ -422,9 +397,7 @@ async def test_aupsert_runs_s3_and_dynamodb_writes_concurrently() -> None: s3_started = asyncio.Event() ddb_started = asyncio.Event() - class CoordinatedVectors( - FakeVectorStore - ): + class CoordinatedVectors(FakeVectorStore): async def put_vectors( self, vectors: list[ @@ -443,9 +416,7 @@ async def put_vectors( timeout=1, ) - class CoordinatedDocuments( - FakeDocumentStore - ): + class CoordinatedDocuments(FakeDocumentStore): async def put_many( self, namespace: str, @@ -503,15 +474,11 @@ async def put_many( async def test_partial_enter_failure_closes_first_store() -> None: vectors = FakeVectorStore() - class FailingDocumentStore( - FakeDocumentStore - ): + class FailingDocumentStore(FakeDocumentStore): async def __aenter__( self, ) -> FailingDocumentStore: - raise RuntimeError( - "DynamoDB open failed" - ) + raise RuntimeError("DynamoDB open failed") client = AsyncDynavec( _config(), diff --git a/tests/test_async_dynamodb.py b/tests/test_async_dynamodb.py index 46b5614..3cb2725 100644 --- a/tests/test_async_dynamodb.py +++ b/tests/test_async_dynamodb.py @@ -115,9 +115,7 @@ def resource( **kwargs: Any, ) -> FakeResourceContext: self.calls.append((service_name, kwargs)) - self.context = FakeResourceContext( - self.resource_instance - ) + self.context = FakeResourceContext(self.resource_instance) return self.context @@ -150,18 +148,14 @@ async def test_put_many_builds_and_writes_items(): ], ) - assert table.batch_writer_calls == [ - {"overwrite_by_pkeys": ["pk"]} - ] + assert table.batch_writer_calls == [{"overwrite_by_pkeys": ["pk"]}] assert len(table.writer.items) == 2 assert table.writer.items[0]["pk"] == "tenant#a" assert table.writer.items[0]["id"] == "a" assert table.writer.items[0]["text"] == "First" - assert table.writer.items[0]["metadata"]["score"] == Decimal( - "0.5" - ) + assert table.writer.items[0]["metadata"]["score"] == Decimal("0.5") async def test_put_many_writes_custom_ttl_attribute(): @@ -342,12 +336,7 @@ async def test_get_many_batches_at_100_keys(): async with AsyncDynamoDBStore(_config(), session) as store: await store.get_many("tenant", ids) - sizes = [ - len( - call["RequestItems"]["docs"]["Keys"] - ) - for call in resource.batch_get_calls - ] + sizes = [len(call["RequestItems"]["docs"]["Keys"]) for call in resource.batch_get_calls] assert sizes == [100, 100, 50] @@ -355,9 +344,7 @@ async def test_get_many_batches_at_100_keys(): async def test_operations_require_open_store(): store = AsyncDynamoDBStore( _config(), - FakeSession( - FakeDynamoResource(FakeTable()) - ), + FakeSession(FakeDynamoResource(FakeTable())), ) with pytest.raises(RuntimeError, match="not open"): diff --git a/tests/test_async_s3vectors.py b/tests/test_async_s3vectors.py index 9a46cba..2316fd2 100644 --- a/tests/test_async_s3vectors.py +++ b/tests/test_async_s3vectors.py @@ -112,10 +112,7 @@ async def test_put_vectors_batches_at_service_limit(): client = FakeS3Client() session = FakeSession(client) - vectors = [ - (f"key-{i}", [0.1, 0.2, 0.3, 0.4], {"tag": "test"}) - for i in range(1200) - ] + vectors = [(f"key-{i}", [0.1, 0.2, 0.3, 0.4], {"tag": "test"}) for i in range(1200)] async with AsyncS3VectorsStore(_config(), session) as store: await store.put_vectors(vectors) @@ -146,10 +143,7 @@ async def fake_put_batch(payload): monkeypatch.setattr(store, "_put_batch", fake_put_batch) - vectors = [ - (f"key-{i}", [0.1, 0.2, 0.3, 0.4], {}) - for i in range(1200) - ] + vectors = [(f"key-{i}", [0.1, 0.2, 0.3, 0.4], {}) for i in range(1200)] await store.put_vectors(vectors) @@ -167,35 +161,23 @@ async def fail(payload): monkeypatch.setattr(store, "_put_batch", fail) with pytest.raises(RuntimeError, match="S3 API failure"): - await store.put_vectors( - [("key", [0.1, 0.2, 0.3, 0.4], {})] - ) + await store.put_vectors([("key", [0.1, 0.2, 0.3, 0.4], {})]) async def test_put_vectors_requires_open_store(): store = AsyncS3VectorsStore(_config(), FakeSession(FakeS3Client())) with pytest.raises(RuntimeError, match="not open"): - await store.put_vectors( - [("key", [0.1, 0.2, 0.3, 0.4], {})] - ) + await store.put_vectors([("key", [0.1, 0.2, 0.3, 0.4], {})]) async def test_query_returns_paginated_results(): pages = [ { - "vectors": [ - {"key": f"k{i}", "distance": i * 0.01} - for i in range(100) - ], + "vectors": [{"key": f"k{i}", "distance": i * 0.01} for i in range(100)], "NextToken": "token-1", }, - { - "vectors": [ - {"key": f"k{i}", "distance": i * 0.01} - for i in range(100, 150) - ] - }, + {"vectors": [{"key": f"k{i}", "distance": i * 0.01} for i in range(100, 150)]}, ] client = FakeS3Client(pages) @@ -215,18 +197,10 @@ async def test_query_returns_paginated_results(): async def test_query_truncates_at_top_k(): pages = [ { - "vectors": [ - {"key": f"k{i}"} - for i in range(100) - ], + "vectors": [{"key": f"k{i}"} for i in range(100)], "NextToken": "token-1", }, - { - "vectors": [ - {"key": f"k{i}"} - for i in range(100, 200) - ] - }, + {"vectors": [{"key": f"k{i}"} for i in range(100, 200)]}, ] client = FakeS3Client(pages) @@ -243,14 +217,7 @@ async def test_query_truncates_at_top_k(): async def test_query_pages_respects_page_size(): - pages = [ - { - "vectors": [ - {"key": f"k{i}"} - for i in range(25) - ] - } - ] + pages = [{"vectors": [{"key": f"k{i}"} for i in range(25)]}] client = FakeS3Client(pages) session = FakeSession(client) diff --git a/tests/test_bedrock_multimodal.py b/tests/test_bedrock_multimodal.py index 707e952..7ffdec8 100644 --- a/tests/test_bedrock_multimodal.py +++ b/tests/test_bedrock_multimodal.py @@ -237,4 +237,3 @@ def test_empty_string_text_is_sent(mock_boto_session): payload = json.loads(client.invoke_model.call_args[1]["body"]) assert payload["inputText"] == "" assert "inputImage" in payload - diff --git a/tests/test_bm25_retriever.py b/tests/test_bm25_retriever.py index 666d16e..8827185 100644 --- a/tests/test_bm25_retriever.py +++ b/tests/test_bm25_retriever.py @@ -108,6 +108,7 @@ def db(monkeypatch): # Tokenizer tests # --------------------------------------------------------------------------- + def test_default_tokenize_basic(): tokens = default_tokenize("The quick brown fox jumps over the lazy dog.") # Stopwords ("the", "over") filtered out @@ -150,6 +151,7 @@ def test_default_tokenize_empty_and_special(): # BM25Index tests # --------------------------------------------------------------------------- + def test_bm25_index_validation(): with pytest.raises(ValueError, match="k1 must be non-negative"): BM25Index(k1=-1.0) @@ -245,12 +247,15 @@ def test_bm25_index_empty_query(): # BM25Retriever tests # --------------------------------------------------------------------------- + def test_bm25_retriever_basic(db): retriever = BM25Retriever(db) - retriever.index_documents([ - Document(id="doc1", text="Dell XPS-13-9310 laptop"), - Document(id="doc2", text="Lenovo ThinkPad laptop"), - ]) + retriever.index_documents( + [ + Document(id="doc1", text="Dell XPS-13-9310 laptop"), + Document(id="doc2", text="Lenovo ThinkPad laptop"), + ] + ) results = retriever.search("XPS-13-9310") assert len(results) >= 1 @@ -266,10 +271,12 @@ def test_bm25_retriever_basic(db): @pytest.mark.asyncio async def test_bm25_retriever_async(db): retriever = BM25Retriever(db) - retriever.index_documents([ - Document(id="d1", text="ConnectionResetError encountered"), - Document(id="d2", text="HTTP 500 internal server error"), - ]) + retriever.index_documents( + [ + Document(id="d1", text="ConnectionResetError encountered"), + Document(id="d2", text="HTTP 500 internal server error"), + ] + ) results = await retriever.asearch("ConnectionResetError") assert len(results) == 1 @@ -290,6 +297,7 @@ def test_bm25_retriever_populate_from_store(db): # BM25HybridRetriever tests # --------------------------------------------------------------------------- + def test_hybrid_retriever_initialization_and_weights(db): retriever = BM25HybridRetriever(db, dense_weight=1.2, sparse_weight=0.6) assert retriever.dense_weight == 1.2 @@ -347,6 +355,7 @@ async def test_hybrid_retriever_asearch(db): # Client & NamespaceView Ergonomics tests # --------------------------------------------------------------------------- + def test_client_retriever_ergonomics(db): # as_bm25_retriever bm25 = db.as_bm25_retriever(populate_from_store=True) diff --git a/tests/test_cache.py b/tests/test_cache.py index d44587b..26bb754 100644 --- a/tests/test_cache.py +++ b/tests/test_cache.py @@ -44,7 +44,7 @@ def test_dissimilar_miss(): def test_namespace_and_topk_scoping(): c = SemanticCache(threshold=0.5) c.put("ns1", [1.0, 0.0], 5, None, _res("a")) - assert c.get("ns2", [1.0, 0.0], 5, None) is None # different namespace + assert c.get("ns2", [1.0, 0.0], 5, None) is None # different namespace assert c.get("ns1", [1.0, 0.0], 10, None) is None # different top_k @@ -216,9 +216,10 @@ def Table(self, table_name): ttl_jitter_seconds=30, ) - with patch("dynavec.cache.time.time", return_value=1_000), patch( - "dynavec.cache.random.randint", side_effect=[0, 30] - ) as randint: + with ( + patch("dynavec.cache.time.time", return_value=1_000), + patch("dynavec.cache.random.randint", side_effect=[0, 30]) as randint, + ): cache.put("ns", [1.0, 0.0], 5, None, _res("a")) cache.put("ns", [0.0, 1.0], 5, None, _res("b")) diff --git a/tests/test_chat_models.py b/tests/test_chat_models.py index 820c1a1..2c58cc0 100644 --- a/tests/test_chat_models.py +++ b/tests/test_chat_models.py @@ -129,9 +129,7 @@ def __enter__(self): ), SimpleNamespace( type="content_block_delta", - delta=SimpleNamespace( - type="input_json_delta", partial_json='{"loc":"NY"}' - ), + delta=SimpleNamespace(type="input_json_delta", partial_json='{"loc":"NY"}'), ), ] diff --git a/tests/test_client_inmemory.py b/tests/test_client_inmemory.py index 5d3ec0f..75c22b6 100644 --- a/tests/test_client_inmemory.py +++ b/tests/test_client_inmemory.py @@ -90,7 +90,11 @@ def query_pages( yield hits[i : i + chunk_size] def get_vectors(self, keys, return_metadata=False): - return {k: {"vector": self._store[k][0], "metadata": self._store[k][1]} for k in keys if k in self._store} + return { + k: {"vector": self._store[k][0], "metadata": self._store[k][1]} + for k in keys + if k in self._store + } def delete_vectors(self, keys): for k in keys: @@ -128,9 +132,12 @@ def get_versioned(self, namespace, doc_id): def get_many(self, namespace, ids): return { - i: {"text": self._store[(namespace, i)]["text"], - "metadata": self._store[(namespace, i)]["metadata"]} - for i in ids if (namespace, i) in self._store + i: { + "text": self._store[(namespace, i)]["text"], + "metadata": self._store[(namespace, i)]["metadata"], + } + for i in ids + if (namespace, i) in self._store } def delete_many(self, namespace, ids): @@ -229,6 +236,7 @@ def db(monkeypatch): ) return Dynavec(cfg, embedder=HashEmbedder(8)) + @pytest.fixture def cross_encoder_db(monkeypatch): monkeypatch.setattr(client_mod, "S3VectorsStore", FakeS3) @@ -244,6 +252,7 @@ def cross_encoder_db(monkeypatch): ) return Dynavec(cfg, embedder=HashEmbedder(8)) + def test_upsert_and_search_roundtrip(db): res = db.upsert( [ @@ -277,6 +286,7 @@ def test_graph_shortest_path_respects_hop_cap(db): assert db.graph_shortest_path("a", "d", hops=2) == [] + def test_graph_shortest_path_returns_minimum_hops(db): db.graph_add_edge("a", "related_to", "c") db.graph_add_edge("c", "related_to", "e") @@ -291,6 +301,7 @@ def test_graph_shortest_path_returns_minimum_hops(db): "d", ] + def test_search_explain_handles_empty_results(db): result = db.search("missing", explain=True) @@ -405,6 +416,7 @@ def test_dimension_mismatch_raises(db): with pytest.raises(DimensionMismatchError): db.upsert([Document(id="x", vector=[0.1, 0.2])]) # wrong dim (2 != 8) + class WrongOutputEmbedder(HashEmbedder): def embed_documents(self, texts): return [[0.1, 0.2] for _ in texts] @@ -516,7 +528,7 @@ def test_ns_tag_present_in_s3(db): db.upsert([Document(id="1", text="hello")], namespace="ns9") # reach into the fake to confirm the namespace tag was written store = db._vectors._store - (_, meta), = (v for k, v in store.items()) + ((_, meta),) = (v for k, v in store.items()) assert meta[NS_METADATA_KEY] == "ns9" @@ -526,7 +538,7 @@ def test_update_metadata_merge_preserves_vector(db): db.update("1", metadata={"rating": 5}) after_meta = db.get(["1"])[0].metadata assert after_meta["cat"] == "food" # preserved - assert after_meta["rating"] == 5 # added + assert after_meta["rating"] == 5 # added # vector unchanged because neither text nor vector was updated assert db._vectors._store["default#1"][0] == before @@ -582,7 +594,10 @@ def read_then_race(namespace, doc_id): snapshot = real_read(namespace, doc_id) # another writer commits after our read but before our write db._docs.put_versioned( - namespace, doc_id, snapshot["text"], {**snapshot["metadata"], "views": 2}, + namespace, + doc_id, + snapshot["text"], + {**snapshot["metadata"], "views": 2}, expected_version=snapshot["version"], ) return snapshot @@ -616,9 +631,7 @@ def test_update_upsert_if_missing_starts_at_version_one(db): def test_search_stream_yields_incrementally(db): - db.upsert( - [Document(id=str(i), text=f"apple item {i}") for i in range(5)] - ) + db.upsert([Document(id=str(i), text=f"apple item {i}") for i in range(5)]) gen = db.search_stream("apple", top_k=4) first = next(gen) assert first.text is not None @@ -688,7 +701,10 @@ def test_search_many_explain_preserves_order_and_namespace(db, explain): ns = db.namespace("kb") ns.upsert([Document(id="1", text="apple"), Document(id="2", text="rocket")]) batches = db.search_many( - ["rocket", "apple"], namespace="kb", top_k=1, explain=explain, + ["rocket", "apple"], + namespace="kb", + top_k=1, + explain=explain, normalize_scores=True, ) assert len(batches) == 2 @@ -724,6 +740,7 @@ def test_graph_search_scopes_to_related_docs(db): hits = db.graph_search("apple", seed_entities=["fruit"], top_k=5) assert {h.id for h in hits} == {"d1", "d2"} # d3 excluded by the graph + def test_hybrid_graph_search_fuses_ann_and_graph_results(db): db.upsert( [ @@ -888,11 +905,12 @@ def test_search_telemetry_marks_cache_hit(db): db._telemetry = rec db._cache = SemanticCache(threshold=0.99) db.upsert([Document(id="1", text="apple pie")]) - db.search("apple pie", top_k=3) # miss -> populates cache - db.search("apple pie", top_k=3) # hit + db.search("apple pie", top_k=3) # miss -> populates cache + db.search("apple pie", top_k=3) # hit hits = [e.cache_hit for e in rec.events()] assert True in hits and False in hits + def test_cross_encoder_rerank_on_toy_data(monkeypatch, cross_encoder_db): class FakeCrossEncoder: instance = None @@ -902,10 +920,7 @@ def __init__(self, model_name): FakeCrossEncoder.instance = self def predict(self, pairs): - return [ - 0.95 if "target document" in document else 0.10 - for _, document in pairs - ] + return [0.95 if "target document" in document else 0.10 for _, document in pairs] monkeypatch.setitem( sys.modules, @@ -913,18 +928,20 @@ def predict(self, pairs): SimpleNamespace(CrossEncoder=FakeCrossEncoder), ) - cross_encoder_db.upsert([ - Document( - id="d1", - text="ordinary document", - vector=[0.1] * 8, - ), - Document( - id="d2", - text="target document", - vector=[0.1] * 8, - ), - ]) + cross_encoder_db.upsert( + [ + Document( + id="d1", + text="ordinary document", + vector=[0.1] * 8, + ), + Document( + id="d2", + text="target document", + vector=[0.1] * 8, + ), + ] + ) results = cross_encoder_db.search( "find the target", @@ -937,18 +954,13 @@ def predict(self, pairs): assert FakeCrossEncoder.instance.model_name == "toy-cross-encoder" -def test_search_explain_records_cross_encoder_rerank_stage( - monkeypatch, cross_encoder_db -): +def test_search_explain_records_cross_encoder_rerank_stage(monkeypatch, cross_encoder_db): class FakeCrossEncoder: def __init__(self, model_name): self.model_name = model_name def predict(self, pairs): - return [ - 0.95 if "target document" in document else 0.10 - for _, document in pairs - ] + return [0.95 if "target document" in document else 0.10 for _, document in pairs] monkeypatch.setitem( sys.modules, @@ -956,18 +968,20 @@ def predict(self, pairs): SimpleNamespace(CrossEncoder=FakeCrossEncoder), ) - cross_encoder_db.upsert([ - Document( - id="d1", - text="ordinary document", - vector=[0.1] * 8, - ), - Document( - id="d2", - text="target document", - vector=[0.1] * 8, - ), - ]) + cross_encoder_db.upsert( + [ + Document( + id="d1", + text="ordinary document", + vector=[0.1] * 8, + ), + Document( + id="d2", + text="target document", + vector=[0.1] * 8, + ), + ] + ) result = cross_encoder_db.search( "find the target", @@ -987,13 +1001,15 @@ def predict(self, pairs): def test_cross_encoder_rerank_requires_text_query(cross_encoder_db): - cross_encoder_db.upsert([ - Document( - id="d1", - text="some document", - vector=[0.1] * 8, - ), - ]) + cross_encoder_db.upsert( + [ + Document( + id="d1", + text="some document", + vector=[0.1] * 8, + ), + ] + ) with pytest.raises(ConfigurationError, match="requires a text query"): cross_encoder_db.search( @@ -1002,6 +1018,7 @@ def test_cross_encoder_rerank_requires_text_query(cross_encoder_db): rerank="cross-encoder", ) + def test_cross_encoder_rerank_requires_document_text(monkeypatch, cross_encoder_db): class FakeCrossEncoder: def __init__(self, model_name): @@ -1013,13 +1030,15 @@ def __init__(self, model_name): SimpleNamespace(CrossEncoder=FakeCrossEncoder), ) - cross_encoder_db.upsert([ - Document( - id="d1", - text=None, - vector=[0.1] * 8, - ), - ]) + cross_encoder_db.upsert( + [ + Document( + id="d1", + text=None, + vector=[0.1] * 8, + ), + ] + ) with pytest.raises(ConfigurationError, match="requires document text"): cross_encoder_db.search( @@ -1033,15 +1052,15 @@ def test_writes_invalidate_query_cache(db): db._cache = SemanticCache(threshold=0.99) db.upsert([Document(id="1", text="apple pie")], namespace="ns") db.upsert([Document(id="o1", text="apple pie")], namespace="other") - db.search("apple pie", top_k=3, namespace="ns") # miss -> cached - db.search("apple pie", top_k=3, namespace="other") # miss -> cached + db.search("apple pie", top_k=3, namespace="ns") # miss -> cached + db.search("apple pie", top_k=3, namespace="other") # miss -> cached assert db._cache.misses == 2 # a write to "ns" evicts only that namespace's entries db.upsert([Document(id="2", text="apple pie tart")], namespace="ns") - db.search("apple pie", top_k=3, namespace="other") # still cached -> hit + db.search("apple pie", top_k=3, namespace="other") # still cached -> hit assert db._cache.hits == 1 - res = db.search("apple pie", top_k=3, namespace="ns") # evicted -> fresh + res = db.search("apple pie", top_k=3, namespace="ns") # evicted -> fresh assert db._cache.misses == 3 assert "2" in {r.id for r in res} @@ -1058,14 +1077,17 @@ def test_writes_keep_cache_when_invalidation_disabled(monkeypatch): monkeypatch.setattr(client_mod, "S3VectorsStore", FakeS3) monkeypatch.setattr(client_mod, "DynamoDBStore", FakeDDB) cfg = DynavecConfig( - vector_bucket="b", index="i", table="t", dimension=8, + vector_bucket="b", + index="i", + table="t", + dimension=8, cache_invalidate_on_write=False, ) db = Dynavec(cfg, embedder=HashEmbedder(8)) db._cache = SemanticCache(threshold=0.99) db.upsert([Document(id="1", text="apple pie")], namespace="ns") - db.search("apple pie", top_k=3, namespace="ns") # cached + db.search("apple pie", top_k=3, namespace="ns") # cached db.upsert([Document(id="2", text="apple pie tart")], namespace="ns") - db.search("apple pie", top_k=3, namespace="ns") # stale hit + db.search("apple pie", top_k=3, namespace="ns") # stale hit assert db._cache.hits == 1 diff --git a/tests/test_crewai_tool.py b/tests/test_crewai_tool.py index 852e575..4a9684c 100644 --- a/tests/test_crewai_tool.py +++ b/tests/test_crewai_tool.py @@ -80,4 +80,4 @@ def test_as_crewai_tool_missing_dependency_raises(): from dynavec.integrations.tools import as_crewai_tool with pytest.raises(MissingDependencyError): - as_crewai_tool(_FakeClient()) \ No newline at end of file + as_crewai_tool(_FakeClient()) diff --git a/tests/test_describe.py b/tests/test_describe.py index 4d2ba4f..1e72411 100644 --- a/tests/test_describe.py +++ b/tests/test_describe.py @@ -128,4 +128,4 @@ def test_get_index_request_matches_s3vectors_service_model(): result = store.get_index() stubber.assert_no_pending_responses() - assert result["index"]["indexName"] == "test-index" \ No newline at end of file + assert result["index"]["indexName"] == "test-index" diff --git a/tests/test_dynamodb.py b/tests/test_dynamodb.py index 917e314..a0ccfc9 100644 --- a/tests/test_dynamodb.py +++ b/tests/test_dynamodb.py @@ -138,7 +138,9 @@ def test_put_versioned_conflict_leaves_item_untouched(moto_store): moto_store.put_versioned("tenant", "a", "loser", {"by": "stale"}, expected_version=1) assert (info.value.doc_id, info.value.namespace, info.value.expected_version) == ( - "a", "tenant", 1, + "a", + "tenant", + 1, ) assert moto_store.get_versioned("tenant", "a")["metadata"] == {"by": "second"} diff --git a/tests/test_embedding_cache.py b/tests/test_embedding_cache.py index f17d8f3..b673cc7 100644 --- a/tests/test_embedding_cache.py +++ b/tests/test_embedding_cache.py @@ -16,13 +16,14 @@ # Fake embedders (same style as test_async_embeddings.py) # --------------------------------------------------------------------------- + class _CountingEmbedder(Embedder): """Deterministic embedder that tracks every text it actually embeds.""" dimension: ClassVar[int] = 4 def __init__(self) -> None: - self.calls: list[list[str]] = [] # one entry per embed_documents call + self.calls: list[list[str]] = [] # one entry per embed_documents call def embed_documents(self, texts: list[str]) -> list[Vector]: self.calls.append(list(texts)) @@ -59,6 +60,7 @@ async def aembed_documents(self, texts: list[str]) -> list[Vector]: # Custom cache backend (proves pluggability) # --------------------------------------------------------------------------- + class _MinimalCache(EmbeddingCache): """Bare-minimum custom cache to verify the ABC contract works.""" @@ -77,19 +79,23 @@ def set(self, key: str, vector: Vector) -> None: # 1. _make_key # --------------------------------------------------------------------------- + def test_make_key_same_inputs_same_key(): assert _make_key("hello", "m1") == _make_key("hello", "m1") + def test_make_key_different_text_different_key(): assert _make_key("hello", "m1") != _make_key("world", "m1") + def test_make_key_different_model_different_key(): assert _make_key("hello", "m1") != _make_key("hello", "m2") + def test_make_key_format(): key = _make_key("abc", "my-model") sha, model = key.split(":") - assert len(sha) == 64 # sha256 hex is always 64 chars + assert len(sha) == 64 # sha256 hex is always 64 chars assert model == "my-model" @@ -97,40 +103,47 @@ def test_make_key_format(): # 2. InMemoryCache # --------------------------------------------------------------------------- + def test_inmemory_miss_returns_none(): assert InMemoryCache().get("nonexistent") is None + def test_inmemory_set_then_get(): cache = InMemoryCache() cache.set("k", [1.0, 2.0]) assert cache.get("k") == [1.0, 2.0] + def test_inmemory_get_for_set_for(): cache = InMemoryCache() cache.set_for("hello", "model-a", [0.5, 0.5]) assert cache.get_for("hello", "model-a") == [0.5, 0.5] - assert cache.get_for("hello", "model-b") is None # different model → miss + assert cache.get_for("hello", "model-b") is None # different model → miss + def test_inmemory_lru_eviction(): cache = InMemoryCache(max_size=2) cache.set("a", [1.0]) cache.set("b", [2.0]) - cache.get("a") # touch "a" → now "b" is LRU - cache.set("c", [3.0]) # should evict "b", not "a" + cache.get("a") # touch "a" → now "b" is LRU + cache.set("c", [3.0]) # should evict "b", not "a" assert cache.get("b") is None assert cache.get("a") == [1.0] assert cache.get("c") == [3.0] + def test_inmemory_max_size_none_is_unlimited(): cache = InMemoryCache(max_size=None) for i in range(10_000): cache.set(str(i), [float(i)]) assert len(cache) == 10_000 + def test_inmemory_negative_max_size_raises(): with pytest.raises(ValueError): InMemoryCache(max_size=-1) + def test_inmemory_len_and_contains(): cache = InMemoryCache() assert len(cache) == 0 @@ -139,22 +152,25 @@ def test_inmemory_len_and_contains(): assert "x" in cache assert "y" not in cache + def test_inmemory_clear(): cache = InMemoryCache() cache.set("x", [0.0]) cache.clear() assert len(cache) == 0 + def test_inmemory_stats_hit_rate(): cache = InMemoryCache() cache.set_for("hi", "m", [1.0]) - cache.get_for("hi", "m") # hit - cache.get_for("bye", "m") # miss + cache.get_for("hi", "m") # hit + cache.get_for("bye", "m") # miss s = cache.stats() assert s["hits"] == 1 assert s["misses"] == 1 assert s["hit_rate"] == 0.5 + def test_inmemory_reset_stats(): cache = InMemoryCache() cache.set_for("hi", "m", [1.0]) @@ -167,28 +183,32 @@ def test_inmemory_reset_stats(): # 3. CachedEmbedder — embed_documents # --------------------------------------------------------------------------- + def test_first_call_hits_inner(): inner = _CountingEmbedder() embedder = CachedEmbedder(inner, cache=InMemoryCache()) embedder.embed_documents(["a", "b"]) assert inner.calls == [["a", "b"]] + def test_second_call_is_full_cache_hit(): inner = _CountingEmbedder() embedder = CachedEmbedder(inner, cache=InMemoryCache()) v1 = embedder.embed_documents(["a", "b"]) v2 = embedder.embed_documents(["a", "b"]) - assert len(inner.calls) == 1 # inner called only once + assert len(inner.calls) == 1 # inner called only once assert v1 == v2 + def test_partial_hit_calls_inner_only_for_misses(): inner = _CountingEmbedder() embedder = CachedEmbedder(inner, cache=InMemoryCache()) - embedder.embed_documents(["a"]) # warm "a" - embedder.embed_documents(["a", "b"]) # "a" hit, "b" miss + embedder.embed_documents(["a"]) # warm "a" + embedder.embed_documents(["a", "b"]) # "a" hit, "b" miss # second call should only have sent "b" to inner assert inner.calls[-1] == ["b"] + def test_order_preserved_with_mixed_hits_misses(): inner = _CountingEmbedder() embedder = CachedEmbedder(inner, cache=InMemoryCache()) @@ -199,12 +219,14 @@ def test_order_preserved_with_mixed_hits_misses(): assert v_zyx[1] == v_xyz[1] assert v_zyx[2] == v_xyz[0] + def test_empty_input_returns_empty(): inner = _CountingEmbedder() embedder = CachedEmbedder(inner, cache=InMemoryCache()) assert embedder.embed_documents([]) == [] assert inner.calls == [] + def test_dimension_delegated_to_inner(): inner = _CountingEmbedder() assert CachedEmbedder(inner).dimension == inner.dimension @@ -214,19 +236,21 @@ def test_dimension_delegated_to_inner(): # 4. CachedEmbedder — embed_query # --------------------------------------------------------------------------- + def test_embed_query_miss_then_hit(): inner = _CountingEmbedder() embedder = CachedEmbedder(inner, cache=InMemoryCache()) v1 = embedder.embed_query("hello") v2 = embedder.embed_query("hello") assert v1 == v2 - assert inner.total_texts_embedded == 1 # embedded only once + assert inner.total_texts_embedded == 1 # embedded only once # --------------------------------------------------------------------------- # 5. Async variants # --------------------------------------------------------------------------- + @pytest.mark.asyncio async def test_aembed_documents_caches(): inner = _CountingEmbedder() @@ -236,6 +260,7 @@ async def test_aembed_documents_caches(): assert len(inner.calls) == 1 assert v1 == v2 + @pytest.mark.asyncio async def test_aembed_documents_partial_hit(): inner = _CountingEmbedder() @@ -244,6 +269,7 @@ async def test_aembed_documents_partial_hit(): await embedder.aembed_documents(["p", "q"]) assert inner.calls[-1] == ["q"] + @pytest.mark.asyncio async def test_aembed_query_caches(): inner = _CountingEmbedder() @@ -253,6 +279,7 @@ async def test_aembed_query_caches(): assert v1 == v2 assert inner.total_texts_embedded == 1 + @pytest.mark.asyncio async def test_aembed_documents_native_async_inner(): """Cache should work when inner has a native async path.""" @@ -260,7 +287,7 @@ async def test_aembed_documents_native_async_inner(): embedder = CachedEmbedder(inner, cache=InMemoryCache()) v1 = await embedder.aembed_documents(["foo", "bar"]) v2 = await embedder.aembed_documents(["foo", "bar"]) - assert len(inner.async_calls) == 1 # native async called only once + assert len(inner.async_calls) == 1 # native async called only once assert v1 == v2 @@ -268,6 +295,7 @@ async def test_aembed_documents_native_async_inner(): # 6. model_key isolation # --------------------------------------------------------------------------- + def test_different_model_keys_dont_share_entries(): shared_cache = InMemoryCache() inner_a = _CountingEmbedder() @@ -275,15 +303,16 @@ def test_different_model_keys_dont_share_entries(): ea = CachedEmbedder(inner_a, cache=shared_cache, model_key="model-A") eb = CachedEmbedder(inner_b, cache=shared_cache, model_key="model-B") ea.embed_documents(["hello"]) - eb.embed_documents(["hello"]) # same text, different key → must be a miss + eb.embed_documents(["hello"]) # same text, different key → must be a miss assert inner_a.total_texts_embedded == 1 - assert inner_b.total_texts_embedded == 1 # inner_b was also called + assert inner_b.total_texts_embedded == 1 # inner_b was also called # --------------------------------------------------------------------------- # 7. Pluggable custom backend # --------------------------------------------------------------------------- + def test_custom_cache_backend_is_used(): custom = _MinimalCache() inner = _CountingEmbedder() @@ -293,6 +322,7 @@ def test_custom_cache_backend_is_used(): assert inner.total_texts_embedded == 1 assert embedder.cache is custom + def test_default_cache_is_inmemory(): embedder = CachedEmbedder(_CountingEmbedder()) assert isinstance(embedder.cache, InMemoryCache) diff --git a/tests/test_eval_trend.py b/tests/test_eval_trend.py index 1bea2d7..3b031c8 100644 --- a/tests/test_eval_trend.py +++ b/tests/test_eval_trend.py @@ -12,6 +12,7 @@ # Helpers # --------------------------------------------------------------------------- + def _retrieval_summary(recall_10: float, ndcg_10: float, mrr: float = 0.8) -> RetrievalEvalSummary: return RetrievalEvalSummary( total_queries=10, @@ -40,6 +41,7 @@ def _rag_summary(faithfulness: float, relevance: float = 0.85) -> EvalSummary: # 1. record() # --------------------------------------------------------------------------- + def test_record_retrieval_summary_returns_eval_run(): store = EvalRunStore() summary = _retrieval_summary(0.85, 0.80) @@ -78,10 +80,11 @@ def test_record_respects_max_runs(): # 2. filter() # --------------------------------------------------------------------------- + def test_filter_by_single_tag(): store = EvalRunStore() store.record(_retrieval_summary(0.80, 0.75), tags={"dataset": "squad", "model": "ada"}) - store.record(_retrieval_summary(0.70, 0.65), tags={"dataset": "nq", "model": "ada"}) + store.record(_retrieval_summary(0.70, 0.65), tags={"dataset": "nq", "model": "ada"}) store.record(_retrieval_summary(0.90, 0.85), tags={"dataset": "squad", "model": "v3"}) squad_runs = store.filter(dataset="squad") @@ -109,6 +112,7 @@ def test_filter_no_match_returns_empty(): # 3. trend() # --------------------------------------------------------------------------- + def test_trend_recall_returns_sorted_time_series(): store = EvalRunStore() # record three runs at explicit timestamps @@ -139,7 +143,7 @@ def test_trend_faithfulness_from_rag_summaries(): def test_trend_with_tag_filter(): store = EvalRunStore() store.record(_retrieval_summary(0.60, 0.55), tags={"model": "ada"}, ts=100.0) - store.record(_retrieval_summary(0.80, 0.75), tags={"model": "v3"}, ts=200.0) + store.record(_retrieval_summary(0.80, 0.75), tags={"model": "v3"}, ts=200.0) store.record(_retrieval_summary(0.65, 0.60), tags={"model": "ada"}, ts=300.0) series = store.trend("recall", k=10, model="ada") @@ -163,6 +167,7 @@ def test_trend_skips_runs_without_metric(): # 4. compare() # --------------------------------------------------------------------------- + def test_compare_two_runs_returns_both(): store = EvalRunStore() r1 = store.record(_retrieval_summary(0.75, 0.70), tags={"model": "ada"}) @@ -186,8 +191,10 @@ def test_compare_missing_id_raises(): # 5. to_dict round-trip # --------------------------------------------------------------------------- + def test_eval_run_to_dict_is_serialisable(): import json + store = EvalRunStore() run = store.record(_retrieval_summary(0.82, 0.78), tags={"dataset": "squad"}) d = run.to_dict() diff --git a/tests/test_export_import.py b/tests/test_export_import.py index 918bb15..b2bc4be 100644 --- a/tests/test_export_import.py +++ b/tests/test_export_import.py @@ -63,7 +63,8 @@ def query(self, query_vector, top_k, filter=None, return_metadata=True, return_d def get_vectors(self, keys, return_metadata=False): return { k: {"vector": self._store[k][0], "metadata": self._store[k][1]} - for k in keys if k in self._store + for k in keys + if k in self._store } def list_pages(self, return_data=False, return_metadata=False, page_size=None): @@ -154,8 +155,14 @@ def test_import_namespace_from_file(monkeypatch, tmp_path): db = _make(monkeypatch) in_file = tmp_path / "import.jsonl" in_file.write_text( - json.dumps({"id": "imp1", "vector": [0.1] * 8, "text": "imported text", "metadata": {"x": 1}}) + "\n" - + json.dumps({"id": "imp2", "vector": [0.2] * 8, "text": "second imported", "metadata": {"x": 2}}) + "\n", + json.dumps( + {"id": "imp1", "vector": [0.1] * 8, "text": "imported text", "metadata": {"x": 1}} + ) + + "\n" + + json.dumps( + {"id": "imp2", "vector": [0.2] * 8, "text": "second imported", "metadata": {"x": 2}} + ) + + "\n", encoding="utf-8", ) diff --git a/tests/test_fusion.py b/tests/test_fusion.py index c2ab0cb..a072039 100644 --- a/tests/test_fusion.py +++ b/tests/test_fusion.py @@ -13,6 +13,7 @@ def _make_result(doc_id: str, score: float) -> SearchResult: # ── helpers ───────────────────────────────────────────────────────────────── + def _two_retriever_fixture(): """Two queries, two retrievers — dense is better.""" labeled = [ @@ -41,6 +42,7 @@ def _three_retriever_fixture(): # ── score ─────────────────────────────────────────────────────────────────── + def test_score_returns_float(): labeled, dense, sparse = _two_retriever_fixture() fitter = RRFWeightFitter(labeled, [dense, sparse]) @@ -50,6 +52,7 @@ def test_score_returns_float(): # ── fit() returns FitResult ───────────────────────────────────────────────── + @pytest.mark.parametrize("method", ["grid", "random"]) def test_fit_returns_fit_result(method): labeled, dense, sparse = _two_retriever_fixture() @@ -65,6 +68,7 @@ def test_fit_returns_fit_result(method): # ── grid search ───────────────────────────────────────────────────────────── + def test_fit_grid_improves_over_equal_weights(): labeled = [LabeledQuery(query="q1", relevant_ids={"doc_a"})] dense = [[_make_result("doc_a", 0.9), _make_result("doc_b", 0.1)]] @@ -91,6 +95,7 @@ def test_fit_grid_three_retrievers(): # ── random search ─────────────────────────────────────────────────────────── + def test_fit_random_three_retrievers(): labeled, r1, r2, r3 = _three_retriever_fixture() fitter = RRFWeightFitter(labeled, [r1, r2, r3]) @@ -102,6 +107,7 @@ def test_fit_random_three_retrievers(): # ── bayesian fitting ──────────────────────────────────────────────────────── + def test_fit_bayesian_two_retrievers(): """Bayesian fitting should return valid FitResult.""" labeled, dense, sparse = _two_retriever_fixture() @@ -128,6 +134,7 @@ def test_fit_bayesian_three_retrievers(): # ── serialisation (save / load) ───────────────────────────────────────────── + def test_save_load_roundtrip(tmp_path): """FitResult should survive a JSON roundtrip.""" original = FitResult( @@ -166,6 +173,7 @@ def test_to_dict(): # ── end-to-end: fit then save/load ────────────────────────────────────────── + def test_fit_and_save_load(tmp_path): """Full workflow: fit → save → load → use weights.""" labeled, dense, sparse = _two_retriever_fixture() @@ -182,6 +190,7 @@ def test_fit_and_save_load(tmp_path): # ── error handling ────────────────────────────────────────────────────────── + def test_empty_labeled_queries_raises(): with pytest.raises(ValueError, match="labeled_queries"): RRFWeightFitter([], [[]]) diff --git a/tests/test_graph_export.py b/tests/test_graph_export.py index 93a3c8a..dc04308 100644 --- a/tests/test_graph_export.py +++ b/tests/test_graph_export.py @@ -55,9 +55,7 @@ def graph(): def test_mermaid_export_renders_the_namespace_subgraph(graph): assert graph.graph_export("mermaid", "corp") == ( - "graph LR\n" - " acme -->|competes_with| globex\n" - " globex -->|supplies| initech\n" + "graph LR\n acme -->|competes_with| globex\n globex -->|supplies| initech\n" ) @@ -84,10 +82,7 @@ def test_cycle_terminates_and_emits_each_edge_once(): g.add_edge("ns", "a", "next", "b") # duplicate adjacency entry assert g.graph_export("mermaid", "ns") == ( - "graph LR\n" - " a -->|next| b\n" - " b -->|next| c\n" - " c -->|next| a\n" + "graph LR\n a -->|next| b\n b -->|next| c\n c -->|next| a\n" ) @@ -98,9 +93,7 @@ def test_cycle_terminates_when_walking_from_a_root(): g.add_node("ns", "orphan") # unreachable from "a", so absent assert g.graph_export("mermaid", "ns", roots=["a"]) == ( - "graph LR\n" - " a -->|next| b\n" - " b -->|next| a\n" + "graph LR\n a -->|next| b\n b -->|next| a\n" ) @@ -108,11 +101,7 @@ def test_self_loop_terminates(): g = FakeGraph() g.add_edge("ns", "a", "cites", "a") - assert g.graph_export("dot", "ns") == ( - "digraph G {\n" - ' "a" -> "a" [label="cites"];\n' - "}\n" - ) + assert g.graph_export("dot", "ns") == ('digraph G {\n "a" -> "a" [label="cites"];\n}\n') # -------------------------------------------------------------------- escaping @@ -156,11 +145,7 @@ def test_mermaid_alias_never_shadows_a_real_id(): g = FakeGraph() g.add_edge("ns", "n0", "points_at", "zz top") # "n0" is a legal bare id - assert g.graph_export("mermaid", "ns") == ( - "graph LR\n" - ' n1["zz top"]\n' - " n0 -->|points_at| n1\n" - ) + assert g.graph_export("mermaid", "ns") == ('graph LR\n n1["zz top"]\n n0 -->|points_at| n1\n') def test_dot_escapes_backslashes(): @@ -189,16 +174,9 @@ def test_isolated_nodes_are_declared(): g.add_node("ns", "lonely") g.add_edge("ns", "a", "next", "b") - assert g.graph_export("mermaid", "ns") == ( - "graph LR\n" - " lonely\n" - " a -->|next| b\n" - ) + assert g.graph_export("mermaid", "ns") == ("graph LR\n lonely\n a -->|next| b\n") assert g.graph_export("dot", "ns") == ( - "digraph G {\n" - ' "lonely";\n' - ' "a" -> "b" [label="next"];\n' - "}\n" + 'digraph G {\n "lonely";\n "a" -> "b" [label="next"];\n}\n' ) @@ -208,8 +186,7 @@ def test_relation_filter_keeps_only_matching_edges(): g.add_edge("ns", "a", "owns", "c") assert g.graph_export("mermaid", "ns", roots=["a"], relation="next") == ( - "graph LR\n" - " a -->|next| b\n" + "graph LR\n a -->|next| b\n" ) diff --git a/tests/test_graph_properties.py b/tests/test_graph_properties.py index 18ab68f..3b404aa 100644 --- a/tests/test_graph_properties.py +++ b/tests/test_graph_properties.py @@ -165,4 +165,3 @@ def test_add_node_partial_update_preserves_unspecified_attributes(store): node_after_ntype = store.get_node("ns", "user") assert node_after_ntype["ntype"] == "Admin" assert node_after_ntype["props"] == {"age": Decimal("31"), "city": "NYC"} - diff --git a/tests/test_haystack.py b/tests/test_haystack.py index b92a5d7..35dc5c0 100644 --- a/tests/test_haystack.py +++ b/tests/test_haystack.py @@ -1,4 +1,5 @@ """Tests for the Haystack integration.""" + # ruff: noqa: E402 import pytest @@ -85,6 +86,7 @@ def test_write_documents_with_embedding(): assert dynavec_document.metadata == {"category": "account"} assert dynavec_document.vector == [0.1, 0.2, 0.3] + def test_write_documents_overwrite_existing(): client = _FakeClient() store = DynavecDocumentStore(client, namespace="kb") @@ -95,9 +97,7 @@ def test_write_documents_overwrite_existing(): meta={"category": "account"}, ) - client.get = lambda ids, **kwargs: [ - type("Result", (), {"id": "doc-1"})() - ] + client.get = lambda ids, **kwargs: [type("Result", (), {"id": "doc-1"})()] count = store.write_documents( [document], @@ -113,42 +113,44 @@ def test_write_documents_overwrite_existing(): assert documents[0].text == "Updated password instructions" assert kwargs == {"namespace": "kb"} + def test_filter_documents(): client = _FakeClient() store = DynavecDocumentStore(client, namespace="kb") - client.list_vectors = lambda **kwargs: iter([ - type( - "Result", - (), - { - "id": "doc-1", - "text": "Password reset", - "metadata": {"category": "account"}, - "vector": None, - }, - )(), - type( - "Result", - (), - { - "id": "doc-2", - "text": "Home loan details", - "metadata": {"category": "loan"}, - "vector": None, - }, - )(), - ]) - - documents = store.filter_documents( - filters={"category": "account"} + client.list_vectors = lambda **kwargs: iter( + [ + type( + "Result", + (), + { + "id": "doc-1", + "text": "Password reset", + "metadata": {"category": "account"}, + "vector": None, + }, + )(), + type( + "Result", + (), + { + "id": "doc-2", + "text": "Home loan details", + "metadata": {"category": "loan"}, + "vector": None, + }, + )(), + ] ) + documents = store.filter_documents(filters={"category": "account"}) + assert len(documents) == 1 assert documents[0].id == "doc-1" assert documents[0].content == "Password reset" assert documents[0].meta == {"category": "account"} + def test_retriever(): client = _FakeClient() @@ -169,6 +171,7 @@ def test_retriever(): assert result["documents"][0].id == "doc-1" assert result["documents"][0].content == "Password reset instructions" + def test_retriever_passes_top_k_and_filters(): client = _FakeClient() client.search_calls = [] @@ -199,6 +202,7 @@ def search(**kwargs): assert call["namespace"] == "support" assert call["filter"] == {"category": "support"} + def test_retriever_overrides_default_top_k(): client = _FakeClient() client.search_calls = [] @@ -218,6 +222,7 @@ def search(**kwargs): assert client.search_calls[0]["top_k"] == 3 + def test_retriever_with_filters(): client = _FakeClient() client.search_calls = [] @@ -270,6 +275,7 @@ def search(**kwargs): assert client.search_calls[0]["filter"] == {"category": "support"} + def test_document_store_serialization(): client = _FakeClient() store = DynavecDocumentStore(client, namespace="kb") @@ -280,6 +286,7 @@ def test_document_store_serialization(): assert data["init_parameters"]["namespace"] == "kb" assert "config" in data["init_parameters"] + def test_document_store_from_dict(monkeypatch): client = _FakeClient() store = DynavecDocumentStore(client, namespace="kb") @@ -295,32 +302,35 @@ def test_document_store_from_dict(monkeypatch): assert restored.namespace == "kb" + def test_filter_documents_with_haystack_filter(): client = _FakeClient() store = DynavecDocumentStore(client, namespace="kb") - client.list_vectors = lambda **kwargs: iter([ - type( - "Result", - (), - { - "id": "doc-1", - "text": "Password reset", - "metadata": {"category": "account"}, - "vector": None, - }, - )(), - type( - "Result", - (), - { - "id": "doc-2", - "text": "Home loan details", - "metadata": {"category": "loan"}, - "vector": None, - }, - )(), - ]) + client.list_vectors = lambda **kwargs: iter( + [ + type( + "Result", + (), + { + "id": "doc-1", + "text": "Password reset", + "metadata": {"category": "account"}, + "vector": None, + }, + )(), + type( + "Result", + (), + { + "id": "doc-2", + "text": "Home loan details", + "metadata": {"category": "loan"}, + "vector": None, + }, + )(), + ] + ) documents = store.filter_documents( filters={ diff --git a/tests/test_hitl_interrupts.py b/tests/test_hitl_interrupts.py index fc6e432..1be0341 100644 --- a/tests/test_hitl_interrupts.py +++ b/tests/test_hitl_interrupts.py @@ -2,7 +2,6 @@ import asyncio from collections.abc import AsyncIterator, Iterator -import pytest from dynavec.agents.base import interrupt, tool from dynavec.agents.react import ReActAgent @@ -30,7 +29,9 @@ def stream(self, messages: list[Message], tools=None, **kwargs) -> Iterator[Chat res = self.invoke(messages, tools=tools, **kwargs) yield res - async def astream(self, messages: list[Message], tools=None, **kwargs) -> AsyncIterator[ChatResult]: + async def astream( + self, messages: list[Message], tools=None, **kwargs + ) -> AsyncIterator[ChatResult]: res = await self.ainvoke(messages, tools=tools, **kwargs) yield res @@ -127,6 +128,7 @@ def test_resume_execution(): def test_async_interrupt_and_resume(): """Test async arun and aresume functionality using asyncio.run.""" + async def _run_async(): initial_responses = [ ChatResult( @@ -157,7 +159,9 @@ async def _run_async(): ) ) ] - resume_agent = ReActAgent(model=MockChatModel(resume_responses), tools=[sensitive_action_tool]) + resume_agent = ReActAgent( + model=MockChatModel(resume_responses), tools=[sensitive_action_tool] + ) final_result = await resume_agent.aresume( messages=result.interrupt_payload["messages"], @@ -169,4 +173,4 @@ async def _run_async(): assert final_result.finished is True assert final_result.termination_reason == "completed" - asyncio.run(_run_async()) \ No newline at end of file + asyncio.run(_run_async()) diff --git a/tests/test_hot_tier.py b/tests/test_hot_tier.py index 73019b0..27163bc 100644 --- a/tests/test_hot_tier.py +++ b/tests/test_hot_tier.py @@ -69,7 +69,8 @@ def query(self, query_vector, top_k, filter=None, return_metadata=True, return_d def get_vectors(self, keys, return_metadata=False): return { k: {"vector": self._store[k][0], "metadata": self._store[k][1]} - for k in keys if k in self._store + for k in keys + if k in self._store } def list_pages(self, return_data=False, return_metadata=False, page_size=None): @@ -117,6 +118,7 @@ def _make(monkeypatch, **cfg_kw): # --------------------------------------------------------------------- matcher + def test_matches_supported_operators(): m = {"cat": "food", "year": 2026, "score": 0.4} assert matches(m, {"cat": "food"}) @@ -135,6 +137,7 @@ def test_matches_rejects_unsupported_operator(): # ---------------------------------------------------------------- integration + def test_disabled_by_default_uses_s3(monkeypatch): db = _make(monkeypatch) assert db._hot is None @@ -161,10 +164,10 @@ def test_warm_serves_from_ram_without_touching_s3_or_ddb(monkeypatch): hits = db.search("apple", top_k=2) assert len(hits) == 2 - assert hits[0].text is not None # text served from RAM + assert hits[0].text is not None # text served from RAM assert hits[0].score >= hits[1].score - assert db._vectors.query_calls == s3_before # no S3 Vectors query - assert db._docs.get_calls == ddb_before # no DynamoDB hydration + assert db._vectors.query_calls == s3_before # no S3 Vectors query + assert db._docs.get_calls == ddb_before # no DynamoDB hydration def test_write_through_keeps_warmed_namespace_current(monkeypatch): @@ -176,7 +179,7 @@ def test_write_through_keeps_warmed_namespace_current(monkeypatch): db.upsert([Document(id="2", text="apple tart")]) hits = db.search("apple tart", top_k=2) assert {h.id for h in hits} == {"1", "2"} - assert db._vectors.query_calls == s3_before # still served from RAM + assert db._vectors.query_calls == s3_before # still served from RAM def test_filter_on_hot_path(monkeypatch): @@ -191,7 +194,7 @@ def test_filter_on_hot_path(monkeypatch): s3_before = db._vectors.query_calls hits = db.search("apple", top_k=5, filter={"cat": "space"}) assert {h.id for h in hits} == {"2"} - assert db._vectors.query_calls == s3_before # filtered in RAM + assert db._vectors.query_calls == s3_before # filtered in RAM def test_unsupported_filter_falls_back_to_s3(monkeypatch): @@ -204,7 +207,7 @@ def test_unsupported_filter_falls_back_to_s3(monkeypatch): assert db._hot.search("default", [0.0] * 8, 5, {"n": {"$regex": "3"}}) is None s3_before = db._vectors.query_calls db.search("apple", top_k=5, filter={"n": {"$regex": "3"}}) - assert db._vectors.query_calls == s3_before + 1 # fell back to S3 + assert db._vectors.query_calls == s3_before + 1 # fell back to S3 def test_delete_removes_from_hot_tier(monkeypatch): @@ -215,7 +218,7 @@ def test_delete_removes_from_hot_tier(monkeypatch): s3_before = db._vectors.query_calls hits = db.search("apple", top_k=5) assert {h.id for h in hits} == {"2"} - assert db._vectors.query_calls == s3_before # still hot, id 1 gone + assert db._vectors.query_calls == s3_before # still hot, id 1 gone def test_capacity_cap_falls_back_to_s3(monkeypatch): @@ -225,7 +228,7 @@ def test_capacity_cap_falls_back_to_s3(monkeypatch): assert loaded == 0 assert not db._hot.is_authoritative("default") hits = db.search("doc", top_k=3) - assert db._vectors.query_calls >= 1 # fell back to S3 + assert db._vectors.query_calls >= 1 # fell back to S3 assert hits @@ -266,7 +269,9 @@ def test_hot_stats(monkeypatch): def test_hot_tier_lru_eviction(monkeypatch): db = _make(monkeypatch, hot_tier=True, hot_tier_max_vectors=3, hot_tier_eviction="lru") - db.upsert([Document(id="a1", text="apple 1"), Document(id="a2", text="apple 2")], namespace="ns_a") + db.upsert( + [Document(id="a1", text="apple 1"), Document(id="a2", text="apple 2")], namespace="ns_a" + ) db.upsert([Document(id="b1", text="banana 1")], namespace="ns_b") db.upsert([Document(id="c1", text="cherry 1")], namespace="ns_c") @@ -292,7 +297,9 @@ def test_hot_tier_lru_eviction(monkeypatch): def test_hot_tier_fifo_eviction(monkeypatch): db = _make(monkeypatch, hot_tier=True, hot_tier_max_vectors=3, hot_tier_eviction="fifo") - db.upsert([Document(id="a1", text="apple 1"), Document(id="a2", text="apple 2")], namespace="ns_a") + db.upsert( + [Document(id="a1", text="apple 1"), Document(id="a2", text="apple 2")], namespace="ns_a" + ) db.upsert([Document(id="b1", text="banana 1")], namespace="ns_b") db.upsert([Document(id="c1", text="cherry 1")], namespace="ns_c") @@ -315,7 +322,9 @@ def test_hot_tier_fifo_eviction(monkeypatch): def test_hot_tier_none_eviction(monkeypatch): db = _make(monkeypatch, hot_tier=True, hot_tier_max_vectors=3, hot_tier_eviction="none") - db.upsert([Document(id="a1", text="apple 1"), Document(id="a2", text="apple 2")], namespace="ns_a") + db.upsert( + [Document(id="a1", text="apple 1"), Document(id="a2", text="apple 2")], namespace="ns_a" + ) db.upsert([Document(id="b1", text="banana 1")], namespace="ns_b") db.upsert([Document(id="c1", text="cherry 1")], namespace="ns_c") @@ -356,7 +365,9 @@ def test_hot_tier_insert_many_evicts_cold_namespace(monkeypatch): # Now upsert 2 new docs into ns_b. Total vectors for ns_b will be 3. # ns_a (1 vector) + ns_b (3 vectors) = 4 > 3. # ns_a should be evicted, and ns_b should remain authoritative with 3 resident vectors. - db.upsert([Document(id="b2", text="banana 2"), Document(id="b3", text="banana 3")], namespace="ns_b") + db.upsert( + [Document(id="b2", text="banana 2"), Document(id="b3", text="banana 3")], namespace="ns_b" + ) assert not db._hot.is_authoritative("ns_a") assert db._hot.is_authoritative("ns_b") assert db.hot_stats()["resident_vectors"] == 3 diff --git a/tests/test_huggingface_embedder.py b/tests/test_huggingface_embedder.py index ac2f35d..49f68d6 100644 --- a/tests/test_huggingface_embedder.py +++ b/tests/test_huggingface_embedder.py @@ -1,6 +1,7 @@ """Unit tests for HFInferenceEmbedder.""" from unittest.mock import MagicMock, patch + import pytest from dynavec.embeddings import HFInferenceEmbedder @@ -44,4 +45,4 @@ def test_embed_documents_batching(mock_post): def test_empty_input(): emb = HFInferenceEmbedder(api_key="fake_token") - assert emb.embed_documents([]) == [] \ No newline at end of file + assert emb.embed_documents([]) == [] diff --git a/tests/test_list_vectors.py b/tests/test_list_vectors.py index 8ffdea4..9b2d7e2 100644 --- a/tests/test_list_vectors.py +++ b/tests/test_list_vectors.py @@ -258,9 +258,7 @@ def test_list_vectors_hydrates_text_and_metadata(): assert result.distance is None assert result.vector is None # hydrating means we do not pay for metadata on the S3 Vectors side - assert db._vectors.calls == [ - dict(return_data=False, return_metadata=False, page_size=None) - ] + assert db._vectors.calls == [dict(return_data=False, return_metadata=False, page_size=None)] def test_list_vectors_without_hydration_uses_s3_metadata(): @@ -278,9 +276,7 @@ def test_list_vectors_without_hydration_uses_s3_metadata(): assert docs.calls == [] # DynamoDB is never touched assert result.text == "mirrored" assert result.metadata == {"cat": "food"} # internal tags stripped - assert db._vectors.calls == [ - dict(return_data=False, return_metadata=True, page_size=None) - ] + assert db._vectors.calls == [dict(return_data=False, return_metadata=True, page_size=None)] def test_list_vectors_without_hydration_tolerates_missing_metadata(): diff --git a/tests/test_markdown_source.py b/tests/test_markdown_source.py index 304febf..7c1c77c 100644 --- a/tests/test_markdown_source.py +++ b/tests/test_markdown_source.py @@ -115,6 +115,7 @@ def test_root_must_be_an_existing_directory(tmp_path): def test_empty_directory_yields_no_records(tmp_path): assert list(MarkdownSource(tmp_path)) == [] + def test_front_matter_dates_normalize_to_iso_strings(tmp_path): (tmp_path / "notes.md").write_text( "---\ntitle: Research Notes\ndate: 2026-09-24\n---\nBody\n", @@ -188,7 +189,6 @@ def test_front_matter_dates_survive_dynamodb_serialization(tmp_path): assert serialized["M"]["date"] == {"S": "2026-09-24"} - def test_markdown_source_flows_through_ingest(tmp_path): (tmp_path / "guide.md").write_text("---\ntopic: aws\n---\nabcdefghij", encoding="utf-8") db = Mock() diff --git a/tests/test_mcp_server.py b/tests/test_mcp_server.py index 564f0b4..74649d1 100644 --- a/tests/test_mcp_server.py +++ b/tests/test_mcp_server.py @@ -111,7 +111,9 @@ def test_client_from_env_embedder_resolution(monkeypatch): monkeypatch.setattr(gemini_mod, "GeminiEmbedder", lambda **kw: MagicMock(dimension=768)) monkeypatch.setattr(voyage_mod, "VoyageEmbedder", lambda **kw: MagicMock(dimension=1024)) monkeypatch.setattr(mistral_mod, "MistralEmbedder", lambda **kw: MagicMock(dimension=1024)) - monkeypatch.setattr(st_mod, "SentenceTransformerEmbedder", lambda **kw: MagicMock(dimension=384)) + monkeypatch.setattr( + st_mod, "SentenceTransformerEmbedder", lambda **kw: MagicMock(dimension=384) + ) base_env = { "DYNAVEC_BUCKET": "b", @@ -175,7 +177,7 @@ def test_mcp_server_search_tool(fake_mcp): rescore="dot", rerank="mmr", ) - assert "[1] id: doc-1 (score: 0.9500 | metadata: {\"topic\": \"greeting\"})" in result_text + assert '[1] id: doc-1 (score: 0.9500 | metadata: {"topic": "greeting"})' in result_text assert "Hello world" in result_text assert "[2] id: doc-2 (score: 0.8200)" in result_text @@ -197,7 +199,9 @@ def test_mcp_server_search_tool_empty_and_error(fake_mcp): def test_mcp_server_graph_search_tool(fake_mcp): mock_db = MagicMock() mock_db.graph_search.return_value = [ - SearchResult(id="doc-graph", score=0.91, text="Graph entity content", metadata={"entity": "acme"}), + SearchResult( + id="doc-graph", score=0.91, text="Graph entity content", metadata={"entity": "acme"} + ), ] mcp = create_mcp_server(mock_db) diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 82b46ca..79db0b0 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -166,14 +166,12 @@ def test_split_text_mirror_optional(): # None and empty filter (None, "default", {NS_METADATA_KEY: "default"}), ({}, "production", {NS_METADATA_KEY: "production"}), - # Bare equality ( {"genre": "scifi"}, "ns1", {"$and": [{"genre": "scifi"}, {NS_METADATA_KEY: "ns1"}]}, ), - # $eq operator ( {"status": {"$eq": "active"}}, @@ -185,7 +183,6 @@ def test_split_text_mirror_optional(): ] }, ), - # $ne operator ( {"archived": {"$ne": True}}, @@ -197,7 +194,6 @@ def test_split_text_mirror_optional(): ] }, ), - # Numeric comparison operators ($gte, $gt, $lte, $lt) ( {"rating": {"$gte": 4.5}}, @@ -229,7 +225,6 @@ def test_split_text_mirror_optional(): ] }, ), - # $in and $nin operators ( {"tag": {"$in": ["python", "aws", "vector"]}}, @@ -251,7 +246,6 @@ def test_split_text_mirror_optional(): ] }, ), - # $or operator combinations ( {"$or": [{"category": "news"}, {"category": "tech"}]}, @@ -295,11 +289,7 @@ def test_split_text_mirror_optional(): "$and": [ { "$or": [ - { - "category": { - "$in": ["electronics", "computers"] - } - }, + {"category": {"$in": ["electronics", "computers"]}}, {"discount": {"$gte": 0.2}}, ] }, @@ -307,7 +297,6 @@ def test_split_text_mirror_optional(): ] }, ), - # Top-level $and list extension (should not double wrap) ( {"$and": [{"a": 1}]}, @@ -330,7 +319,6 @@ def test_split_text_mirror_optional(): ] }, ), - # Complex nested $and + $or + $in + $gte combinations ( { diff --git a/tests/test_metric_properties.py b/tests/test_metric_properties.py index a5c26dc..9257573 100644 --- a/tests/test_metric_properties.py +++ b/tests/test_metric_properties.py @@ -197,7 +197,9 @@ def test_euclidean_and_manhattan_distance_monotonicity(vec_pair, scale_a, delta) for metric in ("euclidean", "manhattan"): scores = score(q, mat, metric) # Closer candidate (near) must have strictly higher similarity score than farther candidate - assert scores[0] > scores[1], f"Metric {metric} failed monotonicity: {scores[0]} <= {scores[1]}" + assert scores[0] > scores[1], ( + f"Metric {metric} failed monotonicity: {scores[0]} <= {scores[1]}" + ) @given( @@ -297,7 +299,9 @@ def test_normalize_scores_bounds_and_order_preservation(scores): scores=arrays( dtype=np.float32, shape=st.integers(min_value=2, max_value=20), - elements=st.floats(min_value=-100.0, max_value=100.0, allow_nan=False, allow_infinity=False), + elements=st.floats( + min_value=-100.0, max_value=100.0, allow_nan=False, allow_infinity=False + ), ), scale=st.floats(min_value=0.1, max_value=50.0, allow_nan=False, allow_infinity=False), shift=st.floats(min_value=-50.0, max_value=50.0, allow_nan=False, allow_infinity=False), diff --git a/tests/test_namespace.py b/tests/test_namespace.py index 12cfece..e3a1434 100644 --- a/tests/test_namespace.py +++ b/tests/test_namespace.py @@ -21,11 +21,15 @@ def embed_documents(self, texts): v[i % self.dimension] += (ord(ch) % 17) / 17.0 out.append(v) return out + + def _cosine_distance(a, b): dot = sum(x * y for x, y in zip(a, b)) na = math.sqrt(sum(x * x for x in a)) or 1e-12 nb = math.sqrt(sum(y * y for y in b)) or 1e-12 return 1.0 - dot / (na * nb) + + class FakeS3(client_mod.S3VectorsStore): def __init__(self, config, boto_session=None): self.config = config @@ -70,12 +74,17 @@ def query_pages( yield hits[i : i + chunk_size] def get_vectors(self, keys, return_metadata=False): - return {k: {"vector": self._store[k][0], "metadata": self._store[k][1]} for k in keys if k in self._store} + return { + k: {"vector": self._store[k][0], "metadata": self._store[k][1]} + for k in keys + if k in self._store + } def delete_vectors(self, keys): for k in keys: self._store.pop(k, None) + class FakeDDB(client_mod.DynamoDBStore): def __init__(self, config, boto_session=None): self.config = config @@ -91,6 +100,8 @@ def get_many(self, namespace, ids): def delete_many(self, namespace, ids): for i in ids: self._store.pop((namespace, i), None) + + class FakeGraph(client_mod.GraphStore): def __init__(self, config, boto_session=None): self.config = config @@ -166,6 +177,8 @@ def get_docs(self, ns, entity_ids): seen.add(d) out.append(d) return out + + @pytest.fixture def db(monkeypatch): monkeypatch.setattr(client_mod, "S3VectorsStore", FakeS3) @@ -174,11 +187,15 @@ def db(monkeypatch): cfg = DynavecConfig(vector_bucket="b", index="i", table="t", dimension=8) return Dynavec(cfg, embedder=HashEmbedder(8)) + + def test_search_many_forwards_multiple_queries_with_namespace_isolation(db): kb = db.namespace("kb") other = db.namespace("other") kb.upsert([Document(id="1", text="apple"), Document(id="2", text="rocket")]) - other.upsert([Document(id="other-apple", text="apple"), Document(id="other-rocket", text="rocket")]) + other.upsert( + [Document(id="other-apple", text="apple"), Document(id="other-rocket", text="rocket")] + ) results = kb.search_many(["apple", "rocket"], top_k=1) @@ -207,6 +224,7 @@ def test_search_many_empty_batch(db): assert kb.search_many([]) == [] + @pytest.mark.parametrize("explain", [False, True]) def test_search_many_dynamic_explain_preserves_namespace_and_order(db, explain): kb = db.namespace("kb") diff --git a/tests/test_openai_embeddings.py b/tests/test_openai_embeddings.py index a3fc8dd..ea5ee0f 100644 --- a/tests/test_openai_embeddings.py +++ b/tests/test_openai_embeddings.py @@ -1,4 +1,5 @@ """Tests for the OpenAI embedder retry behavior.""" + from types import SimpleNamespace import pytest @@ -32,9 +33,7 @@ def __init__(self, error: Exception): @pytest.mark.parametrize("dimensions", [None, 2]) -def test_openai_embedder_retries_rate_limit_and_respects_retry_after( - monkeypatch, dimensions -): +def test_openai_embedder_retries_rate_limit_and_respects_retry_after(monkeypatch, dimensions): class RateLimitError(Exception): status_code = 429 @@ -64,6 +63,7 @@ def __init__(self): expected_request["dimensions"] = dimensions assert fake_client.embeddings.requests == [expected_request, expected_request] + def test_openai_embedder_retries_server_error(monkeypatch): class ServerError(Exception): status_code = 500 @@ -81,6 +81,7 @@ class ServerError(Exception): assert fake_client.embeddings.calls == 2 assert len(delays) == 1 + def test_openai_embedder_does_not_retry_client_error(monkeypatch): class ClientError(Exception): status_code = 400 @@ -99,6 +100,7 @@ class ClientError(Exception): assert fake_client.embeddings.calls == 1 assert delays == [] + def _make_embedder(error): fake_client = _FakeClient(error) diff --git a/tests/test_provisioning.py b/tests/test_provisioning.py index f5c908a..397cc7f 100644 --- a/tests/test_provisioning.py +++ b/tests/test_provisioning.py @@ -25,16 +25,10 @@ def test_ensure_table_creates_expected_key_schema(): ensure_table(config, session) - table = session.client("dynamodb").describe_table( - TableName="test-table" - )["Table"] + table = session.client("dynamodb").describe_table(TableName="test-table")["Table"] - assert table["KeySchema"] == [ - {"AttributeName": "pk", "KeyType": "HASH"} - ] - assert table["AttributeDefinitions"] == [ - {"AttributeName": "pk", "AttributeType": "S"} - ] + assert table["KeySchema"] == [{"AttributeName": "pk", "KeyType": "HASH"}] + assert table["AttributeDefinitions"] == [{"AttributeName": "pk", "AttributeType": "S"}] def test_ensure_index_marks_text_mirror_non_filterable(): @@ -60,9 +54,7 @@ def test_ensure_index_marks_text_mirror_non_filterable(): dataType="float32", dimension=4, distanceMetric="cosine", - metadataConfiguration={ - "nonFilterableMetadataKeys": [TEXT_METADATA_KEY] - }, + metadataConfiguration={"nonFilterableMetadataKeys": [TEXT_METADATA_KEY]}, ) diff --git a/tests/test_quantization.py b/tests/test_quantization.py index a7946a1..e460cc4 100644 --- a/tests/test_quantization.py +++ b/tests/test_quantization.py @@ -196,6 +196,7 @@ def test_load_corrupted_file_raises(tmp_path): with pytest.raises(ValueError): ProductQuantizer.load(corrupt) + def test_opq_rotation_fit(): rng = np.random.default_rng(42) diff --git a/tests/test_retrievers.py b/tests/test_retrievers.py index 8c30af3..553206e 100644 --- a/tests/test_retrievers.py +++ b/tests/test_retrievers.py @@ -365,7 +365,9 @@ def test_hyde_requires_embedder(monkeypatch): monkeypatch.setattr(client_mod, "S3VectorsStore", FakeS3) monkeypatch.setattr(client_mod, "DynamoDBStore", FakeDDB) monkeypatch.setattr(client_mod, "GraphStore", FakeGraph) - bare = Dynavec(DynavecConfig(vector_bucket="b", index="i", table="t", dimension=4, region="us-east-1")) + bare = Dynavec( + DynavecConfig(vector_bucket="b", index="i", table="t", dimension=4, region="us-east-1") + ) with pytest.raises(ConfigurationError): HyDERetriever(bare, lambda q: "x") diff --git a/tests/test_s3vectors.py b/tests/test_s3vectors.py index 32e1797..a320e1a 100644 --- a/tests/test_s3vectors.py +++ b/tests/test_s3vectors.py @@ -8,38 +8,37 @@ class FakeS3VectorsClient: - def __init__(self): - self.calls = [] + def __init__(self): + self.calls = [] - def get_vectors(self, **kwargs): - self.calls.append(kwargs) - return { - "vectors": [ - {"key": key, "data": {"float32": [1.0, 2.0]}} - for key in kwargs["keys"] - ] - } + def get_vectors(self, **kwargs): + self.calls.append(kwargs) + return { + "vectors": [{"key": key, "data": {"float32": [1.0, 2.0]}} for key in kwargs["keys"]] + } def test_get_vectors_batches_keys_and_merges_results(): - fake = FakeS3VectorsClient() - store = S3VectorsStore.__new__(S3VectorsStore) - store._config = DynavecConfig(vector_bucket="bucket", index="index", table="table", dimension=2) - store._client = fake + fake = FakeS3VectorsClient() + store = S3VectorsStore.__new__(S3VectorsStore) + store._config = DynavecConfig(vector_bucket="bucket", index="index", table="table", dimension=2) + store._client = fake - keys = [f"key-{i}" for i in range(250)] - result = store.get_vectors(keys) + keys = [f"key-{i}" for i in range(250)] + result = store.get_vectors(keys) - assert [len(call["keys"]) for call in fake.calls] == [100, 100, 50] - assert all(call["returnData"] is True for call in fake.calls) - assert len(result) == 250 - assert set(result) == set(keys) + assert [len(call["keys"]) for call in fake.calls] == [100, 100, 50] + assert all(call["returnData"] is True for call in fake.calls) + assert len(result) == 250 + assert set(result) == set(keys) def _config(**kwargs): - defaults = dict(vector_bucket="test-bucket", index="test-index", table="test-table", dimension=4) - defaults.update(kwargs) - return DynavecConfig(**defaults) + defaults = dict( + vector_bucket="test-bucket", index="test-index", table="test-table", dimension=4 + ) + defaults.update(kwargs) + return DynavecConfig(**defaults) def _store_with_pages(pages, config=None): @@ -56,17 +55,18 @@ def _store_with_pages(pages, config=None): def test_query_single_page(): - page = [{"key": f"k{i}", "distance": 0.01 * i} for i in range(50)] - store, paginator = _store_with_pages([{"vectors": page}]) + page = [{"key": f"k{i}", "distance": 0.01 * i} for i in range(50)] + store, paginator = _store_with_pages([{"vectors": page}]) - results = store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=50) + results = store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=50) + + assert len(results) == 50 + assert results[0]["key"] == "k0" + assert results[-1]["key"] == "k49" + store._client.get_paginator.assert_called_once_with("query_vectors") + assert paginator.paginate.call_args.kwargs["PaginationConfig"] == {"MaxItems": 50} + assert paginator.paginate.call_args.kwargs["topK"] == 50 - assert len(results) == 50 - assert results[0]["key"] == "k0" - assert results[-1]["key"] == "k49" - store._client.get_paginator.assert_called_once_with("query_vectors") - assert paginator.paginate.call_args.kwargs["PaginationConfig"] == {"MaxItems": 50} - assert paginator.paginate.call_args.kwargs["topK"] == 50 def test_query_multi_page_pagination(): pages = [ @@ -98,47 +98,46 @@ def test_query_multi_page_pagination(): "k249", ] + def test_query_fewer_results_than_top_k(): - store, _ = _store_with_pages([{"vectors": [{"key": "k0"}, {"key": "k1"}]}]) + store, _ = _store_with_pages([{"vectors": [{"key": "k0"}, {"key": "k1"}]}]) - results = store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=50) + results = store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=50) - assert len(results) == 2 + assert len(results) == 2 def test_query_zero_or_negative_top_k(): - store, _ = _store_with_pages([{"vectors": [{"key": "k0"}]}]) + store, _ = _store_with_pages([{"vectors": [{"key": "k0"}]}]) - assert store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=0) == [] - assert store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=-5) == [] - store._client.get_paginator.assert_not_called() + assert store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=0) == [] + assert store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=-5) == [] + store._client.get_paginator.assert_not_called() def test_query_exceeding_service_limit(): - store, _ = _store_with_pages([]) + store, _ = _store_with_pages([]) - with pytest.raises(ValueError, match="exceeds Amazon S3 Vectors maximum limit"): - store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=_MAX_TOP_K + 1) + with pytest.raises(ValueError, match="exceeds Amazon S3 Vectors maximum limit"): + store.query(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=_MAX_TOP_K + 1) def test_query_pages_invalid_page_size(): - store, _ = _store_with_pages([]) + store, _ = _store_with_pages([]) - with pytest.raises(ValueError, match="page_size must be a positive integer"): - list(store.query_pages(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=10, page_size=0)) + with pytest.raises(ValueError, match="page_size must be a positive integer"): + list(store.query_pages(query_vector=[0.1, 0.2, 0.3, 0.4], top_k=10, page_size=0)) def test_config_rejects_non_positive_top_k_page_size(): - with pytest.raises(ValueError, match="top_k_page_size must be a positive integer"): - _config(top_k_page_size=0) + with pytest.raises(ValueError, match="top_k_page_size must be a positive integer"): + _config(top_k_page_size=0) + def test_put_vectors_parallelization(): store, _ = _store_with_pages([]) - vectors = [ - (f"key_{i}", [0.1] * 128, {"tag": "test"}) - for i in range(1500) - ] + vectors = [(f"key_{i}", [0.1] * 128, {"tag": "test"}) for i in range(1500)] def mock_put_batch(payload): time.sleep(0.1) @@ -156,10 +155,7 @@ def mock_put_batch(payload): def test_put_vectors_error_propagation(): store, _ = _store_with_pages([]) - vectors = [ - (f"key_{i}", [0.1] * 128, {"tag": "test"}) - for i in range(1000) - ] + vectors = [(f"key_{i}", [0.1] * 128, {"tag": "test"}) for i in range(1000)] def mock_put_batch_with_error(payload): if payload[0]["key"] == "key_0": @@ -242,9 +238,7 @@ def test_query_pages_raw_chunks(): def test_query_pages_configurable_page_size(): - store, _ = _store_with_pages( - [{"vectors": [{"key": f"k{i}"} for i in range(25)]}] - ) + store, _ = _store_with_pages([{"vectors": [{"key": f"k{i}"} for i in range(25)]}]) result_pages = list( store.query_pages( @@ -324,4 +318,4 @@ def test_query_pages_does_not_rate_limit_after_paginator_exhaustion(): ) assert [len(page) for page in result_pages] == [1] - assert limiter.acquire.call_count == 1 \ No newline at end of file + assert limiter.acquire.call_count == 1 diff --git a/tests/test_telemetry.py b/tests/test_telemetry.py index 8d52c9a..8880a8e 100644 --- a/tests/test_telemetry.py +++ b/tests/test_telemetry.py @@ -78,7 +78,7 @@ def test_aggregate_cache_hit_rate_and_errors(): for e in evs: e.ts = now - 1 agg = aggregate(evs, window_seconds=3600, now=now) - assert agg["cache_total"] == 3 # None excluded + assert agg["cache_total"] == 3 # None excluded assert agg["cache_hits"] == 2 assert agg["cache_hit_rate"] == round(100 * 2 / 3, 1) assert agg["error_rate"] == round(100 * 1 / 5, 2) diff --git a/tests/test_utils_transforms.py b/tests/test_utils_transforms.py index b6931a1..567e664 100644 --- a/tests/test_utils_transforms.py +++ b/tests/test_utils_transforms.py @@ -76,6 +76,7 @@ def test_as_pipeline_coercions(): multi = as_pipeline([lambda c: c, lambda c: c]) assert len(multi) == 2 + def test_retry_uses_retry_delay(monkeypatch): calls = {"n": 0} delays = [] diff --git a/tools/build_docs.py b/tools/build_docs.py index 1006bdd..2e86970 100644 --- a/tools/build_docs.py +++ b/tools/build_docs.py @@ -13,42 +13,60 @@ # Sidebar structure: (group title, [(slug, nav label)]) NAV = [ - ("Getting started", [ - ("index", "Overview"), - ("installation", "Installation"), - ("quickstart", "Quickstart"), - ("configuration", "Configuration"), - ("faq", "FAQ"), - ]), - ("Writing data", [ - ("embeddings", "Embeddings"), - ("upsert", "Upsert"), - ("update-and-lambda", "Update & Lambda"), - ("ingestion", "Ingestion & MCP"), - ]), - ("Searching", [ - ("search", "Search"), - ("metrics-and-rerank", "Metrics & rerank"), - ("namespaces", "Namespaces"), - ("streaming", "Streaming"), - ("caching", "Caching"), - ]), - ("Advanced", [ - ("knowledge-graph", "Knowledge graph"), - ("hot-tier", "In-memory hot tier"), - ("quantization", "Product quantization"), - ("concurrency", "Concurrency"), - ("credentials", "Credentials & IAM"), - ]), - ("Ecosystem", [ - ("ecosystem", "Ecosystem overview"), - ("integrations", "Framework integrations"), - ("dashboard", "Telemetry dashboard"), - ("benchmarking", "Benchmarking"), - ]), - ("About", [ - ("release-notes", "Release notes"), - ]), + ( + "Getting started", + [ + ("index", "Overview"), + ("installation", "Installation"), + ("quickstart", "Quickstart"), + ("configuration", "Configuration"), + ("faq", "FAQ"), + ], + ), + ( + "Writing data", + [ + ("embeddings", "Embeddings"), + ("upsert", "Upsert"), + ("update-and-lambda", "Update & Lambda"), + ("ingestion", "Ingestion & MCP"), + ], + ), + ( + "Searching", + [ + ("search", "Search"), + ("metrics-and-rerank", "Metrics & rerank"), + ("namespaces", "Namespaces"), + ("streaming", "Streaming"), + ("caching", "Caching"), + ], + ), + ( + "Advanced", + [ + ("knowledge-graph", "Knowledge graph"), + ("hot-tier", "In-memory hot tier"), + ("quantization", "Product quantization"), + ("concurrency", "Concurrency"), + ("credentials", "Credentials & IAM"), + ], + ), + ( + "Ecosystem", + [ + ("ecosystem", "Ecosystem overview"), + ("integrations", "Framework integrations"), + ("dashboard", "Telemetry dashboard"), + ("benchmarking", "Benchmarking"), + ], + ), + ( + "About", + [ + ("release-notes", "Release notes"), + ], + ), ] # flat order for prev/next @@ -105,7 +123,8 @@ def code(src: str) -> str: # ---- page bodies (slug -> (title, subtitle, html)) ---- PAGES: dict[str, tuple[str, str, str]] = {} -PAGES["index"] = ("dynavec documentation", +PAGES["index"] = ( + "dynavec documentation", "A serverless hybrid vector database that runs inside your own AWS account.", """

dynavec fuses Amazon DynamoDB (single-digit-millisecond @@ -130,13 +149,16 @@ def code(src: str) -> str:

Benchmarking

Recall, latency, and cost modeled across dimensions and scale.

Ecosystem

The four pillars — dynavec, dynaflow, dynalogs, dynaevals — and the roadmap.

 -""") +""", +) -PAGES["installation"] = ("Installation", +PAGES["installation"] = ( + "Installation", "Base install is boto3 + numpy. Everything else is an optional extra.", """

Install with pip or uv — both pull from PyPI.

-""" + code("""# base (boto3 + numpy only) +""" + + code("""# base (boto3 + numpy only) pip install dynavec uv add dynavec @@ -146,7 +168,8 @@ def code(src: str) -> str: # everything (all embedders + framework adapters) pip install "dynavec[all]" -""") + """ +""") + + """

Optional extras

@@ -159,14 +182,17 @@ def code(src: str) -> str:
ExtraAdds
dynavec needs AWS credentials with permission for S3 Vectors and DynamoDB. See Credentials & IAM.
-""") +""", +) -PAGES["quickstart"] = ("Quickstart", +PAGES["quickstart"] = ( + "Quickstart", "Provision, upsert, and search in a dozen lines.", """

With auto_provision=True, dynavec creates the S3 vector bucket, the vector index, and the DynamoDB table on first use.

-""" + code("""from dynavec import Dynavec, DynavecConfig, Document +""" + + code("""from dynavec import Dynavec, DynavecConfig, Document from dynavec.embeddings import OpenAIEmbedder cfg = DynavecConfig( @@ -186,17 +212,21 @@ def code(src: str) -> str: for hit in db.search("how do cells make energy?", top_k=3): print(hit.score, hit.id, hit.text) -""") + """ +""") + + """
S3 Vectors is eventually consistent right after ingest — allow a few seconds before querying freshly written vectors.
-""") +""", +) -PAGES["configuration"] = ("Configuration", +PAGES["configuration"] = ( + "Configuration", "Everything the client needs, in one frozen dataclass.", """

DynavecConfig is an immutable description of your resources and tuning. It never holds secrets — credentials are passed separately (see Credentials & IAM).

-""" + code("""from dynavec import DynavecConfig +""" + + code("""from dynavec import DynavecConfig cfg = DynavecConfig( vector_bucket="my-vectors", # S3 vector bucket name @@ -214,7 +244,8 @@ def code(src: str) -> str: structured_logging=True, # opt-in structured JSON logging across stores auto_provision=True, ) -""") + """ +""") + + """

Key parameters

@@ -229,9 +260,11 @@ def code(src: str) -> str:
FieldMeaning
structured_loggingOpt-in structured JSON logging across DynamoDB and S3 Vectors store operations with automatic secret redaction. Defaults to False.
log_levelMinimum log level ("DEBUG", "INFO", etc.). Defaults to "INFO".
-""") +""", +) -PAGES["faq"] = ("Frequently Asked Questions", +PAGES["faq"] = ( + "Frequently Asked Questions", "Common questions about regions, limits, consistency, costs, and architecture.", """

Regions & Availability

@@ -298,14 +331,17 @@ def code(src: str) -> str:

How does multi-tenancy work?

dynavec provides native Namespaces. A single S3 vector bucket and DynamoDB table can host many independent tenants. DynamoDB partition keys are cleanly isolated via escaped "{namespace}#{id}" prefixes, and vector queries are automatically scoped so data never leaks across namespaces.

-""") +""", +) -PAGES["embeddings"] = ("Embeddings", +PAGES["embeddings"] = ( + "Embeddings", "Pluggable, bring-your-own-key — or bring your own vectors.", """

Choose an embedder and supply your own API key, or skip the embedder entirely and pass pre-computed vectors. Embedder backends are imported lazily, so the base install stays light.

-""" + code("""from dynavec.embeddings import ( +""" + + code("""from dynavec.embeddings import ( OpenAIEmbedder, GeminiEmbedder, MistralEmbedder, BedrockEmbedder, SentenceTransformerEmbedder, ) @@ -318,33 +354,40 @@ def code(src: str) -> str: # in-account (no third party) or fully local / offline emb = BedrockEmbedder(model_id="amazon.titan-embed-text-v2:0", region="us-east-1") emb = SentenceTransformerEmbedder(model="all-MiniLM-L6-v2") # 384-d, free -""") + """ +""") + + """

Bring your own vectors

No embedder needed — pass vectors directly and query with a vector.

-""" + code("""from dynavec import Dynavec, DynavecConfig, Document +""" + + code("""from dynavec import Dynavec, DynavecConfig, Document db = Dynavec(cfg) # no embedder db.upsert([Document(id="x", vector=my_1536d_vector, metadata={"lang": "en"})]) db.search(vector=my_query_vector, top_k=5) -""") + """ +""") + + """
Compliance tip: use BedrockEmbedder or SentenceTransformerEmbedder to keep embedding in-account or offline — no data leaves your environment.
-""") +""", +) -PAGES["upsert"] = ("Upsert", +PAGES["upsert"] = ( + "Upsert", "Write documents to both stores in one call.", """

Each document carries an id, either text (which gets embedded) or a vector, and optional metadata. dynavec splits metadata: a small filterable subset goes to S3 Vectors, the full copy plus text goes to DynamoDB.

-""" + code("""from dynavec import Document +""" + + code("""from dynavec import Document db.upsert([ Document(id="1", text="apple pie recipe", metadata={"cat": "food", "rating": 5}), Document(id="2", text="rocket launch schedule", metadata={"cat": "space"}), ], namespace="kb", auto_metadata=True) -""") + """ +""") + + """

The metadata switch