ingestion: tweak strategies

Author: frankie567

src/polar_sdk/ingestion/_base.py

@@ -10,7 +10,7 @@
 from polar_sdk import Polar
 
 if TYPE_CHECKING:
-    from polar_sdk import EventsModelTypedDict
+    from polar_sdk import EventCreateCustomerMetadataTypedDict, EventsModelTypedDict
 
     from .strategies._base import S
 
@@ -59,14 +59,20 @@ def __init__(
 
         atexit.register(self.close)
 
-    def strategy(self, strategy_class: type["S"], event_name: str) -> "S":
+    def strategy(
+        self,
+        strategy_class: type["S"],
+        event_name: str,
+        **metadata: "EventCreateCustomerMetadataTypedDict",
+    ) -> "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.
+        :param metadata: Extra metadata to include with all events.
         """
-        return strategy_class(event_name, self)
+        return strategy_class(event_name, self, metadata=metadata)
 
     def ingest(self, event: "EventsModelTypedDict") -> None:
         """

src/polar_sdk/ingestion/strategies/_base.py

@@ -1,19 +1,29 @@
-from typing import Any, TypeVar
+from typing import TYPE_CHECKING, Any, TypeVar, Union
 
 from .._base import Ingestion
 
+if TYPE_CHECKING:
+    from polar_sdk import EventCreateCustomerMetadataTypedDict
+
 
 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.
+    :param metadata: Extra metadata to include with all events.
     """
 
-    def __init__(self, event_name: str, ingestion: Ingestion) -> None:
+    def __init__(
+        self,
+        event_name: str,
+        ingestion: Ingestion,
+        metadata: Union[dict[str, "EventCreateCustomerMetadataTypedDict"], None] = None,
+    ) -> None:
         self.event_name = event_name
         self._ingestion = ingestion
+        self._metadata = metadata or {}
 
     def ingest(self, external_customer_id: str, *args: Any, **kwargs: Any) -> None:
         """

src/polar_sdk/ingestion/strategies/pydantic_ai.py

@@ -1,21 +1,31 @@
 from typing import Any
 
 from pydantic_ai.agent import AgentRunResult
+from pydantic_ai.models import TYPE_CHECKING
 
 from ._base import BaseStrategy
 
+if TYPE_CHECKING:
+    from polar_sdk import EventCreateCustomerMetadataTypedDict
+
 
 class PydanticAIStrategy(BaseStrategy):
     """
     An ingestion strategy to report usage from PydanticAI agent runs.
     """
 
-    def ingest(self, external_customer_id: str, result: AgentRunResult[Any]) -> None:
+    def ingest(
+        self,
+        external_customer_id: str,
+        result: AgentRunResult[Any],
+        **metadata: "EventCreateCustomerMetadataTypedDict",
+    ) -> 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.
+        :param metadata: Extra metadata to include with the event.
         """
         usage = result.usage()
         self._ingestion.ingest(
@@ -27,6 +37,8 @@ def ingest(self, external_customer_id: str, result: AgentRunResult[Any]) -> None
                     "response_tokens": usage.response_tokens or 0,
                     "total_tokens": usage.total_tokens or 0,
                     "requests": usage.requests or 0,
+                    **self._metadata,
+                    **metadata,
                 },
             }
         )