payi 0.1.0a63__py3-none-any.whl → 0.1.0a65__py3-none-any.whl

payi/lib/instrument.py

@@ -21,10 +21,25 @@ from .helpers import PayiCategories
 from .Stopwatch import Stopwatch
 
 
+class _ProviderRequest:
+    def __init__(self, instrumentor: '_PayiInstrumentor'):
+        self._instrumentor: '_PayiInstrumentor' = instrumentor
+        self._estimated_prompt_tokens: Optional[int] = None
+        self._ingest: IngestUnitsParams
+
+    def process_request(self, _kwargs: Any) -> None:
+        return
+
+    def process_chunk(self, _chunk: Any) -> bool:
+        return True
+
+    def process_synchronous_response(self, response: Any, log_prompt_and_response: bool, kwargs: Any) -> Optional[object]:  # noqa: ARG002
+        return None
+
 class PayiInstrumentConfig(TypedDict, total=False):
     proxy: bool
+    global_instrumentation_enabled: bool
     limit_ids: Optional["list[str]"]
-    request_tags: Optional["list[str]"]
     experience_name: Optional[str]
     experience_id: Optional[str]
     use_case_name: Optional[str]
@@ -40,26 +55,14 @@ class _Context(TypedDict, total=False):
     use_case_id: Optional[str]
     use_case_version: Optional[int]
     limit_ids: Optional['list[str]']
-    request_tags: Optional['list[str]']
     user_id: Optional[str]
 
-class _ParentState(TypedDict, total=False):
-    experience_name: Optional[str]
-    experience_id: Optional[str]
-    use_case_name: Optional[str]
-    use_case_id: Optional[str]
-    use_case_version: Optional[int]
-    limit_ids: Optional['list[str]']
-    request_tags: Optional['list[str]']
-
 class _IsStreaming(Enum):
     false = 0
     true = 1 
     kwargs = 2
 
 class _PayiInstrumentor:
-    estimated_prompt_tokens: str = "estimated_prompt_tokens"
-
     def __init__(
         self,
         payi: Optional[Payi],
@@ -86,11 +89,23 @@ class _PayiInstrumentor:
         else:
             self._instrument_specific(instruments)
 
-        if global_config and len(global_config) > 0:
+        global_instrumentation_enabled = global_config.pop("global_instrumentation_enabled", True) if global_config else True
+
+        if global_instrumentation_enabled:
+            if global_config is None:
+                global_config = {}
+            if "proxy" not in global_config:
+                global_config["proxy"] = False
+
+            # Use default clients if not provided for global ingest instrumentation
+            if not self._payi and not self._apayi and global_config.get("proxy") == False:
+                self._payi = Payi()
+                self._apayi = AsyncPayi()
+
             context: _Context = {}
             self._context_stack.append(context)            
             # init_context will update the currrent context stack location
-            self._init_context(context=context, parentState={}, **global_config)
+            self._init_context(context=context, parentContext={}, **global_config) # type: ignore
 
     def _instrument_all(self) -> None:
         self._instrument_openai()
@@ -252,31 +267,24 @@ class _PayiInstrumentor:
 
     def _setup_call_func(
         self
-        ) -> 'tuple[_Context, _ParentState]':
+        ) -> 'tuple[_Context, _Context]':
         context: _Context = {}
-        parentState: _ParentState = {}
+        parentContext: _Context = {}
 
         if len(self._context_stack) > 0:
             # copy current context into the upcoming context
             context = self._context_stack[-1].copy()
             context.pop("proxy")
-            parentState["experience_name"] = context.get("experience_name", None)
-            parentState["experience_id"] = context.get("experience_id", None)
-            parentState["use_case_name"] = context.get("use_case_name", None)
-            parentState["use_case_id"] = context.get("use_case_id", None)
-            parentState["use_case_version"] = context.get("use_case_version", None)
-            parentState["limit_ids"] = context.get("limit_ids", None)
-            parentState["request_tags"] = context.get("request_tags", None)
+            parentContext = {**context}
 
-        return (context, parentState)
+        return (context, parentContext)
 
     def _init_context(
         self,
         context: _Context,
-        parentState: _ParentState,
+        parentContext: _Context,
         proxy: bool,
         limit_ids: Optional["list[str]"] = None,
-        request_tags: Optional["list[str]"] = None,
         experience_name: Optional[str] = None,
         experience_id: Optional[str] = None,
         use_case_name: Optional[str]= None,
@@ -286,56 +294,72 @@ class _PayiInstrumentor:
         ) -> None:
         context["proxy"] = proxy
 
-        # TODO use case what if caller specified epxerience / use_case ID and no name?
+        parent_experience_name = parentContext.get("experience_name", None)
+        parent_experience_id = parentContext.get("experience_id", None)
 
-        # Handle experience name and ID logic
-        if not experience_name:
+        if experience_name is None:
             # If no experience_name specified, use previous values
-            context["experience_name"] = parentState.get("experience_name", None)
-            context["experience_id"] = parentState.get("experience_id", None)
+            context["experience_name"] = parent_experience_name
+            context["experience_id"] = parent_experience_id
+        elif len(experience_name) == 0:
+            # Empty string explicitly blocks inheriting from the parent state
+            context["experience_name"] = None
+            context["experience_id"] = None
         else:
-            previous_experience_name = parentState.get("experience_name", None)
-            previous_experience_id = parentState.get("experience_id", None)
-
-            # If experience_name is specified
-            if experience_name == previous_experience_name:
+            # Check if experience_name is the same as the previous one
+            if experience_name == parent_experience_name:
                 # Same experience name, use previous ID unless new one specified
                 context["experience_name"] = experience_name
-                context["experience_id"] = experience_id if experience_id else previous_experience_id
+                context["experience_id"] = experience_id if experience_id else parent_experience_id
             else:
                 # Different experience name, use specified ID or generate one
                 context["experience_name"] = experience_name
                 context["experience_id"] = experience_id if experience_id else str(uuid.uuid4())
 
-        # Handle use case name and ID logic
-        if not use_case_name: # TODO use case
+        parent_use_case_name = parentContext.get("use_case_name", None)
+        parent_use_case_id = parentContext.get("use_case_id", None)
+        parent_use_case_version = parentContext.get("use_case_version", None)
+
+        if use_case_name is None:
             # If no use_case_name specified, use previous values
-            context["use_case_name"] = parentState.get("use_case_name", None)
-            context["use_case_id"] = parentState.get("use_case_id", None)
-            context["use_case_version"] = parentState.get("use_case_version", None)
+            context["use_case_name"] = parent_use_case_name
+            context["use_case_id"] = parent_use_case_id
+            context["use_case_version"] = parent_use_case_version
+        elif len(use_case_name) == 0:
+            # Empty string explicitly blocks inheriting from the parent state
+            context["use_case_name"] = None
+            context["use_case_id"] = None
+            context["use_case_version"] = None
         else:
-            previous_use_case_name = parentState.get("use_case_name", None)
-            previous_use_case_id = parentState.get("use_case_id", None)
-            previous_use_case_version = parentState.get("use_case_version", None)
-
-            # If use_case_name is specified
-            if use_case_name == previous_use_case_name:
+            if use_case_name == parent_use_case_name:
                 # Same use case name, use previous ID unless new one specified
                 context["use_case_name"] = use_case_name
-                context["use_case_id"] = use_case_id if use_case_id else previous_use_case_id
-                context["use_case_version"] = use_case_version if use_case_version else previous_use_case_version
+                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
             else:
-                # Different experience name, use specified ID or generate one
+                # Different use case name, use specified ID or generate one
                 context["use_case_name"] = use_case_name
                 context["use_case_id"] = use_case_id if use_case_id else str(uuid.uuid4())
-                context["use_case_version"] = use_case_version
-
-        # set any values explicitly passed by the caller, otherwise use what is already in the context
-        if limit_ids:
-            context["limit_ids"] = limit_ids
-        if request_tags:
-            context["request_tags"] = request_tags
-        if user_id:
+                context["use_case_version"] = use_case_version if use_case_version else None
+
+        parent_limit_ids = parentContext.get("limit_ids", None)
+        if limit_ids is None:
+            # use the parent limit_ids if it exists
+            context["limit_ids"] = parent_limit_ids
+        elif len(limit_ids) == 0:
+            # caller passing an empty array explicitly blocks inheriting from the parent state
+            context["limit_ids"] = None
+        else:
+            # union of new and parent lists if the parent context contains limit ids
+            context["limit_ids"] = list(set(limit_ids) | set(parent_limit_ids)) if parent_limit_ids else limit_ids
+
+        if user_id is None:
+            # use the parent user_id if it exists
+            context["user_id"] = parentContext.get("user_id", None)
+        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
 
         self.set_context(context)
@@ -345,7 +369,6 @@ class _PayiInstrumentor:
         func: Any,
         proxy: bool,
         limit_ids: Optional["list[str]"],
-        request_tags: Optional["list[str]"],
         experience_name: Optional[str],
         experience_id: Optional[str],
         use_case_name: Optional[str],
@@ -355,15 +378,14 @@ class _PayiInstrumentor:
         *args: Any,
         **kwargs: Any,
     ) -> Any:
-        context, parentState = self._setup_call_func()
+        context, parentContext = self._setup_call_func()
 
         with self:
             self._init_context(
                 context,
-                parentState,
+                parentContext,
                 proxy, 
                 limit_ids, 
-                request_tags,
                 experience_name,
                 experience_id,
                 use_case_name,
@@ -377,7 +399,6 @@ class _PayiInstrumentor:
         func: Any,
         proxy: bool,
         limit_ids: Optional["list[str]"],
-        request_tags: Optional["list[str]"],
         experience_name: Optional[str],
         experience_id: Optional[str],
         use_case_name: Optional[str],
@@ -387,15 +408,14 @@ class _PayiInstrumentor:
         *args: Any,
         **kwargs: Any,
     ) -> Any:
-        context, parentState = self._setup_call_func()
+        context, parentContext = self._setup_call_func()
 
         with self:
             self._init_context(
                 context,
-                parentState,
+                parentContext,
                 proxy, 
                 limit_ids, 
-                request_tags,
                 experience_name,
                 experience_id,
                 use_case_name,
@@ -477,9 +497,7 @@ class _PayiInstrumentor:
     async def achat_wrapper(
         self,
         category: str,
-        process_chunk: Optional[Callable[[Any, IngestUnitsParams], None]],
-        process_request: Optional[Callable[[IngestUnitsParams, Any, Any], None]],
-        process_synchronous_response: Any,
+        provider: _ProviderRequest,
         is_streaming: _IsStreaming,
         wrapped: Any,
         instance: Any,
@@ -496,7 +514,7 @@ class _PayiInstrumentor:
 
         # after _udpate_headers, all metadata to add to ingest is in extra_headers, keyed by the xproxy-xxx header name
         extra_headers = kwargs.get("extra_headers", {})
-        self._update_headers(context, extra_headers)
+        self._update_extra_headers(context, extra_headers)
 
         if context.get("proxy", True):
             if "extra_headers" not in kwargs:
@@ -504,8 +522,8 @@ class _PayiInstrumentor:
 
             return await wrapped(*args, **kwargs)
 
-        ingest: IngestUnitsParams = {"category": category, "units": {}} # type: ignore
-        ingest["resource"] = kwargs.get("model", "")
+        provider._ingest = {"category": category, "units": {}} # type: ignore
+        provider._ingest["resource"] = kwargs.get("model", "")
 
         if category == PayiCategories.openai and instance and hasattr(instance, "_client"):
             from .OpenAIInstrumentor import OpenAiInstrumentor # noqa: I001
@@ -523,21 +541,20 @@ class _PayiInstrumentor:
                         logging.error("Azure OpenAI invalid resource scope, not ingesting")
                         return wrapped(*args, **kwargs)
 
-                    ingest["resource_scope"] = resource_scope
+                    provider._ingest["resource_scope"] = resource_scope
 
                 category = PayiCategories.azure_openai
 
-                ingest["category"] = category
-                ingest["resource"] = route_as_resource
+                provider._ingest["category"] = category
+                provider._ingest["resource"] = route_as_resource
 
         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
 
-        ingest['properties'] = { 'system.stack_trace': json.dumps(stack) }
+        provider._ingest['properties'] = { 'system.stack_trace': json.dumps(stack) }
 
-        if process_request:
-            process_request(ingest, (), instance)
+        provider.process_request(kwargs)
 
         sw = Stopwatch()
         stream: bool = False
@@ -550,7 +567,7 @@ class _PayiInstrumentor:
             stream = False
 
         try:
-            self._prepare_ingest(ingest, extra_headers, **kwargs)
+            self._prepare_ingest(provider._ingest, extra_headers, **kwargs)
             sw.start()
             response = await wrapped(*args, **kwargs)
 
@@ -568,9 +585,8 @@ class _PayiInstrumentor:
                 instance=instance,
                 instrumentor=self,
                 log_prompt_and_response=self._log_prompt_and_response,
-                ingest=ingest,
                 stopwatch=sw,
-                process_chunk=process_chunk,
+                provider=provider,
                 is_bedrock=False,
             )
 
@@ -578,28 +594,25 @@ class _PayiInstrumentor:
 
         sw.stop()
         duration = sw.elapsed_ms_int()
-        ingest["end_to_end_latency_ms"] = duration
-        ingest["http_status_code"] = 200
+        provider._ingest["end_to_end_latency_ms"] = duration
+        provider._ingest["http_status_code"] = 200
 
-        if process_synchronous_response:
-            return_result: Any = process_synchronous_response(
-                response=response,
-                ingest=ingest,
-                log_prompt_and_response=self._log_prompt_and_response,
-                instrumentor=self)
-            if return_result:
-                return return_result
+        return_result: Any = provider.process_synchronous_response(
+            response=response,
+            log_prompt_and_response=self._log_prompt_and_response,
+            kwargs=kwargs)
+
+        if return_result:
+            return return_result
 
-        await self._aingest_units(ingest)
+        await self._aingest_units(provider._ingest)
 
         return response
 
     def chat_wrapper(
         self,
         category: str,
-        process_chunk: Optional[Callable[[Any, IngestUnitsParams], None]],
-        process_request: Optional[Callable[[IngestUnitsParams, Any, Any], None]],
-        process_synchronous_response: Any,
+        provider: _ProviderRequest,
         is_streaming: _IsStreaming,
         wrapped: Any,
         instance: Any,
@@ -620,7 +633,7 @@ class _PayiInstrumentor:
 
         # after _udpate_headers, all metadata to add to ingest is in extra_headers, keyed by the xproxy-xxx header name
         extra_headers = kwargs.get("extra_headers", {})
-        self._update_headers(context, extra_headers)
+        self._update_extra_headers(context, extra_headers)
 
         if context.get("proxy", True):
             if "extra_headers" not in kwargs:
@@ -628,13 +641,13 @@ class _PayiInstrumentor:
 
             return wrapped(*args, **kwargs)
         
-        ingest: IngestUnitsParams = {"category": category, "units": {}} # type: ignore
+        provider._ingest = {"category": category, "units": {}} # type: ignore
         if is_bedrock:
             # boto3 doesn't allow extra_headers
             kwargs.pop("extra_headers", None)
-            ingest["resource"] = kwargs.get("modelId", "")
+            provider._ingest["resource"] = kwargs.get("modelId", "")
         else:
-            ingest["resource"] = kwargs.get("model", "")
+            provider._ingest["resource"] = kwargs.get("model", "")
 
         if category == PayiCategories.openai and instance and hasattr(instance, "_client"):
             from .OpenAIInstrumentor import OpenAiInstrumentor # noqa: I001
@@ -652,21 +665,20 @@ class _PayiInstrumentor:
                         logging.error("Azure OpenAI invalid resource scope, not ingesting")
                         return wrapped(*args, **kwargs)
 
-                    ingest["resource_scope"] = resource_scope
+                    provider._ingest["resource_scope"] = resource_scope
 
                 category = PayiCategories.azure_openai
 
-                ingest["category"] = category
-                ingest["resource"] = route_as_resource
+                provider._ingest["category"] = category
+                provider._ingest["resource"] = route_as_resource
 
         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
 
-        ingest['properties'] = { 'system.stack_trace': json.dumps(stack) }
+        provider._ingest['properties'] = { 'system.stack_trace': json.dumps(stack) }
 
-        if process_request:
-            process_request(ingest, (), kwargs)
+        provider.process_request(kwargs)
 
         sw = Stopwatch()
         stream: bool = False
@@ -679,7 +691,7 @@ class _PayiInstrumentor:
             stream = False
 
         try:
-            self._prepare_ingest(ingest, extra_headers, **kwargs)
+            self._prepare_ingest(provider._ingest, extra_headers, **kwargs)
             sw.start()
             response = wrapped(*args, **kwargs)
 
@@ -697,9 +709,8 @@ class _PayiInstrumentor:
                 instance=instance,
                 instrumentor=self,
                 log_prompt_and_response=self._log_prompt_and_response,
-                ingest=ingest,
                 stopwatch=sw,
-                process_chunk=process_chunk,
+                provider=provider,
                 is_bedrock=is_bedrock,
             )
 
@@ -714,114 +725,97 @@ class _PayiInstrumentor:
 
         sw.stop()
         duration = sw.elapsed_ms_int()
-        ingest["end_to_end_latency_ms"] = duration
-        ingest["http_status_code"] = 200
+        provider._ingest["end_to_end_latency_ms"] = duration
+        provider._ingest["http_status_code"] = 200
 
-        if process_synchronous_response:
-            return_result: Any = process_synchronous_response(
-                response=response,
-                ingest=ingest,
-                log_prompt_and_response=self._log_prompt_and_response,
-                instrumentor=self)
-            if return_result:
-                return return_result
+        return_result: Any = provider.process_synchronous_response(
+            response=response,
+            log_prompt_and_response=self._log_prompt_and_response,
+            kwargs=kwargs)
+        if return_result:
+            return return_result
 
-        self._ingest_units(ingest)
+        self._ingest_units(provider._ingest)
 
         return response
 
     @staticmethod
-    def _update_headers(
+    def _update_extra_headers(
         context: _Context,
         extra_headers: "dict[str, str]",
     ) -> None:
-        limit_ids: Optional[list[str]] = context.get("limit_ids")
-        request_tags: Optional[list[str]] = context.get("request_tags")
-        experience_name: Optional[str] = context.get("experience_name")
-        experience_id: Optional[str] = context.get("experience_id")
-        use_case_name: Optional[str] = context.get("use_case_name")
-        use_case_id: Optional[str] = context.get("use_case_id")
-        use_case_version: Optional[int] = context.get("use_case_version")
-        user_id: Optional[str] = context.get("user_id")
-
-        # Merge limits from the decorator and extra headers
-        if limit_ids is not None:
-            existing_limit_ids = extra_headers.get(PayiHeaderNames.limit_ids, None)
-            
-            if not existing_limit_ids:
-                extra_headers[PayiHeaderNames.limit_ids] = ",".join(limit_ids)
+        context_limit_ids: Optional[list[str]] = context.get("limit_ids")
+        context_experience_name: Optional[str] = context.get("experience_name")
+        context_experience_id: Optional[str] = context.get("experience_id")
+        context_use_case_name: Optional[str] = context.get("use_case_name")
+        context_use_case_id: Optional[str] = context.get("use_case_id")
+        context_use_case_version: Optional[int] = context.get("use_case_version")
+        context_user_id: Optional[str] = context.get("user_id")
+
+        # headers_limit_ids = extra_headers.get(PayiHeaderNames.limit_ids, None)
+    
+        # 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:
+                # headers_limit_ids is empty, remove it from extra_headers
+                extra_headers.pop(PayiHeaderNames.limit_ids, None)
+            else:   
+                # leave the value in extra_headers
+                ...
+        elif context_limit_ids:
+            extra_headers[PayiHeaderNames.limit_ids] = ",".join(context_limit_ids)
+
+        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:
+                # headers_user_id is empty, remove it from extra_headers
+                extra_headers.pop(PayiHeaderNames.user_id, None)
             else:
-                existing_ids = existing_limit_ids.split(',')
-                combined_ids = list(set(existing_ids + limit_ids))
-                extra_headers[PayiHeaderNames.limit_ids] = ",".join(combined_ids)
-
-        # Merge request from the decorator and extra headers
-        if request_tags is not None:
-            existing_request_tags = extra_headers.get(PayiHeaderNames.request_tags, None)
-
-            if not existing_request_tags:
-                extra_headers[PayiHeaderNames.request_tags] = ",".join(request_tags)
+                # leave the value in extra_headers
+                ...
+        elif context_user_id:
+            extra_headers[PayiHeaderNames.user_id] = context_user_id
+
+        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:
+                # 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)
+                extra_headers.pop(PayiHeaderNames.use_case_version, None)
             else:
-                existing_tags = existing_request_tags.split(',')
-                combined_tags = list(set(existing_tags + request_tags))
-                extra_headers[PayiHeaderNames.request_tags] = ",".join(combined_tags)
-
-        # inner extra_headers user_id takes precedence over outer decorator user_id
-        if user_id is not None and extra_headers.get(PayiHeaderNames.user_id, None) is None:
-            extra_headers[PayiHeaderNames.user_id] = user_id   
-
-        # inner extra_headers experience_name and experience_id take precedence over outer decorator experience_name and experience_id
-        # if either inner value is specified, ignore outer decorator values
-        if PayiHeaderNames.experience_name not in extra_headers and PayiHeaderNames.experience_id not in extra_headers:
-
-            # use both decorator values
-            if experience_name is not None:
-                extra_headers[PayiHeaderNames.experience_name] = experience_name
-            if experience_id is not None:
-                extra_headers[PayiHeaderNames.experience_id] = experience_id
-        
-        elif PayiHeaderNames.experience_id in extra_headers and PayiHeaderNames.experience_name not in extra_headers:
-            # use the decorator experience name and the inner experience id
-            if experience_name is not None:
-                extra_headers[PayiHeaderNames.experience_name] = experience_name
-        
-        else:
-            # use the inner experience name and id as-is
-            ...
-
-        # inner extra_headers use_casee_name and use_case_id take precedence over outer decorator use_case_name and use_case_id
-        # if either inner value is specified, ignore outer decorator values
-        if PayiHeaderNames.use_case_name not in extra_headers and PayiHeaderNames.use_case_id not in extra_headers:
-
-            # use decorator values
-            if use_case_name is not None:
-                extra_headers[PayiHeaderNames.use_case_name] = use_case_name
-            if use_case_id is not None:
-                extra_headers[PayiHeaderNames.use_case_id] = use_case_id
-            if use_case_version is not None:
-                extra_headers[PayiHeaderNames.use_case_version] = str(use_case_version)
-        
-        elif PayiHeaderNames.use_case_id in extra_headers and PayiHeaderNames.use_case_name not in extra_headers:
-            # use the decorator experience name and the inner experience id
-            if use_case_name is not None:
-                extra_headers[PayiHeaderNames.use_case_name] = use_case_name
-            
-            # use the decorator experience version and the inner experience id
-            if use_case_version is not None:
-                extra_headers[PayiHeaderNames.use_case_version] = str(use_case_version) # TODO use case 
-        
-        else:
-            # use the inner experience name and id as-is
-            ...
+                # leave the value in extra_headers
+                ...
+        elif context_use_case_name:
+            extra_headers[PayiHeaderNames.use_case_name] = context_use_case_name
+            if context_use_case_id is not None:
+                extra_headers[PayiHeaderNames.use_case_id] = context_use_case_id
+            if context_use_case_version is not None:
+                extra_headers[PayiHeaderNames.use_case_version] = str(context_use_case_version)
+
+        if PayiHeaderNames.experience_name in extra_headers:
+            headers_experience_name = extra_headers.get(PayiHeaderNames.experience_name, None)
+            if headers_experience_name is None or len(headers_experience_name) == 0:
+                # headers_experience_name is empty, remove all experience related headers
+                extra_headers.pop(PayiHeaderNames.experience_name, None)
+                extra_headers.pop(PayiHeaderNames.experience_id, None)
+            else:
+                # leave the value in extra_headers
+                ...
+        elif context_experience_name is not None:
+            extra_headers[PayiHeaderNames.experience_name] = context_experience_name
+            if context_experience_id is not None:
+                extra_headers[PayiHeaderNames.experience_id] = context_experience_id
 
     @staticmethod
-    def update_for_vision(input: int, units: 'dict[str, Units]') -> int:
-        if _PayiInstrumentor.estimated_prompt_tokens in units:
-            prompt_token_estimate: int = units.pop(_PayiInstrumentor.estimated_prompt_tokens)["input"] # type: ignore
-            vision = input - prompt_token_estimate
+    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 = prompt_token_estimate
+                input = estimated_prompt_tokens
         
         return input
 
@@ -863,16 +857,15 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
         response: Any,
         instance: Any,
         instrumentor: _PayiInstrumentor,
-        ingest: IngestUnitsParams,
         stopwatch: Stopwatch,
-        process_chunk: Optional[Callable[[Any, IngestUnitsParams], None]] = None,
+        provider: _ProviderRequest,
         log_prompt_and_response: bool = True,
         is_bedrock: bool = False,
     ) -> None:
 
         bedrock_from_stream: bool = False
         if is_bedrock:
-            ingest["provider_response_id"] = response["ResponseMetadata"]["RequestId"]
+            provider._ingest["provider_response_id"] = response["ResponseMetadata"]["RequestId"]
             stream = response.get("stream", None)
 
             if stream:
@@ -889,11 +882,10 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
 
         self._instrumentor = instrumentor
         self._stopwatch: Stopwatch = stopwatch
-        self._ingest: IngestUnitsParams = ingest
         self._log_prompt_and_response: bool = log_prompt_and_response
         self._responses: list[str] = []
 
-        self._process_chunk: Optional[Callable[[Any, IngestUnitsParams], None]] = process_chunk
+        self._provider: _ProviderRequest = provider
 
         self._first_token: bool = True
         self._is_bedrock: bool = is_bedrock
@@ -913,7 +905,7 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
 
     def __iter__(self) -> Any:  
         if self._is_bedrock:
-            # MUST be reside in a separate function so that the yield statement doesn't implicitly return its own iterator and overriding self
+            # 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
             return self._iter_bedrock()
         return self
 
@@ -942,7 +934,9 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
                 self._stop_iteration()
             raise e
         else:
-            self._evaluate_chunk(chunk)
+            if self._evaluate_chunk(chunk) == False:
+                return self.__next__()
+            
             return chunk
 
     async def __anext__(self) -> Any:
@@ -953,35 +947,35 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
                 await self._astop_iteration()
             raise e
         else:
-            self._evaluate_chunk(chunk)
+            if self._evaluate_chunk(chunk) == False:
+                return await self.__anext__()
             return chunk
 
-    def _evaluate_chunk(self, chunk: Any) -> None:
+    def _evaluate_chunk(self, chunk: Any) -> bool:
         if self._first_token:
-            self._ingest["time_to_first_token_ms"] = self._stopwatch.elapsed_ms_int()
+            self._provider._ingest["time_to_first_token_ms"] = self._stopwatch.elapsed_ms_int()
             self._first_token = False
 
         if self._log_prompt_and_response:
             self._responses.append(self.chunk_to_json(chunk))
 
-        if self._process_chunk:
-            self._process_chunk(chunk, self._ingest)
+        return self._provider.process_chunk(chunk)
 
     def _process_stop_iteration(self) -> None:
         self._stopwatch.stop()
-        self._ingest["end_to_end_latency_ms"] = self._stopwatch.elapsed_ms_int()
-        self._ingest["http_status_code"] = 200
+        self._provider._ingest["end_to_end_latency_ms"] = self._stopwatch.elapsed_ms_int()
+        self._provider._ingest["http_status_code"] = 200
 
         if self._log_prompt_and_response:
-            self._ingest["provider_response_json"] = self._responses
+            self._provider._ingest["provider_response_json"] = self._responses
 
     async def _astop_iteration(self) -> None:
         self._process_stop_iteration()
-        await self._instrumentor._aingest_units(self._ingest)
+        await self._instrumentor._aingest_units(self._provider._ingest)
 
     def _stop_iteration(self) -> None:
         self._process_stop_iteration()
-        self._instrumentor._ingest_units(self._ingest)
+        self._instrumentor._ingest_units(self._provider._ingest)
 
     @staticmethod
     def chunk_to_json(chunk: Any) -> str:
@@ -1023,18 +1017,6 @@ def payi_instrument(
             elif isinstance(p, AsyncPayi): # type: ignore
                 apayi_param = p
     
-    global_context_ingest: Optional[bool] = None
-
-    if config is not None:
-        if "proxy" not in config:
-            config["proxy"] = False
-        global_context_ingest = config.get("proxy") == False
-
-    # Use default clients if not provided for global ingest instrumentation
-    if not payi_param and not apayi_param and global_context_ingest:
-        payi_param = Payi()
-        apayi_param = AsyncPayi()
-    
     # allow for both payi and apayi to be None for the @proxy case
     _instrumentor = _PayiInstrumentor(
         payi=payi_param,
@@ -1047,7 +1029,6 @@ def payi_instrument(
 
 def ingest(
     limit_ids: Optional["list[str]"] = None,
-    request_tags: Optional["list[str]"] = None,
     experience_name: Optional[str] = None,
     experience_id: Optional[str] = None,
     use_case_name: Optional[str] = None,
@@ -1066,7 +1047,6 @@ def ingest(
                     func,
                     False,
                     limit_ids,
-                    request_tags,
                     experience_name,
                     experience_id,
                     use_case_name,
@@ -1085,7 +1065,6 @@ def ingest(
                     func,
                     False,
                     limit_ids,
-                    request_tags,
                     experience_name,
                     experience_id,
                     use_case_name,
@@ -1100,7 +1079,6 @@ def ingest(
 
 def proxy(
     limit_ids: Optional["list[str]"] = None,
-    request_tags: Optional["list[str]"] = None,
     experience_name: Optional[str] = None,
     experience_id: Optional[str] = None,
     use_case_id: Optional[str] = None,
@@ -1118,7 +1096,6 @@ def proxy(
                     func,
                     True,
                     limit_ids,
-                    request_tags,
                     experience_name,
                     experience_id,
                     use_case_name,
@@ -1138,7 +1115,6 @@ def proxy(
                     func,
                     True,
                     limit_ids,
-                    request_tags,
                     experience_name,
                     experience_id,
                     use_case_name,