From f93d53b8ef9c13887ffaab01ba1ea29d702a90bf Mon Sep 17 00:00:00 2001
From: Joschka Braun <joschka.braun@gmail.com>
Date: Sat, 18 May 2024 17:05:56 -0400
Subject: [PATCH] feat: auto trace marvin

---
 parea/client.py                |  8 ++++++++
 parea/cookbook/trace_marvin.py | 21 +++++++++++++++++++++
 pyproject.toml                 |  2 +-
 3 files changed, 30 insertions(+), 1 deletion(-)
 create mode 100644 parea/cookbook/trace_marvin.py

diff --git a/parea/client.py b/parea/client.py
index 972e29c7..df8c0768 100644
--- a/parea/client.py
+++ b/parea/client.py
@@ -120,6 +120,14 @@ def auto_trace_openai_clients(self, integration: Optional[str]) -> None:
         if integration:
             self._client.add_integration(integration)
 
+        if integration == "marvin":
+            import marvin
+
+            marvin.utilities.openai.Client = patch_openai_client_classes(marvin.utilities.openai.Client, self)
+            marvin.utilities.openai.AsyncClient = patch_openai_client_classes(marvin.utilities.openai.AsyncClient, self)
+            marvin.utilities.openai.AzureOpenAI = patch_openai_client_classes(marvin.utilities.openai.AzureOpenAI, self)
+            marvin.utilities.openai.AsyncAzureOpenAI = patch_openai_client_classes(marvin.utilities.openai.AsyncAzureOpenAI, self)
+
     def trace_dspy(self):
         from parea.utils.trace_integrations.dspy import DSPyInstrumentor
 
diff --git a/parea/cookbook/trace_marvin.py b/parea/cookbook/trace_marvin.py
new file mode 100644
index 00000000..ba934da5
--- /dev/null
+++ b/parea/cookbook/trace_marvin.py
@@ -0,0 +1,21 @@
+import os
+
+import marvin
+from dotenv import load_dotenv
+from pydantic import BaseModel, Field
+
+from parea import Parea
+
+load_dotenv()
+
+p = Parea(api_key=os.getenv("PAREA_API_KEY"))
+p.auto_trace_openai_clients("marvin")
+
+
+class Location(BaseModel):
+    city: str
+    state: str = Field(description="2-letter abbreviation")
+
+
+result = marvin.cast("the big apple", Location)
+print(result)
diff --git a/pyproject.toml b/pyproject.toml
index 2e4d9d1c..01143a7f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry]
 name = "parea-ai"
 packages = [{ include = "parea" }]
-version = "0.2.155"
+version = "0.2.156"
 description = "Parea python sdk"
 readme = "README.md"
 authors = ["joel-parea-ai <joel@parea.ai>"]