@cocoindex.flow_def(name="TextEmbedding")
def text_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope):
    """Định nghĩa luồng xử lý embed text vào vector DB"""
    data_scope["documents"] = flow_builder.add_source(cocoindex.sources.LocalFile(path="markdown_files"))
    doc_embeddings = data_scope.add_collector()

with data_scope["documents"].row() as doc:
    doc["chunks"] = doc["content"].transform(
        cocoindex.functions.SplitRecursively(), 
        language="markdown", 
        chunk_size=2000, 
        chunk_overlap=500
    )

@cocoindex.transform_flow()
def text_to_embedding(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[list[float]]:
    """Chuyển văn bản thành vector embedding với SentenceTransformer"""
    return text.transform(
        cocoindex.functions.SentenceTransformerEmbed(
            model="sentence-transformers/all-MiniLM-L6-v2"
        )
    )

with doc["chunks"].row() as chunk:
    chunk["embedding"] = text_to_embedding(chunk["text"])
    doc_embeddings.collect(
        filename=doc["filename"],
        location=chunk["location"],
        text=chunk["text"],
        embedding=chunk["embedding"]
    )

doc_embeddings.export(
    "doc_embeddings", 
    cocoindex.storages.Postgres(),
    primary_key_fields=["filename", "location"],
    vector_indexes=[
        cocoindex.VectorIndexDef(
            field_name="embedding",
            metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY
        )
    ]
)

def search(pool: ConnectionPool, query: str, top_k: int = 5):
    table_name = cocoindex.utils.get_target_storage_default_name(text_embedding_flow, "doc_embeddings")
    query_vector = text_to_embedding.eval(query)
    with pool.connection() as conn:
        with conn.cursor() as cur:
            cur.execute(f"""
                SELECT filename, text, embedding <=> %s::vector AS distance
                FROM {table_name}
                ORDER BY distance LIMIT %s
            """, (query_vector, top_k))
            return [
                {"filename": row[0], "text": row[1], "score": 1.0 - row[2]} 
                for row in cur.fetchall()
            ]

def _main():
    pool = ConnectionPool(os.getenv("COCOINDEX_DATABASE_URL"))

    while True:
        query = input("Enter search query (or Enter to quit): ")
        if query == '':
            break
        results = search(pool, query)
        print("\nSearch results:")
        for result in results:
            print(f"[{result['score']:.3f}] {result['filename']}")
            print(f"  {result['text']}")
            print("---")
        print()

if __name__ == "__main__":
    load_dotenv()
    cocoindex.init()
    _main()

python main.py cocoindex setup
python main.py cocoindex update

python main.py