monocle-apptrace 0.5.0b1__py3-none-any.whl → 0.5.1b1__py3-none-any.whl

monocle_apptrace/exporters/file_exporter.py

@@ -3,6 +3,7 @@
 from os import linesep, path
 from io import TextIOWrapper
 from datetime import datetime
+import os
 from typing import Optional, Callable, Sequence, Dict, Tuple
 from opentelemetry.sdk.trace import ReadableSpan
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
@@ -32,7 +33,7 @@ class FileSpanExporter(SpanExporterBase):
         self.file_handles: Dict[int, Tuple[TextIOWrapper, str, datetime, bool]] = {}
         self.formatter = formatter
         self.service_name = service_name
-        self.output_path = out_path
+        self.output_path = os.getenv("MONOCLE_TRACE_OUTPUT_PATH", out_path)
         self.file_prefix = file_prefix
         self.time_format = time_format
         self.task_processor = task_processor

monocle_apptrace/instrumentation/common/__init__.py

@@ -2,16 +2,18 @@ from .instrumentor import (
     setup_monocle_telemetry, 
     start_trace, 
     stop_trace, 
-    start_scope, 
-    stop_scope, 
     http_route_handler, 
-    monocle_trace_scope, 
-    amonocle_trace_scope,
-    monocle_trace_scope_method, 
     monocle_trace, 
     amonocle_trace,
     monocle_trace_method,
     monocle_trace_http_route,
     is_valid_trace_id_uuid
 )
+from .scope_wrapper import (
+    start_scope, 
+    stop_scope, 
+    monocle_trace_scope, 
+    amonocle_trace_scope,
+    monocle_trace_scope_method
+)
 from .utils import MonocleSpanException

monocle_apptrace/instrumentation/common/constants.py

@@ -26,7 +26,7 @@ service_type_map = {
     AZURE_APP_SERVICE_ENV_NAME: AZURE_APP_SERVICE_NAME,
     AZURE_FUNCTION_WORKER_ENV_NAME: AZURE_FUNCTION_NAME,
     AWS_LAMBDA_ENV_NAME: AWS_LAMBDA_SERVICE_NAME,
-    GITHUB_CODESPACE_ENV_NAME: GITHUB_CODESPACE_SERVICE_NAME
+    GITHUB_CODESPACE_ENV_NAME: GITHUB_CODESPACE_SERVICE_NAME,
 }
 
 # Env variables to identify infra service name
@@ -35,7 +35,7 @@ service_name_map = {
     AZURE_FUNCTION_NAME: AZURE_FUNCTION_IDENTIFIER_ENV_NAME,
     AZURE_ML_SERVICE_NAME: AZURE_ML_ENDPOINT_ENV_NAME,
     AWS_LAMBDA_SERVICE_NAME: AWS_LAMBDA_FUNCTION_IDENTIFIER_ENV_NAME,
-    GITHUB_CODESPACE_SERVICE_NAME: GITHUB_CODESPACE_IDENTIFIER_ENV_NAME
+    GITHUB_CODESPACE_SERVICE_NAME: GITHUB_CODESPACE_IDENTIFIER_ENV_NAME,
 }
 
 
@@ -49,11 +49,11 @@ llm_type_map = {
     "sagemakerllm": "aws_sagemaker",
     "chatbedrock": "aws_bedrock",
     "openaigenerator": "openai",
-    "bedrockruntime":"aws_bedrock",
-    "sagemakerruntime":"aws_sagemaker",
+    "bedrockruntime": "aws_bedrock",
+    "sagemakerruntime": "aws_sagemaker",
     "anthropic": "anthropic",
-    "chatanthropic":"anthropic",
-    "anthropicchatgenerator":"anthropic",
+    "chatanthropic": "anthropic",
+    "anthropicchatgenerator": "anthropic",
     "chatcompletionsclient": "azure_ai_inference",
     "embeddingsclient": "azure_ai_inference",
     "imageembeddingsclient": "azure_ai_inference",
@@ -61,6 +61,8 @@ llm_type_map = {
     "googleaigeminichatgenerator": "gemini",
     "gemini": "gemini",
     "chatgooglegenerativeai": "gemini",
+    "azurechatcompletion": "azure_openai",
+    "openaichatcompletion": "openai",
 }
 
 MONOCLE_INSTRUMENTOR = "monocle_apptrace"
@@ -73,11 +75,11 @@ QUERY = "input"
 RESPONSE = "response"
 SESSION_PROPERTIES_KEY = "session"
 INFRA_SERVICE_KEY = "infra_service_name"
-META_DATA = 'metadata'
+META_DATA = "metadata"
 MONOCLE_SCOPE_NAME_PREFIX = "monocle.scope."
-SCOPE_METHOD_LIST = 'MONOCLE_SCOPE_METHODS'
-SCOPE_METHOD_FILE = 'monocle_scopes.json'
-SCOPE_CONFIG_PATH = 'MONOCLE_SCOPE_CONFIG_PATH'
+SCOPE_METHOD_LIST = "MONOCLE_SCOPE_METHODS"
+SCOPE_METHOD_FILE = "monocle_scopes.json"
+SCOPE_CONFIG_PATH = "MONOCLE_SCOPE_CONFIG_PATH"
 TRACE_PROPOGATION_URLS = "MONOCLE_TRACE_PROPAGATATION_URLS"
 WORKFLOW_TYPE_KEY = "monocle.workflow_type"
 ADD_NEW_WORKFLOW = "monocle.add_new_workflow"
@@ -85,5 +87,94 @@ WORKFLOW_TYPE_GENERIC = "workflow.generic"
 MONOCLE_SDK_VERSION = "monocle_apptrace.version"
 MONOCLE_SDK_LANGUAGE = "monocle_apptrace.language"
 MONOCLE_DETECTED_SPAN_ERROR = "monocle_apptrace.detected_span_error"
-HTTP_SUCCESS_CODES = ('200', '201', '202', '204', '205', '206')
-CHILD_ERROR_CODE = "child.error.code"
+HTTP_SUCCESS_CODES = ("200", "201", "202", "204", "205", "206")
+CHILD_ERROR_CODE = "child.error.code"
+
+AGENT_PREFIX_KEY = "monocle.agent.prefix"
+
+# agentic sub types
+INFERENCE_AGENT_DELEGATION = "delegation"
+INFERENCE_TOOL_CALL = "tool_call"
+INFERENCE_TURN_END = "turn_end"
+
+SCOPE_NAME = "scope_name"
+AGENT_INVOCATION_SPAN_NAME = "agentic.invocation"
+AGENT_REQUEST_SPAN_NAME = "agentic.request"
+
+AGENTIC_SPANS = [AGENT_INVOCATION_SPAN_NAME, AGENT_REQUEST_SPAN_NAME]
+
+# Span sub types
+
+## OPTIONAL right next to span.type,  span.subtype:
+## subtype is one perspective , are non overlapping, limitations: only one classification scheme for subtypes
+# 1 planning
+SPAN_SUBTYPE_PLANNING = "planning"
+
+# 2 routing and selection INFERENCE_TOOL_CALL, INFERENCE_AGENT_DELEGATION
+SPAN_SUBTYPE_ROUTING = "routing"
+
+# 3 content processing
+SPAN_SUBTYPE_CONTENT_PROCESSING = "content_processing"
+
+# 4 content generation
+SPAN_SUBTYPE_CONTENT_GENERATION = "content_generation"
+
+# 5 communication INFERENCE_TURN_END
+SPAN_SUBTYPE_COMMUNICATION = "communication"
+
+# 6 transformations , if structured output
+SPAN_SUBTYPE_TRANSFORMATIONS = "transformations"
+
+# 7 domain specific,
+SPAN_SUBTYPE_DOMAIN_SPECIFIC = "domain_specific"
+
+# 8 generic (we may skip this property)
+SPAN_SUBTYPE_GENERIC = "generic"
+
+class SPAN_TYPES:
+    GENERIC = "generic"
+    AGENTIC_DELEGATION = "agentic.delegation"
+    AGENTIC_TOOL_INVOCATION = "agentic.tool.invocation"
+    AGENTIC_INVOCATION = "agentic.invocation"
+    AGENTIC_MCP_INVOCATION = "agentic.mcp.invocation"
+    AGENTIC_REQUEST = "agentic.request"
+
+    # http.process
+    HTTP_PROCESS = "http.process"
+    HTTP_SEND = "http.send"
+
+    RETRIEVAL = "retrieval"
+    INFERENCE = "inference"
+    INFERENCE_FRAMEWORK = "inference.framework"
+
+
+class SPAN_SUBTYPES:
+    PLANNING = SPAN_SUBTYPE_PLANNING
+    ROUTING = SPAN_SUBTYPE_ROUTING
+    CONTENT_PROCESSING = SPAN_SUBTYPE_CONTENT_PROCESSING
+    CONTENT_GENERATION = SPAN_SUBTYPE_CONTENT_GENERATION
+    COMMUNICATION = SPAN_SUBTYPE_COMMUNICATION
+    TRANSFORMATIONS = SPAN_SUBTYPE_TRANSFORMATIONS
+    DOMAIN_SPECIFIC = SPAN_SUBTYPE_DOMAIN_SPECIFIC
+    GENERIC = SPAN_SUBTYPE_GENERIC
+
+
+MAP_ATTRIBUTES_TO_SPAN_SUBTYPE = {
+    # inference attributes
+    INFERENCE_AGENT_DELEGATION: SPAN_SUBTYPES.ROUTING,
+    INFERENCE_TOOL_CALL: SPAN_SUBTYPES.ROUTING,
+    INFERENCE_TURN_END: SPAN_SUBTYPES.COMMUNICATION,
+
+    # agentic span.types
+    SPAN_TYPES.AGENTIC_DELEGATION: SPAN_SUBTYPES.ROUTING,
+    SPAN_TYPES.AGENTIC_TOOL_INVOCATION: SPAN_SUBTYPES.ROUTING,
+    SPAN_TYPES.AGENTIC_INVOCATION: SPAN_SUBTYPES.ROUTING,
+    SPAN_TYPES.AGENTIC_MCP_INVOCATION: SPAN_SUBTYPES.ROUTING,
+    
+    # MAYBE?
+    # agentic span.types
+    SPAN_TYPES.AGENTIC_REQUEST: SPAN_SUBTYPES.PLANNING,
+    
+    
+    
+}

monocle_apptrace/instrumentation/common/instrumentor.py

@@ -235,17 +235,12 @@ def is_valid_trace_id_uuid(traceId: str) -> bool:
     return False
 
 from monocle_apptrace.instrumentation.common.method_wrappers import (
-    monocle_trace_scope_method,
-    amonocle_trace_scope,
     monocle_trace,
     amonocle_trace,
     monocle_trace_method,
     monocle_trace_http_route,
-    start_scope,
-    stop_scope,
     start_trace,
     stop_trace,
-    http_route_handler,
-    monocle_trace_scope,
+    http_route_handler
 )
 

monocle_apptrace/instrumentation/common/method_wrappers.py

@@ -5,14 +5,13 @@ from functools import wraps
 import inspect
 from opentelemetry.context import attach, get_current, detach
 from opentelemetry.sdk.trace import Span
-from opentelemetry.sdk.trace import Span
+from opentelemetry.trace.span import INVALID_SPAN
 from opentelemetry.trace import get_tracer
-from opentelemetry.trace.propagation import set_span_in_context, _SPAN_KEY
 from contextlib import contextmanager, asynccontextmanager
 from monocle_apptrace.instrumentation.common.span_handler import SpanHandler
-from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper, task_wrapper
+from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper, get_current_monocle_span, set_monocle_span_in_context, task_wrapper
 from monocle_apptrace.instrumentation.common.utils import (
-    set_scope, remove_scope, http_route_handler, http_async_route_handler
+    http_route_handler, http_async_route_handler
 )
 from monocle_apptrace.instrumentation.common.constants import MONOCLE_INSTRUMENTOR
 from monocle_apptrace.instrumentation.common.instrumentor import get_tracer_provider
@@ -48,7 +47,7 @@ def start_trace(
         tracer = get_tracer(instrumenting_module_name= MONOCLE_INSTRUMENTOR, tracer_provider= get_tracer_provider())
         span_name = span_name or "custom_span"
         span = tracer.start_span(name=span_name)
-        updated_span_context = set_span_in_context(span=span)
+        updated_span_context = set_monocle_span_in_context(span=span)
         
         # Set default monocle attributes
         SpanHandler.set_default_monocle_attributes(span)
@@ -57,7 +56,7 @@ def start_trace(
         
         # Set custom attributes and events using common method
         _setup_span_attributes_and_events(span, attributes, events)
-        
+
         token = attach(updated_span_context)
         return token
     except Exception as e:
@@ -81,62 +80,16 @@ def stop_trace(
         None
     """
     try:
-        _parent_span_context = get_current()
-        if _parent_span_context is not None:
-            parent_span: Span = _parent_span_context.get(_SPAN_KEY, None)
-            if parent_span is not None:
-                # Set final attributes and events using common method
-                _setup_span_attributes_and_events(parent_span, final_attributes, final_events)
-                
-                parent_span.end()
+        parent_span = get_current_monocle_span()
+        if parent_span is not None and parent_span != INVALID_SPAN:
+            # Set final attributes and events using common method
+            _setup_span_attributes_and_events(parent_span, final_attributes, final_events)
+            parent_span.end()
         if token is not None:
             detach(token)
     except Exception as e:
         logger.warning(f"Failed to stop trace: {e}")
 
-def start_scope(
-    scope_name: str, 
-    scope_value: Optional[str] = None
-) -> object:
-    """
-    Start a new scope with the given name and optional value. If no value is provided, a random UUID will be generated.
-    All the spans, across traces created after this call will have the scope attached until the scope is stopped.
-    
-    Args:
-        scope_name: The name of the scope.
-        scope_value: Optional value of the scope. If None, a random UUID will be generated.
-    
-    Returns:
-        Token: A token representing the attached context for the scope. This token is to be used later to stop the current scope.
-    """
-    try:
-        # Set the scope using existing utility
-        token = set_scope(scope_name, scope_value)
-        return token
-    except Exception as e:
-        logger.warning(f"Failed to start scope: {e}")
-        return None
-
-def stop_scope(
-    token: object
-) -> None:
-    """
-    Stop the active scope. All the spans created after this will not have the scope attached.
-    
-    Args:
-        token: The token that was returned when the scope was started.
-    
-    Returns:
-        None
-    """
-    try:
-        # Remove the scope
-        remove_scope(token)
-    except Exception as e:
-        logger.warning(f"Failed to stop scope: {e}")
-    return
-
-
 @contextmanager
 def monocle_trace(
     span_name: Optional[str] = None,
@@ -211,74 +164,6 @@ async def amonocle_trace(
         logger.warning(f"Failed in amonocle_trace: {e}")
         yield  # Still yield to not break the context manager
 
-@contextmanager
-def monocle_trace_scope(
-    scope_name: str, 
-    scope_value: Optional[str] = None
-):
-    """
-    Context manager to start and stop a scope. All the spans, across traces created within the encapsulated code will have the scope attached.
-    
-    Args:
-        scope_name: The name of the scope.
-        scope_value: Optional value of the scope. If None, a random UUID will be generated.
-    """
-    token = start_scope(scope_name, scope_value)
-    try:
-        yield
-    finally:
-        stop_scope(token)
-
-@asynccontextmanager
-async def amonocle_trace_scope(
-    scope_name: str, 
-    scope_value: Optional[str] = None
-):
-    """
-    Async context manager to start and stop a scope. All the spans, across traces created within the encapsulated code will have the scope attached.
-    
-    Args:
-        scope_name: The name of the scope.
-        scope_value: Optional value of the scope. If None, a random UUID will be generated.
-    """
-    token = start_scope(scope_name, scope_value)
-    try:
-        yield
-    finally:
-        stop_scope(token)
-    
-def monocle_trace_scope_method(
-    scope_name: str, 
-    scope_value: Optional[str] = None
-):
-    """
-    Decorator to start and stop a scope for a method. All the spans, across traces created in the method will have the scope attached.
-    
-    Args:
-        scope_name: The name of the scope.
-        scope_value: Optional value of the scope. If None, a random UUID will be generated.
-    """
-    def decorator(func):
-        if inspect.iscoroutinefunction(func):
-            @wraps(func)
-            async def wrapper(*args, **kwargs):
-                async with amonocle_trace_scope(
-                    scope_name, scope_value
-                ):
-                    result = await func(*args, **kwargs)
-                    return result
-            return wrapper
-        else:
-            @wraps(func)
-            def wrapper(*args, **kwargs):
-                with monocle_trace_scope(
-                    scope_name, scope_value
-                ):
-                    result = func(*args, **kwargs)
-                    return result
-            return wrapper
-    return decorator
-
 def monocle_trace_method(
     span_name: Optional[str] = None
 ):

monocle_apptrace/instrumentation/common/scope_wrapper.py

@@ -0,0 +1,126 @@
+import logging
+import inspect
+from typing import Dict, List, Optional, Any
+from functools import wraps
+from contextlib import contextmanager, asynccontextmanager
+from opentelemetry.context import Context
+from monocle_apptrace.instrumentation.common.utils import (
+    set_scope, remove_scope, http_route_handler, http_async_route_handler
+)
+
+logger = logging.getLogger(__name__)
+
+def start_scope(
+    scope_name: str, 
+    scope_value: Optional[str] = None,
+    context:  Optional[Context] = None
+) -> object:
+    """
+    Start a new scope with the given name and optional value. If no value is provided, a random UUID will be generated.
+    All the spans, across traces created after this call will have the scope attached until the scope is stopped.
+    
+    Args:
+        scope_name: The name of the scope.
+        scope_value: Optional value of the scope. If None, a random UUID will be generated.
+    
+    Returns:
+        Token: A token representing the attached context for the scope. This token is to be used later to stop the current scope.
+    """
+    try:
+        # Set the scope using existing utility
+        token = set_scope(scope_name, scope_value, context)
+        return token
+    except Exception as e:
+        logger.warning(f"Failed to start scope: {e}")
+        return None
+
+def stop_scope(
+    token: object
+) -> None:
+    """
+    Stop the active scope. All the spans created after this will not have the scope attached.
+    
+    Args:
+        token: The token that was returned when the scope was started.
+    
+    Returns:
+        None
+    """
+    try:
+        # Remove the scope
+        remove_scope(token)
+    except Exception as e:
+        logger.warning(f"Failed to stop scope: {e}")
+    return
+
+
+
+@contextmanager
+def monocle_trace_scope(
+    scope_name: str, 
+    scope_value: Optional[str] = None
+):
+    """
+    Context manager to start and stop a scope. All the spans, across traces created within the encapsulated code will have the scope attached.
+    
+    Args:
+        scope_name: The name of the scope.
+        scope_value: Optional value of the scope. If None, a random UUID will be generated.
+    """
+    token = None
+    if scope_name:
+        token = start_scope(scope_name, scope_value)
+    try:
+        yield
+    finally:
+        stop_scope(token)
+
+@asynccontextmanager
+async def amonocle_trace_scope(
+    scope_name: str, 
+    scope_value: Optional[str] = None
+):
+    """
+    Async context manager to start and stop a scope. All the spans, across traces created within the encapsulated code will have the scope attached.
+    
+    Args:
+        scope_name: The name of the scope.
+        scope_value: Optional value of the scope. If None, a random UUID will be generated.
+    """
+    token = start_scope(scope_name, scope_value)
+    try:
+        yield
+    finally:
+        stop_scope(token)
+    
+def monocle_trace_scope_method(
+    scope_name: str, 
+    scope_value: Optional[str] = None
+):
+    """
+    Decorator to start and stop a scope for a method. All the spans, across traces created in the method will have the scope attached.
+    
+    Args:
+        scope_name: The name of the scope.
+        scope_value: Optional value of the scope. If None, a random UUID will be generated.
+    """
+    def decorator(func):
+        if inspect.iscoroutinefunction(func):
+            @wraps(func)
+            async def wrapper(*args, **kwargs):
+                async with amonocle_trace_scope(
+                    scope_name, scope_value
+                ):
+                    result = await func(*args, **kwargs)
+                    return result
+            return wrapper
+        else:
+            @wraps(func)
+            def wrapper(*args, **kwargs):
+                with monocle_trace_scope(
+                    scope_name, scope_value
+                ):
+                    result = func(*args, **kwargs)
+                    return result
+            return wrapper
+    return decorator

monocle_apptrace/instrumentation/common/span_handler.py

@@ -25,6 +25,7 @@ WORKFLOW_TYPE_MAP = {
     "openai": "workflow.openai",
     "anthropic": "workflow.anthropic",
     "gemini": "workflow.gemini",
+    "litellm": "workflow.litellm",
 }
 
 FRAMEWORK_WORKFLOW_LIST = [
@@ -32,6 +33,7 @@ FRAMEWORK_WORKFLOW_LIST = [
     "workflow.langchain",
     "workflow.haystack",
     "workflow.teams_ai",
+    "workflow.litellm",
 ]
 class SpanHandler:
 
@@ -63,11 +65,16 @@ class SpanHandler:
             span.set_attribute("span.type", span_type)
         else:
             logger.warning("type of span not found or incorrect written in entity json")
+        if "subtype" in output_processor:
+            span.set_attribute("span.subtype", output_processor["subtype"])
         return span_type
 
     def pre_task_processing(self, to_wrap, wrapped, instance, args,kwargs, span):
-        if "pipeline" in to_wrap['package']:
-            set_attribute(QUERY, args[0]['prompt_builder']['question'])
+        try:
+            if "pipeline" in to_wrap['package']:
+                set_attribute(QUERY, args[0]['prompt_builder']['question'])
+        except Exception as e:
+            logger.warning("Warning: Error occurred in pre_task_processing: %s", str(e))
 
     @staticmethod
     def set_default_monocle_attributes(span: Span, source_path = "" ):
@@ -94,6 +101,17 @@ class SpanHandler:
     def post_task_processing(self, to_wrap, wrapped, instance, args, kwargs, result, ex, span:Span, parent_span:Span):
         pass
 
+    def should_skip(self, processor, instance, args, kwargs) -> bool:
+        should_skip = False
+        accessor = processor.get('should_skip')
+        if accessor:
+            arguments = {"instance":instance, "args":args, "kwargs":kwargs}
+            should_skip = accessor(arguments)
+            if not isinstance(should_skip, bool):
+                logger.warning("Warning: 'should_skip' accessor did not return a boolean value")
+                return True
+        return should_skip
+
     def hydrate_span(self, to_wrap, wrapped, instance, args, kwargs, result, span, parent_span = None, ex:Exception = None) -> bool:
         try:
             detected_error_in_attribute = self.hydrate_attributes(to_wrap, wrapped, instance, args, kwargs, result, span, parent_span)
@@ -145,7 +163,7 @@ class SpanHandler:
             span.set_attribute("entity.count", span_index)
         return detected_error
 
-    def hydrate_events(self, to_wrap, wrapped, instance, args, kwargs, ret_result, span, parent_span=None, ex:Exception=None) -> bool:
+    def hydrate_events(self, to_wrap, wrapped, instance, args, kwargs, ret_result, span: Span, parent_span=None, ex:Exception=None) -> bool:
         detected_error:bool = False
         if 'output_processor' in to_wrap and to_wrap["output_processor"] is not None:
             output_processor=to_wrap['output_processor']
@@ -181,10 +199,16 @@ class SpanHandler:
                             except Exception as e:
                                 logger.debug(f"Error evaluating accessor for attribute '{attribute_key}': {e}")
                     matching_timestamp = getattr(ret_result, "timestamps", {}).get(event_name, None)
-                    if isinstance(matching_timestamp, int):
-                        span.add_event(name=event_name, attributes=event_attributes, timestamp=matching_timestamp)
-                    else:
-                        span.add_event(name=event_name, attributes=event_attributes)
+                    alreadyExist = False
+                    for existing_event in span.events:
+                        if event_name == existing_event.name:
+                            existing_event.attributes._dict.update(event_attributes)
+                            alreadyExist = True
+                    if not alreadyExist:
+                        if isinstance(matching_timestamp, int):
+                            span.add_event(name=event_name, attributes=event_attributes, timestamp=matching_timestamp)
+                        else:
+                            span.add_event(name=event_name, attributes=event_attributes)
         return detected_error
 
     @staticmethod
@@ -215,7 +239,7 @@ class SpanHandler:
         if to_wrap is not None:
             package_name = to_wrap.get('package')
             for (package, framework_workflow_type) in WORKFLOW_TYPE_MAP.items():
-                if (package_name is not None and package in package_name):
+                if (package_name is not None and package_name.startswith(package)):
                     workflow_type = framework_workflow_type
                     break
         return workflow_type

monocle_apptrace/instrumentation/common/utils.py

@@ -11,6 +11,8 @@ from opentelemetry.propagate import extract
 from opentelemetry import baggage
 from monocle_apptrace.instrumentation.common.constants import MONOCLE_SCOPE_NAME_PREFIX, SCOPE_METHOD_FILE, SCOPE_CONFIG_PATH, llm_type_map, MONOCLE_SDK_VERSION, ADD_NEW_WORKFLOW
 from importlib.metadata import version
+from opentelemetry.trace.span import INVALID_SPAN
+_MONOCLE_SPAN_KEY = "monocle" + _SPAN_KEY
 
 T = TypeVar('T')
 U = TypeVar('U')
@@ -188,8 +190,8 @@ def __generate_scope_id() -> str:
     global scope_id_generator
     return f"{hex(scope_id_generator.generate_trace_id())}"
 
-def set_scope(scope_name: str, scope_value:str = None) -> object:
-    return set_scopes({scope_name: scope_value})
+def set_scope(scope_name: str, scope_value:str = None, context:Context = None) -> object:
+    return set_scopes({scope_name: scope_value}, context)
 
 def set_scopes(scopes:dict[str, object], baggage_context:Context = None) -> object:
     if baggage_context is None:
@@ -428,4 +430,33 @@ def patch_instance_method(obj, method_name, func):
     new_cls = type(f"Patched{cls.__name__}", (cls,), {
         method_name: func
     })
-    obj.__class__ = new_cls
+    obj.__class__ = new_cls
+
+
+def set_monocle_span_in_context(
+    span: Span, context: Optional[Context] = None
+) -> Context:
+    """Set the span in the given context.
+
+    Args:
+        span: The Span to set.
+        context: a Context object. if one is not passed, the
+            default current context is used instead.
+    """
+    ctx = set_value(_MONOCLE_SPAN_KEY, span, context=context)
+    return ctx
+
+def get_current_monocle_span(context: Optional[Context] = None) -> Span:
+    """Retrieve the current span.
+
+    Args:
+        context: A Context object. If one is not passed, the
+            default current context is used instead.
+
+    Returns:
+        The Span set in the context if it exists. INVALID_SPAN otherwise.
+    """
+    span = get_value(_MONOCLE_SPAN_KEY, context=context)
+    if span is None or not isinstance(span, Span):
+        return INVALID_SPAN
+    return span