Source code for aidputils.agents.tools.prompt
import oci
import pprint
import re
from aidputils.agents.toolkit.metrics_util import MetricUtility
from aidputils.agents.tools.base_tool import BaseTool
from aidputils.agents.auth.client.generative_ai_inference_v2_client import GenerativeAiInferenceV2Client
from aidputils.agents.tools.utils import Constants, SystemUtils
from aidputils.agents.toolkit import chat_context
# Global client cache: shared across all PromptTool instances and uses
_genai_client_cache = {}
CUSTOM_ENDPOINT_PREFIX = "ocid1.generativeaiendpoint"
[docs]
@BaseTool.register
class PromptTool(BaseTool):
"""
A tool for submitting prompts to OCI GenAI LLM, with configurable prompt templates and runtime parameters.
"""
metricsUtil = MetricUtility("PROMPT_TOOL")
success_counter = metricsUtil.create_counter("PROMPT_TOOL_SUCCESS_COUNTER")
failure_counter = metricsUtil.create_counter("PROMPT_TOOL_FAILURE_COUNTER")
# List of substrings to check for prompt injection, as a static class-level variable
_injection_keywords = ["ignore above", "forget", "reset", "```", "#", "[", "]"]
@classmethod
def _invoke_tool(cls, conf, runtime_params, **context_vars):
is_success = False
try:
tool_name = chat_context.tool_context_var.get()
cls.logger().info(f"context variable for tool name {tool_name}")
llm_conf = conf.get("llm", {})
auth_type = llm_conf.get(Constants.AUTH_TYPE, Constants.AUTH_TYPE_REMOTE)
auth_profile = llm_conf.get(Constants.AUTH_PROFILE, Constants.AUTH_PROFILE_DEFAULT)
signer = SystemUtils.get_signer(auth_type.lower(), config_profile=auth_profile, **context_vars)
prompt_template = conf.get("prompt_template", "")
model_id = llm_conf.get("model_id")
debug_info = {}
# Merge session variables into a copy of runtime_params (do not mutate original)
session_vars = chat_context.session_context_var.get()
merged_params = dict(runtime_params) if runtime_params else {}
if isinstance(session_vars, dict):
safe_dict = {k: v for k, v in session_vars.items() if not (isinstance(v, dict) and v.get("shouldLog") is False)}
cls.logger().info(f"loggable session_vars(filtered)={safe_dict}")
# Expect new structure: {"sessionvariables.foo": {"value": "...", "name": "...", "isRequired": "..."}}
# Normalize to use the same keys (including dots) and take the 'value' field only.
normalized = {sk: sv.get("value") for sk, sv in session_vars.items()}
# Let explicit runtime params override session variables on key collisions
merged_params = {**normalized, **merged_params}
# Fill in the template with parameters
msg_content = cls._build_prompt(prompt_template, merged_params)
cls.logger(tool_name="agent-app-prompt").info(f"Invoking prompt tool for {msg_content}")
debug_info["message"] = msg_content
model_provider = llm_conf.get("model_provider")
if not model_provider or not model_provider.strip():
if model_id and not re.match(r"^ocid\d+\.", model_id):
model_provider = model_id.split(".")[0]
else:
raise ValueError(f"Model provider parameter cannot be empty")
message = cls._build_message(model_provider=model_provider, msg_content=msg_content)
model_args = llm_conf.get("model_args") or {}
debug_info["model_args"] = model_args
chat_detail = oci.generative_ai_inference.models.ChatDetails()
chat_request = cls._build_chat_request(model_provider, message, model_args)
if model_id.startswith(CUSTOM_ENDPOINT_PREFIX):
serving_mode = oci.generative_ai_inference.models.OnDemandServingMode(endpoint_id=model_id)
else:
serving_mode = oci.generative_ai_inference.models.OnDemandServingMode(model_id=model_id)
chat_detail.serving_mode = serving_mode
chat_detail.chat_request = chat_request
chat_detail.compartment_id = llm_conf.get("compartment_id")
cls.logger().info(f"Invoking chat request on model {model_id}")
if llm_conf.get("gen_ai_client") is None:
generative_ai_inference_client = cls._get_genai_client(
endpoint=llm_conf.get("endpoint"),
signer=signer
)
else:
# client injection for testing
cls.logger().info("mock client available via conf")
generative_ai_inference_client = llm_conf.get("gen_ai_client")
chat_response = generative_ai_inference_client.chat(chat_detail)
output = cls._parse_chat_response(model_provider, chat_response=chat_response)
cls.logger(tool_name="agent-app-prompt").info(f"Got prompt tool response {output}")
is_success = True
return {
"output": output,
"debug_info": debug_info
}
except Exception as e:
# Extended error logging for debugging: log exception, headers, response, and chat_response if present
cls.logger().error(f"Exception in PromptTool.invoke: {e}", exc_info=True)
is_success = False
if hasattr(e, "response"):
resp = getattr(e, 'response', None)
cls.logger().error(f"Error response: {resp.headers} {resp.text}")
return {
"output": None,
"debug_info": {"error": f"{resp.headers} {resp.text}"}
}
return {
"output": None,
"debug_info": {"error": str(e)}
}
finally:
if is_success:
cls.metricsUtil.increment_counter(cls.success_counter)
else:
cls.metricsUtil.increment_counter(cls.failure_counter)
@classmethod
def _format_output_as_mcp(cls, response: dict) -> dict:
"""
Standardizes the response as an MCP-compatible output.
Includes output, debug_info, and error (if any).
"""
return {
"type": "text",
"text": response.get("output"),
"debug_info": response.get("debug_info")
}
@classmethod
def _get_genai_client(cls, endpoint, signer):
"""
Return a cached GenerativeAiInferenceClient for the given endpoint and signer, or construct one if needed.
Uses a module-level cache shared across all PromptTool uses.
"""
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
@classmethod
def _build_prompt(cls, prompt_template, runtime_params):
"""
Renders the prompt message using the provided template and parameters.
Substitution behavior:
- Supports ONLY Jinja-style placeholders: {{KEY}}
- Does NOT support single-brace {KEY} placeholders, because prompt templates often
contain literal JSON blocks with many '{' and '}', and Python str.format would
misinterpret them and raise KeyError.
Performs sanitization and basic prompt injection checks.
"""
if runtime_params:
sanitized_params = {}
for k, v in runtime_params.items():
# Convert all param values to sanitized strings
val = str(v).strip().replace('\n', ' ').replace('\r', '')
# Simple control char stripping
val = ''.join(c for c in val if 32 <= ord(c) <= 126)
# Check for prompt injection indicators
lower_val = val.lower()
if any(bad in lower_val for bad in PromptTool._injection_keywords):
raise ValueError(f"Prompt injection detected in parameter '{k}': '{v}'")
sanitized_params[k] = val
# only substitute when {{KEY}} placeholders exist.
# Allow optional whitespace inside the placeholder, e.g. {{KEY }}, {{ KEY}}, {{ KEY }}.
runtime_params_set = set(re.findall(r'\{\{\s*([\w.]+)\s*\}\}', prompt_template))
# If no placeholders, return the template literally (avoid .format on JSON)
if not runtime_params_set:
return prompt_template
# We cannot call Python's str.format() on a template containing literal JSON,
# because JSON uses single braces and str.format will treat them as fields.
# Instead, do targeted replacement of only the {{KEY}} placeholders.
rendered = prompt_template
for param in runtime_params_set:
if param in sanitized_params:
# Replace any whitespace-variant of the placeholder for this param.
rendered = re.sub(r'\{\{\s*' + re.escape(param) + r'\s*\}\}', sanitized_params[param], rendered)
return rendered
return prompt_template
@staticmethod
def _set_attr_if_present(obj, key, dict_obj):
if key in dict_obj:
setattr(obj, key, dict_obj[key])
[docs]
@classmethod
def format_query(cls, query_template: str):
"""
Convert all double-brace variables (e.g. {{user_id}}) to single-brace ({user_id}).
Example:
Input: "SELECT * FROM users WHERE id = {{user_id}} AND name = {name}"
Output: "SELECT * FROM users WHERE id = {user_id} AND name = {name}"
"""
return re.sub(r'\{\{(\w+)\}\}', r'{\1}', query_template)
@classmethod
def _build_chat_request(cls, model_provider, message, model_args):
if cls._is_model_provider_not_cohere(model_provider):
cls.logger().debug("Generating a generic chat request")
chat_request = oci.generative_ai_inference.models.GenericChatRequest()
chat_request.messages = [message]
else:
cls.logger().debug("Generating a cohere chat request")
chat_request = oci.generative_ai_inference.models.CohereChatRequest()
chat_request.message = message
for attr in [
"frequency_penalty",
"presence_penalty",
"max_tokens",
"temperature",
"top_p",
"top_k",
]:
cls._set_attr_if_present(chat_request, attr, model_args)
return chat_request
@classmethod
def _build_message(cls, model_provider, msg_content):
if cls._is_model_provider_not_cohere(model_provider):
content = oci.generative_ai_inference.models.TextContent()
content.text = msg_content
message = oci.generative_ai_inference.models.Message()
message.role = "USER"
message.content = [content]
cls.logger().debug("generic message : " + pprint.pformat(message))
return message
else:
cls.logger().debug("cohere message : " + pprint.pformat(msg_content))
return msg_content
@classmethod
def _parse_chat_response(cls, model_provider, chat_response):
if cls._is_model_provider_not_cohere(model_provider):
return (chat_response['data']['chat_response']['choices'][0]['message']['content'][0]['text'])
else:
return (
chat_response.get('data', {})
.get('chat_response', {})
.get('text', None)
)
@classmethod
def _is_model_provider_not_cohere(cls, model_provider):
return model_provider is None or model_provider.lower() != "cohere"