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

payi/lib/OpenAIInstrumentor.py

@@ -1,6 +1,7 @@
 import json
 import logging
-from typing import Any, Union
+from typing import Any, Union, Optional
+from typing_extensions import override
 from importlib.metadata import version
 
 import tiktoken  # type: ignore
@@ -9,7 +10,7 @@ from wrapt import wrap_function_wrapper  # type: ignore
 from payi.types import IngestUnitsParams
 from payi.types.ingest_units_params import Units
 
-from .instrument import _IsStreaming, _PayiInstrumentor
+from .instrument import _IsStreaming, _ProviderRequest, _PayiInstrumentor
 
 
 class OpenAiInstrumentor:
@@ -63,9 +64,7 @@ def embeddings_wrapper(
 ) -> Any:
     return instrumentor.chat_wrapper(
         "system.openai",
-        None, # process_chat_chunk,
-        None, # process_chat_request,
-        process_ebmeddings_synchronous_response,
+        _OpenAiEmbeddingsProviderRequest(instrumentor),
         _IsStreaming.false,
         wrapped,
         instance,
@@ -83,9 +82,7 @@ async def aembeddings_wrapper(
 ) -> Any:
     return await instrumentor.achat_wrapper(
         "system.openai",
-        None, # process_chat_chunk,
-        None, # process_chat_request,
-        process_ebmeddings_synchronous_response,
+        _OpenAiEmbeddingsProviderRequest(instrumentor),
         _IsStreaming.false,
         wrapped,
         instance,
@@ -103,9 +100,7 @@ def chat_wrapper(
 ) -> Any:
     return instrumentor.chat_wrapper(
         "system.openai",
-        process_chat_chunk,
-        process_chat_request,
-        process_chat_synchronous_response,
+        _OpenAiChatProviderRequest(instrumentor),
         _IsStreaming.kwargs,
         wrapped,
         instance,
@@ -123,9 +118,7 @@ async def achat_wrapper(
 ) -> Any:
     return await instrumentor.achat_wrapper(
         "system.openai",
-        process_chat_chunk,
-        process_chat_request,
-        process_chat_synchronous_response,
+        _OpenAiChatProviderRequest(instrumentor),
         _IsStreaming.kwargs,
         wrapped,
         instance,
@@ -133,13 +126,89 @@ async def achat_wrapper(
         kwargs,
     )
 
-def process_ebmeddings_synchronous_response(response: str, ingest: IngestUnitsParams, log_prompt_and_response: bool, **kwargs: Any) -> Any: #  noqa: ARG001
-    return process_chat_synchronous_response(response, ingest, log_prompt_and_response, **kwargs)
+class _OpenAiEmbeddingsProviderRequest(_ProviderRequest):
+    @override
+    def process_synchronous_response(
+        self,
+        response: Any,
+        log_prompt_and_response: bool,
+        kwargs: Any) -> Any:
+        return process_chat_synchronous_response(response, self._ingest, log_prompt_and_response, self._estimated_prompt_tokens)
+
+class _OpenAiChatProviderRequest(_ProviderRequest):
+    def __init__(self, instrumentor: _PayiInstrumentor):
+        super().__init__(instrumentor)
+        self._include_usage_added = False
+
+    @override
+    def process_chunk(self, chunk: Any) -> bool:
+        model = model_to_dict(chunk)
+        
+        if "provider_response_id" not in self._ingest:
+            response_id = model.get("id", None)
+            if response_id:
+                self._ingest["provider_response_id"] = response_id
+
+        send_chunk_to_client = True
+
+        usage = model.get("usage")
+        if usage:
+            add_usage_units(usage, self._ingest["units"], self._estimated_prompt_tokens)
+
+            # If we aded "include_usage" in the request on behalf of the client, do not return the extra 
+            # packet which contains the usage to the client as they are not expecting the data
+            if self._include_usage_added:
+                send_chunk_to_client = False
 
-def process_chat_synchronous_response(response: str, ingest: IngestUnitsParams, log_prompt_and_response: bool, **kwargs: Any) -> Any: #  noqa: ARG001
+        return send_chunk_to_client
+
+    @override
+    def process_request(self, kwargs: Any) -> None: # noqa: ARG001
+        messages = kwargs.get("messages", None)
+        if not messages or len(messages) == 0:
+            return
+        
+        estimated_token_count = 0 
+        has_image = False
+
+        try: 
+            enc = tiktoken.encoding_for_model(kwargs.get("model")) # type: ignore
+        except KeyError:
+            enc = tiktoken.get_encoding("o200k_base") # type: ignore
+        
+        for message in messages:
+            msg_has_image, msg_prompt_tokens = has_image_and_get_texts(enc, message.get('content', ''))
+            if msg_has_image:
+                has_image = True
+                estimated_token_count += msg_prompt_tokens
+        
+        if has_image and estimated_token_count > 0:
+            self._estimated_prompt_tokens = estimated_token_count
+
+        stream: bool = kwargs.get("stream", False)
+        if stream:
+            add_include_usage = True
+
+            stream_options: dict[str, Any] = kwargs.get("stream_options", None)
+            if stream_options and "include_usage" in stream_options:
+                add_include_usage = stream_options["include_usage"] == False
+
+            if add_include_usage:
+                kwargs['stream_options'] = {"include_usage": True}
+                self._include_usage_added = True
+
+    @override
+    def process_synchronous_response(
+        self,
+        response: Any,
+        log_prompt_and_response: bool,
+        kwargs: Any) -> Any:
+        process_chat_synchronous_response(response, self._ingest, log_prompt_and_response, self._estimated_prompt_tokens)
+
+def process_chat_synchronous_response(response: str, ingest: IngestUnitsParams, log_prompt_and_response: bool, estimated_prompt_tokens: Optional[int]) -> Any:
     response_dict = model_to_dict(response)
 
-    add_usage_units(response_dict.get("usage", {}), ingest["units"])
+    add_usage_units(response_dict.get("usage", {}), ingest["units"], estimated_prompt_tokens)
 
     if log_prompt_and_response:
         ingest["provider_response_json"] = [json.dumps(response_dict)]
@@ -149,19 +218,6 @@ def process_chat_synchronous_response(response: str, ingest: IngestUnitsParams,
 
     return None
 
-def process_chat_chunk(chunk: Any, ingest: IngestUnitsParams) -> None:
-    model = model_to_dict(chunk)
-    
-    if "provider_response_id" not in ingest:
-        response_id = model.get("id", None)
-        if response_id:
-            ingest["provider_response_id"] = response_id
-
-    usage = model.get("usage")
-    if usage:
-        add_usage_units(usage, ingest["units"])
-
-
 def model_to_dict(model: Any) -> Any:
     if version("pydantic") < "2.0.0":
         return model.dict()
@@ -173,7 +229,7 @@ def model_to_dict(model: Any) -> Any:
         return model
 
 
-def add_usage_units(usage: "dict[str, Any]", units: "dict[str, Units]") -> None:
+def add_usage_units(usage: "dict[str, Any]", units: "dict[str, Units]", estimated_prompt_tokens: Optional[int]) -> None:
     input = usage["prompt_tokens"] if "prompt_tokens" in usage else 0
     output = usage["completion_tokens"] if "completion_tokens" in usage else 0
     input_cache = 0
@@ -184,7 +240,7 @@ def add_usage_units(usage: "dict[str, Any]", units: "dict[str, Units]") -> None:
         if input_cache != 0:
             units["text_cache_read"] = Units(input=input_cache, output=0)
 
-    input = _PayiInstrumentor.update_for_vision(input - input_cache, units)
+    input = _PayiInstrumentor.update_for_vision(input - input_cache, units, estimated_prompt_tokens)
 
     units["text"] = Units(input=input, output=output)
 
@@ -197,28 +253,4 @@ def has_image_and_get_texts(encoding: tiktoken.Encoding, content: Union[str, 'li
             return has_image, 0
         
         token_count = sum(len(encoding.encode(item.get("text", ""))) for item in content if item.get("type") == "text")
-        return has_image, token_count
-
-def process_chat_request(ingest: IngestUnitsParams, *args: Any, **kwargs: Any) -> None: # noqa: ARG001
-    messages = kwargs.get("messages")
-    if not messages or len(messages) == 0:
-        return
-    
-    estimated_token_count = 0 
-    has_image = False
-
-    try: 
-        enc = tiktoken.encoding_for_model(kwargs.get("model")) # type: ignore
-    except KeyError:
-        enc = tiktoken.get_encoding("o200k_base") # type: ignore
-    
-    for message in messages:
-        msg_has_image, msg_prompt_tokens = has_image_and_get_texts(enc, message.get('content', ''))
-        if msg_has_image:
-            has_image = True
-            estimated_token_count += msg_prompt_tokens
-    
-    if not has_image or estimated_token_count == 0:
-        return
-
-    ingest["units"][_PayiInstrumentor.estimated_prompt_tokens] = Units(input=estimated_token_count, output=0)
+        return has_image, token_count

payi/lib/helpers.py

@@ -49,6 +49,7 @@ def create_headers(
     use_case_id: Union[str, None] = None,
     use_case_name: Union[str, None] = None,
     use_case_version: Union[int, None] = None,
+    route_as_resource: Union[str, None] = None,
 ) -> Dict[str, str]:
     headers: Dict[str, str] = {}
 
@@ -68,6 +69,8 @@ def create_headers(
         headers.update({ PayiHeaderNames.use_case_name: use_case_name})
     if use_case_version:
         headers.update({ PayiHeaderNames.use_case_version: str(use_case_version)})
+    if route_as_resource:
+        headers.update({ PayiHeaderNames.route_as_resource: route_as_resource})
     return headers
 
 def _resolve_payi_base_url(payi_base_url: Union[str, None]) -> str: