diff --git a/examples/openai_assistant_tool.py b/examples/openai_assistant_tool.py new file mode 100644 index 0000000..a43707e --- /dev/null +++ b/examples/openai_assistant_tool.py @@ -0,0 +1,90 @@ +"""Use dynavec as an OpenAI Assistant or Function Calling tool. + +This adapter allows an OpenAI Assistant or Chat Completions model to perform +vector retrieval over your Dynavec knowledge base directly without paying for +OpenAI's hosted vector store files. + +Requirements: + pip install "dynavec[sentence-transformers]" + python examples/openai_assistant_tool.py +""" + +import json +from types import SimpleNamespace + +from dynavec import Document, Dynavec, DynavecConfig +from dynavec.embeddings import SentenceTransformerEmbedder +from dynavec.integrations.tools import as_openai_tool + + +def main() -> None: + # 1. Initialize Dynavec with a local/in-memory configuration + embedder = SentenceTransformerEmbedder(model="all-MiniLM-L6-v2") + + cfg = DynavecConfig( + vector_bucket="dynavec-demo", + index="openai-assistant-docs", + table="dynavec_assistant_kb", + dimension=embedder.dimension, + region="us-east-1", + auto_provision=True, + ) + + db = Dynavec(cfg, embedder=embedder) + + # 2. Populate knowledge base + db.upsert( + [ + Document( + id="doc-1", + text="Dynavec is a serverless vector database running on AWS DynamoDB and S3 Vectors.", + ), + Document( + id="doc-2", + text="Dynavec provides sub-50ms hybrid keyword and vector retrieval at zero idle cost.", + ), + Document( + id="doc-3", + text="OpenAI Assistants can call custom external retrieval tools during run loops.", + ), + ], + namespace="docs", + ) + + # 3. Create the OpenAI Assistant Tool adapter + tool = as_openai_tool( + db, + name="dynavec_search", + description="Search documentation and return relevant knowledge base passages.", + namespace="docs", + top_k=2, + ) + + print("=== OpenAI Tool Schema ===") + print(json.dumps(tool.schema, indent=2)) + print() + + # 4. Simulate an Assistant Run loop requesting a tool call + # In real usage with OpenAI SDK: + # run = client.beta.threads.runs.create_and_poll(...) + # if run.status == "requires_action": + # tool_outputs = tool.submit_tool_outputs(run.required_action.submit_tool_outputs.tool_calls) + # client.beta.threads.runs.submit_tool_outputs(..., tool_outputs=tool_outputs) + + simulated_tool_call = SimpleNamespace( + id="call_mock_12345", + type="function", + function=SimpleNamespace( + name="dynavec_search", + arguments=json.dumps({"query": "serverless vector database cost"}), + ), + ) + + print("=== Simulating Tool Execution for Assistant Run ===") + outputs = tool.submit_tool_outputs([simulated_tool_call]) + print("Submitted Tool Output:") + print(json.dumps(outputs, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/src/dynavec/integrations/__init__.py b/src/dynavec/integrations/__init__.py index 7a2e205..87e9c90 100644 --- a/src/dynavec/integrations/__init__.py +++ b/src/dynavec/integrations/__init__.py @@ -1 +1,17 @@ -"""Framework integrations (LangChain today; LlamaIndex/CrewAI to follow).""" +"""Framework integrations (OpenAI Assistants, LangChain, CrewAI, DSPy, LlamaIndex).""" + +from .tools import ( + OpenAIAssistantTool, + as_crewai_tool, + as_langchain_tool, + as_openai_tool, + make_retriever_fn, +) + +__all__ = [ + "OpenAIAssistantTool", + "as_openai_tool", + "as_langchain_tool", + "as_crewai_tool", + "make_retriever_fn", +] diff --git a/src/dynavec/integrations/tools.py b/src/dynavec/integrations/tools.py index 50736d3..1f82719 100644 --- a/src/dynavec/integrations/tools.py +++ b/src/dynavec/integrations/tools.py @@ -1,13 +1,15 @@ """Framework-agnostic retrieval tool factory. -Most agent frameworks (LangGraph, CrewAI, Strands, plain function-calling) just -need a **callable** that takes a query string and returns text. This builds one -from a dynavec client / namespace, plus thin wrappers that register it as a -native tool in LangChain and CrewAI when those are installed. +Most agent frameworks (OpenAI Assistants, LangGraph, CrewAI, Strands, plain +function-calling) just need a **callable** that takes a query string and +returns text. This module provides ``make_retriever_fn`` plus first-class +adapters for OpenAI Assistants, LangChain, and CrewAI. """ from __future__ import annotations +import json +from collections.abc import Sequence from typing import Any, Callable from ..client import Dynavec @@ -53,6 +55,166 @@ def retrieve(query: str) -> str: return retrieve +class OpenAIAssistantTool: + """OpenAI Assistants and Function Calling tool adapter backed by Dynavec. + + Provides OpenAI-compatible function calling schemas, single-tool-call execution, + and batch execution for Assistant run loops without requiring external SDKs. + """ + + def __init__( + self, + retriever_fn: Callable[[str], str], + name: str = "dynavec_search", + description: str | None = None, + ) -> None: + self.retriever_fn = retriever_fn + self.name = name + self.description = description or ( + "Search the knowledge base and return the most relevant passages for a query." + ) + + @property + def schema(self) -> dict[str, Any]: + """Return the standard OpenAI Tool definition dictionary.""" + return { + "type": "function", + "function": { + "name": self.name, + "description": self.description, + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query to retrieve relevant passages.", + } + }, + "required": ["query"], + "additionalProperties": False, + }, + }, + } + + def to_openai_tool(self) -> dict[str, Any]: + """Return the OpenAI Tool definition dictionary (alias for ``schema``).""" + return self.schema + + def to_dict(self) -> dict[str, Any]: + """Return the OpenAI Tool definition dictionary.""" + return self.schema + + def __call__(self, query: str) -> str: + """Directly invoke the retriever with a query string.""" + return self.retriever_fn(query) + + def handle_tool_call(self, tool_call: Any) -> dict[str, str]: + """Execute retrieval for a single OpenAI ToolCall and return output dict. + + Accepts either an OpenAI SDK ``ToolCall`` instance (or Pydantic/Namespace model) + or a dictionary format: + ``{"id": "call_123", "function": {"name": "dynavec_search", "arguments": "{\\"query\\": \\"...\\"}"}}`` + + Returns: + A dictionary formatted for OpenAI ``submit_tool_outputs``: + ``{"tool_call_id": "", "output": ""}`` + """ + call_id = getattr(tool_call, "id", None) + if call_id is None and isinstance(tool_call, dict): + call_id = tool_call.get("id") or tool_call.get("tool_call_id") or "" + call_id = str(call_id or "") + + fn_obj = getattr(tool_call, "function", None) + if fn_obj is None and isinstance(tool_call, dict): + fn_obj = tool_call.get("function") + + raw_args = getattr(fn_obj, "arguments", None) + if raw_args is None and isinstance(fn_obj, dict): + raw_args = fn_obj.get("arguments") + if raw_args is None and isinstance(tool_call, dict): + raw_args = tool_call.get("arguments") + + try: + if isinstance(raw_args, str): + try: + parsed_args = json.loads(raw_args) + except (json.JSONDecodeError, TypeError): + parsed_args = {"query": raw_args} + elif isinstance(raw_args, dict): + parsed_args = raw_args + elif raw_args is None: + parsed_args = {} + else: + parsed_args = {"query": str(raw_args)} + + if parsed_args is None: + query = "" + elif isinstance(parsed_args, dict): + query = parsed_args.get("query") + if query is None: + for fallback_key in ("q", "input", "search_query", "text", "prompt"): + if fallback_key in parsed_args: + query = parsed_args[fallback_key] + break + if query is None and parsed_args: + query = next((v for v in parsed_args.values() if isinstance(v, str)), "") + query = str(query or "") + elif isinstance(parsed_args, (list, tuple, set)): + query = " ".join(str(x) for x in parsed_args) + else: + query = str(parsed_args) + + output = self.retriever_fn(query) + except Exception as exc: + output = f"Error executing tool '{self.name}': {exc}" + + return {"tool_call_id": call_id, "output": output} + + def submit_tool_outputs(self, tool_calls: Sequence[Any] | None) -> list[dict[str, str]]: + """Batch-process multiple tool calls from an assistant run loop. + + Filters calls matching ``self.name``, executes retrieval for each, and + returns a list of ``{"tool_call_id": ..., "output": ...}`` objects + ready to be passed to ``client.beta.threads.runs.submit_tool_outputs()``. + """ + if not tool_calls: + return [] + + outputs: list[dict[str, str]] = [] + for call in tool_calls: + fn_obj = getattr(call, "function", None) + if fn_obj is None and isinstance(call, dict): + fn_obj = call.get("function") + + fn_name = getattr(fn_obj, "name", None) + if fn_name is None and isinstance(fn_obj, dict): + fn_name = fn_obj.get("name") + if fn_name is None and isinstance(call, dict): + fn_name = call.get("name") + + if fn_name and fn_name != self.name: + continue + + outputs.append(self.handle_tool_call(call)) + return outputs + + +def as_openai_tool( + source: Dynavec | NamespaceView, + *, + name: str = "dynavec_search", + description: str | None = None, + **kw: Any, +) -> OpenAIAssistantTool: + """Wrap a Dynavec client or NamespaceView as an OpenAI Assistant tool. + + Requires zero third-party dependencies. Returns an ``OpenAIAssistantTool`` + ready to be registered on an Assistant or passed into chat completions. + """ + fn = make_retriever_fn(source, **kw) + return OpenAIAssistantTool(fn, name=name, description=description) + + def as_langchain_tool(source, *, name: str = "dynavec_search", **kw) -> Any: """Wrap the retriever as a LangChain ``StructuredTool`` (requires langchain-core).""" from ..exceptions import MissingDependencyError @@ -63,9 +225,7 @@ def as_langchain_tool(source, *, name: str = "dynavec_search", **kw) -> Any: raise MissingDependencyError("as_langchain_tool", "langchain-core", "langchain") from exc fn = make_retriever_fn(source, **kw) - return StructuredTool.from_function( - func=fn, name=name, description=fn.__doc__ - ) + return StructuredTool.from_function(func=fn, name=name, description=fn.__doc__) def as_crewai_tool(source, *, name: str = "dynavec_search", **kw) -> Any: diff --git a/tests/test_openai_tool.py b/tests/test_openai_tool.py new file mode 100644 index 0000000..9d9bb31 --- /dev/null +++ b/tests/test_openai_tool.py @@ -0,0 +1,293 @@ +"""Tests for the OpenAI Assistant and Function Calling tool adapter. + +These tests run 100% offline with zero external network or API key dependencies. +""" + +from types import SimpleNamespace + +from dynavec.integrations.tools import OpenAIAssistantTool, as_openai_tool +from dynavec.models import SearchResult +from dynavec.namespace import NamespaceView + + +class _FakeClient: + """Stands in for Dynavec client .search().""" + + def __init__(self): + self.calls = [] + + def search(self, query, **kwargs): + self.calls.append((query, kwargs)) + return [ + SearchResult(id="doc-1", score=0.92, text="Dynavec stores vectors on DynamoDB & S3."), + SearchResult(id="doc-2", score=0.85, text="Serverless hybrid search is fast."), + ] + + +def test_schema_structure_default(): + client = _FakeClient() + tool = as_openai_tool(client) + + schema = tool.schema + assert schema["type"] == "function" + assert schema["function"]["name"] == "dynavec_search" + assert "Search the knowledge base" in schema["function"]["description"] + + params = schema["function"]["parameters"] + assert params["type"] == "object" + assert "query" in params["properties"] + assert params["properties"]["query"]["type"] == "string" + assert params["required"] == ["query"] + assert params["additionalProperties"] is False + + # Check aliases + assert tool.to_openai_tool() == schema + assert tool.to_dict() == schema + + +def test_schema_structure_custom_name_and_desc(): + client = _FakeClient() + tool = as_openai_tool( + client, + name="custom_search", + description="Search customer knowledge base", + ) + + schema = tool.schema + assert schema["function"]["name"] == "custom_search" + assert schema["function"]["description"] == "Search customer knowledge base" + + +def test_direct_callable_execution(): + client = _FakeClient() + tool = as_openai_tool(client, top_k=2, namespace="kb") + + output = tool("vector indexing") + assert "Dynavec stores vectors" in output + assert "Serverless hybrid search" in output + assert client.calls == [ + ( + "vector indexing", + {"top_k": 2, "namespace": "kb", "filter": None, "rescore": None}, + ) + ] + + +def test_handle_tool_call_object_sdk_style(): + client = _FakeClient() + tool = as_openai_tool(client) + + # Mimics openai.types.beta.threads.runs.RequiredActionFunctionToolCall + call_obj = SimpleNamespace( + id="call_abc123", + type="function", + function=SimpleNamespace( + name="dynavec_search", + arguments='{"query": "dynamodb storage"}', + ), + ) + + output = tool.handle_tool_call(call_obj) + assert output["tool_call_id"] == "call_abc123" + assert "Dynavec stores vectors" in output["output"] + + +def test_handle_tool_call_dict_format(): + client = _FakeClient() + tool = as_openai_tool(client) + + call_dict = { + "id": "call_xyz789", + "type": "function", + "function": { + "name": "dynavec_search", + "arguments": '{"query": "pricing model"}', + }, + } + + output = tool.handle_tool_call(call_dict) + assert output["tool_call_id"] == "call_xyz789" + assert "Dynavec stores vectors" in output["output"] + + +def test_handle_tool_call_dict_arguments(): + """Some orchestrators (LiteLLM, Autogen) pass pre-parsed dicts in arguments.""" + client = _FakeClient() + tool = as_openai_tool(client) + + call_dict = { + "id": "call_parsed", + "function": { + "name": "dynavec_search", + "arguments": {"query": "preparsed query"}, + }, + } + + output = tool.handle_tool_call(call_dict) + assert output["tool_call_id"] == "call_parsed" + assert "Dynavec stores vectors" in output["output"] + + +def test_handle_tool_call_fallback_argument_keys(): + client = _FakeClient() + tool = as_openai_tool(client) + + call_dict = { + "id": "call_alt", + "function": { + "name": "dynavec_search", + "arguments": '{"search_query": "alternative key test"}', + }, + } + + output = tool.handle_tool_call(call_dict) + assert output["tool_call_id"] == "call_alt" + assert client.calls[-1][0] == "alternative key test" + + +def test_handle_tool_call_malformed_json_graceful(): + client = _FakeClient() + tool = as_openai_tool(client) + + call_dict = { + "id": "call_bad_json", + "function": { + "name": "dynavec_search", + "arguments": "unquoted raw string query", + }, + } + + # Should gracefully use the raw string rather than crashing + output = tool.handle_tool_call(call_dict) + assert output["tool_call_id"] == "call_bad_json" + assert client.calls[-1][0] == "unquoted raw string query" + + +def test_handle_tool_call_retriever_exception_safe(): + def failing_retriever(query: str) -> str: + raise RuntimeError("Database connection timeout") + + tool = OpenAIAssistantTool(failing_retriever, name="dynavec_search") + call_dict = { + "id": "call_err", + "function": {"name": "dynavec_search", "arguments": '{"query": "error"}'}, + } + + output = tool.handle_tool_call(call_dict) + assert output["tool_call_id"] == "call_err" + assert "Error executing tool 'dynavec_search': Database connection timeout" in output["output"] + + +def test_submit_tool_outputs_batch(): + client = _FakeClient() + tool = as_openai_tool(client, name="dynavec_search") + + tool_calls = [ + SimpleNamespace( + id="call_1", + function=SimpleNamespace(name="dynavec_search", arguments='{"query": "first query"}'), + ), + SimpleNamespace( + id="call_2", + function=SimpleNamespace(name="dynavec_search", arguments='{"query": "second query"}'), + ), + ] + + outputs = tool.submit_tool_outputs(tool_calls) + assert len(outputs) == 2 + assert outputs[0]["tool_call_id"] == "call_1" + assert outputs[1]["tool_call_id"] == "call_2" + + +def test_submit_tool_outputs_filters_foreign_tools(): + client = _FakeClient() + tool = as_openai_tool(client, name="dynavec_search") + + tool_calls = [ + {"id": "call_weather", "function": {"name": "get_weather", "arguments": "{}"}}, + { + "id": "call_dynavec", + "function": {"name": "dynavec_search", "arguments": '{"query": "architecture"}'}, + }, + {"id": "call_calc", "function": {"name": "calculator", "arguments": "{}"}}, + ] + + outputs = tool.submit_tool_outputs(tool_calls) + assert len(outputs) == 1 + assert outputs[0]["tool_call_id"] == "call_dynavec" + + +def test_namespace_view_support(): + client = _FakeClient() + ns_view = NamespaceView(client, "billing") + tool = as_openai_tool(ns_view, top_k=3, include_scores=True) + + output = tool("invoice questions") + assert "[0.920] Dynavec stores vectors on DynamoDB & S3." in output + assert client.calls == [ + ( + "invoice questions", + {"top_k": 3, "namespace": "billing", "filter": None, "rescore": None}, + ) + ] + + +def test_handle_tool_call_null_and_none_arguments(): + client = _FakeClient() + tool = as_openai_tool(client) + + call_dict = { + "id": "call_null", + "function": {"name": "dynavec_search", "arguments": "null"}, + } + output = tool.handle_tool_call(call_dict) + assert output["tool_call_id"] == "call_null" + assert client.calls[-1][0] == "" + + +def test_handle_tool_call_list_arguments(): + client = _FakeClient() + tool = as_openai_tool(client) + + call_dict = { + "id": "call_list", + "function": {"name": "dynavec_search", "arguments": '["serverless", "database"]'}, + } + output = tool.handle_tool_call(call_dict) + assert output["tool_call_id"] == "call_list" + assert client.calls[-1][0] == "serverless database" + + +def test_submit_tool_outputs_none_and_empty(): + client = _FakeClient() + tool = as_openai_tool(client) + + assert tool.submit_tool_outputs(None) == [] + assert tool.submit_tool_outputs([]) == [] + + +def test_as_openai_tool_forwarding_kwargs(): + client = _FakeClient() + rescore_marker = object() + tool = as_openai_tool( + client, + top_k=7, + namespace="custom_ns", + filter={"type": "pdf"}, + rescore=rescore_marker, + join=" --- ", + include_scores=True, + ) + + output = tool("search test") + assert " --- " in output + assert "[0.920]" in output + assert client.calls[-1] == ( + "search test", + { + "top_k": 7, + "namespace": "custom_ns", + "filter": {"type": "pdf"}, + "rescore": rescore_marker, + }, + )