chore: import upstream snapshot with attribution

application/agents/base.py

@@ -0,0 +1,691 @@
+import json
+import logging
+import uuid
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Generator, List, Optional
+
+from application.agents.tool_executor import (
+    ToolExecutor,
+    result_status,
+    truncate_tool_result,
+)
+from application.core.json_schema_utils import (
+    JsonSchemaValidationError,
+    normalize_json_schema_payload,
+)
+from application.core.settings import settings
+from application.llm.handlers.base import ToolCall
+from application.llm.handlers.handler_creator import LLMHandlerCreator
+from application.llm.llm_creator import LLMCreator
+from application.logging import build_stack_data, log_activity, LogContext
+
+logger = logging.getLogger(__name__)
+
+
+class BaseAgent(ABC):
+    def __init__(
+        self,
+        endpoint: str,
+        llm_name: str,
+        model_id: str,
+        api_key: str,
+        agent_id: Optional[str] = None,
+        user_api_key: Optional[str] = None,
+        prompt: str = "",
+        chat_history: Optional[List[Dict]] = None,
+        retrieved_docs: Optional[List[Dict]] = None,
+        decoded_token: Optional[Dict] = None,
+        attachments: Optional[List[Dict]] = None,
+        json_schema: Optional[Dict] = None,
+        json_schema_strict: bool = True,
+        json_object: bool = False,
+        llm_params: Optional[Dict] = None,
+        multimodal_content: Optional[List] = None,
+        limited_token_mode: Optional[bool] = False,
+        token_limit: Optional[int] = settings.DEFAULT_AGENT_LIMITS["token_limit"],
+        limited_request_mode: Optional[bool] = False,
+        request_limit: Optional[int] = settings.DEFAULT_AGENT_LIMITS["request_limit"],
+        compressed_summary: Optional[str] = None,
+        llm=None,
+        llm_handler=None,
+        tool_executor: Optional[ToolExecutor] = None,
+        backup_models: Optional[List[str]] = None,
+        model_user_id: Optional[str] = None,
+    ):
+        self.endpoint = endpoint
+        self.llm_name = llm_name
+        self.model_id = model_id
+        self.api_key = api_key
+        self.agent_id = agent_id
+        self.user_api_key = user_api_key
+        self.prompt = prompt
+        self.decoded_token = decoded_token or {}
+        self.user: str = self.decoded_token.get("sub")
+        # BYOM-resolution scope: owner for shared agents, caller for
+        # caller-owned BYOM, None for built-ins. Falls back to self.user
+        # for worker/legacy callers that don't thread model_user_id.
+        self.model_user_id = model_user_id
+        self.tools: List[Dict] = []
+        self.chat_history: List[Dict] = chat_history if chat_history is not None else []
+
+        if llm is not None:
+            self.llm = llm
+        else:
+            self.llm = LLMCreator.create_llm(
+                llm_name,
+                api_key=api_key,
+                user_api_key=user_api_key,
+                decoded_token=decoded_token,
+                model_id=model_id,
+                agent_id=agent_id,
+                backup_models=backup_models,
+                model_user_id=model_user_id,
+            )
+
+        # For BYOM, registry id (UUID) differs from upstream model id
+        # (e.g. ``mistral-large-latest``). LLMCreator resolved this onto
+        # the LLM instance; cache it for subsequent gen calls.
+        self.upstream_model_id = (
+            getattr(self.llm, "model_id", None) or model_id
+        )
+
+        self.retrieved_docs = retrieved_docs or []
+
+        if llm_handler is not None:
+            self.llm_handler = llm_handler
+        else:
+            self.llm_handler = LLMHandlerCreator.create_handler(
+                llm_name if llm_name else "default"
+            )
+
+        # Tool executor — injected or created
+        if tool_executor is not None:
+            self.tool_executor = tool_executor
+        else:
+            self.tool_executor = ToolExecutor(
+                user_api_key=user_api_key,
+                user=self.user,
+                decoded_token=decoded_token,
+                agent_id=agent_id,
+            )
+
+        self.attachments = attachments or []
+        self.json_schema = None
+        if json_schema is not None:
+            try:
+                self.json_schema = normalize_json_schema_payload(json_schema)
+            except JsonSchemaValidationError as exc:
+                logger.warning("Ignoring invalid JSON schema payload: %s", exc)
+        # Per-request structured-output controls (OpenAI-compatible):
+        # ``json_schema_strict`` mirrors response_format.json_schema.strict;
+        # ``json_object`` mirrors response_format {"type":"json_object"}.
+        self.json_schema_strict = json_schema_strict
+        self.json_object = json_object
+        # OpenAI sampling params forwarded from the request (temperature,
+        # max_tokens, top_p, ...). Empty when the caller sent none.
+        self.llm_params = llm_params or {}
+        # Full OpenAI content array (text + image_url parts) for the current
+        # user turn, when the request was multimodal; None otherwise.
+        self.multimodal_content = multimodal_content
+        self.limited_token_mode = limited_token_mode
+        self.token_limit = token_limit
+        self.limited_request_mode = limited_request_mode
+        self.request_limit = request_limit
+        self.compressed_summary = compressed_summary
+        self.current_token_count = 0
+        self.context_limit_reached = False
+        self.conversation_id: Optional[str] = None
+        self.initial_user_id: Optional[str] = None
+
+    @log_activity()
+    def gen(
+        self, query: str, log_context: LogContext = None
+    ) -> Generator[Dict, None, None]:
+        yield from self._gen_inner(query, log_context)
+        yield from self._emit_responses_metadata()
+
+    def _emit_responses_metadata(self) -> Generator[Dict, None, None]:
+        """Surface the latest Responses API id so the route can persist it in
+        message metadata for previous_response_id chaining across turns."""
+        if not settings.OPENAI_RESPONSES_STORE:
+            return
+        response_id = getattr(self.llm, "_last_response_id", None)
+        if response_id:
+            yield {"metadata": {"response_id": response_id}}
+
+    def _previous_response_id(self) -> Optional[str]:
+        """Most recent stored Responses API id from chat history, if any."""
+        for turn in reversed(self.chat_history or []):
+            if not isinstance(turn, dict):
+                continue
+            meta = turn.get("metadata")
+            if isinstance(meta, dict) and meta.get("response_id"):
+                return meta["response_id"]
+        return None
+
+    @abstractmethod
+    def _gen_inner(
+        self, query: str, log_context: LogContext
+    ) -> Generator[Dict, None, None]:
+        pass
+
+    def gen_continuation(
+        self,
+        messages: List[Dict],
+        tools_dict: Dict,
+        pending_tool_calls: List[Dict],
+        tool_actions: List[Dict],
+        reasoning_content: str = "",
+    ) -> Generator[Dict, None, None]:
+        """Resume generation after tool actions are resolved.
+
+        Processes the client-provided *tool_actions* (approvals, denials,
+        or client-side results), appends the resulting messages, then
+        hands back to the LLM to continue the conversation.
+
+        Args:
+            messages: The saved messages array from the pause point.
+            tools_dict: The saved tools dictionary.
+            pending_tool_calls: The pending tool call descriptors from the pause.
+            tool_actions: Client-provided actions resolving the pending calls.
+        """
+        self._prepare_tools(tools_dict)
+
+        actions_by_id = {a["call_id"]: a for a in tool_actions}
+
+        # Build a single assistant message containing all tool calls so
+        # the message history matches the format LLM providers expect
+        # (one assistant message with N tool_calls, followed by N tool results).
+        tc_objects: List[Dict[str, Any]] = []
+        for pending in pending_tool_calls:
+            call_id = pending["call_id"]
+            args = pending["arguments"]
+            args_str = (
+                json.dumps(args) if isinstance(args, dict) else (args or "{}")
+            )
+            tc_obj: Dict[str, Any] = {
+                "id": call_id,
+                "type": "function",
+                "function": {
+                    "name": pending["name"],
+                    "arguments": args_str,
+                },
+            }
+            if pending.get("thought_signature"):
+                tc_obj["thought_signature"] = pending["thought_signature"]
+            tc_objects.append(tc_obj)
+
+        resumed_assistant: Dict[str, Any] = {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": tc_objects,
+        }
+        if reasoning_content:
+            resumed_assistant["reasoning_content"] = reasoning_content
+        messages.append(resumed_assistant)
+
+        # Now process each pending call and append tool result messages
+        for pending in pending_tool_calls:
+            call_id = pending["call_id"]
+            args = pending["arguments"]
+            action = actions_by_id.get(call_id)
+            if not action:
+                action = {
+                    "call_id": call_id,
+                    "decision": "denied",
+                    "comment": "No response provided",
+                }
+
+            if action.get("decision") == "approved":
+                # Execute the tool server-side
+                tc = ToolCall(
+                    id=call_id,
+                    name=pending["name"],
+                    arguments=(
+                        json.dumps(args) if isinstance(args, dict) else args
+                    ),
+                )
+                tool_gen = self._execute_tool_action(tools_dict, tc)
+                tool_response = None
+                while True:
+                    try:
+                        event = next(tool_gen)
+                        yield event
+                    except StopIteration as e:
+                        tool_response, _ = e.value
+                        break
+                messages.append(
+                    self.llm_handler.create_tool_message(tc, tool_response)
+                )
+
+            elif action.get("decision") == "denied":
+                comment = action.get("comment", "")
+                denial = (
+                    f"Tool execution denied by user. Reason: {comment}"
+                    if comment
+                    else "Tool execution denied by user."
+                )
+                tc = ToolCall(
+                    id=call_id, name=pending["name"], arguments=args
+                )
+                messages.append(
+                    self.llm_handler.create_tool_message(tc, denial)
+                )
+                yield {
+                    "type": "tool_call",
+                    "data": {
+                        "tool_name": pending.get("tool_name", "unknown"),
+                        "call_id": call_id,
+                        "action_name": pending.get("llm_name", pending["name"]),
+                        "arguments": args,
+                        "status": "denied",
+                    },
+                }
+
+            elif "result" in action:
+                result = action["result"]
+                result_str = (
+                    json.dumps(result)
+                    if not isinstance(result, str)
+                    else result
+                )
+                tc = ToolCall(
+                    id=call_id, name=pending["name"], arguments=args
+                )
+                messages.append(
+                    self.llm_handler.create_tool_message(tc, result_str)
+                )
+                yield {
+                    "type": "tool_call",
+                    "data": {
+                        "tool_name": pending.get("tool_name", "unknown"),
+                        "call_id": call_id,
+                        "action_name": pending.get("llm_name", pending["name"]),
+                        "arguments": args,
+                        "result": truncate_tool_result(result_str),
+                        "status": result_status(result),
+                    },
+                }
+
+        # Resume the LLM loop with the updated messages
+        llm_response = self._llm_gen(messages)
+        yield from self._handle_response(
+            llm_response, tools_dict, messages, None
+        )
+
+        yield {"sources": self.retrieved_docs}
+        yield {"tool_calls": self._get_truncated_tool_calls()}
+        yield from self._emit_responses_metadata()
+
+    # ---- Tool delegation (thin wrappers around ToolExecutor) ----
+
+    @property
+    def tool_calls(self) -> List[Dict]:
+        return self.tool_executor.tool_calls
+
+    @tool_calls.setter
+    def tool_calls(self, value: List[Dict]):
+        self.tool_executor.tool_calls = value
+
+    def _get_tools(self, api_key: str = None) -> Dict[str, Dict]:
+        return self.tool_executor._get_tools_by_api_key(api_key or self.user_api_key)
+
+    def _get_user_tools(self, user="local"):
+        return self.tool_executor._get_user_tools(user)
+
+    def _build_tool_parameters(self, action):
+        return self.tool_executor._build_tool_parameters(action)
+
+    def _prepare_tools(self, tools_dict):
+        self.tools = self.tool_executor.prepare_tools_for_llm(tools_dict)
+
+    def _execute_tool_action(self, tools_dict, call):
+        # Mirror the request's attachments onto the executor so sandbox tools
+        # can lazily bridge a referenced chat attachment to a conversation
+        # artifact; only the caller's own (user-scoped) attachments are passed.
+        self.tool_executor.attachments = self.attachments
+        return self.tool_executor.execute(
+            tools_dict, call, self.llm.__class__.__name__
+        )
+
+    def _get_truncated_tool_calls(self):
+        return self.tool_executor.get_truncated_tool_calls()
+
+    # ---- Context / token management ----
+
+    def _calculate_current_context_tokens(self, messages: List[Dict]) -> int:
+        from application.api.answer.services.compression.token_counter import (
+            TokenCounter,
+        )
+        return TokenCounter.count_message_tokens(messages)
+
+    def _check_context_limit(self, messages: List[Dict]) -> bool:
+        from application.core.model_utils import get_token_limit
+
+        try:
+            current_tokens = self._calculate_current_context_tokens(messages)
+            self.current_token_count = current_tokens
+            context_limit = get_token_limit(
+                self.model_id, user_id=self.model_user_id or self.user
+            )
+            threshold = int(context_limit * settings.COMPRESSION_THRESHOLD_PERCENTAGE)
+
+            if current_tokens >= threshold:
+                logger.warning(
+                    f"Context limit approaching: {current_tokens}/{context_limit} tokens "
+                    f"({(current_tokens/context_limit)*100:.1f}%)"
+                )
+                return True
+            return False
+        except Exception as e:
+            logger.error(f"Error checking context limit: {str(e)}", exc_info=True)
+            return False
+
+    def _validate_context_size(self, messages: List[Dict]) -> None:
+        from application.core.model_utils import get_token_limit
+
+        current_tokens = self._calculate_current_context_tokens(messages)
+        self.current_token_count = current_tokens
+        context_limit = get_token_limit(
+            self.model_id, user_id=self.model_user_id or self.user
+        )
+        percentage = (current_tokens / context_limit) * 100
+
+        if current_tokens >= context_limit:
+            logger.warning(
+                f"Context at limit: {current_tokens:,}/{context_limit:,} tokens "
+                f"({percentage:.1f}%). Model: {self.model_id}"
+            )
+        elif current_tokens >= int(
+            context_limit * settings.COMPRESSION_THRESHOLD_PERCENTAGE
+        ):
+            logger.info(
+                f"Context approaching limit: {current_tokens:,}/{context_limit:,} tokens "
+                f"({percentage:.1f}%)"
+            )
+
+    def _truncate_text_middle(self, text: str, max_tokens: int) -> str:
+        from application.utils import num_tokens_from_string
+
+        current_tokens = num_tokens_from_string(text)
+        if current_tokens <= max_tokens:
+            return text
+
+        chars_per_token = len(text) / current_tokens if current_tokens > 0 else 4
+        target_chars = int(max_tokens * chars_per_token * 0.95)
+
+        if target_chars <= 0:
+            return ""
+
+        start_chars = int(target_chars * 0.4)
+        end_chars = int(target_chars * 0.4)
+
+        truncation_marker = "\n\n[... content truncated to fit context limit ...]\n\n"
+        truncated = text[:start_chars] + truncation_marker + text[-end_chars:]
+
+        logger.info(
+            f"Truncated text from {current_tokens:,} to ~{max_tokens:,} tokens "
+            f"(removed middle section)"
+        )
+        return truncated
+
+    # ---- Message building ----
+
+    def _build_messages(
+        self,
+        system_prompt: str,
+        query: str,
+    ) -> List[Dict]:
+        """Build messages using pre-rendered system prompt"""
+        from application.core.model_utils import get_token_limit
+        from application.utils import num_tokens_from_string
+
+        if self.compressed_summary:
+            compression_context = (
+                "\n\n---\n\n"
+                "This session is being continued from a previous conversation that "
+                "has been compressed to fit within context limits. "
+                "The conversation is summarized below:\n\n"
+                f"{self.compressed_summary}"
+            )
+            system_prompt = system_prompt + compression_context
+
+        context_limit = get_token_limit(
+            self.model_id, user_id=self.model_user_id or self.user
+        )
+        system_tokens = num_tokens_from_string(system_prompt)
+
+        safety_buffer = int(context_limit * 0.1)
+        available_after_system = context_limit - system_tokens - safety_buffer
+
+        max_query_tokens = int(available_after_system * 0.8)
+        query_tokens = num_tokens_from_string(query)
+
+        if query_tokens > max_query_tokens:
+            query = self._truncate_text_middle(query, max_query_tokens)
+            query_tokens = num_tokens_from_string(query)
+
+        available_for_history = max(available_after_system - query_tokens, 0)
+
+        working_history = self._truncate_history_to_fit(
+            self.chat_history,
+            available_for_history,
+        )
+
+        messages = [{"role": "system", "content": system_prompt}]
+
+        for i in working_history:
+            if "prompt" in i and "response" in i:
+                messages.append({"role": "user", "content": i["prompt"]})
+                asst_msg: Dict[str, Any] = {
+                    "role": "assistant",
+                    "content": i["response"],
+                }
+                # Persisted thought from the prior turn rides along as
+                # reasoning_content so providers that require it on the
+                # follow-up call (DeepSeek thinking mode) accept the
+                # request. Other OpenAI-compatible APIs ignore the field.
+                if i.get("thought"):
+                    asst_msg["reasoning_content"] = i["thought"]
+                messages.append(asst_msg)
+            if "tool_calls" in i:
+                for tool_call in i["tool_calls"]:
+                    call_id = tool_call.get("call_id") or str(uuid.uuid4())
+                    args = tool_call.get("arguments")
+                    args_str = (
+                        json.dumps(args)
+                        if isinstance(args, dict)
+                        else (args or "{}")
+                    )
+                    messages.append({
+                        "role": "assistant",
+                        "content": None,
+                        "tool_calls": [{
+                            "id": call_id,
+                            "type": "function",
+                            "function": {
+                                "name": tool_call.get("action_name", ""),
+                                "arguments": args_str,
+                            },
+                        }],
+                    })
+                    result = tool_call.get("result")
+                    result_str = (
+                        json.dumps(result)
+                        if not isinstance(result, str)
+                        else (result or "")
+                    )
+                    messages.append({
+                        "role": "tool",
+                        "tool_call_id": call_id,
+                        "content": result_str,
+                    })
+        # When the request was multimodal, send the full content array (text +
+        # image_url parts) so images reach the model; the text-only `query` above
+        # is used only for token budgeting / retrieval.
+        user_content = (
+            self.multimodal_content
+            if getattr(self, "multimodal_content", None)
+            else query
+        )
+        messages.append({"role": "user", "content": user_content})
+        return messages
+
+    def _truncate_history_to_fit(
+        self,
+        history: List[Dict],
+        max_tokens: int,
+    ) -> List[Dict]:
+        from application.utils import num_tokens_from_string
+
+        if not history or max_tokens <= 0:
+            return []
+
+        truncated = []
+        current_tokens = 0
+
+        for message in reversed(history):
+            message_tokens = 0
+
+            if "prompt" in message and "response" in message:
+                message_tokens += num_tokens_from_string(message["prompt"])
+                message_tokens += num_tokens_from_string(message["response"])
+
+            if "tool_calls" in message:
+                for tool_call in message["tool_calls"]:
+                    tool_str = (
+                        f"Tool: {tool_call.get('tool_name')} | "
+                        f"Action: {tool_call.get('action_name')} | "
+                        f"Args: {tool_call.get('arguments')} | "
+                        f"Response: {tool_call.get('result')}"
+                    )
+                    message_tokens += num_tokens_from_string(tool_str)
+
+            if current_tokens + message_tokens <= max_tokens:
+                current_tokens += message_tokens
+                truncated.insert(0, message)
+            else:
+                break
+
+        if len(truncated) < len(history):
+            logger.info(
+                f"Truncated chat history from {len(history)} to {len(truncated)} messages "
+                f"to fit within {max_tokens:,} token budget"
+            )
+
+        return truncated
+
+    # ---- LLM generation ----
+
+    def _llm_gen(self, messages: List[Dict], log_context: Optional[LogContext] = None):
+        self._validate_context_size(messages)
+
+        # Use the upstream id resolved by LLMCreator (see __init__).
+        # Built-in models: same as self.model_id. BYOM: the user's
+        # typed model name, not the internal UUID.
+        gen_kwargs = {"model": self.upstream_model_id, "messages": messages}
+        if self.attachments:
+            gen_kwargs["_usage_attachments"] = self.attachments
+
+        if (
+            hasattr(self.llm, "_supports_tools")
+            and self.llm._supports_tools
+            and self.tools
+        ):
+            gen_kwargs["tools"] = self.tools
+        if (
+            self.json_schema
+            and hasattr(self.llm, "_supports_structured_output")
+            and self.llm._supports_structured_output()
+        ):
+            structured_format = self.llm.prepare_structured_output_format(
+                self.json_schema, strict=getattr(self, "json_schema_strict", True)
+            )
+            if structured_format:
+                if self.llm_name == "openai":
+                    gen_kwargs["response_format"] = structured_format
+                elif self.llm_name == "google":
+                    gen_kwargs["response_schema"] = structured_format
+        elif (
+            getattr(self, "json_object", False)
+            and self.llm_name == "openai"
+            and hasattr(self.llm, "_supports_structured_output")
+            and self.llm._supports_structured_output()
+        ):
+            # OpenAI json_object mode: guarantee valid JSON, no schema enforcement.
+            gen_kwargs["response_format"] = {"type": "json_object"}
+        if (
+            settings.OPENAI_RESPONSES_STORE
+            and hasattr(self.llm, "_uses_responses_api")
+            and self.llm._uses_responses_api()
+        ):
+            previous_response_id = self._previous_response_id()
+            if previous_response_id:
+                gen_kwargs["previous_response_id"] = previous_response_id
+
+        # Forward OpenAI sampling params (temperature, max_tokens, top_p, ...).
+        if self.llm_params:
+            gen_kwargs.update(self.llm_params)
+        resp = self.llm.gen_stream(**gen_kwargs)
+
+        if log_context:
+            data = build_stack_data(self.llm, exclude_attributes=["client"])
+            log_context.stacks.append({"component": "llm", "data": data})
+        return resp
+
+    def _llm_handler(
+        self,
+        resp,
+        tools_dict: Dict,
+        messages: List[Dict],
+        log_context: Optional[LogContext] = None,
+        attachments: Optional[List[Dict]] = None,
+    ):
+        resp = self.llm_handler.process_message_flow(
+            self, resp, tools_dict, messages, attachments, True
+        )
+        if log_context:
+            data = build_stack_data(self.llm_handler, exclude_attributes=["tool_calls"])
+            log_context.stacks.append({"component": "llm_handler", "data": data})
+        return resp
+
+    def _handle_response(self, response, tools_dict, messages, log_context):
+        is_structured_output = (
+            self.json_schema is not None
+            and hasattr(self.llm, "_supports_structured_output")
+            and self.llm._supports_structured_output()
+        )
+
+        if isinstance(response, str):
+            answer_data = {"answer": response}
+            if is_structured_output:
+                answer_data["structured"] = True
+                answer_data["schema"] = self.json_schema
+            yield answer_data
+            return
+        if hasattr(response, "message") and getattr(response.message, "content", None):
+            answer_data = {"answer": response.message.content}
+            if is_structured_output:
+                answer_data["structured"] = True
+                answer_data["schema"] = self.json_schema
+            yield answer_data
+            return
+        processed_response_gen = self._llm_handler(
+            response, tools_dict, messages, log_context, self.attachments
+        )
+
+        for event in processed_response_gen:
+            if isinstance(event, str):
+                answer_data = {"answer": event}
+                if is_structured_output:
+                    answer_data["structured"] = True
+                    answer_data["schema"] = self.json_schema
+                yield answer_data
+            elif hasattr(event, "message") and getattr(event.message, "content", None):
+                answer_data = {"answer": event.message.content}
+                if is_structured_output:
+                    answer_data["structured"] = True
+                    answer_data["schema"] = self.json_schema
+                yield answer_data
+            elif isinstance(event, dict) and "type" in event:
+                yield event