[WIP] Support for Capella-hosted embedding models (#68)

.github/workflows/tests.yaml

@@ -33,6 +33,9 @@ jobs:
         run: |
           poetry install --with dev
 
+      - name: Export OPENAI API KEY secret
+        run: echo "AGENT_CATALOG_EMBEDDING_MODEL_AUTH=${{ secrets.AGENT_CATALOG_EMBEDDING_MODEL_AUTH }}" >> $GITHUB_ENV
+
       - run: echo "🔥 Starting the core tests (🤞 please don't break!)."
       - name: Verify Pytest installation
         run: poetry show pytest

docs/source/env.rst

@@ -23,7 +23,7 @@ Mandatory Environment Variables
     The username of the account/database access key used to access your Couchbase cluster.
 
 ``AGENT_CATALOG_PASSWORD``
-    The password of the account/database access key used to access your Couchbase cluster
+    The password of the account/database access key used to access your Couchbase cluster.
 
 ``AGENT_CATALOG_BUCKET``
     The name of the bucket where your catalog and all audit logs are/will be stored.
@@ -34,9 +34,8 @@ Optional Environment Variables
 
 ``AGENT_CATALOG_CONN_ROOT_CERTIFICATE``
     Path to the `TLS <https://en.wikipedia.org/wiki/Transport_Layer_Security>`_ Root Certificate associated with your
-    Couchbase cluster for secure connection establishment.
-
-    Instructions for Couchbase Server certificates can be found `here <https://docs.couchbase.com/server/current/learn/security/certificates.html>`_.
+    Couchbase cluster.
+    More information about Couchbase Server certificates can be found `here <https://docs.couchbase.com/server/current/learn/security/certificates.html>`_.
 
 ``AGENT_CATALOG_ACTIVITY``
     The location on your filesystem that denotes where the local audit logs are stored.
@@ -68,17 +67,39 @@ Optional Environment Variables
     The location + filename of the audit logs that the :python:`agentc.Auditor` will write to.
     By default, the :python:`agentc.Auditor` class will write and rotate logs in the :file:`./agent-activity` directory.
 
-``AGENT_CATALOG_EMBEDDING_MODEL``
+``AGENT_CATALOG_EMBEDDING_MODEL_NAME``
     The embedding model that Agent Catalog will use when indexing and querying tools and prompts.
     This *must* be a valid embedding model that is supported by the :python:`sentence_transformers.SentenceTransformer`
-    class.
+    class *or* the name of a model that can be used from the endpoint specified in the environment variable
+    ``AGENT_CATALOG_EMBEDDING_MODEL_URL``.
     By default, the ``sentence-transformers/all-MiniLM-L12-v2`` model is used.
 
+``AGENT_CATALOG_EMBEDDING_MODEL_URL``
+    An OpenAI-standard client base URL whose ``/embeddings`` endpoint will be used to generate embeddings for Agent
+    Catalog tools and prompts.
+    The specified endpoint *must* host the embedding model given in ``AGENT_CATALOG_EMBEDDING_MODEL_NAME``.
+    If this variable is specified, Agent Catalog will assume the model given in ``AGENT_CATALOG_EMBEDDING_MODEL_NAME``
+    should be accessed through an OpenAI-standard interface.
+    This variable *must* be specified with ``AGENT_CATALOG_EMBEDDING_MODEL_AUTH``.
+    By default, this variable is not set (thus, a locally hosted SentenceTransformers is used).
+
+``AGENT_CATALOG_EMBEDDING_MODEL_AUTH``
+    The field used in the authorization header of all OpenAI-standard client embedding requests.
+    For embedding models hosted by OpenAI, this field refers to the API key.
+    For embedding models hosted by Capella, this field refers to the Base64-encoded value of
+    ``MY_USERNAME.MY_PASSWORD``.
+    If this variable is specified, Agent Catalog will assume the model given in ``AGENT_CATALOG_EMBEDDING_MODEL_NAME``
+    should be accessed through an OpenAI-standard interface.
+    This variable *must* be specified with ``AGENT_CATALOG_EMBEDDING_MODEL_URL``.
+    By default, this variable is not set (thus, a locally hosted SentenceTransformers is used).
+
 ``AGENT_CATALOG_INDEX_PARTITION``
-    Required for advanced vector index definition. This is an integer that defines the number of index partitions on your node.
-    If not set, this value is ``2 * number of nodes with 'search' service`` on your cluster.
+    The number of index partitions associated with your cluster.
+    This variable is used during the creation of vector indexes for semantic catalog search.
+    By default, this value is set to ``2 * number of nodes with 'search' service on your cluster``.
     More information on index partitioning can be found `here <https://docs.couchbase.com/server/current/n1ql/n1ql-language-reference/index-partitioning.html>`_.
 
 ``AGENT_CATALOG_MAX_SOURCE_PARTITION``
-    Required for advanced vector index definition. This is an integer that defines the maximum number of source partitions.
-    If not set, this value is 1024.
+    The maximum number of source partitions associated with your cluster.
+    This variable is used during the creation of vector indexes for semantic catalog search.
+    By default, this value is set to 1024.

libs/agentc_cli/agentc_cli/cmds/env.py

@@ -26,7 +26,9 @@ def cmd_env(ctx: Context = None):
         "AGENT_CATALOG_SNAPSHOT": os.getenv("AGENT_CATALOG_SNAPSHOT", LATEST_SNAPSHOT_VERSION),
         "AGENT_CATALOG_PROVIDER_OUTPUT": os.getenv("AGENT_CATALOG_PROVIDER_OUTPUT", None),
         "AGENT_CATALOG_AUDITOR_OUTPUT": os.getenv("AGENT_CATALOG_AUDITOR_OUTPUT", None),
-        "AGENT_CATALOG_EMBEDDING_MODEL": os.getenv("AGENT_CATALOG_EMBEDDING_MODEL", DEFAULT_EMBEDDING_MODEL),
+        "AGENT_CATALOG_EMBEDDING_MODEL_NAME": os.getenv("AGENT_CATALOG_EMBEDDING_MODEL_NAME", DEFAULT_EMBEDDING_MODEL),
+        "AGENT_CATALOG_EMBEDDING_MODEL_URL": os.getenv("AGENT_CATALOG_EMBEDDING_MODEL_URL", None),
+        "AGENT_CATALOG_EMBEDDING_MODEL_AUTH": os.getenv("AGENT_CATALOG_EMBEDDING_MODEL_AUTH", None),
     }
     for line in json.dumps(environment_dict, indent=4).split("\n"):
         if re.match(r'\s*"AGENT_CATALOG_.*": (?!null)', line):

libs/agentc_cli/agentc_cli/cmds/index.py

@@ -28,6 +28,7 @@ def cmd_index(
     source_dirs: list[str | os.PathLike],
     kinds: list[typing.Literal["tool", "prompt"]],
     embedding_model_name: str = DEFAULT_EMBEDDING_MODEL,
+    embedding_model_url: str = None,
     dry_run: bool = False,
     ctx: Context = None,
     **_,
@@ -57,6 +58,7 @@ def cmd_index(
     repo, get_path_version = load_repository(pathlib.Path(os.getcwd()))
     embedding_model = EmbeddingModel(
         embedding_model_name=embedding_model_name,
+        embedding_model_url=embedding_model_url,
         catalog_path=pathlib.Path(ctx.catalog),
     )
 

libs/agentc_cli/agentc_cli/main.py

@@ -75,7 +75,8 @@ def resolve_command(self, ctx, args):
 
 @click.group(
     cls=AliasedGroup,
-    epilog="See: https://docs.couchbase.com or https://couchbaselabs.github.io/agent-catalog/index.html# for more information.",
+    epilog="See: https://docs.couchbase.com or "
+    "https://couchbaselabs.github.io/agent-catalog/index.html# for more information.",
     context_settings=dict(max_content_width=800),
 )
 @click.option(
@@ -527,9 +528,15 @@ def find(
 )
 @click.option(
     "-em",
-    "--embedding-model",
+    "--embedding-model-name",
     default=DEFAULT_EMBEDDING_MODEL,
-    help="Embedding model used when indexing source files into the local catalog.",
+    help="Name of the embedding model used when indexing source files into the local catalog.",
+    show_default=True,
+)
+@click.option(
+    "--embedding-model-url",
+    default=None,
+    help="Base URL of an OpenAI-standard endpoint that exposes an embedding model.",
     show_default=True,
 )
 @click.option(
@@ -540,7 +547,7 @@ def find(
     show_default=True,
 )
 @click.pass_context
-def index(ctx, source_dirs, tools, prompts, embedding_model, dry_run):
+def index(ctx, source_dirs, tools, prompts, embedding_model_name, embedding_model_url, dry_run):
     """Walk the source directory trees (SOURCE_DIRS) to index source files into the local catalog.
     Source files that will be scanned include *.py, *.sqlpp, *.yaml, etc."""
 
@@ -565,7 +572,8 @@ def index(ctx, source_dirs, tools, prompts, embedding_model, dry_run):
         ctx=ctx.obj,
         source_dirs=source_dirs,
         kinds=kinds,
-        embedding_model_name=embedding_model,
+        embedding_model_name=embedding_model_name,
+        embedding_model_url=embedding_model_url,
         dry_run=dry_run,
     )
 

libs/agentc_core/agentc_core/catalog/descriptor.py

@@ -12,6 +12,7 @@
 from ..tool.descriptor.models import SemanticSearchToolDescriptor
 from ..tool.descriptor.models import SQLPPQueryToolDescriptor
 from ..version import VersionDescriptor
+from agentc_core.learned.model import EmbeddingModel
 
 
 class CatalogKind(enum.StrEnum):
@@ -47,10 +48,8 @@ class CatalogDescriptor(pydantic.BaseModel):
 
     kind: CatalogKind = pydantic.Field(description="The type of items within the catalog.")
 
-    embedding_model: str = pydantic.Field(
-        description="The sentence-transformers embedding model used to generate the vector representations "
-        "of each catalog entry.",
-        examples=["sentence-transformers/all-MiniLM-L12-v2"],
+    embedding_model: EmbeddingModel = pydantic.Field(
+        description="Embedding model used to generate embedding for tool/prompt descriptions to store in the catalogs.",
     )
 
     version: VersionDescriptor = pydantic.Field(

libs/agentc_core/agentc_core/catalog/index.py

@@ -9,6 +9,7 @@
 from ..indexer import AllIndexers
 from ..indexer import vectorize_descriptor
 from ..learned.embedding import EmbeddingModel
+from ..learned.model import EmbeddingModel as CatalogDescriptorEmbeddingModel
 from ..record.descriptor import RecordDescriptor
 from .catalog.mem import CatalogMem
 from .descriptor import CatalogDescriptor
@@ -167,13 +168,20 @@ def index_catalog_start(
         logger.error("Encountered error(s) while crawling source directories: " + "\n".join([str(e) for e in all_errs]))
         raise all_errs[0]
 
+    catalog_descriptor_embedding_model = (
+        CatalogDescriptorEmbeddingModel(name=embedding_model.name, base_url=None, kind="sentence-transformers")
+        if embedding_model.embedding_model_url is None
+        else CatalogDescriptorEmbeddingModel(
+            name=embedding_model.name, base_url=embedding_model.embedding_model_url, kind="openai"
+        )
+    )
     next_catalog = CatalogMem(
         embedding_model=embedding_model,
         catalog_descriptor=CatalogDescriptor(
             schema_version=meta_version.schema_version,
             library_version=meta_version.library_version,
             version=catalog_version,
-            embedding_model=embedding_model.name,
+            embedding_model=catalog_descriptor_embedding_model,
             kind=kind,
             source_dirs=source_dirs,
             items=all_descriptors,