adapter.py

import json
import uuid
import time
import asyncio
import httpx
from typing import Optional, AsyncGenerator
from dotenv import load_dotenv
import os

load_dotenv()

from sign import sign_request, client_metadata
from tool_dsml import build_dsml_tool_prompt, has_dsml_content, parse_dsml_invoke, strip_dsml_tags
from tool_sieve import StreamSieve, SieveEvent


class MiniMaxAdapter:
    def __init__(self):
        self.base_url = os.getenv("TARGET_URL", "https://agent.minimaxi.com").rstrip("/")
        self.stream_base_url = os.getenv("STREAM_URL", "https://agent-stream.minimaxi.com").rstrip("/")
        self.agent_id = os.getenv("AGENT_ID", "")
        self.token = os.getenv("TOKEN", "")

        self.browser_extra = {}
        for k in ("uuid", "device_id", "os_name", "browser_name",
                  "device_memory", "cpu_core_num", "browser_language", "browser_platform"):
            v = os.getenv(k.upper(), "")
            if v:
                self.browser_extra[k] = v

        self._user_id = int(os.getenv("USER_ID", "0"))
        self._screen_width = int(os.getenv("SCREEN_WIDTH", "0"))
        self._screen_height = int(os.getenv("SCREEN_HEIGHT", "0"))
        self._session_id = None
        self._http = httpx.AsyncClient(timeout=60, follow_redirects=True)

        provider = os.getenv("MODEL_PROVIDER", "minimax")
        model_id = os.getenv("MODEL_ID", "MiniMax-M3")
        model_variant = os.getenv("MODEL_VARIANT", "thinking")
        self.model_config = {
            "provider_id": provider,
            "model_id": model_id,
            "variant": model_variant,
        }
        self.use_worktree = os.getenv("USE_WORKTREE", "false").lower() == "true"
        self._dsml_prefix = ""

    def _build_all(self, body_str: str, url_path: str) -> tuple:
        ts = int(time.time())
        sig = sign_request(
            body_str=body_str, token=self.token, url_path=url_path, timestamp_sec=ts,
            uuid=self.browser_extra.get("uuid", ""),
            device_id=self.browser_extra.get("device_id", ""),
            user_id=self._user_id,
            screen_width=self._screen_width, screen_height=self._screen_height,
            device_memory=int(self.browser_extra.get("device_memory", "0")),
            cpu_core_num=int(self.browser_extra.get("cpu_core_num", "0")),
            browser_language=self.browser_extra.get("browser_language", ""),
            browser_platform=self.browser_extra.get("browser_platform", ""),
        )
        import urllib.parse
        meta = client_metadata(
            timestamp_ms=ts * 1000, token=self.token,
            uuid=self.browser_extra.get("uuid", ""),
            device_id=self.browser_extra.get("device_id", ""),
            user_id=self._user_id,
            screen_width=self._screen_width, screen_height=self._screen_height,
        )
        for k, v in self.browser_extra.items():
            if k not in meta:
                meta[k] = v
        from sign import js_string
        qs = urllib.parse.urlencode({k: js_string(v) for k, v in meta.items()})
        full_url = f"{self.base_url}{url_path}?{qs}&token={self.token}"
        headers = {
            "Content-Type": "application/json", "Accept": "*/*",
            "Origin": self.base_url, "Referer": f"{self.base_url}/",
            "User-Agent": os.getenv("USER_AGENT", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"),
            "token": self.token,
            "x-timestamp": sig["x-timestamp"], "x-signature": sig["x-signature"],
            "yy": sig["yy"],
        }
        return full_url, headers

    async def ensure_session(self):
        if self._session_id:
            return self._session_id
        if not self.agent_id:
            raise RuntimeError("AGENT_ID is not configured")
        sess_path = f"/archon/api/v1/agent/{self.agent_id}/session"
        body = json.dumps({
            "team_mode_off": True,
            "model": f"{self.model_config['provider_id']}/{self.model_config['model_id']}",
        }, separators=(",", ":"), ensure_ascii=False)
        url, headers = self._build_all(body, sess_path)
        resp = await self._http.post(url, content=body, headers=headers)
        if resp.status_code != 200:
            raise RuntimeError(f"Session creation failed: {resp.status_code} {resp.text[:300]}")
        data = resp.json()
        if data.get("base_resp", {}).get("status_code") != 0:
            raise RuntimeError(f"Session creation error: {data}")
        self._session_id = data.get("session_id")
        if not self._session_id:
            raise RuntimeError(f"No session_id in response: {data}")
        return self._session_id

    async def _poll_session_idle(self, session_id: str, max_retries: int = 5) -> bool:
        poll_headers = {"token": self.token, "Accept": "*/*",
                        "Origin": self.base_url, "Referer": f"{self.base_url}/"}
        for _ in range(max_retries):
            await asyncio.sleep(1)
            url = f"{self.base_url}/archon/api/v1/session/{session_id}?token={self.token}"
            resp = await self._http.get(url, headers=poll_headers)
            if resp.status_code == 200:
                try:
                    st = resp.json().get("session", {}).get("status", {}).get("type")
                    if st == 0:
                        return True
                except:
                    pass
        return False

    def _inject_dsml_prompt(self, messages: list, tools: list, tool_choice: str | None = None):
        dsml_prompt = build_dsml_tool_prompt(tools, tool_choice)
        if dsml_prompt:
            self._dsml_prefix = dsml_prompt

    def _build_msg_body(self, messages: list) -> str:
        last_content = ""
        if messages:
            last = messages[-1]
            if isinstance(last, dict):
                content = last.get("content", "")
                if isinstance(content, list):
                    texts = [p.get("text", "") for p in content
                             if isinstance(p, dict) and p.get("type") == "text"]
                    last_content = " ".join(texts)
                elif isinstance(content, str):
                    last_content = content
        # Prepend system + DSML context to the user message
        ctx_parts = []
        for m in messages[:-1]:
            r = m.get("role", "")
            c = m.get("content", "")
            if c:
                if r == "system":
                    ctx_parts.append(f"系统指令: {c}")
                elif r == "user":
                    ctx_parts.append(f"用户: {c}")
                elif r == "assistant":
                    ctx_parts.append(f"助手: {c}")
        if self._dsml_prefix:
            ctx_parts.append(self._dsml_prefix)
        prefix = "\n\n".join(ctx_parts)
        if prefix:
            last_content = prefix + "\n\n用户说: " + last_content
        return json.dumps({
            "content": last_content,
            "model": self.model_config,
            "turn_id": str(uuid.uuid4()),
            "worktreeMode": self.use_worktree,
        }, separators=(",", ":"), ensure_ascii=False)

    def _extract_content_from_chunk(self, data: dict) -> Optional[str]:
        if not isinstance(data, dict):
            return None
        if data.get("type") == 6:
            chunk = data.get("agent_message_chunk", {})
            if isinstance(chunk, dict):
                return chunk.get("msg_content")
        return None

    def _is_finish_chunk(self, data: dict) -> bool:
        if not isinstance(data, dict):
            return False
        ev_type = data.get("type")
        if ev_type == 6:
            chunk = data.get("agent_message_chunk", {})
            if isinstance(chunk, dict) and chunk.get("finish") is True:
                return True
        if ev_type == 2:
            msg = data.get("agent_message", {})
            if isinstance(msg, dict) and msg.get("finish_reason"):
                return True
        return False

    async def _send_message(self, messages: list) -> str:
        session_id = await self.ensure_session()
        await self._poll_session_idle(session_id)
        msg_body = self._build_msg_body(messages)
        msg_path = f"/archon/api/v1/session/{session_id}/message"
        url, headers = self._build_all(msg_body, msg_path)
        url = url.replace(self.base_url, self.stream_base_url)
        resp = await self._http.post(url, content=msg_body, headers=headers)
        if resp.status_code == 409:
            self._session_id = None
            raise RuntimeError(f"Session busy, retry. {resp.text[:200]}")
        if resp.status_code != 200:
            raise RuntimeError(f"Message failed: {resp.status_code} {resp.text[:300]}")
        full_content = ""
        fallback_content = ""
        for line in resp.text.split("\n"):
            line = line.strip()
            if not line or not line.startswith("data:"):
                continue
            try:
                ev_data = json.loads(line[5:])
            except json.JSONDecodeError:
                continue
            content = self._extract_content_from_chunk(ev_data)
            if content:
                full_content += content
            if ev_data.get("type") == 2:
                msg = ev_data.get("agent_message", {})
                if isinstance(msg, dict) and msg.get("role") == "assistant":
                    fallback_content = msg.get("msg_content", "")
            if self._is_finish_chunk(ev_data):
                break
        if not full_content:
            full_content = fallback_content
        return full_content

    async def stream_chat(self, messages: list, tools: list | None = None,
                          tool_choice: str | None = None, **kwargs) -> AsyncGenerator[dict, None]:
        try:
            if tools:
                self._inject_dsml_prompt(messages, tools, tool_choice)
            raw = await self._send_message(messages)

            sieve = StreamSieve()
            events = sieve.feed(raw)
            events.extend(sieve.flush())

            text_parts = []
            tool_calls_result = None
            for ev in events:
                if ev.type == "text":
                    text_parts.append(ev.data)
                elif ev.type == "tool_calls":
                    tool_calls_result = ev.data

            if tool_calls_result:
                yield {
                    "choices": [{
                        "delta": {"role": "assistant", "content": None},
                        "index": 0,
                    }]
                }
                for tc in tool_calls_result:
                    yield {
                        "choices": [{
                            "delta": {
                                "tool_calls": [{
                                    "index": 0,
                                    "id": tc["id"],
                                    "type": "function",
                                    "function": {
                                        "name": tc["function"]["name"],
                                        "arguments": tc["function"]["arguments"],
                                    },
                                }]
                            },
                            "index": 0,
                        }]
                    }
                yield {
                    "choices": [{"delta": {}, "finish_reason": "tool_calls", "index": 0}]
                }
            else:
                text = "".join(text_parts)
                if not text:
                    text = raw
                text = strip_dsml_tags(text)
                yield {"choices": [{"delta": {"role": "assistant"}, "index": 0}]}
                if text:
                    yield {"choices": [{"delta": {"content": text}, "index": 0}]}
                yield {"choices": [{"delta": {}, "finish_reason": "stop", "index": 0}]}

        except Exception as e:
            yield {"choices": [{"delta": {"content": f"[ERROR] {e}"}, "index": 0}]}
            yield {"choices": [{"delta": {}, "finish_reason": "stop", "index": 0}]}

    async def non_stream_chat(self, messages: list, tools: list | None = None,
                              tool_choice: str | None = None, **kwargs) -> dict:
        try:
            if tools:
                self._inject_dsml_prompt(messages, tools, tool_choice)
            raw = await self._send_message(messages)

            sieve = StreamSieve()
            events = sieve.feed(raw)
            events.extend(sieve.flush())

            text_parts = []
            tool_calls_result = None
            for ev in events:
                if ev.type == "text":
                    text_parts.append(ev.data)
                elif ev.type == "tool_calls":
                    tool_calls_result = ev.data

            if tool_calls_result:
                return {
                    "choices": [{
                        "message": {
                            "role": "assistant",
                            "content": None,
                            "tool_calls": [{
                                "id": tc["id"],
                                "type": "function",
                                "function": {
                                    "name": tc["function"]["name"],
                                    "arguments": tc["function"]["arguments"],
                                },
                            } for tc in tool_calls_result],
                        },
                        "finish_reason": "tool_calls",
                    }]
                }
            else:
                text = "".join(text_parts)
                if not text:
                    text = raw
                text = strip_dsml_tags(text)
                return {"choices": [{"message": {"role": "assistant", "content": text}}]}

        except Exception as e:
            return {"choices": [{"message": {"role": "assistant", "content": f"[ERROR] {e}"}}]}

    async def close(self):
        await self._http.aclose()