#!/usr/bin/env python3
"""
Example usage of the ChromaDB MCP Server
This script demonstrates how to use the MCP tools programmatically.
In a real scenario, these tools would be called by an MCP client.
"""
import asyncio
import json
from app.chroma_mcp_server import (
list_collections,
create_collection,
add_documents,
query_collection,
get_collection,
update_document,
delete_documents,
collection_stats,
peek_collection
)
class Document:
def __init__(self, id: str, content: str, metadata: dict = None, embedding: list = None):
self.id = id
self.content = content
self.metadata = metadata
self.embedding = embedding
async def main():
"""Example usage of ChromaDB MCP tools"""
print("š ChromaDB MCP Server Example")
print("=" * 40)
# 1. List existing collections
print("\n1. Listing collections...")
collections = await list_collections()
print(f"Collections: {collections}")
# 2. Create a new collection
print("\n2. Creating collection 'example_docs'...")
result = await create_collection(
"example_docs",
{"description": "Example document collection for MCP demo"}
)
print(result)
# 3. Add some sample documents
print("\n3. Adding sample documents...")
documents = [
Document(
"doc1",
"This is a document about artificial intelligence and machine learning.",
{"category": "tech", "year": 2024, "tags": ["AI", "ML"]}
),
Document(
"doc2",
"Python is a popular programming language for data science.",
{"category": "programming", "year": 2024, "tags": ["Python", "data science"]}
),
Document(
"doc3",
"Vector databases are useful for semantic search and similarity matching.",
{"category": "databases", "year": 2024, "tags": ["vectors", "search"]}
)
]
result = await add_documents("example_docs", documents)
print(result)
# 4. Get collection statistics
print("\n4. Collection statistics...")
stats = await collection_stats("example_docs")
print(stats)
# 5. Query the collection
print("\n5. Querying for 'artificial intelligence'...")
query_request = {
"collection_name": "example_docs",
"query_texts": ["artificial intelligence"],
"n_results": 2
}
results = await query_collection(type('QueryRequest', (), query_request)())
print(json.dumps(json.loads(results), indent=2))
# 6. Peek at the collection
print("\n6. Peeking at first 2 documents...")
peek = await peek_collection("example_docs", 2)
print(peek)
# 7. Update a document
print("\n7. Updating document 'doc1'...")
update_request = {
"collection_name": "example_docs",
"document_id": "doc1",
"content": "This is an updated document about artificial intelligence, machine learning, and deep learning.",
"metadata": {"category": "tech", "year": 2024, "tags": ["AI", "ML", "deep learning"], "updated": True}
}
result = await update_document(type('UpdateRequest', (), update_request)())
print(result)
# 8. Query again to see the update
print("\n8. Querying after update...")
query_request["query_texts"] = ["deep learning"]
results = await query_collection(type('QueryRequest', (), query_request)())
print(json.dumps(json.loads(results), indent=2))
# 9. Clean up - delete documents
print("\n9. Cleaning up...")
result = await delete_documents("example_docs", ["doc1", "doc2", "doc3"])
print(result)
# 10. Delete the collection
print("\n10. Deleting collection...")
result = await delete_collection("example_docs")
print(result)
print("\nā
Example completed!")
if __name__ == "__main__":
asyncio.run(main())
