new-qwen/serv/llm.py

from __future__ import annotations

import json
from typing import Any, Callable

from config import ServerConfig
from model_router import CompletionRequest, ModelRouter
from tools import ToolError, ToolRegistry


DEFAULT_SYSTEM_PROMPT = """You are new-qwen-serv, a server-side coding agent.
You help the user through a remote client.
Use tools when they are necessary.
Be concise, precise and action-oriented.
When you modify files, explain what changed.
Do not claim to have executed tools unless a tool result confirms it."""


class QwenAgent:
    def __init__(self, config: ServerConfig, router: ModelRouter, tools: ToolRegistry) -> None:
        self.config = config
        self.router = router
        self.tools = tools

    def run(
        self,
        history: list[dict[str, Any]],
        user_message: str,
        on_event: Callable[[dict[str, Any]], None] | None = None,
        approval_callback: Callable[[str, dict[str, Any]], dict[str, Any]] | None = None,
        is_cancelled: Callable[[], bool] | None = None,
        preferred_provider: str | None = None,
    ) -> dict[str, Any]:
        emit = on_event or (lambda _event: None)
        cancel_check = is_cancelled or (lambda: False)

        def ensure_not_cancelled() -> None:
            if cancel_check():
                raise ToolError("Job canceled by operator")

        system_prompt = self.config.system_prompt or DEFAULT_SYSTEM_PROMPT
        messages: list[dict[str, Any]] = [{"role": "system", "content": system_prompt}]
        messages.extend(history)
        messages.append({"role": "user", "content": user_message})
        events: list[dict[str, Any]] = []
        selected_provider: str | None = None
        selected_model: str | None = None

        for _ in range(self.config.max_tool_rounds):
            ensure_not_cancelled()
            completion = self.router.complete(
                CompletionRequest(
                    messages=messages,
                    tools=self.tools.schemas(),
                    tool_choice="auto",
                    preferred_provider=preferred_provider,
                    require_tools=True,
                )
            )
            emit(
                {
                    "type": "model_request",
                    "message": "Запрашиваю ответ модели",
                    "provider": completion.provider_name,
                    "model": completion.model_name,
                    "selection_reason": completion.selection_reason,
                    "attempted": completion.attempted,
                }
            )
            selected_provider = completion.provider_name
            selected_model = completion.model_name
            response = completion.payload
            ensure_not_cancelled()
            choice = response["choices"][0]["message"]
            tool_calls = choice.get("tool_calls") or []
            content = choice.get("content")
            if content:
                assistant_event = {"type": "assistant", "content": content}
                events.append(assistant_event)
                emit(assistant_event)

            if not tool_calls:
                return {
                    "answer": content or "",
                    "events": events,
                    "usage": response.get("usage"),
                    "provider": selected_provider,
                    "model": selected_model,
                    "messages": messages + [{"role": "assistant", "content": content or ""}],
                }

            messages.append(
                {
                    "role": "assistant",
                    "content": content or "",
                    "tool_calls": tool_calls,
                }
            )

            for call in tool_calls:
                ensure_not_cancelled()
                tool_name = call["function"]["name"]
                try:
                    arguments = json.loads(call["function"]["arguments"] or "{}")
                except json.JSONDecodeError:
                    arguments = {}

                tool_call_event = {"type": "tool_call", "name": tool_name, "arguments": arguments}
                events.append(tool_call_event)
                emit(tool_call_event)
                if approval_callback and self.tools.requires_approval(tool_name):
                    approval_result = approval_callback(tool_name, arguments)
                    approval_event = {
                        "type": "approval_result",
                        "tool_name": tool_name,
                        "approval_id": approval_result["approval_id"],
                        "status": approval_result["status"],
                        "actor": approval_result.get("actor"),
                    }
                    events.append(approval_event)
                    emit(approval_event)
                    ensure_not_cancelled()
                    if approval_result["status"] != "approved":
                        result = {"error": f"Tool '{tool_name}' was rejected by operator"}
                        tool_result_event = {"type": "tool_result", "name": tool_name, "result": result}
                        events.append(tool_result_event)
                        emit(tool_result_event)
                        messages.append(
                            {
                                "role": "tool",
                                "tool_call_id": call["id"],
                                "content": self.tools.encode_result(result),
                            }
                        )
                        continue
                try:
                    ensure_not_cancelled()
                    result = self.tools.execute(tool_name, arguments)
                except ToolError as exc:
                    result = {"error": str(exc)}
                except Exception as exc:
                    result = {"error": f"Unexpected tool failure: {exc}"}

                tool_result_event = {"type": "tool_result", "name": tool_name, "result": result}
                events.append(tool_result_event)
                emit(tool_result_event)
                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": call["id"],
                        "content": self.tools.encode_result(result),
                    }
                )

        final_message = (
            "Остановлено по лимиту tool rounds. Попробуй сузить задачу или продолжить отдельным сообщением."
        )
        final_event = {"type": "assistant", "content": final_message}
        events.append(final_event)
        emit(final_event)
        return {
            "answer": final_message,
            "events": events,
            "usage": None,
            "provider": selected_provider,
            "model": selected_model,
            "messages": messages + [{"role": "assistant", "content": final_message}],
        }