src/agent_server/llm/agent.py - master - man4j/agent-server

src/agent_server/llm/agent.py
220 строк6 KB
Vladimir
initial
18 апр 2026, 19:36
18 апр 2026, 19:36a95f341
import datetime
import decimal
import json
import uuid
from typing import Any, Awaitable, Callable

from agent_server.chat.types import ChatMessage, ToolCall, ToolExecutionRecord, ToolResultPayload
from agent_server.llm.completion import (
    LlmRoundResult,
    build_chat_completion_kwargs,
    build_summary_completion_kwargs,
)
from agent_server.config.profiles import LlmProfile

import chainlit as cl
from openai import OpenAI


def ensure_tool_call(acc: dict[int, ToolCall], idx: int) -> ToolCall:
    if idx not in acc:
        acc[idx] = {
            "id": None,
            "type": "function",
            "function": {
                "name": None,
                "arguments": "",
            },
        }
    return acc[idx]


def finalize_tool_calls(acc: dict[int, ToolCall]) -> list[ToolCall]:
    return [tc for _, tc in sorted(acc.items(), key=lambda x: x[0])]


def json_default(o):
    if isinstance(o, (datetime.datetime, datetime.date, datetime.time)):
        return o.isoformat()
    if isinstance(o, decimal.Decimal):
        return float(o)
    if isinstance(o, uuid.UUID):
        return str(o)
    return str(o)


async def one_llm_round(
    *,
    client: OpenAI,
    profile: LlmProfile,
    messages: list[ChatMessage],
    tools: list[dict[str, Any]],
    thinking_step: cl.Step,
    answer_msg: cl.Message,
):
    tool_calls_acc: dict[int, ToolCall] = {}
    assistant_content = ""
    reasoning_buf = ""
    usage = None
    finish_reason = None

    stream = client.chat.completions.create(
        **build_chat_completion_kwargs(
            profile=profile,
            messages=messages,
            stream=True,
            tools=tools,
        )
    )

    for chunk in stream:
        if not chunk.choices:
            if getattr(chunk, "usage", None):
                usage = chunk.usage
            continue

        choice = chunk.choices[0]
        delta = choice.delta
        finish_reason = choice.finish_reason or finish_reason

        rc = getattr(delta, "reasoning_content", None)
        if rc:
            reasoning_buf += rc
            await thinking_step.stream_token(rc)

        cc = getattr(delta, "content", None)
        if cc:
            assistant_content += cc
            await answer_msg.stream_token(cc)

        tcs = getattr(delta, "tool_calls", None)
        if tcs:
            for t in tcs:
                idx = getattr(t, "index", 0)
                acc = ensure_tool_call(tool_calls_acc, idx)

                if getattr(t, "id", None):
                    acc["id"] = t.id
                if getattr(t, "type", None):
                    acc["type"] = t.type

                fn = getattr(t, "function", None)
                if fn:
                    if getattr(fn, "name", None):
                        acc["function"]["name"] = fn.name
                    args_part = getattr(fn, "arguments", None)
                    if args_part:
                        acc["function"]["arguments"] += args_part

        if getattr(chunk, "usage", None):
            usage = chunk.usage

    return LlmRoundResult(
        assistant_content=assistant_content,
        reasoning_content=reasoning_buf,
        tool_calls=finalize_tool_calls(tool_calls_acc),
        usage=usage,
        finish_reason=finish_reason,
    )


async def handle_tool_calls(
    *,
    tool_calls: list[ToolCall],
    call_tool: Callable[[str, dict[str, Any]], Awaitable[ToolResultPayload]],
) -> tuple[list[ChatMessage], list[ToolExecutionRecord]]:
    tool_messages: list[ChatMessage] = []
    tool_results: list[ToolExecutionRecord] = []

    for tc in tool_calls:
        fn = tc.get("function", {})
        name = fn.get("name")
        raw_args = fn.get("arguments", "") or ""
        tool_call_id = tc.get("id")

        async with cl.Step(name=f"Tool: {name}", type="tool") as tool_step:
            tool_step.input = raw_args

            try:
                tool_input = json.loads(raw_args) if raw_args.strip() else {}
            except Exception as e:
                result: ToolResultPayload = {"error": f"Invalid JSON arguments: {type(e).__name__}: {e}"}
                content = json.dumps(result, ensure_ascii=False, default=json_default)
                tool_step.output = content
                await tool_step.update()

                tool_messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_call_id,
                        "content": content,
                    }
                )
                tool_results.append(
                    {
                        "tool_name": name,
                        "tool_call_id": tool_call_id,
                        "tool_input": None,
                        "result": result,
                    }
                )
                continue

            try:
                result = await call_tool(name, tool_input)
            except Exception as e:
                result = {"error": f"{type(e).__name__}: {e}"}

            content = build_tool_message_content(result)

            tool_step.output = content
            await tool_step.update()

        tool_messages.append(
            {
                "role": "tool",
                "tool_call_id": tool_call_id,
                "content": content,
            }
        )
        tool_results.append(
            {
                "tool_name": name,
                "tool_call_id": tool_call_id,
                "tool_input": tool_input,
                "result": result,
            }
        )

    return tool_messages, tool_results


async def summarize_messages(
    *,
    client: OpenAI,
    profile: LlmProfile,
    messages: list[ChatMessage],
) -> str:
    response = client.chat.completions.create(
        **build_summary_completion_kwargs(
            profile=profile,
            messages=messages,
        )
    )
    if not response.choices:
        return ""

    message = response.choices[0].message
    return (getattr(message, "content", None) or "").strip()


def build_tool_message_content(result: ToolResultPayload) -> str:
    if isinstance(result, dict) and result.get("model_visible") is False:
        safe_payload = {
            "ok": bool(result.get("ok")),
            "ui_type": result.get("ui_type"),
            "message": result.get("model_message") or "UI result was rendered for the user.",
        }
        return json.dumps(safe_payload, ensure_ascii=False, default=json_default)

    return json.dumps(result, ensure_ascii=False, default=json_default)