PhoneWork/orchestrator/agent.py

"""LangChain orchestration agent backed by ZhipuAI (OpenAI-compatible API).

Uses LangChain 1.x tool-calling pattern: bind_tools + manual agentic loop.
"""

from __future__ import annotations

import asyncio
import json
import logging
import re
from collections import defaultdict
from typing import Dict, List, Optional

from langchain_core.messages import (
    AIMessage,
    BaseMessage,
    HumanMessage,
    SystemMessage,
    ToolMessage,
)
from langchain_openai import ChatOpenAI

from agent.manager import manager
from config import OPENAI_API_KEY, OPENAI_BASE_URL, OPENAI_MODEL, WORKING_DIR
from orchestrator.tools import TOOLS, set_current_user

logger = logging.getLogger(__name__)

SYSTEM_PROMPT_TEMPLATE = """You are PhoneWork, an AI assistant that helps users control Claude Code \
from their phone via Feishu (飞书).

You manage Claude Code sessions. Each session has a conv_id and runs in a project directory.

Base working directory: {working_dir}
Users refer to projects by subfolder name (e.g. "todo_app") or relative path. \
Pass these names directly to `create_conversation` — the tool resolves them automatically.

{active_session_line}

Your responsibilities:
1. NEW session: call `create_conversation` with the project name/path. \
   If the user's message also contains a task, pass it as `initial_message` too.
2. Follow-up to ACTIVE session: call `send_to_conversation` with the active conv_id shown above.
3. List sessions: call `list_conversations`.
4. Close session: call `close_conversation`.

Guidelines:
- Relay Claude Code's output verbatim.
- If no active session and the user sends a task without naming a directory, ask them which project.
- Keep your own words brief — let Claude Code's output speak.
- Reply in the same language the user uses (Chinese or English).
"""

MAX_ITERATIONS = 10
_TOOL_MAP = {t.name: t for t in TOOLS}

COMMAND_PATTERN = re.compile(r"^/(new|list|close|switch|retry|help)", re.IGNORECASE)


def _looks_like_command(text: str) -> bool:
    return bool(COMMAND_PATTERN.match(text.strip()))


class OrchestrationAgent:
    """Per-user agent with conversation history and active session tracking."""

    def __init__(self) -> None:
        llm = ChatOpenAI(
            base_url=OPENAI_BASE_URL,
            api_key=OPENAI_API_KEY,
            model=OPENAI_MODEL,
            temperature=0.0,
        )
        self._llm_with_tools = llm.bind_tools(TOOLS)

        # user_id -> list[BaseMessage]
        self._history: Dict[str, List[BaseMessage]] = defaultdict(list)
        # user_id -> most recently active conv_id
        self._active_conv: Dict[str, Optional[str]] = defaultdict(lambda: None)
        # user_id -> asyncio.Lock (prevents concurrent processing per user)
        self._user_locks: Dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)

    def _build_system_prompt(self, user_id: str) -> str:
        conv_id = self._active_conv[user_id]
        if conv_id:
            active_line = f"ACTIVE SESSION: conv_id={conv_id!r}  ← use this for all follow-up messages"
        else:
            active_line = "ACTIVE SESSION: none"
        return SYSTEM_PROMPT_TEMPLATE.format(
            working_dir=WORKING_DIR,
            active_session_line=active_line,
        )

    def get_active_conv(self, user_id: str) -> Optional[str]:
        return self._active_conv.get(user_id)

    async def run(self, user_id: str, text: str) -> str:
        """Process a user message and return the agent's reply."""
        async with self._user_locks[user_id]:
            return await self._run_locked(user_id, text)

    async def _run_locked(self, user_id: str, text: str) -> str:
        """Internal implementation, must be called with user lock held."""
        set_current_user(user_id)
        active_conv = self._active_conv[user_id]
        short_uid = user_id[-8:]
        logger.info(">>> user=...%s  conv=%s  msg=%r", short_uid, active_conv, text[:80])
        logger.debug("    history_len=%d", len(self._history[user_id]))

        # Passthrough mode: if active session and not a command, bypass LLM
        if active_conv and not _looks_like_command(text):
            try:
                reply = await manager.send(active_conv, text, user_id=user_id)
                logger.info("<<< [passthrough] reply: %r", reply[:120])
                return reply
            except KeyError:
                logger.warning("Session %s no longer exists, clearing active_conv", active_conv)
                self._active_conv[user_id] = None
            except Exception as exc:
                logger.exception("Passthrough error for user=%s", user_id)
                return f"[Error] {exc}"

        messages: List[BaseMessage] = (
            [SystemMessage(content=self._build_system_prompt(user_id))]
            + self._history[user_id]
            + [HumanMessage(content=text)]
        )

        reply = ""
        try:
            for iteration in range(MAX_ITERATIONS):
                logger.debug("    LLM call #%d", iteration)
                ai_msg: AIMessage = await self._llm_with_tools.ainvoke(messages)
                messages.append(ai_msg)

                if not ai_msg.tool_calls:
                    reply = ai_msg.content or ""
                    logger.debug("    → done (no tool call)")
                    break

                for tc in ai_msg.tool_calls:
                    tool_name = tc["name"]
                    tool_args = tc["args"]
                    tool_id = tc["id"]

                    args_summary = ", ".join(
                        f"{k}={str(v)[:50]!r}" for k, v in tool_args.items()
                    )
                    logger.info("    ⚙  %s(%s)", tool_name, args_summary)

                    tool_obj = _TOOL_MAP.get(tool_name)
                    if tool_obj is None:
                        result = f"Unknown tool: {tool_name}"
                        logger.warning("    unknown tool: %s", tool_name)
                    else:
                        try:
                            result = await tool_obj.arun(tool_args)
                        except Exception as exc:
                            result = f"Tool error: {exc}"
                            logger.error("    tool %s error: %s", tool_name, exc)

                    logger.debug("    ←  %s: %r", tool_name, str(result)[:120])

                    if tool_name == "create_conversation":
                        try:
                            data = json.loads(result)
                            if "conv_id" in data:
                                self._active_conv[user_id] = data["conv_id"]
                                logger.info("    ✓  active session → %s", data["conv_id"])
                        except Exception:
                            pass

                    messages.append(
                        ToolMessage(content=str(result), tool_call_id=tool_id)
                    )
            else:
                reply = "[Max iterations reached]"
                logger.warning("    max iterations reached")

        except Exception as exc:
            logger.exception("agent error for user=%s", user_id)
            reply = f"[Error] {exc}"

        logger.info("<<< reply: %r", reply[:120])

        # Update history
        self._history[user_id].append(HumanMessage(content=text))
        self._history[user_id].append(AIMessage(content=reply))

        if len(self._history[user_id]) > 40:
            self._history[user_id] = self._history[user_id][-40:]

        return reply


agent = OrchestrationAgent()