Improve DuckDB Store

README.md

@@ -35,19 +35,18 @@ We differentiate between observers and stores. Observers wrap generative AI APIs
 To get started you can run the code below. It sends requests to a HF serverless endpoint and log the interactions into a Hub dataset, using the default store `DatasetsStore`. The dataset will be pushed to your personal workspace (http://hf.co/{your_username}). To learn how to configure stores, go to the next section.
 
 ```python
-from observers.observers import wrap_openai
-from observers.stores import DuckDBStore
 from openai import OpenAI
-
-store = DuckDBStore()
+from observers import wrap_openai
 
 openai_client = OpenAI()
-client = wrap_openai(openai_client, store=store)
+
+client = wrap_openai(openai_client)
 
 response = client.chat.completions.create(
     model="gpt-4o",
     messages=[{"role": "user", "content": "Tell me a joke."}],
 )
+print(response)
 ```
 
 ## Observers

examples/models/openai_example.py

@@ -1,17 +1,19 @@
-import os
-
 from openai import OpenAI
 
 from observers import wrap_openai
 
-openai_client = OpenAI(
-    base_url="https://api-inference.huggingface.co/v1/", api_key=os.getenv("HF_TOKEN")
-)
+
+openai_client = OpenAI()
 
 client = wrap_openai(openai_client)
 
 response = client.chat.completions.create(
-    model="Qwen/Qwen2.5-72B-Instruct",
-    messages=[{"role": "user", "content": "Tell me a joke."}],
+    model="gpt-4o",
+    messages=[{"role": "user", "content": "Tell me a joke in the voice of a pirate."}],
+    stream=True,
+    temperature=0.5,
 )
-print(response)
+
+for chunk in response:
+    if chunk.choices[0].delta.content:
+        print(chunk.choices[0].delta.content, end="", flush=True)

src/observers/base.py

@@ -2,15 +2,34 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional
+from typing_extensions import Literal
 
 if TYPE_CHECKING:
     from argilla import Argilla
 
 
+@dataclass
+class Function:
+    """Function tool call information"""
+
+    name: str
+    arguments: str
+
+
+@dataclass
+class ToolCall:
+    """Tool call information"""
+
+    id: str
+    type: Literal["function"]
+    function: Function
+
+
 @dataclass
 class Message:
     role: Literal["system", "user", "assistant", "function"]
     content: str
+    tool_calls: Optional[List[ToolCall]] = None
 
 
 @dataclass

src/observers/models/aisuite.py

@@ -46,7 +46,7 @@ def wrap_aisuite(
     return ChatCompletionObserver(
         client=client,
         create=client.chat.completions.create,
-        format_input=lambda inputs, **kwargs: {"messages": inputs, **kwargs},
+        format_input=lambda messages, **kwargs: kwargs | {"messages": messages},
         parse_response=AisuiteRecord.from_response,
         store=store,
         tags=tags,

src/observers/models/base.py

@@ -22,6 +22,7 @@ class ChatCompletionRecord(Record):
 
     model: str = None
     timestamp: str = field(default_factory=lambda: datetime.datetime.now().isoformat())
+    arguments: Optional[Dict[str, Any]] = None
 
     messages: List[Message] = None
     assistant_message: Optional[str] = None
@@ -55,7 +56,7 @@ def table_columns(self):
             "properties",
             "error",
             "raw_response",
-            "synced_at",
+            "arguments",
         ]
 
     @property
@@ -65,7 +66,7 @@ def duckdb_schema(self):
             id VARCHAR PRIMARY KEY,
             model VARCHAR,
             timestamp TIMESTAMP,
-            messages STRUCT(role VARCHAR, content VARCHAR)[],
+            messages JSON,
             assistant_message TEXT,
             completion_tokens INTEGER,
             prompt_tokens INTEGER,
@@ -77,7 +78,7 @@ def duckdb_schema(self):
             properties JSON,
             error VARCHAR,
             raw_response JSON,
-            synced_at TIMESTAMP
+            arguments JSON,
         )
         """
 
@@ -160,7 +161,14 @@ def table_name(self):
 
     @property
     def json_fields(self):
-        return ["tool_calls", "function_call", "tags", "properties", "raw_response"]
+        return [
+            "tool_calls",
+            "function_call",
+            "tags",
+            "properties",
+            "raw_response",
+            "arguments",
+        ]
 
     @property
     def image_fields(self):
@@ -223,13 +231,17 @@ def chat(self) -> Self:
     def completions(self) -> Self:
         return self
 
-    def _log_record(self, response, error=None, model=None):
+    def _log_record(
+        self, response, error=None, model=None, messages=None, arguments=None
+    ):
         record = self.parse_response(
             response,
             error=error,
             model=model,
+            messages=messages,
             tags=self.tags,
             properties=self.properties,
+            arguments=arguments,
         )
         if random.random() < self.logging_rate:
             self.store.add(record)
@@ -255,30 +267,47 @@ def create(
         """
         response = None
         kwargs = self.handle_kwargs(kwargs)
+        excluded_args = {"model", "messages", "tags", "properties"}
+        arguments = {k: v for k, v in kwargs.items() if k not in excluded_args}
         model = kwargs.get("model")
         input_data = self.format_input(messages, **kwargs)
 
         if kwargs.get("stream", False):
 
             def stream_responses():
-                response = []
+                response_buffer = []
                 try:
                     for chunk in self.create_fn(**input_data):
                         yield chunk
-                        response.append(chunk)
-                    self._log_record(response, model=model)
+                        response_buffer.append(chunk)
+                    self._log_record(
+                        response_buffer,
+                        model=model,
+                        messages=messages,
+                        arguments=arguments,
+                    )
                 except Exception as e:
-                    self._log_record(response, error=e, model=model)
+                    self._log_record(
+                        response_buffer,
+                        error=e,
+                        model=model,
+                        messages=messages,
+                        arguments=arguments,
+                    )
                     raise
 
             return stream_responses()
 
         try:
             response = self.create_fn(**input_data)
-            self._log_record(response, model=model)
+            self._log_record(
+                response, model=model, messages=messages, arguments=arguments
+            )
             return response
         except Exception as e:
-            self._log_record(response, error=e, model=model)
+            self._log_record(
+                response, error=e, model=model, messages=messages, arguments=arguments
+            )
             raise
 
     def handle_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
@@ -318,13 +347,17 @@ class AsyncChatCompletionObserver(ChatCompletionObserver):
             The logging rate to use for logging, defaults to 1
     """
 
-    async def _log_record_async(self, response, error=None, model=None):
+    async def _log_record_async(
+        self, response, error=None, model=None, messages=None, arguments=None
+    ):
         record = self.parse_response(
             response,
             error=error,
             model=model,
+            messages=messages,
             tags=self.tags,
             properties=self.properties,
+            arguments=arguments,
         )
         if random.random() < self.logging_rate:
             await self.store.add_async(record)
@@ -346,30 +379,47 @@ async def create(
         """
         response = None
         kwargs = self.handle_kwargs(kwargs)
-        input_data = self.format_input(messages, **kwargs)
+        excluded_args = {"model", "messages", "tags", "properties"}
+        arguments = {k: v for k, v in kwargs.items() if k not in excluded_args}
         model = kwargs.get("model")
+        input_data = self.format_input(messages, **kwargs)
 
         if kwargs.get("stream", False):
 
             async def stream_responses():
-                response = []
+                response_buffer = []
                 try:
                     async for chunk in await self.create_fn(**input_data):
                         yield chunk
-                        response.append(chunk)
-                    await self._log_record_async(response, model=model)
+                        response_buffer.append(chunk)
+                    await self._log_record_async(
+                        response_buffer,
+                        model=model,
+                        messages=messages,
+                        arguments=arguments,
+                    )
                 except Exception as e:
-                    await self._log_record_async(response, error=e, model=model)
+                    await self._log_record_async(
+                        response_buffer,
+                        error=e,
+                        model=model,
+                        messages=messages,
+                        arguments=arguments,
+                    )
                     raise
 
             return stream_responses()
 
         try:
             response = await self.create_fn(**input_data)
-            await self._log_record_async(response, model=model)
+            await self._log_record_async(
+                response, model=model, messages=messages, arguments=arguments
+            )
             return response
         except Exception as e:
-            await self._log_record_async(response, error=e, model=model)
+            await self._log_record_async(
+                response, error=e, model=model, messages=messages, arguments=arguments
+            )
             raise
 
     async def __aenter__(self) -> "AsyncChatCompletionObserver":

src/observers/models/openai.py

@@ -1,6 +1,6 @@
 import uuid
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
-
+from observers.stores.duckdb import DuckDBStore
 from openai import AsyncOpenAI, OpenAI
 from typing_extensions import Self
 
@@ -14,7 +14,6 @@
     from openai.types.chat import ChatCompletion, ChatCompletionChunk
 
     from observers.stores.datasets import DatasetsStore
-    from observers.stores.duckdb import DuckDBStore
 
 
 class OpenAIRecord(ChatCompletionRecord):
@@ -25,11 +24,14 @@ def from_response(
         cls,
         response: Union[List["ChatCompletionChunk"], "ChatCompletion"] = None,
         error=None,
+        messages=None,
         **kwargs,
     ) -> Self:
         """Create a response record from an API response or error"""
         if not response:
-            return cls(finish_reason="error", error=str(error), **kwargs)
+            return cls(
+                finish_reason="error", error=str(error), messages=messages, **kwargs
+            )
 
         # Handle streaming responses
         if isinstance(response, list):
@@ -56,6 +58,7 @@ def from_response(
 
             return cls(
                 id=first_dump.get("id") or str(uuid.uuid4()),
+                messages=messages,
                 completion_tokens=completion_tokens,
                 prompt_tokens=prompt_tokens,
                 total_tokens=total_tokens,
@@ -73,6 +76,7 @@ def from_response(
         usage = response_dump.get("usage", {}) or {}
         return cls(
             id=response.id or str(uuid.uuid4()),
+            messages=messages,
             completion_tokens=usage.get("completion_tokens"),
             prompt_tokens=usage.get("prompt_tokens"),
             total_tokens=usage.get("total_tokens"),
@@ -87,7 +91,7 @@ def from_response(
 
 def wrap_openai(
     client: Union["OpenAI", "AsyncOpenAI"],
-    store: Optional[Union["DuckDBStore", "DatasetsStore"]] = None,
+    store: Optional[Union["DuckDBStore", "DatasetsStore"]] = DuckDBStore(),
     tags: Optional[List[str]] = None,
     properties: Optional[Dict[str, Any]] = None,
     logging_rate: Optional[float] = 1,
@@ -114,7 +118,7 @@ def wrap_openai(
     observer_args = {
         "client": client,
         "create": client.chat.completions.create,
-        "format_input": lambda inputs, **kwargs: {"messages": inputs, **kwargs},
+        "format_input": lambda messages, **kwargs: kwargs | {"messages": messages},
         "parse_response": OpenAIRecord.from_response,
         "store": store,
         "tags": tags,

src/observers/stores/datasets.py

@@ -189,7 +189,6 @@ def add(self, record: "Record"):
         with self._scheduler.lock:
             with (self._scheduler.folder_path / self._filename).open("a") as f:
                 record_dict = asdict(record)
-                record_dict["synced_at"] = None
 
                 # Handle JSON fields
                 for json_field in record.json_fields: