mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-09-18 07:37:40 +02:00
95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
"""Main agent loop — orchestrates the Claude Agent SDK query loop.
|
||
|
||
Heavy logic is delegated to sibling modules:
|
||
- agent_mock – mock-agent fallback, streaming helpers, session analytics
|
||
- agent_hooks – SDK hook factories (approval, permissions, post-tool)
|
||
- agent_options – MCP server construction & ClaudeAgentOptions building
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
import logging
|
||
from typeguard import typechecked
|
||
|
||
from claude_agent_sdk import (
|
||
query, ClaudeAgentOptions, AssistantMessage,
|
||
)
|
||
from claude_agent_sdk.types import StreamEvent
|
||
from backend.apps.agents.manager.HaikFix.helpers.handle_stream_event import handle_stream_event
|
||
from backend.apps.agents.manager.HaikFix.helpers.handle_assistant_message import handle_assistant_message
|
||
|
||
from backend.apps.agents.manager.HaikFix.PromptChunks import ImageChunk, ImageChunkDict, TextChunk, TextChunkDict
|
||
from typing import List, Dict, Literal, Union, Optional
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@typechecked
|
||
def build_image_prompt_content(prompt: str, images: List[ImageChunk]) -> List[TextChunk | ImageChunk]:
|
||
content: List[Union[ImageChunkDict, TextChunkDict]] = [TextChunk(text=prompt).to_dict()]
|
||
for img in images:
|
||
content.append(img.to_dict())
|
||
return content
|
||
|
||
PromptMsgDict = Dict[
|
||
Literal["type", "message"],
|
||
Dict[
|
||
Literal["role", "content"],
|
||
List[
|
||
Union[ImageChunkDict, TextChunkDict]
|
||
]
|
||
]
|
||
]
|
||
|
||
@typechecked
|
||
def build_prompt_msg(prompt: str, images: Optional[List[ImageChunk]]) -> PromptMsgDict:
|
||
content = build_image_prompt_content(prompt, images)
|
||
return {
|
||
"type": "user",
|
||
"message": {
|
||
"role": "user",
|
||
"content": content
|
||
}
|
||
}
|
||
|
||
@typechecked
|
||
async def run_agent_loop(
|
||
prompt: str,
|
||
images: Optional[List[ImageChunk]] = None,
|
||
options: ClaudeAgentOptions | None = None,
|
||
branch_id: str | None = None,
|
||
):
|
||
"""Run the Claude Agent SDK query loop for a session."""
|
||
|
||
prompt_msg = build_prompt_msg(prompt=prompt, images=images)
|
||
|
||
async def prompt_stream():
|
||
yield prompt_msg
|
||
|
||
stream_text_msg_id: Optional[str] = None
|
||
stream_tool_msg_ids_ordered: List[str] = []
|
||
stream_block_index_map: Dict[int, str] = {}
|
||
_turn_number: int = 0
|
||
_first_event: bool = True
|
||
|
||
async for message in query(prompt=prompt_stream(), options=options):
|
||
|
||
if isinstance(message, StreamEvent):
|
||
stream_text_msg_id = await handle_stream_event(
|
||
session_id=options.session_id,
|
||
event=message.event,
|
||
stream_text_msg_id=stream_text_msg_id,
|
||
stream_tool_ids=stream_tool_msg_ids_ordered,
|
||
block_map=stream_block_index_map,
|
||
)
|
||
|
||
elif isinstance(message, AssistantMessage):
|
||
stream_text_msg_id, stream_tool_msg_ids_ordered, stream_block_index_map = (
|
||
await handle_assistant_message(
|
||
session_id=options.session_id,
|
||
branch_id=branch_id,
|
||
message=message,
|
||
stream_text_msg_id=stream_text_msg_id,
|
||
stream_tool_ids=stream_tool_msg_ids_ordered,
|
||
)
|
||
)
|
||
_turn_number += 1 |