ee7b62bc5b
All checks were successful
Tests / test (push) Successful in 13m33s
- Add search_all_knowledge_bases MCP tool + POST /api/kb/search-all endpoint + query_documents_multi_kb() in sqlite_manager with dedup and score merge - Add inline health check to search_knowledge_base (24h cache, appends warnings) - Rewrite CLAUDE.md template: 351→128 lines, project content first, session tips - Simplify session command: 625→224 lines, 4-field state file - Update README.md What's New to v2.5, remove Discord section - Add cross-KB search tool to CLAUDE.md MCP tools table - Add 5 new tests for query_documents_multi_kb (all 593 tests pass) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
693 lines
25 KiB
Python
693 lines
25 KiB
Python
"""
|
|
Tests for SQLite database manager operations.
|
|
"""
|
|
|
|
import pytest
|
|
import tempfile
|
|
import json
|
|
from pathlib import Path
|
|
from unittest.mock import patch, MagicMock
|
|
import numpy as np
|
|
|
|
from app.core.sqlite_manager import SQLiteManager, generate_slug, get_sqlite_manager
|
|
from app.core.kb_types import KBType
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSQLiteManager:
|
|
"""Test SQLite manager basic operations."""
|
|
|
|
def test_generate_slug(self):
|
|
"""Test slug generation function."""
|
|
test_cases = [
|
|
("My Agent OS", "my-agent-os"),
|
|
("My Code Base!", "my-code-base"),
|
|
("Test_KB 123", "test-kb-123"),
|
|
("Hello World", "hello-world"),
|
|
("Multiple Spaces", "multiple-spaces"),
|
|
("Special@#$%Chars", "specialchars"),
|
|
("", ""),
|
|
]
|
|
|
|
for input_name, expected_slug in test_cases:
|
|
result = generate_slug(input_name)
|
|
assert result == expected_slug, f"Failed for input: {input_name}"
|
|
|
|
def test_sqlite_manager_initialization(self):
|
|
"""Test SQLite manager initialization."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "test.db"
|
|
manager = SQLiteManager(str(db_path))
|
|
|
|
assert manager.db_path == str(db_path)
|
|
assert Path(db_path).exists()
|
|
|
|
def test_create_collection(self, clean_db):
|
|
"""Test creating a knowledge base collection."""
|
|
result = clean_db.create_collection(
|
|
name="test_collection",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test collection",
|
|
tags=["test", "example"]
|
|
)
|
|
|
|
assert result["name"] == "test_collection"
|
|
assert result["kb_type"] == "generic"
|
|
assert result["description"] == "Test collection"
|
|
assert "id" in result
|
|
assert "slug" in result
|
|
assert "created_at" in result
|
|
|
|
def test_create_collection_duplicate_name(self, clean_db):
|
|
"""Test creating collection with duplicate name."""
|
|
# Create first collection
|
|
clean_db.create_collection("duplicate_test", KBType.GENERIC)
|
|
|
|
# Try to create duplicate
|
|
with pytest.raises(ValueError, match="already exists"):
|
|
clean_db.create_collection("duplicate_test", KBType.CODE)
|
|
|
|
def test_create_collection_slug_conflict(self, clean_db):
|
|
"""Test creating collection with slug conflict."""
|
|
# Create first collection
|
|
clean_db.create_collection("Test Collection", KBType.GENERIC)
|
|
|
|
# Try to create collection with similar slug
|
|
with pytest.raises(ValueError, match="slug conflict"):
|
|
clean_db.create_collection("test-collection", KBType.CODE)
|
|
|
|
def test_list_collections(self, clean_db):
|
|
"""Test listing all collections."""
|
|
# Create a sample kb first
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
collections = clean_db.list_collections()
|
|
|
|
assert isinstance(collections, list)
|
|
assert len(collections) >= 1
|
|
assert any(col["name"] == kb_data["name"] for col in collections)
|
|
|
|
def test_get_collection_metadata(self, clean_db):
|
|
"""Test getting collection metadata."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
metadata = clean_db.get_collection_metadata(kb_data["name"])
|
|
|
|
assert metadata["kb_type"] == KBType.GENERIC.value
|
|
assert "description" in metadata
|
|
assert "created_at" in metadata
|
|
|
|
def test_get_collection_metadata_not_found(self, clean_db):
|
|
"""Test getting metadata for non-existent collection."""
|
|
with pytest.raises(ValueError, match="not found"):
|
|
clean_db.get_collection_metadata("nonexistent_collection")
|
|
|
|
def test_delete_collection(self, clean_db):
|
|
"""Test deleting a collection."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
# Verify collection exists
|
|
assert clean_db.collection_exists(kb_data["name"])
|
|
|
|
# Delete it
|
|
result = clean_db.delete_collection(kb_data["name"])
|
|
assert result is True
|
|
|
|
# Verify it's gone
|
|
assert not clean_db.collection_exists(kb_data["name"])
|
|
|
|
def test_delete_collection_not_found(self, clean_db):
|
|
"""Test deleting non-existent collection."""
|
|
result = clean_db.delete_collection("nonexistent_collection")
|
|
assert result is False
|
|
|
|
def test_collection_exists(self, clean_db):
|
|
"""Test checking if collection exists."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
assert clean_db.collection_exists(kb_data["name"]) is True
|
|
assert clean_db.collection_exists("nonexistent_collection") is False
|
|
|
|
def test_get_collection_count(self, clean_db):
|
|
"""Test getting document count for collection."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
# Add some documents
|
|
documents = ["Test document 1", "Test document 2"]
|
|
embeddings = [[0.1] * 768, [0.2] * 768]
|
|
metadatas = [{"filename": "doc1.txt"}, {"filename": "doc2.txt"}]
|
|
ids = ["doc1", "doc2"]
|
|
|
|
clean_db.add_documents(
|
|
kb_name=kb_data["name"],
|
|
documents=documents,
|
|
embeddings=embeddings,
|
|
metadatas=metadatas,
|
|
ids=ids
|
|
)
|
|
|
|
count = clean_db.get_collection_count(kb_data["name"])
|
|
assert count == 2
|
|
|
|
def test_get_collection_count_empty(self, clean_db):
|
|
"""Test getting document count for empty collection."""
|
|
# Create empty collection
|
|
clean_db.create_collection("empty_collection", KBType.GENERIC)
|
|
|
|
count = clean_db.get_collection_count("empty_collection")
|
|
assert count == 0
|
|
|
|
def test_get_collection_by_id(self, clean_db):
|
|
"""Test getting collection by ID."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
collection = clean_db.get_collection_by_id(kb_data["id"])
|
|
assert collection is not None
|
|
assert collection["name"] == kb_data["name"]
|
|
assert collection["id"] == kb_data["id"]
|
|
|
|
def test_get_collection_by_id_not_found(self, clean_db):
|
|
"""Test getting non-existent collection by ID."""
|
|
collection = clean_db.get_collection_by_id(99999)
|
|
assert collection is None
|
|
|
|
def test_list_collections_by_type(self, clean_db):
|
|
"""Test listing collections filtered by type."""
|
|
# Create collections of different types
|
|
clean_db.create_collection("generic_kb", KBType.GENERIC)
|
|
clean_db.create_collection("code_kb", KBType.CODE)
|
|
clean_db.create_collection("docs_kb", KBType.DOCUMENTATION)
|
|
|
|
# Get only code collections
|
|
code_collections = clean_db.list_collections_by_type(KBType.CODE)
|
|
assert len(code_collections) == 1
|
|
assert code_collections[0]["metadata"]["kb_type"] == "code"
|
|
|
|
def test_get_kb_by_slug(self, clean_db):
|
|
"""Test getting KB by slug."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
kb_name = clean_db.get_kb_by_slug(kb_data["slug"])
|
|
assert kb_name == kb_data["name"]
|
|
|
|
def test_get_kb_by_slug_not_found(self, clean_db):
|
|
"""Test getting non-existent KB by slug."""
|
|
kb_name = clean_db.get_kb_by_slug("nonexistent-slug")
|
|
assert kb_name is None
|
|
|
|
def test_slug_exists(self, clean_db):
|
|
"""Test checking if slug exists."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
assert clean_db.slug_exists(kb_data["slug"]) is True
|
|
assert clean_db.slug_exists("nonexistent-slug") is False
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestSQLiteManagerDocuments:
|
|
"""Test SQLite manager document operations."""
|
|
|
|
def test_add_documents(self, clean_db):
|
|
"""Test adding documents to collection."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
documents = ["Test document 1", "Test document 2"]
|
|
embeddings = [[0.1] * 768, [0.2] * 768]
|
|
metadatas = [{"filename": "doc1.txt"}, {"filename": "doc2.txt"}]
|
|
ids = ["doc1", "doc2"]
|
|
|
|
clean_db.add_documents(
|
|
kb_name=kb_data["name"],
|
|
documents=documents,
|
|
embeddings=embeddings,
|
|
metadatas=metadatas,
|
|
ids=ids
|
|
)
|
|
|
|
# Verify documents were added
|
|
count = clean_db.get_collection_count(kb_data["name"])
|
|
assert count == 2
|
|
|
|
def test_add_documents_nonexistent_collection(self, clean_db):
|
|
"""Test adding documents to non-existent collection."""
|
|
with pytest.raises(ValueError, match="not found"):
|
|
clean_db.add_documents(
|
|
kb_name="nonexistent_collection",
|
|
documents=["test"],
|
|
embeddings=[[0.1] * 768],
|
|
metadatas=[{}],
|
|
ids=["test"]
|
|
)
|
|
|
|
def test_get_documents_by_metadata(self, clean_db):
|
|
"""Test retrieving documents by metadata filter."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
# Add documents with metadata
|
|
for i in range(5):
|
|
np.random.seed(42 + i)
|
|
embedding = np.random.randn(768).tolist()
|
|
clean_db.add_documents(
|
|
kb_name=kb_data["name"],
|
|
documents=[f"This is test document {i}"],
|
|
embeddings=[embedding],
|
|
metadatas=[{"filename": f"test_{i}.txt", "chunk_index": i}],
|
|
ids=[f"node_doc_{i}"]
|
|
)
|
|
|
|
# Get all documents
|
|
all_docs = clean_db.get_documents_by_metadata(
|
|
kb_name=kb_data["name"],
|
|
where={}
|
|
)
|
|
assert len(all_docs) == 5
|
|
|
|
# Filter by filename
|
|
filtered_docs = clean_db.get_documents_by_metadata(
|
|
kb_name=kb_data["name"],
|
|
where={"filename": "test_0.txt"}
|
|
)
|
|
assert len(filtered_docs) == 1
|
|
assert filtered_docs[0]["metadata"]["filename"] == "test_0.txt"
|
|
|
|
def test_query_documents(self, clean_db):
|
|
"""Test querying documents by vector similarity."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
# Add documents
|
|
for i in range(5):
|
|
np.random.seed(42 + i)
|
|
embedding = np.random.randn(768).tolist()
|
|
clean_db.add_documents(
|
|
kb_name=kb_data["name"],
|
|
documents=[f"This is test document {i}"],
|
|
embeddings=[embedding],
|
|
metadatas=[{"filename": f"test_{i}.txt"}],
|
|
ids=[f"node_doc_{i}"]
|
|
)
|
|
|
|
query_embedding = [0.1] * 768
|
|
results = clean_db.query_documents(
|
|
kb_name=kb_data["name"],
|
|
query_embedding=query_embedding,
|
|
n_results=3
|
|
)
|
|
|
|
assert "ids" in results
|
|
assert "documents" in results
|
|
assert "metadatas" in results
|
|
assert "distances" in results
|
|
assert len(results["ids"][0]) <= 3
|
|
|
|
def test_query_documents_nonexistent_collection(self, clean_db):
|
|
"""Test querying non-existent collection."""
|
|
with pytest.raises(ValueError, match="not found"):
|
|
clean_db.query_documents(
|
|
kb_name="nonexistent_collection",
|
|
query_embedding=[0.1] * 768
|
|
)
|
|
|
|
def test_query_similar(self, clean_db):
|
|
"""Test querying similar documents by KB ID."""
|
|
kb_data = clean_db.create_collection(
|
|
name="test_kb",
|
|
kb_type=KBType.GENERIC,
|
|
description="Test knowledge base"
|
|
)
|
|
|
|
# Add documents
|
|
for i in range(5):
|
|
np.random.seed(42 + i)
|
|
embedding = np.random.randn(768).tolist()
|
|
clean_db.add_documents(
|
|
kb_name=kb_data["name"],
|
|
documents=[f"This is test document {i}"],
|
|
embeddings=[embedding],
|
|
metadatas=[{"filename": f"test_{i}.txt"}],
|
|
ids=[f"node_doc_{i}"]
|
|
)
|
|
|
|
query_embedding = [0.1] * 768
|
|
results = clean_db.query_similar(
|
|
kb_id=kb_data["id"],
|
|
query_embedding=query_embedding,
|
|
top_k=3
|
|
)
|
|
|
|
assert isinstance(results, list)
|
|
assert len(results) <= 3
|
|
|
|
if results: # If we have results
|
|
for result in results:
|
|
assert "doc_id" in result
|
|
assert "text" in result
|
|
assert "metadata" in result
|
|
assert "similarity" in result
|
|
assert isinstance(result["similarity"], float)
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestSQLiteManagerProjects:
|
|
"""Test SQLite manager project operations."""
|
|
|
|
def test_create_project(self, clean_db):
|
|
"""Test creating a project."""
|
|
result = clean_db.create_project(
|
|
name="test_project",
|
|
path="/path/to/project",
|
|
description="Test project",
|
|
metadata={"type": "test"}
|
|
)
|
|
|
|
assert result["name"] == "test_project"
|
|
assert result["path"] == "/path/to/project"
|
|
assert result["description"] == "Test project"
|
|
assert result["metadata"]["type"] == "test"
|
|
assert "id" in result
|
|
assert "created_at" in result
|
|
|
|
def test_create_project_duplicate(self, clean_db):
|
|
"""Test creating duplicate project."""
|
|
clean_db.create_project("duplicate_project", "/path/1")
|
|
|
|
with pytest.raises(ValueError, match="already exists"):
|
|
clean_db.create_project("duplicate_project", "/path/2")
|
|
|
|
def test_get_project(self, clean_db):
|
|
"""Test getting project by ID."""
|
|
created = clean_db.create_project("test_project", "/path/to/project")
|
|
retrieved = clean_db.get_project(created["id"])
|
|
|
|
assert retrieved is not None
|
|
assert retrieved["name"] == "test_project"
|
|
assert retrieved["path"] == "/path/to/project"
|
|
|
|
def test_get_project_not_found(self, clean_db):
|
|
"""Test getting non-existent project."""
|
|
project = clean_db.get_project(99999)
|
|
assert project is None
|
|
|
|
def test_list_projects(self, clean_db):
|
|
"""Test listing all projects."""
|
|
# Create multiple projects
|
|
clean_db.create_project("project1", "/path/1")
|
|
clean_db.create_project("project2", "/path/2")
|
|
|
|
projects = clean_db.list_projects()
|
|
assert len(projects) >= 2
|
|
|
|
project_names = [p["name"] for p in projects]
|
|
assert "project1" in project_names
|
|
assert "project2" in project_names
|
|
|
|
def test_assign_kb_to_project(self, clean_db):
|
|
"""Test assigning KB to project."""
|
|
# Create project and KB
|
|
project = clean_db.create_project("test_project", "/path")
|
|
kb = clean_db.create_collection("test_kb", KBType.GENERIC)
|
|
|
|
# Assign KB to project
|
|
assignment = clean_db.assign_kb_to_project(
|
|
project_id=project["id"],
|
|
kb_id=kb["id"],
|
|
mcp_type="knowledge_docs"
|
|
)
|
|
|
|
assert assignment["project_id"] == project["id"]
|
|
assert assignment["kb_id"] == kb["id"]
|
|
assert assignment["mcp_type"] == "knowledge_docs"
|
|
|
|
def test_get_project_kbs(self, clean_db):
|
|
"""Test getting KB assignments for project."""
|
|
# Create project and KBs
|
|
project = clean_db.create_project("test_project", "/path")
|
|
kb1 = clean_db.create_collection("kb1", KBType.GENERIC)
|
|
kb2 = clean_db.create_collection("kb2", KBType.CODE)
|
|
|
|
# Assign KBs
|
|
clean_db.assign_kb_to_project(project["id"], kb1["id"], "knowledge_docs")
|
|
clean_db.assign_kb_to_project(project["id"], kb2["id"], "project_profile")
|
|
|
|
# Get assignments
|
|
kbs = clean_db.get_project_kbs(project["id"])
|
|
assert kbs["knowledge_docs"] == kb1["id"]
|
|
assert kbs["project_profile"] == kb2["id"]
|
|
|
|
def test_get_project_mcps_detailed(self, clean_db):
|
|
"""Test getting detailed MCP info for project."""
|
|
# Create project and KB
|
|
project = clean_db.create_project("test_project", "/path")
|
|
kb = clean_db.create_collection("test_kb", KBType.GENERIC)
|
|
|
|
# Assign KB
|
|
clean_db.assign_kb_to_project(project["id"], kb["id"], "knowledge_docs")
|
|
|
|
# Get detailed info
|
|
mcps = clean_db.get_project_mcps_detailed(project["id"])
|
|
assert "knowledge_docs" in mcps
|
|
assert mcps["knowledge_docs"]["kb_id"] == kb["id"]
|
|
assert mcps["knowledge_docs"]["kb_name"] == "test_kb"
|
|
|
|
def test_set_kb_folder(self, clean_db):
|
|
"""Test setting KB folder configuration."""
|
|
project = clean_db.create_project("test_project", "/path")
|
|
|
|
folder_config = clean_db.set_kb_folder(
|
|
project_id=project["id"],
|
|
mcp_type="knowledge_docs",
|
|
folder_path="/path/to/docs",
|
|
auto_sync=True
|
|
)
|
|
|
|
assert folder_config["project_id"] == project["id"]
|
|
assert folder_config["mcp_type"] == "knowledge_docs"
|
|
assert folder_config["folder_path"] == "/path/to/docs"
|
|
assert folder_config["auto_sync"] is True
|
|
|
|
def test_get_kb_folders(self, clean_db):
|
|
"""Test getting KB folder configurations."""
|
|
project = clean_db.create_project("test_project", "/path")
|
|
|
|
# Set multiple folder configs
|
|
clean_db.set_kb_folder(project["id"], "knowledge_docs", "/docs", True)
|
|
clean_db.set_kb_folder(project["id"], "project_profile", "/profile", False)
|
|
|
|
# Get all configs
|
|
folders = clean_db.get_kb_folders(project["id"])
|
|
assert folders["knowledge_docs"]["folder_path"] == "/docs"
|
|
assert folders["knowledge_docs"]["auto_sync"] is True
|
|
assert folders["project_profile"]["folder_path"] == "/profile"
|
|
assert folders["project_profile"]["auto_sync"] is False
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestSQLiteManagerSingleton:
|
|
"""Test SQLite manager singleton pattern."""
|
|
|
|
def test_get_sqlite_manager_singleton(self):
|
|
"""Test that get_sqlite_manager returns singleton."""
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
db_path = Path(temp_dir) / "test.db"
|
|
|
|
manager1 = get_sqlite_manager(str(db_path))
|
|
manager2 = get_sqlite_manager(str(db_path))
|
|
|
|
# Should be the same instance
|
|
assert manager1 is manager2
|
|
|
|
def test_get_sqlite_manager_default_path(self):
|
|
"""Test get_sqlite_manager with default path."""
|
|
# This should use the default path from Config
|
|
manager = get_sqlite_manager()
|
|
assert manager is not None
|
|
assert hasattr(manager, 'db_path')
|
|
|
|
|
|
@pytest.mark.unit
|
|
class TestQueryDocumentsMultiKB:
|
|
"""Test cross-KB search functionality."""
|
|
|
|
def _add_docs(self, db, kb_name, docs):
|
|
"""Helper to add documents with embeddings to a KB."""
|
|
documents = [d["text"] for d in docs]
|
|
embeddings = [d["embedding"] for d in docs]
|
|
metadatas = [d.get("metadata", {}) for d in docs]
|
|
ids = [d["id"] for d in docs]
|
|
db.add_documents(kb_name=kb_name, documents=documents,
|
|
embeddings=embeddings, metadatas=metadatas, ids=ids)
|
|
|
|
def test_multi_kb_basic(self, clean_db):
|
|
"""Test searching across two KBs returns merged results."""
|
|
clean_db.create_collection("kb_a", KBType.GENERIC)
|
|
clean_db.create_collection("kb_b", KBType.GENERIC)
|
|
|
|
np.random.seed(10)
|
|
base = np.random.randn(768).astype(np.float32)
|
|
base = base / np.linalg.norm(base)
|
|
|
|
# KB_A: one doc very similar to query
|
|
perturbed_a = base + np.random.randn(768).astype(np.float32) * 0.01
|
|
perturbed_a = perturbed_a / np.linalg.norm(perturbed_a)
|
|
self._add_docs(clean_db, "kb_a", [
|
|
{"id": "a1", "text": "Document from KB A", "embedding": perturbed_a.tolist(),
|
|
"metadata": {"filename": "a1.md"}}
|
|
])
|
|
|
|
# KB_B: one doc less similar
|
|
perturbed_b = base + np.random.randn(768).astype(np.float32) * 0.05
|
|
perturbed_b = perturbed_b / np.linalg.norm(perturbed_b)
|
|
self._add_docs(clean_db, "kb_b", [
|
|
{"id": "b1", "text": "Document from KB B", "embedding": perturbed_b.tolist(),
|
|
"metadata": {"filename": "b1.md"}}
|
|
])
|
|
|
|
result = clean_db.query_documents_multi_kb(
|
|
kb_names=["kb_a", "kb_b"],
|
|
query_embedding=base.tolist(),
|
|
n_results=10
|
|
)
|
|
|
|
assert len(result["ids"]) == 2
|
|
assert len(result["documents"]) == 2
|
|
assert len(result["scores"]) == 2
|
|
# First result should have higher score
|
|
assert result["scores"][0] >= result["scores"][1]
|
|
# Each result should have _source_kb
|
|
kbs_found = {m["_source_kb"] for m in result["metadatas"]}
|
|
assert "kb_a" in kbs_found
|
|
assert "kb_b" in kbs_found
|
|
|
|
def test_multi_kb_dedup(self, clean_db):
|
|
"""Test that duplicate content across KBs is deduplicated."""
|
|
clean_db.create_collection("kb_x", KBType.GENERIC)
|
|
clean_db.create_collection("kb_y", KBType.GENERIC)
|
|
|
|
np.random.seed(20)
|
|
emb = np.random.randn(768).astype(np.float32)
|
|
emb = emb / np.linalg.norm(emb)
|
|
|
|
same_text = "Exact same content in both KBs for dedup testing"
|
|
self._add_docs(clean_db, "kb_x", [
|
|
{"id": "x1", "text": same_text, "embedding": emb.tolist(), "metadata": {}}
|
|
])
|
|
self._add_docs(clean_db, "kb_y", [
|
|
{"id": "y1", "text": same_text, "embedding": emb.tolist(), "metadata": {}}
|
|
])
|
|
|
|
result = clean_db.query_documents_multi_kb(
|
|
kb_names=["kb_x", "kb_y"],
|
|
query_embedding=emb.tolist(),
|
|
n_results=10
|
|
)
|
|
|
|
# Should only get 1 result after dedup
|
|
assert len(result["ids"]) == 1
|
|
|
|
def test_multi_kb_empty(self, clean_db):
|
|
"""Test searching with no matching KBs returns empty."""
|
|
result = clean_db.query_documents_multi_kb(
|
|
kb_names=["nonexistent_kb"],
|
|
query_embedding=[0.1] * 768,
|
|
n_results=5
|
|
)
|
|
assert result["ids"] == []
|
|
assert result["documents"] == []
|
|
assert result["scores"] == []
|
|
|
|
def test_multi_kb_respects_n_results(self, clean_db):
|
|
"""Test that n_results limit is respected."""
|
|
clean_db.create_collection("kb_limit", KBType.GENERIC)
|
|
|
|
np.random.seed(30)
|
|
base = np.random.randn(768).astype(np.float32)
|
|
base = base / np.linalg.norm(base)
|
|
|
|
docs = []
|
|
for i in range(5):
|
|
perturbed = base + np.random.randn(768).astype(np.float32) * 0.01
|
|
perturbed = perturbed / np.linalg.norm(perturbed)
|
|
docs.append({
|
|
"id": f"lim{i}", "text": f"Document {i} for limit test",
|
|
"embedding": perturbed.tolist(), "metadata": {}
|
|
})
|
|
self._add_docs(clean_db, "kb_limit", docs)
|
|
|
|
result = clean_db.query_documents_multi_kb(
|
|
kb_names=["kb_limit"],
|
|
query_embedding=base.tolist(),
|
|
n_results=2
|
|
)
|
|
assert len(result["ids"]) == 2
|
|
|
|
def test_multi_kb_score_ordering(self, clean_db):
|
|
"""Test that results are sorted by descending similarity score."""
|
|
clean_db.create_collection("kb_order", KBType.GENERIC)
|
|
|
|
np.random.seed(40)
|
|
base = np.random.randn(768).astype(np.float32)
|
|
base = base / np.linalg.norm(base)
|
|
|
|
# Add docs with varying similarity
|
|
docs = []
|
|
for i, noise in enumerate([0.01, 0.1, 0.5]):
|
|
perturbed = base + np.random.randn(768).astype(np.float32) * noise
|
|
perturbed = perturbed / np.linalg.norm(perturbed)
|
|
docs.append({
|
|
"id": f"ord{i}", "text": f"Document {i} noise={noise}",
|
|
"embedding": perturbed.tolist(), "metadata": {}
|
|
})
|
|
self._add_docs(clean_db, "kb_order", docs)
|
|
|
|
result = clean_db.query_documents_multi_kb(
|
|
kb_names=["kb_order"],
|
|
query_embedding=base.tolist(),
|
|
n_results=10
|
|
)
|
|
|
|
scores = result["scores"]
|
|
for i in range(len(scores) - 1):
|
|
assert scores[i] >= scores[i + 1], "Scores should be in descending order"
|