payi 0.1.0a110__py3-none-any.whl → 0.1.0a137__py3-none-any.whl

payi/lib/instrument.py

@@ -1,150 +1,57 @@
+from __future__ import annotations
+
 import os
+import copy
 import json
 import time
 import uuid
+import atexit
 import asyncio
 import inspect
 import logging
+import threading
 import traceback
-from abc import abstractmethod
 from enum import Enum
-from typing import Any, Set, Union, Callable, Optional, Sequence, TypedDict
+from typing import Any, Set, Union, Optional, Sequence, TypedDict, cast
 from datetime import datetime, timezone
-from dataclasses import dataclass
 
 import nest_asyncio  # type: ignore
-from wrapt import ObjectProxy  # type: ignore
+from wrapt import wrap_function_wrapper  # type: ignore
 
 from payi import Payi, AsyncPayi, APIStatusError, APIConnectionError, __version__ as _payi_version
 from payi.types import IngestUnitsParams
-from payi.lib.helpers import PayiHeaderNames
+from payi.lib.helpers import PayiHeaderNames, _compact_json
 from payi.types.shared import XproxyResult
 from payi.types.ingest_response import IngestResponse
-from payi.types.ingest_units_params import Units, ProviderResponseFunctionCall
 from payi.types.shared.xproxy_error import XproxyError
 from payi.types.pay_i_common_models_api_router_header_info_param import PayICommonModelsAPIRouterHeaderInfoParam
 
 from .helpers import PayiCategories
 from .Stopwatch import Stopwatch
+from .StreamWrappers import _GeneratorWrapper, _StreamManagerWrapper, _StreamIteratorWrapper
+from .ProviderRequest import PriceAs, _StreamingType, _ProviderRequest
 
 global _g_logger
 _g_logger: logging.Logger = logging.getLogger("payi.instrument")
 
-@dataclass
-class _ChunkResult:
-    send_chunk_to_caller: bool
-    ingest: bool = False
-    
-class _ProviderRequest:
-    def __init__(
-            self, 
-            instrumentor: '_PayiInstrumentor',
-            category: str,
-            streaming_type: '_StreamingType',
-            module_name: str,
-            module_version: str,
-            is_aws_client: Optional[bool] = None,
-            is_google_vertex_or_genai_client: Optional[bool] = None,
-            ) -> None:
-        self._instrumentor: '_PayiInstrumentor' = instrumentor
-        self._module_name: str = module_name
-        self._module_version: str = module_version  
-        self._estimated_prompt_tokens: Optional[int] = None
-        self._category: str = category
-        self._ingest: IngestUnitsParams = { "category": category, "units": {} } # type: ignore
-        self._streaming_type: '_StreamingType' = streaming_type
-        self._is_aws_client: Optional[bool] = is_aws_client
-        self._is_google_vertex_or_genai_client: Optional[bool] = is_google_vertex_or_genai_client
-        self._function_call_builder: Optional[dict[int, ProviderResponseFunctionCall]] = None
-        self._building_function_response: bool = False
-        self._function_calls: Optional[list[ProviderResponseFunctionCall]] = None
-
-    def process_chunk(self, _chunk: Any) -> _ChunkResult:
-        return _ChunkResult(send_chunk_to_caller=True)
-
-    def process_synchronous_response(self, response: Any, log_prompt_and_response: bool, kwargs: Any) -> Optional[object]:  # noqa: ARG002
-        return None
-    
-    @abstractmethod
-    def process_request(self, instance: Any, extra_headers: 'dict[str, str]', args: Sequence[Any], kwargs: Any) -> bool:
-        ...
-    
-    def process_request_prompt(self, prompt: 'dict[str, Any]', args: Sequence[Any], kwargs: 'dict[str, Any]') -> None:
-        ...
-    
-    def process_initial_stream_response(self, response: Any) -> None:
-        pass
-
-    def remove_inline_data(self, prompt: 'dict[str, Any]') -> bool:# noqa: ARG002
-        return False
-
-    @property
-    def is_aws_client(self) -> bool:
-        return self._is_aws_client if self._is_aws_client is not None else False
-
-    @property
-    def is_google_vertex_or_genai_client(self) -> bool:
-        return self._is_google_vertex_or_genai_client if self._is_google_vertex_or_genai_client is not None else False
-
-    def process_exception(self, exception: Exception, kwargs: Any, ) -> bool: # noqa: ARG002
-        self.exception_to_semantic_failure(exception)
-        return True
-    
-    @property
-    def supports_extra_headers(self) -> bool:
-        return not self.is_aws_client and not self.is_google_vertex_or_genai_client
-    
-    @property
-    def streaming_type(self) -> '_StreamingType':
-        return self._streaming_type
-
-    def exception_to_semantic_failure(self, e: Exception) -> None:
-        exception_str = f"{type(e).__name__}"
-    
-        fields: list[str] = []
-    
-        for attr in dir(e):
-            if not attr.startswith("__"):
-                try:
-                    value = getattr(e, attr)
-                    if value and not inspect.ismethod(value) and not inspect.isfunction(value) and not callable(value):
-                        fields.append(f"{attr}={value}")
-                except Exception as _ex:
-                    pass
- 
-        existing_properties = self._ingest.get("properties", None)
-        if not existing_properties:
-            existing_properties = {}
-
-        existing_properties['system.failure'] = exception_str
-        if fields:
-            failure_description = ",".join(fields)
-            existing_properties["system.failure.description"] = failure_description[:128]
-
-        self._ingest["properties"] = existing_properties
+class PayiInstrumentModelMapping(TypedDict, total=False):
+    model: str
+    price_as_category: Optional[str]
+    price_as_resource: Optional[str]
+    # "global", "datazone", "region", "region.<region_name>"
+    resource_scope: Optional[str]   
 
-        if "http_status_code" not in self._ingest:
-            # use a non existent http status code so when presented to the user, the origin is clear
-            self._ingest["http_status_code"] = 299
+class PayiInstrumentAwsBedrockConfig(TypedDict, total=False):
+    guardrail_trace: Optional[bool]
+    add_streaming_xproxy_result: Optional[bool]
+    model_mappings: Optional[Sequence[PayiInstrumentModelMapping]]
 
-    def add_streaming_function_call(self, index: int, name: Optional[str], arguments: Optional[str]) -> None:
-        if not self._function_call_builder:
-            self._function_call_builder = {}
+class PayiInstrumentAzureOpenAiConfig(TypedDict, total=False):
+    # map deployment name known model
+    model_mappings: Sequence[PayiInstrumentModelMapping] 
 
-        if not index in self._function_call_builder:
-            self._function_call_builder[index] = ProviderResponseFunctionCall(name=name or "", arguments=arguments or "")
-        else:
-            function = self._function_call_builder[index]
-            if name:
-                function["name"] = function["name"] + name
-            if arguments:
-                function["arguments"] = (function.get("arguments", "") or "") + arguments
-
-    def add_synchronous_function_call(self, name: str, arguments: Optional[str]) -> None:
-        if not self._function_calls:
-            self._function_calls = []
-            self._ingest["provider_response_function_calls"] = self._function_calls
-        self._function_calls.append(ProviderResponseFunctionCall(name=name, arguments=arguments))
+class PayiInstrumentOfflineInstrumentationConfig(TypedDict, total=False):
+    file_name: str
 
 class PayiInstrumentConfig(TypedDict, total=False):
     proxy: bool
@@ -155,40 +62,55 @@ class PayiInstrumentConfig(TypedDict, total=False):
     use_case_name: Optional[str]
     use_case_id: Optional[str]
     use_case_version: Optional[int]
-    use_case_properties: Optional["dict[str, str]"]
+    use_case_properties: Optional["dict[str, Optional[str]]"]
     user_id: Optional[str]
     account_name: Optional[str]
     request_tags: Optional["list[str]"]
-    request_properties: Optional["dict[str, str]"]
+    request_properties: Optional["dict[str, Optional[str]]"]
+    aws_config: Optional[PayiInstrumentAwsBedrockConfig]
+    azure_openai_config: Optional[PayiInstrumentAzureOpenAiConfig]
+    offline_instrumentation: Optional[PayiInstrumentOfflineInstrumentationConfig]
 
 class PayiContext(TypedDict, total=False):
     use_case_name: Optional[str]
     use_case_id: Optional[str]
     use_case_version: Optional[int]
     use_case_step: Optional[str]
-    use_case_properties: Optional["dict[str, str]"]
+    use_case_properties: Optional["dict[str, Optional[str]]"]
     limit_ids: Optional['list[str]']
     user_id: Optional[str]
     account_name: Optional[str]
     request_tags: Optional["list[str]"]
-    request_properties: Optional["dict[str, str]"]
+    request_properties: Optional["dict[str, Optional[str]]"]
     price_as_category: Optional[str]
     price_as_resource: Optional[str]
     resource_scope: Optional[str]
     last_result: Optional[Union[XproxyResult, XproxyError]]
 
+class PayiInstanceDefaultContext(TypedDict, total=False):
+    use_case_name: Optional[str]
+    use_case_id: Optional[str]
+    use_case_version: Optional[int]
+    use_case_properties: Optional["dict[str, str]"]
+    limit_ids: Optional['list[str]']
+    user_id: Optional[str]
+    account_name: Optional[str]
+    request_properties: Optional["dict[str, str]"]
+    price_as_category: Optional[str]
+    price_as_resource: Optional[str]
+    resource_scope: Optional[str]
+
 class _Context(TypedDict, total=False):
     proxy: Optional[bool]
     use_case_name: Optional[str]
     use_case_id: Optional[str]
     use_case_version: Optional[int]
     use_case_step: Optional[str]
-    use_case_properties: Optional["dict[str, str]"]
+    use_case_properties: Optional["dict[str, Optional[str]]"]
     limit_ids: Optional['list[str]']
     user_id: Optional[str]
     account_name: Optional[str]
-    request_tags: Optional["list[str]"]
-    request_properties: Optional["dict[str, str]"]
+    request_properties: Optional["dict[str, Optional[str]]"]
     price_as_category: Optional[str]
     price_as_resource: Optional[str]
     resource_scope: Optional[str]
@@ -198,10 +120,14 @@ class _IsStreaming(Enum):
     true = 1 
     kwargs = 2
 
-class _StreamingType(Enum):
-    generator = 0
-    iterator = 1
-    stream_manager = 2
+class _ThreadLocalContextStorage(threading.local):
+    """
+    Thread-local storage for context stacks. Each thread gets its own context stack.
+    
+    Note: We don't use __init__ because threading.local's __init__ semantics are tricky.
+    Instead, we lazily initialize the context_stack attribute in the property accessor.
+    """
+    context_stack: "list[_Context]"
 
 class _InternalTrackContext:
     def __init__(
@@ -233,9 +159,6 @@ class _PayiInstrumentor:
         instruments: Union[Set[str], None] = None,
         log_prompt_and_response: bool = True,
         logger: Optional[logging.Logger] = None,
-        prompt_and_response_logger: Optional[
-            Callable[[str, "dict[str, str]"], None]
-        ] = None,  # (request id, dict of data to store) -> None
         global_config: PayiInstrumentConfig = {},
         caller_filename: str = ""
     ):
@@ -252,14 +175,16 @@ class _PayiInstrumentor:
         if self._apayi:
             _g_logger.debug(f"Pay-i instrumentor initialized with AsyncPayi instance: {self._apayi}")            
 
-        self._context_stack: list[_Context] = []  # Stack of context dictionaries
+        # Thread-local storage for context stacks - each thread gets its own stack
+        self._thread_local_storage = _ThreadLocalContextStorage()
+        
         self._log_prompt_and_response: bool = log_prompt_and_response
-        self._prompt_and_response_logger: Optional[Callable[[str, dict[str, str]], None]] = prompt_and_response_logger
 
         self._blocked_limits: set[str] = set()
         self._exceeded_limits: set[str] = set()
 
-        self._api_connection_error_last_log_time: float = time.time()
+        # by not setting to time.time() the first connection error is always logged
+        self._api_connection_error_last_log_time: float = 0
         self._api_connection_error_count: int = 0
         self._api_connection_error_window: int = global_config.get("connection_error_logging_window", 60)
         if self._api_connection_error_window < 0:
@@ -272,21 +197,43 @@ class _PayiInstrumentor:
 
         self._last_result: Optional[Union[XproxyResult, XproxyError]] = None
 
+        self._offline_instrumentation = global_config.pop("offline_instrumentation", None)
+        self._offline_ingest_packets: list[IngestUnitsParams] = []
+        self._offline_instrumentation_file_name: Optional[str] = None
+        
+        if self._offline_instrumentation is not None:
+            timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
+            self._offline_instrumentation_file_name = self._offline_instrumentation.get("file_name", f"payi_instrumentation_{timestamp}.json")
+            
+            # Register exit handler to write packets when process exits
+            atexit.register(lambda: self._write_offline_ingest_packets())
+
         global_instrumentation = global_config.pop("global_instrumentation", True)
 
+        # configure first, then instrument
+        aws_config = global_config.get("aws_config", None)
+        if aws_config:
+            from .BedrockInstrumentor import BedrockInstrumentor
+            BedrockInstrumentor.configure(aws_config=aws_config)
+
+        azure_openai_config = global_config.get("azure_openai_config", None)
+        if azure_openai_config:
+            from .OpenAIInstrumentor import OpenAiInstrumentor
+            OpenAiInstrumentor.configure(azure_openai_config=azure_openai_config)
+
         if instruments is None or "*" in instruments:
             self._instrument_all()
         else:
-            self._instrument_specific(instruments)
+            self._instrument_specific(instruments=instruments)
+
+        self._instrument_futures()
 
         if global_instrumentation:
             if "proxy" not in global_config:
                 global_config["proxy"] = self._proxy_default
 
             # Use default clients if not provided for global ingest instrumentation
-            if not self._payi and not self._apayi:
-                self._payi = Payi()
-                self._apayi = AsyncPayi()
+            self._ensure_payi_clients()
 
             if "use_case_name" not in global_config and caller_filename:
                 description = f"Default use case for {caller_filename}.py"
@@ -295,23 +242,53 @@ class _PayiInstrumentor:
                         self._payi.use_cases.definitions.create(name=caller_filename, description=description)
                     elif self._apayi:
                         self._call_async_use_case_definition_create(use_case_name=caller_filename, use_case_description=description)
+                    else:
+                        # in the case of _local_instrumentation is not None
+                        pass
                     global_config["use_case_name"] = caller_filename
                 except Exception as e:
                     self._logger.error(f"Error creating default use case definition based on file name {caller_filename}: {e}")
 
             self.__enter__()
 
-            # _init_current_context will update the currrent context stack location
+            # _init_current_context will update the current context stack location
             context: _Context = {}
+
             # Copy allowed keys from global_config into context
             # Dynamically use keys from _Context TypedDict
             context_keys = list(_Context.__annotations__.keys()) if hasattr(_Context, '__annotations__') else []
             for key in context_keys:
                 if key in global_config:
-                    context[key] = global_config[key] # type: ignore
+                    context[key] = global_config[key] # type: ignore[literal-required]
+
+            self._init_current_context(**context)
+
+    def _ensure_payi_clients(self) -> None:
+        if self._offline_instrumentation is not None:
+            return
+
+        if not self._payi and not self._apayi:
+            self._payi = Payi()
+            self._apayi = AsyncPayi()
+
+    def _instrument_futures(self) -> None:
+        """Install hooks for all common concurrent execution patterns."""
+        def _thread_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
+            return self._thread_submit_wrapper(wrapped, instance, args, kwargs)
 
-            self._init_current_context(**context) 
+        async def _task_wrapper(wrapped: Any, instance: Any, args: Any, kwargs: Any) -> Any:
+            return await self._create_task_wrapper(wrapped, instance, args, kwargs)
 
+        try:
+            wrap_function_wrapper("concurrent.futures", "ThreadPoolExecutor.submit", _thread_wrapper)
+        except Exception as e:
+            self._logger.debug(f"Error wrapping ThreadPoolExecutor.submit: {e}")
+
+        try:
+            wrap_function_wrapper("asyncio", "create_task", _task_wrapper)
+        except Exception as e:
+            self._logger.debug(f"Error wrapping asyncio.create_task: {e}")
+        
     def _instrument_all(self) -> None:
         self._instrument_openai()
         self._instrument_anthropic()
@@ -375,6 +352,93 @@ class _PayiInstrumentor:
         except Exception as e:
             self._logger.error(f"Error instrumenting Google GenAi: {e}")
 
+    def _thread_submit_wrapper(
+        self,
+        wrapped: Any,
+        _instance: Any,
+        args: Any,
+        kwargs: Any,
+    ) -> Any:
+        if len(args) > 0:
+            fn = args[0]
+            fn_args = args[1:]
+            captured_context = copy.deepcopy(self._context_safe)
+
+            def context_wrapper(*inner_args: Any, **inner_kwargs: Any) -> Any:
+                with self:
+                    # self._context_stack[-1].update(captured_context)
+                    self._init_current_context(**captured_context)
+                    return fn(*inner_args, **inner_kwargs)
+
+            return wrapped(context_wrapper, *fn_args, **kwargs)
+        return wrapped(*args, **kwargs)
+
+    async def _create_task_wrapper(
+        self,
+        wrapped: Any,
+        _instance: Any,
+        args: Any,
+        kwargs: Any,
+    ) -> Any:
+        if len(args) > 0:
+            coro = args[0]
+            captured_context = copy.deepcopy(self._context_safe)
+
+            async def context_wrapper() -> Any:
+                with self:
+                    # self._context_stack[-1].update(captured_context)
+                    self._init_current_context(**captured_context)
+                    return await coro
+            
+            return wrapped(context_wrapper(), *args[1:], **kwargs)
+        return wrapped(*args, **kwargs)
+
+    @staticmethod
+    def _model_mapping_to_context_dict(model_mappings: Sequence[PayiInstrumentModelMapping]) -> 'dict[str, _Context]':
+        context: dict[str, _Context] = {}
+        for mapping in model_mappings:
+            model = mapping.get("model", "")
+            if not model:
+                continue
+
+            price_as_category = mapping.get("price_as_category", None)
+            price_as_resource = mapping.get("price_as_resource", None)
+            resource_scope = mapping.get("resource_scope", None)
+
+            if not price_as_category and not price_as_resource:
+                continue
+
+            context[model] = _Context(
+                price_as_category=price_as_category,
+                price_as_resource=price_as_resource,
+                resource_scope=resource_scope,
+            )
+        return context
+
+    def _write_offline_ingest_packets(self) -> None:
+        if not self._offline_instrumentation_file_name or not self._offline_ingest_packets:
+            return
+            
+        try:
+            # Convert datetime objects to ISO strings for JSON serialization
+            serializable_packets: list[IngestUnitsParams] = []
+            for packet in self._offline_ingest_packets:
+                serializable_packet = packet.copy()
+                
+                # Convert datetime fields to ISO format strings
+                if 'event_timestamp' in serializable_packet and isinstance(serializable_packet['event_timestamp'], datetime):
+                    serializable_packet['event_timestamp'] = serializable_packet['event_timestamp'].isoformat()
+                    
+                serializable_packets.append(serializable_packet)
+            
+            with open(self._offline_instrumentation_file_name, 'w', encoding='utf-8') as f:
+                json.dump(serializable_packets, f)
+                
+            self._logger.debug(f"Written {len(self._offline_ingest_packets)} ingest packets to {self._offline_instrumentation_file_name}")
+            
+        except Exception as e:
+            self._logger.error(f"Error writing offline ingest packets to {self._offline_instrumentation_file_name}: {e}")
+
     @staticmethod
     def _create_logged_ingest_units(
         ingest_units: IngestUnitsParams,
@@ -391,10 +455,10 @@ class _PayiInstrumentor:
 
         return log_ingest_units
         
-    def _process_ingest_units(
-            self,
-            request: _ProviderRequest, log_data: 'dict[str, str]',
-            extra_headers: 'dict[str, str]') -> None:
+    def _after_invoke_update_request(
+        self,
+        request: _ProviderRequest,
+        extra_headers: 'dict[str, str]') -> None:
         ingest_units = request._ingest
 
         if request._module_version:
@@ -404,9 +468,14 @@ class _PayiInstrumentor:
             # convert the function call builder to a list of function calls
             ingest_units["provider_response_function_calls"] = list(request._function_call_builder.values())
 
+        if "provider_response_id" not in ingest_units or not ingest_units["provider_response_id"]:
+            ingest_units["provider_response_id"] = f"payi_{uuid.uuid4()}"
+
         if 'resource' not in ingest_units or ingest_units['resource'] == '':
             ingest_units['resource'] = "system.unknown_model"
 
+        request.merge_internal_request_properties()
+
         request_json = ingest_units.get('provider_request_json', "")
         if request_json and self._instrument_inline_data is False:
             try:
@@ -414,7 +483,7 @@ class _PayiInstrumentor:
                 if request.remove_inline_data(prompt_dict):
                     self._logger.debug(f"Removed inline data from provider_request_json")
                     # store the modified dict back as JSON string
-                    ingest_units['provider_request_json'] = json.dumps(prompt_dict)
+                    ingest_units['provider_request_json'] = _compact_json(prompt_dict)
 
             except Exception as e:
                 self._logger.error(f"Error serializing provider_request_json: {e}")
@@ -424,19 +493,6 @@ class _PayiInstrumentor:
             if not units or all(unit.get("input", 0) == 0 and unit.get("output", 0) == 0 for unit in units.values()):
                 self._logger.info('ingesting with no token counts')
 
-        if self._log_prompt_and_response and self._prompt_and_response_logger:
-            response_json = ingest_units.pop("provider_response_json", None)
-            request_json = ingest_units.pop("provider_request_json", None)
-            stack_trace = ingest_units.get("properties", {}).pop("system.stack_trace", None)  # type: ignore
-
-            if response_json is not None:
-                # response_json is a list of strings, convert a single json string
-                log_data["provider_response_json"] = json.dumps(response_json)
-            if request_json is not None:
-                log_data["provider_request_json"] = request_json
-            if stack_trace is not None:
-                log_data["stack_trace"] = stack_trace
-
     def _process_ingest_units_response(self, ingest_response: IngestResponse) -> None:
         if ingest_response.xproxy_result.limits:
             for limit_id, state in ingest_response.xproxy_result.limits.items():
@@ -479,11 +535,8 @@ class _PayiInstrumentor:
 
         self._logger.debug(f"_aingest_units")
 
-        # return early if there are no units to ingest and on a successul ingest request
-        log_data: 'dict[str,str]' = {}
         extra_headers: 'dict[str, str]' = {}
-
-        self._process_ingest_units(request, log_data=log_data, extra_headers=extra_headers)
+        self._after_invoke_update_request(request, extra_headers=extra_headers)
 
         try:
             if self._logger.isEnabledFor(logging.DEBUG):
@@ -493,6 +546,18 @@ class _PayiInstrumentor:
                 ingest_response = await self._apayi.ingest.units(**ingest_units, extra_headers=extra_headers)
             elif self._payi:
                 ingest_response = self._payi.ingest.units(**ingest_units, extra_headers=extra_headers)
+            elif self._offline_instrumentation is not None:
+                self._offline_ingest_packets.append(ingest_units.copy())
+
+                # simulate a successful ingest for local instrumentation
+                now=datetime.now(timezone.utc)
+                ingest_response = IngestResponse(
+                    event_timestamp=now,
+                    ingest_timestamp=now,
+                    request_id="local_instrumentation",
+                    xproxy_result=XproxyResult(request_id="local_instrumentation"))
+                pass
+                
             else:
                 self._logger.error("No payi instance to ingest units")
                 return XproxyError(code="configuration_error", message="No Payi or AsyncPayi instance configured for ingesting units")
@@ -502,10 +567,6 @@ class _PayiInstrumentor:
             if ingest_response:
                 self._process_ingest_units_response(ingest_response)
 
-                if ingest_response and self._log_prompt_and_response and self._prompt_and_response_logger:
-                    request_id = ingest_response.xproxy_result.request_id
-                    self._prompt_and_response_logger(request_id, log_data)  # type: ignore
-
             return ingest_response.xproxy_result
 
         except APIConnectionError as api_ex:
@@ -564,7 +625,7 @@ class _PayiInstrumentor:
             # Try to get the response body as JSON
             body = e.body
             if body is None:
-                self._logger.error("APIStatusError response has no body attribute")
+                self._logger.warning(f"Pay-i ingest exception {e}, status {e.status_code} has no body")
                 return XproxyError(code="unknown_error", message=str(e))
 
             # If body is bytes, decode to string
@@ -578,8 +639,9 @@ class _PayiInstrumentor:
             if not body_dict:
                 try:
                     body_dict = json.loads(body)  # type: ignore
-                except Exception as json_ex:
-                    self._logger.error(f"Failed to parse response body as JSON: {json_ex}")
+                except Exception:
+                    body_type = type(body).__name__ # type: ignore
+                    self._logger.warning(f"Pay-i ingest exception {e}, status {e.status_code} cannot parse response JSON body for body type {body_type}")
                     return XproxyError(code="invalid_json", message=str(e))
 
             xproxy_error = body_dict.get("xproxy_error", {})
@@ -588,7 +650,7 @@ class _PayiInstrumentor:
             return XproxyError(code=code, message=message)
 
         except Exception as ex:
-            self._logger.error(f"Exception in _process_api_status_error: {ex}")
+            self._logger.warning(f"Pay-i ingest exception {e}, status {e.status_code} processing handled exception {ex}")
             return XproxyError(code="exception", message=str(ex))
 
     def _ingest_units_worker(self, request: _ProviderRequest) -> Optional[Union[XproxyResult, XproxyError]]:
@@ -597,10 +659,8 @@ class _PayiInstrumentor:
 
         self._logger.debug(f"_ingest_units")
 
-        # return early if there are no units to ingest and on a successul ingest request
-        log_data: 'dict[str,str]' = {}
         extra_headers: 'dict[str, str]' = {}
-        self._process_ingest_units(request, log_data=log_data, extra_headers=extra_headers)
+        self._after_invoke_update_request(request, extra_headers=extra_headers)
 
         try:
             if self._payi:
@@ -612,16 +672,18 @@ class _PayiInstrumentor:
 
                 self._process_ingest_units_response(ingest_response)
 
-                if self._log_prompt_and_response and self._prompt_and_response_logger:
-                    request_id = ingest_response.xproxy_result.request_id
-                    self._prompt_and_response_logger(request_id, log_data)  # type: ignore
-
                 return ingest_response.xproxy_result
             elif self._apayi:
                 # task runs async. aingest_units will invoke the callback and post process
                 sync_response = self._call_aingest_sync(request)
                 self._logger.debug(f"_ingest_units: apayi success ({sync_response})")
                 return sync_response
+            elif self._offline_instrumentation is not None:
+                self._offline_ingest_packets.append(ingest_units.copy())
+
+                # simulate a successful ingest for local instrumentation
+                return XproxyResult(request_id="local_instrumentation")
+
             else:
                 self._logger.error("No payi instance to ingest units")
                 return XproxyError(code="configuration_error", message="No Payi or AsyncPayi instance configured for ingesting units")
@@ -639,6 +701,21 @@ class _PayiInstrumentor:
     def _ingest_units(self, request: _ProviderRequest) -> Optional[Union[XproxyResult, XproxyError]]:
         return self.set_xproxy_result(self._ingest_units_worker(request))
 
+    @property
+    def _context_stack(self) -> "list[_Context]":
+        """
+        Get the thread-local context stack. On first access per thread, 
+        initializes with the current state of the main thread's context stack.
+        """
+        # Lazy-initialize the context_stack for this thread if it doesn't exist
+        if not hasattr(self._thread_local_storage, 'context_stack'):
+            self._thread_local_storage.context_stack = []
+        
+        stack = self._thread_local_storage.context_stack
+
+        
+        return stack
+
     def _setup_call_func(
         self
         ) -> _Context:
@@ -649,6 +726,31 @@ class _PayiInstrumentor:
 
         return {}
 
+    @staticmethod
+    def _valid_str_or_none(value: Optional[str], default: Optional[str] = None) -> Optional[str]:
+        if value is None:
+            return default
+        elif len(value) == 0:
+            # an empty string explicitly blocks the default value
+            return None
+        else:
+            return value
+        
+    @staticmethod
+    def _valid_properties_or_none(value: Optional["dict[str, Optional[str]]"], default: Optional["dict[str, Optional[str]]"] = None) -> Optional["dict[str, Optional[str]]"]:
+        if value is None:
+            return default.copy() if default else None
+        elif len(value) == 0:
+            # an empty dictionary explicitly blocks the default value
+            return None
+        elif default:
+            # merge dictionaries, child overrides parent keys
+            merged = default.copy()
+            merged.update(value)
+            return merged
+        else:
+            return value.copy()
+
     def _init_current_context(
         self,
         proxy: Optional[bool] = None,
@@ -659,16 +761,15 @@ class _PayiInstrumentor:
         use_case_step: Optional[str]= None,
         user_id: Optional[str]= None,
         account_name: Optional[str]= None,
-        request_tags: Optional["list[str]"] = None,
-        request_properties: Optional["dict[str, str]"] = None,
-        use_case_properties: Optional["dict[str, str]"] = None,
+        request_properties: Optional["dict[str, Optional[str]]"] = None,
+        use_case_properties: Optional["dict[str, Optional[str]]"] = None,
         price_as_category: Optional[str] = None,
         price_as_resource: Optional[str] = None,
         resource_scope: Optional[str] = None,
         ) -> None:
 
         # there will always be a current context
-        context: _Context = self.get_context() # type: ignore
+        context: _Context = self._context # type: ignore
         parent_context: _Context = self._context_stack[-2] if len(self._context_stack) > 1 else {}
 
         parent_proxy = parent_context.get("proxy", self._proxy_default)
@@ -709,26 +810,12 @@ class _PayiInstrumentor:
             assign_use_case_values = True
 
         if assign_use_case_values:
-            context["use_case_id"] = use_case_id if use_case_id else parent_use_case_id
-            context["use_case_version"] = use_case_version if use_case_version else parent_use_case_version
-            context["use_case_step"] = use_case_step if use_case_step else parent_use_case_step
+            context["use_case_version"] = use_case_version if use_case_version is not None else parent_use_case_version
+            context["use_case_id"] =  self._valid_str_or_none(use_case_id, parent_use_case_id)
+            context["use_case_step"] = self._valid_str_or_none(use_case_step, parent_use_case_step)
 
             parent_use_case_properties = parent_context.get("use_case_properties", None)
-            if use_case_properties is not None:
-                if not use_case_properties:
-                    # an empty dictionary explicitly blocks inheriting from the parent state
-                    context["use_case_properties"] = None
-                else:
-                    if parent_use_case_properties:
-                        # merge dictionaries, child overrides parent keys
-                        merged = parent_use_case_properties.copy()
-                        merged.update(use_case_properties)
-                        context["use_case_properties"] = merged
-                    else:
-                        context["use_case_properties"] = use_case_properties.copy()
-            elif parent_use_case_properties:
-                # use the parent use_case_properties if it exists
-                context["use_case_properties"] = parent_use_case_properties.copy()
+            context["use_case_properties"] = self._valid_properties_or_none(use_case_properties, parent_use_case_properties)
 
         parent_limit_ids = parent_context.get("limit_ids", None)
         if limit_ids is None:
@@ -742,56 +829,13 @@ class _PayiInstrumentor:
             context["limit_ids"] = list(set(limit_ids) | set(parent_limit_ids)) if parent_limit_ids else limit_ids.copy()
 
         parent_user_id = parent_context.get("user_id", None)
-        if user_id is None:
-            # use the parent user_id if it exists
-            context["user_id"] = parent_user_id
-        elif len(user_id) == 0:
-            # caller passing an empty string explicitly blocks inheriting from the parent state
-            context["user_id"] = None
-        else:
-            context["user_id"] = user_id
+        context["user_id"] = self._valid_str_or_none(user_id, parent_user_id)
 
         parent_account_name = parent_context.get("account_name", None)
-        if account_name is None:
-            # use the parent account_name if it exists
-            context["account_name"] = parent_account_name
-        elif len(account_name) == 0:
-            # caller passing an empty string explicitly blocks inheriting from the parent state
-            context["account_name"] = None
-        else:
-            context["account_name"] = account_name
-
-        parent_request_tags = parent_context.get("request_tags", None)
-        if request_tags is not None:
-            if len(request_tags) == 0:
-                # caller passing an empty list explicitly blocks inheriting from the parent state
-                context["request_tags"] = None
-            else:
-                if parent_request_tags:
-                    # union of new and parent lists if the parent context contains request tags
-                    context["request_tags"] = list(set(request_tags) | set(parent_request_tags))
-                else:
-                    context["request_tags"] = request_tags.copy()
-        elif parent_request_tags:
-            # use the parent request_tags if it exists
-            context["request_tags"] = parent_request_tags.copy()
+        context["account_name"] = self._valid_str_or_none(account_name, parent_account_name)
 
         parent_request_properties = parent_context.get("request_properties", None)
-        if request_properties is not None:
-            if not request_properties:
-                # an empty dictionary explicitly blocks inheriting from the parent state
-                context["request_properties"] = None
-            else:
-                if parent_request_properties:
-                    # merge dictionaries, child overrides parent keys
-                    merged = parent_request_properties.copy()
-                    merged.update(request_properties)
-                    context["request_properties"] = merged
-                else:
-                    context["request_properties"] = request_properties.copy()
-        elif parent_request_properties:
-            # use the parent request_properties if it exists
-            context["request_properties"] = parent_request_properties.copy()
+        context["request_properties"] = self._valid_properties_or_none(request_properties, parent_request_properties)
 
         if price_as_category:
             context["price_as_category"] = price_as_category
@@ -810,9 +854,8 @@ class _PayiInstrumentor:
         use_case_version: Optional[int],        
         user_id: Optional[str],
         account_name: Optional[str],
-        request_tags: Optional["list[str]"] = None,
-        request_properties: Optional["dict[str, str]"] = None,
-        use_case_properties: Optional["dict[str, str]"] = None,
+        request_properties: Optional["dict[str, Optional[str]]"] = None,
+        use_case_properties: Optional["dict[str, Optional[str]]"] = None,
         *args: Any,
         **kwargs: Any,
     ) -> Any:
@@ -825,7 +868,6 @@ class _PayiInstrumentor:
                 use_case_version=use_case_version,
                 user_id=user_id,
                 account_name=account_name,
-                request_tags=request_tags,
                 request_properties=request_properties,
                 use_case_properties=use_case_properties
             )
@@ -841,9 +883,8 @@ class _PayiInstrumentor:
         use_case_version: Optional[int],        
         user_id: Optional[str],
         account_name: Optional[str],
-        request_tags: Optional["list[str]"] = None,
-        request_properties: Optional["dict[str, str]"] = None,
-        use_case_properties: Optional["dict[str, str]"] = None,
+        request_properties: Optional["dict[str, Optional[str]]"] = None,
+        use_case_properties: Optional["dict[str, Optional[str]]"] = None,
         *args: Any,
         **kwargs: Any,
     ) -> Any:
@@ -856,40 +897,51 @@ class _PayiInstrumentor:
                 use_case_version=use_case_version,
                 user_id=user_id,
                 account_name=account_name,
-                request_tags=request_tags,
                 request_properties=request_properties,
                 use_case_properties=use_case_properties)
             return func(*args, **kwargs)
 
     def __enter__(self) -> Any:
-        # Push a new context dictionary onto the stack
+        # Push a new context dictionary onto the thread-local stack
         self._context_stack.append({})
         return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        # Pop the current context off the stack
+        # Pop the current context off the thread-local stack
         if self._context_stack:
             self._context_stack.pop()
 
-    def get_context(self) -> Optional[_Context]:
+    @property
+    def _context(self) -> Optional[_Context]:
         # Return the current top of the stack
         return self._context_stack[-1] if self._context_stack else None
 
-    def get_context_safe(self) -> _Context:
+    @property
+    def _context_safe(self) -> _Context:
         # Return the current top of the stack
-        return self.get_context() or {}
+        return self._context or {}
+
+    def _extract_price_as(self, extra_headers: "dict[str, str]") -> PriceAs:
+        context = self._context_safe
+
+        return PriceAs(
+            category=extra_headers.pop(PayiHeaderNames.price_as_category, None) or context.get("price_as_category", None),
+            resource=extra_headers.pop(PayiHeaderNames.price_as_resource, None) or context.get("price_as_resource", None),
+            resource_scope=extra_headers.pop(PayiHeaderNames.resource_scope, None) or context.get("resource_scope", None),
+        )
 
-    def _prepare_ingest(
+    def _before_invoke_update_request(
         self,
         request: _ProviderRequest,
-        context: _Context,
         ingest_extra_headers: "dict[str, str]", # do not conflict with potential kwargs["extra_headers"]
         args: Sequence[Any],
         kwargs: 'dict[str, Any]',
     ) -> None:
 
+        # pop and ignore the request tags header since it is no longer processed
+        ingest_extra_headers.pop(PayiHeaderNames.request_tags, None)
+
         limit_ids = ingest_extra_headers.pop(PayiHeaderNames.limit_ids, None)
-        request_tags = ingest_extra_headers.pop(PayiHeaderNames.request_tags, None)
 
         use_case_name = ingest_extra_headers.pop(PayiHeaderNames.use_case_name, None)
         use_case_id = ingest_extra_headers.pop(PayiHeaderNames.use_case_id, None)
@@ -899,10 +951,11 @@ class _PayiInstrumentor:
         user_id = ingest_extra_headers.pop(PayiHeaderNames.user_id, None)
         account_name = ingest_extra_headers.pop(PayiHeaderNames.account_name, None)
 
+        request_properties = ingest_extra_headers.pop(PayiHeaderNames.request_properties, "")
+        use_case_properties = ingest_extra_headers.pop(PayiHeaderNames.use_case_properties, "")
+
         if limit_ids:
             request._ingest["limit_ids"] = limit_ids.split(",")
-        if request_tags:
-            request._ingest["request_tags"] = request_tags.split(",")
         if use_case_name:
             request._ingest["use_case_name"] = use_case_name
         if use_case_id:
@@ -915,14 +968,10 @@ class _PayiInstrumentor:
             request._ingest["user_id"] = user_id
         if account_name:
             request._ingest["account_name"] = account_name
-
-        request_properties = context.get("request_properties", None)
         if request_properties:
-            request._ingest["properties"] = request_properties
-
-        use_case_properties = context.get("use_case_properties", None)
+            request._ingest["properties"] = json.loads(request_properties)
         if use_case_properties:
-            request._ingest["use_case_properties"] = use_case_properties
+            request._ingest["use_case_properties"] = json.loads(use_case_properties)
 
         if len(ingest_extra_headers) > 0:
             request._ingest["provider_request_headers"] = [PayICommonModelsAPIRouterHeaderInfoParam(name=k, value=v) for k, v in ingest_extra_headers.items()]
@@ -946,10 +995,10 @@ class _PayiInstrumentor:
         request.process_request_prompt(provider_prompt, args, kwargs)
 
         if self._log_prompt_and_response:
-            request._ingest["provider_request_json"] = json.dumps(provider_prompt)
+            request._ingest["provider_request_json"] = _compact_json(provider_prompt)
 
         request._ingest["event_timestamp"] = datetime.now(timezone.utc)
-        
+
     async def async_invoke_wrapper(
         self,
         request: _ProviderRequest,
@@ -961,7 +1010,7 @@ class _PayiInstrumentor:
     ) -> Any:
         self._logger.debug(f"async_invoke_wrapper: instance {instance}, category {request._category}")
 
-        context = self.get_context()
+        context = self._context
 
         # Bedrock client does not have an async method
 
@@ -973,23 +1022,29 @@ class _PayiInstrumentor:
 
         # after _udpate_headers, all metadata to add to ingest is in extra_headers, keyed by the xproxy-xxx header name
         extra_headers: Optional[dict[str, str]] = kwargs.get("extra_headers")
-        if extra_headers is None:
-            extra_headers = {}
+        extra_headers = (extra_headers or {}).copy()
         self._update_extra_headers(context, extra_headers)
 
         if context.get("proxy", self._proxy_default):
-            if "extra_headers" not in kwargs and extra_headers:
+            if not request.supports_extra_headers:
+                kwargs.pop("extra_headers", None)
+            elif extra_headers:
+                # Pass the copy to the wrapped function. Assumes anthropic and openai clients
                 kwargs["extra_headers"] = extra_headers
 
             self._logger.debug(f"async_invoke_wrapper: sending proxy request")
 
             return await wrapped(*args, **kwargs)
+        
+        request._price_as = self._extract_price_as(extra_headers)
+        if not request.supports_extra_headers and "extra_headers" in kwargs:
+            kwargs.pop("extra_headers", None)
 
         current_frame = inspect.currentframe()
         # f_back excludes the current frame, strip() cleans up whitespace and newlines
         stack = [frame.strip() for frame in traceback.format_stack(current_frame.f_back)]  # type: ignore
 
-        request._ingest['properties'] = { 'system.stack_trace': json.dumps(stack) }
+        request._ingest['properties'] = { 'system.stack_trace': _compact_json(stack) }
 
         if request.process_request(instance, extra_headers, args, kwargs) is False:
             self._logger.debug(f"async_invoke_wrapper: calling wrapped instance")
@@ -1006,9 +1061,13 @@ class _PayiInstrumentor:
             stream = False
 
         try:
-            self._prepare_ingest(request, context, extra_headers, args, kwargs)
+            self._before_invoke_update_request(request, extra_headers, args, kwargs)
             self._logger.debug(f"async_invoke_wrapper: calling wrapped instance (stream={stream})")
 
+            if "extra_headers" in kwargs:
+                # replace the original extra_headers with the updated copy which has all of the Pay-i headers removed
+                kwargs["extra_headers"] = extra_headers
+
             sw.start()
             response = await wrapped(*args, **kwargs)
 
@@ -1055,6 +1114,8 @@ class _PayiInstrumentor:
         request._ingest["end_to_end_latency_ms"] = duration
         request._ingest["http_status_code"] = 200
 
+        request.add_instrumented_response_headers(response)
+
         return_result: Any = request.process_synchronous_response(
             response=response,
             log_prompt_and_response=self._log_prompt_and_response,
@@ -1064,7 +1125,8 @@ class _PayiInstrumentor:
             self._logger.debug(f"async_invoke_wrapper: process sync response return")
             return return_result
 
-        await self._aingest_units(request)
+        xproxy_result = await self._aingest_units(request)
+        request.assign_xproxy_result(response, xproxy_result)
 
         self._logger.debug(f"async_invoke_wrapper: finished")
         return response
@@ -1080,7 +1142,7 @@ class _PayiInstrumentor:
     ) -> Any:
         self._logger.debug(f"invoke_wrapper: instance {instance}, category {request._category}")
 
-        context = self.get_context()
+        context = self._context
 
         if not context:
             if not request.supports_extra_headers:
@@ -1093,26 +1155,29 @@ class _PayiInstrumentor:
 
         # after _udpate_headers, all metadata to add to ingest is in extra_headers, keyed by the xproxy-xxx header name
         extra_headers: Optional[dict[str, str]] = kwargs.get("extra_headers")
-        if extra_headers is None:
-            extra_headers = {}
+        extra_headers = (extra_headers or {}).copy()
         self._update_extra_headers(context, extra_headers)
 
         if context.get("proxy", self._proxy_default):
             if not request.supports_extra_headers:
                 kwargs.pop("extra_headers", None)
-            elif "extra_headers" not in kwargs and extra_headers:
-                # assumes anthropic and openai clients
+            elif extra_headers:
+                # Pass the copy to the wrapped function. Assumes anthropic and openai clients
                 kwargs["extra_headers"] = extra_headers
 
             self._logger.debug(f"invoke_wrapper: sending proxy request")
 
             return wrapped(*args, **kwargs)
+
+        request._price_as = self._extract_price_as(extra_headers)
+        if not request.supports_extra_headers and "extra_headers" in kwargs:
+            kwargs.pop("extra_headers", None)
         
         current_frame = inspect.currentframe()
         # f_back excludes the current frame, strip() cleans up whitespace and newlines
         stack = [frame.strip() for frame in traceback.format_stack(current_frame.f_back)]  # type: ignore
 
-        request._ingest['properties'] = { 'system.stack_trace': json.dumps(stack) }
+        request._ingest['properties'] = { 'system.stack_trace': _compact_json(stack) }
 
         if request.process_request(instance, extra_headers, args, kwargs) is False:
             self._logger.debug(f"invoke_wrapper: calling wrapped instance")
@@ -1129,9 +1194,13 @@ class _PayiInstrumentor:
             stream = False
 
         try:
-            self._prepare_ingest(request, context, extra_headers, args, kwargs)
+            self._before_invoke_update_request(request, extra_headers, args, kwargs)
             self._logger.debug(f"invoke_wrapper: calling wrapped instance (stream={stream})")
 
+            if "extra_headers" in kwargs:
+                # replace the original extra_headers with the updated copy which has all of the Pay-i headers removed
+                kwargs["extra_headers"] = extra_headers
+
             sw.start()
             response = wrapped(*args, **kwargs)
             
@@ -1188,15 +1257,19 @@ class _PayiInstrumentor:
         request._ingest["end_to_end_latency_ms"] = duration
         request._ingest["http_status_code"] = 200
 
+        request.add_instrumented_response_headers(response)
+
         return_result: Any = request.process_synchronous_response(
             response=response,
             log_prompt_and_response=self._log_prompt_and_response,
             kwargs=kwargs)
+
         if return_result:
             self._logger.debug(f"invoke_wrapper: process sync response return")
             return return_result
 
-        self._ingest_units(request)
+        xproxy_result = self._ingest_units(request)
+        request.assign_xproxy_result(response, xproxy_result)
 
         self._logger.debug(f"invoke_wrapper: finished")
         return response
@@ -1205,7 +1278,7 @@ class _PayiInstrumentor:
         self
     ) -> 'dict[str, str]':
         extra_headers: dict[str, str] = {}
-        context = self.get_context()
+        context = self._context
         if context:
             self._update_extra_headers(context, extra_headers)
 
@@ -1229,19 +1302,43 @@ class _PayiInstrumentor:
 
         context_user_id: Optional[str] = context.get("user_id")
         context_account_name: Optional[str] = context.get("account_name")
-        context_request_tags: Optional[list[str]] = context.get("request_tags")
 
         context_price_as_category: Optional[str] = context.get("price_as_category")
         context_price_as_resource: Optional[str] = context.get("price_as_resource")
         context_resource_scope: Optional[str] = context.get("resource_scope")
 
-        # headers_limit_ids = extra_headers.get(PayiHeaderNames.limit_ids, None)
-    
+        context_request_properties: Optional[dict[str, Optional[str]]] = context.get("request_properties")
+        context_use_case_properties: Optional[dict[str, Optional[str]]] = context.get("use_case_properties")
+
+        if PayiHeaderNames.request_properties in extra_headers:
+            headers_request_properties = extra_headers.get(PayiHeaderNames.request_properties, None)
+
+            if not headers_request_properties:
+                # headers_request_properties is empty, remove it from extra_headers
+                extra_headers.pop(PayiHeaderNames.request_properties, None)
+            else:
+                # leave the value in extra_headers
+                ...
+        elif context_request_properties:
+            extra_headers[PayiHeaderNames.request_properties] = _compact_json(context_request_properties)
+
+        if PayiHeaderNames.use_case_properties in extra_headers:
+            headers_use_case_properties = extra_headers.get(PayiHeaderNames.use_case_properties, None)
+
+            if not headers_use_case_properties:
+                # headers_use_case_properties is empty, remove it from extra_headers
+                extra_headers.pop(PayiHeaderNames.use_case_properties, None)
+            else:
+                # leave the value in extra_headers
+                ...
+        elif context_use_case_properties:
+            extra_headers[PayiHeaderNames.use_case_properties] = _compact_json(context_use_case_properties)
+
         # If the caller specifies limit_ids in extra_headers, it takes precedence over the decorator
         if PayiHeaderNames.limit_ids in extra_headers:
             headers_limit_ids = extra_headers.get(PayiHeaderNames.limit_ids)
 
-            if headers_limit_ids is None or len(headers_limit_ids) == 0:
+            if not headers_limit_ids:
                 # headers_limit_ids is empty, remove it from extra_headers
                 extra_headers.pop(PayiHeaderNames.limit_ids, None)
             else:   
@@ -1252,7 +1349,7 @@ class _PayiInstrumentor:
 
         if PayiHeaderNames.user_id in extra_headers:
             headers_user_id = extra_headers.get(PayiHeaderNames.user_id, None)
-            if headers_user_id is None or len(headers_user_id) == 0:
+            if not headers_user_id:
                 # headers_user_id is empty, remove it from extra_headers
                 extra_headers.pop(PayiHeaderNames.user_id, None)
             else:
@@ -1263,7 +1360,7 @@ class _PayiInstrumentor:
 
         if PayiHeaderNames.account_name in extra_headers:
             headers_account_name = extra_headers.get(PayiHeaderNames.account_name, None)
-            if headers_account_name is None or len(headers_account_name) == 0:
+            if not headers_account_name:
                 # headers_account_name is empty, remove it from extra_headers
                 extra_headers.pop(PayiHeaderNames.account_name, None)
             else:
@@ -1274,7 +1371,7 @@ class _PayiInstrumentor:
 
         if PayiHeaderNames.use_case_name in extra_headers:
             headers_use_case_name = extra_headers.get(PayiHeaderNames.use_case_name, None)
-            if headers_use_case_name is None or len(headers_use_case_name) == 0:
+            if not headers_use_case_name:
                 # headers_use_case_name is empty, remove all use case related headers
                 extra_headers.pop(PayiHeaderNames.use_case_name, None)
                 extra_headers.pop(PayiHeaderNames.use_case_id, None)
@@ -1290,10 +1387,7 @@ class _PayiInstrumentor:
             if context_use_case_version is not None:
                 extra_headers[PayiHeaderNames.use_case_version] = str(context_use_case_version)
             if context_use_case_step is not None:
-                extra_headers[PayiHeaderNames.use_case_step] = str(context_use_case_step)
-
-        if PayiHeaderNames.request_tags not in extra_headers and context_request_tags:
-            extra_headers[PayiHeaderNames.request_tags] = ",".join(context_request_tags)
+                extra_headers[PayiHeaderNames.use_case_step] = context_use_case_step
 
         if PayiHeaderNames.price_as_category not in extra_headers and context_price_as_category:
             extra_headers[PayiHeaderNames.price_as_category] = context_price_as_category
@@ -1304,16 +1398,6 @@ class _PayiInstrumentor:
         if PayiHeaderNames.resource_scope not in extra_headers and context_resource_scope:
             extra_headers[PayiHeaderNames.resource_scope] = context_resource_scope
 
-    @staticmethod
-    def update_for_vision(input: int, units: 'dict[str, Units]', estimated_prompt_tokens: Optional[int]) -> int:
-        if estimated_prompt_tokens:
-            vision = input - estimated_prompt_tokens
-            if (vision > 0):
-                units["vision"] = Units(input=vision, output=0)
-                input = estimated_prompt_tokens
-        
-        return input
-
     @staticmethod
     def payi_wrapper(func: Any) -> Any:
         def _payi_wrapper(o: Any) -> Any:
@@ -1346,351 +1430,6 @@ class _PayiInstrumentor:
 
         return _payi_awrapper
 
-class _StreamIteratorWrapper(ObjectProxy):  # type: ignore
-    def __init__(
-        self,
-        response: Any,
-        instance: Any,
-        instrumentor: _PayiInstrumentor,
-        stopwatch: Stopwatch,
-        request: _ProviderRequest,
-    ) -> None:
-
-        instrumentor._logger.debug(f"StreamIteratorWrapper: instance {instance}, category {request._category}")
-
-        request.process_initial_stream_response(response)
-
-        bedrock_from_stream: bool = False
-        if request.is_aws_client:
-            stream = response.get("stream", None)
-
-            if stream:
-                response = stream
-                bedrock_from_stream = True
-            else:
-                response = response.get("body")
-                bedrock_from_stream = False
-
-        super().__init__(response)  # type: ignore
-
-        self._response = response
-        self._instance = instance
-
-        self._instrumentor = instrumentor
-        self._stopwatch: Stopwatch = stopwatch
-        self._responses: list[str] = []
-
-        self._request: _ProviderRequest = request
-
-        self._first_token: bool = True
-        self._bedrock_from_stream: bool = bedrock_from_stream
-        self._ingested: bool = False
-        self._iter_started: bool = False
-
-    def __enter__(self) -> Any:
-        self._instrumentor._logger.debug(f"StreamIteratorWrapper: __enter__")
-        return self
-
-    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: 
-        self._instrumentor._logger.debug(f"StreamIteratorWrapper: __exit__")
-        self.__wrapped__.__exit__(exc_type, exc_val, exc_tb)  # type: ignore
-
-    async def __aenter__(self) -> Any:
-        self._instrumentor._logger.debug(f"StreamIteratorWrapper: __aenter__")
-        return self
-
-    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        self._instrumentor._logger.debug(f"StreamIteratorWrapper: __aexit__")
-        await self.__wrapped__.__aexit__(exc_type, exc_val, exc_tb)  # type: ignore
-
-    def __iter__(self) -> Any:  
-        self._iter_started = True
-        if self._request.is_aws_client:
-            # MUST reside in a separate function so that the yield statement (e.g. the generator) doesn't implicitly return its own iterator and overriding self
-            self._instrumentor._logger.debug(f"StreamIteratorWrapper: bedrock __iter__")
-            return self._iter_bedrock()
-
-        self._instrumentor._logger.debug(f"StreamIteratorWrapper: __iter__")
-        return self
-
-    def _iter_bedrock(self) -> Any:
-        # botocore EventStream doesn't have a __next__ method so iterate over the wrapped object in place
-        for event in self.__wrapped__: # type: ignore
-            result: Optional[_ChunkResult] = None
-
-            if (self._bedrock_from_stream):
-                result = self._evaluate_chunk(event)
-            else:
-                chunk = event.get('chunk') # type: ignore
-                if chunk:
-                    decode = chunk.get('bytes').decode() # type: ignore
-                    result = self._evaluate_chunk(decode)
-
-            if result and result.ingest:
-                self._stop_iteration()
-
-            yield event
-
-        self._instrumentor._logger.debug(f"StreamIteratorWrapper: bedrock iter finished")
-
-        self._stop_iteration()
-
-    def __aiter__(self) -> Any:
-        self._iter_started = True
-        self._instrumentor._logger.debug(f"StreamIteratorWrapper: __aiter__")
-        return self
-
-    def __next__(self) -> object:
-        try:
-            chunk: object = self.__wrapped__.__next__()  # type: ignore
-
-            if self._ingested:
-                self._instrumentor._logger.debug(f"StreamIteratorWrapper: __next__ already ingested, not processing chunk {chunk}")
-                return chunk # type: ignore
-
-            result = self._evaluate_chunk(chunk)
-
-            if result.ingest:
-                self._stop_iteration()
-
-            if result.send_chunk_to_caller:
-                return chunk # type: ignore
-            else:
-                return self.__next__()
-        except Exception as e:
-            if isinstance(e, StopIteration):
-                self._stop_iteration()
-            else:
-                self._instrumentor._logger.debug(f"StreamIteratorWrapper: __next__ exception {e}")
-            raise e
-
-    async def __anext__(self) -> object:
-        try:
-            chunk: object = await self.__wrapped__.__anext__()  # type: ignore
-
-            if self._ingested:
-                self._instrumentor._logger.debug(f"StreamIteratorWrapper: __next__ already ingested, not processing chunk {chunk}")
-                return chunk # type: ignore
-
-            result = self._evaluate_chunk(chunk)
-
-            if result.ingest:
-                await self._astop_iteration()
-
-            if  result.send_chunk_to_caller:
-                return chunk # type: ignore
-            else:
-                return await self.__anext__()
-
-        except Exception as e:
-            if isinstance(e, StopAsyncIteration):
-                await self._astop_iteration()
-            else:
-                self._instrumentor._logger.debug(f"StreamIteratorWrapper: __anext__ exception {e}")
-            raise e
-
-    def _evaluate_chunk(self, chunk: Any) -> _ChunkResult:
-        if self._first_token:
-            self._request._ingest["time_to_first_token_ms"] = self._stopwatch.elapsed_ms_int()
-            self._first_token = False
-
-        if self._instrumentor._log_prompt_and_response:
-            self._responses.append(self.chunk_to_json(chunk))
-
-        return self._request.process_chunk(chunk)
-
-    def _process_stop_iteration(self) -> None:
-        self._instrumentor._logger.debug(f"StreamIteratorWrapper: process stop iteration")
-
-        self._stopwatch.stop()
-        self._request._ingest["end_to_end_latency_ms"] = self._stopwatch.elapsed_ms_int()
-        self._request._ingest["http_status_code"] = 200
-
-        if self._instrumentor._log_prompt_and_response:
-            self._request._ingest["provider_response_json"] = self._responses
-
-    async def _astop_iteration(self) -> None:
-        if self._ingested:
-            self._instrumentor._logger.debug(f"StreamIteratorWrapper: astop iteration already ingested, skipping")
-            return
-
-        self._process_stop_iteration()
-
-        await self._instrumentor._aingest_units(self._request)
-        self._ingested = True
-
-    def _stop_iteration(self) -> None:
-        if self._ingested:
-            self._instrumentor._logger.debug(f"StreamIteratorWrapper: stop iteration already ingested, skipping")
-            return
-
-        self._process_stop_iteration()
-        self._instrumentor._ingest_units(self._request)
-        self._ingested = True
-
-    @staticmethod
-    def chunk_to_json(chunk: Any) -> str:
-        if hasattr(chunk, "to_json"):
-            return str(chunk.to_json())
-        elif isinstance(chunk, bytes):
-            return chunk.decode()
-        elif isinstance(chunk, str):
-            return chunk
-        else:
-            # assume dict
-            return json.dumps(chunk)
-
-class _StreamManagerWrapper(ObjectProxy):  # type: ignore
-    def __init__(
-        self,
-        stream_manager: Any,  # type: ignore
-        instance: Any,
-        instrumentor: _PayiInstrumentor, 
-        stopwatch: Stopwatch,
-        request: _ProviderRequest,
-    ) -> None:
-        instrumentor._logger.debug(f"StreamManagerWrapper: instance {instance}, category {request._category}")
-
-        super().__init__(stream_manager)  # type: ignore
-
-        self._stream_manager = stream_manager  
-        self._instance = instance
-        self._instrumentor = instrumentor
-        self._stopwatch: Stopwatch = stopwatch
-        self._responses: list[str] = []
-        self._request: _ProviderRequest = request
-        self._first_token: bool = True
-
-    def __enter__(self) -> _StreamIteratorWrapper:
-        self._instrumentor._logger.debug(f"_StreamManagerWrapper: __enter__")
-
-        return _StreamIteratorWrapper(
-            response=self.__wrapped__.__enter__(),  # type: ignore
-            instance=self._instance,
-            instrumentor=self._instrumentor,
-            stopwatch=self._stopwatch,
-            request=self._request,
-        )
-
-class _GeneratorWrapper:  # type: ignore
-    def __init__(
-        self,
-        generator: Any,
-        instance: Any,
-        instrumentor: _PayiInstrumentor, 
-        stopwatch: Stopwatch,
-        request: _ProviderRequest,
-    ) -> None:
-        instrumentor._logger.debug(f"GeneratorWrapper: instance {instance}, category {request._category}")
-
-        super().__init__()  # type: ignore
-        
-        self._generator = generator
-        self._instance = instance
-        self._instrumentor = instrumentor
-        self._stopwatch: Stopwatch = stopwatch
-        self._log_prompt_and_response: bool = instrumentor._log_prompt_and_response
-        self._responses: list[str] = []
-        self._request: _ProviderRequest = request
-        self._first_token: bool = True
-        self._ingested: bool = False
-        self._iter_started: bool = False
-
-    def __iter__(self) -> Any:
-        self._iter_started = True
-        self._instrumentor._logger.debug(f"GeneratorWrapper: __iter__")
-        return self
-        
-    def __aiter__(self) -> Any:
-        self._instrumentor._logger.debug(f"GeneratorWrapper: __aiter__")
-        return self
-
-    def _process_chunk(self, chunk: Any) -> _ChunkResult:
-        if self._first_token:
-            self._request._ingest["time_to_first_token_ms"] = self._stopwatch.elapsed_ms_int()
-            self._first_token = False
-            
-        if self._log_prompt_and_response:
-            dict = self._chunk_to_dict(chunk) 
-            self._responses.append(json.dumps(dict))
-                
-        return self._request.process_chunk(chunk)
-    
-    def __next__(self) -> Any:
-        try:
-            chunk = next(self._generator)
-            result = self._process_chunk(chunk)
-
-            if result.ingest:
-                self._stop_iteration()
-
-            # ignore result.send_chunk_to_caller:
-            return chunk
-
-        except Exception as e:
-            if isinstance(e, StopIteration):
-                self._stop_iteration()
-            else:
-                self._instrumentor._logger.debug(f"GeneratorWrapper: __next__ exception {e}")            
-            raise e
-
-    async def __anext__(self) -> Any:
-        try:
-            chunk = await anext(self._generator) # type: ignore
-            result = self._process_chunk(chunk)
-
-            if result.ingest:
-                await self._astop_iteration()
-
-            # ignore result.send_chunk_to_caller:
-            return chunk # type: ignore
-
-        except Exception as e:
-            if isinstance(e, StopAsyncIteration):
-                await self._astop_iteration()
-            else:
-                self._instrumentor._logger.debug(f"GeneratorWrapper: __anext__ exception {e}")
-            raise e
-
-    @staticmethod
-    def _chunk_to_dict(chunk: Any) -> 'dict[str, object]':
-        if hasattr(chunk, "to_dict"):
-            return chunk.to_dict() # type: ignore
-        elif hasattr(chunk, "to_json_dict"):  
-            return chunk.to_json_dict() # type: ignore
-        else:
-            return {}
-
-    def _stop_iteration(self) -> None:
-        if self._ingested:
-            self._instrumentor._logger.debug(f"GeneratorWrapper: stop iteration already ingested, skipping")
-            return
-
-        self._process_stop_iteration()
-
-        self._instrumentor._ingest_units(self._request)
-        self._ingested = True
-
-    async def _astop_iteration(self) -> None:
-        if self._ingested:
-            self._instrumentor._logger.debug(f"GeneratorWrapper: astop iteration already ingested, skipping")
-            return
-
-        self._process_stop_iteration()
-        
-        await self._instrumentor._aingest_units(self._request)
-        self._ingested = True
-
-    def _process_stop_iteration(self) -> None:
-        self._instrumentor._logger.debug(f"GeneratorWrapper: stop iteration")
-
-        self._stopwatch.stop()
-        self._request._ingest["end_to_end_latency_ms"] = self._stopwatch.elapsed_ms_int()
-        self._request._ingest["http_status_code"] = 200
-            
-        if self._log_prompt_and_response:
-            self._request._ingest["provider_response_json"] = self._responses
-
 global _instrumentor
 _instrumentor: Optional[_PayiInstrumentor] = None
 
@@ -1699,7 +1438,6 @@ def payi_instrument(
     payi: Optional[Union[Payi, AsyncPayi, 'list[Union[Payi, AsyncPayi]]']] = None,
     instruments: Optional[Set[str]] = None,
     log_prompt_and_response: bool = True,
-    prompt_and_response_logger: Optional[Callable[[str, "dict[str, str]"], None]] = None,
     config: Optional[PayiInstrumentConfig] = None,
     logger: Optional[logging.Logger] = None,
 ) -> None:
@@ -1732,7 +1470,6 @@ def payi_instrument(
         instruments=instruments,
         log_prompt_and_response=log_prompt_and_response,
         logger=logger,
-        prompt_and_response_logger=prompt_and_response_logger,
         global_config=config if config else PayiInstrumentConfig(),
         caller_filename=caller_filename
     )
@@ -1750,9 +1487,9 @@ def track(
     use_case_properties: Optional["dict[str, str]"] = None,
     proxy: Optional[bool] = None,
 ) -> Any:
+    _ = request_tags
 
     def _track(func: Any) -> Any:
-        import asyncio
         if asyncio.iscoroutinefunction(func):
             async def awrapper(*args: Any, **kwargs: Any) -> Any:
                 if not _instrumentor:
@@ -1770,9 +1507,8 @@ def track(
                     use_case_version,
                     user_id,
                     account_name,
-                    request_tags,
-                    request_properties,
-                    use_case_properties,
+                    cast(Optional['dict[str, Optional[str]]'], request_properties),
+                    cast(Optional['dict[str, Optional[str]]'], use_case_properties),
                     *args,
                     **kwargs,
                 )
@@ -1794,9 +1530,8 @@ def track(
                     use_case_version,
                     user_id,
                     account_name,
-                    request_tags,
-                    request_properties,
-                    use_case_properties,
+                    cast(Optional['dict[str, Optional[str]]'], request_properties),
+                    cast(Optional['dict[str, Optional[str]]'], use_case_properties),
                     *args,
                     **kwargs,
                 )
@@ -1834,14 +1569,15 @@ def track_context(
 
     context["user_id"] = user_id
     context["account_name"] = account_name
-    context["request_tags"] = request_tags
 
     context["price_as_category"] = price_as_category
     context["price_as_resource"] = price_as_resource
     context["resource_scope"] = resource_scope
 
-    context["request_properties"] = request_properties
-    context["use_case_properties"] = use_case_properties
+    context["request_properties"] = cast(Optional['dict[str, Optional[str]]'], request_properties)
+    context["use_case_properties"] = cast(Optional['dict[str, Optional[str]]'], use_case_properties)
+
+    _ = request_tags
 
     return _InternalTrackContext(context)
 
@@ -1852,7 +1588,7 @@ def get_context() -> PayiContext:
     """
     if not _instrumentor:
         return PayiContext()
-    internal_context = _instrumentor.get_context() or {}
+    internal_context = _instrumentor._context_safe
 
     context_dict = {
         key: value