"""Shared base for all Enterprise Brain agents. By default agents use the gateway model string directly (no extra provider). Pass ``api_base`` + optional ``api_key`` to route through a LiteLLM proxy instead. ``EnterpriseAgent`` subclasses ``PydanticAIResponder`` (pylogue), inheriting streaming, tool-rendering, system-prompt management and conversation history out of the box. Callers can use it directly as a pylogue responder, or call ``agent.run(...)`` / ``agent.tool_plain(...)`` as before. """ import contextvars import os import re from typing import Any, Callable import logfire from genai_prices import data_snapshot from genai_prices.types import ClauseContains, ClauseEquals, ModelInfo, Provider from pydantic_ai import Agent from pydantic_ai.models.openai import OpenAIChatModel from pydantic_ai.providers.litellm import LiteLLMProvider from pydantic_ai.providers.openai import OpenAIProvider from pylogue.integrations.pydantic_ai import PydanticAIResponder DEFAULT_MODEL_ID = "gateway/openai:gpt-5-mini" DEFAULT_LITELLM_MODEL_ID = os.environ.get("LITELLM_PROVIDER_MODEL_NAME", "gemini/gemini-2.5-flash") DEFAULT_MLX_MODEL = os.environ.get("MLX_MODEL", "") # Keep it as mlx-community/Qwen3.5-9B-MLX-4bit in your env variable DEFAULT_MLX_BASE_URL = os.environ.get("MLX_BASE_URL", "http://localhost:8080/v1") # --------------------------------------------------------------------------- # Pylogue session detection # --------------------------------------------------------------------------- # Set to True for the exact duration of a PydanticAIResponder.__call__ invocation # (i.e. a pylogue WebSocket message). All other call paths (agent.run() from REST # handlers) never touch this var so it stays False. _pylogue_active: contextvars.ContextVar[bool] = contextvars.ContextVar( "_pylogue_active", default=False ) def _register_litellm_model_prices() -> None: """Register a LiteLLM wrapper provider so genai-prices can cost custom deployments. LiteLLM deployments are typically named like ``"divami-gemini/gemini-2.5-flash"``, wrapping a well-known base model (``"gemini-2.5-flash"`` — the part after the last ``/``). This function: 1. Strips the base model name from the deployment string. 2. Looks it up in the bundled genai-prices database. 3. Registers a custom provider (matched by ``provider_id="litellm"``) that claims any model name containing the base name and reuses its pricing. No prices are hardcoded — they are always pulled from the bundled data. """ deployment = DEFAULT_LITELLM_MODEL_ID # e.g. "gemini/gemini-2.5-flash" base_model = deployment.split("/")[-1].lower() # e.g. "gemini-2.5-flash" # Resolve pricing from the bundled genai-prices data for the base model name. snap = data_snapshot.get_snapshot() resolved_prices = None for provider in snap.providers: model = provider.find_model(base_model) if model is not None: resolved_prices = model.prices break if resolved_prices is None: import warnings warnings.warn( f"genai-prices: could not resolve pricing for base model '{base_model}' " f"(derived from '{deployment}'). Cost will not be tracked.", stacklevel=2, ) return # Match any model name that contains the base model name — this covers both # the request-side name (e.g. "gemini/gemini-2.5-flash") and the response-side # deployment name LiteLLM echoes back (e.g. "divami-gemini/gemini-2.5-flash"). model_match = ClauseContains(contains=base_model) api_base = os.environ.get("LITELLM_API_BASE", "") litellm_provider = Provider( id="litellm-wrapper", name="LiteLLM Wrapper", api_pattern=re.escape(api_base) if api_base else r"(?!)", # LiteLLMProvider.name always returns "litellm" — match on that. provider_match=ClauseEquals(equals="litellm"), model_match=model_match, models=[ ModelInfo( id=deployment, match=model_match, prices=resolved_prices, ) ], ) data_snapshot.set_custom_snapshot( data_snapshot.DataSnapshot( providers=[litellm_provider, *snap.providers], from_auto_update=snap.from_auto_update, ) ) _register_litellm_model_prices() class MLXChatModel(OpenAIChatModel): """``OpenAIChatModel`` variant for ``mlx_lm.server``. ``mlx_lm.server`` requires the system message to be strictly the *first* message in the list and rejects requests that have system messages scattered through conversation history (which pydantic-ai produces when message_history is used across turns). This subclass post-processes the mapped messages to: 1. Collect all ``role="system"`` messages. 2. Keep only the last one (the most-current instructions). 3. Hoist it to position 0. """ async def _map_messages(self, messages, model_request_parameters): # type: ignore[override] openai_messages = await super()._map_messages(messages, model_request_parameters) system_msgs = [m for m in openai_messages if m.get("role") == "system"] non_system_msgs = [m for m in openai_messages if m.get("role") != "system"] if system_msgs: # Keep the last system message — it carries the current instructions. return [system_msgs[-1], *non_system_msgs] return non_system_msgs def make_mlx_model( model_name: str = "mlx-community/Qwen3.5-9B-MLX-4bit", base_url: str = DEFAULT_MLX_BASE_URL, ) -> MLXChatModel: """Create an MLX-backed model via mlx_lm.server's OpenAI-compatible endpoint. Prerequisite: an instance of ``mlx_lm.server`` running and serving the desired model at the specified URL (default: ``http://localhost:8080/v1``). You can run it locally using `mlx_lm.server --model mlx-community/Qwen3.5-9B-MLX-4bit --host 0.0.0.0` once you `pip install mlx_lm`. The Qwen model will be automatically downloaded on first run if not present locally. """ return MLXChatModel(model_name=model_name, provider=OpenAIProvider(base_url=base_url)) class EnterpriseAgent(PydanticAIResponder): """Pydantic-AI agent with LiteLLM routing and full pylogue responder support. Inherits from ``PydanticAIResponder``, so every instance is directly usable as a pylogue streaming responder (``await agent("prompt", context=ctx)``), with streaming, tool-rendering, system-prompt management, and conversation history all included. ``pylogue_instructions`` (mermaid/HTML hints) are injected **only** when the agent is invoked via the pylogue ``/chat`` WebSocket — i.e. when ``set_context()`` has been called by pylogue's session machinery. REST API calls from ``/query`` and ``/query/stream`` bypass ``set_context`` entirely, so they receive a clean system prompt with no mermaid or HTML instructions. Args: instructions: The agent's persona / system instructions. model_id: Model identifier when NOT routing through LiteLLM or MLX (default: ``"gateway/openai:gpt-5-mini"``). api_base: When provided, requests are routed through the LiteLLM proxy at this URL. api_key: API key for the proxy; falls back to ``LITELLM_API_KEY`` env var. mlx_model: When provided (or ``MLX_MODEL`` env var is set), requests are served by a local ``mlx_lm.server`` instance. Takes precedence over ``model_id`` but not ``api_base``. mlx_base_url: mlx_lm.server OpenAI-compatible endpoint; falls back to ``MLX_BASE_URL`` env var or ``http://localhost:8080/v1``. logfire_env: When provided, logfire is configured with this environment name and pydantic-ai instrumentation is enabled automatically. agent_deps: Optional deps object forwarded to the underlying pydantic_ai.Agent runs (passed to PydanticAIResponder). show_tool_details: Whether to render full tool call/result details in the stream (passed to PydanticAIResponder). retries: Number of retries for the underlying agent. """ def _compose_system_prompt(self, user=None) -> str: """Inject pylogue instructions only during pylogue WebSocket sessions. ``_pylogue_active`` is a ContextVar set to ``True`` for the exact duration of ``__call__`` (the pylogue WS message handler). REST API calls via ``agent.run()`` never enter ``__call__``, so the var stays ``False`` and they receive a clean system prompt. """ from pylogue.integrations.common import ( PYLOGUE_INSTRUCTIONS, compose_system_prompt as _csp, ) return _csp( base_prompt=self._prompt_state.get("base_prompt", ""), additional_instructions=list(self._prompt_state.get("additional", [])), user=user, pylogue_instructions=PYLOGUE_INSTRUCTIONS if _pylogue_active.get() else "", ) async def __call__(self, text: str, context=None): # type: ignore[override] """Pylogue WebSocket message handler. Sets ``_pylogue_active`` for the exact duration of this call so that ``_compose_system_prompt`` injects pylogue instructions for WS sessions but not for REST API calls via ``agent.run()``. """ token = _pylogue_active.set(True) try: async for chunk in super().__call__(text, context): yield chunk finally: _pylogue_active.reset(token) def __init__( self, instructions: str, *, model_id: str = DEFAULT_MODEL_ID, api_base: str | None = None, api_key: str | None = None, mlx_model: str | None = None, mlx_base_url: str = DEFAULT_MLX_BASE_URL, logfire_env: str | None = None, service_name: str | None = None, agent_deps: Any = None, show_tool_details: bool = True, retries: int = 1, deps_type: type | None = None, output_type: type | None = None, capabilities: list | None = None, ) -> None: if logfire_env: logfire.configure(environment=logfire_env, service_name=service_name, scrubbing=False) logfire.instrument_pydantic_ai() if effective_mlx := (mlx_model or DEFAULT_MLX_MODEL): model = make_mlx_model(effective_mlx, mlx_base_url) elif api_base: resolved_key = api_key or os.getenv("LITELLM_API_KEY") model = OpenAIChatModel( DEFAULT_LITELLM_MODEL_ID, provider=LiteLLMProvider(api_base=api_base, api_key=resolved_key), ) else: model = model_id # pydantic-ai resolves the string model directly pydantic_agent = Agent( model, instructions=instructions, retries=retries, **(dict(deps_type=deps_type) if deps_type is not None else {}), **(dict(output_type=output_type) if output_type is not None else {}), **(dict(capabilities=capabilities) if capabilities is not None else {}), ) # For pylogue (non-A2A) calls: provide an empty deps instance so tools # that use RunContext[RequestContext] don't receive None. if deps_type is not None and agent_deps is None: agent_deps = deps_type() # PydanticAIResponder takes over from here: stores self.agent, # registers the dynamic system-prompt function, etc. super().__init__(pydantic_agent, agent_deps=agent_deps, show_tool_details=show_tool_details) # ------------------------------------------------------------------ # Explicit agent-level helpers (not on PydanticAIResponder) # ------------------------------------------------------------------ def tool_plain(self, func: Callable) -> Callable: """Register a plain (no RunContext) tool on the underlying agent.""" return self.agent.tool_plain(func) def tool(self, func: Callable) -> Callable: """Register a tool with RunContext on the underlying agent.""" return self.agent.tool(func) def system_prompt(self, func: Callable) -> Callable: """Register a dynamic system-prompt function on the underlying agent.""" return self.agent.system_prompt(func) async def run(self, *args: Any, **kwargs: Any): """Run the agent with automatic conversation-history threading. Passes ``self.message_history`` into the underlying ``Agent.run`` call (unless the caller supplies their own ``message_history`` kwarg) and updates it from the result, so successive ``run()`` calls share context. """ kwargs.setdefault("message_history", self.message_history) result = await self.agent.run(*args, **kwargs) self.message_history = result.all_messages() return result def to_a2a(self, **kwargs: Any): """Expose the agent as an A2A ASGI app (delegates to ``Agent.to_a2a``).""" return self.agent.to_a2a(**kwargs)