mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
zeekay
898d818183
- Make MemoryResponse.memory field optional (used by legacy /v1/remember endpoint) - Fix datetime.utcnow() deprecation warning (use datetime.now(timezone.utc)) - Make LanceDB import optional in backend_registry.py - Remove duplicate imports in db/__init__.py - Add singleton get_sqlite_client() function - Fix import path (embeddings vs embedding) in mcp/server.py - Move module-level MCPMemoryServer instantiation to main() - Add 7 missing abstract method implementations to InfinityClient - Fix test patching target (get_db_client not get_client) - Add proper HTTPException re-raise in exception handlers - Add reset_db_client() and reset_memory_service() to test fixtures - Add list/DataFrame type handling in search_memories() - Delete stale build directory and lancedb_client.py - Fix nested workspace issue (remove hanzo-agent from uv.workspace) - Add test dependencies (polars, httpx, fastapi, fastembed) - Update CI workflow to install test dependencies All 84 hanzo-memory tests pass, 3156 core SDK tests pass.
229 lines
7.2 KiB
Python
229 lines
7.2 KiB
Python
"""Tests for Knowledge Base API endpoints."""
|
|
|
|
import polars as pl
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
from hanzo_memory.server import app
|
|
|
|
|
|
class TestKnowledgeAPI:
|
|
"""Test Knowledge Base API endpoints."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def setup(self, mock_auth, mock_db_client, mock_services):
|
|
"""Set up test client and mocks."""
|
|
self.client = TestClient(app)
|
|
self.headers = {"Authorization": "Bearer test-token"}
|
|
self.mock_db = mock_db_client
|
|
self.mock_embedding_service = mock_services["embedding"]
|
|
|
|
def test_create_knowledge_base(self):
|
|
"""Test creating a knowledge base."""
|
|
self.mock_db.create_knowledge_base.return_value = {
|
|
"kb_id": "kb123",
|
|
"name": "Test KB",
|
|
}
|
|
|
|
response = self.client.post(
|
|
"/v1/kb/create",
|
|
json={
|
|
"userid": "user123",
|
|
"name": "Test KB",
|
|
"kb_id": "kb123",
|
|
},
|
|
headers=self.headers,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["kb_id"] == "kb123"
|
|
assert "Test KB" in data["message"]
|
|
|
|
# Verify DB calls
|
|
self.mock_db.create_knowledge_base.assert_called_once()
|
|
call_args = self.mock_db.create_knowledge_base.call_args[1]
|
|
assert call_args["kb_id"] == "kb123"
|
|
assert call_args["user_id"] == "user123"
|
|
assert call_args["name"] == "Test KB"
|
|
|
|
def test_create_knowledge_base_default_project(self):
|
|
"""Test creating KB with default project."""
|
|
# Mock project creation to raise (already exists)
|
|
self.mock_db.create_project.side_effect = Exception("Already exists")
|
|
self.mock_db.create_knowledge_base.return_value = {"kb_id": "kb123"}
|
|
|
|
response = self.client.post(
|
|
"/v1/kb/create",
|
|
json={
|
|
"userid": "user123",
|
|
"name": "Test KB",
|
|
},
|
|
headers=self.headers,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
|
|
# Verify default project was attempted
|
|
self.mock_db.create_project.assert_called_once()
|
|
project_args = self.mock_db.create_project.call_args[1]
|
|
assert project_args["project_id"] == "default-user123"
|
|
assert project_args["user_id"] == "user123"
|
|
|
|
def test_list_knowledge_bases(self):
|
|
"""Test listing knowledge bases."""
|
|
response = self.client.get(
|
|
"/v1/kb/list",
|
|
params={"userid": "user123"},
|
|
headers=self.headers,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["userid"] == "user123"
|
|
assert "knowledge_bases" in data
|
|
assert data["total"] == 0
|
|
|
|
def test_add_facts(self):
|
|
"""Test adding facts to a knowledge base."""
|
|
self.mock_embedding_service.embed_text.return_value = [[0.1] * 384]
|
|
self.mock_db.add_fact.return_value = {"fact_id": "fact123"}
|
|
|
|
response = self.client.post(
|
|
"/v1/kb/facts/add",
|
|
json={
|
|
"userid": "user123",
|
|
"kb_id": "kb123",
|
|
"facts": [
|
|
{
|
|
"content": "Test fact 1",
|
|
"metadata": {"type": "test"},
|
|
},
|
|
{
|
|
"content": "Test fact 2",
|
|
"parent_id": "fact123",
|
|
},
|
|
],
|
|
},
|
|
headers=self.headers,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["kb_id"] == "kb123"
|
|
assert data["facts_added"] == 2
|
|
assert len(data["facts"]) == 2
|
|
|
|
# Verify embeddings were generated
|
|
assert self.mock_embedding_service.embed_text.call_count == 2
|
|
|
|
# Verify facts were added
|
|
assert self.mock_db.add_fact.call_count == 2
|
|
|
|
def test_get_facts_with_query(self):
|
|
"""Test searching facts with a query."""
|
|
# Mock search results
|
|
mock_df = pl.DataFrame(
|
|
{
|
|
"fact_id": ["fact1", "fact2"],
|
|
"content": ["Fact 1", "Fact 2"],
|
|
"parent_id": ["", "parent1"],
|
|
"metadata": ['{"type": "test"}', '{"type": "example"}'],
|
|
"_similarity": [0.95, 0.85],
|
|
}
|
|
)
|
|
|
|
self.mock_db.search_facts.return_value = mock_df
|
|
self.mock_embedding_service.embed_text.return_value = [[0.2] * 384]
|
|
|
|
response = self.client.post(
|
|
"/v1/kb/facts/get",
|
|
json={
|
|
"userid": "user123",
|
|
"kb_id": "kb123",
|
|
"query": "test query",
|
|
"limit": 10,
|
|
},
|
|
headers=self.headers,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["kb_id"] == "kb123"
|
|
assert len(data["facts"]) == 2
|
|
assert data["facts"][0]["fact_id"] == "fact1"
|
|
assert data["facts"][0]["similarity_score"] == 0.95
|
|
|
|
# Verify search was called
|
|
self.mock_db.search_facts.assert_called_once()
|
|
search_args = self.mock_db.search_facts.call_args[1]
|
|
assert search_args["kb_id"] == "kb123"
|
|
assert search_args["limit"] == 10
|
|
|
|
def test_get_facts_without_query(self):
|
|
"""Test getting facts without a search query."""
|
|
# API requires query parameter
|
|
response = self.client.post(
|
|
"/v1/kb/facts/get",
|
|
json={
|
|
"userid": "user123",
|
|
"kb_id": "kb123",
|
|
},
|
|
headers=self.headers,
|
|
)
|
|
assert response.status_code == 400
|
|
|
|
# With query, should return empty results
|
|
response = self.client.post(
|
|
"/v1/kb/facts/get",
|
|
json={
|
|
"userid": "user123",
|
|
"kb_id": "kb123",
|
|
"query": "test",
|
|
},
|
|
headers=self.headers,
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["kb_id"] == "kb123"
|
|
assert data["facts"] == []
|
|
assert data["total"] == 0
|
|
|
|
def test_delete_fact(self):
|
|
"""Test deleting a fact."""
|
|
response = self.client.post(
|
|
"/v1/kb/facts/delete",
|
|
json={
|
|
"userid": "user123",
|
|
"kb_id": "kb123",
|
|
"fact_id": "fact123",
|
|
"cascade": True,
|
|
},
|
|
headers=self.headers,
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["kb_id"] == "kb123"
|
|
assert data["fact_id"] == "fact123"
|
|
assert data["deleted"] is True
|
|
assert data["cascade"] is True
|
|
|
|
def test_knowledge_base_error_handling(self):
|
|
"""Test error handling in knowledge base operations."""
|
|
self.mock_db.create_knowledge_base.side_effect = Exception("DB Error")
|
|
|
|
response = self.client.post(
|
|
"/v1/kb/create",
|
|
json={
|
|
"userid": "user123",
|
|
"name": "Test KB",
|
|
},
|
|
headers=self.headers,
|
|
)
|
|
|
|
assert response.status_code == 500
|
|
data = response.json()
|
|
assert "detail" in data
|
|
assert "DB Error" in data["detail"]
|