|
| 1 | +from dotenv import load_dotenv |
| 2 | +import datetime |
| 3 | +import cocoindex |
| 4 | +import math |
| 5 | +import cocoindex.targets.lancedb as coco_lancedb |
| 6 | + |
| 7 | +# Define LanceDB connection constants |
| 8 | +LANCEDB_URI = "./lancedb_data" # Local directory for LanceDB |
| 9 | +LANCEDB_TABLE = "TextEmbedding" |
| 10 | + |
| 11 | + |
| 12 | +@cocoindex.transform_flow() |
| 13 | +def text_to_embedding( |
| 14 | + text: cocoindex.DataSlice[str], |
| 15 | +) -> cocoindex.DataSlice[list[float]]: |
| 16 | + """ |
| 17 | + Embed the text using a SentenceTransformer model. |
| 18 | + This is a shared logic between indexing and querying, so extract it as a function. |
| 19 | + """ |
| 20 | + return text.transform( |
| 21 | + cocoindex.functions.SentenceTransformerEmbed( |
| 22 | + model="sentence-transformers/all-MiniLM-L6-v2" |
| 23 | + ) |
| 24 | + ) |
| 25 | + |
| 26 | + |
| 27 | +@cocoindex.flow_def(name="TextEmbeddingWithLanceDB") |
| 28 | +def text_embedding_flow( |
| 29 | + flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope |
| 30 | +) -> None: |
| 31 | + """ |
| 32 | + Define an example flow that embeds text into a vector database. |
| 33 | + """ |
| 34 | + data_scope["documents"] = flow_builder.add_source( |
| 35 | + cocoindex.sources.LocalFile(path="markdown_files"), |
| 36 | + refresh_interval=datetime.timedelta(seconds=5), |
| 37 | + ) |
| 38 | + |
| 39 | + doc_embeddings = data_scope.add_collector() |
| 40 | + |
| 41 | + with data_scope["documents"].row() as doc: |
| 42 | + doc["chunks"] = doc["content"].transform( |
| 43 | + cocoindex.functions.SplitRecursively(), |
| 44 | + language="markdown", |
| 45 | + chunk_size=500, |
| 46 | + chunk_overlap=100, |
| 47 | + ) |
| 48 | + |
| 49 | + with doc["chunks"].row() as chunk: |
| 50 | + chunk["embedding"] = text_to_embedding(chunk["text"]) |
| 51 | + doc_embeddings.collect( |
| 52 | + id=cocoindex.GeneratedField.UUID, |
| 53 | + filename=doc["filename"], |
| 54 | + location=chunk["location"], |
| 55 | + text=chunk["text"], |
| 56 | + # 'text_embedding' is the name of the vector we've created the LanceDB table with. |
| 57 | + text_embedding=chunk["embedding"], |
| 58 | + ) |
| 59 | + |
| 60 | + doc_embeddings.export( |
| 61 | + "doc_embeddings", |
| 62 | + coco_lancedb.LanceDB(db_uri=LANCEDB_URI, table_name=LANCEDB_TABLE), |
| 63 | + primary_key_fields=["id"], |
| 64 | + # We cannot enable it when the table has no data yet, as LanceDB requires data to train the index. |
| 65 | + # See: https://github.com/lancedb/lance/issues/4034 |
| 66 | + # |
| 67 | + # vector_indexes=[ |
| 68 | + # cocoindex.VectorIndexDef( |
| 69 | + # "text_embedding", cocoindex.VectorSimilarityMetric.L2_DISTANCE |
| 70 | + # ), |
| 71 | + # ], |
| 72 | + ) |
| 73 | + |
| 74 | + |
| 75 | +@text_embedding_flow.query_handler( |
| 76 | + result_fields=cocoindex.QueryHandlerResultFields( |
| 77 | + embedding=["embedding"], |
| 78 | + score="score", |
| 79 | + ), |
| 80 | +) |
| 81 | +async def search(query: str) -> cocoindex.QueryOutput: |
| 82 | + print("Searching...", query) |
| 83 | + db = await coco_lancedb.connect_async(LANCEDB_URI) |
| 84 | + table = await db.open_table(LANCEDB_TABLE) |
| 85 | + |
| 86 | + # Get the embedding for the query |
| 87 | + query_embedding = await text_to_embedding.eval_async(query) |
| 88 | + |
| 89 | + search = await table.search(query_embedding, vector_column_name="text_embedding") |
| 90 | + search_results = await search.limit(5).to_list() |
| 91 | + |
| 92 | + print(search_results) |
| 93 | + |
| 94 | + return cocoindex.QueryOutput( |
| 95 | + results=[ |
| 96 | + { |
| 97 | + "filename": result["filename"], |
| 98 | + "text": result["text"], |
| 99 | + "embedding": result["text_embedding"], |
| 100 | + # Qdrant's L2 "distance" is squared, so we take the square root to align with normal L2 distance |
| 101 | + "score": math.sqrt(result["_distance"]), |
| 102 | + } |
| 103 | + for result in search_results |
| 104 | + ], |
| 105 | + query_info=cocoindex.QueryInfo( |
| 106 | + embedding=query_embedding, |
| 107 | + similarity_metric=cocoindex.VectorSimilarityMetric.L2_DISTANCE, |
| 108 | + ), |
| 109 | + ) |
0 commit comments