78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
from dataclasses import dataclass
|
|
from typing import Any, get_origin, get_args
|
|
|
|
from pydantic import BaseModel
|
|
|
|
|
|
@dataclass
|
|
class Usage:
|
|
"""Token usage statistics."""
|
|
|
|
input_tokens: int = 0
|
|
cached_input_tokens: int = 0
|
|
output_tokens: int = 0
|
|
|
|
@property
|
|
def total_tokens(self) -> int:
|
|
"""Total tokens used (input + output)."""
|
|
return self.input_tokens + self.output_tokens
|
|
|
|
def __repr__(self) -> str:
|
|
return (
|
|
f"Usage(input={self.input_tokens}, "
|
|
f"cached={self.cached_input_tokens}, "
|
|
f"output={self.output_tokens}, "
|
|
f"total={self.total_tokens})"
|
|
)
|
|
|
|
|
|
def is_pydantic_model(type_hint: Any) -> bool:
|
|
"""Check if a type hint is a Pydantic model."""
|
|
try:
|
|
return isinstance(type_hint, type) and issubclass(type_hint, BaseModel)
|
|
except TypeError:
|
|
return False
|
|
|
|
|
|
def get_json_schema(pydantic_model: type[BaseModel]) -> dict[str, Any]:
|
|
"""Generate JSON schema from Pydantic model.
|
|
|
|
Sets additionalProperties to false to match Codex behavior.
|
|
"""
|
|
schema = pydantic_model.model_json_schema()
|
|
|
|
# Recursively set additionalProperties: false for all objects
|
|
def set_additional_properties(obj: dict[str, Any]) -> None:
|
|
if isinstance(obj, dict):
|
|
if obj.get("type") == "object":
|
|
obj["additionalProperties"] = False
|
|
for value in obj.values():
|
|
if isinstance(value, dict):
|
|
set_additional_properties(value)
|
|
elif isinstance(value, list):
|
|
for item in value:
|
|
if isinstance(item, dict):
|
|
set_additional_properties(item)
|
|
|
|
set_additional_properties(schema)
|
|
return schema
|
|
|
|
|
|
def parse_json_response(response: str, pydantic_model: type[BaseModel]) -> BaseModel:
|
|
"""Parse JSON response into Pydantic model.
|
|
|
|
Raises:
|
|
json.JSONDecodeError: If response is not valid JSON
|
|
pydantic.ValidationError: If JSON doesn't match model schema
|
|
"""
|
|
return pydantic_model.model_validate_json(response)
|
|
|
|
|
|
def extract_text_from_response(response: str) -> str:
|
|
"""Extract plain text from response.
|
|
|
|
For string outputs, we just return the text as-is.
|
|
Claude Code may wrap responses in markdown or other formatting.
|
|
"""
|
|
return response.strip()
|