@cocoindex.flow_def(name="CodeEmbedding")
def code_embedding_flow(flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope):
    data_scope["files"] = flow_builder.add_source(
        cocoindex.sources.LocalFile(
            path="../..",
            included_patterns=["*.py", "*.rs", "*.toml", "*.md", "*.mdx"],
            excluded_patterns=[".*", "target", "**/node_modules"]
        )
    )
    code_embeddings = data_scope.add_collector()

@cocoindex.op.function()
def extract_extension(filename: str) -> str:
    return os.path.splitext(filename)[1]

with data_scope["files"].row() as file:
    file["extension"] = file["filename"].transform(extract_extension)

with data_scope["files"].row() as file:
    file["chunks"] = file["content"].transform(
        cocoindex.functions.SplitRecursively(),
        language=file["extension"],
        chunk_size=1000,
        chunk_overlap=300
    )

@cocoindex.transform_flow()
def code_to_embedding(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[list[float]]:
    return text.transform(
        cocoindex.functions.SentenceTransformerEmbed(model="sentence-transformers/all-MiniLM-L6-v2")
    )

with data_scope["files"].row() as file:
    with file["chunks"].row() as chunk:
        chunk["embedding"] = chunk["text"].call(code_to_embedding)
    code_embeddings.collect(
        filename=file["filename"],
        location=chunk["location"],
        code=chunk["text"],
        embedding=chunk["embedding"]
    )

code_embeddings.export(
    "code_embeddings",
    cocoindex.storages.Postgres(),
    primary_key_fields=["filename", "location"],
    vector_indexes=[
        cocoindex.VectorIndex("embedding", cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY)
    ]
)

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

@cocoindex.main_fn()
def _run():
    pool = ConnectionPool(os.getenv("COCOINDEX_DATABASE_URL"))
    while True:
        try:
            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['code']}")
                print("---\n")
        except KeyboardInterrupt:
            break

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

python main.py