124 lines
3.8 KiB
Python
124 lines
3.8 KiB
Python
"""Patch OpenAI chat completion calls in simulation scripts for centralized cost logging."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
from typing import Any, Dict, Optional
|
|
|
|
from app.utils.llm_cost import record_llm_cost
|
|
|
|
|
|
_PATCHED = False
|
|
_PATCH_CONTEXT: Dict[str, Any] = {}
|
|
|
|
|
|
def _build_metadata(model: str) -> Dict[str, Any]:
|
|
metadata = copy.deepcopy(_PATCH_CONTEXT)
|
|
metadata.setdefault("component", "scripts.simulation")
|
|
metadata.setdefault("phase", "simulation_run")
|
|
metadata.setdefault("model", model)
|
|
return metadata
|
|
|
|
|
|
def _normalize_messages_for_qwen(messages: Any) -> Any:
|
|
"""Ensure Qwen-compatible ordering: a single system message at index 0."""
|
|
if not isinstance(messages, list):
|
|
return messages
|
|
|
|
if not messages:
|
|
return messages
|
|
|
|
system_contents = []
|
|
non_system_messages = []
|
|
|
|
for item in messages:
|
|
if isinstance(item, dict) and str(item.get("role", "")).lower() == "system":
|
|
content = item.get("content")
|
|
if content is not None and str(content).strip():
|
|
system_contents.append(str(content))
|
|
continue
|
|
non_system_messages.append(item)
|
|
|
|
# No system messages detected, keep original payload.
|
|
if not system_contents:
|
|
return messages
|
|
|
|
merged_system = {
|
|
"role": "system",
|
|
"content": "\n\n".join(system_contents),
|
|
}
|
|
return [merged_system] + non_system_messages
|
|
|
|
|
|
def install_openai_cost_patch(
|
|
*,
|
|
simulation_id: Optional[str],
|
|
project_id: Optional[str],
|
|
platform: str,
|
|
component: str,
|
|
phase: str = "simulation_run",
|
|
) -> None:
|
|
"""
|
|
Install monkey-patch for OpenAI SDK calls used indirectly by CAMEL/OASIS.
|
|
|
|
The patch is process-scoped and idempotent.
|
|
"""
|
|
global _PATCHED
|
|
global _PATCH_CONTEXT
|
|
|
|
_PATCH_CONTEXT = {
|
|
"simulation_id": simulation_id,
|
|
"project_id": project_id,
|
|
"platform": platform,
|
|
"component": component,
|
|
"phase": phase,
|
|
}
|
|
|
|
if _PATCHED:
|
|
return
|
|
|
|
try:
|
|
from openai.resources.chat.completions.completions import Completions, AsyncCompletions
|
|
except Exception as exc: # pragma: no cover
|
|
print(f"[llm_cost_patch] Failed to import OpenAI completion classes: {exc}")
|
|
return
|
|
|
|
original_sync_create = Completions.create
|
|
original_async_create = AsyncCompletions.create
|
|
|
|
def sync_create_wrapper(self, *args, **kwargs):
|
|
model = kwargs.get("model", "unknown_model")
|
|
if "qwen" in str(model).lower() and "messages" in kwargs:
|
|
kwargs = dict(kwargs)
|
|
kwargs["messages"] = _normalize_messages_for_qwen(kwargs.get("messages"))
|
|
response = original_sync_create(self, *args, **kwargs)
|
|
try:
|
|
record_llm_cost(
|
|
model_name=model,
|
|
usage=getattr(response, "usage", None),
|
|
metadata=_build_metadata(model),
|
|
)
|
|
except Exception as exc:
|
|
print(f"[llm_cost_patch] Failed to record sync cost: {exc}")
|
|
return response
|
|
|
|
async def async_create_wrapper(self, *args, **kwargs):
|
|
model = kwargs.get("model", "unknown_model")
|
|
if "qwen" in str(model).lower() and "messages" in kwargs:
|
|
kwargs = dict(kwargs)
|
|
kwargs["messages"] = _normalize_messages_for_qwen(kwargs.get("messages"))
|
|
response = await original_async_create(self, *args, **kwargs)
|
|
try:
|
|
record_llm_cost(
|
|
model_name=model,
|
|
usage=getattr(response, "usage", None),
|
|
metadata=_build_metadata(model),
|
|
)
|
|
except Exception as exc:
|
|
print(f"[llm_cost_patch] Failed to record async cost: {exc}")
|
|
return response
|
|
|
|
Completions.create = sync_create_wrapper
|
|
AsyncCompletions.create = async_create_wrapper
|
|
_PATCHED = True
|