"""Helper methods for initializing Oracle Cloud GenAI LLMs."""
from contextlib import contextmanager
from typing import Any, AsyncIterator, AsyncGenerator, Iterable, List, Dict, Optional, Tuple, Union
from aidputils.agents.auth.client.generative_ai_inference_v2_client import GenerativeAiInferenceV2Client
from aidputils.agents.auth.util import auth_utils
from aidputils.agents.toolkit.configs import OCIAIConf, InvokeConfig, LangGraphRunnableConfig
from aidputils.agents.guardrails.guarded_chat_oci_genai import GuardedChatOCIGenAI
from aidputils.agents.tools.mcp.mcp_client import MCPHTTPClient
from aidputils.agents.tools.utils import SystemUtils
from aidputils.agents.auth.util.agent_util import _run_in_event_loop
from aidputils.agents.toolkit import chat_context
from langchain.agents.middleware import dynamic_prompt, ModelRequest
import logging
_genai_client_cache = {}
logger = logging.getLogger('agent_helper')
def _is_handoff_marker_message(msg: object) -> bool:
"""Best-effort detection of supervisor/worker handoff marker messages.
We add these messages in aidputils.agents.toolkit.handoff.create_handoff_back_messages
and they include either special response_metadata or well-known text.
"""
try:
from aidputils.agents.toolkit.handoff import METADATA_KEY_IS_HANDOFF_BACK
except Exception:
METADATA_KEY_IS_HANDOFF_BACK = "__is_handoff_back" # type: ignore
# langchain_core BaseMessage-ish objects typically have: content, response_metadata, name
content = getattr(msg, "content", None)
response_metadata = getattr(msg, "response_metadata", None)
if isinstance(response_metadata, dict) and response_metadata.get(METADATA_KEY_IS_HANDOFF_BACK):
return True
if isinstance(content, str):
# Defensive: UI output shows this exact text
if "Transferring back to" in content:
return True
if "Successfully transferred back to" in content:
return True
if "Successfully transferred to" in content:
return True
return False
def _message_text_key(msg: object) -> str:
"""Create a stable-ish key for deduping UI output."""
name = getattr(msg, "name", None)
content = getattr(msg, "content", None)
if isinstance(content, list):
# Rare: some providers use multi-part content; stringify for keying
content = str(content)
return f"{name}:{content}" # good enough for UI streaming
class Constants:
AUTH_TYPE = "auth_type"
# Tool and memory config
DH_USER_PRINCIPAL_KEY = "dh-user-principal"
DH_USER_PRINCIPAL_KEY_ALTERNATE = "dh_user_principal"
DH_USER_PRINCIPAL_ENV_KEY = "DH_USER_PRINCIPAL"
SESSION_ID = "session_id"
THREAD_ID = "thread_id"
CHECKPOINT_NS = "checkpoint_ns"
SESSION_VARIABLES = "session_variables"
[docs]
def init_oci_llm(llm_conf: OCIAIConf):
"""
Initialize GuardedChatOCIGenAI language model instance using the provided OCIAIConf.
Args:
llm_conf (OCIAIConf): The configuration object specifying model, provider, endpoint, compartment, authentication profile, model args, and optional guardrails config.
Returns:
GuardedChatOCIGenAI: Instantiated language model object ready for use with the specified settings.
"""
print (llm_conf.model_dump)
logger.info("Initializing GuardedChatOCIGenAI with guardrails configuration")
chat = GuardedChatOCIGenAI(
guardrails_config=llm_conf.guardrails_config,
auth_type=llm_conf.auth_type,
model_id=llm_conf.model_id,
provider=llm_conf.model_provider,
service_endpoint=llm_conf.endpoint,
compartment_id=llm_conf.compartment_id,
auth_profile=llm_conf.auth_profile,
client=get_client(llm_conf=llm_conf),
model_kwargs=llm_conf.model_args,
is_stream=True
)
return chat
[docs]
def get_client(llm_conf):
auth_type = llm_conf.auth_type
signer = SystemUtils.get_signer(
signer_type=auth_type.lower(),
config_profile=llm_conf.auth_profile
)
endpoint = llm_conf.endpoint
key = (endpoint, id(signer))
global _genai_client_cache
if key in _genai_client_cache:
return _genai_client_cache[key]
client = GenerativeAiInferenceV2Client(endpoint=endpoint, signer=signer)
_genai_client_cache[key] = client
return client
[docs]
def pre_invoke_setup(**kwargs) -> dict:
invoke_config = InvokeConfig(
thread_id=kwargs.get(Constants.THREAD_ID),
checkpoint_ns=kwargs.get(Constants.CHECKPOINT_NS),
)
# Get session variables from chat_context and filter by shouldLog flag
session_vars = chat_context.session_context_var.get()
if isinstance(session_vars, dict):
filtered = {
k: v
for k, v in session_vars.items()
if not (isinstance(v, dict) and v.get("shouldLog") is False)
}
invoke_config.session_variables = filtered
return LangGraphRunnableConfig(configurable=invoke_config).model_dump()
[docs]
def post_tool_setup(token, mcp_clients: List[MCPHTTPClient] | None = None, kwargs: dict | None = None):
"""
Synchronous cleanup after tool execution.
- Always resets the auth context using the provided token.
- If mcp_clients and session_id are provided, synchronously stops each MCP client session,
bridging async calls via _run_in_event_loop to work whether an event loop is running or not.
"""
if mcp_clients and kwargs is not None and kwargs.get(Constants.THREAD_ID) is not None:
session_id = kwargs.get(Constants.THREAD_ID)
for mcp_client in mcp_clients:
try:
_run_in_event_loop(mcp_client.stop_session(session_id=session_id))
except Exception:
# Best-effort cleanup; ignore individual client stop failures
logger.exception("Failed to stop MCP client session for session %s", session_id)
pass
[docs]
async def parse_stream_response(stream):
"""
Parse and yield message chunks from a LangGraph streaming response.
Yields only discrete message chunks (BaseMessage or equivalent) from events.
"""
async for event in stream:
if isinstance(event, tuple) and len(event) == 2:
chunk, metadata = event
yield chunk
continue
# Fallback: dict format {node_name: {"messages": [...]}}
if isinstance(event, dict):
for node_name, node_out in event.items():
if isinstance(node_out, dict) and "messages" in node_out:
msgs = node_out["messages"]
if isinstance(msgs, list):
for msg in msgs:
yield msg
continue
# Ignore unknown event types
[docs]
async def parse_stream_response_for_supervisor(
stream: AsyncIterator[object],
*,
supervisor_name: str,
include_intermediate: bool = False,
) -> AsyncIterator[object]:
"""UI-friendly stream parser.
The default LangGraph `stream_mode="messages"` will emit *all* messages from *all* nodes,
including worker outputs and handoff markers. Many chat UIs then concatenate those into
one transcript, causing duplicated/verbose output and slow rendering.
This generator filters the stream to yield only end-user facing messages.
Rules:
- Drop ToolMessage (tool chatter)
- Drop any handoff marker messages
- By default, only emit AI messages from the supervisor node (by `.name`)
- Deduplicate repeated messages
Args:
stream: The raw async stream from `agent.astream(...)`
supervisor_name: The graph node name used for the supervisor agent.
include_intermediate: If True, also allow non-supervisor AI messages (still filtered
for handoff/tool chatter). Default False.
"""
try:
from langchain_core.messages import AIMessage, ToolMessage
except Exception: # pragma: no cover
AIMessage = object # type: ignore
ToolMessage = object # type: ignore
seen: set[str] = set()
last_supervisor_msg: object | None = None
async for msg in parse_stream_response(stream):
# Filter tool chatter early
if isinstance(msg, ToolMessage):
continue
if _is_handoff_marker_message(msg):
continue
# Only expose supervisor AI outputs by default.
# Some providers/versions may not populate `name` on AIMessage. In that
# case, we fall back to "last AIMessage wins" at the end.
if not include_intermediate:
if isinstance(msg, AIMessage):
if (getattr(msg, "name", None) or "") == supervisor_name:
last_supervisor_msg = msg
else:
# keep tracking a last AI message, but don't emit it immediately
# unless it matches supervisor_name
pass
if not isinstance(msg, AIMessage):
continue
if (getattr(msg, "name", None) or "") != supervisor_name:
continue
key = _message_text_key(msg)
if key in seen:
continue
seen.add(key)
yield msg
# Fallback: if no supervisor-named message was emitted, yield the last AIMessage
# we saw (still filtered of tool/handoff). This prevents "no messages" failure
# modes when message.name isn't set by the backend.
if not seen and last_supervisor_msg is not None:
yield last_supervisor_msg
[docs]
async def stream_messages(
self,
input,
config: dict | None = None,
kwargs: dict | None = None,
) -> AsyncGenerator[Any, None]:
"""
Stream messages as they are generated by the agent.
Yields BaseMessage objects as new messages are added by nodes.
Ensures the auth context (ContextVar) remains active during streaming and is cleaned up after.
"""
logger.info("Starting _stream_messages")
kwargs = kwargs or {}
config = config or pre_invoke_setup(**kwargs)
token = pre_tool_setup(**kwargs)
try:
# Use a single user-facing attribute: `self.agent`.
# It may be either:
# - A LangGraph compiled agent (has `.astream`)
# - An AgentShim (has `.adapter`)
agent = getattr(self, "agent", None)
if agent is None:
raise RuntimeError("Agent not initialized. Call setup() first.")
# If the agent is an AgentShim, ensure its adapter is initialized and
# stream from the underlying framework agent.
if hasattr(agent, "adapter"):
if not getattr(agent, "_setup_complete", False) and getattr(agent, "adapter", None) is not None:
await agent.adapter.setup()
agent._setup_complete = True
agent = getattr(getattr(agent, "adapter", None), "agent", None)
if agent is None:
raise RuntimeError("AgentShim adapter not initialized. Call setup() first.")
# Determine streaming mode based on availability of StreamingData class
try:
from aidputils.agents.auth.client.generative_ai_inference_v2_client import StreamingData
stream = agent.astream(input=input, config=config, stream_mode="messages")
except ImportError:
stream = agent.astream(input=input, config=config)
async for chunk in parse_stream_response(stream):
yield chunk
except Exception:
logger.exception("Streaming error")
raise
finally:
# Cleanup auth context after streaming completes or errors out
# Backward compatible: kwargs may (or may not) contain mcp_clients.
post_tool_setup(token, mcp_clients=kwargs.get("mcp_clients"), kwargs=kwargs)
# --------------------------------------------------------------------------------------
# Backward-compatible merged setup helper
# --------------------------------------------------------------------------------------
[docs]
class FlowInvokeSetup:
"""Backward-compatible wrapper for pre/post helper methods.
This class exists to provide a single, higher-level API that composes:
- pre_invoke_setup(**kwargs) -> config
- pre_tool_setup(**kwargs) -> token
- post_tool_setup(token, mcp_clients=..., kwargs=...) + auth context reset
Existing functions remain unchanged for backward compatibility.
"""
def __init__(self, *, token, config: dict, kwargs: dict):
self.token = token
self.config = config
self._kwargs = kwargs
self._cleaned_up = False
[docs]
def cleanup(self, mcp_clients: List[MCPHTTPClient] | None = None):
"""Run the same cleanup as post_tool_setup and reset auth context.
Idempotent: safe to call multiple times.
"""
if self._cleaned_up:
return
try:
post_tool_setup(self.token, mcp_clients=mcp_clients, kwargs=self._kwargs)
finally:
# Ensure auth context is reset even if MCP cleanup fails.
try:
auth_utils.reset_auth_context(self.token)
except Exception:
logger.exception("Failed to reset auth context")
self._cleaned_up = True
[docs]
@contextmanager
def flow_setup(*, mcp_clients: List[MCPHTTPClient] | None = None, **kwargs):
"""Context-manager wrapper around setup_tool_invoke().
This enables simplified user code:
with flow_setup(mcp_clients=mcp_clients, **kwargs) as config:
agent_response = await agent.ainvoke(input=message, config=config)
Yields:
dict: The invocation config (same as pre_invoke_setup).
"""
setup = setup_tool_invoke(**kwargs)
try:
yield setup.config
finally:
setup.cleanup(mcp_clients=mcp_clients)
[docs]
def fill_template(template: str, variables: Dict | None) -> str:
import re
# Support {{...}} placeholders (with dotted keys, e.g., {{sessionvariables.tone}})
return re.sub(
r"\{\{\s*([^{}]+?)\s*\}\}",
lambda m: str((variables or {}).get(m.group(1).strip(), m.group(0))),
template,
)
[docs]
def fill_template_from_session_vars(template: str, variables: Dict | None = None) -> str:
"""
Build a prompt by merging variables from session context with provided variables and defaults.
- template uses Python str.format placeholders, supports dotted keys like {sessionvariables.tone}
- variables is optional additional override dict; its keys take precedence over session vars
"""
session_vars = chat_context.session_context_var.get()
merged: Dict = {}
if isinstance(session_vars, dict):
# Normalize stored session variables to their 'value' field (mirror PromptTool behavior)
normalized = {sk: (sv.get("value") if isinstance(sv, dict) else sv) for sk, sv in session_vars.items()}
merged.update(normalized)
if variables:
for k, v in variables.items():
if isinstance(v, (str, int, float, bool)):
merged[k] = v
return fill_template(template, merged)
[docs]
def create_dynamic_sys_prompt_middleware(template: str):
# This closure captures 'template' so the middleware can use it later
@dynamic_prompt
def sys_prompt_middleware(req: ModelRequest) -> str:
# Use your existing rendering logic here
return fill_template_from_session_vars(template)
return sys_prompt_middleware
[docs]
def create_get_dynamic_messages(system_prompt_template: str):
"""
Factory that returns a get_dynamic_messages(state) function compatible with LangGraph's prompt callback.
The returned function:
- Renders a SystemMessage from system_prompt_template using session vars merged with provided defaults
- Prepends that SystemMessage to the existing history in state['messages']
- Returns [SystemMessage] + history
"""
def get_dynamic_messages(state: Dict):
# Render system prompt text using session variables + provided defaults
rendered = fill_template_from_session_vars(system_prompt_template)
# Construct a SystemMessage if langchain is available; otherwise provide a dict fallback
try:
from langchain_core.messages import SystemMessage as _SM
sys_msg = _SM(content=rendered)
except Exception:
sys_msg = {"role": "system", "content": rendered}
# Retrieve existing history list from state
history = state.get("messages", []) if isinstance(state, dict) else []
if not isinstance(history, list):
history = []
# Prepend system message to history
return [sys_msg] + list(history)
return get_dynamic_messages