[FEATURE]: Add support for embedding models in OpenRouter

[prrao87]

Description of the changes

Adds support for OpenRouter embedding models

• Updates the LlmApiType::OpenRouter module to accept the same parameters as Gemini/OpenAI models (task_type and output_dimension to allow embedding generation)
• Update the documentation page to show that OpenRouter now supports embedding models, and include an example

No behavior changes for existing OpenRouter functionality; I only extended the existing embedding factory to route OpenRouter through the same embedding client used for OpenAI.

Motivation and context

OpenRouter now supports a few embedding models (type in "embedding" into the search box here. Users who are using OpenRouter LLMs for generative tasks can also use embedding models to generate embeddings in a CocoIndex flow, by selecting the appropriate model from the list of supported models in OpenRouter.

Note if it's a breaking change

• Yes
• No

How I tested

I built the library from source and ran the following code to generate embeddings in LanceDB (note the update to the text_to_embedding flow transform, where OpenRouter is called instead of sentence-transformers). The flow runs through to completion and the embeddings are persisted to LanceDB 😄.

import os
import datetime
import cocoindex
import math
import cocoindex.targets.lancedb as coco_lancedb

# Define LanceDB connection constants
LANCEDB_URI = "./lancedb_data"  # Local directory for LanceDB
LANCEDB_TABLE = "TextEmbedding"


@cocoindex.transform_flow()
def text_to_embedding(
    text: cocoindex.DataSlice[str],
) -> cocoindex.DataSlice[list[float]]:
    """
    Embed the text using an OpenRouter embedding model.
    This is a shared logic between indexing and querying, so extract it as a function.
    """
    return text.transform(
        cocoindex.functions.EmbedText(
            api_type=cocoindex.LlmApiType.OPEN_ROUTER,
            model="openai/text-embedding-3-small",
            # Optional, use the default task type if not specified
            task_type="SEMANTICS_SIMILARITY",
            output_dimension=1536,
        )
    )


@cocoindex.flow_def(name="TextEmbeddingWithLanceDB")
def text_embedding_flow(
    flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope
) -> None:
    """
    Define an example flow that embeds text into a vector database.
    """
    ENABLE_LANCEDB_VECTOR_INDEX = os.environ.get(
        "ENABLE_LANCEDB_VECTOR_INDEX", "0"
    ).lower() in ("true", "1")

    data_scope["documents"] = flow_builder.add_source(
        cocoindex.sources.LocalFile(path="markdown_files"),
        refresh_interval=datetime.timedelta(seconds=5),
    )

    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=500,
            chunk_overlap=100,
        )

        with doc["chunks"].row() as chunk:
            chunk["embedding"] = text_to_embedding(chunk["text"])
            doc_embeddings.collect(
                id=cocoindex.GeneratedField.UUID,
                filename=doc["filename"],
                location=chunk["location"],
                text=chunk["text"],
                # 'text_embedding' is the name of the vector we've created the LanceDB table with.
                text_embedding=chunk["embedding"],
            )

    # We cannot enable index when the table has no data yet, as LanceDB requires data to train the index.
    # See: https://github.com/lancedb/lance/issues/4034
    # Guard it with ENABLE_LANCEDB_VECTOR_INDEX environment variable.
    vector_indexes = []
    if ENABLE_LANCEDB_VECTOR_INDEX:
        vector_indexes.append(
            cocoindex.VectorIndexDef(
                "text_embedding", cocoindex.VectorSimilarityMetric.L2_DISTANCE
            )
        )
    doc_embeddings.export(
        "doc_embeddings",
        coco_lancedb.LanceDB(db_uri=LANCEDB_URI, table_name=LANCEDB_TABLE),
        primary_key_fields=["id"],
        vector_indexes=vector_indexes,
        fts_indexes=[
            cocoindex.FtsIndexDef(
                field_name="text", parameters={"tokenizer_name": "simple"}
            )
        ],
    )


@text_embedding_flow.query_handler(
    result_fields=cocoindex.QueryHandlerResultFields(
        embedding=["embedding"],
        score="score",
    ),
)
async def search(query: str) -> cocoindex.QueryOutput:
    print("Searching...", query)
    db = await coco_lancedb.connect_async(LANCEDB_URI)
    table = await db.open_table(LANCEDB_TABLE)

    # Get the embedding for the query
    query_embedding = await text_to_embedding.eval_async(query)

    search = await table.search(query_embedding, vector_column_name="text_embedding")
    search_results = await search.limit(5).to_list()

    return cocoindex.QueryOutput(
        results=[
            {
                "filename": result["filename"],
                "text": result["text"],
                "embedding": result["text_embedding"],
                # Qdrant's L2 "distance" is squared, so we take the square root to align with normal L2 distance
                "score": math.sqrt(result["_distance"]),
            }
            for result in search_results
        ],
        query_info=cocoindex.QueryInfo(
            embedding=query_embedding,
            similarity_metric=cocoindex.VectorSimilarityMetric.L2_DISTANCE,
        ),
    )

[prrao87]

@cocoindex-dev @georgeh0 hi! I added support for OpenRouter embedding models in this PR. I notice that the tests for OpenAI embedding models use the #[ignore] flag, so no new tests were added or modified in this PR. Hope that makes sense, but let me know if you have any feedback/changes required.

[georgeh0]

Thanks a lot for making this!