"""Chat request/response DTOs."""
from datetime import datetime
from uuid import UUID
from pydantic import BaseModel, Field
from veupath_chatbot.platform.types import (
JSONObject,
JSONValue,
ModelProvider,
ReasoningEffort,
)
from veupath_chatbot.services.chat.mention_context import MentionType
from veupath_chatbot.transport.http.schemas.optimization import (
OptimizationProgressEventData,
)
[docs]
class ChatMention(BaseModel):
"""A reference to a strategy or experiment included via @-mention."""
type: MentionType
id: str
display_name: str = Field(alias="displayName")
model_config = {"populate_by_name": True}
[docs]
class ChatRequest(BaseModel):
"""Request to send a chat message."""
strategy_id: UUID | None = Field(default=None, alias="strategyId")
site_id: str = Field(alias="siteId")
message: str = Field(min_length=1, max_length=200_000)
# Per-request model overrides (optional; falls back to server defaults).
provider: ModelProvider | None = Field(default=None)
model_id: str | None = Field(default=None, alias="model")
reasoning_effort: ReasoningEffort | None = Field(
default=None, alias="reasoningEffort"
)
# Thesis experiment controls.
disable_rag: bool = Field(default=False, alias="disableRag")
temperature: float | None = Field(default=None)
seed: int | None = Field(default=None)
# Per-model tuning overrides from user settings.
context_size: int | None = Field(default=None, alias="contextSize")
response_tokens: int | None = Field(default=None, alias="responseTokens")
reasoning_budget: int | None = Field(default=None, alias="reasoningBudget")
# Tools the user has disabled in the UI.
disabled_tools: list[str] = Field(default_factory=list, alias="disabledTools")
# @-mention references to strategies and experiments.
mentions: list[ChatMention] = Field(default_factory=list)
model_config = {"populate_by_name": True}
[docs]
class SubKaniTokenUsageResponse(BaseModel):
"""Token usage for a single sub-kani agent."""
prompt_tokens: int = Field(alias="promptTokens")
completion_tokens: int = Field(alias="completionTokens")
llm_call_count: int = Field(default=0, alias="llmCallCount")
estimated_cost_usd: float = Field(default=0.0, alias="estimatedCostUsd")
model_config = {"populate_by_name": True}
[docs]
class SubKaniActivityResponse(BaseModel):
"""Sub-kani tool call activity."""
calls: dict[str, list[ToolCallResponse]]
status: dict[str, str]
models: dict[str, str] | None = None
token_usage: dict[str, SubKaniTokenUsageResponse] | None = Field(
default=None, alias="tokenUsage"
)
model_config = {"populate_by_name": True}
[docs]
class ThinkingResponse(BaseModel):
"""In-progress tool call state."""
tool_calls: list[ToolCallResponse] | None = Field(default=None, alias="toolCalls")
last_tool_calls: list[ToolCallResponse] | None = Field(
default=None, alias="lastToolCalls"
)
sub_kani_calls: dict[str, list[ToolCallResponse]] | None = Field(
default=None, alias="subKaniCalls"
)
sub_kani_status: dict[str, str] | None = Field(default=None, alias="subKaniStatus")
reasoning: str | None = None
updated_at: datetime | None = Field(default=None, alias="updatedAt")
model_config = {"populate_by_name": True}
[docs]
class TokenUsageResponse(BaseModel):
"""Token usage statistics for a message turn."""
prompt_tokens: int = Field(alias="promptTokens")
completion_tokens: int = Field(alias="completionTokens")
total_tokens: int = Field(alias="totalTokens")
cached_tokens: int = Field(default=0, alias="cachedTokens")
tool_call_count: int = Field(alias="toolCallCount")
registered_tool_count: int = Field(alias="registeredToolCount")
llm_call_count: int = Field(default=0, alias="llmCallCount")
sub_kani_prompt_tokens: int = Field(default=0, alias="subKaniPromptTokens")
sub_kani_completion_tokens: int = Field(default=0, alias="subKaniCompletionTokens")
sub_kani_call_count: int = Field(default=0, alias="subKaniCallCount")
estimated_cost_usd: float = Field(default=0.0, alias="estimatedCostUsd")
model_id: str = Field(default="", alias="modelId")
model_config = {"populate_by_name": True}
[docs]
class CitationResponse(BaseModel):
"""Citation from research tools."""
id: str
source: str # "pubmed" | "arxiv" | "google_scholar" | "web"
tag: str | None = None
title: str
url: str | None = None
authors: list[str] | None = None
year: int | None = None
doi: str | None = None
pmid: str | None = None
snippet: str | None = None
accessed_at: str | None = Field(default=None, alias="accessedAt")
model_config = {"populate_by_name": True}
[docs]
class PlanningArtifactResponse(BaseModel):
"""Strategy planning artifact."""
id: str
title: str
summary_markdown: str = Field(alias="summaryMarkdown")
assumptions: list[str] = Field(default_factory=list)
parameters: dict[str, JSONValue] = Field(default_factory=dict)
proposed_strategy_plan: JSONObject | None = Field(
default=None, alias="proposedStrategyPlan"
)
created_at: str = Field(alias="createdAt")
model_config = {"populate_by_name": True}
[docs]
class MessageResponse(BaseModel):
"""Chat message."""
role: str
content: str
model_id: str | None = Field(default=None, alias="modelId")
tool_calls: list[ToolCallResponse] | None = Field(default=None, alias="toolCalls")
sub_kani_activity: SubKaniActivityResponse | None = Field(
default=None, alias="subKaniActivity"
)
citations: list[CitationResponse] | None = None
planning_artifacts: list[PlanningArtifactResponse] | None = Field(
default=None, alias="planningArtifacts"
)
reasoning: str | None = Field(default=None)
optimization_progress: OptimizationProgressEventData | None = Field(
default=None, alias="optimizationProgress"
)
token_usage: TokenUsageResponse | None = Field(default=None, alias="tokenUsage")
timestamp: datetime
model_config = {"populate_by_name": True, "extra": "ignore"}