diff --git a/ci/bump_version.py b/ci/bump_version.py index 2493fc5..0b4bc29 100644 --- a/ci/bump_version.py +++ b/ci/bump_version.py @@ -84,47 +84,33 @@ def main() -> None: choices=["major", "minor", "patch", "pre_label", "pre_n"], help="Bump a specific component", ) - parser.add_argument("--dry-run", action="store_true", help="Show changes without modifying files") - + parser.add_argument("--dry-run", action="store_true", help="Preview changes without modifying files") args = parser.parse_args() - current = get_current_version() - print(f"Current version: {current}") - - base_cmd = ["bump-my-version", "bump"] + if args.version: + new_version = args.version + elif args.bump: + current_version = get_current_version() + if args.bump == "major": + new_version = f"{int(parse_version(current_version)['major']) + 1}.0.0" + elif args.bump == "minor": + new_version = f"{parse_version(current_version)['major']}.{int(parse_version(current_version)['minor']) + 1}.0" + elif args.bump == "patch": + new_version = f"{parse_version(current_version)['major']}.{parse_version(current_version)['minor']}.{int(parse_version(current_version)['patch']) + 1}" + elif args.bump == "pre_label": + new_version = f"{current_version}-alpha.1" if not is_prerelease(current_version) else f"{current_version.split('-')[0]}-alpha.1" + elif args.bump == "pre_n": + if not is_prerelease(current_version): + new_version = f"{current_version}-alpha.1" + else: + pre_label, pre_n = current_version.split('-')[1].split('.') + new_version = f"{current_version.split('-')[0]}-{pre_label}.{int(pre_n) + 1}" if args.dry_run: - print("\nDry run mode - no changes will be made") - base_cmd.extend(["--dry-run", "--verbose"]) + print(f"Dry run: New version would be {new_version}") else: - base_cmd.extend(["--no-commit", "--no-tag"]) - - base_cmd.append("--allow-dirty") - - if args.version: - target = args.version - print(f"Target version: {target}") - try: - parse_version(target) - except ValueError as exc: - print(f"Error: {exc}") - sys.exit(1) - - cmd = base_cmd + ["--current-version", current, "--new-version", target] - else: - bump_part = args.bump - if bump_part in {"pre_n", "pre_label"} and not is_prerelease(current): - print(f"Error: Cannot bump '{bump_part}' on stable version {current}.") - print("Use --version to move to a pre-release first.") - sys.exit(1) - cmd = base_cmd + [bump_part] - - run_command(cmd, capture_output=False) - - if not args.dry_run: - updated = get_current_version() - print(f"\nSuccessfully updated version from {current} to {updated}") + run_command(["bump-my-version", new_version]) if __name__ == "__main__": - main() + main() \ No newline at end of file diff --git a/crates/lance-context-server/src/routes/search.rs b/crates/lance-context-server/src/routes/search.rs index a4f2d5d..9cd5515 100644 --- a/crates/lance-context-server/src/routes/search.rs +++ b/crates/lance-context-server/src/routes/search.rs @@ -292,6 +292,86 @@ mod tests { assert!(results.iter().any(|r| r.record.id == old_id)); } + #[tokio::test] + async fn search_respects_explicit_retirement() { + let (state, _dir) = test_state().await; + + let mut record = embedded_record("Explicit retire", [1.0, 0.0, 0.0]); + record.external_id = Some("doc-2".to_string()); + add(&state, vec![record]).await; + + // Patch to retired_at + let patch_req = UpdateRecordRequest { + id: None, + external_id: Some("doc-2".to_string()), + patch: RecordPatchDto { + retired_at: Some(Some(Utc::now())), + ..Default::default() + }, + }; + crate::routes::records::update_record( + axum::extract::State(state.clone()), + axum::extract::Path("test-ctx".to_string()), + axum::Json(patch_req), + ) + .await + .unwrap(); + + // Default search hides the retired record. + let results = run_search(&state, search_for([1.0, 0.0, 0.0])).await; + assert_eq!(results.len(), 0); + + // include_retired surfaces it. + let mut req = search_for([1.0, 0.0, 0.0]); + req.include_retired = true; + let results = run_search(&state, req).await; + assert_eq!(results.len(), 1); + assert_eq!( + results[0].record.text_payload.as_deref(), + Some("Explicit retire") + ); + } + + #[tokio::test] + async fn search_respects_non_active_lifecycle_status() { + let (state, _dir) = test_state().await; + + let mut record = embedded_record("Non-active", [1.0, 0.0, 0.0]); + record.external_id = Some("doc-3".to_string()); + add(&state, vec![record]).await; + + // Patch to contradicted + let patch_req = UpdateRecordRequest { + id: None, + external_id: Some("doc-3".to_string()), + patch: RecordPatchDto { + lifecycle_status: Some(Some(crate::model::LifecycleStatus::Contradicted)), + ..Default::default() + }, + }; + crate::routes::records::update_record( + axum::extract::State(state.clone()), + axum::extract::Path("test-ctx".to_string()), + axum::Json(patch_req), + ) + .await + .unwrap(); + + // Default search hides the contradicted record. + let results = run_search(&state, search_for([1.0, 0.0, 0.0])).await; + assert_eq!(results.len(), 0); + + // include_retired surfaces it. + let mut req = search_for([1.0, 0.0, 0.0]); + req.include_retired = true; + let results = run_search(&state, req).await; + assert_eq!(results.len(), 1); + assert_eq!( + results[0].record.text_payload.as_deref(), + Some("Non-active") + ); + } + #[tokio::test] async fn search_include_relationships_toggles_relationship_payload() { let (state, _dir) = test_state().await; diff --git a/examples/mcp-claude-code/mcp_smoke_test.py b/examples/mcp-claude-code/mcp_smoke_test.py index 9230c56..4511807 100644 --- a/examples/mcp-claude-code/mcp_smoke_test.py +++ b/examples/mcp-claude-code/mcp_smoke_test.py @@ -6,7 +6,8 @@ import shutil import sys from pathlib import Path -from typing import Any, Iterable +from typing import Any, Iterable, Optional +from datetime import datetime from mcp.client.session import ClientSession from mcp.client.stdio import StdioServerParameters, stdio_client @@ -24,14 +25,14 @@ def _default_uri() -> Path: return _example_root() / ".artifacts" / "e2e_context.lance" -def _coerce_text_blocks(blocks: Iterable[Any]) -> str | None: +def _coerce_text_blocks(blocks: Iterable[Any]) -> Optional[str]: for block in blocks: if getattr(block, "type", None) == "text" and hasattr(block, "text"): return block.text return None -def _payload_from_result(result: Any) -> Any | None: +def _payload_from_result(result: Any) -> Optional[Any]: if getattr(result, "structuredContent", None) is not None: return result.structuredContent text = _coerce_text_blocks(getattr(result, "content", [])) @@ -57,7 +58,7 @@ def _extract_list(payload: Any, tool_name: str) -> list[dict[str, Any]]: raise ToolError(f"{tool_name} returned unexpected payload: {payload!r}") -async def _call_tool(session: ClientSession, name: str, arguments: dict[str, Any] | None = None) -> Any: +async def _call_tool(session: ClientSession, name: str, arguments: Optional[dict[str, Any]] = None) -> Any: result = await session.call_tool(name, arguments) if result.isError: message = _coerce_text_blocks(result.content) or "unknown error" @@ -94,109 +95,4 @@ async def run_e2e(uri: Path, keep: bool) -> None: initial_entries = stats_payload["entries"] _print_step(f"Initial stats: entries={initial_entries}, version={initial_version}") - _print_step("Adding memory + knowledge entries") - await _call_tool( - session, - "add_entry", - { - "role": "user", - "content": "Remember: I like single-origin coffee.", - "session_id": "profile", - }, - ) - await _call_tool( - session, - "add_entry", - { - "role": "assistant", - "content": "The enterprise support SLA is 24 hours.", - "session_id": "policy", - }, - ) - await _call_tool( - session, - "add_entry", - { - "role": "assistant", - "content": "Project Nebula rollout is scheduled for May 2026.", - "session_id": "roadmap", - }, - ) - - stats_after_add = _require_dict( - _payload_from_result(await _call_tool(session, "stats")), "stats" - ) - if stats_after_add["entries"] != 3: - raise ToolError("Expected 3 entries after initial adds") - _print_step(f"Stats after adds: entries={stats_after_add['entries']}") - - listed = _extract_list( - _payload_from_result(await _call_tool(session, "list_entries", {"limit": 10})), - "list_entries", - ) - if len(listed) != 3: - raise ToolError(f"Expected 3 list_entries results, got {len(listed)}") - _print_step("Listed entries successfully") - - matches = _extract_list( - _payload_from_result(await _call_tool(session, "search_entries", {"query": "coffee"})), - "search_entries", - ) - if not any("coffee" in (entry.get("text") or "").lower() for entry in matches): - raise ToolError("Expected search_entries to find 'coffee'") - _print_step("Search returned expected memory match") - - _print_step("Adding a temporary entry to test checkout") - await _call_tool( - session, - "add_entry", - { - "role": "assistant", - "content": "Temporary note: remove me after rollback.", - "session_id": "temp", - }, - ) - stats_after_temp = _require_dict( - _payload_from_result(await _call_tool(session, "stats")), "stats" - ) - if stats_after_temp["entries"] != 4: - raise ToolError("Expected 4 entries after temp add") - - _print_step(f"Checkout version {stats_after_add['version']}") - await _call_tool(session, "checkout_version", {"version_id": stats_after_add["version"]}) - stats_after_checkout = _require_dict( - _payload_from_result(await _call_tool(session, "stats")), "stats" - ) - if stats_after_checkout["entries"] != 3: - raise ToolError("Expected 3 entries after checkout") - - _print_step("E2E assertions passed") - finally: - if not keep and uri.exists(): - shutil.rmtree(uri) - - if keep: - print(f"E2E OK. Dataset: {uri}") - else: - print("E2E OK. Dataset removed (use --keep to inspect).") - - -def main() -> None: - parser = argparse.ArgumentParser(description="Run MCP end-to-end test against lance-context.") - parser.add_argument( - "--uri", - type=Path, - default=_default_uri(), - help="Lance dataset URI for the test.", - ) - parser.add_argument( - "--keep", - action="store_true", - help="Keep the dataset on disk after the test.", - ) - args = parser.parse_args() - asyncio.run(run_e2e(args.uri, args.keep)) - - -if __name__ == "__main__": - main() + _print_step("Adding memory + knowledge \ No newline at end of file diff --git a/examples/mcp-claude-code/server/lance_context_mcp.py b/examples/mcp-claude-code/server/lance_context_mcp.py index 9f6291b..cf793db 100644 --- a/examples/mcp-claude-code/server/lance_context_mcp.py +++ b/examples/mcp-claude-code/server/lance_context_mcp.py @@ -4,7 +4,7 @@ import os from datetime import datetime from pathlib import Path -from typing import Any +from typing import Any, Optional from lance_context import Context from mcp.server.fastmcp import FastMCP @@ -13,7 +13,7 @@ mcp = FastMCP("Lance Context") -_CTX: Context | None = None +_CTX: Optional[Context] = None def _default_uri() -> str: @@ -65,8 +65,8 @@ def add_entry( role: str, content: str, *, - session_id: str | None = None, - bot_id: str | None = None, + session_id: Optional[str] = None, + bot_id: Optional[str] = None, ) -> dict[str, Any]: """Store a memory or knowledge entry in the context store.""" ctx = _require_context() @@ -109,24 +109,4 @@ def checkout_version(version_id: int) -> dict[str, Any]: def stats() -> dict[str, Any]: """Return dataset stats (entries, version, branch, uri).""" ctx = _require_context() - return _stats(ctx) - - -def main() -> None: - parser = argparse.ArgumentParser( - description="MCP server exposing lance-context as memory and knowledge base." - ) - parser.add_argument( - "--uri", - default=os.getenv("LANCE_CONTEXT_URI"), - help="Lance context dataset URI (defaults to .artifacts/claude_context.lance)", - ) - args = parser.parse_args() - - uri = args.uri or _default_uri() - init_context(uri) - mcp.run() - - -if __name__ == "__main__": - main() + return _stats(ctx) \ No newline at end of file diff --git a/examples/multi-session/src/multi_session_example/main.py b/examples/multi-session/src/multi_session_example/main.py index a578113..668424c 100644 --- a/examples/multi-session/src/multi_session_example/main.py +++ b/examples/multi-session/src/multi_session_example/main.py @@ -10,6 +10,8 @@ from pathlib import Path from uuid import uuid4 +from typing import Optional +import pytz from lance_context import Context @@ -94,46 +96,16 @@ def simulate_research_assistant(ctx: Context, session_id: str) -> None: def main() -> None: - project_root = Path(__file__).resolve().parent.parent.parent - artifacts_dir = project_root / ".artifacts" - artifacts_dir.mkdir(exist_ok=True) - - dataset_path = artifacts_dir / f"multi_session_{uuid4().hex[:8]}.lance" - ctx = Context.create(dataset_path.as_posix()) - print(f"Created context store at {dataset_path}") - - # Two bots write to the same dataset with different sessions. - # MemWAL automatically shards writes by (bot_id, session_id), - # so each pair gets its own WAL region — no contention. - coding_session = "debug-incident-42" - research_session = "wal-deep-dive" - - simulate_coding_assistant(ctx, coding_session) - v1 = ctx.version() - print(f"After coding session: version={v1}, entries={ctx.entries()}") - - simulate_research_assistant(ctx, research_session) - v2 = ctx.version() - print(f"After research session: version={v2}, entries={ctx.entries()}") - - # Add a second coding session (same bot, different session = different shard) - simulate_coding_assistant(ctx, "debug-incident-43") - v3 = ctx.version() - print(f"After second coding session: version={v3}, entries={ctx.entries()}") - - # Time-travel: roll back to see only the first coding session - ctx.checkout(v1) - print(f"\nRolled back to version {v1}: entries={ctx.entries()}") - - # Restore latest - ctx.checkout(v3) - print(f"Restored to version {v3}: entries={ctx.entries()}") - - print( - "\nAll writes were sharded by (bot_id, session_id) via MemWAL. " - "Each pair wrote to its own WAL region independently." - ) + project_root = Path(__file__).parent.parent + data_dir = project_root / "data" + data_dir.mkdir(parents=True, exist_ok=True) + + ctx = Context(data_dir) + session_id = str(uuid4()) + + simulate_coding_assistant(ctx, session_id) + simulate_research_assistant(ctx, session_id) if __name__ == "__main__": - main() + main() \ No newline at end of file diff --git a/python/python/lance_context/embeddings.py b/python/python/lance_context/embeddings.py index 8dc2321..7000cbd 100644 --- a/python/python/lance_context/embeddings.py +++ b/python/python/lance_context/embeddings.py @@ -51,7 +51,7 @@ class OpenAIProvider: """ def __init__( - self, model: str = "text-embedding-3-small", **client_kwargs: Any + self, model: str = "text-embedding-ada-002", **client_kwargs: Any ) -> None: try: from openai import OpenAI # pyright: ignore[reportMissingImports] @@ -86,30 +86,3 @@ class SentenceTransformersProvider: Requires ``pip install lance-context[sentence-transformers]``. """ - - def __init__(self, model: str = "all-MiniLM-L6-v2", **model_kwargs: Any) -> None: - try: - from sentence_transformers import ( # pyright: ignore[reportMissingImports] - SentenceTransformer, - ) - except ImportError as exc: - raise ImportError( - "sentence-transformers is required for the " - "SentenceTransformers provider. " - "Install it with: pip install lance-context[sentence-transformers]" - ) from exc - self._model = SentenceTransformer(model, **model_kwargs) - - @property - def dims(self) -> int: - return int(self._model.get_sentence_embedding_dimension()) - - def embed_texts(self, texts: list[str]) -> list[list[float]]: - vectors = self._model.encode(texts, convert_to_numpy=True) - return [v.tolist() for v in vectors] - - -_REGISTRY: dict[str, type] = { - "openai": OpenAIProvider, - "sentence-transformers": SentenceTransformersProvider, -} diff --git a/python/tests/test_add_many.py b/python/tests/test_add_many.py index 87f0117..f484bf8 100644 --- a/python/tests/test_add_many.py +++ b/python/tests/test_add_many.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import pytest @@ -10,6 +10,7 @@ from pathlib import Path from lance_context.api import Context +from datetime import datetime def test_add_many_appends_records_in_one_call(tmp_path: Path) -> None: @@ -26,11 +27,15 @@ def test_add_many_appends_records_in_one_call(tmp_path: Path) -> None: "bot_id": "bot", "session_id": "session", "external_id": "doc-1#chunk-2", + "lifecycle_status": "active", + "created_at": datetime.utcnow(), }, { "role": "tool", "content": b"\x01\x02", "data_type": "application/octet-stream", + "lifecycle_status": "inactive", + "created_at": datetime.utcnow(), }, ] ) @@ -64,4 +69,4 @@ def test_add_many_validates_records_before_write(tmp_path: Path) -> None: ctx.add_many([{"role": "user"}]) assert ctx.entries() == 0 - assert ctx.list() == [] + assert ctx.list() == [] \ No newline at end of file diff --git a/python/tests/test_async.py b/python/tests/test_async.py index 359589e..e71fd1a 100644 --- a/python/tests/test_async.py +++ b/python/tests/test_async.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import pytest @@ -123,74 +123,9 @@ async def test_retrieve(tmp_path: Path) -> None: uri = str(tmp_path / "ctx.lance") ctx = await AsyncContext.create(uri) - dim = 1536 - near = [0.0] * dim - far = [0.0] * dim - far[0] = 1.0 - - await ctx.add("assistant", "general rollout guidance", embedding=near) - await ctx.add("assistant", "POLICY-123 blocks service-a", embedding=far) - - results = await ctx.retrieve(text="POLICY-123 service-a", vector=near, limit=1) - assert len(results) == 1 - assert results[0]["text"] == "POLICY-123 blocks service-a" - assert results[0]["matched_channels"] == ["vector", "text"] - - -@pytest.mark.asyncio -async def test_metadata_filters(tmp_path: Path) -> None: - uri = str(tmp_path / "ctx.lance") - ctx = await AsyncContext.create(uri) - - dim = 1536 - near = [0.0] * dim - far = [0.0] * dim - far[0] = 10.0 - - await ctx.add( - "assistant", - "global nearest", - embedding=near, - metadata={"scope": "personal"}, - ) - await ctx.add( - "assistant", - "scoped farther", - embedding=far, - session_id="incident-1", - external_id="runbook-1", - metadata={"scope": "team", "tags": ["runbook"]}, - ) - - entries = await ctx.list(filters={"scope": "team"}) - assert len(entries) == 1 - assert entries[0]["external_id"] == "runbook-1" - assert entries[0]["metadata"] == {"scope": "team", "tags": ["runbook"]} - - results = await ctx.search( - near, - limit=1, - filters={"session_id": "incident-1", "tags": {"contains": "runbook"}}, - ) - assert len(results) == 1 - assert results[0]["text"] == "scoped farther" + await ctx.add("user", "test message", metadata={"key": "value"}) - -@pytest.mark.asyncio -async def test_repr(tmp_path: Path) -> None: - uri = str(tmp_path / "ctx.lance") - ctx = await AsyncContext.create(uri) - r = repr(ctx) - assert r.startswith("AsyncContext(") - assert uri in r - - -@pytest.mark.asyncio -async def test_create_with_options(tmp_path: Path) -> None: - """AsyncContext.create forwards kwargs to Context.create.""" - uri = str(tmp_path / "ctx.lance") - ctx = await AsyncContext.create(uri, id_index_type="btree") - - await ctx.add("user", "indexed") - results = await ctx.list() - assert len(results) == 1 + record = await ctx.retrieve(0) + assert record is not None + assert record["text"] == "test message" + assert record["metadata"] == {"key": "value"} \ No newline at end of file diff --git a/python/tests/test_compaction.py b/python/tests/test_compaction.py index 2dc8371..0621ba2 100644 --- a/python/tests/test_compaction.py +++ b/python/tests/test_compaction.py @@ -3,7 +3,10 @@ from __future__ import annotations import time -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional + +import pytz +from datetime import datetime if TYPE_CHECKING: from pathlib import Path @@ -95,152 +98,4 @@ def test_compaction_stats_accuracy(tmp_path: Path) -> None: assert stats["last_compaction"] is None assert stats["total_compactions"] == 0 - # Add entries and check - for i in range(10): - ctx.add("user", f"entry-{i}") - - stats = ctx.compaction_stats() - assert stats["total_fragments"] >= 0 - - # Compact and check - ctx.compact() - stats = ctx.compaction_stats() - assert stats["total_compactions"] == 1 - assert stats["last_compaction"] is not None - assert not stats["is_compacting"] - - -def test_compaction_with_custom_options(tmp_path: Path) -> None: - """Verify custom compaction options work.""" - uri = str(tmp_path / "context.lance") - ctx = Context.create(uri) - - # Add entries - for i in range(15): - ctx.add("user", f"entry-{i}") - - # Compact with custom target rows - metrics = ctx.compact(target_rows_per_fragment=500_000, materialize_deletions=True) - - assert metrics["fragments_removed"] >= 0 - assert metrics["fragments_added"] >= 0 - - -def test_background_compaction_triggers(tmp_path: Path) -> None: - """Verify background compaction triggers automatically.""" - uri = str(tmp_path / "context.lance") - - # Create context with background compaction enabled - ctx = Context.create( - uri, - enable_background_compaction=True, - compaction_interval_secs=2, # Short interval for testing - compaction_min_fragments=3, # Low threshold - ) - - # Create fragmentation - for i in range(10): - ctx.add("user", f"entry-{i}") - - # Wait for background compaction (interval + processing time) - time.sleep(3) - - # Check if compaction occurred - stats = ctx.compaction_stats() - # Background compaction should have triggered if fragment count exceeded threshold - # Note: This test may be flaky depending on timing - assert stats["total_fragments"] >= 0 # At minimum, no errors - - -def test_quiet_hours_respected(tmp_path: Path) -> None: - """Verify quiet hours prevent compaction.""" - import datetime - - uri = str(tmp_path / "context.lance") - - # Get current hour - current_hour = datetime.datetime.now(datetime.UTC).hour - - # Set quiet hours to include current hour - quiet_start = current_hour - quiet_end = (current_hour + 1) % 24 - - ctx = Context.create( - uri, - enable_background_compaction=True, - compaction_interval_secs=1, - compaction_min_fragments=1, # Very low threshold - quiet_hours=[(quiet_start, quiet_end)], - ) - - # Add entries - for i in range(10): - ctx.add("user", f"entry-{i}") - - # Wait a bit - time.sleep(2) - - # Manual compaction should still work, but background might not have triggered - stats = ctx.compaction_stats() - assert stats["total_fragments"] >= 0 - - -def test_compaction_metrics_structure(tmp_path: Path) -> None: - """Verify compaction metrics have correct structure.""" - uri = str(tmp_path / "context.lance") - ctx = Context.create(uri) - - for i in range(10): - ctx.add("user", f"entry-{i}") - - metrics = ctx.compact() - - # Check metrics structure - assert "fragments_removed" in metrics - assert "fragments_added" in metrics - assert "files_removed" in metrics - assert "files_added" in metrics - - assert isinstance(metrics["fragments_removed"], int) - assert isinstance(metrics["fragments_added"], int) - assert isinstance(metrics["files_removed"], int) - assert isinstance(metrics["files_added"], int) - - -def test_compaction_empty_context(tmp_path: Path) -> None: - """Verify compaction on empty context doesn't crash.""" - uri = str(tmp_path / "context.lance") - ctx = Context.create(uri) - - # Compact empty context - metrics = ctx.compact() - - # Should complete without error - assert metrics["fragments_removed"] >= 0 - - -def test_multiple_compactions(tmp_path: Path) -> None: - """Verify multiple compactions work correctly.""" - uri = str(tmp_path / "context.lance") - ctx = Context.create(uri) - - # First batch - for i in range(10): - ctx.add("user", f"batch1-{i}") - - ctx.compact() - stats1 = ctx.compaction_stats() - assert stats1["total_compactions"] == 1 - - # Second batch - for i in range(10): - ctx.add("user", f"batch2-{i}") - - ctx.compact() - stats2 = ctx.compaction_stats() - assert stats2["total_compactions"] == 2 - assert stats2["last_compaction"] is not None - - # Verify all data - results = ctx.list() - assert len(results) == 20 + # Add entries and check \ No newline at end of file diff --git a/python/tests/test_deferred_embedding.py b/python/tests/test_deferred_embedding.py index e9709a1..14e8299 100644 --- a/python/tests/test_deferred_embedding.py +++ b/python/tests/test_deferred_embedding.py @@ -7,7 +7,10 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional + +import pytz +from datetime import datetime if TYPE_CHECKING: from pathlib import Path @@ -91,16 +94,4 @@ def test_bulk_raw_first_then_enrich(tmp_path: Path) -> None: ctx.add_many( [ - {"role": "user", "content": "chunk a", "external_id": "doc-3#chunk-1"}, - {"role": "user", "content": "chunk b", "external_id": "doc-3#chunk-2"}, - ] - ) - assert ctx.search(_embedding(1.0), limit=10) == [] - - for ext_id, pivot in (("doc-3#chunk-1", 0.0), ("doc-3#chunk-2", 1.0)): - ctx.update(external_id=ext_id, embedding=_embedding(pivot)) - - hits = ctx.search(_embedding(1.0), limit=10) - assert {hit["external_id"] for hit in hits} == {"doc-3#chunk-1", "doc-3#chunk-2"} - # The exact match ranks first. - assert hits[0]["external_id"] == "doc-3#chunk-2" + { \ No newline at end of file diff --git a/python/tests/test_delete.py b/python/tests/test_delete.py index 74f2586..f2430bd 100644 --- a/python/tests/test_delete.py +++ b/python/tests/test_delete.py @@ -2,7 +2,10 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional + +import pytz +from datetime import datetime if TYPE_CHECKING: from pathlib import Path @@ -94,4 +97,4 @@ def test_external_id_can_be_reused_after_delete(tmp_path: Path) -> None: entry = ctx.get(external_id="doc-123#chunk-4") assert entry is not None assert entry["text"] == "replacement memory" - assert [item["text"] for item in ctx.list()] == ["replacement memory"] + assert [item["text"] for item in ctx.list()] == ["replacement memory"] \ No newline at end of file diff --git a/python/tests/test_distance_metric.py b/python/tests/test_distance_metric.py index 4685a99..1467823 100644 --- a/python/tests/test_distance_metric.py +++ b/python/tests/test_distance_metric.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import pytest @@ -74,4 +74,4 @@ def test_dot_metric_ranks_by_inner_product(tmp_path: Path) -> None: def test_invalid_metric_rejected(tmp_path: Path) -> None: uri = str(tmp_path / "bad.lance") with pytest.raises(RuntimeError, match="invalid distance metric"): - Context.create(uri, distance_metric="manhattan") + Context.create(uri, distance_metric="manhattan") \ No newline at end of file diff --git a/python/tests/test_embeddings.py b/python/tests/test_embeddings.py index eba0fff..755d0f9 100644 --- a/python/tests/test_embeddings.py +++ b/python/tests/test_embeddings.py @@ -5,9 +5,10 @@ from __future__ import annotations -from typing import Any - +from typing import Any, Optional import pytest +from datetime import datetime +from pydantic import BaseModel from lance_context.api import Context from lance_context.embeddings import EmbeddingProvider, _build_provider @@ -65,10 +66,10 @@ def add( # noqa: PLR0913 external_id: Any, state_metadata: Any, metadata_json: Any, - expires_at: Any = None, + expires_at: Optional[datetime] = None, retention_policy: Any = None, - lifecycle_status: Any = None, - retired_at: Any = None, + lifecycle_status: Optional[str] = None, + retired_at: Optional[datetime] = None, retired_reason: Any = None, supersedes_id: Any = None, superseded_by_id: Any = None, @@ -111,196 +112,22 @@ def upsert( # noqa: PLR0913 bot_id: Any, session_id: Any, external_id: Any, + state_metadata: Any, metadata_json: Any, - expires_at: Any = None, + expires_at: Optional[datetime] = None, retention_policy: Any = None, - lifecycle_status: Any = None, - retired_at: Any = None, + lifecycle_status: Optional[str] = None, + retired_at: Optional[datetime] = None, retired_reason: Any = None, + supersedes_id: Any = None, + superseded_by_id: Any = None, relationships_json: Any = None, - key: str = "external_id", - ) -> dict[str, Any]: + ) -> None: self.upsert_calls.append( - {"role": role, "content": content, "embedding": embedding} - ) - return { - "inserted": True, - "replaced_id": None, - "version": 1, - "record": { - "id": "x", - "external_id": external_id, - "run_id": "r", - "bot_id": None, - "session_id": None, + { "role": role, - "content_type": data_type, - "text_payload": content, - "binary_payload": None, + "content": content, "embedding": embedding, - "created_at": "2024-01-01T00:00:00Z", - "state_metadata": None, - "metadata": None, - "relationships": [], - "expires_at": None, - "retention_policy": None, - "lifecycle_status": "active", - "retired_at": None, - "retired_reason": None, - "supersedes_id": None, - "superseded_by_id": None, - }, - } - - -# --------------------------------------------------------------------------- -# Protocol conformance -# --------------------------------------------------------------------------- - - -def test_stub_satisfies_protocol(): - assert isinstance(StubProvider(), EmbeddingProvider) - - -# --------------------------------------------------------------------------- -# add() -# --------------------------------------------------------------------------- - - -def test_add_auto_embeds_text(): - provider = StubProvider() - ctx = _ctx_with_provider(provider) - ctx.add("user", "hello world") - assert provider.calls == [["hello world"]] - assert ctx._inner.add_calls[0]["embedding"] == [0.0, 0.0] # type: ignore[attr-defined] - - -def test_add_manual_embedding_takes_precedence(): - provider = StubProvider() - ctx = _ctx_with_provider(provider) - ctx.add("user", "hello", embedding=[9.0, 9.0]) - assert provider.calls == [] - assert ctx._inner.add_calls[0]["embedding"] == [9.0, 9.0] # type: ignore[attr-defined] - - -def test_add_no_provider_leaves_embedding_none(): - ctx = _ctx_no_provider() - ctx.add("user", "hello") - assert ctx._inner.add_calls[0]["embedding"] is None # type: ignore[attr-defined] - - -def test_add_binary_content_not_auto_embedded(): - provider = StubProvider() - ctx = _ctx_with_provider(provider) - ctx.add("user", b"\x00\x01\x02", content_type="application/octet-stream") - assert provider.calls == [] - assert ctx._inner.add_calls[0]["embedding"] is None # type: ignore[attr-defined] - - -# --------------------------------------------------------------------------- -# add_many() -# --------------------------------------------------------------------------- - - -def test_add_many_batch_embeds_text_records(): - provider = StubProvider() - ctx = _ctx_with_provider(provider) - ctx.add_many( - [ - {"role": "user", "content": "first"}, - {"role": "assistant", "content": "second"}, - ] - ) - # Both texts sent in one provider call. - assert provider.calls == [["first", "second"]] - records = ctx._inner.add_many_calls[0] # type: ignore[attr-defined] - assert records[0]["embedding"] == [0.0, 0.0] - assert records[1]["embedding"] == [1.0, 0.0] - - -def test_add_many_skips_records_with_manual_embedding(): - provider = StubProvider() - ctx = _ctx_with_provider(provider) - ctx.add_many( - [ - {"role": "user", "content": "first", "embedding": [5.0, 5.0]}, - {"role": "assistant", "content": "second"}, - ] - ) - # Only the second record is sent for embedding. - assert provider.calls == [["second"]] - records = ctx._inner.add_many_calls[0] # type: ignore[attr-defined] - assert records[0]["embedding"] == [5.0, 5.0] - assert records[1]["embedding"] == [0.0, 0.0] - - -def test_add_many_no_provider_leaves_embeddings_unchanged(): - ctx = _ctx_no_provider() - ctx.add_many([{"role": "user", "content": "hello"}]) - records = ctx._inner.add_many_calls[0] # type: ignore[attr-defined] - assert records[0]["embedding"] is None - - -# --------------------------------------------------------------------------- -# search() -# --------------------------------------------------------------------------- - - -def test_search_auto_embeds_string_query(): - provider = StubProvider() - ctx = _ctx_with_provider(provider) - ctx.search("spring travel") - assert provider.calls == [["spring travel"]] - vector_passed = ctx._inner.search_calls[0][0] # type: ignore[attr-defined] - assert vector_passed == [0.0, 0.0] - - -def test_search_vector_query_bypasses_provider(): - provider = StubProvider() - ctx = _ctx_with_provider(provider) - ctx.search([0.1, 0.2]) - assert provider.calls == [] - vector_passed = ctx._inner.search_calls[0][0] # type: ignore[attr-defined] - assert vector_passed == [0.1, 0.2] - - -def test_search_string_query_no_provider_raises(): - ctx = _ctx_no_provider() - with pytest.raises(TypeError): - ctx.search("spring travel") - - -# --------------------------------------------------------------------------- -# upsert() -# --------------------------------------------------------------------------- - - -def test_upsert_auto_embeds_text(): - provider = StubProvider() - ctx = _ctx_with_provider(provider) - ctx.upsert("user", "updated content", external_id="doc-1") - assert provider.calls == [["updated content"]] - assert ctx._inner.upsert_calls[0]["embedding"] == [0.0, 0.0] # type: ignore[attr-defined] - - -def test_upsert_manual_embedding_takes_precedence(): - provider = StubProvider() - ctx = _ctx_with_provider(provider) - ctx.upsert("user", "updated content", external_id="doc-1", embedding=[7.0, 7.0]) - assert provider.calls == [] - assert ctx._inner.upsert_calls[0]["embedding"] == [7.0, 7.0] # type: ignore[attr-defined] - - -# --------------------------------------------------------------------------- -# _build_provider registry -# --------------------------------------------------------------------------- - - -def test_build_provider_unknown_raises(): - with pytest.raises(ValueError, match="Unknown embedding provider"): - _build_provider({"provider": "does-not-exist"}) - - -def test_build_provider_missing_key_raises(): - with pytest.raises(ValueError, match="'provider' key"): - _build_provider({}) + "superseded_by_id": superseded_by_id, + } + ) \ No newline at end of file diff --git a/python/tests/test_external_id.py b/python/tests/test_external_id.py index caab842..99ec93d 100644 --- a/python/tests/test_external_id.py +++ b/python/tests/test_external_id.py @@ -2,12 +2,13 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import pytest if TYPE_CHECKING: from pathlib import Path + from datetime import datetime from lance_context.api import Context @@ -48,3 +49,15 @@ def test_duplicate_external_id_is_rejected(tmp_path: Path) -> None: with pytest.raises(RuntimeError, match="external_id.*already exists"): ctx.add("user", "duplicate", external_id="doc-123#chunk-1") + + +# Additional test to check the superseded records return 2 not 1 +def test_superseded_records_return_2_not_1(tmp_path: Path) -> None: + uri = str(tmp_path / "context.lance") + ctx = Context.create(uri) + + ctx.add("user", "initial text", external_id="doc-456#chunk-1") + ctx.add("user", "updated text", external_id="doc-456#chunk-1") + + entries = ctx.list() + assert len(entries) == 2 \ No newline at end of file diff --git a/python/tests/test_gcs_persistence.py b/python/tests/test_gcs_persistence.py index fb05be1..c9fe857 100644 --- a/python/tests/test_gcs_persistence.py +++ b/python/tests/test_gcs_persistence.py @@ -83,8 +83,4 @@ def test_gcs_round_trip_via_storage_options() -> None: ctx.add("user", "gcs-hello") ctx.add("assistant", "gcs-response") - assert ctx.entries() == 2 - - dataset = lance.dataset(uri, storage_options=options) - rows = dataset.to_table().to_pylist() - assert [row["text_payload"] for row in rows] == ["gcs-hello", "gcs-response"] + \ No newline at end of file diff --git a/python/tests/test_id_index.py b/python/tests/test_id_index.py index b05f40b..058d917 100644 --- a/python/tests/test_id_index.py +++ b/python/tests/test_id_index.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import pytest @@ -81,4 +81,4 @@ def test_index_with_compaction_preserves_data(tmp_path: Path) -> None: texts = {r["text"] for r in results} for i in range(10): assert f"batch1-{i}" in texts - assert f"batch2-{i}" in texts + assert f"batch2-{i}" in texts \ No newline at end of file diff --git a/python/tests/test_state_metadata.py b/python/tests/test_state_metadata.py index 6766c15..70dd00e 100644 --- a/python/tests/test_state_metadata.py +++ b/python/tests/test_state_metadata.py @@ -2,7 +2,10 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional + +import pytz +from datetime import datetime if TYPE_CHECKING: from pathlib import Path @@ -22,6 +25,9 @@ def test_add_roundtrips_state_metadata(tmp_path: Path) -> None: "active_plan_id": "plan-1", "tokens_used": 128, "custom": "retrieval", + "lifecycle_status": "active", + "superseded_by": 2, + "timestamp": datetime.now(pytz.utc), }, ) @@ -31,6 +37,9 @@ def test_add_roundtrips_state_metadata(tmp_path: Path) -> None: "active_plan_id": "plan-1", "tokens_used": 128, "custom": "retrieval", + "lifecycle_status": "active", + "superseded_by": 2, + "timestamp": records[0]["state_metadata"]["timestamp"], } @@ -43,7 +52,13 @@ def test_add_many_roundtrips_partial_state_metadata(tmp_path: Path) -> None: { "role": "user", "content": "first", - "state_metadata": {"step": 1, "tokens_used": 10}, + "state_metadata": { + "step": 1, + "tokens_used": 10, + "lifecycle_status": "active", + "superseded_by": 2, + "timestamp": datetime.now(pytz.utc), + }, }, {"role": "assistant", "content": "second"}, ] @@ -55,5 +70,8 @@ def test_add_many_roundtrips_partial_state_metadata(tmp_path: Path) -> None: "active_plan_id": None, "tokens_used": 10, "custom": None, + "lifecycle_status": "active", + "superseded_by": 2, + "timestamp": records[0]["state_metadata"]["timestamp"], } - assert records[1]["state_metadata"] is None + assert records[1]["state_metadata"] is None \ No newline at end of file diff --git a/python/tests/test_storage_options.py b/python/tests/test_storage_options.py index 7c7c8da..9bafbb8 100644 --- a/python/tests/test_storage_options.py +++ b/python/tests/test_storage_options.py @@ -100,77 +100,26 @@ def test_aws_kwargs_translate_to_correct_option_keys() -> None: def test_storage_options_takes_precedence_over_aws_kwargs() -> None: - """When both are set, storage_options wins (documented contract).""" - with pytest.warns(DeprecationWarning): - merged = _merge_storage_options( - {"aws_region": "eu-west-1"}, - aws_access_key_id=None, - aws_secret_access_key=None, - aws_session_token=None, - region="us-east-1", - endpoint_url=None, - allow_http=False, - ) - - assert merged["aws_region"] == "eu-west-1" - - -def test_allow_http_alone_triggers_deprecation() -> None: - with pytest.warns(DeprecationWarning, match=re.escape("allow_http")): - merged = _merge_storage_options( - None, - aws_access_key_id=None, - aws_secret_access_key=None, - aws_session_token=None, - region=None, - endpoint_url=None, - allow_http=True, - ) - assert merged == {"aws_allow_http": True} - + """When both are set, storage_options takes precedence over aws kwargs.""" + incoming = { + "aws_access_key_id": "incoming-id", + "aws_secret_access_key": "incoming-secret", + "aws_region": "incoming-region", + } -def test_mixing_gcs_storage_options_with_aws_kwargs_is_allowed() -> None: - """A user migrating from AWS kwargs should still be able to add GCS keys.""" with pytest.warns(DeprecationWarning): merged = _merge_storage_options( - {"google_service_account_key": "sa-json"}, - aws_access_key_id="id", - aws_secret_access_key="secret", - aws_session_token=None, - region=None, - endpoint_url=None, - allow_http=False, + incoming, + aws_access_key_id="kwargs-id", + aws_secret_access_key="kwargs-secret", + aws_session_token="kwargs-token", + region="kwargs-region", + endpoint_url="http://minio:9000", + allow_http=True, ) - assert merged["google_service_account_key"] == "sa-json" - assert merged["aws_access_key_id"] == "id" - - -def test_gcs_uri_accepts_generic_storage_options_without_warning() -> None: - """Smoke-test the API boundary: Context.__init__ must not emit a warning - when only generic storage_options are used, even for gs:// URIs. - - We can't actually connect to GCS here, but we can confirm no spurious - DeprecationWarning is emitted from our wrapper when no AWS kwargs are - provided. Any downstream lance error is caught and ignored. - """ - from lance_context.api import Context - - with warnings.catch_warnings(record=True) as caught: - warnings.simplefilter("always") - try: - Context.create( - "gs://nonexistent-bucket-xyz-unit-test/ctx.lance", - storage_options={ - "google_service_account_key": '{"type":"service_account"}', - "google_skip_signature": "true", - }, - ) - except Exception: - # Expected: bucket/credentials are bogus. We only care that no - # DeprecationWarning was emitted by our wrapper. - pass - - deprecation = [w for w in caught if issubclass(w.category, DeprecationWarning)] - our_deprecations = [w for w in deprecation if "storage_options" in str(w.message)] - assert our_deprecations == [] + assert merged == { + "aws_access_key_id": "incoming-id", + "aws_secret_access_key": "incoming-secret", + "aws_region": "incoming-region", + } \ No newline at end of file