refactor(tracer): split large tracer file into multiple specific trac…

…er files

TASK: IL-297

Concepts.md

@@ -31,7 +31,7 @@ class Task(ABC, Generic[Input, Output]):
 
 `Input` and `Output` are normal Python datatypes that can be serialized from and to JSON. For this the Intelligence
 Layer relies on [Pydantic](https://docs.pydantic.dev/). The types that can actually be used are defined in form
-of the type-alias [`PydanticSerializable`](src/intelligence_layer/core/tracer.py#L44).
+of the type-alias [`PydanticSerializable`](src/intelligence_layer/core/tracer/tracer.py#L44).
 
 The second parameter `task_span` is used for [tracing](#Trace) which is described below.
 

run.py

@@ -3,8 +3,7 @@
 from dotenv import load_dotenv
 from fastapi import Depends, FastAPI
 
-from intelligence_layer.core.model import ControlModel, LuminousControlModel
-from intelligence_layer.core.tracer import NoOpTracer
+from intelligence_layer.core import ControlModel, LuminousControlModel, NoOpTracer
 from intelligence_layer.use_cases.classify.classify import (
     ClassifyInput,
     SingleLabelClassifyOutput,

src/intelligence_layer/core/__init__.py

@@ -1,9 +1,8 @@
+from intelligence_layer.core.intelligence_app import AuthService as AuthService
 from intelligence_layer.core.intelligence_app import (
     AuthenticatedIntelligenceApp as AuthenticatedIntelligenceApp,
 )
-from intelligence_layer.core.intelligence_app import AuthService as AuthService
 from intelligence_layer.core.intelligence_app import IntelligenceApp as IntelligenceApp
-
 from .chunk import Chunk as Chunk
 from .chunk import ChunkInput as ChunkInput
 from .chunk import ChunkOutput as ChunkOutput
@@ -16,12 +15,13 @@
 from .echo import EchoInput as EchoInput
 from .echo import EchoOutput as EchoOutput
 from .echo import EchoTask as EchoTask
+from .echo import TokenWithLogProb as TokenWithLogProb
 from .instruct import Instruct as Instruct
 from .instruct import InstructInput as InstructInput
+from .intelligence_app import AuthService as AuthService
 from .intelligence_app import (
     AuthenticatedIntelligenceApp as AuthenticatedIntelligenceApp,
 )
-from .intelligence_app import AuthService as AuthService
 from .intelligence_app import IntelligenceApp as IntelligenceApp
 from .model import AlephAlphaModel as AlephAlphaModel
 from .model import CompleteInput as CompleteInput
@@ -37,24 +37,26 @@
 from .task import Input as Input
 from .task import Output as Output
 from .task import Task as Task
+from .task import Token as Token
 from .text_highlight import ScoredTextHighlight as ScoredTextHighlight
 from .text_highlight import TextHighlight as TextHighlight
 from .text_highlight import TextHighlightInput as TextHighlightInput
 from .text_highlight import TextHighlightOutput as TextHighlightOutput
-from .tracer import CompositeTracer as CompositeTracer
-from .tracer import FileSpan as FileSpan
-from .tracer import FileTaskSpan as FileTaskSpan
-from .tracer import FileTracer as FileTracer
-from .tracer import InMemorySpan as InMemorySpan
-from .tracer import InMemoryTaskSpan as InMemoryTaskSpan
-from .tracer import InMemoryTracer as InMemoryTracer
-from .tracer import LogEntry as LogEntry
-from .tracer import NoOpTracer as NoOpTracer
-from .tracer import OpenTelemetryTracer as OpenTelemetryTracer
-from .tracer import PydanticSerializable as PydanticSerializable
-from .tracer import Span as Span
-from .tracer import TaskSpan as TaskSpan
-from .tracer import Tracer as Tracer
-from .tracer import utc_now as utc_now
+from .tracer.composite_tracer import CompositeTracer as CompositeTracer
+from .tracer.in_memory_tracer import InMemorySpan as InMemorySpan
+from .tracer.in_memory_tracer import InMemoryTaskSpan as InMemoryTaskSpan
+from .tracer.in_memory_tracer import InMemoryTracer as InMemoryTracer
+from .tracer.open_telemetry_tracer import OpenTelemetryTracer as OpenTelemetryTracer
+from .tracer.tracer import FileSpan as FileSpan
+from .tracer.tracer import FileTaskSpan as FileTaskSpan
+from .tracer.tracer import FileTracer as FileTracer
+from .tracer.tracer import JsonSerializer as JsonSerializer
+from .tracer.tracer import LogEntry as LogEntry
+from .tracer.tracer import NoOpTracer as NoOpTracer
+from .tracer.tracer import PydanticSerializable as PydanticSerializable
+from .tracer.tracer import Span as Span
+from .tracer.tracer import TaskSpan as TaskSpan
+from .tracer.tracer import Tracer as Tracer
+from .tracer.tracer import utc_now as utc_now
 
 __all__ = [symbol for symbol in dir()]

src/intelligence_layer/core/tracer/composite_tracer.py

@@ -0,0 +1,108 @@
+from datetime import datetime
+from typing import Generic, Optional, Sequence
+
+from intelligence_layer.core import (
+    Chunk,
+    InMemoryTracer,
+    PydanticSerializable,
+    Span,
+    TaskSpan,
+    Tracer,
+    utc_now,
+)
+from intelligence_layer.core.tracer.tracer import SpanVar, TracerVar
+from intelligence_layer.use_cases import ClassifyInput, PromptBasedClassify
+
+
+class CompositeTracer(Tracer, Generic[TracerVar]):
+    """A :class:`Tracer` that allows for recording to multiple tracers simultaneously.
+
+    Each log-entry and span will be forwarded to all subtracers.
+
+    Args:
+        tracers: tracers that will be forwarded all subsequent log and span calls.
+
+    Example:
+        >>> from intelligence_layer.core import InMemoryTracer, FileTracer, CompositeTracer, Chunk
+        >>> from intelligence_layer.use_cases import PromptBasedClassify, ClassifyInput
+
+        >>> tracer_1 = InMemoryTracer()
+        >>> tracer_2 = InMemoryTracer()
+        >>> tracer = CompositeTracer([tracer_1, tracer_2])
+        >>> task = PromptBasedClassify()
+        >>> response = task.run(ClassifyInput(chunk=Chunk("Cool"), labels=frozenset({"label", "other label"})), tracer)
+    """
+
+    def __init__(self, tracers: Sequence[TracerVar]) -> None:
+        assert len(tracers) > 0
+        self.tracers = tracers
+
+    def span(
+        self,
+        name: str,
+        timestamp: Optional[datetime] = None,
+        trace_id: Optional[str] = None,
+    ) -> "CompositeSpan[Span]":
+        timestamp = timestamp or utc_now()
+        trace_id = self.ensure_id(trace_id)
+        return CompositeSpan(
+            [tracer.span(name, timestamp, trace_id) for tracer in self.tracers]
+        )
+
+    def task_span(
+        self,
+        task_name: str,
+        input: PydanticSerializable,
+        timestamp: Optional[datetime] = None,
+        trace_id: Optional[str] = None,
+    ) -> "CompositeTaskSpan":
+        timestamp = timestamp or utc_now()
+        trace_id = self.ensure_id(trace_id)
+        return CompositeTaskSpan(
+            [
+                tracer.task_span(task_name, input, timestamp, trace_id)
+                for tracer in self.tracers
+            ]
+        )
+
+
+class CompositeSpan(Generic[SpanVar], CompositeTracer[SpanVar], Span):
+    """A :class:`Span` that allows for recording to multiple spans simultaneously.
+
+    Each log-entry and span will be forwarded to all subspans.
+
+    Args:
+        tracers: spans that will be forwarded all subsequent log and span calls.
+    """
+
+    def id(self) -> str:
+        return self.tracers[0].id()
+
+    def log(
+        self,
+        message: str,
+        value: PydanticSerializable,
+        timestamp: Optional[datetime] = None,
+    ) -> None:
+        timestamp = timestamp or utc_now()
+        for tracer in self.tracers:
+            tracer.log(message, value, timestamp)
+
+    def end(self, timestamp: Optional[datetime] = None) -> None:
+        timestamp = timestamp or utc_now()
+        for tracer in self.tracers:
+            tracer.end(timestamp)
+
+
+class CompositeTaskSpan(CompositeSpan[TaskSpan], TaskSpan):
+    """A :class:`TaskSpan` that allows for recording to multiple TaskSpans simultaneously.
+
+    Each log-entry and span will be forwarded to all subspans.
+
+    Args:
+        tracers: task spans that will be forwarded all subsequent log and span calls.
+    """
+
+    def record_output(self, output: PydanticSerializable) -> None:
+        for tracer in self.tracers:
+            tracer.record_output(output)

src/intelligence_layer/core/tracer/in_memory_tracer.py

@@ -0,0 +1,134 @@
+from datetime import datetime
+from typing import Optional, Union
+
+from pydantic import BaseModel, Field, SerializeAsAny
+from rich.tree import Tree
+
+from intelligence_layer.core import (
+    LogEntry,
+    PydanticSerializable,
+    Span,
+    TaskSpan,
+    Tracer,
+    utc_now,
+)
+from intelligence_layer.core.tracer.tracer import _render_log_value
+
+
+class InMemoryTracer(BaseModel, Tracer):
+    """Collects log entries in a nested structure, and keeps them in memory.
+
+    If desired, the structure is serializable with Pydantic, so you can write out the JSON
+    representation to a file, or return via an API, or something similar.
+
+    Attributes:
+        name: A descriptive name of what the tracer contains log entries about.
+        entries: A sequential list of log entries and/or nested InMemoryTracers with their own
+            log entries.
+    """
+
+    entries: list[Union[LogEntry, "InMemoryTaskSpan", "InMemorySpan"]] = []
+
+    def span(
+        self,
+        name: str,
+        timestamp: Optional[datetime] = None,
+        trace_id: Optional[str] = None,
+    ) -> "InMemorySpan":
+        child = InMemorySpan(
+            name=name,
+            start_timestamp=timestamp or utc_now(),
+            trace_id=self.ensure_id(trace_id),
+        )
+        self.entries.append(child)
+        return child
+
+    def task_span(
+        self,
+        task_name: str,
+        input: PydanticSerializable,
+        timestamp: Optional[datetime] = None,
+        trace_id: Optional[str] = None,
+    ) -> "InMemoryTaskSpan":
+        child = InMemoryTaskSpan(
+            name=task_name,
+            input=input,
+            start_timestamp=timestamp or utc_now(),
+            trace_id=self.ensure_id(trace_id),
+        )
+        self.entries.append(child)
+        return child
+
+    def _rich_render_(self) -> Tree:
+        """Renders the trace via classes in the `rich` package"""
+        tree = Tree(label="Trace")
+
+        for log in self.entries:
+            tree.add(log._rich_render_())
+
+        return tree
+
+    def _ipython_display_(self) -> None:
+        """Default rendering for Jupyter notebooks"""
+        from rich import print
+
+        print(self._rich_render_())
+
+
+class InMemorySpan(InMemoryTracer, Span):
+    name: str
+    start_timestamp: datetime = Field(default_factory=datetime.utcnow)
+    end_timestamp: Optional[datetime] = None
+    trace_id: str
+
+    def id(self) -> str:
+        return self.trace_id
+
+    def log(
+        self,
+        message: str,
+        value: PydanticSerializable,
+        timestamp: Optional[datetime] = None,
+    ) -> None:
+        self.entries.append(
+            LogEntry(
+                message=message,
+                value=value,
+                timestamp=timestamp or utc_now(),
+                trace_id=self.id(),
+            )
+        )
+
+    def end(self, timestamp: Optional[datetime] = None) -> None:
+        if not self.end_timestamp:
+            self.end_timestamp = timestamp or utc_now()
+
+    def _rich_render_(self) -> Tree:
+        """Renders the trace via classes in the `rich` package"""
+        tree = Tree(label=self.name)
+
+        for log in self.entries:
+            tree.add(log._rich_render_())
+
+        return tree
+
+
+class InMemoryTaskSpan(InMemorySpan, TaskSpan):
+    input: SerializeAsAny[PydanticSerializable]
+    output: Optional[SerializeAsAny[PydanticSerializable]] = None
+
+    def record_output(self, output: PydanticSerializable) -> None:
+        self.output = output
+
+    def _rich_render_(self) -> Tree:
+        """Renders the trace via classes in the `rich` package"""
+        tree = Tree(label=self.name)
+
+        tree.add(_render_log_value(self.input, "Input"))
+
+        for log in self.entries:
+            tree.add(log._rich_render_())
+
+        tree.add(_render_log_value(self.output, "Output"))
+
+        return tree