|
| 1 | +from dotenv import load_dotenv |
| 2 | +from psycopg_pool import ConnectionPool |
| 3 | +import cocoindex |
| 4 | +import os |
| 5 | +from typing import Any |
| 6 | + |
| 7 | + |
| 8 | +@cocoindex.transform_flow() |
| 9 | +def text_to_embedding( |
| 10 | + text: cocoindex.DataSlice[str], |
| 11 | +) -> cocoindex.DataSlice[list[float]]: |
| 12 | + """ |
| 13 | + Embed the text using a SentenceTransformer model. |
| 14 | + This is a shared logic between indexing and querying, so extract it as a function. |
| 15 | + """ |
| 16 | + return text.transform( |
| 17 | + cocoindex.functions.SentenceTransformerEmbed( |
| 18 | + model="sentence-transformers/all-MiniLM-L6-v2" |
| 19 | + ) |
| 20 | + ) |
| 21 | + |
| 22 | + |
| 23 | +@cocoindex.flow_def(name="AzureBlobTextEmbedding") |
| 24 | +def azure_blob_text_embedding_flow( |
| 25 | + flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope |
| 26 | +) -> None: |
| 27 | + """ |
| 28 | + Define an example flow that embeds text from Azure Blob Storage into a vector database. |
| 29 | + """ |
| 30 | + account_name = os.environ["AZURE_STORAGE_ACCOUNT_NAME"] |
| 31 | + container_name = os.environ["AZURE_BLOB_CONTAINER_NAME"] |
| 32 | + prefix = os.environ.get("AZURE_BLOB_PREFIX", None) |
| 33 | + |
| 34 | + data_scope["documents"] = flow_builder.add_source( |
| 35 | + cocoindex.sources.AzureBlob( |
| 36 | + account_name=account_name, |
| 37 | + container_name=container_name, |
| 38 | + prefix=prefix, |
| 39 | + included_patterns=["*.md", "*.mdx", "*.txt", "*.docx"], |
| 40 | + binary=False, |
| 41 | + ) |
| 42 | + ) |
| 43 | + |
| 44 | + doc_embeddings = data_scope.add_collector() |
| 45 | + |
| 46 | + with data_scope["documents"].row() as doc: |
| 47 | + doc["chunks"] = doc["content"].transform( |
| 48 | + cocoindex.functions.SplitRecursively(), |
| 49 | + language="markdown", |
| 50 | + chunk_size=2000, |
| 51 | + chunk_overlap=500, |
| 52 | + ) |
| 53 | + |
| 54 | + with doc["chunks"].row() as chunk: |
| 55 | + chunk["embedding"] = text_to_embedding(chunk["text"]) |
| 56 | + doc_embeddings.collect( |
| 57 | + filename=doc["filename"], |
| 58 | + location=chunk["location"], |
| 59 | + text=chunk["text"], |
| 60 | + embedding=chunk["embedding"], |
| 61 | + ) |
| 62 | + |
| 63 | + doc_embeddings.export( |
| 64 | + "doc_embeddings", |
| 65 | + cocoindex.targets.Postgres(), |
| 66 | + primary_key_fields=["filename", "location"], |
| 67 | + vector_indexes=[ |
| 68 | + cocoindex.VectorIndexDef( |
| 69 | + field_name="embedding", |
| 70 | + metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY, |
| 71 | + ) |
| 72 | + ], |
| 73 | + ) |
| 74 | + |
| 75 | + |
| 76 | +def search(pool: ConnectionPool, query: str, top_k: int = 5) -> list[dict[str, Any]]: |
| 77 | + # Get the table name, for the export target in the azure_blob_text_embedding_flow above. |
| 78 | + table_name = cocoindex.utils.get_target_default_name( |
| 79 | + azure_blob_text_embedding_flow, "doc_embeddings" |
| 80 | + ) |
| 81 | + # Evaluate the transform flow defined above with the input query, to get the embedding. |
| 82 | + query_vector = text_to_embedding.eval(query) |
| 83 | + # Run the query and get the results. |
| 84 | + with pool.connection() as conn: |
| 85 | + with conn.cursor() as cur: |
| 86 | + cur.execute( |
| 87 | + f""" |
| 88 | + SELECT filename, text, embedding <=> %s::vector AS distance |
| 89 | + FROM {table_name} ORDER BY distance LIMIT %s |
| 90 | + """, |
| 91 | + (query_vector, top_k), |
| 92 | + ) |
| 93 | + return [ |
| 94 | + {"filename": row[0], "text": row[1], "score": 1.0 - row[2]} |
| 95 | + for row in cur.fetchall() |
| 96 | + ] |
| 97 | + |
| 98 | + |
| 99 | +def _main() -> None: |
| 100 | + # Initialize the database connection pool. |
| 101 | + pool = ConnectionPool(os.getenv("COCOINDEX_DATABASE_URL")) |
| 102 | + |
| 103 | + azure_blob_text_embedding_flow.setup() |
| 104 | + update_stats = azure_blob_text_embedding_flow.update() |
| 105 | + print(update_stats) |
| 106 | + |
| 107 | + # Run queries in a loop to demonstrate the query capabilities. |
| 108 | + while True: |
| 109 | + query = input("Enter search query (or Enter to quit): ") |
| 110 | + if query == "": |
| 111 | + break |
| 112 | + # Run the query function with the database connection pool and the query. |
| 113 | + results = search(pool, query) |
| 114 | + print("\nSearch results:") |
| 115 | + for result in results: |
| 116 | + print(f"[{result['score']:.3f}] {result['filename']}") |
| 117 | + print(f" {result['text']}") |
| 118 | + print("---") |
| 119 | + print() |
| 120 | + |
| 121 | + |
| 122 | +if __name__ == "__main__": |
| 123 | + load_dotenv() |
| 124 | + cocoindex.init() |
| 125 | + _main() |
0 commit comments