new-qwen/serv/llm.py

from __future__ import annotations

import json
from typing import Any, Callable
from urllib import error, request

from config import ServerConfig
from oauth import OAuthError, QwenOAuthManager
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, oauth: QwenOAuthManager, tools: ToolRegistry) -> None:
        self.config = config
        self.oauth = oauth
        self.tools = tools

    def _request_completion(self, messages: list[dict[str, Any]]) -> dict[str, Any]:
        creds = self.oauth.get_valid_credentials()
        base_url = self.oauth.get_openai_base_url(creds)
        payload = {
            "model": self.config.model,
            "messages": messages,
            "tools": self.tools.schemas(),
            "tool_choice": "auto",
        }
        data = json.dumps(payload).encode("utf-8")
        req = request.Request(
            f"{base_url}/chat/completions",
            data=data,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {creds['access_token']}",
            },
            method="POST",
        )
        try:
            with request.urlopen(req, timeout=180) as response:
                return json.loads(response.read().decode("utf-8"))
        except error.HTTPError as exc:
            body = exc.read().decode("utf-8", errors="replace")
            raise OAuthError(f"LLM request failed with HTTP {exc.code}: {body}") from exc

    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,
    ) -> dict[str, Any]:
        emit = on_event or (lambda _event: None)
        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]] = []

        for _ in range(self.config.max_tool_rounds):
            emit({"type": "model_request", "message": "Запрашиваю ответ модели"})
            response = self._request_completion(messages)
            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"),
                    "messages": messages + [{"role": "assistant", "content": content or ""}],
                }

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

            for call in tool_calls:
                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)
                    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:
                    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,
            "messages": messages + [{"role": "assistant", "content": final_message}],
        }