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 9fea9b2..ee90c48 100644 --- a/src/dynavec/agents/base.py +++ b/src/dynavec/agents/base.py @@ -6,7 +6,7 @@ import json from collections.abc import Callable from dataclasses import dataclass, field -from typing import Any +from typing import Any, overload from ..chat.base import Tool, ToolCall @@ -28,9 +28,12 @@ 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 @@ -116,11 +119,12 @@ def __init__( parameters: dict[str, Any] | None = None, ) -> None: self.fn = fn - self.name = name or fn.__name__ - 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) + # 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) def to_chat_tool(self) -> Tool: """Convert to a dynavec.chat.Tool schema.""" @@ -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(): @@ -155,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 @@ -167,6 +171,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 +182,71 @@ 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]: +) -> Callable[[Callable[..., Any]], AgentTool]: ... + + +def tool( + fn_or_name: Callable[..., Any] | str | None = None, + *, + name: 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.""" + custom_name = name if name is not None else (fn_or_name if isinstance(fn_or_name, str) else None) + + if callable(fn_or_name): + fn = fn_or_name + tool_name = custom_name or getattr(fn, "__name__", "unnamed_tool") + return AgentTool( + fn=fn, + name=tool_name, + description=description, + parameters=parameters, + ) 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..ce06fe3 100644 --- a/src/dynavec/agents/react.py +++ b/src/dynavec/agents/react.py @@ -3,10 +3,12 @@ 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 ..exceptions import NodeInterrupt from .base import AgentResult, AgentStep, AgentTool DEFAULT_REACT_SYSTEM_PROMPT = """You are a helpful and precise reasoning agent. @@ -36,21 +38,52 @@ 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 +93,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 +109,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 +153,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 +165,45 @@ 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 +233,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/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 6db7f6a..1752908 100644 --- a/src/dynavec/exceptions.py +++ b/src/dynavec/exceptions.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import Any + class DynavecError(Exception): """Base class for all dynavec errors.""" @@ -68,3 +70,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[str, Any] | 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/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 4ee91cb..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", ) @@ -74,26 +72,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_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 new file mode 100644 index 0000000..c48f281 --- /dev/null +++ b/tests/test_hitl_interrupts.py @@ -0,0 +1,171 @@ +"""Tests for Human-in-the-loop / interrupts functionality in dynavec agents.""" +from __future__ import annotations + +import asyncio +import json +from collections.abc import AsyncIterator, Iterator +from typing import Any + +from dynavec.agents.base import interrupt, tool +from dynavec.agents.react import ReActAgent +from dynavec.chat.base import ChatChunk, ChatModel, ChatResult, Message, ToolCall + + +class DummyChatModel(ChatModel): + """Simple mock chat model for HITL tests.""" + + def __init__(self, responses: list[ChatResult]) -> None: + super().__init__() + self.responses = responses + self.call_count = 0 + + def invoke( + self, messages: list[Message], tools: list[Any] | None = None, **kwargs: Any + ) -> ChatResult: + res = self.responses[self.call_count] + self.call_count += 1 + return res + + async def ainvoke( + self, messages: list[Message], tools: list[Any] | None = None, **kwargs: Any + ) -> ChatResult: + return self.invoke(messages, tools, **kwargs) + + def stream( + self, messages: list[Message], tools: list[Any] | None = None, **kwargs: Any + ) -> Iterator[ChatChunk]: + yield ChatChunk(content="") + + async def astream( + self, messages: list[Message], tools: list[Any] | None = None, **kwargs: Any + ) -> AsyncIterator[ChatChunk]: + yield ChatChunk(content="") + + +@tool +def sensitive_action_tool(action: str) -> str: + """Action requiring human approval.""" + interrupt("approval_node", f"Action '{action}' requires approval. Confirm?") + return f"Action '{action}' executed successfully." + + +def test_interrupt_raised_and_caught() -> None: + mock_model = DummyChatModel( + responses=[ + ChatResult( + message=Message( + role="assistant", + content="", + tool_calls=[ + ToolCall( + id="tc_1", + name="sensitive_action_tool", + arguments=json.dumps({"action": "delete_database"}), + ) + ], + ) + ) + ] + ) + + agent = ReActAgent(model=mock_model, tools=[sensitive_action_tool]) + result = agent.run("Please delete the database.") + + assert result.finished is False + assert result.termination_reason == "interrupted" + assert result.interrupt_payload is not None + + +def test_resume_execution() -> None: + mock_model = DummyChatModel( + responses=[ + ChatResult( + message=Message( + role="assistant", + content="", + tool_calls=[ + ToolCall( + id="tc_1", + name="sensitive_action_tool", + arguments=json.dumps({"action": "delete_database"}), + ) + ], + ) + ), + ChatResult( + message=Message( + role="assistant", + content="Database action completed.", + ) + ), + ] + ) + + agent = ReActAgent(model=mock_model, tools=[sensitive_action_tool]) + result = agent.run("Please delete the database.") + + assert result.finished is False + assert result.interrupt_payload is not None + + payload = result.interrupt_payload + saved_messages = payload["messages"] + start_step = payload["start_step"] + + resume_agent = ReActAgent(model=mock_model, tools=[sensitive_action_tool]) + final_result = resume_agent.resume( + saved_messages, + start_step=start_step, + human_decision="approved", + ) + + assert final_result.finished is True + assert final_result.termination_reason == "completed" + assert "Database action completed" in final_result.output + + +def test_async_interrupt_and_resume() -> None: + async def _run_async() -> None: + mock_model = DummyChatModel( + responses=[ + ChatResult( + message=Message( + role="assistant", + content="", + tool_calls=[ + ToolCall( + id="tc_1", + name="sensitive_action_tool", + arguments=json.dumps({"action": "delete_database"}), + ) + ], + ) + ), + ChatResult( + message=Message( + role="assistant", + content="Deployment completed.", + ) + ), + ] + ) + + agent = ReActAgent(model=mock_model, tools=[sensitive_action_tool]) + result = await agent.arun("Deploy latest release.") + + assert result.finished is False + assert result.interrupt_payload is not None + + payload = result.interrupt_payload + assert payload is not None + + resume_agent = ReActAgent(model=mock_model, tools=[sensitive_action_tool]) + final_result = await resume_agent.aresume( + payload["messages"], + start_step=payload["start_step"], + human_decision="approved", + ) + + assert final_result.finished is True + assert final_result.termination_reason == "completed" + + 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:
Recall, latency, and cost modeled across dimensions and scale. The four pillars — dynavec, dynaflow, dynalogs, dynaevals — and the roadmap.Benchmarking
Ecosystem
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]" -""") + """ +""") + + """| Extra | Adds |
|---|
With auto_provision=True, dynavec creates the S3 vector bucket, the vector index, and the
DynamoDB table on first use.
DynavecConfig is an immutable description of your resources and tuning. It never holds secrets —
credentials are passed separately (see Credentials & IAM).
| Field | Meaning |
|---|---|
structured_logging | Opt-in structured JSON logging across DynamoDB and S3 Vectors store operations with automatic secret redaction. Defaults to False. |
log_level | Minimum log level ("DEBUG", "INFO", etc.). Defaults to "INFO". |
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.
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 -""") + """ +""") + + """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) -""") + """ +""") + + """BedrockEmbedder or
SentenceTransformerEmbedder to keep embedding in-account or offline — no data leaves your
environment.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.
ItemTooLargeError (with doc_id, namespace,
size_bytes, and limit_bytes) if one is too big. The whole call fails, so S3 Vectors
and DynamoDB never end up with half a batch.
-""" + code("""from dynavec import ItemTooLargeError
+"""
+ + code("""from dynavec import ItemTooLargeError
from dynavec.ingest import chunk_text
try:
@@ -372,72 +416,88 @@ def code(src: str) -> str:
for i, chunk in enumerate(chunk_text(long_text, chunk_size=2000))],
namespace="kb",
)
-"""))
+"""),
+)
-PAGES["update-and-lambda"] = ("Update & Lambda transforms",
+PAGES["update-and-lambda"] = (
+ "Update & Lambda transforms",
"Change text, vector, or metadata — and transform data in-account.",
"""
update() is a read-modify-write: metadata merges by default, and the vector is only
re-derived when the text changes or you pass a new vector.
Transforms are plain callables run on each document before it is written — for enrichment, redaction, or deriving vectors elsewhere.
-""" + code("""from dynavec.transforms import TransformPipeline +""" + + code("""from dynavec.transforms import TransformPipeline def redact(ctx): ctx.metadata["pii"] = False return ctx db.upsert(docs, transform=redact) # or transform=TransformPipeline([...]) -""") + """ +""") + + """LambdaTransform invokes a Lambda you own with the document payload and applies whatever it
returns — keeping custom logic in-account.
lambda:InvokeFunction on that function — see
Credentials & IAM.A source is any iterable of records. ingest() chunks, embeds, and upserts them.
MarkdownSource reads UTF-8 .md and .txt files recursively.
Install dynavec[ingest] for YAML front matter support.
A Markdown file may begin with a YAML mapping between two --- lines:
Front matter becomes metadata and is removed from the text before chunking. Use storage-compatible values (strings, numbers, booleans, lists, and mappings); quote dates to keep them as strings. Malformed or unclosed front matter raises an error naming the file. Text files are read verbatim.
@@ -450,20 +510,25 @@ def redact(ctx):MCPResourceSource turns an MCP server's resources (Notion, Confluence, Drive, your
own) into an embeddable corpus — no per-source code.
Chunk ids are "{record_id}#chunk{n}" with source_id / chunk
metadata, so you can group or delete a whole document later.
Provide a query string (embedded for you) or a raw vector. S3 Vectors returns the
nearest keys; dynavec hydrates the full documents from DynamoDB via BatchGetItem.
Filters use the S3 Vectors dialect — bare {"k": v} is equality; operators like
$gte, $in, $and, $or are supported. Only keys in
filterable_keys can be filtered.
Results come back as SearchResult with id, score (higher = more
similar), distance, text, and metadata. Refine ordering with
metrics & rerank, speed up repeats with
caching.
The S3 Vectors index metric is cosine or euclidean. On top of the returned candidates, dynavec can rescore with any metric — or a weighted combination — client-side.
-""" + code("""# single metric +""" + + code("""# single metric db.search("q", top_k=5, rescore="manhattan") # weighted combination (normalized per candidate set) db.search("q", top_k=5, rescore={"cosine": 0.7, "dot": 0.3}) -""") + """ +""") + + """Maximal Marginal Relevance balances relevance against diversity so results are not near-duplicates.
-""" + code("""db.search("q", top_k=5, rerank="mmr", mmr_lambda=0.5) # 1=relevance, 0=diversity -""") + """ +""" + + code("""db.search("q", top_k=5, rerank="mmr", mmr_lambda=0.5) # 1=relevance, 0=diversity +""") + + """Both over-fetch top_k * over_fetch candidates first. You can also fuse multiple result lists
with Reciprocal Rank Fusion:
A namespace gives each tenant an isolated view of the same dynavec index. Every vector is tagged with its @@ -515,7 +592,8 @@ def redact(ctx):
Use db.namespace("tenant-id") to create a lightweight handle bound to one tenant. The same
document ID can exist in multiple namespaces because the namespace is part of the storage key.
A typical RAG application maps the authenticated tenant to a namespace, retrieves context from that namespace, and passes only that context to the LLM.
-""" + code("""tenant = db.namespace(tenant_id) +""" + + code("""tenant = db.namespace(tenant_id) hits = tenant.search(user_query, top_k=5) context = "\\n\\n".join(hit.text for hit in hits) @@ -535,37 +615,47 @@ def redact(ctx): answer = llm.generate( f"Answer using only the following context:\\n{context}\\n\\nQuestion: {user_query}" ) -""") + """ +""") + + """The complete runnable recipe is available in
examples/multi_tenant_rag.py. It demonstrates two tenants using the same index, including identical
document IDs with tenant-specific content.
search_stream() is a generator: it yields one hit at a time as S3 Vectors paginates, so an agent can start
consuming the first results before the full set returns. Amazon S3 Vectors returns at most 100 vectors per response page;
dynavec follows nextToken up to the service limit of 10,000 results.
page_size (or DynavecConfig(top_k_page_size=50)) only changes how many hits are hydrated from DynamoDB per batch. It does not change the S3 Vectors page size (fixed at 100) or time-to-first-result.
Attach a cache and repeated queries are served without hitting S3 Vectors. Three backends:
-""" + code("""from dynavec import Dynavec, SemanticCache, DynamoDBCache, RedisCache +""" + + code("""from dynavec import Dynavec, SemanticCache, DynamoDBCache, RedisCache # 1) in-process semantic cache — also serves near-duplicate queries db = Dynavec(cfg, embedder=emb, cache=SemanticCache(threshold=0.97)) @@ -579,7 +669,8 @@ def redact(ctx): # 3) sub-millisecond shared cache on Redis / AWS ElastiCache db = Dynavec(cfg, embedder=emb, cache=RedisCache("redis://my-elasticache:6379/0")) -""") + """ +""") + + """| Backend | Best for |
|---|---|
SemanticCache | single process; tolerant of near-duplicate hits; zero infra |
RedisCache | many workers/hosts; lowest latency; AWS ElastiCache |
Force a fresh search per call with db.search(..., use_cache=False).
Alongside the vector index, dynavec keeps a lightweight entity-relationship graph in DynamoDB. Entities link to documents; you can traverse the graph first (cheap key lookups) to gather a candidate set, then rank only those against the query embedding. That is the DynamoDB → S3 Vectors reference join.
-""" + code("""# build the graph +""" + + code("""# build the graph db.graph_add_edge("acme", "competes_with", "globex", namespace="kb") db.graph_link("acme", ["doc-1", "doc-2"], namespace="kb") @@ -607,27 +701,32 @@ def redact(ctx): top_k=10, namespace="kb", ) -""") + """ +""") + + """Traversal helpers: graph_add_node, graph_add_edge, graph_link,
graph_neighbors.
Both deletes are idempotent — removing something that is already gone returns 0 — and both
return the number of edges removed.
graph_delete_node leaves linked documents and their embeddings in place; delete those with
db.delete(...) if you want them gone. Nothing indexes inbound edges, so finding them scans the
namespace (dynamodb:Scan) — fine for occasional cleanup, not for a hot path. Edge removal is a
conditional write that retries if another writer changes the adjacency list concurrently.
Amazon S3 Vectors is cheap and serverless, but it is an object-backed ANN: its per-query @@ -643,7 +742,8 @@ def redact(ctx): ever make queries faster, never wrong.
S3 Vectors stores float32 and manages its own layout, so PQ does not change what it stores. PQ compresses
the vectors dynavec caches — the in-memory hot tier and local candidate caches — turning a
dim × 4 byte vector into m bytes.
| Param | Meaning |
|---|---|
m | Number of subspaces; must divide the vector dimension. |
nbits | Bits per subquantizer (8 → 256 centroids, uint8 codes). |
dynavec's workload is I/O-bound (network calls to AWS), and Python releases the GIL during those calls — so
a thread pool gives real parallelism without an async rewrite. Batched writes fan out across threads, and
search_many runs several queries concurrently.
aioboto3) is on the roadmap for very high
concurrency.dynavec uses the standard boto3 credential chain, so exported env vars just work. You can also pass keys explicitly or assume a cross-account role.
-""" + code("""from dynavec import Dynavec, AWSCredentials +""" + + code("""from dynavec import Dynavec, AWSCredentials # explicit keys / profile / cross-account role creds = AWSCredentials( @@ -733,7 +845,8 @@ def redact(ctx): # assume_role_arn="arn:aws:iam::OTHER_ACCOUNT:role/dynavec", ) db = Dynavec(cfg, credentials=creds) -""") + """ +""") + + """dynavec (programmatic access only — no console sign-in needed)..env file (below), then run any example.ap-south-1. A typo like ap-soute-1
makes the ARN invalid and shows two errors. The s3vectors block uses "*", so it is
not affected.Load it before running — set -a && . ./.env && set +a — or use
python-dotenv. dynavec then picks up the credentials automatically.
.env (gitignored) locally,
GitHub Secrets in CI, and prefer an IAM role over long-lived keys in production.A framework-agnostic retriever tool — just a callable that takes a query and returns text.
-""" + code("""from dynavec.integrations.tools import make_retriever_fn +""" + + code("""from dynavec.integrations.tools import make_retriever_fn retrieve = make_retriever_fn(db, top_k=4) # fn(query: str) -> str # also: as_langchain_tool(db), as_crewai_tool(db) -""") + """ +""") + + """Expose dynavec as an MCP server with dynavec_search and dynavec_graph_search tools. Configure via environment variables and launch over stdio:
Add to your Claude Desktop / Cursor configuration (claude_desktop_config.json):
Programmatic initialization is also supported:
-""" + code("""from dynavec.mcp import create_mcp_server +""" + + code("""from dynavec.mcp import create_mcp_server mcp = create_mcp_server(db) mcp.run(transport="stdio") -""")) +"""), +) -PAGES["benchmarking"] = ("Benchmarking", +PAGES["benchmarking"] = ( + "Benchmarking", "Recall, latency, and cost — with tables and charts.", """The suite measures recall@k and latency against a labeled dataset and models cost versus Pinecone, OpenSearch, Qdrant, Weaviate, and Milvus across dimensions and scale.
-""" + code("""pip install "dynavec[benchmark]" +""" + + code("""pip install "dynavec[benchmark]" # recall + latency (local control, or real AWS) python -m benchmarks.run_benchmark --backend local --n 50000 --dim 384 @@ -835,17 +968,22 @@ def redact(ctx): # cost comparison + charts (dimensions 384-3072, 100K -> 1B vectors) python -m benchmarks.report --qpm 1_000_000 -""") + """ +""") + + """dynavec follows semantic versioning. Upgrade with:
-""" + code("""pip install --upgrade dynavec -uv pip install --upgrade dynavec""") + """ +""" + + code("""pip install --upgrade dynavec +uv pip install --upgrade dynavec""") + + """The full machine-readable history lives in CHANGELOG.md, and every version is a @@ -977,10 +1115,12 @@ def redact(ctx):
Initial public release: the hybrid Amazon DynamoDB + Amazon S3 Vectors store, pluggable embedders, namespace RAG, product quantization, RRF fusion, MMR rerank, the GraphRAG layer, caching backends, framework adapters, and one-shot provisioning.
-""") +""", +) -PAGES["dashboard"] = ("Telemetry dashboard", +PAGES["dashboard"] = ( + "Telemetry dashboard", "A native, in-your-brand observability dashboard — a Langfuse-style view of real query telemetry.", """Attach a recorder to your client and every search is captured with latency, cache outcome, @@ -988,19 +1128,23 @@ def redact(ctx):
Attach a recorder to your client and serve the API:
-""" + code("""from dynavec import Dynavec, DynavecConfig, SemanticCache +""" + + code("""from dynavec import Dynavec, DynavecConfig, SemanticCache from dynavec.telemetry import TelemetryRecorder 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""") + + """Points at that API; falls back to sample data if unset:
-""" + code("""cd dashboard +""" + + code("""cd dashboard npm install -NEXT_PUBLIC_DYNAVEC_API=http://127.0.0.1:8779 npm run dev # http://localhost:3000""") + """ +NEXT_PUBLIC_DYNAVEC_API=http://127.0.0.1:8779 npm run dev # http://localhost:3000""") + + """No AWS? python examples/dashboard_demo.py runs real searches against in-memory
stand-ins and serves the API on :8779 for the dashboard to read.
Click any row to open a detail drawer with per-call similarity scores, filter state, and error details.
-""") +""", +) + def issue(n: int) -> str: return f'#{n}' -PAGES["ecosystem"] = ("Ecosystem overview", +PAGES["ecosystem"] = ( + "Ecosystem overview", "Four projects that together cover remembering, running, observing and measuring an AI application.", """dynavec started as a vector database, and it's growing into a full agent stack that runs inside your own cloud account. The plan has four parts. Each one answers a question about an AI application: what does it know, what does it do, what happened when it ran, and how good were the results?
-Some of the newer pieces have already landed in the dynavec package.
One interface for OpenAI, Anthropic and Amazon Bedrock. With Bedrock, the model calls stay in your AWS account like the rest of dynavec.
-""" + code("""from dynavec.chat import BedrockChatModel, Message +""" + + code("""from dynavec.chat import BedrockChatModel, Message model = BedrockChatModel(region_name="us-east-1") # or OpenAIChatModel() result = model.invoke([ Message(role="system", content="Answer in one sentence."), Message(role="user", content="What is a vector database?"), ]) -print(result.message.content)""") + """ +print(result.message.content)""") + + """Every model also has stream(), ainvoke() and astream(), and accepts
tools= for function calling.
Check how well search finds the documents you expect:
-""" + code("""from dynavec.eval import compute_mrr, compute_ndcg_at_k, compute_recall_at_k +""" + + code("""from dynavec.eval import compute_mrr, compute_ndcg_at_k, compute_recall_at_k retrieved = ["d3", "d1", "d7", "d2"] # ids returned by search, best first relevant = {"d1", "d2"} # ids you know are correct compute_recall_at_k(retrieved, relevant, k=3) # 0.5 compute_mrr(retrieved, relevant) # 0.5 -compute_ndcg_at_k(retrieved, relevant, k=3) # 0.3869""") + """ +compute_ndcg_at_k(retrieved, relevant, k=3) # 0.3869""") + + """And use an LLM as a judge to check whether an answer sticks to the retrieved context and actually answers the question:
-""" + code("""from dynavec.eval import OpenAIJudge, evaluate_rag +""" + + code("""from dynavec.eval import OpenAIJudge, evaluate_rag result = evaluate_rag( query="What powers the cell?", @@ -1109,7 +1275,8 @@ def issue(n: int) -> str: answer="The mitochondria.", judge=OpenAIJudge(model="gpt-4o-mini"), # or BedrockJudge, GeminiJudge, CustomJudge ) -print(result.faithfulness.score, result.answer_relevance.score)""") + """ +print(result.faithfulness.score, result.answer_relevance.score)""") + + """Attach a TelemetryRecorder to your client to capture latency, cache hits and scores for every
search, then view them in the telemetry dashboard.
Want to help? Pick an unassigned issue, comment on it to get assigned, and follow CONTRIBUTING.md.
-""") +""", +) def render(slug: str) -> str: @@ -1154,7 +1354,9 @@ def render(slug: str) -> str: # sidebar side = ['', '