Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 29 additions & 10 deletions contextual_orchestrator/cost_ledger.py
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,8 @@ class UsageRecord:

def as_dict(self) -> Dict[str, Any]:
"""Flatten the record (attribution inlined) for JSON + SQL storage."""
# Execution identity is evidence of what ran — never a client-chosen tag.
# Account/service/team/group/company remain descriptive attribution.
row = {
"usage_record_id": self.usage_record_id,
"created_at": self.created_at,
Expand Down Expand Up @@ -583,12 +585,12 @@ def _seed_dimension_catalog(self) -> None:
ph = self._placeholder()
cur = self._conn.cursor()
for order, (name, label, _column) in enumerate(ATTRIBUTION_DIMENSION_CATALOG):
cur.execute(
cur.execute( # nosemgrep -- sqlalchemy-execute-raw-query FP: only the DB-API placeholder char is interpolated; the value is bound.
f"SELECT 1 FROM cost_attribution_dimensions WHERE dimension_name = {ph}", # nosec B608 - ph is a DB-API placeholder.
(name,),
)
if cur.fetchone() is None:
cur.execute(
cur.execute( # nosemgrep -- sqlalchemy-execute-raw-query FP: only DB-API placeholder chars are interpolated; values are bound.
"INSERT INTO cost_attribution_dimensions "
f"(dimension_name, dimension_label, dimension_order) VALUES ({ph}, {ph}, {ph})", # nosec B608 - ph is a DB-API placeholder.
(name, label, order),
Expand All @@ -602,7 +604,7 @@ def append(self, record: UsageRecord) -> None:
placeholders = ", ".join(ph for _ in _USAGE_COLUMNS)
columns = ", ".join(_USAGE_COLUMNS)
cur = self._conn.cursor()
cur.execute(
cur.execute( # nosemgrep -- sqlalchemy-execute-raw-query FP: columns are the fixed _USAGE_COLUMNS constant; values are bound.
f"INSERT INTO llm_usage_records ({columns}) VALUES ({placeholders})", # nosec B608 - columns are fixed _USAGE_COLUMNS.
tuple(row.get(column) for column in _USAGE_COLUMNS),
)
Expand All @@ -622,7 +624,7 @@ def query(self, start: Optional[int] = None, end: Optional[int] = None) -> List[
where = f" WHERE {' AND '.join(clauses)}" if clauses else ""
columns = ", ".join(_USAGE_COLUMNS)
cur = self._conn.cursor()
cur.execute(f"SELECT {columns} FROM llm_usage_records{where}", tuple(params)) # nosec B608 - columns and clauses are fixed.
cur.execute(f"SELECT {columns} FROM llm_usage_records{where}", tuple(params)) # nosec B608 - columns and clauses are fixed. # nosemgrep -- sqlalchemy-execute-raw-query FP: fixed columns and clause templates; all values are bound.
return [dict(zip(_USAGE_COLUMNS, values)) for values in cur.fetchall()]


Expand Down Expand Up @@ -681,14 +683,31 @@ def record_usage(
) -> UsageRecord:
"""Compute cost, build a :class:`UsageRecord`, persist it, and return it."""
if isinstance(attribution, dict) or attribution is None:
dims = AttributionDimensions.from_mapping(attribution)
# Strip caller-controlled execution identity before mapping so a
# client cannot spoof model/provider rollups (buyer-bill honesty).
if isinstance(attribution, dict):
cleaned = {
key: value
for key, value in attribution.items()
if key not in {"model_name", "provider", "upstream_api"}
}
else:
cleaned = None
dims = AttributionDimensions.from_mapping(cleaned)
else:
dims = attribution
# Keep the model_name dimension aligned with the served model unless the
# caller pinned it explicitly, and default the provider dimension too.
if dims.model_name == UNATTRIBUTED and model:
dims = AttributionDimensions(
account=attribution.account,
service=attribution.service,
upstream_api=UNATTRIBUTED,
model_name=UNATTRIBUTED,
team=attribution.team,
group=attribution.group,
company=attribution.company,
)
# Execution identity always wins — descriptive dimensions stay as-is.
if model:
dims.model_name = model
if dims.upstream_api == UNATTRIBUTED and provider:
if provider:
dims.upstream_api = provider

cost_amount, currency = self.price_book.compute_cost(
Expand Down
159 changes: 150 additions & 9 deletions contextual_orchestrator/orchestrator.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,8 @@
from .credentials import NotConfigured, get_credential


ChatMessage = dict[str, str]
# content is usually str; multimodal vision messages use OpenAI content-parts lists.
ChatMessage = dict[str, Any]

class BudgetExceededError(RuntimeError):
"""Raised when an operator-configured spend budget is already exhausted."""
Expand Down Expand Up @@ -215,6 +216,10 @@ def __init__(
) -> None:
self.timeout = timeout
self.max_output_tokens = max_output_tokens
self.default_temperature = 0.2
self.default_top_p: float | None = None
self.default_presence_penalty: float | None = None
self.default_frequency_penalty: float | None = None
self.max_retries = max_retries
self.retry_backoff = retry_backoff
self.retry_backoff_cap = retry_backoff_cap
Expand All @@ -230,7 +235,7 @@ def __init__(
@staticmethod
def _build_ssl_context(ca_bundle: str | None, verify_tls: bool) -> ssl.SSLContext:
if not verify_tls:
return ssl._create_unverified_context() # nosec B323 - explicit dev-only provider TLS opt-out.
return ssl._create_unverified_context() # nosec B323 - explicit dev-only provider TLS opt-out. # nosemgrep -- unverified-ssl-context: intentional, default-secure (verify_tls defaults True) dev-only opt-out for self-signed endpoints.
if ca_bundle:
if not os.path.isfile(ca_bundle):
raise ValueError(f"provider CA bundle does not exist: {ca_bundle}")
Expand All @@ -246,9 +251,29 @@ def take_usage(self) -> dict[str, Any] | None:
self._local.usage = None
return usage

def chat(self, agent: ModelAgent, messages: list[ChatMessage], temperature: float = 0.2) -> str:
"""Send messages to a mock or OpenAI-compatible chat endpoint with retries."""
def chat(
self,
agent: ModelAgent,
messages: list[ChatMessage],
temperature: float | None = None,
top_p: float | None = None,
) -> str:
"""Send messages to a mock or OpenAI-compatible chat endpoint with retries.

When ``temperature``/``top_p`` are omitted, ``default_temperature`` and
``default_top_p`` are used so request-scoped Completions sampling can be
applied without threading kwargs through every orchestrator hop.
"""
self._local.usage = None
# Expose the effective sampling knobs for request-path tests / diagnostics.
effective_temperature = self.default_temperature if temperature is None else temperature
effective_top_p = self.default_top_p if top_p is None else top_p
effective_presence = self.default_presence_penalty
effective_frequency = self.default_frequency_penalty
self._local.last_temperature = effective_temperature
self._local.last_top_p = effective_top_p
self._local.last_presence_penalty = effective_presence
self._local.last_frequency_penalty = effective_frequency
if agent.base_url.startswith("mock://"):
return self._mock(agent, messages)

Expand All @@ -262,10 +287,16 @@ def chat(self, agent: ModelAgent, messages: list[ChatMessage], temperature: floa
payload = { # pragma: no cover
"model": agent.model,
"messages": messages,
"temperature": temperature,
"temperature": effective_temperature,
"stream": False,
"max_tokens": self.max_output_tokens,
}
if effective_top_p is not None: # pragma: no cover
payload["top_p"] = effective_top_p
if effective_presence is not None: # pragma: no cover
payload["presence_penalty"] = effective_presence
if effective_frequency is not None: # pragma: no cover
payload["frequency_penalty"] = effective_frequency
return self._send_with_retry(agent, payload)

def _send_with_retry(self, agent: ModelAgent, payload: dict[str, Any]) -> str:
Expand Down Expand Up @@ -307,7 +338,7 @@ def _send(self, agent: ModelAgent, payload: dict[str, Any]) -> str:

def _open_provider(self, request: urllib.request.Request) -> Any:
"""Open a provider request built from a validated provider URL."""
return urllib.request.urlopen( # nosec B310 - request URL comes from _provider_url after provider validation.
return urllib.request.urlopen( # nosec B310 - request URL comes from _provider_url after provider validation. # nosemgrep -- dynamic-urllib-use: URL is built by _provider_url after scheme/host validation; egress to loopback/private/reserved is blocked.
request,
timeout=self.timeout,
context=self._ssl_context,
Expand Down Expand Up @@ -634,6 +665,9 @@ def _coerce_input_text(value: Any) -> str:
if isinstance(item, str):
parts.append(item)
elif isinstance(item, dict):
# OpenAI content-parts: {"type": "text", "text": "..."}
if isinstance(item.get("text"), str):
parts.append(item["text"])
content = item.get("content")
if isinstance(content, str):
parts.append(content)
Expand All @@ -644,6 +678,15 @@ def _coerce_input_text(value: Any) -> str:
return " ".join(parts)


def _coerce_message_content_text(content: Any) -> str:
"""Best-effort plain text from chat message content (string or content-parts)."""
if isinstance(content, str):
return content
if isinstance(content, list):
return _coerce_input_text(content)
return ""


def load_agents(path: str) -> list[ModelAgent]: # pragma: no cover
"""Load model agent definitions from an agents JSON file."""
with open(path, encoding="utf-8") as handle:
Expand Down Expand Up @@ -879,7 +922,14 @@ def _reload_state(self) -> None:

# Orchestration-only body keys that must not be forwarded to the provider.
_ORCHESTRATION_ONLY_KEYS = frozenset(
{"orchestration", "orchestration_mode", "mode", "include_orchestration_trace"}
{
"orchestration",
"orchestration_mode",
"mode",
"include_orchestration_trace",
"attribution",
"routing",
}
)

def proxy_completion(
Expand All @@ -897,7 +947,23 @@ def proxy_completion(
text = self._latest_user_text(messages)
else:
text = _coerce_input_text(body.get("input"))
agent = self._select_agent(text, "worker")
requested_model = body.get("model")
# When the client names a model, resolve a pool agent that actually serves
# that model id. Silent rewrite to an unrelated agent.model is a commercial
# honesty failure for OpenAI SDKs (passthrough tools/Responses paths).
if isinstance(requested_model, str) and requested_model.strip():
matched = [
agent
for agent in self.agents
if not getattr(agent, "disabled", False) and agent.model == requested_model
]
if not matched:
raise ValueError(
f"model {requested_model!r} is not available in the agent pool"
)
agent = matched[0]
else:
agent = self._select_agent(text, "worker")
upstream = {
key: value
for key, value in body.items()
Expand Down Expand Up @@ -1600,7 +1666,13 @@ def _needs_workflow(self, text: str) -> bool:
return hits >= self.policy.conduct_hint_threshold or len(text) > 700

def _latest_user_text(self, messages: list[ChatMessage]) -> str:
return next((m.get("content", "") for m in reversed(messages) if m.get("role") == "user"), "") # pragma: no cover
for message in reversed(messages):
if message.get("role") != "user":
continue
text = _coerce_message_content_text(message.get("content", ""))
if text:
return text
return "" # pragma: no cover

def _model_judge_verification(self, task: str, fallback: dict[str, Any]) -> dict[str, Any]:
"""Ask a model to judge the verifier report (fixes term-matching false negatives).
Expand Down Expand Up @@ -1695,6 +1767,52 @@ def list_agents(self, page_number: int = 1, page_size: int = 10) -> list[dict[st
end = start + page_size
return [self._agent_to_admin_payload(agent) for agent in self.agents[start:end]]

def list_openai_models(self) -> dict[str, Any]:
"""Return an OpenAI-compatible ``/v1/models`` list from the agent pool.

Buyers discover selectable model ids without admin-scope agent pool access.
Each enabled agent model appears once; gateway default
``contextual-orchestrator`` is always first.
"""
created = 1_700_000_000 # stable epoch so list responses are deterministic
data: list[dict[str, Any]] = [
{
"id": "contextual-orchestrator",
"object": "model",
"created": created,
"owned_by": "contextual-orchestrator",
}
]
seen: set[str] = {"contextual-orchestrator"}
for agent in self.agents:
if agent.disabled:
continue
model_id = str(agent.model).strip()
if not model_id or model_id in seen:
continue
seen.add(model_id)
data.append(
{
"id": model_id,
"object": "model",
"created": created,
"owned_by": agent.provider_name
or self._infer_provider_name(agent.base_url)
or "agent_pool",
}
)
return {"object": "list", "data": data}

def get_openai_model(self, model_id: str) -> dict[str, Any]:
"""Return one OpenAI model object or raise ``KeyError`` when unknown."""
wanted = (model_id or "").strip()
if not wanted:
raise KeyError(model_id)
for item in self.list_openai_models()["data"]:
if item["id"] == wanted:
return item
raise KeyError(model_id)

def list_recent_runs(self, page_number: int = 1, page_size: int = 10) -> list[dict[str, Any]]:
"""Return a paginated list of recent workflow run records."""
if page_number < 1 or page_size < 1: # pragma: no cover
Expand Down Expand Up @@ -8517,6 +8635,29 @@ def chat_completion_response(
}


def text_completion_response(
result: dict[str, Any],
model: str = "contextual-orchestrator",
usage: dict[str, int] | None = None,
) -> dict[str, Any]: # pragma: no cover
"""Wrap orchestration output as OpenAI legacy ``text_completion`` (``/v1/completions``)."""
return {
"id": f"cmpl-{int(time.time() * 1000)}",
"object": "text_completion",
"created": int(time.time()),
"model": model,
"choices": [
{
"index": 0,
"text": result["answer"],
"logprobs": None,
"finish_reason": "stop",
}
],
"usage": usage or {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
}


_STREAM_CHUNK_SIZE = 32


Expand Down
Loading
Loading