aip-agents-binary 0.0.0__py3-none-any.whl

aip_agents/agent/model_invocation_adapter.py

@@ -0,0 +1,346 @@
+"""Private model invocation adapters for LangGraph ReAct agents."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Awaitable, Callable, Sequence
+from dataclasses import dataclass
+from typing import Any, cast
+
+from gllm_core.event import EventEmitter  # type: ignore[import-untyped]
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import AIMessage, BaseMessage, SystemMessage
+from langchain_core.tools import BaseTool
+
+from aip_agents.context.runtime import detect_context_overflow
+from aip_agents.schema.model_id import ModelId, ModelProvider
+from aip_agents.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+BaseMessageToDict = Callable[[BaseMessage], dict[str, Any]]
+DictToBaseMessage = Callable[[dict[str, Any]], BaseMessage]
+ModifyModelRequestHook = Callable[[list[Any], str, dict[str, Any]], tuple[list[Any], str]]
+CompactConversation = Callable[[list[dict[str, Any]], str], Awaitable[tuple[list[dict[str, Any]], str]]]
+ApplyCompaction = Callable[..., Awaitable[tuple[list[dict[str, Any]], Any]]]
+ConvertLangchainMessages = Callable[[list[BaseMessage], str], Any]
+ConvertLmOutput = Callable[[Any], AIMessage]
+PrepareLmInvokerTools = Callable[[Sequence[BaseTool]], list[Any]]
+SanitizeMessages = Callable[[list[BaseMessage]], list[BaseMessage]]
+EnsureInstructionSync = Callable[[list[BaseMessage], str], list[BaseMessage]]
+ProviderAndModelName = Callable[[str | ModelId], tuple[str | None, str]]
+ProviderAndModelNameForInvoker = Callable[[Any], tuple[str | None, str]]
+PatchInvoker = Callable[[Any], None]
+IsOpenAIReasoningModel = Callable[[str], bool]
+
+
+@dataclass(slots=True)
+class ModelInvocationRequest:
+    """Inputs needed to invoke one ReAct model step."""
+
+    current_messages: Sequence[BaseMessage]
+    state: dict[str, Any]
+    state_with_config: dict[str, Any]
+    config: dict[str, Any] | None
+    enhanced_instruction: str
+
+
+@dataclass(slots=True)
+class ModelInvocationResult:
+    """Normalized output from a model invocation adapter."""
+
+    ai_message: AIMessage
+    model_provider: str | None
+    model_name: str | None
+    model_id: Any
+
+
+@dataclass(slots=True)
+class ModelInvocationContext:
+    """Narrow callback surface adapters need from the ReAct agent."""
+
+    agent_name: str
+    resolved_tools: Sequence[BaseTool]
+    compaction_config: dict[str, Any] | None
+    event_emitter: EventEmitter | None
+    model_provider: str | None
+    execute_modify_model_request_hook: ModifyModelRequestHook
+    compact_conversation_with_recall: CompactConversation
+    apply_compaction: ApplyCompaction
+    base_message_to_dict: BaseMessageToDict
+    dict_to_base_message: DictToBaseMessage
+    prepare_tools_for_lm_invoker: PrepareLmInvokerTools
+    sanitize_messages: SanitizeMessages
+    ensure_instruction_sync: EnsureInstructionSync
+    convert_langchain_messages_to_gllm_messages: ConvertLangchainMessages
+    convert_lm_output_to_langchain_message: ConvertLmOutput
+
+
+@dataclass(slots=True)
+class LMInvokerRuntimePolicy:
+    """Provider-specific runtime patch policy for LM invoker execution."""
+
+    provider_and_model_name: ProviderAndModelName
+    provider_and_model_name_for_invoker: ProviderAndModelNameForInvoker
+    is_openai_reasoning_model: IsOpenAIReasoningModel
+    patch_openai_reasoning_thinking_extraction: PatchInvoker
+    patch_anthropic_cache_usage: PatchInvoker
+
+    def apply_to_built_invoker(self, model: str | ModelId, lm_invoker: Any) -> None:
+        """Apply runtime patches to an invoker built from a model identifier."""
+        provider, model_name = self.provider_and_model_name(model)
+        self._apply(provider, model_name, lm_invoker)
+
+    def apply_to_provided_invoker(self, lm_invoker: Any) -> None:
+        """Apply runtime patches to a caller-provided invoker instance."""
+        provider, model_name = self.provider_and_model_name_for_invoker(lm_invoker)
+        self._apply(provider, model_name, lm_invoker)
+
+    def _apply(self, provider: str | None, model_name: str, lm_invoker: Any) -> None:
+        if provider == ModelProvider.OPENAI and self.is_openai_reasoning_model(model_name):
+            self.patch_openai_reasoning_thinking_extraction(lm_invoker)
+        if provider == ModelProvider.ANTHROPIC:
+            self.patch_anthropic_cache_usage(lm_invoker)
+
+
+class LMInvokerModelInvocationAdapter:
+    """Invoke the existing gllm-inference LM invoker path."""
+
+    def __init__(self, context: ModelInvocationContext, lm_invoker: Any):
+        """Initialize the adapter with the shared context and invoker instance."""
+        self._context = context
+        self._lm_invoker = lm_invoker
+
+    async def invoke(self, request: ModelInvocationRequest) -> ModelInvocationResult:
+        """Invoke the LM invoker and return a normalized result."""
+        modified_messages, enhanced_instruction = self._context.execute_modify_model_request_hook(
+            list(request.current_messages),
+            request.enhanced_instruction,
+            request.state_with_config,
+        )
+        base_messages = cast("list[BaseMessage]", modified_messages)
+
+        if self._context.compaction_config is not None:
+            message_dicts = [self._context.base_message_to_dict(message) for message in base_messages]
+            compacted_conversation, enhanced_instruction = await self._context.compact_conversation_with_recall(
+                message_dicts,
+                enhanced_instruction,
+            )
+            if compacted_conversation != message_dicts:
+                base_messages = [self._context.dict_to_base_message(message) for message in compacted_conversation]
+
+        messages = self._context.convert_langchain_messages_to_gllm_messages(
+            self._context.sanitize_messages(list(base_messages)),
+            enhanced_instruction,
+        )
+        effective_event_emitter = request.state.get("event_emitter") or self._context.event_emitter
+
+        if self._context.resolved_tools:
+            self._lm_invoker.set_tools(self._context.prepare_tools_for_lm_invoker(self._context.resolved_tools))
+
+        lm_output = await self._invoke(messages, effective_event_emitter, base_messages, enhanced_instruction)
+        ai_message = self._context.convert_lm_output_to_langchain_message(lm_output)
+        return ModelInvocationResult(
+            ai_message=ai_message,
+            model_provider=_clean_string(getattr(self._lm_invoker, "model_provider", None)),
+            model_name=_clean_string(getattr(self._lm_invoker, "model_name", None)),
+            model_id=getattr(self._lm_invoker, "model_id", None),
+        )
+
+    async def _invoke(
+        self,
+        messages: Any,
+        event_emitter: EventEmitter | None,
+        modified_messages: list[BaseMessage],
+        enhanced_instruction: str,
+    ) -> Any:
+        start_time = time.perf_counter()
+        logger.info(
+            "Agent '%s': LLM invoke start (tools=%d)",
+            self._context.agent_name,
+            len(self._context.resolved_tools),
+        )
+        total_chars = sum(len(str(getattr(message, "content", ""))) for message in messages)
+        logger.info("Prompt size: %d messages, %d chars", len(messages), total_chars)
+
+        try:
+            lm_output = await self._lm_invoker.invoke(
+                messages=messages,
+                event_emitter=event_emitter,
+            )
+        except Exception as exc:
+            if not detect_context_overflow(exc) or self._context.compaction_config is None:
+                raise
+            logger.info(
+                "Agent '%s': context overflow detected, retrying with compacted messages.",
+                self._context.agent_name,
+            )
+            message_dicts = [self._context.base_message_to_dict(message) for message in modified_messages]
+            compacted_dicts, _retry_prep = await self._context.apply_compaction(
+                message_dicts,
+                provider_overflow=True,
+            )
+            retry_messages = self._context.convert_langchain_messages_to_gllm_messages(
+                self._context.sanitize_messages(
+                    [self._context.dict_to_base_message(message) for message in compacted_dicts]
+                ),
+                enhanced_instruction,
+            )
+            lm_output = await self._lm_invoker.invoke(
+                messages=retry_messages,
+                event_emitter=event_emitter,
+            )
+
+        logger.info(
+            "Agent '%s': LLM invoke finished in %.3fs", self._context.agent_name, time.perf_counter() - start_time
+        )
+        return lm_output
+
+
+class LangChainModelInvocationAdapter:
+    """Invoke the existing LangChain BaseChatModel path."""
+
+    def __init__(self, context: ModelInvocationContext, model: BaseChatModel):
+        """Initialize the adapter with the shared context and chat model."""
+        self._context = context
+        self._model = model
+
+    async def invoke(self, request: ModelInvocationRequest) -> ModelInvocationResult:
+        """Invoke the LangChain chat model and return a normalized result."""
+        langchain_prompt: list[BaseMessage] = [SystemMessage(content=request.enhanced_instruction)] + list(
+            request.current_messages
+        )
+        langchain_prompt, enhanced_instruction = self._context.execute_modify_model_request_hook(
+            langchain_prompt,
+            request.enhanced_instruction,
+            request.state_with_config,
+        )
+        langchain_prompt = self._context.ensure_instruction_sync(
+            cast("list[BaseMessage]", langchain_prompt),
+            enhanced_instruction,
+        )
+
+        if self._context.compaction_config is not None:
+            langchain_prompt = await self._compact_prompt(langchain_prompt, enhanced_instruction)
+
+        model_with_tools = (
+            self._model.bind_tools(self._context.resolved_tools) if self._context.resolved_tools else self._model
+        )
+        ai_message = await self._invoke(model_with_tools, langchain_prompt, request.config)
+        return ModelInvocationResult(
+            ai_message=ai_message,
+            model_provider=_langchain_model_provider(self._model) or self._context.model_provider,
+            model_name=_langchain_model_name(self._model),
+            model_id=_langchain_model_id(self._model),
+        )
+
+    async def _compact_prompt(
+        self, langchain_prompt: list[BaseMessage], enhanced_instruction: str
+    ) -> list[BaseMessage]:
+        conversation_messages = [message for message in langchain_prompt if not isinstance(message, SystemMessage)]
+        message_dicts = [self._context.base_message_to_dict(message) for message in conversation_messages]
+        system_content = langchain_prompt[0].content if langchain_prompt else enhanced_instruction
+
+        compacted_conversation, system_content = await self._context.compact_conversation_with_recall(
+            message_dicts,
+            str(system_content),
+        )
+        if compacted_conversation != message_dicts:
+            return [SystemMessage(content=system_content)] + [
+                self._context.dict_to_base_message(message) for message in compacted_conversation
+            ]
+
+        langchain_prompt[0] = SystemMessage(content=system_content)
+        return langchain_prompt
+
+    async def _invoke(
+        self,
+        model_with_tools: BaseChatModel,
+        langchain_prompt: list[BaseMessage],
+        config: dict[str, Any] | None,
+    ) -> AIMessage:
+        start_time = time.perf_counter()
+        logger.info(
+            "Agent '%s': LangChain model invoke start (tools=%d)",
+            self._context.agent_name,
+            len(self._context.resolved_tools),
+        )
+
+        try:
+            ai_message = await model_with_tools.ainvoke(langchain_prompt, config)
+        except Exception as exc:
+            if not detect_context_overflow(exc) or self._context.compaction_config is None:
+                raise
+            logger.info(
+                "Agent '%s': context overflow detected, retrying with compacted messages.",
+                self._context.agent_name,
+            )
+            conversation_messages = [message for message in langchain_prompt if not isinstance(message, SystemMessage)]
+            message_dicts = [self._context.base_message_to_dict(message) for message in conversation_messages]
+            compacted_dicts, _retry_prep = await self._context.apply_compaction(
+                message_dicts,
+                provider_overflow=True,
+            )
+            retry_prompt = [langchain_prompt[0]] + [
+                self._context.dict_to_base_message(message) for message in compacted_dicts
+            ]
+            ai_message = await model_with_tools.ainvoke(retry_prompt, config)
+
+        logger.info(
+            "Agent '%s': LangChain model invoke finished in %.3fs",
+            self._context.agent_name,
+            time.perf_counter() - start_time,
+        )
+        return ai_message
+
+
+def _clean_string(value: Any) -> str | None:
+    return value.strip() or None if isinstance(value, str) else None
+
+
+def _langchain_model_provider(model: BaseChatModel) -> str | None:
+    provider = _clean_string(getattr(model, "model_provider", None))
+    if provider:
+        return provider
+
+    model_id = getattr(model, "model_id", None)
+    if isinstance(model_id, ModelId):
+        return _clean_string(model_id.provider)  # pragma: no cover
+    if isinstance(model_id, str) and "/" in model_id:
+        return _clean_string(model_id.split("/", 1)[0])  # pragma: no cover
+
+    module_and_name = f"{type(model).__module__}.{type(model).__name__}".lower()
+    if "azure" in module_and_name and "openai" in module_and_name:
+        return ModelProvider.AZURE_OPENAI  # pragma: no cover
+    if "openai" in module_and_name:
+        return ModelProvider.OPENAI  # pragma: no cover
+    if "anthropic" in module_and_name:
+        return ModelProvider.ANTHROPIC  # pragma: no cover
+    if "google" in module_and_name or "gemini" in module_and_name:
+        return ModelProvider.GOOGLE_GENAI  # pragma: no cover
+    return None
+
+
+def _langchain_model_name(model: BaseChatModel) -> str | None:
+    model_name = _clean_string(getattr(model, "model_name", None) or getattr(model, "model", None))
+    if model_name:
+        return model_name
+
+    model_id = getattr(model, "model_id", None)
+    if isinstance(model_id, ModelId):
+        return _clean_string(model_id.name)  # pragma: no cover
+    if isinstance(model_id, str):
+        if "/" in model_id:  # pragma: no cover
+            return _clean_string(model_id.rsplit("/", 1)[-1])  # pragma: no cover
+        if ":" in model_id:  # pragma: no cover
+            return _clean_string(model_id.rsplit(":", 1)[-1])  # pragma: no cover
+        return _clean_string(model_id)  # pragma: no cover
+    return None
+
+
+def _langchain_model_id(model: BaseChatModel) -> Any:
+    model_id = getattr(model, "model_id", None)
+    if isinstance(model_id, ModelId | str):
+        return model_id  # pragma: no cover
+    return None

aip_agents/agent/system_instruction_context.py

@@ -0,0 +1,34 @@
+"""Utilities for constructing agent system instructions.
+
+Authors:
+    Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
+"""
+
+from __future__ import annotations
+
+from aip_agents.utils.constants import DefaultTimezone
+from aip_agents.utils.datetime import get_timezone_aware_now
+
+__all__ = ["get_current_date_context", "DefaultTimezone"]
+
+
+def get_current_date_context(timezone: str = DefaultTimezone.JAKARTA) -> str:
+    """Generate current date context for system prompts.
+
+    Args:
+        timezone: IANA timezone name for date formatting.
+
+    Returns:
+        Formatted date context string for inclusion in system prompts.
+    """
+    now, timezone_label, used_fallback = get_timezone_aware_now(timezone)
+    display_timezone = "UTC" if used_fallback else timezone_label
+    current_date = now.strftime("%Y-%m-%d")
+
+    return (
+        f"Very important: The user's timezone is {display_timezone}. "
+        f"The current date is {current_date}\n\n"
+        "Any dates before this are in the past, and any dates after this are in the future. "
+        "When the user asks for the 'latest', 'most recent', 'today's', etc. don't assume\n"
+        "your knowledge is up to date."
+    )