dex-connector-cohere 0.1.0__tar.gz

dex_connector_cohere-0.1.0/.gitignore

@@ -0,0 +1,6 @@
+/.dex-local/dexcli
+/.dex-local/dex-server-zd
+/.dex-local/data/
+/.dex-local/cli.log
+/.dex-local/visualize.json
+/.dex-local/visualize.err

dex_connector_cohere-0.1.0/PKG-INFO

@@ -0,0 +1,51 @@
+Metadata-Version: 2.5
+Name: dex-connector-cohere
+Version: 0.1.0
+Summary: Cohere connector for Dex Flows
+License-Expression: MIT
+Requires-Python: >=3.11
+Requires-Dist: dex-connectors-sdk==0.1.0
+Requires-Dist: dex-python-sdk>=1.28.0
+Requires-Dist: httpx>=0.27
+Description-Content-Type: text/markdown
+
+# Cohere connector for Python
+
+Install with `pip install dex-connector-cohere==0.1.0`. The package requires
+`dex-connectors-sdk==0.1.0` and `dex-python-sdk>=1.28.0`, and implements the
+manifest Queries `embed` and `rerank` with the same branches, defaults, and
+failure classification as the Go module.
+
+## Operations
+
+| Operation | Factory | Happy branch | Optional branches |
+| --- | --- | --- | --- |
+| `embed` | `new_embed_step` | `embedded` | `providerRejected`, `invalidResponse`, `defect` |
+| `rerank` | `new_rerank_step` | `ranked` | `providerRejected`, `invalidResponse`, `defect` |
+
+`embed` takes a model, 1 to 96 nonblank texts, and an input type
+(`search_document`, `search_query`, `classification`, or `clustering`) and
+returns one float vector per text. `rerank` takes a model, a query, 1 to 1000
+nonblank documents, and an optional `top_n`, and returns each document's index
+and relevance score in ranked order. Both calls may be billed, so the manifest
+allows three attempts within five minutes. 408, 409, 429, 5xx, and transport
+failures retry; 401 and 498 select `providerRejected` with an authentication
+kind, and 403 with an authorization kind.
+
+## Authorization
+
+Create a Dex Web connection named `cohere` as described in the
+[connector README](../README.md#connection). The API key travels only in the
+`Authorization: Bearer` header and is resolved for every attempt.
+
+## Example
+
+[`examples/cohere_search`](examples/cohere_search/flow.py) embeds a query and
+then reranks candidate documents. Run the Worker with
+`uv run python -m examples.cohere_search.worker` after `dexcli dev` has loaded
+a project containing the `cohere` connection.
+
+## Tests
+
+`coheremock.Mock` scripts both Queries and defaults to the manifest answers.
+Run `uv run pytest tests` for provider parity and manifest contract checks.

dex_connector_cohere-0.1.0/README.md

@@ -0,0 +1,40 @@
+# Cohere connector for Python
+
+Install with `pip install dex-connector-cohere==0.1.0`. The package requires
+`dex-connectors-sdk==0.1.0` and `dex-python-sdk>=1.28.0`, and implements the
+manifest Queries `embed` and `rerank` with the same branches, defaults, and
+failure classification as the Go module.
+
+## Operations
+
+| Operation | Factory | Happy branch | Optional branches |
+| --- | --- | --- | --- |
+| `embed` | `new_embed_step` | `embedded` | `providerRejected`, `invalidResponse`, `defect` |
+| `rerank` | `new_rerank_step` | `ranked` | `providerRejected`, `invalidResponse`, `defect` |
+
+`embed` takes a model, 1 to 96 nonblank texts, and an input type
+(`search_document`, `search_query`, `classification`, or `clustering`) and
+returns one float vector per text. `rerank` takes a model, a query, 1 to 1000
+nonblank documents, and an optional `top_n`, and returns each document's index
+and relevance score in ranked order. Both calls may be billed, so the manifest
+allows three attempts within five minutes. 408, 409, 429, 5xx, and transport
+failures retry; 401 and 498 select `providerRejected` with an authentication
+kind, and 403 with an authorization kind.
+
+## Authorization
+
+Create a Dex Web connection named `cohere` as described in the
+[connector README](../README.md#connection). The API key travels only in the
+`Authorization: Bearer` header and is resolved for every attempt.
+
+## Example
+
+[`examples/cohere_search`](examples/cohere_search/flow.py) embeds a query and
+then reranks candidate documents. Run the Worker with
+`uv run python -m examples.cohere_search.worker` after `dexcli dev` has loaded
+a project containing the `cohere` connection.
+
+## Tests
+
+`coheremock.Mock` scripts both Queries and defaults to the manifest answers.
+Run `uv run pytest tests` for provider parity and manifest contract checks.

dex_connector_cohere-0.1.0/examples/cohere_search/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) 2026 Super Durable
+# SPDX-License-Identifier: MIT
+"""Cohere search example Flow."""

dex_connector_cohere-0.1.0/examples/cohere_search/flow.py

@@ -0,0 +1,129 @@
+# Copyright (c) 2026 Super Durable
+# SPDX-License-Identifier: MIT
+"""Runnable Cohere search Flow: embed the query, then rerank candidate documents."""
+
+from dataclasses import dataclass
+
+import dex
+from dex_connectors.sdk import QueryResult
+
+from dex_connectors.cohere import (
+    Connection,
+    EmbedInput,
+    EmbedOutput,
+    RankedDocument,
+    RerankInput,
+    RerankOutput,
+    new_embed_step,
+    new_rerank_step,
+)
+
+FLOW_TYPE = "CohereSearch"
+CONNECTION_NAME = "cohere"
+_EMBED_STEP_TYPE = "EmbedSearchQuery"
+_RERANK_STEP_TYPE = "RerankSearchDocuments"
+_ANNOTATION_GROUP = {"groupId": "cohere-search", "groupLabel": "Cohere search"}
+
+
+@dataclass(frozen=True, slots=True)
+class SearchRequest:
+    """Search query and at least two candidate documents."""
+
+    query: str
+    documents: tuple[str, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class SearchOutcome:
+    """Candidate documents in Cohere's ranked order."""
+
+    results: tuple[RankedDocument, ...] = ()
+
+
+_REQUEST = dex.Attribute("cohere-search-request", SearchRequest)
+
+
+def map_to_embed_input(request: SearchRequest) -> EmbedInput:
+    """Embed the query as a search query."""
+    return EmbedInput(model="embed-v4.0", texts=(request.query,), input_type="search_query")
+
+
+def map_to_rerank_input(request: SearchRequest) -> RerankInput:
+    """Rerank every candidate document against the query."""
+    return RerankInput(model="rerank-v5.0-pro", query=request.query, documents=request.documents)
+
+
+# dex:group group-id:cohere-search group-label:"Cohere search"
+# dex:explanation text:"Check and store the search request before calling Cohere."
+class _StartRequest(dex.Step[SearchRequest]):
+    def get_step_type(self) -> str:
+        return "StartCohereSearch"
+
+    def execute(self, context: dex.Context, input: SearchRequest) -> dex.StepDecision:
+        if not input.query.strip() or len(input.documents) < 2:
+            return dex.force_fail("query and at least two documents are required")
+        _REQUEST.set(context, input)
+        return dex.go_to(dex.StepRef(_EMBED_STEP_TYPE), input)
+
+
+# dex:group group-id:cohere-search group-label:"Cohere search"
+# dex:explanation text:"Continue from the query embedding to the rerank call."
+class _AfterEmbedding(dex.Step[QueryResult[EmbedOutput]]):
+    def get_step_type(self) -> str:
+        return "AfterCohereEmbedding"
+
+    def execute(self, context: dex.Context, input: QueryResult[EmbedOutput]) -> dex.StepDecision:
+        return dex.go_to(dex.StepRef(_RERANK_STEP_TYPE), _REQUEST.get(context))
+
+
+# dex:group group-id:cohere-search group-label:"Cohere search"
+# dex:explanation text:"Complete with the ranked documents."
+class _CompleteSearch(dex.Step[QueryResult[RerankOutput]]):
+    def get_step_type(self) -> str:
+        return "CompleteCohereSearch"
+
+    def execute(self, context: dex.Context, input: QueryResult[RerankOutput]) -> dex.StepDecision:
+        return dex.graceful_complete(SearchOutcome(input.value.results))
+
+
+class SearchFlow(dex.Flow[SearchRequest]):
+    """Embed a query and rerank documents through the static cohere connection."""
+
+    def __init__(self, connection: Connection) -> None:
+        if connection.reference.name != CONNECTION_NAME:
+            raise ValueError("Cohere search requires the static cohere connection")
+        self._start = _StartRequest()
+        self._after_embedding = _AfterEmbedding()
+        self._complete = _CompleteSearch()
+        self._embed = new_embed_step(
+            step_type=_EMBED_STEP_TYPE,
+            step_input_type=SearchRequest,
+            annotations={**_ANNOTATION_GROUP, "explanation": "Embed the search query."},
+            connection=connection,
+            connection_name=CONNECTION_NAME,
+            map_to_operation_input=map_to_embed_input,
+            embedded=self._after_embedding,
+        )
+        self._rerank = new_rerank_step(
+            step_type=_RERANK_STEP_TYPE,
+            step_input_type=SearchRequest,
+            annotations={**_ANNOTATION_GROUP, "explanation": "Rank the candidate documents."},
+            connection=connection,
+            connection_name=CONNECTION_NAME,
+            map_to_operation_input=map_to_rerank_input,
+            ranked=self._complete,
+        )
+
+    def get_flow_type(self) -> str:
+        """Return the stable Flow identity used in Dex Web."""
+        return FLOW_TYPE
+
+    def get_steps(self) -> dex.StepList[SearchRequest]:
+        """Register the start, both Cohere calls, and their result Steps."""
+        return dex.StepList.start_step(self._start).other_steps(
+            self._embed, self._after_embedding, self._rerank, self._complete
+        )
+
+    def get_persistence_schema(self) -> dex.PersistenceSchema:
+        """Persist the search request between the two calls."""
+        return dex.PersistenceSchema.of(_REQUEST)

dex_connector_cohere-0.1.0/examples/cohere_search/worker.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2026 Super Durable
+# SPDX-License-Identifier: MIT
+"""Worker entrypoint for the Cohere search Flow."""
+
+import asyncio
+import os
+import tempfile
+from pathlib import Path
+
+import dex
+from dex_connectors.sdk.projectconfig import load_from_environment
+
+from dex_connectors.cohere import new_project_connection
+
+from .flow import CONNECTION_NAME, SearchFlow
+
+
+async def run_worker() -> None:
+    """Load the named Dex Web connection and serve searches."""
+    project = await load_from_environment()
+    connection = None
+    cache = None
+    try:
+        connection = new_project_connection(project, CONNECTION_NAME)
+        registry = dex.Registry([SearchFlow(connection)], allow_async_handlers=True)
+        cache_directory = os.environ.get(
+            "DEX_BLOB_CACHE_DIR", str(Path(tempfile.gettempdir()) / "dex-cohere-search-blobs")
+        )
+        cache = dex.open_blob_cache(dex.BlobCacheConfig(cache_directory, 1 << 30))
+        worker = dex.AsyncWorker(
+            registry,
+            cache,
+            dex.WorkerOptions(server_address=os.environ.get("DEX_FLOW_SERVICE_ADDRESS", "localhost:8801")),
+        )
+        async with worker:
+            await worker.start()
+    finally:
+        if connection is not None:
+            await connection.aclose()
+        if cache is not None:
+            cache.close()
+        await project.close()
+
+
+def main() -> None:
+    """Run until the Worker stops or the process receives an interrupt."""
+    asyncio.run(run_worker())
+
+
+if __name__ == "__main__":
+    main()

dex_connector_cohere-0.1.0/pyproject.toml

@@ -0,0 +1,37 @@
+[build-system]
+requires = ["hatchling>=1.27"]
+build-backend = "hatchling.build"
+
+[project]
+name = "dex-connector-cohere"
+version = "0.1.0"
+description = "Cohere connector for Dex Flows"
+readme = "README.md"
+requires-python = ">=3.11"
+license = "MIT"
+dependencies = ["dex-connectors-sdk==0.1.0", "dex-python-sdk>=1.28.0", "httpx>=0.27"]
+
+[dependency-groups]
+dev = ["pytest>=8", "pytest-asyncio>=0.24", "mypy>=1.13", "ruff>=0.7", "pyyaml>=6", "types-PyYAML>=6"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/dex_connectors"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = [".", "src"]
+asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "function"
+
+[tool.mypy]
+python_version = "3.11"
+strict = true
+mypy_path = "src"
+explicit_package_bases = true
+
+[tool.ruff]
+target-version = "py311"
+line-length = 110
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP", "B"]

dex_connector_cohere-0.1.0/src/dex_connectors/cohere/__init__.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2026 Super Durable
+# SPDX-License-Identifier: MIT
+"""Cohere embedding and rerank Queries for Dex connector Steps."""
+
+from .client import (
+    CONNECTOR_ID,
+    DEFECT,
+    EMBED_DEFINITION,
+    EMBEDDED,
+    INPUT_TYPES,
+    INVALID_RESPONSE,
+    PROVIDER_REJECTED,
+    RANKED,
+    RERANK_DEFINITION,
+    Client,
+    Config,
+    Credentials,
+    EmbedOperation,
+    RerankOperation,
+)
+from .connection import Connection, new_connection, new_connection_with_operations, new_project_connection
+from .models import EmbedInput, EmbedOutput, RankedDocument, RerankInput, RerankOutput
+from .steps import new_embed_step, new_rerank_step
+
+__all__ = [
+    "Client",
+    "Config",
+    "Connection",
+    "Credentials",
+    "EmbedInput",
+    "EmbedOperation",
+    "EmbedOutput",
+    "RankedDocument",
+    "RerankInput",
+    "RerankOperation",
+    "RerankOutput",
+    "new_connection",
+    "new_connection_with_operations",
+    "new_embed_step",
+    "new_project_connection",
+    "new_rerank_step",
+    "CONNECTOR_ID",
+    "DEFECT",
+    "EMBED_DEFINITION",
+    "EMBEDDED",
+    "INPUT_TYPES",
+    "INVALID_RESPONSE",
+    "PROVIDER_REJECTED",
+    "RANKED",
+    "RERANK_DEFINITION",
+]

dex_connector_cohere-0.1.0/src/dex_connectors/cohere/_generated/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) 2026 Super Durable
+# SPDX-License-Identifier: MIT

dex_connector_cohere-0.1.0/src/dex_connectors/cohere/_generated/config.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2026 Super Durable
+# SPDX-License-Identifier: MIT
+# ruff: noqa
+
+from dataclasses import dataclass
+
+from dex_connectors.sdk import SecretString
+
+
+@dataclass(frozen=True, slots=True)
+class Config:
+
+    base_url: str = "https://api.cohere.com"
+    timeout_seconds: float = 30
+    max_response_bytes: int = 4194304
+
+setattr(Config, '__module__', 'dex_connectors.cohere.client')
+
+
+@dataclass(frozen=True, slots=True)
+class Credentials:
+
+    api_key: SecretString
+
+setattr(Credentials, '__module__', 'dex_connectors.cohere.client')

dex_connector_cohere-0.1.0/src/dex_connectors/cohere/_generated/connection.py

@@ -0,0 +1,131 @@
+# Copyright (c) 2026 Super Durable
+# SPDX-License-Identifier: MIT
+# ruff: noqa
+
+from typing import Any
+
+import httpx
+from dex_connectors.sdk import ConnectionRef, Query, Call, SecretString
+from dex_connectors.sdk.projectconfig import ConnectionKey, LoadedProject
+
+from ..client import Client, Config, Credentials
+from .definitions import CONNECTOR_ID, EMBED_DEFINITION, RERANK_DEFINITION
+from ..models import EmbedInput, EmbedOutput, RerankInput, RerankOutput
+
+class Connection:
+
+    __slots__ = ('_embed', '_rerank', '_client', 'reference')
+
+    def __init__(self,
+        reference: ConnectionRef,
+        embed: Query[EmbedInput, EmbedOutput],
+        rerank: Query[RerankInput, RerankOutput],
+        client: Client | None = None,
+    ) -> None:
+        reference.validate()
+        if reference.provider != CONNECTOR_ID or embed is None or rerank is None:
+            raise ValueError("Cohere connection provider and operations are required")
+        if embed.definition().operation != EMBED_DEFINITION.operation:
+            raise ValueError("Cohere connection operation identity differs")
+        if rerank.definition().operation != RERANK_DEFINITION.operation:
+            raise ValueError("Cohere connection operation identity differs")
+        self.reference = reference
+        self._embed = embed
+        self._rerank = rerank
+        self._client = client
+
+    def embed(self) -> Query[EmbedInput, EmbedOutput]:
+        return self._embed
+
+    def rerank(self) -> Query[RerankInput, RerankOutput]:
+        return self._rerank
+
+    async def aclose(self) -> None:
+        if self._client is not None:
+            await self._client.aclose()
+
+    def __repr__(self) -> str:
+        return "cohere.Connection([REDACTED])"
+
+    def __reduce__(self) -> tuple[object, ...]:
+        raise TypeError("Cohere connections cannot be serialized")
+
+setattr(Connection, '__module__', 'dex_connectors.cohere.connection')
+
+
+def new_connection(client: Client, reference: ConnectionRef) -> Connection:
+    if client is None:
+        raise ValueError("Cohere client is required")
+    return Connection(reference, client.embed(), client.rerank(), client)
+
+setattr(new_connection, '__module__', 'dex_connectors.cohere.connection')
+
+
+def new_connection_with_operations(
+    embed: Query[EmbedInput, EmbedOutput],
+    rerank: Query[RerankInput, RerankOutput],
+    reference: ConnectionRef,
+) -> Connection:
+    return Connection(reference, embed, rerank)
+
+setattr(new_connection_with_operations, '__module__', 'dex_connectors.cohere.connection')
+
+
+class _ProjectCredentialSource:
+
+    def __init__(self, project: LoadedProject, key: ConnectionKey) -> None:
+        self.project = project
+        self.key = key
+
+    async def resolve(self, call: Call) -> Credentials:
+        material = await self.project.resolve_credential(self.key, str(call.id))
+        api_key = material.get("api_key")
+        if not isinstance(api_key, str) or not api_key:
+            raise ValueError("Cohere credential api_key is required")
+        return Credentials(SecretString(api_key))
+
+
+def new_project_connection(
+    project: LoadedProject, connection_name: str, http_client: httpx.AsyncClient | None = None
+) -> Connection:
+    if project is None:
+        raise ValueError("Cohere connection requires loaded project configuration")
+    reference = ConnectionRef(CONNECTOR_ID, connection_name)
+    reference.validate()
+    key = ConnectionKey(CONNECTOR_ID, connection_name)
+    matching = [
+        entry
+        for entry in project.configuration.connections
+        if entry.connector_id == CONNECTOR_ID and entry.connection_name == connection_name
+    ]
+    if len(matching) != 1:
+        raise ValueError("Cohere project connection is missing or duplicated")
+    if matching[0].provider != CONNECTOR_ID:
+        raise ValueError("Cohere project connection provider differs")
+    config = _decode_config(matching[0].configuration)
+    client = Client(config, _ProjectCredentialSource(project, key), http_client)
+    return new_connection(client, reference)
+
+setattr(new_project_connection, '__module__', 'dex_connectors.cohere.connection')
+
+
+def _decode_config(value: dict[str, Any]) -> Config:
+    if set(value) - {"baseUrl", "timeout", "maxResponseBytes"}:
+        raise ValueError("Cohere configuration has unknown fields")
+    config_arguments: dict[str, Any] = {}
+    if "baseUrl" in value:
+        base_url = value["baseUrl"]
+        if not (isinstance(base_url, str)):
+            raise ValueError("Cohere configuration baseUrl has an invalid type")
+        config_arguments['base_url'] = base_url
+    if "timeout" in value:
+        timeout = value["timeout"]
+        if not (type(timeout) is int):
+            raise ValueError("Cohere configuration timeout has an invalid type")
+        config_arguments['timeout_seconds'] = timeout / 1_000_000_000
+    if "maxResponseBytes" in value:
+        max_response_bytes = value["maxResponseBytes"]
+        if not (type(max_response_bytes) is int):
+            raise ValueError("Cohere configuration maxResponseBytes has an invalid type")
+        config_arguments['max_response_bytes'] = max_response_bytes
+    return Config(**config_arguments)

dex_connector_cohere-0.1.0/src/dex_connectors/cohere/_generated/definitions.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2026 Super Durable
+# SPDX-License-Identifier: MIT
+# ruff: noqa
+
+from datetime import timedelta
+
+from dex_connectors.sdk import (
+    BranchDefinition,
+    MutationDefinition,
+    OperationRef,
+    QueryDefinition,
+    RetryPolicy,
+    StepDefaults,
+    default_connector_step_execute_retry_policy,
+)
+
+CONNECTOR_ID = "cohere"
+PROVIDER = "cohere"
+EMBED_OPERATION_ID = "embed"
+RERANK_OPERATION_ID = "rerank"
+DEFECT = "defect"
+EMBEDDED = "embedded"
+INVALID_RESPONSE = "invalidResponse"
+PROVIDER_REJECTED = "providerRejected"
+RANKED = "ranked"
+
+EMBED_DEFINITION = QueryDefinition(
+    OperationRef(CONNECTOR_ID, EMBED_OPERATION_ID),
+    (
+        BranchDefinition(EMBEDDED, "Cohere returned one float vector per text.", False),
+        BranchDefinition(PROVIDER_REJECTED, "Cohere conclusively rejected the embedding request.", True),
+        BranchDefinition(INVALID_RESPONSE, "Cohere returned malformed or oversized embedding data.", True),
+        BranchDefinition(DEFECT, "Local input or connection credentials are invalid.", True),
+    ),
+    StepDefaults(
+        timedelta(seconds=45),
+        2,
+        execute_retry=RetryPolicy(timedelta(seconds=2), 2, timedelta(seconds=30), 3, timedelta(seconds=300)),
+    ),
+)
+
+RERANK_DEFINITION = QueryDefinition(
+    OperationRef(CONNECTOR_ID, RERANK_OPERATION_ID),
+    (
+        BranchDefinition(RANKED, "Cohere returned ranked document indices and scores.", False),
+        BranchDefinition(PROVIDER_REJECTED, "Cohere conclusively rejected the rerank request.", True),
+        BranchDefinition(INVALID_RESPONSE, "Cohere returned malformed or oversized rerank data.", True),
+        BranchDefinition(DEFECT, "Local input or connection credentials are invalid.", True),
+    ),
+    StepDefaults(
+        timedelta(seconds=45),
+        2,
+        execute_retry=RetryPolicy(timedelta(seconds=2), 2, timedelta(seconds=30), 3, timedelta(seconds=300)),
+    ),
+)