hkuds--deeptutor/deeptutor/co_writer/edit_agent.py

"""
EditAgent - Co-writer editing agent.
Inherits from unified BaseAgent.
"""

import asyncio
from datetime import datetime
import json
from typing import Any, Literal
import uuid

from deeptutor.agents.base_agent import BaseAgent
from deeptutor.co_writer.storage import _atomic_write_json
from deeptutor.runtime.registry.tool_registry import get_tool_registry
from deeptutor.services.llm import clean_thinking_tags
from deeptutor.services.path_service import get_path_service
from deeptutor.tools.rag_tool import rag_search
from deeptutor.tools.web_search import web_search


# Resolved per-call so a per-user PathService (set after auth) routes
# co-writer history/tool-call files under the caller's own workspace.
def _user_dir():
    return get_path_service().get_co_writer_dir()


def _history_file():
    return get_path_service().get_co_writer_history_file()


def tool_calls_dir():
    return get_path_service().get_co_writer_tool_calls_dir()


def ensure_dirs():
    """Ensure directories exist"""
    _user_dir().mkdir(parents=True, exist_ok=True)
    tool_calls_dir().mkdir(parents=True, exist_ok=True)


# History is a debugging/audit trail, not a primary store: cap the entry
# count and clip stored texts so a long-lived workspace can't grow the file
# (rewritten in full on every operation) without bound.
_HISTORY_MAX_ENTRIES = 200
_HISTORY_TEXT_LIMIT = 20_000


def load_history() -> list:
    """Load history"""
    ensure_dirs()
    history_file = _history_file()
    if history_file.exists():
        try:
            with open(history_file, encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            return []
    return []


def save_history(history: list):
    """Save history (atomic write; a crash mid-write must not corrupt it)."""
    ensure_dirs()
    _atomic_write_json(_history_file(), history)


def _clip_history_value(value: Any) -> Any:
    if isinstance(value, str) and len(value) > _HISTORY_TEXT_LIMIT:
        return value[:_HISTORY_TEXT_LIMIT] + "…[truncated]"
    if isinstance(value, dict):
        return {k: _clip_history_value(v) for k, v in value.items()}
    return value


def append_history(record: dict) -> None:
    """Append one operation record, clipping texts and capping entries."""
    history = load_history()
    history.append(_clip_history_value(record))
    if len(history) > _HISTORY_MAX_ENTRIES:
        history = history[-_HISTORY_MAX_ENTRIES:]
    save_history(history)


def save_tool_call(call_id: str, tool_type: str, data: dict[str, Any]) -> str:
    """Save tool call result, return file path"""
    ensure_dirs()
    filename = f"{call_id}_{tool_type}.json"
    filepath = tool_calls_dir() / filename
    with open(filepath, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    return str(filepath)


class EditAgent(BaseAgent):
    """Co-writer editing agent using unified BaseAgent."""

    def __init__(
        self,
        language: str = "en",
        enabled_tools: list[str] | None = None,
        **kwargs: Any,
    ):
        """
        Initialize EditAgent.

        Args:
            language: Language setting ('en' | 'zh'), default 'en'

        Note: LLM configuration (api_key, base_url, model, etc.) is loaded
        automatically from the unified config service. Use refresh_config()
        to pick up configuration changes made in Settings.
        """
        super().__init__(
            module_name="co_writer",
            agent_name="edit_agent",
            language=language,
            **kwargs,
        )
        self.enabled_tools = enabled_tools or ["rag", "web_search"]
        self._tool_registry = get_tool_registry()

    async def process(
        self,
        text: str,
        instruction: str,
        action: Literal["rewrite", "shorten", "expand"] = "rewrite",
        source: Literal["rag", "web"] | None = None,
        kb_name: str | None = None,
    ) -> dict[str, Any]:
        """
        Process edit request

        Returns:
            Dict containing:
                - edited_text: Edited text
                - operation_id: Operation ID
        """
        operation_id = datetime.now().strftime("%Y%m%d_%H%M%S") + "_" + uuid.uuid4().hex[:6]

        context = ""
        tool_call_file = None

        if source:
            context, tool_call_file = await self.gather_context(
                source=source,
                query=instruction,
                kb_name=kb_name,
                operation_id=operation_id,
            )
            if not context:
                source = None

        # Build prompts
        system_template = self.get_prompt(
            "system",
            "You are an expert editor and writing assistant.\n\nAvailable reference tools:\n{available_tools}",
        )
        system_prompt = system_template.format(available_tools=self._build_available_tools_text())

        action_verbs = {"rewrite": "Rewrite", "shorten": "Shorten", "expand": "Expand"}
        action_verb = action_verbs.get(action, "Rewrite")

        action_template = self.get_prompt(
            "action_template",
            "{action_verb} the following text based on the user's instruction.\n\nUser Instruction: {instruction}\n\n",
        )
        user_prompt = action_template.format(action_verb=action_verb, instruction=instruction)

        if context:
            context_template = self.get_prompt(
                "context_template", "Reference Context ({source_label}):\n{context}\n\n"
            )
            user_prompt += context_template.format(
                context=context,
                source_label=self._get_source_label(source),
            )

        text_template = self.get_prompt(
            "user_template",
            "Target Text to Edit:\n{text}\n\nOutput only the edited text, without quotes or explanations.",
        )
        user_prompt += text_template.format(text=text)

        # Call LLM using inherited method
        self.logger.info(f"Calling LLM for {action}...")
        _chunks: list[str] = []
        async for _c in self.stream_llm(
            user_prompt=user_prompt,
            system_prompt=system_prompt,
            stage=f"edit_{action}",
        ):
            _chunks.append(_c)
        response = clean_thinking_tags("".join(_chunks), self.binding, self.get_model())

        # Record operation history
        append_history(
            {
                "id": operation_id,
                "timestamp": datetime.now().isoformat(),
                "action": action,
                "source": source,
                "kb_name": kb_name,
                "input": {"original_text": text, "instruction": instruction},
                "output": {"edited_text": response},
                "tool_call_file": tool_call_file,
                "model": self.get_model(),
            }
        )

        self.logger.info(f"Operation {operation_id} recorded successfully")

        return {"edited_text": response, "operation_id": operation_id}

    async def gather_context(
        self,
        *,
        source: Literal["rag", "web"],
        query: str,
        kb_name: str | None = None,
        operation_id: str = "",
    ) -> tuple[str, str | None]:
        """Fetch reference context for an edit.

        Returns ``(context, tool_call_file)``; empty context means the tool
        was unavailable or failed — callers degrade to a plain edit.
        """
        if source == "rag":
            if "rag" not in self.enabled_tools:
                self.logger.warning("RAG source requested but tool is not enabled")
                return "", None
            if not kb_name:
                self.logger.warning("RAG source selected but no kb_name provided")
                return "", None
            self.logger.info(f"Searching RAG in KB: {kb_name} for: {query}")
            try:
                search_result = await rag_search(
                    query=query, kb_name=kb_name, only_need_context=True
                )
                context = search_result.get("answer", "")
                self.logger.info(f"RAG context found: {len(context)} chars")
                tool_call_file = save_tool_call(
                    operation_id,
                    "rag",
                    {
                        "type": "rag",
                        "timestamp": datetime.now().isoformat(),
                        "operation_id": operation_id,
                        "query": query,
                        "kb_name": kb_name,
                        "mode": "naive",
                        "context": context,
                        "raw_result": search_result,
                    },
                )
                return context, tool_call_file
            except Exception as e:
                self.logger.error(f"RAG search failed: {e}, continuing without context")
                return "", None

        if "web_search" not in self.enabled_tools:
            self.logger.warning("Web source requested but tool is not enabled")
            return "", None
        self.logger.info(f"Searching Web for: {query}")
        try:
            # web_search is synchronous network I/O; keep it off the
            # event loop so concurrent requests aren't stalled.
            search_result = await asyncio.to_thread(web_search, query)
            context = search_result.get("answer", "")
            self.logger.info(f"Web context found: {len(context)} chars")
            tool_call_file = save_tool_call(
                operation_id,
                "web",
                {
                    "type": "web_search",
                    "timestamp": datetime.now().isoformat(),
                    "operation_id": operation_id,
                    "query": query,
                    "answer": context,
                    "citations": search_result.get("citations", []),
                    "search_results": search_result.get("search_results", []),
                    "usage": search_result.get("usage", {}),
                },
            )
            return context, tool_call_file
        except Exception as e:
            self.logger.error(f"Web search failed: {e}, continuing without context")
            return "", None

    async def auto_mark(self, text: str) -> dict[str, Any]:
        """
        AI auto-marking feature - Add annotation tags to text

        Returns:
            Dict containing:
                - marked_text: Text with annotations
                - operation_id: Operation ID
        """
        operation_id = datetime.now().strftime("%Y%m%d_%H%M%S") + "_" + uuid.uuid4().hex[:6]

        system_prompt = self.get_prompt("auto_mark_system", "")
        user_template = self.get_prompt(
            "auto_mark_user_template", "Process the following text:\n{text}"
        )
        user_prompt = user_template.format(text=text)

        self.logger.info("Calling LLM for auto-mark...")
        _chunks: list[str] = []
        async for _c in self.stream_llm(
            user_prompt=user_prompt,
            system_prompt=system_prompt,
            stage="auto_mark",
        ):
            _chunks.append(_c)
        response = clean_thinking_tags("".join(_chunks), self.binding, self.get_model())

        # Record operation history
        append_history(
            {
                "id": operation_id,
                "timestamp": datetime.now().isoformat(),
                "action": "automark",
                "source": None,
                "kb_name": None,
                "input": {"original_text": text, "instruction": "AI Auto Mark"},
                "output": {"edited_text": response},
                "tool_call_file": None,
                "model": self.get_model(),
            }
        )

        self.logger.info(f"Auto-mark operation {operation_id} recorded successfully")

        return {"marked_text": response, "operation_id": operation_id}

    def _build_available_tools_text(self) -> str:
        tool_names = [name for name in self.enabled_tools if name in {"rag", "web_search"}]
        if not tool_names:
            return (
                "（当前未启用外部参考工具）"
                if str(self.language).lower().startswith("zh")
                else "(no external reference tools enabled)"
            )
        return self._tool_registry.build_prompt_text(
            tool_names,
            format="list",
            language=self.language,
        )

    def _get_source_label(self, source: Literal["rag", "web"] | None) -> str:
        labels = {
            "en": {"rag": "knowledge base", "web": "web search"},
            "zh": {"rag": "知识库", "web": "网页搜索"},
        }
        lang = "zh" if str(self.language).lower().startswith("zh") else "en"
        if source in labels[lang]:
            return labels[lang][source]
        return "reference" if lang == "en" else "参考资料"


# Legacy compatibility - export get_stats pointing to BaseAgent's stats
def get_stats():
    """Get shared stats tracker for co_writer module."""
    return BaseAgent.get_stats("co_writer")


def reset_stats():
    """Reset shared stats for co_writer module."""
    BaseAgent.reset_stats("co_writer")


def print_stats():
    """Print stats summary for co_writer module."""
    BaseAgent.print_stats("co_writer")