forked from testcontainers/testcontainers-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathexample_basic.py
More file actions
65 lines (50 loc) · 2.05 KB
/
example_basic.py
File metadata and controls
65 lines (50 loc) · 2.05 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
import chromadb
from chromadb.config import Settings
from testcontainers.chroma import ChromaContainer
def basic_example():
with ChromaContainer() as chroma:
# Get connection URL
connection_url = chroma.get_connection_url()
# Create Chroma client
client = chromadb.HttpClient(host=connection_url, settings=Settings(allow_reset=True))
# Create a collection
collection_name = "test_collection"
collection = client.create_collection(name=collection_name)
print(f"Created collection: {collection_name}")
# Add documents and embeddings
documents = [
"This is a test document about AI",
"Machine learning is a subset of AI",
"Deep learning uses neural networks",
]
embeddings = [
[0.1, 0.2, 0.3], # Simple example embeddings
[0.2, 0.3, 0.4],
[0.3, 0.4, 0.5],
]
ids = ["doc1", "doc2", "doc3"]
metadatas = [
{"source": "test1", "category": "AI"},
{"source": "test2", "category": "ML"},
{"source": "test3", "category": "DL"},
]
collection.add(documents=documents, embeddings=embeddings, ids=ids, metadatas=metadatas)
print("Added documents to collection")
# Query the collection
results = collection.query(query_embeddings=[[0.1, 0.2, 0.3]], n_results=2)
print("\nQuery results:")
print(f"Documents: {results['documents'][0]}")
print(f"Distances: {results['distances'][0]}")
print(f"Metadatas: {results['metadatas'][0]}")
# Get collection info
collection_info = client.get_collection(collection_name)
print("\nCollection info:")
print(f"Name: {collection_info.name}")
print(f"Count: {collection_info.count()}")
# List all collections
collections = client.list_collections()
print("\nAvailable collections:")
for coll in collections:
print(f"- {coll.name}")
if __name__ == "__main__":
basic_example()