monocle-apptrace 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

monocle_apptrace/README.md

@@ -1,37 +1,59 @@
-#Monocle User Guide
-
 ## Monocle Concepts
+
 ### Traces
-Traces are the full view of a single end-to-end application KPI eg Chatbot application to provide a response to end user’s question. Traces consists of various metadata about the application run including status, start time, duration, input/outputs etc. It also includes a list of individual steps aka “spans with details about that step.
-It’s typically the workflow code components of an application that generate the traces for application runs. 
+Traces are the full view of a single end-to-end application execution. 
+
+Examples of traces include one response to end user’s question by a chatbot app. Traces consists of various metadata about the application run including status, start time, duration, input/outputs etc. It also includes a list of individual steps aka “spans with details about that step.It’s typically the workflow code components of an application that generate the traces for application runs. 
+
+Traces are collections of spans. 
+
 ### Spans
-Spans are the individual steps executed by the application to perform a GenAI related task” eg app retrieving vectors from DB, app querying LLM for inference etc. The span includes the type of operation, start time, duration and metadata relevant to that step eg Model name, parameters and model endpoint/server for an inference request.
-It’s typically the workflow code components of an application that generate the traces for application runs.
+Spans are the individual steps executed by the application to perform a GenAI related task.
+
+Examples of spans include app retrieving vectors from DB, app querying LLM for inference etc. The span includes the type of operation, start time, duration and metadata relevant to that step eg Model name, parameters and model endpoint/server for an inference request.
 
-## Setup Monocle
-- You can download Monocle library releases from Pypi
+## Contribute to Monocle 
+
+Monocle includes:
+- Methods for instrumentation of app code 
+  - Base code for wrapping methods of interest in included in current folder
+  - Framework specific code is organized in a folder with the framework name
+- Metamodel for how attributes and events for GenAI components are represented in OpenTelemety format
+  - See [metamodel](./metamodel/README.md) for supported GenAI entities, how functions operating on those entities map to spans and format of spans 
+- Exporters to send trace data to various locations. See [exporters](./exporters)
+
+See [Monocle committer guide](/Monocle_committer_guide.md). 
+
+## Get Monocle
+
+Option 1 - Download released packages from Pypi
 ``` 
-    > python3 -m pip install pipenv
-    > pip install monocle-observability
+    python3 -m pip install pipenv
+    pip install monocle-apptrace
 ```
-- You can locally build and install Monocle library from source
+
+Option 2 - Build and install locally from source
 ```
-> pip install .
-> pip install -e ".[dev]"
+    pip install .
+    pip install -e ".[dev]"
 
-> python3 -m pip install pipenv
-> pipenv install build
+    python3 -m pip install pipenv
+    pipenv install build
 ```
 
-## Examples 
-### Enable Monocle tracing in your application
+## Examples of app instrumentation with Monocle
+ 
+### apps written using LLM orchestration frameworks 
+
 ```python
-from monocle_apptrace.instrumentor import setup_monocle_telemetry
-from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
 from langchain.chains import LLMChain
 from langchain_openai import OpenAI
 from langchain.prompts import PromptTemplate
 
+# Import the monocle_apptrace instrumentation method 
+from monocle_apptrace.instrumentor import setup_monocle_telemetry
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
+
 # Call the setup Monocle telemetry method
 setup_monocle_telemetry(workflow_name = "simple_math_app",
         span_processors=[BatchSpanProcessor(ConsoleSpanExporter())])
@@ -42,19 +64,19 @@ prompt = PromptTemplate.from_template("1 + {number} = ")
 chain = LLMChain(llm=llm, prompt=prompt)
 chain.invoke({"number":2})
 
-# Request callbacks: Finally, let's use the request `callbacks` to achieve the same result
-chain = LLMChain(llm=llm, prompt=prompt)
-chain.invoke({"number":2}, {"callbacks":[handler]})
-    
+# Trace is generated when invoke() method is called
+
 ```
 
-### Monitoring custom methods with Monocle
+### apps with custom methods
 
 ```python
+
+# Import the monocle_apptrace instrumentation method
 from monocle_apptrace.wrapper import WrapperMethod,task_wrapper,atask_wrapper
 from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter
 
-# extend the default wrapped methods list as follows
+# Extend the default wrapped methods list as follows
 app_name = "simple_math_app"
 setup_monocle_telemetry(
         workflow_name=app_name,
@@ -74,4 +96,6 @@ setup_monocle_telemetry(
                 wrapper=atask_wrapper)
         ])
 
-```
+# Trace is generated when the invoke() method is called in langchain.schema.runnable package
+
+```

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 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="us-east-1"):
+        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:
+                if region_name == "us-east-1":
+                    self.s3_client.create_bucket(Bucket=self.bucket_name)
+                else:
+                    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:
+            if asyncio.get_event_loop().is_running():
+                task = asyncio.create_task(self._retry_with_backoff(self.__upload_to_s3, serialized_data))
+                await task
+            else:
+                await self._retry_with_backoff(self.__upload_to_s3, serialized_data)
+
+        except Exception as e:
+            logger.error(f"Failed to upload span batch: {e}")
+
+    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,128 @@
+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:
+            if asyncio.get_event_loop().is_running():
+                task = asyncio.create_task(self._retry_with_backoff(self.__upload_to_blob, serialized_data))
+                await task
+            else:
+                await self._retry_with_backoff(self.__upload_to_blob, serialized_data)
+        except Exception as e:
+            logger.error(f"Failed to upload span batch: {e}")
+
+    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,47 @@
+import time
+import random
+import logging
+from abc import ABC, abstractmethod
+from azure.core.exceptions import ServiceRequestError, ClientAuthenticationError
+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
+
+    async def _retry_with_backoff(self, func, *args, **kwargs):
+        """Handle retries with exponential backoff."""
+        attempt = 0
+        while attempt < self.max_retries:
+            try:
+                return func(*args, **kwargs)
+            except ServiceRequestError as e:
+                logger.warning(f"Network connectivity error: {e}. Retrying in {self.backoff_factor ** attempt} seconds...")
+                sleep_time = self.backoff_factor * (2 ** attempt) + random.uniform(0, 1)
+                await asyncio.sleep(sleep_time)
+                attempt += 1
+            except ClientAuthenticationError as e:
+                logger.error(f"Failed to authenticate: {str(e)}")
+                break
+
+        logger.error("Max retries exceeded.")
+        raise ServiceRequestError(message="Max retries exceeded.")

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, 'wrapper_config', '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,7 +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.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__)
 
@@ -17,6 +18,10 @@ def wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
     attach(set_value("workflow_name", name))
     inputs = set()
     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:
         span.set_attribute(PROMPT_INPUT_KEY, workflow_input)
@@ -44,3 +49,15 @@ def get_workflow_input(args, inputs):
 def set_workflow_attributes(span, workflow_name):
     span.set_attribute("workflow_name",workflow_name)
     span.set_attribute("workflow_type", WORKFLOW_TYPE_MAP["haystack"])
+
+def get_embedding_model(instance):
+    try:
+        if hasattr(instance, 'get_component'):
+            text_embedder = instance.get_component('text_embedder')
+            if text_embedder and hasattr(text_embedder, 'model'):
+                # Set the embedding model attribute
+                return text_embedder.model
+    except:
+        pass
+
+    return None