refactor: move events to audit module (#58)

Author: micpst

docs/how-to/create_custom_event_handler.md

@@ -10,8 +10,7 @@ In this guide we will implement a simple [Event Handler](../reference/event_hand
 First, we need to create a new class that inherits from `EventHandler` and implements the all abstract methods.
 
 ```python
-from dbally.audit import EventHandler
-from dbally.data_models.audit import RequestStart, RequestEnd
+from dbally.audit import EventHandler, RequestStart, RequestEnd
 
 class FileEventHandler(EventHandler):
 

docs/reference/event_handlers/index.md

@@ -10,10 +10,10 @@ db-ally provides an `EventHandler` abstract class that can be used to log the ru
 Each run of [dbally.Collection.ask][dbally.Collection.ask] will trigger all instances of EventHandler that were passed to the Collection's constructor (or the [dbally.create_collection][dbally.create_collection] function).
 
 
-1. `EventHandler.request_start` is called with [RequestStart][dbally.data_models.audit.RequestStart], it can return a context object that will be passed to next calls.
+1. `EventHandler.request_start` is called with [RequestStart][dbally.audit.events.RequestStart], it can return a context object that will be passed to next calls.
 2. For each event that occurs during the run, `EventHandler.event_start` is called with the context object returned by `EventHandler.request_start` and an Event object. It can return context for the `EventHandler.event_end` method.
 3. When the event ends `EventHandler.event_end` is called with the context object returned by `EventHandler.event_start` and an Event object.
-4. On the end of the run `EventHandler.request_end` is called with the context object returned by `EventHandler.request_start` and the [RequestEnd][dbally.data_models.audit.RequestEnd].
+4. On the end of the run `EventHandler.request_end` is called with the context object returned by `EventHandler.request_start` and the [RequestEnd][dbally.audit.events.RequestEnd].
 
 
 ``` mermaid
@@ -42,8 +42,14 @@ Currently handled events:
 
 ::: dbally.audit.EventHandler
 
-::: dbally.data_models.audit.RequestStart
+::: dbally.audit.events.RequestStart
 
-::: dbally.data_models.audit.RequestEnd
+::: dbally.audit.events.RequestEnd
 
-::: dbally.data_models.audit.LLMEvent
+::: dbally.audit.events.Event
+
+::: dbally.audit.events.LLMEvent
+
+::: dbally.audit.events.SimilarityEvent
+
+::: dbally.audit.spans.EventSpan

src/dbally/audit/__init__.py

@@ -7,8 +7,19 @@
 except ImportError:
     pass
 
+from .event_tracker import EventTracker
+from .events import Event, LLMEvent, RequestEnd, RequestStart, SimilarityEvent
+from .spans import EventSpan
+
 __all__ = [
     "CLIEventHandler",
     "LangSmithEventHandler",
+    "Event",
     "EventHandler",
+    "EventTracker",
+    "EventSpan",
+    "LLMEvent",
+    "RequestEnd",
+    "RequestStart",
+    "SimilarityEvent",
 ]

src/dbally/audit/event_handlers/base.py

@@ -1,8 +1,8 @@
 import abc
 from abc import ABC
-from typing import Generic, TypeVar, Union
+from typing import Generic, Optional, TypeVar
 
-from dbally.data_models.audit import LLMEvent, RequestEnd, RequestStart, SimilarityEvent
+from dbally.audit.events import Event, RequestEnd, RequestStart
 
 RequestCtx = TypeVar("RequestCtx")
 EventCtx = TypeVar("EventCtx")
@@ -26,7 +26,7 @@ async def request_start(self, user_request: RequestStart) -> RequestCtx:
         """
 
     @abc.abstractmethod
-    async def event_start(self, event: Union[LLMEvent, SimilarityEvent], request_context: RequestCtx) -> EventCtx:
+    async def event_start(self, event: Event, request_context: RequestCtx) -> EventCtx:
         """
         Function that is called during every event execution.
 
@@ -40,9 +40,7 @@ async def event_start(self, event: Union[LLMEvent, SimilarityEvent], request_con
         """
 
     @abc.abstractmethod
-    async def event_end(
-        self, event: Union[None, LLMEvent, SimilarityEvent], request_context: RequestCtx, event_context: EventCtx
-    ) -> None:
+    async def event_end(self, event: Optional[Event], request_context: RequestCtx, event_context: EventCtx) -> None:
         """
         Function that is called during every event execution.
 

src/dbally/audit/event_handlers/cli_event_handler.py

@@ -1,7 +1,7 @@
 import re
 from io import StringIO
 from sys import stdout
-from typing import Optional, Union
+from typing import Optional
 
 try:
     from rich import print as pprint
@@ -15,7 +15,7 @@
     pprint = print  # type: ignore
 
 from dbally.audit.event_handlers.base import EventHandler
-from dbally.data_models.audit import LLMEvent, RequestEnd, RequestStart, SimilarityEvent
+from dbally.audit.events import Event, LLMEvent, RequestEnd, RequestStart, SimilarityEvent
 
 _RICH_FORMATING_KEYWORD_SET = {"green", "orange", "grey", "bold", "cyan"}
 _RICH_FORMATING_PATTERN = rf"\[.*({'|'.join(_RICH_FORMATING_KEYWORD_SET)}).*\]"
@@ -40,14 +40,14 @@ class CLIEventHandler(EventHandler):
     ![Example output from CLIEventHandler](../../assets/event_handler_example.png)
     """
 
-    def __init__(self, buffer: StringIO = None) -> None:
+    def __init__(self, buffer: Optional[StringIO] = None) -> None:
         super().__init__()
 
         self.buffer = buffer
         out = self.buffer if buffer else stdout
         self._console = Console(file=out, record=True) if RICH_OUTPUT else None
 
-    def _print_syntax(self, content: str, lexer: str = None) -> None:
+    def _print_syntax(self, content: str, lexer: Optional[str] = None) -> None:
         if self._console:
             if lexer:
                 console_content = Syntax(content, lexer, word_wrap=True)
@@ -69,7 +69,7 @@ async def request_start(self, user_request: RequestStart) -> None:
         self._print_syntax("[grey53]\n=======================================")
         self._print_syntax("[grey53]=======================================\n")
 
-    async def event_start(self, event: Union[LLMEvent, SimilarityEvent], request_context: None) -> None:
+    async def event_start(self, event: Event, request_context: None) -> None:
         """
         Displays information that event has started, then all messages inside the prompt
 
@@ -98,9 +98,7 @@ async def event_start(self, event: Union[LLMEvent, SimilarityEvent], request_con
                 f"[cyan bold]FETCHER: [grey53]{event.fetcher}\n"
             )
 
-    async def event_end(
-        self, event: Union[None, LLMEvent, SimilarityEvent], request_context: None, event_context: None
-    ) -> None:
+    async def event_end(self, event: Optional[Event], request_context: None, event_context: None) -> None:
         """
         Displays the response from the LLM.
 

src/dbally/audit/event_handlers/langsmith_event_handler.py

@@ -1,12 +1,12 @@
 import socket
 from getpass import getuser
-from typing import Optional, Union
+from typing import Optional
 
 from langsmith.client import Client
 from langsmith.run_trees import RunTree
 
 from dbally.audit.event_handlers.base import EventHandler
-from dbally.data_models.audit import LLMEvent, RequestEnd, RequestStart, SimilarityEvent
+from dbally.audit.events import Event, LLMEvent, RequestEnd, RequestStart, SimilarityEvent
 
 
 class LangSmithEventHandler(EventHandler[RunTree, RunTree]):
@@ -47,7 +47,7 @@ async def request_start(self, user_request: RequestStart) -> RunTree:
 
         return run_tree
 
-    async def event_start(self, event: Union[None, LLMEvent, SimilarityEvent], request_context: RunTree) -> RunTree:
+    async def event_start(self, event: Event, request_context: RunTree) -> RunTree:
         """
         Log the start of the event.
 
@@ -79,9 +79,7 @@ async def event_start(self, event: Union[None, LLMEvent, SimilarityEvent], reque
 
         raise ValueError("Unsupported event")
 
-    async def event_end(
-        self, event: Union[None, LLMEvent, SimilarityEvent], request_context: RunTree, event_context: RunTree
-    ) -> None:
+    async def event_end(self, event: Optional[Event], request_context: RunTree, event_context: RunTree) -> None:
         """
         Log the end of the event.
 

src/dbally/audit/event_span.py

@@ -1,9 +1,9 @@
 from contextlib import asynccontextmanager
-from typing import AsyncIterator, Dict, List, Optional, Union
+from typing import AsyncIterator, Dict, List, Optional
 
 from dbally.audit.event_handlers.base import EventHandler
-from dbally.audit.event_span import EventSpan
-from dbally.data_models.audit import LLMEvent, RequestEnd, RequestStart, SimilarityEvent
+from dbally.audit.events import Event, RequestEnd, RequestStart
+from dbally.audit.spans import EventSpan
 
 
 class EventTracker:
@@ -69,7 +69,7 @@ def subscribe(self, event_handler: EventHandler) -> None:
         self._handlers.append(event_handler)
 
     @asynccontextmanager
-    async def track_event(self, event: Union[LLMEvent, SimilarityEvent]) -> AsyncIterator[EventSpan]:
+    async def track_event(self, event: Event) -> AsyncIterator[EventSpan]:
         """
         Context manager for processing an event.
 

src/dbally/audit/event_tracker.py

@@ -1,21 +1,20 @@
+from abc import ABC
 from dataclasses import dataclass
-from enum import Enum
 from typing import Optional, Union
 
 from dbally.collection.results import ExecutionResult
 from dbally.prompts import ChatFormat
 
 
-class EventType(Enum):
+@dataclass
+class Event(ABC):
     """
-    Enum for event types.
+    Base class for all events.
     """
 
-    LLM = "LLM"
-
 
 @dataclass
-class LLMEvent:
+class LLMEvent(Event):
     """
     Class for LLM event.
     """
@@ -30,7 +29,7 @@ class LLMEvent:
 
 
 @dataclass
-class SimilarityEvent:
+class SimilarityEvent(Event):
     """
     SimilarityEvent is fired when a SimilarityIndex lookup is performed.
     """

src/dbally/data_models/audit.py renamed to src/dbally/audit/events.py

@@ -0,0 +1,21 @@
+from typing import Optional
+
+from dbally.audit.events import Event
+
+
+class EventSpan:
+    """
+    Helper class for logging events.
+    """
+
+    def __init__(self) -> None:
+        self.data: Optional[Event] = None
+
+    def __call__(self, data: Event) -> None:
+        """
+        Call method for logging events.
+
+        Args:
+            data: Event data.
+        """
+        self.data = data