File size: 554 Bytes
a228306
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3c2b266
a228306
 
 
 
 
 
 
 
 
 
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
import uuid

from config import client, COLLECTION
from embeddings import embed


def index_doc(text, payload):

    vector = embed(text)

    client.upsert(
        collection_name=COLLECTION,
        points=[{
            "id": str(uuid.uuid4()),
            "vector": vector,
            "payload": payload | {"text": text}
        }]
    )


def search(query, k=3):

    vector = embed(query)

    results = client.search(
        collection_name=COLLECTION,
        query_vector=vector,
        limit=k
    )

    return [r.payload for r in results]