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
17 changes: 12 additions & 5 deletions contextual_orchestrator/cost_ledger.py
Original file line number Diff line number Diff line change
Expand Up @@ -222,14 +222,21 @@ class UsageRecord:

def as_dict(self) -> Dict[str, Any]:
"""Flatten the record (attribution inlined) for JSON + SQL storage."""
# Prefer an explicit attribution model_name (client tag) for rollups;
# otherwise the served model id on the record.
rollup_model = (
self.attribution.model_name
if self.attribution.model_name != UNATTRIBUTED
else self.model_name
)
row = {
"usage_record_id": self.usage_record_id,
"created_at": self.created_at,
"workflow_run_id": self.workflow_run_id,
"request_channel": self.request_channel,
"route_mode": self.route_mode,
"provider_name": self.provider_name,
"model_name": self.model_name,
"model_name": rollup_model,
"prompt_tokens": self.prompt_tokens,
"completion_tokens": self.completion_tokens,
"total_tokens": self.total_tokens,
Expand Down Expand Up @@ -583,12 +590,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 +609,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 +629,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
81 changes: 75 additions & 6 deletions contextual_orchestrator/orchestrator.py
Original file line number Diff line number Diff line change
Expand Up @@ -215,6 +215,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 +234,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 +250,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 +286,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 +337,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 @@ -897,7 +927,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 @@ -8517,6 +8563,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