crewplus 0.3.4__tar.gz → 0.3.6__tar.gz

{crewplus-0.3.4 → crewplus-0.3.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: crewplus
-Version: 0.3.4
+Version: 0.3.6
 Summary: Base services for CrewPlus AI applications
 Author-Email: Tim Liu <tim@opsmateai.com>
 License: MIT
@@ -16,6 +16,8 @@ Requires-Dist: anthropic[vertex]>=0.76.0
 Requires-Dist: langchain-milvus<0.3.0,>=0.2.1
 Requires-Dist: langfuse<4.0.0,>=3.1.3
 Requires-Dist: pymilvus<=2.6.7,>=2.5.7
+Requires-Dist: httpx<1.0,>=0.27
+Requires-Dist: opentelemetry-api>=1.20
 Description-Content-Type: text/markdown
 
 # CrewPlus

{crewplus-0.3.4 → crewplus-0.3.6}/crewplus/services/__init__.py

@@ -7,6 +7,13 @@ from .azure_chat_model import TracedAzureChatOpenAI
 from .feedback_manager import LangfuseFeedbackManager
 from .schemas.feedback import FeedbackIn, FeedbackUpdate, FeedbackOut
 from .tracing_manager import TracingManager, TracingContext
+from .rerankers import (
+    CohereFoundryReranker,
+    RerankError,
+    RerankItem,
+    RerankResponse,
+    create_reranker,
+)
 from .capabilities import (
     BaseCapability,
     CapabilityConfig,
@@ -29,6 +36,11 @@ __all__ = [
     "FeedbackOut",
     "TracingManager",
     "TracingContext",
+    "CohereFoundryReranker",
+    "RerankError",
+    "RerankItem",
+    "RerankResponse",
+    "create_reranker",
     "BaseCapability",
     "CapabilityConfig",
     "CapabilityRegistry",

{crewplus-0.3.4 → crewplus-0.3.6}/crewplus/services/model_load_balancer.py

@@ -9,6 +9,7 @@ from .gemini_chat_model import GeminiChatModel
 from .claude_chat_model import ClaudeChatModel
 from .azure_chat_model import TracedAzureChatOpenAI
 from .oai_claude_chat_model import OAIClaudeChatModel
+from .rerankers import RERANK_PROVIDERS, create_reranker
 
 
 class ModelLoadBalancer:
@@ -132,7 +133,9 @@ class ModelLoadBalancer:
 
         Args:
             provider: The model provider (e.g., 'azure-openai', 'google-genai').
-            model_type: The type of model (e.g., 'inference', 'embedding', 'embedding-large').
+            model_type: The type of model (e.g., 'inference', 'embedding', 'embedding-large', 'rerank').
+                A 'rerank' entry (provider 'azure-foundry-cohere') returns a rerank client,
+                not a LangChain model.
             deployment_name: The unique name for the model deployment.
             with_metadata: If True, returns a tuple of (model, deployment_name).
             selection_strategy: The selection strategy ('random', 'round_robin', or 'least_used'). Defaults to 'random'.
@@ -358,6 +361,10 @@ class ModelLoadBalancer:
             # Custom OpenAI-compatible endpoint for Claude models
             # Uses OpenAI-compatible API with capability support (web_search, bash, etc.)
             return self._instantiate_oai_claude_model(model_config, disable_streaming)
+        elif provider in RERANK_PROVIDERS:
+            # Unlike embeddings, a bad rerank entry raises: the host decides whether
+            # running without rerank is acceptable, not the balancer.
+            return create_reranker(model_config)
         else:
             self.logger.error(f"Unsupported provider: {provider}")
             raise ValueError(f"Unsupported provider: {provider}")

crewplus-0.3.6/crewplus/services/rerankers/__init__.py

@@ -0,0 +1,19 @@
+"""
+Rerank model clients (cross-encoders) configured from MODELS_CONFIG.
+
+Only transport lives here. Retrieval policy — candidate selection, document text,
+fallback to vector order — belongs to the caller (crewplus-search).
+"""
+from .base import RerankError, RerankErrorKind, RerankItem, RerankResponse
+from .cohere_foundry import CohereFoundryReranker
+from .factory import RERANK_PROVIDERS, create_reranker
+
+__all__ = [
+    "RerankError",
+    "RerankErrorKind",
+    "RerankItem",
+    "RerankResponse",
+    "CohereFoundryReranker",
+    "RERANK_PROVIDERS",
+    "create_reranker",
+]

crewplus-0.3.6/crewplus/services/rerankers/base.py

@@ -0,0 +1,51 @@
+"""
+Shared result and error types for rerank model clients.
+
+Callers (e.g. crewplus-search) depend on these shapes structurally, not by import,
+so the fields here are a contract: renaming one breaks every caller silently.
+"""
+from dataclasses import dataclass
+from typing import List, Literal, Optional
+
+RerankErrorKind = Literal[
+    "timeout",
+    "rate_limited",
+    "http_4xx",
+    "http_5xx",
+    "bad_response",
+    "transport",
+    "circuit_open",
+]
+
+
+@dataclass(frozen=True)
+class RerankItem:
+    """One reranked document: its position in the input list and its relevance."""
+    index: int
+    relevance_score: float
+
+
+@dataclass(frozen=True)
+class RerankResponse:
+    """Reranked documents, sorted by relevance_score descending."""
+    items: List[RerankItem]
+    search_units: Optional[int]
+    latency_ms: float
+    model: str
+
+
+class RerankError(Exception):
+    """
+    A rerank call that produced no usable ranking.
+
+    ``kind`` is what callers branch on (and log as the fallback reason); the
+    message is for humans and never contains the query, documents or credentials.
+    """
+
+    def __init__(self, kind: RerankErrorKind, message: str, *, status: Optional[int] = None):
+        super().__init__(message)
+        self.kind: RerankErrorKind = kind
+        self.status = status
+
+    def __repr__(self) -> str:
+        return f"RerankError(kind={self.kind!r}, status={self.status!r}, message={str(self)!r})"

crewplus-0.3.6/crewplus/services/rerankers/cohere_foundry.py

@@ -0,0 +1,337 @@
+"""
+Cohere Rerank client for models deployed on Azure AI Foundry.
+
+Transport only: one HTTP call per ``arerank``, classified errors, a process-local
+circuit breaker. Deciding what to do when a call fails (e.g. falling back to vector
+order) is the caller's job — this client never retries and never degrades silently.
+"""
+import asyncio
+import json
+import logging
+import math
+import threading
+import time
+from typing import Any, Callable, Dict, List, Optional, Sequence
+
+import httpx
+from opentelemetry import trace
+from opentelemetry.trace import Status, StatusCode
+
+from .base import RerankError, RerankErrorKind, RerankItem, RerankResponse
+
+logger = logging.getLogger(__name__)
+_tracer = trace.get_tracer(__name__)
+
+RERANK_PATH = "/providers/cohere/v2/rerank"
+
+# Errors that say the service is unhealthy or overloaded. A 4xx or a malformed body
+# means a config or contract problem that no amount of waiting fixes, so those do
+# not trip the breaker.
+_BREAKER_COUNTED_KINDS = frozenset({"timeout", "transport", "http_5xx", "rate_limited"})
+
+
+class _CircuitBreaker:
+    """Opens after N consecutive counted failures; lets one probe through after a cool-down."""
+
+    def __init__(self, threshold: int, open_seconds: float, clock: Callable[[], float]):
+        self._threshold = threshold
+        self._open_seconds = open_seconds
+        self._clock = clock
+        self._lock = threading.Lock()
+        self._failures = 0
+        self._state = "closed"
+        self._opened_at = 0.0
+        self._probe_in_flight = False
+
+    @property
+    def state(self) -> str:
+        return self._state
+
+    def allow(self) -> bool:
+        if self._threshold <= 0:
+            return True
+        with self._lock:
+            if self._state == "closed":
+                return True
+            if self._state == "open":
+                if self._clock() - self._opened_at < self._open_seconds:
+                    return False
+                self._state = "half_open"
+                self._probe_in_flight = False
+            if self._probe_in_flight:
+                return False
+            self._probe_in_flight = True
+            return True
+
+    def record(self, error_kind: Optional[str]) -> None:
+        if self._threshold <= 0:
+            return
+        with self._lock:
+            self._probe_in_flight = False
+            if error_kind not in _BREAKER_COUNTED_KINDS:
+                # The service answered (successfully or with a client-side error).
+                if self._state != "closed":
+                    logger.info("[Rerank] circuit closed after a successful probe")
+                self._failures = 0
+                self._state = "closed"
+                return
+            self._failures += 1
+            if self._state == "half_open" or self._failures >= self._threshold:
+                if self._state != "open":
+                    logger.warning(
+                        "[Rerank] circuit opened after %d consecutive failures (last=%s)",
+                        self._failures, error_kind,
+                    )
+                self._state = "open"
+                self._opened_at = self._clock()
+
+
+class CohereFoundryReranker:
+    """
+    Calls ``POST {api_base}/providers/cohere/v2/rerank`` on an Azure AI Foundry resource.
+
+    One ``httpx.AsyncClient`` is kept per event loop and reused, because most of a
+    rerank call's latency is the network round trip and a fresh TLS handshake per
+    call would roughly double it.
+
+    Args:
+        api_base: Foundry resource URL, e.g. ``https://<resource>.services.ai.azure.com``.
+        api_key: Resource key, sent as the ``api-key`` header.
+        deployment_name: Model deployment name, sent as ``model``.
+        timeout_s: Whole-request deadline (connect + send + wait + read together).
+        max_tokens_per_doc: Per-document truncation applied by the service.
+        breaker_threshold: Consecutive counted failures that open the circuit; 0 disables it.
+        breaker_open_seconds: How long the circuit stays open before one probe is allowed.
+        transport: Test hook; a custom httpx transport.
+        clock: Test hook; monotonic clock used by the breaker.
+    """
+
+    def __init__(
+        self,
+        *,
+        api_base: str,
+        api_key: str,
+        deployment_name: str,
+        timeout_s: float = 3.0,
+        max_tokens_per_doc: int = 4096,
+        breaker_threshold: int = 5,
+        breaker_open_seconds: float = 30.0,
+        transport: Optional[httpx.AsyncBaseTransport] = None,
+        clock: Callable[[], float] = time.monotonic,
+    ):
+        if not api_base or not api_key or not deployment_name:
+            raise ValueError("api_base, api_key and deployment_name are required")
+        if timeout_s <= 0:
+            raise ValueError("timeout_s must be positive")
+        if max_tokens_per_doc <= 0:
+            raise ValueError("max_tokens_per_doc must be positive")
+
+        self._url = api_base.rstrip("/") + RERANK_PATH
+        self._api_key = api_key
+        self._deployment = deployment_name
+        self._timeout_s = timeout_s
+        self._max_tokens_per_doc = max_tokens_per_doc
+        self._transport = transport
+        self._breaker = _CircuitBreaker(breaker_threshold, breaker_open_seconds, clock)
+        self._clients: Dict[asyncio.AbstractEventLoop, httpx.AsyncClient] = {}
+        self._clients_lock = threading.Lock()
+
+    @property
+    def model(self) -> str:
+        return self._deployment
+
+    @property
+    def circuit_state(self) -> str:
+        return self._breaker.state
+
+    def __repr__(self) -> str:
+        return f"CohereFoundryReranker(deployment={self._deployment!r}, url={self._url!r})"
+
+    async def arerank(
+        self,
+        query: str,
+        documents: Sequence[str],
+        *,
+        top_n: Optional[int] = None,
+    ) -> RerankResponse:
+        """
+        Rerank ``documents`` against ``query``.
+
+        Returns items sorted by relevance, each pointing back into ``documents`` by index.
+
+        Raises:
+            RerankError: on any failure; ``kind`` classifies it.
+        """
+        if not documents:
+            return RerankResponse(items=[], search_units=None, latency_ms=0.0, model=self._deployment)
+        if top_n is not None and top_n <= 0:
+            raise ValueError("top_n must be positive")
+
+        body: Dict[str, Any] = {
+            "model": self._deployment,
+            "query": query,
+            "documents": list(documents),
+            "max_tokens_per_doc": self._max_tokens_per_doc,
+        }
+        if top_n is not None:
+            body["top_n"] = min(top_n, len(documents))
+        # Encoded once so the payload size can be logged: production latency is judged
+        # from these lines, and size is the first thing to correlate it with.
+        content = json.dumps(body, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
+        stats = {"docs": len(documents), "top_n": top_n, "payload_bytes": len(content)}
+
+        if not self._breaker.allow():
+            self._log_outcome(stats, 0.0, None, "circuit_open")
+            raise RerankError("circuit_open", "rerank circuit is open; call skipped")
+
+        with _tracer.start_as_current_span("rerank.cohere") as span:
+            span.set_attribute("rerank.deployment", self._deployment)
+            span.set_attribute("rerank.documents", len(documents))
+            span.set_attribute("rerank.payload_bytes", len(content))
+            if top_n is not None:
+                span.set_attribute("rerank.top_n", top_n)
+            started = time.perf_counter()
+            try:
+                response = await self._post(content)
+                result = self._parse(response, len(documents), started)
+            except RerankError as e:
+                latency_ms = (time.perf_counter() - started) * 1000
+                self._breaker.record(e.kind)
+                span.set_attribute("error.kind", e.kind)
+                span.set_attribute("rerank.latency_ms", latency_ms)
+                span.set_status(Status(StatusCode.ERROR, e.kind))
+                self._log_outcome(stats, latency_ms, None, e.kind, status=e.status)
+                raise
+            self._breaker.record(None)
+            span.set_attribute("rerank.latency_ms", result.latency_ms)
+            if result.search_units is not None:
+                span.set_attribute("rerank.search_units", result.search_units)
+            self._log_outcome(
+                stats, result.latency_ms, result.search_units, None,
+                top1=result.items[0].index if result.items else None,
+                top1_score=result.items[0].relevance_score if result.items else None,
+            )
+            return result
+
+    async def aclose(self) -> None:
+        """Close every per-loop HTTP client this instance opened."""
+        with self._clients_lock:
+            clients = list(self._clients.values())
+            self._clients.clear()
+        for client in clients:
+            try:
+                await client.aclose()
+            except Exception as e:  # a client bound to an already-closed loop
+                logger.debug("[Rerank] ignoring error while closing client: %s", e)
+
+    def _client(self) -> httpx.AsyncClient:
+        loop = asyncio.get_running_loop()
+        with self._clients_lock:
+            for stale in [lp for lp in self._clients if lp.is_closed()]:
+                self._clients.pop(stale, None)
+            client = self._clients.get(loop)
+            if client is None:
+                client = httpx.AsyncClient(
+                    timeout=httpx.Timeout(self._timeout_s),
+                    transport=self._transport,
+                )
+                self._clients[loop] = client
+            return client
+
+    async def _post(self, content: bytes) -> httpx.Response:
+        # httpx timeouts are per phase (connect / write / read), so a slow upload plus a
+        # slow read could take several times timeout_s; wait_for makes it one deadline.
+        try:
+            response = await asyncio.wait_for(self._send(content), timeout=self._timeout_s)
+        except asyncio.TimeoutError as e:
+            raise RerankError("timeout", f"rerank request exceeded the {self._timeout_s}s deadline") from e
+
+        status = response.status_code
+        if status == 429:
+            raise RerankError("rate_limited", "rerank rate limited (429)", status=status)
+        if status >= 500:
+            raise RerankError("http_5xx", f"rerank server error ({status})", status=status)
+        if status >= 300:
+            raise RerankError("http_4xx", f"rerank request rejected ({status})", status=status)
+        return response
+
+    async def _send(self, content: bytes) -> httpx.Response:
+        try:
+            return await self._client().post(
+                self._url,
+                headers={"api-key": self._api_key, "Content-Type": "application/json"},
+                content=content,
+            )
+        except httpx.TimeoutException as e:
+            raise RerankError("timeout", f"rerank request timed out after {self._timeout_s}s") from e
+        except httpx.HTTPError as e:
+            raise RerankError("transport", f"rerank transport error: {type(e).__name__}") from e
+
+    def _parse(self, response: httpx.Response, n_docs: int, started: float) -> RerankResponse:
+        try:
+            payload = response.json()
+        except ValueError as e:
+            raise RerankError("bad_response", "rerank response is not JSON") from e
+
+        results = payload.get("results") if isinstance(payload, dict) else None
+        if not isinstance(results, list):
+            raise RerankError("bad_response", "rerank response has no results list")
+
+        items: List[RerankItem] = []
+        seen = set()
+        for entry in results:
+            if not isinstance(entry, dict):
+                raise RerankError("bad_response", "rerank result entry is not an object")
+            index = entry.get("index")
+            score = entry.get("relevance_score")
+            if isinstance(index, bool) or not isinstance(index, int) or not 0 <= index < n_docs or index in seen:
+                raise RerankError("bad_response", f"rerank result has an invalid index: {index!r}")
+            if isinstance(score, bool) or not isinstance(score, (int, float)) or math.isnan(score):
+                raise RerankError("bad_response", "rerank result has an invalid relevance_score")
+            seen.add(index)
+            items.append(RerankItem(index=index, relevance_score=float(score)))
+        items.sort(key=lambda item: item.relevance_score, reverse=True)
+
+        search_units = None
+        meta = payload.get("meta")
+        if isinstance(meta, dict) and isinstance(meta.get("billed_units"), dict):
+            units = meta["billed_units"].get("search_units")
+            if isinstance(units, (int, float)) and not isinstance(units, bool):
+                search_units = int(units)
+
+        return RerankResponse(
+            items=items,
+            search_units=search_units,
+            latency_ms=(time.perf_counter() - started) * 1000,
+            model=self._deployment,
+        )
+
+    def _log_outcome(
+        self,
+        stats: Dict[str, Any],
+        latency_ms: float,
+        search_units: Optional[int],
+        error_kind: Optional[RerankErrorKind],
+        *,
+        status: Optional[int] = None,
+        top1: Optional[int] = None,
+        top1_score: Optional[float] = None,
+    ) -> None:
+        # One line per call, same key=value layout either way, so a log query can group
+        # by deployment / outcome and plot latency against payload size. Never log the
+        # query, the documents or the key: they may carry customer data.
+        if error_kind is None:
+            logger.info(
+                "[Rerank] deployment=%s outcome=ok docs=%d top_n=%s payload_bytes=%d latency_ms=%.1f "
+                "search_units=%s top1_index=%s top1_score=%s circuit=%s",
+                self._deployment, stats["docs"], stats["top_n"], stats["payload_bytes"], latency_ms,
+                search_units, top1, f"{top1_score:.4f}" if top1_score is not None else None,
+                self._breaker.state,
+            )
+        else:
+            logger.warning(
+                "[Rerank] deployment=%s outcome=%s docs=%d top_n=%s payload_bytes=%d latency_ms=%.1f "
+                "status=%s timeout_s=%s circuit=%s",
+                self._deployment, error_kind, stats["docs"], stats["top_n"], stats["payload_bytes"],
+                latency_ms, status, self._timeout_s, self._breaker.state,
+            )

crewplus-0.3.6/crewplus/services/rerankers/factory.py

@@ -0,0 +1,61 @@
+"""
+Build rerank clients from MODELS_CONFIG entries.
+
+The same factory serves ModelLoadBalancer (agent-executor) and hosts that derive a
+single entry from MODELS_CONFIG without running a balancer (sandbox), so both get
+identical clients from identical config.
+"""
+from typing import Any, Mapping
+
+from .cohere_foundry import CohereFoundryReranker
+
+RERANK_PROVIDERS = frozenset({"azure-foundry-cohere"})
+
+_REQUIRED_FIELDS = ("api_base", "api_key", "deployment_name")
+
+
+def create_reranker(entry: Mapping[str, Any]) -> CohereFoundryReranker:
+    """
+    Build a reranker from one MODELS_CONFIG entry.
+
+    Expected entry::
+
+        {"provider": "azure-foundry-cohere", "type": "rerank",
+         "api_base": "https://<resource>.services.ai.azure.com",
+         "deployment_name": "Cohere-rerank-v4.0-pro", "api_key": "...",
+         "timeout_ms": 3000, "max_tokens_per_doc": 4096,
+         "breaker": {"consecutive_failures": 5, "open_seconds": 30}}
+
+    Raises:
+        ValueError: unknown provider or a missing/invalid field. The message names
+            the field only, never its value.
+    """
+    provider = entry.get("provider")
+    if provider not in RERANK_PROVIDERS:
+        raise ValueError(f"Unsupported rerank provider: {provider!r}")
+
+    missing = [field for field in _REQUIRED_FIELDS if not entry.get(field)]
+    if missing:
+        raise ValueError(f"Rerank model config is missing required field(s): {', '.join(missing)}")
+
+    breaker = entry.get("breaker") or {}
+    if not isinstance(breaker, Mapping):
+        raise ValueError("Rerank model config field 'breaker' must be an object")
+
+    try:
+        timeout_s = float(entry.get("timeout_ms", 3000)) / 1000.0
+        max_tokens_per_doc = int(entry.get("max_tokens_per_doc", 4096))
+        breaker_threshold = int(breaker.get("consecutive_failures", 5))
+        breaker_open_seconds = float(breaker.get("open_seconds", 30))
+    except (TypeError, ValueError) as e:
+        raise ValueError(f"Rerank model config has a non-numeric tuning field: {e}") from None
+
+    return CohereFoundryReranker(
+        api_base=str(entry["api_base"]),
+        api_key=str(entry["api_key"]),
+        deployment_name=str(entry["deployment_name"]),
+        timeout_s=timeout_s,
+        max_tokens_per_doc=max_tokens_per_doc,
+        breaker_threshold=breaker_threshold,
+        breaker_open_seconds=breaker_open_seconds,
+    )