ingestion: implement strategy architecture and PydanticAI strategy

Author: frankie567

src/polar_sdk/ingestion/_base.py

@@ -1,5 +1,6 @@
 import atexit
 import contextlib
+import datetime
 import logging
 import queue
 import threading
@@ -12,6 +13,8 @@
 if TYPE_CHECKING:
     from polar_sdk import EventsModelTypedDict
 
+    from .strategies._base import S
+
 logger = logging.getLogger("polar_sdk.ingestion")
 
 
@@ -20,7 +23,7 @@ class Ingestion:
     Event ingestion client for Polar.
 
     This class handles the ingestion of events into the Polar API without blocking
-    the main thread, by using background thread sending them in batches.
+    the main thread, by using a background thread sending them in batches.
 
     :param access_token: The access_token required for authentication
     :param server: The server by name to use for all methods
@@ -58,12 +61,23 @@ def __init__(
 
         atexit.register(self.close)
 
+    def strategy(self, strategy_class: type["S"], event_name: str) -> "S":
+        """
+        Instantiate a strategy tied to this ingestion client.
+
+        :param strategy_class: The strategy class to instantiate.
+        :param event_name: The name of the events that'll be reported by the strategy.
+        """
+        return strategy_class(event_name, self)
+
     def ingest(self, event: "EventsModelTypedDict") -> None:
         """
         Send an event to the ingestion queue.
 
         :param event: The event to send.
         """
+        if event.get("timestamp") is None:
+            event["timestamp"] = datetime.datetime.now(datetime.timezone.utc)
         self._queue.put(event, block=False)
 
     def flush(self, max_batch_size: Union[int, None] = None) -> None:

src/polar_sdk/ingestion/strategies/__init__.py

@@ -0,0 +1,4 @@
+from ._base import BaseStrategy
+from .pydantic_ai import PydanticAIStrategy
+
+__all__ = ["BaseStrategy", "PydanticAIStrategy"]

src/polar_sdk/ingestion/strategies/_base.py

@@ -0,0 +1,27 @@
+from typing import Any, TypeVar
+
+from .._base import Ingestion
+
+
+class BaseStrategy:
+    """
+    Base structure of an event ingestion strategy.
+
+    :param event_name: The name of the events that'll be reported by the strategy.
+    :param ingestion: The ingestion client to use for sending events.
+    """
+
+    def __init__(self, event_name: str, ingestion: Ingestion) -> None:
+        self.event_name = event_name
+        self._ingestion = ingestion
+
+    def ingest(self, external_customer_id: str, *args: Any, **kwargs: Any) -> None:
+        """
+        Send an event to the ingestion queue.
+
+        :param external_customer_id: The external customer ID to associate with the event.
+        """
+        raise NotImplementedError()
+
+
+S = TypeVar("S", bound=BaseStrategy)

src/polar_sdk/ingestion/strategies/pydantic_ai.py

@@ -0,0 +1,32 @@
+from typing import Any
+
+from pydantic_ai.agent import AgentRunResult
+
+from ._base import BaseStrategy
+
+
+class PydanticAIStrategy(BaseStrategy):
+    """
+    An ingestion strategy to report usage from PydanticAI agent runs.
+    """
+
+    def ingest(self, external_customer_id: str, result: AgentRunResult[Any]) -> None:
+        """
+        Report usage from a PydanticAI agent run.
+
+        :param external_customer_id: The external customer ID to associate with the event.
+        :param result: The result of the agent run.
+        """
+        usage = result.usage()
+        self._ingestion.ingest(
+            {
+                "name": self.event_name,
+                "external_customer_id": external_customer_id,
+                "metadata": {
+                    "request_tokens": usage.request_tokens or 0,
+                    "response_tokens": usage.response_tokens or 0,
+                    "total_tokens": usage.total_tokens or 0,
+                    "requests": usage.requests or 0,
+                },
+            }
+        )