Pythonvector_search.py
Copy
1
import chromadb
2
from openai import OpenAI
3
 
4
client = OpenAI()
5
chroma = chromadb.Client()
6
collection = chroma.create_collection("products")
7
 
8
PRODUCTS = [
9
    {"id": "p1", "name": "AirPods Pro", "category": "audio", "description": "Wireless earbuds with active noise cancellation"},
10
    {"id": "p2", "name": "Sony WH-1000XM5", "category": "audio", "description": "Over-ear headphones with industry-leading noise cancelling"},
11
    {"id": "p3", "name": "MacBook Pro 14", "category": "laptop", "description": "Professional laptop with M3 chip and Liquid Retina display"},
12
    {"id": "p4", "name": "iPad Air", "category": "tablet", "description": "Versatile tablet for creativity and productivity"},
13
    {"id": "p5", "name": "Bose QuietComfort 45", "category": "audio", "description": "Premium wireless headphones with balanced sound"},
14
]
15
 
16
def embed(text: str) -> list[float]:
17
    return client.embeddings.create(model="text-embedding-3-small", input=text).data[0].embedding
18
 
19
# Index products
20
for p in PRODUCTS:
21
    collection.add(
22
        ids=[p["id"]],
23
        documents=[p["description"]],
24
        embeddings=[embed(p["description"])],
25
        metadatas=[{"name": p["name"], "category": p["category"]}],
26
    )
27
 
28
def semantic_search(query: str, category: str | None = None, top_k: int = 3) -> list[dict]:
29
    where = {"category": category} if category else None
30
    results = collection.query(
31
        query_embeddings=[embed(query)],
32
        n_results=top_k,
33
        where=where,
34
    )
35
    return [
36
        {"name": m["name"], "category": m["category"], "description": d}
37
        for m, d in zip(results["metadatas"][0], results["documents"][0])
38
    ]
39
 
40
# Test
41
print("Query: 'wireless audio device'")
42
for r in semantic_search("wireless audio device"):
43
    print(f"  - {r['name']} ({r['category']}): {r['description']}")
44
 
45
print("\nQuery: 'portable computer', category=laptop")
46
for r in semantic_search("portable computer", category="laptop"):
47
    print(f"  - {r['name']}: {r['description']}")
48