monocle-apptrace 0.1.1__py3-none-any.whl → 0.3.0b1__py3-none-any.whl

monocle_apptrace/botocore/__init__.py

@@ -0,0 +1,9 @@
+import os
+from monocle_apptrace.utils import get_wrapper_methods_config
+
+parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+BOTOCORE_METHODS = get_wrapper_methods_config(
+    wrapper_methods_config_path=os.path.join(parent_dir, 'metamodel', 'maps', 'botocore_methods.json'),
+    attributes_config_base_path=os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+

monocle_apptrace/constants.py

@@ -3,20 +3,34 @@ AZURE_ML_ENDPOINT_ENV_NAME = "AZUREML_ENTRY_SCRIPT"
 AZURE_FUNCTION_WORKER_ENV_NAME = "FUNCTIONS_WORKER_RUNTIME"
 AZURE_APP_SERVICE_ENV_NAME = "WEBSITE_SITE_NAME"
 AWS_LAMBDA_ENV_NAME = "AWS_LAMBDA_RUNTIME_API"
+GITHUB_CODESPACE_ENV_NAME = "CODESPACES"
 
+AWS_LAMBDA_FUNCTION_IDENTIFIER_ENV_NAME = "AWS_LAMBDA_FUNCTION_NAME"
+AZURE_FUNCTION_IDENTIFIER_ENV_NAME = "WEBSITE_SITE_NAME"
+AZURE_APP_SERVICE_IDENTIFIER_ENV_NAME = "WEBSITE_DEPLOYMENT_ID"
+GITHUB_CODESPACE_IDENTIFIER_ENV_NAME = "GITHUB_REPOSITORY"
 # Azure naming reference can be found here
 # https://learn.microsoft.com/en-us/azure/cloud-adoption-framework/ready/azure-best-practices/resource-abbreviations
 AZURE_FUNCTION_NAME = "azure.func"
 AZURE_APP_SERVICE_NAME = "azure.asp"
 AZURE_ML_SERVICE_NAME = "azure.mlw"
 AWS_LAMBDA_SERVICE_NAME = "aws.lambda"
+GITHUB_CODESPACE_SERVICE_NAME = "github_codespace"
 
-azure_service_map = {
+# Env variables to identify infra service type
+service_type_map = {
     AZURE_ML_ENDPOINT_ENV_NAME: AZURE_ML_SERVICE_NAME,
     AZURE_APP_SERVICE_ENV_NAME: AZURE_APP_SERVICE_NAME,
-    AZURE_FUNCTION_WORKER_ENV_NAME: AZURE_FUNCTION_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
 }
 
-aws_service_map = {
-    AWS_LAMBDA_ENV_NAME: AWS_LAMBDA_SERVICE_NAME
+# Env variables to identify infra service name
+service_name_map = {
+    AZURE_APP_SERVICE_NAME: AZURE_APP_SERVICE_IDENTIFIER_ENV_NAME,
+    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
 }

monocle_apptrace/exporters/aws/s3_exporter.py

@@ -0,0 +1,158 @@
+import os
+import time
+import random
+import datetime
+import logging
+import asyncio
+import boto3
+from botocore.exceptions import ClientError
+from botocore.exceptions import (
+    BotoCoreError,
+    ConnectionClosedError,
+    ConnectTimeoutError,
+    EndpointConnectionError,
+    ReadTimeoutError,
+)
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+from monocle_apptrace.exporters.base_exporter import SpanExporterBase
+from typing import Sequence
+import json
+logger = logging.getLogger(__name__)
+
+class S3SpanExporter(SpanExporterBase):
+    def __init__(self, bucket_name=None, region_name=None):
+        super().__init__()
+        # Use environment variables if credentials are not provided
+        DEFAULT_FILE_PREFIX = "monocle_trace_"
+        DEFAULT_TIME_FORMAT = "%Y-%m-%d_%H.%M.%S"
+        self.max_batch_size = 500
+        self.export_interval = 1
+        self.s3_client = boto3.client(
+            's3',
+            aws_access_key_id=os.getenv('AWS_ACCESS_KEY_ID'),
+            aws_secret_access_key=os.getenv('AWS_SECRET_ACCESS_KEY'),
+            region_name=region_name,
+        )
+        self.bucket_name = bucket_name or os.getenv('MONOCLE_S3_BUCKET_NAME','default-bucket')
+        self.file_prefix = DEFAULT_FILE_PREFIX
+        self.time_format = DEFAULT_TIME_FORMAT
+        self.export_queue = []
+        self.last_export_time = time.time()
+
+        # Check if bucket exists or create it
+        if not self.__bucket_exists(self.bucket_name):
+            try:
+                self.s3_client.create_bucket(
+                Bucket=self.bucket_name,
+                CreateBucketConfiguration={'LocationConstraint': region_name}
+                 )
+                logger.info(f"Bucket {self.bucket_name} created successfully.")
+            except ClientError as e:
+                logger.error(f"Error creating bucket {self.bucket_name}: {e}")
+                raise e
+
+    def __bucket_exists(self, bucket_name):
+        try:
+            # Check if the bucket exists by calling head_bucket
+            self.s3_client.head_bucket(Bucket=bucket_name)
+            return True
+        except ClientError as e:
+            error_code = e.response['Error']['Code']
+            if error_code == '404':
+                # Bucket not found
+                logger.error(f"Bucket {bucket_name} does not exist (404).")
+                return False
+            elif error_code == '403':
+                # Permission denied
+                logger.error(f"Access to bucket {bucket_name} is forbidden (403).")
+                raise PermissionError(f"Access to bucket {bucket_name} is forbidden.")
+            elif error_code == '400':
+                # Bad request or malformed input
+                logger.error(f"Bad request for bucket {bucket_name} (400).")
+                raise ValueError(f"Bad request for bucket {bucket_name}.")
+            else:
+                # Other client errors
+                logger.error(f"Unexpected error when accessing bucket {bucket_name}: {e}")
+                raise e
+        except TypeError as e:
+            # Handle TypeError separately
+            logger.error(f"Type error while checking bucket existence: {e}")
+            raise e
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        """Synchronous export method that internally handles async logic."""
+        try:
+            # Run the asynchronous export logic in an event loop
+            asyncio.run(self.__export_async(spans))
+            return SpanExportResult.SUCCESS
+        except Exception as e:
+            logger.error(f"Error exporting spans: {e}")
+            return SpanExportResult.FAILURE
+
+    async def __export_async(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        try:
+            # Add spans to the export queue
+            for span in spans:
+                self.export_queue.append(span)
+                # If the queue reaches MAX_BATCH_SIZE, export the spans
+                if len(self.export_queue) >= self.max_batch_size:
+                    await self.__export_spans()
+
+            # Check if it's time to force a flush
+            current_time = time.time()
+            if current_time - self.last_export_time >= self.export_interval:
+                await self.__export_spans()  # Export spans if time interval has passed
+                self.last_export_time = current_time  # Reset the last export time
+
+            return SpanExportResult.SUCCESS
+        except Exception as e:
+            logger.error(f"Error exporting spans: {e}")
+            return SpanExportResult.FAILURE
+
+    def __serialize_spans(self, spans: Sequence[ReadableSpan]) -> str:
+        try:
+            # Serialize spans to JSON or any other format you prefer
+            valid_json_list = []
+            for span in spans:
+                try:
+                    valid_json_list.append(span.to_json(indent=0).replace("\n", ""))
+                except json.JSONDecodeError as e:
+                    logger.warning(f"Invalid JSON format in span data: {span.context.span_id}. Error: {e}")
+                    continue
+            ndjson_data = "\n".join(valid_json_list) + "\n"
+            return ndjson_data
+        except Exception as e:
+            logger.warning(f"Error serializing spans: {e}")
+
+
+    async def __export_spans(self):
+        if len(self.export_queue) == 0:
+            return
+
+        # Take a batch of spans from the queue
+        batch_to_export = self.export_queue[:self.max_batch_size]
+        serialized_data = self.__serialize_spans(batch_to_export)
+        self.export_queue = self.export_queue[self.max_batch_size:]
+        try:
+            self.__upload_to_s3(serialized_data)
+        except Exception as e:
+            logger.error(f"Failed to upload span batch: {e}")
+
+    @SpanExporterBase.retry_with_backoff(exceptions=(EndpointConnectionError, ConnectionClosedError, ReadTimeoutError, ConnectTimeoutError))
+    def __upload_to_s3(self, span_data_batch: str):
+        current_time = datetime.datetime.now().strftime(self.time_format)
+        file_name = f"{self.file_prefix}{current_time}.ndjson"
+        self.s3_client.put_object(
+            Bucket=self.bucket_name,
+            Key=file_name,
+            Body=span_data_batch
+        )
+        logger.info(f"Span batch uploaded to AWS S3 as {file_name}.")
+
+    async def force_flush(self, timeout_millis: int = 30000) -> bool:
+        await self.__export_spans()  # Export any remaining spans in the queue
+        return True
+
+    def shutdown(self) -> None:
+        logger.info("S3SpanExporter has been shut down.")

monocle_apptrace/exporters/azure/blob_exporter.py

@@ -0,0 +1,125 @@
+import os
+import time
+import random
+import datetime
+import logging
+import asyncio
+from azure.storage.blob import BlobServiceClient, BlobClient, ContainerClient
+from azure.core.exceptions import ResourceNotFoundError, ClientAuthenticationError, ServiceRequestError
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+from typing import Sequence
+from monocle_apptrace.exporters.base_exporter import SpanExporterBase
+import json
+logger = logging.getLogger(__name__)
+
+class AzureBlobSpanExporter(SpanExporterBase):
+    def __init__(self, connection_string=None, container_name=None):
+        super().__init__()
+        DEFAULT_FILE_PREFIX = "monocle_trace_"
+        DEFAULT_TIME_FORMAT = "%Y-%m-%d_%H.%M.%S"
+        self.max_batch_size = 500
+        self.export_interval = 1
+        # Use default values if none are provided
+        if not connection_string:
+            connection_string = os.getenv('MONOCLE_BLOB_CONNECTION_STRING')
+            if not connection_string:
+                raise ValueError("Azure Storage connection string is not provided or set in environment variables.")
+
+        if not container_name:
+            container_name = os.getenv('MONOCLE_BLOB_CONTAINER_NAME', 'default-container')
+
+        self.blob_service_client = BlobServiceClient.from_connection_string(connection_string)
+        self.container_name = container_name
+        self.file_prefix = DEFAULT_FILE_PREFIX
+        self.time_format = DEFAULT_TIME_FORMAT
+
+        # Check if container exists or create it
+        if not self.__container_exists(container_name):
+            try:
+                self.blob_service_client.create_container(container_name)
+                logger.info(f"Container {container_name} created successfully.")
+            except Exception as e:
+                logger.error(f"Error creating container {container_name}: {e}")
+                raise e
+
+    def __container_exists(self, container_name):
+        try:
+            container_client = self.blob_service_client.get_container_client(container_name)
+            container_client.get_container_properties()
+            return True
+        except ResourceNotFoundError:
+            logger.error(f"Container {container_name} not found (404).")
+            return False
+        except ClientAuthenticationError:
+            logger.error(f"Access to container {container_name} is forbidden (403).")
+            raise PermissionError(f"Access to container {container_name} is forbidden.")
+        except Exception as e:
+            logger.error(f"Unexpected error when checking if container {container_name} exists: {e}")
+            raise e
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        """Synchronous export method that internally handles async logic."""
+        try:
+            # Run the asynchronous export logic in an event loop
+            asyncio.run(self._export_async(spans))
+            return SpanExportResult.SUCCESS
+        except Exception as e:
+            logger.error(f"Error exporting spans: {e}")
+            return SpanExportResult.FAILURE
+
+    async def _export_async(self, spans: Sequence[ReadableSpan]):
+        """The actual async export logic is run here."""
+        # Add spans to the export queue
+        for span in spans:
+            self.export_queue.append(span)
+            if len(self.export_queue) >= self.max_batch_size:
+                await self.__export_spans()
+
+        # Force a flush if the interval has passed
+        current_time = time.time()
+        if current_time - self.last_export_time >= self.export_interval:
+            await self.__export_spans()
+            self.last_export_time = current_time
+
+    def __serialize_spans(self, spans: Sequence[ReadableSpan]) -> str:
+        try:
+            valid_json_list = []
+            for span in spans:
+                try:
+                    valid_json_list.append(span.to_json(indent=0).replace("\n", ""))
+                except json.JSONDecodeError as e:
+                    logger.warning(f"Invalid JSON format in span data: {span.context.span_id}. Error: {e}")
+                    continue
+
+            ndjson_data = "\n".join(valid_json_list) + "\n"
+            return ndjson_data
+        except Exception as e:
+            logger.warning(f"Error serializing spans: {e}")
+
+    async def __export_spans(self):
+        if len(self.export_queue) == 0:
+            return
+
+        batch_to_export = self.export_queue[:self.max_batch_size]
+        serialized_data = self.__serialize_spans(batch_to_export)
+        self.export_queue = self.export_queue[self.max_batch_size:]
+        try:
+            self.__upload_to_blob(serialized_data)
+        except Exception as e:
+            logger.error(f"Failed to upload span batch: {e}")
+
+    @SpanExporterBase.retry_with_backoff(exceptions=(ResourceNotFoundError, ClientAuthenticationError, ServiceRequestError))
+    def __upload_to_blob(self, span_data_batch: str):
+        current_time = datetime.datetime.now().strftime(self.time_format)
+        file_name = f"{self.file_prefix}{current_time}.ndjson"
+        blob_client = self.blob_service_client.get_blob_client(container=self.container_name, blob=file_name)
+        blob_client.upload_blob(span_data_batch, overwrite=True)
+        logger.info(f"Span batch uploaded to Azure Blob Storage as {file_name}.")
+
+    async def force_flush(self, timeout_millis: int = 30000) -> bool:
+        await self.__export_spans()
+        return True
+
+    def shutdown(self) -> None:
+        logger.info("AzureBlobSpanExporter has been shut down.")

monocle_apptrace/exporters/base_exporter.py

@@ -0,0 +1,48 @@
+import time
+import random
+import logging
+from abc import ABC, abstractmethod
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+from typing import Sequence
+import asyncio
+
+logger = logging.getLogger(__name__)
+
+class SpanExporterBase(ABC):
+    def __init__(self):
+        self.backoff_factor = 2
+        self.max_retries = 10
+        self.export_queue = []
+        self.last_export_time = time.time()
+
+    @abstractmethod
+    async def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        pass
+
+    @abstractmethod
+    async def force_flush(self, timeout_millis: int = 30000) -> bool:
+        pass
+
+    def shutdown(self) -> None:
+        pass
+
+    @staticmethod
+    def retry_with_backoff(retries=3, backoff_in_seconds=1, max_backoff_in_seconds=32, exceptions=(Exception,)):
+        def decorator(func):
+            def wrapper(*args, **kwargs):
+                attempt = 0
+                while attempt < retries:
+                    try:
+                        return func(*args, **kwargs)
+                    except exceptions as e:
+                        attempt += 1
+                        sleep_time = min(max_backoff_in_seconds, backoff_in_seconds * (2 ** (attempt - 1)))
+                        sleep_time = sleep_time * (1 + random.uniform(-0.1, 0.1))  # Add jitter
+                        logger.warning(f"Network connectivity error, Attempt {attempt} failed: {e}. Retrying in {sleep_time:.2f} seconds...")
+                        time.sleep(sleep_time)
+                raise Exception(f"Failed after {retries} attempts")
+
+            return wrapper
+
+        return decorator

monocle_apptrace/exporters/exporter_processor.py

@@ -0,0 +1,19 @@
+from abc import ABC, abstractmethod
+import logging
+from typing import Callable
+
+logger = logging.getLogger(__name__)
+
+class ExportTaskProcessor(ABC):
+
+    @abstractmethod
+    def start(self):
+        return
+
+    @abstractmethod
+    def stop(self):
+        return
+
+    @abstractmethod
+    def queue_task(self, async_task: Callable[[Callable, any], any] = None, args: any = None):
+        return

monocle_apptrace/exporters/monocle_exporters.py

@@ -0,0 +1,27 @@
+from typing import Dict, Any
+import os, warnings
+from importlib import import_module
+from opentelemetry.sdk.trace.export import SpanExporter, ConsoleSpanExporter
+from monocle_apptrace.exporters.file_exporter import FileSpanExporter
+
+monocle_exporters:Dict[str, Any] = {
+    "s3": {"module": "monocle_apptrace.exporters.aws.s3_exporter", "class": "S3SpanExporter"},
+    "blob": {"module":"monocle_apptrace.exporters.azure.blob_exporter", "class": "AzureBlobSpanExporter"},
+    "okahu": {"module":"monocle_apptrace.exporters.okahu.okahu_exporter", "class": "OkahuSpanExporter"},
+    "file": {"module":"monocle_apptrace.exporters.file_exporter", "class": "FileSpanExporter"}
+}
+
+def get_monocle_exporter() -> SpanExporter:
+    exporter_name = os.environ.get("MONOCLE_EXPORTER", "file")
+    try:
+        exporter_class_path  = monocle_exporters[exporter_name]
+    except Exception as ex:
+        warnings.warn(f"Unsupported Monocle span exporter setting {exporter_name}, using default FileSpanExporter.")
+        return FileSpanExporter()
+    try:
+        exporter_module = import_module(exporter_class_path.get("module"))
+        exporter_class = getattr(exporter_module, exporter_class_path.get("class"))
+        return exporter_class()
+    except Exception as ex:
+        warnings.warn(f"Unable to set Monocle span exporter to {exporter_name}, error {ex}. Using ConsoleSpanExporter")
+        return ConsoleSpanExporter()

monocle_apptrace/exporters/okahu/okahu_exporter.py

@@ -0,0 +1,115 @@
+import json
+import logging
+import os
+from typing import Callable, Optional, Sequence
+import requests
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult, ConsoleSpanExporter
+from requests.exceptions import ReadTimeout
+
+from monocle_apptrace.exporters.exporter_processor import ExportTaskProcessor
+
+REQUESTS_SUCCESS_STATUS_CODES = (200, 202)
+OKAHU_PROD_INGEST_ENDPOINT = "https://ingest.okahu.co/api/v1/trace/ingest"
+
+logger = logging.getLogger(__name__)
+
+
+class OkahuSpanExporter(SpanExporter):
+    def __init__(
+            self,
+            endpoint: Optional[str] = None,
+            timeout: Optional[int] = None,
+            session: Optional[requests.Session] = None,
+            task_processor: ExportTaskProcessor = None
+    ):
+        """Okahu exporter."""
+        okahu_endpoint: str = os.environ.get("OKAHU_INGESTION_ENDPOINT", OKAHU_PROD_INGEST_ENDPOINT)
+        self.endpoint = endpoint or okahu_endpoint
+        api_key: str = os.environ.get("OKAHU_API_KEY")
+        self._closed = False
+        if not api_key:
+            raise ValueError("OKAHU_API_KEY not set.")
+        self.timeout = timeout or 15
+        self.session = session or requests.Session()
+        self.session.headers.update(
+            {"Content-Type": "application/json", "x-api-key": api_key}
+        )
+
+        self.task_processor = task_processor or None
+        if task_processor is not None:
+            task_processor.start()
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        # After the call to Shutdown subsequent calls to Export are
+        # not allowed and should return a Failure result
+        if not hasattr(self, 'session'):
+            return self.exporter.export(spans)
+
+        if self._closed:
+            logger.warning("Exporter already shutdown, ignoring batch")
+            return SpanExportResult.FAILUREencoder
+        if len(spans) == 0:
+            return
+
+        span_list = {
+            "batch": []
+        }
+
+        # append the batch object with all the spans object
+        for span in spans:
+            # create a object from serialized span
+            obj = json.loads(span.to_json())
+            if obj["parent_id"] is None:
+                obj["parent_id"] = "None"
+            else:
+                obj["parent_id"] = remove_0x_from_start(obj["parent_id"])
+            if obj["context"] is not None:
+                obj["context"]["trace_id"] = remove_0x_from_start(obj["context"]["trace_id"])
+                obj["context"]["span_id"] = remove_0x_from_start(obj["context"]["span_id"])
+            span_list["batch"].append(obj)
+
+        def send_spans_to_okahu(span_list_local=None):
+            try:
+                result = self.session.post(
+                    url=self.endpoint,
+                    data=json.dumps(span_list_local),
+                    timeout=self.timeout,
+                )
+                if result.status_code not in REQUESTS_SUCCESS_STATUS_CODES:
+                    logger.error(
+                        "Traces cannot be uploaded; status code: %s, message %s",
+                        result.status_code,
+                        result.text,
+                    )
+                    return SpanExportResult.FAILURE
+                logger.warning("spans successfully exported to okahu")
+                return SpanExportResult.SUCCESS
+            except ReadTimeout as e:
+                logger.warning("Trace export timed out: %s", str(e))
+                return SpanExportResult.FAILURE
+
+        # if async task function is present, then push the request to asnc task
+
+        if self.task_processor is not None and callable(self.task_processor.queue_task):
+            self.task_processor.queue_task(send_spans_to_okahu, span_list)
+            return SpanExportResult.SUCCESS
+        return send_spans_to_okahu(span_list)
+
+    def shutdown(self) -> None:
+        if self._closed:
+            logger.warning("Exporter already shutdown, ignoring call")
+            return
+        if hasattr(self, 'session'):
+            self.session.close()
+        self._closed = True
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
+
+
+# only removes the first occurrence of 0x from the string
+def remove_0x_from_start(my_str: str):
+    if my_str.startswith("0x"):
+        return my_str.replace("0x", "", 1)
+    return my_str

monocle_apptrace/haystack/__init__.py

@@ -1,9 +1,9 @@
-
 import os
 import logging
-from monocle_apptrace.utils import load_wrapper_from_config
+from monocle_apptrace.utils import get_wrapper_methods_config
 
 logger = logging.getLogger(__name__)
 parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
-HAYSTACK_METHODS = load_wrapper_from_config(
-    os.path.join(parent_dir, 'metamodel', 'maps', 'haystack_methods.json'))
+HAYSTACK_METHODS = get_wrapper_methods_config(
+    wrapper_methods_config_path=os.path.join(parent_dir, 'metamodel', 'maps', 'haystack_methods.json'),
+    attributes_config_base_path=os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

monocle_apptrace/haystack/wrap_pipeline.py

@@ -4,8 +4,8 @@ from opentelemetry.context import attach, set_value
 from opentelemetry.instrumentation.utils import (
     _SUPPRESS_INSTRUMENTATION_KEY,
 )
-from monocle_apptrace.wrap_common import PROMPT_INPUT_KEY, PROMPT_OUTPUT_KEY, WORKFLOW_TYPE_MAP, with_tracer_wrapper
-from monocle_apptrace.utils import set_embedding_model
+from monocle_apptrace.wrap_common import PROMPT_INPUT_KEY, PROMPT_OUTPUT_KEY, WORKFLOW_TYPE_MAP, with_tracer_wrapper, DATA_INPUT_KEY
+from monocle_apptrace.utils import set_embedding_model, set_attribute
 
 logger = logging.getLogger(__name__)
 
@@ -20,6 +20,7 @@ def wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
     workflow_input = get_workflow_input(args, inputs)
     embedding_model = get_embedding_model(instance)
     set_embedding_model(embedding_model)
+    set_attribute(DATA_INPUT_KEY, workflow_input)
 
 
     with tracer.start_as_current_span(f"{name}.workflow") as span:

monocle_apptrace/instrumentor.py

@@ -1,7 +1,5 @@
-
-
-import logging
-from typing import Collection,List
+import logging, os
+from typing import Collection, List
 from wrapt import wrap_function_wrapper
 from opentelemetry.trace import get_tracer
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
@@ -11,17 +9,16 @@ from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanProcessor
 from opentelemetry.sdk.resources import SERVICE_NAME, Resource
 from opentelemetry import trace
 from opentelemetry.context import get_value, attach, set_value
+from monocle_apptrace.utils import process_wrapper_method_config
 from monocle_apptrace.wrap_common import SESSION_PROPERTIES_KEY
 from monocle_apptrace.wrapper import INBUILT_METHODS_LIST, WrapperMethod
-from monocle_apptrace.exporters.file_exporter import FileSpanExporter
-
+from monocle_apptrace.exporters.monocle_exporters import get_monocle_exporter
 
 logger = logging.getLogger(__name__)
 
 _instruments = ()
 
 class MonocleInstrumentor(BaseInstrumentor):
-
     workflow_name: str = ""
     user_wrapper_methods: list[WrapperMethod] = []
     instrumented_method_list: list[object] = []
@@ -36,8 +33,8 @@ class MonocleInstrumentor(BaseInstrumentor):
         return _instruments
 
     def _instrument(self, **kwargs):
-        tracer_provider = kwargs.get("tracer_provider")
-        tracer = get_tracer(instrumenting_module_name= __name__, tracer_provider= tracer_provider)
+        tracer_provider: TracerProvider = kwargs.get("tracer_provider")
+        tracer = get_tracer(instrumenting_module_name="monocle_apptrace", tracer_provider=tracer_provider)
 
         user_method_list = [
             {
@@ -46,8 +43,9 @@ class MonocleInstrumentor(BaseInstrumentor):
                 "method": method.method,
                 "span_name": method.span_name,
                 "wrapper": method.wrapper,
+                "output_processor": method.output_processor
             } for method in self.user_wrapper_methods]
-
+        process_wrapper_method_config(user_method_list)
         final_method_list = user_method_list + INBUILT_METHODS_LIST
 
         for wrapped_method in final_method_list:
@@ -69,7 +67,6 @@ class MonocleInstrumentor(BaseInstrumentor):
                                 object:{wrap_object},
                                 method:{wrap_method}""")
 
-
     def _uninstrument(self, **kwargs):
         for wrapped_method in self.instrumented_method_list:
             try:
@@ -93,8 +90,9 @@ def setup_monocle_telemetry(
     resource = Resource(attributes={
         SERVICE_NAME: workflow_name
     })
-    span_processors = span_processors or [BatchSpanProcessor(FileSpanExporter())]
+    span_processors = span_processors or [BatchSpanProcessor(get_monocle_exporter())]
     trace_provider = TracerProvider(resource=resource)
+    attach(set_value("workflow_name", workflow_name))
     tracer_provider_default = trace.get_tracer_provider()
     provider_type = type(tracer_provider_default).__name__
     is_proxy_provider = "Proxy" in provider_type
@@ -102,15 +100,14 @@ def setup_monocle_telemetry(
         processor.on_start = on_processor_start
         if not is_proxy_provider:
             tracer_provider_default.add_span_processor(processor)
-        else :
+        else:
             trace_provider.add_span_processor(processor)
-    if is_proxy_provider :
+    if is_proxy_provider:
         trace.set_tracer_provider(trace_provider)
     instrumentor = MonocleInstrumentor(user_wrapper_methods=wrapper_methods or [])
     # instrumentor.app_name = workflow_name
     if not instrumentor.is_instrumented_by_opentelemetry:
-        instrumentor.instrument()
-
+        instrumentor.instrument(trace_provider=trace_provider)
 
 def on_processor_start(span: Span, parent_context):
     context_properties = get_value(SESSION_PROPERTIES_KEY)
@@ -121,4 +118,4 @@ def on_processor_start(span: Span, parent_context):
             )
 
 def set_context_properties(properties: dict) -> None:
-    attach(set_value(SESSION_PROPERTIES_KEY, properties))
+    attach(set_value(SESSION_PROPERTIES_KEY, properties))