mseep-agentops 0.4.18__py3-none-any.whl → 0.4.23__py3-none-any.whl

agentops/instrumentation/providers/google_genai/stream_wrapper.py

@@ -0,0 +1,238 @@
+"""Google Generative AI stream wrapper implementation.
+
+This module provides wrappers for Google Generative AI's streaming functionality,
+focusing on the generate_content_stream method for both sync and async operations.
+It instruments streams to collect telemetry data for monitoring and analysis.
+"""
+
+import logging
+from typing import TypeVar
+
+from opentelemetry import context as context_api
+from opentelemetry.trace import SpanKind, Status, StatusCode
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+
+from agentops.semconv import SpanAttributes, LLMRequestTypeValues, CoreAttributes, MessageAttributes
+from agentops.instrumentation.common.wrappers import _with_tracer_wrapper
+from agentops.instrumentation.providers.google_genai.attributes.model import (
+    get_generate_content_attributes,
+    get_stream_attributes,
+)
+from agentops.instrumentation.providers.google_genai.attributes.common import (
+    extract_request_attributes,
+)
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+@_with_tracer_wrapper
+def generate_content_stream_wrapper(tracer, wrapped, instance, args, kwargs):
+    """Wrapper for the GenerativeModel.generate_content_stream method.
+
+    This wrapper creates spans for tracking stream performance and processes
+    the streaming responses to collect telemetry data.
+
+    Args:
+        tracer: The OpenTelemetry tracer to use
+        wrapped: The original stream method
+        instance: The instance the method is bound to
+        args: Positional arguments to the method
+        kwargs: Keyword arguments to the method
+
+    Returns:
+        A wrapped generator that captures telemetry data
+    """
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return wrapped(*args, **kwargs)
+
+    span = tracer.start_span(
+        "gemini.generate_content_stream",
+        kind=SpanKind.CLIENT,
+        attributes={SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.CHAT.value},
+    )
+
+    # Extract request parameters and custom config
+    request_attributes = get_generate_content_attributes(args=args, kwargs=kwargs)
+    for key, value in request_attributes.items():
+        span.set_attribute(key, value)
+
+    # Mark as streaming request
+    span.set_attribute(SpanAttributes.LLM_REQUEST_STREAMING, True)
+
+    # Extract custom parameters from config (if present)
+    if "config" in kwargs:
+        config_attributes = extract_request_attributes({"config": kwargs["config"]})
+        for key, value in config_attributes.items():
+            span.set_attribute(key, value)
+
+    try:
+        stream = wrapped(*args, **kwargs)
+
+        # Extract model information if available
+        stream_attributes = get_stream_attributes(stream)
+        for key, value in stream_attributes.items():
+            span.set_attribute(key, value)
+
+        def instrumented_stream():
+            """Generator that wraps the original stream with instrumentation.
+
+            Yields:
+                Items from the original stream with added instrumentation
+            """
+            full_text = ""
+            last_chunk_with_metadata = None
+
+            try:
+                for chunk in stream:
+                    # Keep track of the last chunk that might have metadata
+                    if hasattr(chunk, "usage_metadata") and chunk.usage_metadata:
+                        last_chunk_with_metadata = chunk
+
+                    # Track token count (approximate by word count if metadata not available)
+                    if hasattr(chunk, "text"):
+                        text_value = getattr(chunk, "text", None)
+                        if text_value:
+                            full_text += text_value
+
+                    yield chunk
+
+                # Set final content when complete
+                if full_text:
+                    span.set_attribute(MessageAttributes.COMPLETION_CONTENT.format(i=0), full_text)
+                    span.set_attribute(MessageAttributes.COMPLETION_ROLE.format(i=0), "assistant")
+
+                # Get token usage from the last chunk if available
+                if last_chunk_with_metadata and hasattr(last_chunk_with_metadata, "usage_metadata"):
+                    metadata = last_chunk_with_metadata.usage_metadata
+                    if hasattr(metadata, "prompt_token_count"):
+                        span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, metadata.prompt_token_count)
+                    if hasattr(metadata, "candidates_token_count"):
+                        span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, metadata.candidates_token_count)
+                    if hasattr(metadata, "total_token_count"):
+                        span.set_attribute(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, metadata.total_token_count)
+
+                span.set_status(Status(StatusCode.OK))
+            except Exception as e:
+                span.record_exception(e)
+                span.set_attribute(CoreAttributes.ERROR_MESSAGE, str(e))
+                span.set_attribute(CoreAttributes.ERROR_TYPE, e.__class__.__name__)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                raise
+            finally:
+                span.end()
+
+        return instrumented_stream()
+    except Exception as e:
+        span.record_exception(e)
+        span.set_attribute(CoreAttributes.ERROR_MESSAGE, str(e))
+        span.set_attribute(CoreAttributes.ERROR_TYPE, e.__class__.__name__)
+        span.set_status(Status(StatusCode.ERROR, str(e)))
+        span.end()
+        raise
+
+
+@_with_tracer_wrapper
+async def generate_content_stream_async_wrapper(tracer, wrapped, instance, args, kwargs):
+    """Wrapper for the async GenerativeModel.generate_content_stream method.
+
+    This wrapper creates spans for tracking async stream performance and processes
+    the streaming responses to collect telemetry data.
+
+    Args:
+        tracer: The OpenTelemetry tracer to use
+        wrapped: The original async stream method
+        instance: The instance the method is bound to
+        args: Positional arguments to the method
+        kwargs: Keyword arguments to the method
+
+    Returns:
+        A wrapped async generator that captures telemetry data
+    """
+    if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
+        return await wrapped(*args, **kwargs)
+
+    span = tracer.start_span(
+        "gemini.generate_content_stream_async",
+        kind=SpanKind.CLIENT,
+        attributes={SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.CHAT.value},
+    )
+
+    # Extract request parameters and custom config
+    request_attributes = get_generate_content_attributes(args=args, kwargs=kwargs)
+    for key, value in request_attributes.items():
+        span.set_attribute(key, value)
+
+    # Mark as streaming request
+    span.set_attribute(SpanAttributes.LLM_REQUEST_STREAMING, True)
+
+    # Extract custom parameters from config (if present)
+    if "config" in kwargs:
+        config_attributes = extract_request_attributes({"config": kwargs["config"]})
+        for key, value in config_attributes.items():
+            span.set_attribute(key, value)
+
+    try:
+        stream = await wrapped(*args, **kwargs)
+
+        # Extract model information if available
+        stream_attributes = get_stream_attributes(stream)
+        for key, value in stream_attributes.items():
+            span.set_attribute(key, value)
+
+        async def instrumented_stream():
+            """Async generator that wraps the original stream with instrumentation.
+
+            Yields:
+                Items from the original stream with added instrumentation
+            """
+            full_text = ""
+            last_chunk_with_metadata = None
+
+            try:
+                async for chunk in stream:
+                    # Keep track of the last chunk that might have metadata
+                    if hasattr(chunk, "usage_metadata") and chunk.usage_metadata:
+                        last_chunk_with_metadata = chunk
+
+                    if hasattr(chunk, "text"):
+                        text_value = getattr(chunk, "text", None)
+                        if text_value:
+                            full_text += text_value
+
+                    yield chunk
+
+                # Set final content when complete
+                if full_text:
+                    span.set_attribute(MessageAttributes.COMPLETION_CONTENT.format(i=0), full_text)
+                    span.set_attribute(MessageAttributes.COMPLETION_ROLE.format(i=0), "assistant")
+
+                # Get token usage from the last chunk if available
+                if last_chunk_with_metadata and hasattr(last_chunk_with_metadata, "usage_metadata"):
+                    metadata = last_chunk_with_metadata.usage_metadata
+                    if hasattr(metadata, "prompt_token_count"):
+                        span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, metadata.prompt_token_count)
+                    if hasattr(metadata, "candidates_token_count"):
+                        span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, metadata.candidates_token_count)
+                    if hasattr(metadata, "total_token_count"):
+                        span.set_attribute(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, metadata.total_token_count)
+
+                span.set_status(Status(StatusCode.OK))
+            except Exception as e:
+                span.record_exception(e)
+                span.set_attribute(CoreAttributes.ERROR_MESSAGE, str(e))
+                span.set_attribute(CoreAttributes.ERROR_TYPE, e.__class__.__name__)
+                span.set_status(Status(StatusCode.ERROR, str(e)))
+                raise
+            finally:
+                span.end()
+
+        return instrumented_stream()
+    except Exception as e:
+        span.record_exception(e)
+        span.set_attribute(CoreAttributes.ERROR_MESSAGE, str(e))
+        span.set_attribute(CoreAttributes.ERROR_TYPE, e.__class__.__name__)
+        span.set_status(Status(StatusCode.ERROR, str(e)))
+        span.end()
+        raise

agentops/instrumentation/providers/ibm_watsonx_ai/__init__.py

@@ -0,0 +1,28 @@
+"""IBM WatsonX AI instrumentation for AgentOps.
+
+This package provides instrumentation for IBM's WatsonX AI foundation models,
+capturing telemetry for model interactions including completions, chat, and streaming responses.
+"""
+
+import logging
+from agentops.instrumentation.common import LibraryInfo
+
+logger = logging.getLogger(__name__)
+
+# Library information
+_library_info = LibraryInfo(
+    name="ibm_watsonx_ai",
+    package_name="ibm-watsonx-ai",
+    default_version="1.3.11",  # Default to known supported version if not found
+)
+LIBRARY_NAME = _library_info.name
+LIBRARY_VERSION = _library_info.version
+
+# Import after defining constants to avoid circular imports
+from agentops.instrumentation.providers.ibm_watsonx_ai.instrumentor import WatsonxInstrumentor  # noqa: E402
+
+__all__ = [
+    "LIBRARY_NAME",
+    "LIBRARY_VERSION",
+    "WatsonxInstrumentor",
+]

agentops/instrumentation/providers/ibm_watsonx_ai/attributes/__init__.py

@@ -0,0 +1,27 @@
+"""Attribute extraction utilities for IBM watsonx.ai instrumentation."""
+
+from agentops.instrumentation.providers.ibm_watsonx_ai.attributes.attributes import (
+    get_generate_attributes,
+    get_chat_attributes,
+    get_tokenize_attributes,
+    get_model_details_attributes,
+)
+from agentops.instrumentation.providers.ibm_watsonx_ai.attributes.common import (
+    extract_params_attributes,
+    convert_params_to_dict,
+    extract_prompt_from_args,
+    extract_messages_from_args,
+    extract_params_from_args,
+)
+
+__all__ = [
+    "get_generate_attributes",
+    "get_chat_attributes",
+    "get_tokenize_attributes",
+    "get_model_details_attributes",
+    "extract_params_attributes",
+    "convert_params_to_dict",
+    "extract_prompt_from_args",
+    "extract_messages_from_args",
+    "extract_params_from_args",
+]

agentops/instrumentation/providers/ibm_watsonx_ai/attributes/attributes.py

@@ -0,0 +1,277 @@
+"""Attributes for IBM watsonx.ai model instrumentation.
+
+This module provides attribute extraction functions for IBM watsonx.ai model operations.
+"""
+
+from typing import Any, Dict, Optional, Tuple
+from agentops.instrumentation.common.attributes import AttributeMap
+from agentops.semconv import SpanAttributes, MessageAttributes
+from agentops.instrumentation.providers.ibm_watsonx_ai.attributes.common import (
+    extract_params_attributes,
+    convert_params_to_dict,
+    extract_prompt_from_args,
+    extract_messages_from_args,
+    extract_params_from_args,
+)
+
+
+def get_generate_attributes(
+    args: Optional[Tuple] = None, kwargs: Optional[Dict] = None, return_value: Optional[Any] = None
+) -> AttributeMap:
+    """Extract token usage attributes from generate method calls."""
+    attributes = {}
+
+    # Extract prompt using helper function
+    prompt = extract_prompt_from_args(args, kwargs)
+    if prompt:
+        attributes[MessageAttributes.PROMPT_ROLE.format(i=0)] = "user"
+        attributes[MessageAttributes.PROMPT_CONTENT.format(i=0)] = prompt
+        attributes[MessageAttributes.PROMPT_TYPE.format(i=0)] = "text"
+
+    # Extract parameters using helper functions
+    params = extract_params_from_args(args, kwargs)
+    if params:
+        params_dict = convert_params_to_dict(params)
+        if params_dict:
+            attributes.update(extract_params_attributes(params_dict))
+
+    # Extract response information
+    if return_value:
+        if isinstance(return_value, dict):
+            # Extract model information
+            if "model_id" in return_value:
+                attributes[SpanAttributes.LLM_REQUEST_MODEL] = return_value["model_id"]
+
+            # Handle results
+            if "results" in return_value:
+                for idx, result in enumerate(return_value["results"]):
+                    # Extract completion
+                    if "generated_text" in result:
+                        attributes[MessageAttributes.COMPLETION_CONTENT.format(i=idx)] = result["generated_text"]
+                        attributes[MessageAttributes.COMPLETION_ROLE.format(i=idx)] = "assistant"
+                        attributes[MessageAttributes.COMPLETION_TYPE.format(i=idx)] = "text"
+
+                    # Extract token usage
+                    if "input_token_count" in result:
+                        attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] = result["input_token_count"]
+                    if "generated_token_count" in result:
+                        attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] = result["generated_token_count"]
+                    if "input_token_count" in result and "generated_token_count" in result:
+                        attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = (
+                            result["input_token_count"] + result["generated_token_count"]
+                        )
+
+                    if "stop_reason" in result:
+                        attributes[SpanAttributes.LLM_RESPONSE_STOP_REASON] = result["stop_reason"]
+
+    return attributes
+
+
+def get_tokenize_attributes(
+    args: Optional[Tuple] = None, kwargs: Optional[Dict] = None, return_value: Optional[Any] = None
+) -> AttributeMap:
+    """Extract attributes from tokenize method calls."""
+    attributes = {}
+
+    # Extract input from args or kwargs using helper function
+    prompt = extract_prompt_from_args(args, kwargs)
+    if prompt:
+        attributes[MessageAttributes.PROMPT_ROLE.format(i=0)] = "user"
+        attributes[MessageAttributes.PROMPT_CONTENT.format(i=0)] = prompt
+        attributes[MessageAttributes.PROMPT_TYPE.format(i=0)] = "text"
+
+    # Extract response information
+    if return_value and isinstance(return_value, dict):
+        if "model_id" in return_value:
+            attributes[SpanAttributes.LLM_REQUEST_MODEL] = return_value["model_id"]
+        if "result" in return_value:
+            attributes["ibm.watsonx.tokenize.result"] = str(return_value["result"])
+            if "token_count" in return_value["result"]:
+                attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] = return_value["result"]["token_count"]
+
+    return attributes
+
+
+def get_model_details_attributes(
+    args: Optional[Tuple] = None, kwargs: Optional[Dict] = None, return_value: Optional[Any] = None
+) -> AttributeMap:
+    """Extract attributes from get_details method calls."""
+    if not isinstance(return_value, dict):
+        return {}
+
+    # Basic model information
+    attributes = {
+        f"ibm.watsonx.model.{key}": value
+        for key, value in return_value.items()
+        if key
+        in [
+            "model_id",
+            "label",
+            "provider",
+            "source",
+            "short_description",
+            "long_description",
+            "number_params",
+            "input_tier",
+            "output_tier",
+        ]
+    }
+
+    # Model functions
+    if "functions" in return_value:
+        attributes["ibm.watsonx.model.functions"] = str([func["id"] for func in return_value["functions"]])
+
+    # Model tasks
+    if "tasks" in return_value:
+        task_info = [
+            {k: v for k, v in task.items() if k in ["id", "ratings", "tags"]} for task in return_value["tasks"]
+        ]
+        attributes["ibm.watsonx.model.tasks"] = str(task_info)
+
+    # Model limits
+    if "model_limits" in return_value:
+        limits = return_value["model_limits"]
+        attributes.update(
+            {
+                f"ibm.watsonx.model.{key}": value
+                for key, value in limits.items()
+                if key in ["max_sequence_length", "max_output_tokens", "training_data_max_records"]
+            }
+        )
+
+    # Service tier limits
+    if "limits" in return_value:
+        for tier, tier_limits in return_value["limits"].items():
+            attributes.update(
+                {
+                    f"ibm.watsonx.model.limits.{tier}.{key}": value
+                    for key, value in tier_limits.items()
+                    if key in ["call_time", "max_output_tokens"]
+                }
+            )
+
+    # Model lifecycle
+    if "lifecycle" in return_value:
+        attributes.update(
+            {
+                f"ibm.watsonx.model.lifecycle.{stage['id']}": stage["start_date"]
+                for stage in return_value["lifecycle"]
+                if "id" in stage and "start_date" in stage
+            }
+        )
+
+    # Training parameters
+    if "training_parameters" in return_value:
+        attributes.update(
+            {
+                f"ibm.watsonx.model.training.{key}": str(value) if isinstance(value, dict) else value
+                for key, value in return_value["training_parameters"].items()
+            }
+        )
+
+    return attributes
+
+
+def get_chat_attributes(
+    args: Optional[Tuple] = None, kwargs: Optional[Dict] = None, return_value: Optional[Any] = None
+) -> AttributeMap:
+    """Extract attributes from chat method calls."""
+    attributes = {}
+
+    # Extract messages using helper function
+    messages = extract_messages_from_args(args, kwargs)
+    if messages:
+        # Process each message in the conversation
+        for i, message in enumerate(messages):
+            if not isinstance(message, dict):
+                continue
+
+            # Extract role and content
+            role = message.get("role", "")
+            content = message.get("content", [])
+
+            # Handle content which can be a list of different types (text, image_url)
+            if isinstance(content, list):
+                # Combine all text content
+                text_content = []
+                image_urls = []
+
+                for content_item in content:
+                    if isinstance(content_item, dict):
+                        if content_item.get("type") == "text":
+                            text_content.append(content_item.get("text", ""))
+                        elif content_item.get("type") == "image_url":
+                            image_url = content_item.get("image_url", {})
+                            if isinstance(image_url, dict) and "url" in image_url:
+                                url = image_url["url"]
+                                # Only store URLs that start with http, otherwise use placeholder
+                                if url and isinstance(url, str) and url.startswith(("http://", "https://")):
+                                    image_urls.append(url)
+                                else:
+                                    image_urls.append("[IMAGE_PLACEHOLDER]")
+
+                # Set text content if any
+                if text_content:
+                    attributes[MessageAttributes.PROMPT_CONTENT.format(i=i)] = " ".join(text_content)
+                    attributes[MessageAttributes.PROMPT_TYPE.format(i=i)] = "text"
+                    attributes[MessageAttributes.PROMPT_ROLE.format(i=i)] = role
+
+                # Set image URLs if any
+                if image_urls:
+                    attributes[f"ibm.watsonx.chat.message.{i}.images"] = str(image_urls)
+            else:
+                # Handle string content
+                attributes[MessageAttributes.PROMPT_CONTENT.format(i=i)] = str(content)
+                attributes[MessageAttributes.PROMPT_TYPE.format(i=i)] = "text"
+                attributes[MessageAttributes.PROMPT_ROLE.format(i=i)] = role
+
+    # Extract parameters using helper functions
+    params = extract_params_from_args(args, kwargs)
+    if params:
+        params_dict = convert_params_to_dict(params)
+        if params_dict:
+            attributes.update(extract_params_attributes(params_dict))
+
+    # Extract response information
+    if return_value and isinstance(return_value, dict):
+        # Extract model information
+        if "model_id" in return_value:
+            attributes[SpanAttributes.LLM_REQUEST_MODEL] = return_value["model_id"]
+        elif "model" in return_value:
+            attributes[SpanAttributes.LLM_REQUEST_MODEL] = return_value["model"]
+
+        # Extract completion from choices
+        if "choices" in return_value:
+            for idx, choice in enumerate(return_value["choices"]):
+                if isinstance(choice, dict) and "message" in choice:
+                    message = choice["message"]
+                    if isinstance(message, dict):
+                        if "content" in message:
+                            attributes[MessageAttributes.COMPLETION_CONTENT.format(i=idx)] = message["content"]
+                            attributes[MessageAttributes.COMPLETION_ROLE.format(i=idx)] = message.get(
+                                "role", "assistant"
+                            )
+                            attributes[MessageAttributes.COMPLETION_TYPE.format(i=idx)] = "text"
+                        if "finish_reason" in choice:
+                            attributes[SpanAttributes.LLM_RESPONSE_STOP_REASON] = choice["finish_reason"]
+
+        # Extract token usage
+        if "usage" in return_value:
+            usage = return_value["usage"]
+            if isinstance(usage, dict):
+                if "prompt_tokens" in usage:
+                    attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] = usage["prompt_tokens"]
+                if "completion_tokens" in usage:
+                    attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] = usage["completion_tokens"]
+                if "total_tokens" in usage:
+                    attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = usage["total_tokens"]
+
+        # Extract additional metadata
+        if "id" in return_value:
+            attributes["ibm.watsonx.chat.id"] = return_value["id"]
+        if "model_version" in return_value:
+            attributes["ibm.watsonx.model.version"] = return_value["model_version"]
+        if "created_at" in return_value:
+            attributes["ibm.watsonx.chat.created_at"] = return_value["created_at"]
+
+    return attributes

agentops/instrumentation/providers/ibm_watsonx_ai/attributes/common.py

@@ -0,0 +1,104 @@
+"""Common utilities and constants for IBM watsonx.ai attribute processing.
+
+This module contains shared constants, attribute mappings, and utility functions for processing
+trace and span attributes in IBM watsonx.ai instrumentation.
+"""
+
+from typing import Any, Dict, Optional, Tuple, List
+from agentops.instrumentation.common.attributes import AttributeMap
+from agentops.semconv import SpanAttributes
+from agentops.logging import logger
+from ibm_watsonx_ai.foundation_models.schema import TextGenParameters, TextChatParameters
+
+# Mapping of generation parameters to their OpenTelemetry attribute names
+GENERATION_PARAM_ATTRIBUTES: AttributeMap = {
+    "max_new_tokens": SpanAttributes.LLM_REQUEST_MAX_TOKENS,
+    "min_new_tokens": "ibm.watsonx.min_new_tokens",
+    "temperature": SpanAttributes.LLM_REQUEST_TEMPERATURE,
+    "top_p": SpanAttributes.LLM_REQUEST_TOP_P,
+    "top_k": "ibm.watsonx.top_k",
+    "repetition_penalty": "ibm.watsonx.repetition_penalty",
+    "time_limit": "ibm.watsonx.time_limit",
+    "random_seed": "ibm.watsonx.random_seed",
+    "stop_sequences": "ibm.watsonx.stop_sequences",
+    "truncate_input_tokens": "ibm.watsonx.truncate_input_tokens",
+    "decoding_method": "ibm.watsonx.decoding_method",
+}
+
+# Mapping of guardrail parameters to their OpenTelemetry attribute names
+GUARDRAIL_PARAM_ATTRIBUTES: AttributeMap = {
+    "guardrails": "ibm.watsonx.guardrails.enabled",
+    "guardrails_hap_params": "ibm.watsonx.guardrails.hap_params",
+    "guardrails_pii_params": "ibm.watsonx.guardrails.pii_params",
+}
+
+
+def extract_prompt_from_args(args: Optional[Tuple] = None, kwargs: Optional[Dict] = None) -> Optional[str]:
+    """Extract prompt from method arguments."""
+    if args and len(args) > 0:
+        return args[0]
+    elif kwargs and "prompt" in kwargs:
+        return kwargs["prompt"]
+    return None
+
+
+def extract_messages_from_args(
+    args: Optional[Tuple] = None, kwargs: Optional[Dict] = None
+) -> Optional[List[Dict[str, Any]]]:
+    """Extract messages from method arguments."""
+    if args and len(args) > 0:
+        return args[0]
+    elif kwargs and "messages" in kwargs:
+        return kwargs["messages"]
+    return None
+
+
+def extract_params_from_args(args: Optional[Tuple] = None, kwargs: Optional[Dict] = None) -> Optional[Any]:
+    """Extract parameters from method arguments."""
+    if args and len(args) > 1:
+        return args[1]
+    elif kwargs and "params" in kwargs:
+        return kwargs["params"]
+    return None
+
+
+def convert_params_to_dict(params: Any) -> Dict[str, Any]:
+    """Convert parameter objects to dictionaries."""
+    if not params:
+        return {}
+
+    if isinstance(params, (TextGenParameters, TextChatParameters)):
+        try:
+            return params.to_dict()
+        except Exception as e:
+            logger.debug(f"Could not convert params object to dict: {e}")
+            return {}
+
+    return params if isinstance(params, dict) else {}
+
+
+def extract_params_attributes(params: Dict[str, Any]) -> AttributeMap:
+    """Extract generation parameters from a params dictionary."""
+    attributes = {}
+
+    # Extract standard generation parameters
+    for param_name, attr_name in GENERATION_PARAM_ATTRIBUTES.items():
+        if param_name in params:
+            value = params[param_name]
+            if isinstance(value, list):
+                value = str(value)
+            attributes[attr_name] = value
+
+    # Extract guardrail parameters
+    for param_name, attr_name in GUARDRAIL_PARAM_ATTRIBUTES.items():
+        if param_name in params:
+            value = params[param_name]
+            if isinstance(value, dict):
+                value = str(value)
+            attributes[attr_name] = value
+
+    # Extract concurrency limit
+    if "concurrency_limit" in params:
+        attributes["ibm.watsonx.concurrency_limit"] = params["concurrency_limit"]
+
+    return attributes