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 json
import logging
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 config import OPENAI_API_KEY, OPENAI_BASE_URL, OPENAI_MODEL, WORKING_DIR
from orchestrator.tools import TOOLS

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}


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)

    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,
        )

    async def run(self, user_id: str, text: str) -> str:
        """Process a user message and return the agent's reply."""
        messages: List[BaseMessage] = (
            [SystemMessage(content=self._build_system_prompt(user_id))]
            + self._history[user_id]
            + [HumanMessage(content=text)]
        )

        reply = ""
        try:
            for _ in range(MAX_ITERATIONS):
                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 ""
                    break

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

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

                    # If a session was just created, record it as the active session
                    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 for %s set to %s",
                                    user_id, data["conv_id"],
                                )
                        except Exception:
                            pass

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

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

        # 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


# Module-level singleton
agent = OrchestrationAgent()