Merge pull request #65 from cfahlgren1/fix-set-client-name

Add client-specific record classes for Aisuite and Litellm

Author: davidberenstein1957

src/observers/models/aisuite.py

@@ -14,6 +14,10 @@
     from observers.stores.duckdb import DuckDBStore
 
 
+class AisuiteRecord(OpenAIRecord):
+    client_name: str = "aisuite"
+
+
 def wrap_aisuite(
     client: "Client",
     store: Optional[Union["DatasetsStore", "DuckDBStore", "ArgillaStore"]] = None,
@@ -43,7 +47,7 @@ def wrap_aisuite(
         client=client,
         create=client.chat.completions.create,
         format_input=lambda inputs, **kwargs: {"messages": inputs, **kwargs},
-        parse_response=OpenAIRecord.from_response,
+        parse_response=AisuiteRecord.from_response,
         store=store,
         tags=tags,
         properties=properties,

src/observers/models/litellm.py

@@ -14,6 +14,10 @@
     from observers.stores.duckdb import DuckDBStore
 
 
+class LitellmRecord(OpenAIRecord):
+    client_name: str = "litellm"
+
+
 def wrap_litellm(
     client: Union["completion", "acompletion"],
     store: Optional[Union["DatasetsStore", "DuckDBStore", "ArgillaStore"]] = None,
@@ -44,7 +48,7 @@ def wrap_litellm(
         "client": client,
         "create": client,
         "format_input": lambda inputs, **kwargs: {"messages": inputs, **kwargs},
-        "parse_response": OpenAIRecord.from_response,
+        "parse_response": LitellmRecord.from_response,
         "store": store,
         "tags": tags,
         "properties": properties,