API Reference¶

class AgentConfig(BaseModel):
    """Agent loop behaviour settings."""

    max_iterations: int = 15
    max_threads_per_query: int = 20
    max_link_depth: int = 2
    total_token_budget: int = 200_000
    parallel_tool_calls: bool = True

class AppConfig(BaseModel):
    """Root application configuration."""

    mattermost: MattermostConfig
    llm: LLMConfig
    agent: AgentConfig = Field(default_factory=AgentConfig)
    output: OutputConfig = Field(default_factory=OutputConfig)
    logging: LoggingConfig = Field(default_factory=LoggingConfig)

class LLMConfig(BaseModel):
    """LLM provider settings."""

    base_url: str = "https://api.openai.com/v1"
    api_key: str
    model: str = "gpt-4o-mini"
    temperature: float = 0.2
    max_tokens_per_response: int = 4000
    request_timeout_seconds: int = 120
    verify_ssl: bool = True

    def __repr__(self) -> str:
        masked_key = self.api_key[:5] + "***" if len(self.api_key) > 5 else "***"
        return (
            f"LLMConfig(base_url={self.base_url!r}, api_key={masked_key!r}, "
            f"model={self.model!r}, temperature={self.temperature}, "
            f"max_tokens_per_response={self.max_tokens_per_response})"
        )

    def __str__(self) -> str:
        return self.__repr__()

class LoggingConfig(BaseModel):
    """Logging settings."""

    level: str = "INFO"

class MattermostConfig(BaseModel):
    """Mattermost connection settings.

    Auth is either a personal access token (``token``) or a login/password pair.
    Exactly one of the two must be provided.
    """

    url: str
    team: str | None = None
    timeout_seconds: int = 30
    rate_limit_rps: int = 10
    verify_ssl: bool = True

    # Token-based auth
    token: str | None = None

    # Login/password auth
    login: str | None = None
    password: str | None = None

    @model_validator(mode="after")
    def _validate_auth(self) -> MattermostConfig:
        has_token = bool(self.token)
        has_login = bool(self.login) and bool(self.password)
        if not has_token and not has_login:
            raise ValueError(
                "Mattermost auth is not configured. Provide either 'token' or both 'login' and 'password'."
            )
        if has_token and has_login:
            raise ValueError("Provide either 'token' or 'login'+'password', not both.")
        return self

    def __repr__(self) -> str:
        if self.token:
            auth = "token=" + (self.token[:5] + "***" if len(self.token) > 5 else "***")
        else:
            auth = f"login={self.login!r}, password=***"
        team = f", team={self.team!r}" if self.team else ""
        return (
            f"MattermostConfig(url={self.url!r}, {auth}"
            f"{team}, timeout_seconds={self.timeout_seconds}, "
            f"rate_limit_rps={self.rate_limit_rps})"
        )

    def __str__(self) -> str:
        return self.__repr__()

class OutputConfig(BaseModel):
    """Output rendering settings."""

    show_thread_tree: bool = True
    show_token_usage: bool = True
    show_timings: bool = True
    format: str = "markdown"  # markdown | json | plain

class MattermostAPIError(Exception):
    """Raised when the Mattermost API returns an error response."""

    def __init__(self, status_code: int, message: str) -> None:
        self.status_code = status_code
        super().__init__(f"Mattermost API error {status_code}: {message}")

class MattermostClient:
    """Async HTTP client for Mattermost REST API v4.

    Usage::

        async with MattermostClient(config) as client:
            hits = await client.search_posts(team_id, "incident", None, 20)
    """

    def __init__(self, config: MattermostConfig) -> None:
        self._config = config
        self._base_url = config.url.rstrip("/")
        self._headers: dict[str, str] = {"Content-Type": "application/json"}
        if config.token:
            self._headers["Authorization"] = f"Bearer {config.token}"
        self._semaphore = asyncio.Semaphore(config.rate_limit_rps)
        self._http: httpx.AsyncClient | None = None

    async def __aenter__(self) -> MattermostClient:
        self._http = httpx.AsyncClient(
            base_url=self._base_url,
            headers=self._headers,
            timeout=self._config.timeout_seconds,
            verify=self._config.verify_ssl,
        )
        if not self._config.token:
            await self._login()
        return self

    async def _login(self) -> None:
        """Authenticate with login/password and store the session token."""
        assert self._http is not None
        response = await self._http.post(
            "/api/v4/users/login",
            json={"login_id": self._config.login, "password": self._config.password},
        )
        if response.status_code >= 400:
            try:
                detail = response.json().get("message", response.text)
            except Exception:
                detail = response.text
            raise MattermostAPIError(response.status_code, detail)

        token = response.headers.get("Token")
        if not token:
            raise MattermostAPIError(0, "Login succeeded but no Token header returned.")

        assert self._http is not None
        self._http.headers["Authorization"] = f"Bearer {token}"
        logger.debug("Authenticated via login/password, session token obtained.")

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        if self._http is not None:
            await self._http.aclose()
            self._http = None

    @property
    def _client(self) -> httpx.AsyncClient:
        if self._http is None:
            raise RuntimeError("MattermostClient must be used as an async context manager.")
        return self._http

    async def _request(
        self,
        method: str,
        path: str,
        **kwargs: Any,
    ) -> Any:
        """Execute a rate-limited, retried HTTP request and return parsed JSON."""

        @retry(
            retry=retry_if_exception(_should_retry),
            stop=stop_after_attempt(3),
            wait=wait_exponential(multiplier=1, min=1, max=8),
            reraise=True,
        )
        async def _inner() -> Any:
            async with self._semaphore:
                logger.debug("%s %s", method.upper(), path)
                response = await self._client.request(method, path, **kwargs)

            if response.status_code >= 400:
                try:
                    detail = response.json().get("message", response.text)
                except Exception:
                    detail = response.text
                raise MattermostAPIError(response.status_code, detail)

            if response.status_code == 204 or not response.content:
                return {}

            return response.json()

        return await _inner()

    # ------------------------------------------------------------------ #
    # Public API methods                                                   #
    # ------------------------------------------------------------------ #

    async def get_my_teams(self) -> list[Team]:
        """Return all teams the current user belongs to."""
        data: list[dict[str, Any]] = await self._request("GET", "/api/v4/users/me/teams")
        return [Team.model_validate(t) for t in data]

    async def validate_connection(self) -> bool:
        """Return True if the token is valid and the server is reachable."""
        try:
            await self._request("GET", "/api/v4/users/me")
        except Exception:
            return False
        else:
            return True

    async def get_team_id(self, team_name_or_id: str) -> str:
        """Resolve a team name or ID to a team ID string."""
        try:
            data: dict[str, Any] = await self._request("GET", f"/api/v4/teams/name/{team_name_or_id}")
            return str(data["id"])
        except MattermostAPIError as exc:
            if exc.status_code == 404:
                # Maybe it was already an ID — validate it
                try:
                    data = await self._request("GET", f"/api/v4/teams/{team_name_or_id}")
                    return str(data["id"])
                except MattermostAPIError:
                    pass
            raise

    async def search_posts(
        self,
        team_id: str,
        query: str,
        channel_id: str | None = None,
        per_page: int = 20,
    ) -> list[SearchHit]:
        """Full-text search for posts within a team."""
        payload: dict[str, Any] = {"terms": query, "is_or_search": False, "per_page": per_page}
        if channel_id:
            payload["channel_id"] = channel_id

        data: dict[str, Any] = await self._request(
            "POST",
            f"/api/v4/teams/{team_id}/posts/search",
            json=payload,
        )

        order: list[str] = data.get("order", [])
        posts_map: dict[str, Any] = data.get("posts", {})

        # Fetch channel names in parallel (one request per unique channel_id)
        unique_channel_ids: set[str] = {posts_map[pid]["channel_id"] for pid in order if pid in posts_map}
        channel_names: dict[str, str] = {}
        if unique_channel_ids:
            results = await asyncio.gather(
                *[self.get_channel(cid) for cid in unique_channel_ids],
                return_exceptions=True,
            )
            for cid, result in zip(unique_channel_ids, results, strict=False):
                if isinstance(result, dict):
                    channel_names[cid] = result.get("name", cid)
                else:
                    channel_names[cid] = cid

        hits: list[SearchHit] = []
        for post_id in order:
            if post_id not in posts_map:
                continue
            raw = posts_map[post_id]
            cid = raw.get("channel_id", "")
            hits.append(
                SearchHit(
                    post_id=post_id,
                    message=raw.get("message", ""),
                    channel_id=cid,
                    channel_name=channel_names.get(cid, cid),
                    user_id=raw.get("user_id", ""),
                )
            )

        return hits

    async def get_thread(self, post_id: str) -> Thread:
        """Fetch the complete thread containing post_id."""
        data: dict[str, Any] = await self._request(
            "GET",
            f"/api/v4/posts/{post_id}/thread",
            params={"skipFetchThreads": "false", "collapsedThreads": "false"},
        )

        posts_map: dict[str, Any] = data.get("posts", {})
        posts: list[Post] = []
        for raw in posts_map.values():
            try:
                posts.append(Post.model_validate(raw))
            except Exception as exc:
                logger.warning("Failed to parse post %s: %s", raw.get("id"), exc)

        posts.sort(key=lambda p: p.create_at)

        # The root post is the one with no root_id
        root_id = next(
            (p.id for p in posts if not p.root_id),
            posts[0].id if posts else post_id,
        )

        return Thread(root_post_id=root_id, posts=posts)

    async def get_user(self, user_id: str) -> User:
        """Fetch a user by ID."""
        data: dict[str, Any] = await self._request("GET", f"/api/v4/users/{user_id}")
        return User.model_validate(data)

    async def get_channel(self, channel_id: str) -> dict[str, str]:
        """Fetch channel metadata by ID."""
        data: dict[str, Any] = await self._request("GET", f"/api/v4/channels/{channel_id}")
        return {k: str(v) for k, v in data.items()}

class AgentLoop:
    """Orchestrates the LLM + tool-calling loop to answer a user's question."""

    def __init__(
        self,
        config: AppConfig,
        client: MattermostClient,
        console: Console,
        verbose: bool = False,
    ) -> None:
        self._config = config
        self._mm_client = client
        self._console = console
        self._verbose = verbose

        self._llm = AsyncOpenAI(
            base_url=config.llm.base_url,
            api_key=config.llm.api_key,
            http_client=httpx.AsyncClient(
                verify=config.llm.verify_ssl,
                timeout=config.llm.request_timeout_seconds,
            ),
        )

    async def run(
        self,
        question: str,
        on_status: Callable[[str], None] | None = None,
    ) -> AskResult:
        """Run the agentic loop for a given question and return AskResult."""
        start = time.monotonic()

        state = AgentState()

        # Resolve team_id once
        if on_status:
            on_status("Connecting to Mattermost...")
        if not self._config.mattermost.team:
            raise ValueError("Mattermost team is not set. Add 'team' to your config or use --team.")
        team_id = await self._mm_client.get_team_id(self._config.mattermost.team)

        system = SYSTEM_PROMPT.format(max_depth=self._config.agent.max_link_depth)

        messages: list[ChatCompletionMessageParam] = [
            {"role": "system", "content": system},
            {"role": "user", "content": question},
        ]

        answer = ""
        iteration = 0

        while iteration < self._config.agent.max_iterations:
            iteration += 1

            if on_status:
                on_status(f"Thinking... (iteration {iteration})")

            logger.debug("LLM call iteration %d", iteration)

            typed_tools = cast(list[ChatCompletionToolParam], TOOL_DEFINITIONS)
            response = await self._llm.chat.completions.create(
                model=self._config.llm.model,
                messages=messages,
                tools=typed_tools,
                tool_choice="auto",
                temperature=self._config.llm.temperature,
                max_tokens=self._config.llm.max_tokens_per_response,
                timeout=self._config.llm.request_timeout_seconds,
            )

            # Track token usage
            if response.usage:
                state.token_usage = state.token_usage.add(
                    TokenUsage(
                        prompt_tokens=response.usage.prompt_tokens,
                        completion_tokens=response.usage.completion_tokens,
                        total_tokens=response.usage.total_tokens,
                    )
                )

            # Check budget
            if state.token_usage.total_tokens >= self._config.agent.total_token_budget:
                logger.warning("Token budget exhausted at iteration %d", iteration)
                if on_status:
                    on_status(f"Token budget exhausted ({state.token_usage.total_tokens:,} tokens). Continue? [y/N] ")
                confirmed = await asyncio.get_event_loop().run_in_executor(None, self._ask_continue)
                if confirmed:
                    # Double the remaining budget and keep going
                    self._config.agent.total_token_budget = (
                        state.token_usage.total_tokens + self._config.agent.total_token_budget
                    )
                    logger.debug("Budget extended to %d", self._config.agent.total_token_budget)
                    continue
                state.incomplete = True
                if on_status:
                    on_status("Stopped by user.")
                break

            choice = response.choices[0]
            message = choice.message

            # Append assistant message to conversation
            messages.append(cast(ChatCompletionMessageParam, message.model_dump(exclude_none=True)))

            # Only handle function tool calls (ignore custom tool calls).
            # Use isinstance to satisfy mypy; MagicMock objects in tests pass this check
            # because MagicMock supports arbitrary attribute access and isinstance checks.
            raw_tool_calls = message.tool_calls or []
            tool_calls: list[ChatCompletionMessageToolCall] = [
                tc  # type: ignore[misc]
                for tc in raw_tool_calls
                if hasattr(tc, "function") and hasattr(tc, "id")
            ]

            if not tool_calls:
                # LLM gave a final text answer (or only unsupported custom tools)
                answer = message.content or ""
                if on_status:
                    on_status("Answer ready.")
                break

            # --- Execute all tool calls (optionally in parallel) ---
            if on_status:
                tool_names = ", ".join(tc.function.name for tc in tool_calls)
                on_status(f"Calling tools: {tool_names}")

            if self._verbose:
                for tc in tool_calls:
                    self._console.print(f"[dim]Tool call: {tc.function.name}({tc.function.arguments})[/dim]")

            tool_tasks = [
                self._run_tool(tc.id, tc.function.name, tc.function.arguments, team_id, state) for tc in tool_calls
            ]

            if self._config.agent.parallel_tool_calls:
                tool_results = await asyncio.gather(*tool_tasks, return_exceptions=True)
            else:
                tool_results = []
                for task in tool_tasks:
                    tool_results.append(await task)

            # Append tool results to message history
            for tc, result in zip(tool_calls, tool_results, strict=False):
                result_str = (
                    orjson.dumps({"error": str(result)}).decode() if isinstance(result, BaseException) else str(result)
                )

                if self._verbose:
                    self._console.print(f"[dim]Tool result ({tc.function.name}): {result_str[:200]}[/dim]")

                tool_msg: ChatCompletionToolMessageParam = {
                    "role": "tool",
                    "tool_call_id": tc.id,
                    "content": result_str,
                }
                messages.append(tool_msg)

        else:
            # Iteration limit reached without a final answer
            state.incomplete = True
            if on_status:
                on_status("Iteration limit reached — stopping early.")
            # Use the last assistant content if any
            for msg in reversed(messages):
                if isinstance(msg, dict) and msg.get("role") == "assistant":
                    answer = str(msg.get("content") or "")
                    break

        elapsed = time.monotonic() - start

        # Collect unique permalinks from explored threads
        permalinks = list({t["permalink"] for t in state.explored_threads if t.get("permalink")})

        return AskResult(
            answer=answer,
            threads_explored=state.threads_fetched,
            tool_calls_made=state.tool_calls_made,
            token_usage=state.token_usage,
            elapsed_seconds=elapsed,
            incomplete=state.incomplete,
            permalinks=permalinks,
        )

    def _ask_continue(self) -> bool:
        """Prompt the user synchronously (runs in a thread executor)."""
        try:
            answer = input("Continue beyond token budget? [y/N]: ").strip().lower()
        except (EOFError, KeyboardInterrupt):
            return False
        else:
            return answer in ("y", "yes", "д", "да")

    async def _run_tool(
        self,
        _tool_call_id: str,
        tool_name: str,
        arguments_json: str,
        team_id: str,
        state: AgentState,
    ) -> str:
        """Parse arguments and delegate to execute_tool."""
        try:
            args: dict[str, Any] = orjson.loads(arguments_json)
        except orjson.JSONDecodeError as exc:
            return orjson.dumps({"error": f"Invalid tool arguments JSON: {exc}"}).decode()

        return await execute_tool(
            tool_name=tool_name,
            tool_args=args,
            client=self._mm_client,
            team_id=team_id,
            state=state,
            config=self._config.agent,
            mm_url=self._config.mattermost.url,
            mm_team=self._config.mattermost.team or "",
        )