import csv
import json
import os
import tempfile
import pytest
from seekdb_mcp.server import (
add_data_to_collection,
ai_complete,
ai_rerank,
create_ai_model,
create_ai_model_endpoint,
create_collection,
delete_collection,
delete_documents,
drop_ai_model,
drop_ai_model_endpoint,
execute_sql,
export_csv_file_from_seekdb,
get_ai_model_endpoints,
get_current_time,
get_registered_ai_models,
hybrid_search,
import_csv_file_to_seekdb,
list_collections,
query_collection,
full_text_search,
has_collection,
seekdb_memory_insert,
seekdb_memory_query,
seekdb_memory_delete,
seekdb_memory_update,
update_collection,
seekdb_memory_collection_name,
)
@pytest.fixture
def test_collection():
"""Fixture that creates a collection before test and deletes it after."""
collection_name = "test_collection"
create_collection(collection_name)
yield collection_name
delete_collection(collection_name)
@pytest.fixture
def test_ai_model():
"""Fixture that creates an AI embedding model before test and drops it after."""
model_name = "seekdb_embed"
model_type = "dense_embedding"
provider_model_name = "BAAI/bge-m3"
create_ai_model(
model_name=model_name, model_type=model_type, provider_model_name=provider_model_name
)
yield model_name
drop_ai_model(model_name)
@pytest.fixture
def test_ai_model_endpoint(test_ai_model):
"""Fixture that creates an AI model endpoint before test and drops it after.
Depends on test_ai_model fixture to ensure the AI model exists first.
"""
endpoint_name = "seekdb_embed_endpoint"
url = "https://api.siliconflow.cn/v1/embeddings"
access_key = os.getenv("siliconflow_access_key")
provider = "siliconflow"
create_ai_model_endpoint(
endpoint_name=endpoint_name,
ai_model_name=test_ai_model,
url=url,
access_key=access_key,
provider=provider,
)
yield endpoint_name
drop_ai_model_endpoint(endpoint_name)
@pytest.fixture
def test_memory():
create_collection(seekdb_memory_collection_name)
yield
delete_collection(seekdb_memory_collection_name)
@pytest.fixture
def student_csv_file():
# Create a temporary directory
with tempfile.TemporaryDirectory() as tmpdir:
file_path = os.path.join(tmpdir, "student.csv")
# Write CSV content
with open(file_path, "w", encoding="utf-8", newline="") as f:
writer = csv.writer(f)
writer.writerow(["name", "age"])
writer.writerow(["张三", "20"])
writer.writerow(["李四", "21"])
yield file_path
# Cleanup: drop the table if it exists
table_name = "student"
drop_sql = f"DROP TABLE IF EXISTS `{table_name}`"
execute_sql(drop_sql)
@pytest.fixture
def documents_csv_file():
# Create a temporary directory
with tempfile.TemporaryDirectory() as tmpdir:
file_path = os.path.join(tmpdir, "documents.csv")
# Write CSV content
with open(file_path, "w", encoding="utf-8", newline="") as f:
writer = csv.writer(f)
writer.writerow(["document", "category"])
writer.writerow(["Machine learning is a subset of artificial intelligence", "AI"])
writer.writerow(["Vector databases enable semantic search", "Database"])
yield file_path
# Cleanup: delete the collection if it exists
collection_name = "documents"
delete_collection(collection_name)
def test_execute_sql():
sql = "select 1"
result = execute_sql(sql)
result_dict = json.loads(result)
assert result_dict["success"] is True
assert result_dict["data"] == [["1"]], f"Expected [['1']], got {result_dict['data']}"
def test_get_current_time():
result = get_current_time()
print(result)
result_dict = json.loads(result)
assert result_dict["success"] is True
def test_create_collection():
collection_name = "test_collection"
result = create_collection(collection_name)
result_dict = json.loads(result)
assert result_dict["success"] is True
delete_collection(collection_name)
def test_has_collection(test_collection):
# Test that collection exists
result = has_collection(test_collection)
result_dict = json.loads(result)
assert result_dict["success"] is True
assert result_dict["exists"] is True, (
f"Expected 'exists' to be True, got {result_dict['exists']}"
)
def test_list_collection(test_collection):
result = list_collections()
result_dict = json.loads(result)
assert result_dict["success"] is True
assert test_collection in result_dict["collections"], (
f"Expected '{test_collection}' in collections, got {result_dict['collections']}"
)
def test_add_data_to_collection(test_collection):
ids = ["1", "2"]
documents = ["I love apple", "I love pear"]
metadatas = [{"category": "preference", "index": 0}, {"category": "preference", "index": 1}]
result = add_data_to_collection(
collection_name=test_collection, ids=ids, documents=documents, metadatas=metadatas
)
result_dict = json.loads(result)
assert result_dict["success"] is True
def test_delete_documents(test_collection):
test_add_data_to_collection(test_collection)
where_document = {"$contains": "apple"}
result = delete_documents(collection_name=test_collection, where_document=where_document)
result_dict = json.loads(result)
assert result_dict["success"] is True, (
f"Expected delete_documents to succeed, got {result_dict}"
)
query_result = query_collection(
collection_name=test_collection, query_texts=["apple"], n_results=10
)
query_result_dict = json.loads(query_result)
# Verify "I love apple" is not in the returned documents
returned_documents = query_result_dict["data"]["documents"]
all_docs = [doc for docs in returned_documents for doc in docs]
assert "I love apple" not in all_docs, (
f"Expected 'I love apple' to be deleted, but found it in {all_docs}"
)
def test_query_collection(test_collection):
# Add data first
test_add_data_to_collection(test_collection)
query_texts = ["pear"]
result = query_collection(collection_name=test_collection, query_texts=query_texts, n_results=1)
result_dict = json.loads(result)
assert result_dict["success"] is True
# Check if the returned documents contain "pear"
documents = result_dict["data"]["documents"]
assert any("pear" in doc for docs in documents for doc in docs), (
f"Expected 'pear' in documents, got {documents}"
)
def test_update_collection(test_collection):
test_add_data_to_collection(test_collection)
ids = ["1"]
documents = ["I love banana"]
result = update_collection(collection_name=test_collection, ids=ids, documents=documents)
result_dict = json.loads(result)
assert result_dict["success"] is True, (
f"Expected update_collection to succeed, got {result_dict}"
)
# Query to verify document with id=1 has been updated to "I love banana"
query_result = query_collection(
collection_name=test_collection, query_texts=["banana"], n_results=1
)
query_result_dict = json.loads(query_result)
returned_documents = query_result_dict["data"]["documents"]
assert returned_documents[0][0] == "I love banana", (
f"Expected 'I love banana', got {returned_documents[0][0]}"
)
def test_delete_collection():
collection_name = "test_collection"
result = delete_collection(collection_name)
result_dict = json.loads(result)
assert result_dict["success"] is True
def test_full_text_search():
drop_table = "drop table if exists sport_data_whole"
execute_sql(drop_table)
create_table = """
CREATE TABLE if not exists sport_data_whole (
event VARCHAR(64),
date VARCHAR(20),
news VARCHAR(65535),
FULLTEXT INDEX ft_idx1_news(news)
WITH PARSER ik PARSER_PROPERTIES = (ik_mode = "max_word")
)
"""
execute_sql(create_table)
insert_data = """insert into sport_data_whole values
('世界杯','2025-12-11 16:00:00','格策 黄牌 冠军 乌龙球 博阿滕 比赛 观众 球员'),
('足球赛','2025-12-11 15:00:00','逆转 博阿滕 犯规 黄牌 进球 球员 精彩 精彩'),
('奥运会','2025-12-11 15:17:00','逆转 红牌 红牌 犯规 逆转 比赛 观众 观众')
"""
execute_sql(insert_data)
table_name = "sport_data_whole"
column_name = "news"
search_expr = "+黄牌 +进球"
result = full_text_search(
table_name=table_name, column_name=column_name, search_expr=search_expr
)
result_dict = json.loads(result)
assert result_dict["success"] is True
assert result_dict["data"][0][2] == "逆转 博阿滕 犯规 黄牌 进球 球员 精彩 精彩", (
f"Expected '逆转 博阿滕 犯规 黄牌 进球 球员 精彩 精彩', got {result_dict['data'][0][2]}"
)
drop_table = "drop table if exists sport_data_whole"
execute_sql(drop_table)
def test_hybrid_search(test_collection):
documents = [
"Machine learning is a subset of artificial intelligence",
"Python is a popular programming language",
"Vector databases enable semantic search",
"Neural networks are inspired by the human brain",
"Natural language processing helps computers understand text",
]
ids = ["id1", "id2", "id3", "id4", "id5"]
metadatas = [
{"category": "AI", "index": 0, "year": 2021},
{"category": "Programming", "index": 1, "year": 2018},
{"category": "Database", "index": 2, "year": 2027},
{"category": "AI", "index": 3, "year": 2019},
{"category": "NLP", "index": 4, "year": 2016},
]
add_data_to_collection(
collection_name=test_collection, ids=ids, documents=documents, metadatas=metadatas
)
fulltext_search_keyword = "machine learning"
fulltext_where = {"category": {"$eq": "AI"}}
knn_query_texts = ["AI research"]
knn_where = {"year": {"$gte": 2020}}
result = hybrid_search(
collection_name=test_collection,
fulltext_search_keyword=fulltext_search_keyword,
fulltext_where=fulltext_where,
knn_query_texts=knn_query_texts,
knn_where=knn_where,
)
result_dict = json.loads(result)
assert result_dict["success"] is True
assert result_dict["data"]["ids"][0] == ["id1", "id3"], (
f"Expected ['id1', 'id3'], got {result_dict['data']['ids'][0]}"
)
def test_create_ai_model():
model_name = "seekdb_embed"
model_type = "dense_embedding"
provider_model_name = "BAAI/bge-m3"
result = create_ai_model(
model_name=model_name, model_type=model_type, provider_model_name=provider_model_name
)
result_dict = json.loads(result)
assert result_dict["success"] is True
drop_ai_model(model_name)
def test_drop_ai_model():
model_name = "seekdb_embed"
result = drop_ai_model(model_name)
result_dict = json.loads(result)
assert result_dict["success"] is True
def test_create_ai_model_endpoint(test_ai_model):
endpoint_name = "seekdb_embed_endpoint"
url = "https://api.siliconflow.cn/v1/embeddings"
access_key = os.getenv("siliconflow_access_key")
provider = "siliconflow"
result = create_ai_model_endpoint(
endpoint_name=endpoint_name,
ai_model_name=test_ai_model,
url=url,
access_key=access_key,
provider=provider,
)
result_dict = json.loads(result)
assert result_dict["success"] is True
drop_ai_model_endpoint(endpoint_name)
def test_drop_ai_model_endpoint():
endpoint_name = "seekdb_embed_endpoint"
result = drop_ai_model_endpoint(endpoint_name)
result_dict = json.loads(result)
assert result_dict["success"] is True
def test_ai_complete():
model_name = "seekdb_complete"
model_type = "completion"
provider_model_name = "THUDM/GLM-4-9B-0414"
create_ai_model(
model_name=model_name, model_type=model_type, provider_model_name=provider_model_name
)
endpoint_name = "seekdb_complete_endpoint"
url = "https://api.siliconflow.cn/v1/chat/completions"
access_key = os.getenv("siliconflow_access_key")
provider = "siliconflow"
create_ai_model_endpoint(
endpoint_name=endpoint_name,
ai_model_name=model_name,
url=url,
access_key=access_key,
provider=provider,
)
prompt = "Translate 'Hello World' to Chinese"
result = ai_complete(model_name, prompt)
result_dict = json.loads(result)
assert result_dict["success"] is True
response_text = result_dict["response"]
assert "你好" in response_text and "世界" in response_text, (
f"Expected '你好，世界' in response, got {response_text}"
)
def test_ai_rerank():
model_name = "seekdb_rerank"
model_type = "rerank"
provider_model_name = "BAAI/bge-reranker-v2-m3"
create_ai_model(
model_name=model_name, model_type=model_type, provider_model_name=provider_model_name
)
endpoint_name = "seekdb_rerank_endpoint"
url = "https://api.siliconflow.cn/v1/rerank"
access_key = os.getenv("siliconflow_access_key")
provider = "siliconflow"
create_ai_model_endpoint(
endpoint_name=endpoint_name,
ai_model_name=model_name,
url=url,
access_key=access_key,
provider=provider,
)
query = "Apple"
documents = ["vegetable", "fruit", "banana", "apple"]
result = ai_rerank(model_name, query, documents)
result_dict = json.loads(result)
assert result_dict["success"] is True
assert result_dict["reranked_documents"] == ["apple", "banana", "fruit", "vegetable"], (
f"Expected ['apple', 'banana', 'fruit', 'vegetable'], got {result_dict['reranked_documents']}"
)
def test_get_registered_ai_model(test_ai_model):
result = get_registered_ai_models()
result_dict = json.loads(result)
assert result_dict["success"] is True
# Check that seekdb_embed is in the returned models
model_names = [model["NAME"] for model in result_dict["models"]]
assert test_ai_model in model_names, (
f"Expected 'seekdb_embed' in model names, got {model_names}"
)
def test_get_ai_model_endpoints(test_ai_model_endpoint):
result = get_ai_model_endpoints()
result_dict = json.loads(result)
assert result_dict["success"] is True
# Check that the endpoint we created is in the returned list
endpoint_names = [endpoint["ENDPOINT_NAME"] for endpoint in result_dict["endpoints"]]
assert test_ai_model_endpoint in endpoint_names, (
f"Expected '{test_ai_model_endpoint}' in endpoint names, got {endpoint_names}"
)
def test_seekdb_memory_insert_and_query(test_memory):
content = "I love apple"
meta = {}
# Insert the memory
insert_result = seekdb_memory_insert(content=content, meta=meta)
insert_result_dict = json.loads(insert_result)
assert insert_result_dict["success"] is True, (
f"Expected insert to succeed, got {insert_result_dict}"
)
# Query to verify insertion
query_result = seekdb_memory_query(query="apple", topk=5)
query_result_dict = json.loads(query_result)
# Find the memory we just inserted
memories = query_result_dict["memories"]
matching_memories = [m for m in memories if m["content"] == content]
assert len(matching_memories) > 0, f"Expected to find '{content}' in memories, got {memories}"
def test_seekdb_memory_update(test_memory):
content = "I love apple"
meta = {}
# Insert the memory
insert_result = seekdb_memory_insert(content=content, meta=meta)
insert_result_dict = json.loads(insert_result)
assert insert_result_dict["success"] is True, (
f"Expected insert to succeed, got {insert_result_dict}"
)
# Query to verify insertion
query_result = seekdb_memory_query(query="apple", topk=5)
query_result_dict = json.loads(query_result)
# Find the memory we just inserted
memories = query_result_dict["memories"]
matching_memories = [m for m in memories if m["content"] == content]
assert len(matching_memories) > 0, f"Expected to find '{content}' in memories, got {memories}"
# Get the mem_id of the memory we want to update
mem_id = matching_memories[0]["mem_id"]
# Update the memory content
new_content = "I love pear"
update_result = seekdb_memory_update(mem_id=mem_id, content=new_content, meta=meta)
update_result_dict = json.loads(update_result)
assert update_result_dict["success"] is True, (
f"Expected update to succeed, got {update_result_dict}"
)
# Query to verify update
query_result_after_update = seekdb_memory_query(query="pear", topk=5)
query_result_after_update_dict = json.loads(query_result_after_update)
# Verify the memory was updated - same mem_id should now have new content
memories_after_update = query_result_after_update_dict["memories"]
matching_updated_memories = [m for m in memories_after_update if m["content"] == new_content]
# Verify it's the same memory record (same mem_id)
assert matching_updated_memories[0]["mem_id"] == mem_id, (
f"Expected mem_id to be '{mem_id}', got '{matching_updated_memories[0]['mem_id']}'"
)
def test_seekdb_memory_delete(test_memory):
content = "I love apple"
meta = {}
# Insert the memory
insert_result = seekdb_memory_insert(content=content, meta=meta)
insert_result_dict = json.loads(insert_result)
assert insert_result_dict["success"] is True, (
f"Expected insert to succeed, got {insert_result_dict}"
)
# Query to verify insertion
query_result = seekdb_memory_query(query="apple", topk=5)
query_result_dict = json.loads(query_result)
# Find the memory we just inserted
memories = query_result_dict["memories"]
matching_memories = [m for m in memories if m["content"] == content]
assert len(matching_memories) > 0, f"Expected to find '{content}' in memories, got {memories}"
# Get the mem_id of the memory we want to delete
mem_id = matching_memories[0]["mem_id"]
# Delete the memory
delete_result = seekdb_memory_delete(mem_id=mem_id)
delete_result_dict = json.loads(delete_result)
assert delete_result_dict["success"] is True, (
f"Expected delete to succeed, got {delete_result_dict}"
)
# Query to verify deletion
query_result_after_delete = seekdb_memory_query(query="apple", topk=5)
query_result_after_delete_dict = json.loads(query_result_after_delete)
# Verify the memory was deleted - same mem_id should no longer exist
memories_after_delete = query_result_after_delete_dict["memories"]
matching_deleted_memories = [m for m in memories_after_delete if m["mem_id"] == mem_id]
assert len(matching_deleted_memories) == 0, (
f"Expected mem_id '{mem_id}' to be deleted, but still found it in memories: {memories_after_delete}"
)
def test_import_csv_file_to_seekdb_without_vector(student_csv_file):
file_path = student_csv_file
table_name = os.path.splitext(os.path.basename(file_path))[0]
# Import CSV file to seekdb (without vector)
result = import_csv_file_to_seekdb(file_path)
result_dict = json.loads(result)
assert result_dict["success"] is True, f"Expected import to succeed, got {result_dict}"
# Verify table exists by querying it
check_table_sql = f"SELECT 1 FROM `{table_name}` LIMIT 1"
check_result = execute_sql(check_table_sql)
check_result_dict = json.loads(check_result)
assert check_result_dict["success"] is True, (
f"Expected table '{table_name}' to exist, but query failed: {check_result_dict}"
)
# Query all data from the table
query_sql = f"SELECT name, age FROM `{table_name}` ORDER BY age"
query_result = execute_sql(query_sql)
query_result_dict = json.loads(query_result)
assert query_result_dict["success"] is True, (
f"Expected query to succeed, got {query_result_dict}"
)
# Verify the data
data = query_result_dict["data"]
assert len(data) == 2, f"Expected 2 rows, got {len(data)}"
# Check first row: 张三, 20
assert data[0][0] == "张三", f"Expected first row name to be '张三', got '{data[0][0]}'"
assert data[0][1] == "20", f"Expected first row age to be '20', got '{data[0][1]}'"
# Check second row: 李四, 21
assert data[1][0] == "李四", f"Expected second row name to be '李四', got '{data[1][0]}'"
assert data[1][1] == "21", f"Expected second row age to be '21', got '{data[1][1]}'"
def test_import_csv_file_to_seekdb_with_vector(documents_csv_file):
file_path = documents_csv_file
collection_name = os.path.splitext(os.path.basename(file_path))[0]
# Import CSV file to seekdb (with vector)
result = import_csv_file_to_seekdb(file_path, 1)
result_dict = json.loads(result)
assert result_dict["success"] is True, f"Expected import to succeed, got {result_dict}"
# Verify the 'documents' collection exists
has_result = has_collection(collection_name)
has_result_dict = json.loads(has_result)
assert has_result_dict["success"] is True
assert has_result_dict["exists"] is True, (
f"Expected collection '{collection_name}' to exist, got {has_result_dict}"
)
# Query the collection and verify the content matches documents.csv
# CSV content:
# document,category
# Machine learning is a subset of artificial intelligence,AI
# Vector databases enable semantic search,Database
# Query for "Machine learning"
query_result1 = query_collection(
collection_name=collection_name, query_texts=["Machine learning"], n_results=2
)
query_result1_dict = json.loads(query_result1)
assert query_result1_dict["success"] is True, (
f"Expected query to succeed, got {query_result1_dict}"
)
# Verify the documents contain expected content
documents = query_result1_dict["data"]["documents"]
all_docs = [doc for docs in documents for doc in docs]
assert any(
"Machine learning is a subset of artificial intelligence" in doc for doc in all_docs
), (
f"Expected 'Machine learning is a subset of artificial intelligence' in documents, got {all_docs}"
)
assert any("Vector databases enable semantic search" in doc for doc in all_docs), (
f"Expected 'Vector databases enable semantic search' in documents, got {all_docs}"
)
# Verify metadatas contain the category information
metadatas = query_result1_dict["data"]["metadatas"]
all_metadatas = [m for metas in metadatas for m in metas]
categories = [m.get("category") for m in all_metadatas if m.get("category")]
assert "AI" in categories, f"Expected 'AI' in categories, got {categories}"
assert "Database" in categories, f"Expected 'Database' in categories, got {categories}"
def test_export_csv_file_from_seekdb_with_table(student_csv_file):
file_path_input = student_csv_file
table_name = os.path.splitext(os.path.basename(file_path_input))[0]
# Create output file in the same temp directory
output_dir = os.path.dirname(file_path_input)
file_path_output = os.path.join(output_dir, "student-output.csv")
# Import CSV file to seekdb
import_result = import_csv_file_to_seekdb(file_path_input)
import_result_dict = json.loads(import_result)
assert import_result_dict["success"] is True, (
f"Expected import to succeed, got {import_result_dict}"
)
# Export data from seekdb to CSV file
export_result = export_csv_file_from_seekdb(table_name, file_path_output)
export_result_dict = json.loads(export_result)
assert export_result_dict["success"] is True, (
f"Expected export to succeed, got {export_result_dict}"
)
# Read input CSV file
with open(file_path_input, "r", encoding="utf-8") as f:
input_reader = csv.reader(f)
input_rows = list(input_reader)
# Read output CSV file
with open(file_path_output, "r", encoding="utf-8") as f:
output_reader = csv.reader(f)
output_rows = list(output_reader)
# Verify headers are the same
assert input_rows[0] == output_rows[0], (
f"Expected headers to match. Input: {input_rows[0]}, Output: {output_rows[0]}"
)
# Verify row count is the same (excluding header)
assert len(input_rows) == len(output_rows), (
f"Expected same number of rows. Input: {len(input_rows)}, Output: {len(output_rows)}"
)
# Verify data rows match (compare as sets to handle potential order differences)
input_data = set(tuple(row) for row in input_rows[1:])
output_data = set(tuple(row) for row in output_rows[1:])
assert input_data == output_data, (
f"Expected data to match. Input: {input_data}, Output: {output_data}"
)
def test_export_csv_file_from_seekdb_with_collection(documents_csv_file):
file_path_input = documents_csv_file
collection_name = os.path.splitext(os.path.basename(file_path_input))[0]
# Create output file in the same temp directory
output_dir = os.path.dirname(file_path_input)
file_path_output = os.path.join(output_dir, "documents-output.csv")
# Import CSV file to seekdb as collection (with vector)
import_result = import_csv_file_to_seekdb(file_path_input, 1)
import_result_dict = json.loads(import_result)
assert import_result_dict["success"] is True, (
f"Expected import to succeed, got {import_result_dict}"
)
# Export data from seekdb collection to CSV file
export_result = export_csv_file_from_seekdb(collection_name, file_path_output)
export_result_dict = json.loads(export_result)
assert export_result_dict["success"] is True, (
f"Expected export to succeed, got {export_result_dict}"
)
# Read input CSV file
with open(file_path_input, "r", encoding="utf-8") as f:
input_reader = csv.reader(f)
input_rows = list(input_reader)
# Read output CSV file
with open(file_path_output, "r", encoding="utf-8") as f:
output_reader = csv.reader(f)
output_rows = list(output_reader)
# Verify headers are the same
assert input_rows[0] == output_rows[0], (
f"Expected headers to match. Input: {input_rows[0]}, Output: {output_rows[0]}"
)
# Verify row count is the same (excluding header)
assert len(input_rows) == len(output_rows), (
f"Expected same number of rows. Input: {len(input_rows)}, Output: {len(output_rows)}"
)
# Verify data rows match (compare as sets to handle potential order differences)
input_data = set(tuple(row) for row in input_rows[1:])
output_data = set(tuple(row) for row in output_rows[1:])
assert input_data == output_data, (
f"Expected data to match. Input: {input_data}, Output: {output_data}"
)
