payi 0.1.0a59__py3-none-any.whl → 0.1.0a61__py3-none-any.whl

payi/lib/instrument.py

@@ -21,20 +21,43 @@ from .helpers import PayiCategories
 from .Stopwatch import Stopwatch
 
 
-class Context(TypedDict, total=False):
+class PayiInstrumentConfig(TypedDict, total=False):
     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]
+    use_case_id: Optional[str]
+    use_case_version: Optional[int]
+    user_id: Optional[str]
+
+class _Context(TypedDict, total=False):
+    proxy: bool
+    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]']
     user_id: Optional[str]
 
-class IsStreaming(Enum):
+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:
+class _PayiInstrumentor:
     estimated_prompt_tokens: str = "estimated_prompt_tokens"
 
     def __init__(
@@ -49,7 +72,12 @@ class PayiInstrumentor:
     ):
         self._payi: Optional[Payi] = payi
         self._apayi: Optional[AsyncPayi] = apayi
-        self._context_stack: list[Context] = []  # Stack of context dictionaries
+
+        if not self._payi and not self._apayi:
+            self._payi = Payi()
+            self._apayi = AsyncPayi()
+
+        self._context_stack: list[_Context] = []  # Stack of context dictionaries
         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
 
@@ -221,40 +249,51 @@ class PayiInstrumentor:
 
     def _setup_call_func(
         self
-        ) -> 'tuple[Context, Optional[str], Optional[str]]':
-        context: Context = {}
+        ) -> 'tuple[_Context, _ParentState]':
+        context: _Context = {}
+        parentState: _ParentState = {}
 
         if len(self._context_stack) > 0:
             # copy current context into the upcoming context
             context = self._context_stack[-1].copy()
             context.pop("proxy")
-            previous_experience_name = context.get("experience_name", None)
-            previous_experience_id = context.get("experience_id", None)
-        else:
-            previous_experience_name = None
-            previous_experience_id = None
-        return (context, previous_experience_name, previous_experience_id)
+            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)
+
+        return (context, parentState)
 
     def _init_context(
         self,
-        context: Context,
-        previous_experience_name: Optional[str],
-        previous_experience_id: Optional[str],
+        context: _Context,
+        parentState: _ParentState,
         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],
+        use_case_id: Optional[str],
+        use_case_version: Optional[int],                
         user_id: Optional[str],
         ) -> None:
         context["proxy"] = proxy
 
+        # TODO use case what if caller specified epxerience / use_case ID and no name?
+
         # Handle experience name and ID logic
         if not experience_name:
             # If no experience_name specified, use previous values
-            context["experience_name"] = previous_experience_name
-            context["experience_id"] = previous_experience_id
+            context["experience_name"] = parentState.get("experience_name", None)
+            context["experience_id"] = parentState.get("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:
                 # Same experience name, use previous ID unless new one specified
@@ -265,6 +304,29 @@ class PayiInstrumentor:
                 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
+            # 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)
+        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:
+                # 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
+            else:
+                # Different experience 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
@@ -283,22 +345,27 @@ class PayiInstrumentor:
         request_tags: Optional["list[str]"],
         experience_name: Optional[str],
         experience_id: Optional[str],
+        use_case_name: Optional[str],
+        use_case_id: Optional[str],
+        use_case_version: Optional[int],        
         user_id: Optional[str],
         *args: Any,
         **kwargs: Any,
     ) -> Any:
-        context, previous_experience_name, previous_experience_id = self._setup_call_func()
+        context, parentState = self._setup_call_func()
 
         with self:
             self._init_context(
                 context,
-                previous_experience_name,
-                previous_experience_id,
+                parentState,
                 proxy, 
                 limit_ids, 
                 request_tags,
                 experience_name,
                 experience_id,
+                use_case_name,
+                use_case_id,
+                use_case_version,
                 user_id)
             return await func(*args, **kwargs)
 
@@ -310,22 +377,27 @@ class PayiInstrumentor:
         request_tags: Optional["list[str]"],
         experience_name: Optional[str],
         experience_id: Optional[str],
+        use_case_name: Optional[str],
+        use_case_id: Optional[str],
+        use_case_version: Optional[int],        
         user_id: Optional[str],
         *args: Any,
         **kwargs: Any,
     ) -> Any:
-        context, previous_experience_name, previous_experience_id = self._setup_call_func()
+        context, parentState = self._setup_call_func()
 
         with self:
             self._init_context(
                 context,
-                previous_experience_name,
-                previous_experience_id,
+                parentState,
                 proxy, 
                 limit_ids, 
                 request_tags,
                 experience_name,
                 experience_id,
+                use_case_name,
+                use_case_id,
+                use_case_version,
                 user_id)
             return func(*args, **kwargs)
 
@@ -339,16 +411,15 @@ class PayiInstrumentor:
         if self._context_stack:
             self._context_stack.pop()
 
-    def set_context(self, context: Context) -> None:
+    def set_context(self, context: _Context) -> None:
         # Update the current top of the stack with the provided context
         if self._context_stack:
             self._context_stack[-1].update(context)
 
-    def get_context(self) -> Optional[Context]:
+    def get_context(self) -> Optional[_Context]:
         # Return the current top of the stack
         return self._context_stack[-1] if self._context_stack else None
 
-
     def _prepare_ingest(
         self,
         ingest: IngestUnitsParams,
@@ -359,6 +430,9 @@ class PayiInstrumentor:
         request_tags = ingest_extra_headers.pop(PayiHeaderNames.request_tags, None)
         experience_name = ingest_extra_headers.pop(PayiHeaderNames.experience_name, None)
         experience_id = ingest_extra_headers.pop(PayiHeaderNames.experience_id, 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)
+        use_case_version = ingest_extra_headers.pop(PayiHeaderNames.use_case_version, None)
         user_id = ingest_extra_headers.pop(PayiHeaderNames.user_id, None)
 
         if limit_ids:
@@ -369,6 +443,12 @@ class PayiInstrumentor:
             ingest["experience_name"] = experience_name
         if experience_id:
             ingest["experience_id"] = experience_id
+        if use_case_name:
+            ingest["use_case_name"] = use_case_name
+        if use_case_id:
+            ingest["use_case_id"] = use_case_id
+        if use_case_version:
+            ingest["use_case_version"] = int(use_case_version)
         if user_id:
             ingest["user_id"] = user_id
 
@@ -397,7 +477,7 @@ class PayiInstrumentor:
         process_chunk: Optional[Callable[[Any, IngestUnitsParams], None]],
         process_request: Optional[Callable[[IngestUnitsParams, Any, Any], None]],
         process_synchronous_response: Any,
-        is_streaming: IsStreaming,
+        is_streaming: _IsStreaming,
         wrapped: Any,
         instance: Any,
         args: Any,
@@ -459,9 +539,9 @@ class PayiInstrumentor:
         sw = Stopwatch()
         stream: bool = False
         
-        if is_streaming == IsStreaming.kwargs:
+        if is_streaming == _IsStreaming.kwargs:
             stream = kwargs.get("stream", False)
-        elif is_streaming == IsStreaming.true:
+        elif is_streaming == _IsStreaming.true:
             stream = True
         else:
             stream = False
@@ -517,7 +597,7 @@ class PayiInstrumentor:
         process_chunk: Optional[Callable[[Any, IngestUnitsParams], None]],
         process_request: Optional[Callable[[IngestUnitsParams, Any, Any], None]],
         process_synchronous_response: Any,
-        is_streaming: IsStreaming,
+        is_streaming: _IsStreaming,
         wrapped: Any,
         instance: Any,
         args: Any,
@@ -588,9 +668,9 @@ class PayiInstrumentor:
         sw = Stopwatch()
         stream: bool = False
         
-        if is_streaming == IsStreaming.kwargs:
+        if is_streaming == _IsStreaming.kwargs:
             stream = kwargs.get("stream", False)
-        elif is_streaming == IsStreaming.true:
+        elif is_streaming == _IsStreaming.true:
             stream = True
         else:
             stream = False
@@ -649,13 +729,16 @@ class PayiInstrumentor:
 
     @staticmethod
     def _update_headers(
-        context: Context,
+        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
@@ -703,10 +786,35 @@ class PayiInstrumentor:
             # 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
+            ...
+
     @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
+        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
             if (vision > 0):
                 units["vision"] = Units(input=vision, output=0)
@@ -751,7 +859,7 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
         self,
         response: Any,
         instance: Any,
-        instrumentor: PayiInstrumentor,
+        instrumentor: _PayiInstrumentor,
         ingest: IngestUnitsParams,
         stopwatch: Stopwatch,
         process_chunk: Optional[Callable[[Any, IngestUnitsParams], None]] = None,
@@ -885,13 +993,14 @@ class ChatStreamWrapper(ObjectProxy):  # type: ignore
             return json.dumps(chunk)
 
 global _instrumentor
-_instrumentor: Optional[PayiInstrumentor] = None
+_instrumentor: Optional[_PayiInstrumentor] = None
 
 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,
 ) -> None:
     global _instrumentor
     if _instrumentor:
@@ -910,10 +1019,21 @@ def payi_instrument(
                 payi_param = p
             elif isinstance(p, AsyncPayi): # type: ignore
                 apayi_param = p
+    
+    global_context: _Context = {}
 
-    # allow for both payi and apayi to be None for the @proxy case
+    if config and len(config) > 0:
+        if "proxy" not in config:
+            config["proxy"] = False
+        global_context = {**config}
+
+    # create the clients on their behalf if not provided for global ingest instrumentation
+    if not payi_param and not apayi_param and len(global_context) > 0 and global_context.get("proxy", False):
+        payi_param = Payi()
+        apayi_param = AsyncPayi()
     
-    _instrumentor = PayiInstrumentor(
+    # allow for both payi and apayi to be None for the @proxy case
+    _instrumentor = _PayiInstrumentor(
         payi=payi_param,
         apayi=apayi_param,
         instruments=instruments,
@@ -921,11 +1041,17 @@ def payi_instrument(
         prompt_and_response_logger=prompt_and_response_logger,
     )
 
+    if len(global_context) > 0:
+        _instrumentor._context_stack.append(global_context)
+
 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,
+    use_case_id: Optional[str] = None,
+    use_case_version: Optional[int] = None,
     user_id: Optional[str] = None,
 ) -> Any:
     def _ingest(func: Any) -> Any:
@@ -942,6 +1068,9 @@ def ingest(
                     request_tags,
                     experience_name,
                     experience_id,
+                    use_case_name,
+                    use_case_id,
+                    use_case_version,
                     user_id,
                     *args,
                     **kwargs,
@@ -958,6 +1087,9 @@ def ingest(
                     request_tags,
                     experience_name,
                     experience_id,
+                    use_case_name,
+                    use_case_id,
+                    use_case_version,
                     user_id,
                     *args,
                     **kwargs,
@@ -970,6 +1102,9 @@ def proxy(
     request_tags: Optional["list[str]"] = None,
     experience_name: Optional[str] = None,
     experience_id: Optional[str] = None,
+    use_case_id: Optional[str] = None,
+    use_case_name: Optional[str] = None,
+    use_case_version: Optional[int] = None,
     user_id: Optional[str] = None,
 ) -> Any:
     def _proxy(func: Any) -> Any:
@@ -985,6 +1120,9 @@ def proxy(
                     request_tags,
                     experience_name,
                     experience_id,
+                    use_case_name,
+                    use_case_id,
+                    use_case_version,
                     user_id,
                     *args,
                     **kwargs
@@ -1002,6 +1140,9 @@ def proxy(
                     request_tags,
                     experience_name,
                     experience_id,
+                    use_case_name,
+                    use_case_id,
+                    use_case_version,
                     user_id,
                     *args,
                     **kwargs

payi/pagination.py

@@ -0,0 +1,50 @@
+# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
+
+from typing import List, Generic, TypeVar, Optional
+from typing_extensions import override
+
+from ._base_client import BasePage, PageInfo, BaseSyncPage, BaseAsyncPage
+
+__all__ = ["SyncCursorPage", "AsyncCursorPage"]
+
+_T = TypeVar("_T")
+
+
+class SyncCursorPage(BaseSyncPage[_T], BasePage[_T], Generic[_T]):
+    items: List[_T]
+    cursor: Optional[str] = None
+
+    @override
+    def _get_page_items(self) -> List[_T]:
+        items = self.items
+        if not items:
+            return []
+        return items
+
+    @override
+    def next_page_info(self) -> Optional[PageInfo]:
+        cursor = self.cursor
+        if not cursor:
+            return None
+
+        return PageInfo(params={"cursor": cursor})
+
+
+class AsyncCursorPage(BaseAsyncPage[_T], BasePage[_T], Generic[_T]):
+    items: List[_T]
+    cursor: Optional[str] = None
+
+    @override
+    def _get_page_items(self) -> List[_T]:
+        items = self.items
+        if not items:
+            return []
+        return items
+
+    @override
+    def next_page_info(self) -> Optional[PageInfo]:
+        cursor = self.cursor
+        if not cursor:
+            return None
+
+        return PageInfo(params={"cursor": cursor})