vet

Mirror of Vet, an AI code review tool
git clone git://git.laack.co/vet.git
Log | Files | Refs | README | LICENSE

message_parser.py (4974B)

      1 from typing import Any
      2 
      3 from vet.imbue_core.agents.agent_api.data_types import AgentAssistantMessage
      4 from vet.imbue_core.agents.agent_api.data_types import AgentContentBlock
      5 from vet.imbue_core.agents.agent_api.data_types import AgentMessage
      6 from vet.imbue_core.agents.agent_api.data_types import AgentResultMessage
      7 from vet.imbue_core.agents.agent_api.data_types import AgentSystemEventType
      8 from vet.imbue_core.agents.agent_api.data_types import AgentSystemMessage
      9 from vet.imbue_core.agents.agent_api.data_types import AgentTextBlock
     10 from vet.imbue_core.agents.agent_api.data_types import AgentThinkingBlock
     11 from vet.imbue_core.agents.agent_api.data_types import AgentToolResultBlock
     12 from vet.imbue_core.agents.agent_api.data_types import AgentToolUseBlock
     13 from vet.imbue_core.agents.agent_api.data_types import AgentUnknownBlock
     14 from vet.imbue_core.agents.agent_api.data_types import AgentUnknownMessage
     15 from vet.imbue_core.agents.agent_api.data_types import AgentUsage
     16 from vet.imbue_core.agents.agent_api.data_types import AgentUserMessage
     17 
     18 
     19 def parse_claude_message(data: dict[str, Any]) -> AgentMessage:
     20     """Parse message from CLI output using unified types.
     21 
     22     Reference:
     23     https://github.com/anthropics/claude-agent-sdk-python/blob/main/src/claude_agent_sdk/_internal/message_parser.py
     24     https://docs.claude.com/en/api/agent-sdk/typescript#sdkmessage
     25     https://docs.claude.com/en/api/agent-sdk/python#message-types
     26     """
     27 
     28     match data["type"]:
     29         case "user":
     30             return AgentUserMessage(content=parse_claude_content_blocks(data), original_message=data)
     31 
     32         case "assistant":
     33             return AgentAssistantMessage(content=parse_claude_content_blocks(data), original_message=data)
     34 
     35         case "system":
     36             # Normalize system event types
     37             event_type = parse_claude_system_event_type(data.get("subtype", ""))
     38             return AgentSystemMessage(
     39                 event_type=event_type,
     40                 session_id=data.get("session_id"),
     41                 error=data.get("error"),
     42                 original_message=data,
     43             )
     44 
     45         case "result":
     46             # Build normalized usage
     47             usage = None
     48             raw_usage = data.get("usage")
     49             if raw_usage or data.get("total_cost_usd"):
     50                 usage = AgentUsage(
     51                     input_tokens=raw_usage.get("input_tokens") if raw_usage else None,
     52                     output_tokens=raw_usage.get("output_tokens") if raw_usage else None,
     53                     cached_tokens=(raw_usage.get("cache_read_input_tokens") if raw_usage else None),
     54                     total_tokens=(
     55                         raw_usage.get("input_tokens", 0) + raw_usage.get("output_tokens", 0) if raw_usage else None
     56                     ),
     57                     total_cost_usd=data.get("total_cost_usd"),
     58                 )
     59 
     60             return AgentResultMessage(
     61                 session_id=data["session_id"],
     62                 is_error=data["is_error"],
     63                 duration_ms=data.get("duration_ms"),
     64                 api_duration_ms=data.get("duration_api_ms"),
     65                 num_turns=data.get("num_turns"),
     66                 usage=usage,
     67                 result=data.get("result"),
     68                 error=data.get("error") if data["is_error"] else None,
     69                 original_message=data,
     70             )
     71 
     72         case _:
     73             return AgentUnknownMessage(raw=data, original_message=data)
     74 
     75 
     76 def parse_claude_system_event_type(subtype: str) -> AgentSystemEventType:
     77     """Parse Claude system event subtype to unified event type."""
     78     subtype_lower = subtype.lower()
     79 
     80     # TODO add other system event types as we find them
     81     #  basically the documentattion doesn't mention any other system event types
     82     #  other than init AFAIKT
     83     if "init" in subtype_lower:
     84         return AgentSystemEventType.SESSION_STARTED
     85     else:
     86         return AgentSystemEventType.OTHER
     87 
     88 
     89 def parse_claude_content_blocks(data: dict[str, Any]) -> list[AgentContentBlock]:
     90     return [parse_claude_content_block(block) for block in data["message"]["content"]]
     91 
     92 
     93 def parse_claude_content_block(block: dict[str, Any]) -> AgentContentBlock:
     94     """Parse content block from CLI output using unified types."""
     95 
     96     match block["type"]:
     97         case "text":
     98             return AgentTextBlock(text=block["text"])
     99 
    100         case "thinking":
    101             # Claude Code thinking blocks
    102             return AgentThinkingBlock(
    103                 content=block.get("thinking", ""),
    104                 thinking_tokens=block.get("thinking_tokens"),
    105             )
    106 
    107         case "tool_use":
    108             return AgentToolUseBlock(id=block["id"], name=block["name"], input=block["input"])
    109 
    110         case "tool_result":
    111             return AgentToolResultBlock(
    112                 tool_use_id=block["tool_use_id"],
    113                 content=block.get("content"),
    114                 is_error=block.get("is_error"),
    115             )
    116 
    117         case _:
    118             return AgentUnknownBlock(raw=block)