Release/0.4.1

.github/workflows/python-tests.yml

@@ -1,56 +1,55 @@
 name: Python SDK Tests
 
 on:
-  pull_request:
+    pull_request:
 
-  # Allows you to run this workflow manually from the Actions tab
-  workflow_dispatch:
+    # Allows you to run this workflow manually from the Actions tab
+    workflow_dispatch:
 
 jobs:
-  test:
-    permissions:
-      pull-requests: write
-
-    runs-on: ubuntu-latest
-
-    strategy:
-      matrix:
-        python-version: ['3.8', '3.9', '3.10', '3.11', '3.12', '3.13']
-
-    steps:
-    - uses: actions/checkout@v3
-
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
-      with:
-        python-version: ${{ matrix.python-version }}
-
-    - name: Install dependencies
-      working-directory: ./
-      run: |
-        python -m pip install --upgrade pip
-        pip install pytest pytest-cov
-        pip install -e .
-        pip install -r dev-requirements.txt
-
-    - name: Run tests
-      working-directory: ./
-      run: |
-        pytest tests/ --cov=basalt --cov-report=term --cov-report=xml
-        echo "PYTHON_VERSION=${{ matrix.python-version }}" >> $GITHUB_OUTPUT
-        echo "COVERAGE_PCT=$(python -c "import xml.etree.ElementTree as ET; tree = ET.parse('coverage.xml'); root = tree.getroot(); print(f'{float(root.attrib[\"line-rate\"]) * 100:.2f}')")" >> $GITHUB_OUTPUT
-        echo "TEST_COUNT=$(python -c "import xml.etree.ElementTree as ET; tree = ET.parse('coverage.xml'); root = tree.getroot(); print(root.find('.//metrics').attrib['tests'])")" >> $GITHUB_OUTPUT
-      id: test_results
-
-    - name: Create result file
-      run: |
-        mkdir -p test-results
-        echo "${{ steps.test_results.outputs.COVERAGE_PCT }}" > test-results/coverage.txt
-        echo "${{ steps.test_results.outputs.TEST_COUNT }}" > test-results/test-count.txt
-
-    - name: Upload test results
-      uses: actions/upload-artifact@v4
-      with:
-        name: test-results-${{ matrix.python-version }}
-        path: test-results/
-
+    test:
+        permissions:
+            pull-requests: write
+
+        runs-on: ubuntu-latest
+
+        strategy:
+            matrix:
+                python-version: ["3.10", "3.11", "3.12", "3.13"]
+
+        steps:
+            - uses: actions/checkout@v3
+
+            - name: Set up Python ${{ matrix.python-version }}
+              uses: actions/setup-python@v4
+              with:
+                  python-version: ${{ matrix.python-version }}
+
+            - name: Install dependencies
+              working-directory: ./
+              run: |
+                  python -m pip install --upgrade pip
+                  pip install pytest pytest-cov
+                  pip install -e .
+                  pip install -r dev-requirements.txt
+
+            - name: Run tests
+              working-directory: ./
+              run: |
+                  pytest tests/ --cov=basalt --cov-report=term --cov-report=xml
+                  echo "PYTHON_VERSION=${{ matrix.python-version }}" >> $GITHUB_OUTPUT
+                  echo "COVERAGE_PCT=$(python -c "import xml.etree.ElementTree as ET; tree = ET.parse('coverage.xml'); root = tree.getroot(); print(f'{float(root.attrib[\"line-rate\"]) * 100:.2f}')")" >> $GITHUB_OUTPUT
+                  echo "TEST_COUNT=$(python -c "import xml.etree.ElementTree as ET; tree = ET.parse('coverage.xml'); root = tree.getroot(); print(root.find('.//metrics').attrib['tests'])")" >> $GITHUB_OUTPUT
+              id: test_results
+
+            - name: Create result file
+              run: |
+                  mkdir -p test-results
+                  echo "${{ steps.test_results.outputs.COVERAGE_PCT }}" > test-results/coverage.txt
+                  echo "${{ steps.test_results.outputs.TEST_COUNT }}" > test-results/test-count.txt
+
+            - name: Upload test results
+              uses: actions/upload-artifact@v4
+              with:
+                  name: test-results-${{ matrix.python-version }}
+                  path: test-results/

basalt/endpoints/monitor/send_trace.py

@@ -88,11 +88,21 @@ def prepare_request(self, dto: Optional[Input] = None) -> Dict[str, Any]:
                 else:
                     processed_log["parentId"] = None
 
-								# Rename ideal output
+                # Rename ideal output
                 if "ideal_output" in processed_log:
                   processed_log["idealOutput"] = processed_log["ideal_output"]
                   del processed_log["ideal_output"]
 
+                # Rename input tokens
+                if "input_tokens" in processed_log:
+                    processed_log["inputTokens"] = processed_log["input_tokens"]
+                    del processed_log["input_tokens"]
+
+                # Rename output tokens
+                if "output_tokens" in processed_log:
+                    processed_log["outputTokens"] = processed_log["output_tokens"]
+                    del processed_log["output_tokens"]
+
                 processed_logs.append(processed_log)
 
         # Create the request body
@@ -124,18 +134,18 @@ def prepare_request(self, dto: Optional[Input] = None) -> Dict[str, Any]:
             "path": "/monitor/trace",
             "body": body
         }
-    
+
     def decode_response(self, response: Any) -> Tuple[Optional[Exception], Optional[Output]]:
         """
         Decodes the response from sending a trace.
-        
+
         Args:
             response (Any): The response from the API.
-            
+
         Returns:
             Tuple[Optional[Exception], Optional[Dict[str, Any]]]: The decoded response.
         """
         if not isinstance(response, dict):
             return Exception("Failed to decode response (invalid body format)"), None
-            
-        return None, response.get("trace", {})
+
+        return None, response.get("trace", {})

basalt/objects/base_log.py

@@ -2,7 +2,7 @@
 from typing import Dict, Optional, Any, List
 import uuid
 
-from ..ressources.monitor.base_log_types import BaseLogParams, LogType
+from ..ressources.monitor.base_log_types import BaseLogParamsWithType, LogType
 from ..ressources.monitor.evaluator_types import Evaluator
 from ..ressources.monitor.trace_types import Trace
 from ..ressources.monitor.log_types import Log
@@ -12,17 +12,19 @@ class BaseLog:
     """
     Base class for logs and generations.
     """
-    def __init__(self, params: BaseLogParams):
+    def __init__(self, params: BaseLogParamsWithType):
         self._id = f"log-{uuid.uuid4().hex[:8]}"
         self._type = params.get("type")
         self._name = params.get("name")
+        self._input = params.get("input")
+        self._output = params.get("output")
+        self._ideal_output = params.get("ideal_output")
         self._start_time = params.get("start_time") if params.get("start_time") is not None else datetime.now()
         self._end_time = params.get("end_time")
         self._metadata = params.get("metadata")
         self._trace = params.get("trace")
         self._parent = params.get("parent")
         self._evaluators = params.get("evaluators")
-        self._ideal_output = params.get("ideal_output")
 
         # Add to trace's logs list if trace exists
         if self._trace:
@@ -53,6 +55,16 @@ def name(self) -> str:
         """Get the log name."""
         return self._name
 
+    @property
+    def input(self) -> Optional['Input']:
+        """Get the generation input."""
+        return self._input
+
+    @property
+    def output(self) -> Optional['Output']:
+        """Get the generation output."""
+        return self._output
+
     @property
     def ideal_output(self) -> Optional[str]:
         """Get the ideal output."""

basalt/objects/generation.py

@@ -1,33 +1,32 @@
-from datetime import datetime
-from typing import Dict, Optional, Any, List, Union
+from typing import Dict, Optional, Any, List
 
 from .base_log import BaseLog
 from ..ressources.monitor.generation_types import GenerationParams
-from ..ressources.monitor.base_log_types import BaseLogParams, LogType
+from ..ressources.monitor.base_log_types import LogType, Input, Output, BaseLogParamsWithType
 
 
 class Generation(BaseLog):
     """
     Class representing a generation in the monitoring system.
     """
     def __init__(self, params: GenerationParams):
-        base_log_params = {
-            "name": params.get("name"),
-            "ideal_output": params.get("ideal_output"),
-            "start_time": params.get("start_time"),
-            "end_time": params.get("end_time"),
-            "metadata": params.get("metadata"),
-            "parent": params.get("parent"),
-            "trace": params.get("trace"),
-            "evaluators": params.get("evaluators"),
-            "type": LogType.GENERATION,
-        }
+        base_log_params = BaseLogParamsWithType(
+            name= params.get("name"),
+            input=params.get("input"),
+            output=params.get("output"),
+            ideal_output=params.get("ideal_output"),
+            start_time=params.get("start_time"),
+            end_time=params.get("end_time"),
+            metadata=params.get("metadata"),
+            parent=params.get("parent"),
+            trace=params.get("trace"),
+            evaluators=params.get("evaluators"),
+            type=LogType.GENERATION
+        )
 
         super().__init__(base_log_params)
 
         self._prompt = params.get("prompt")
-        self._input = params.get("input")
-        self._output = params.get("output")
         self._input_tokens = params.get("input_tokens")
         self._output_tokens = params.get("output_tokens")
         self._cost = params.get("cost")
@@ -51,16 +50,6 @@ def prompt(self) -> Optional[Dict[str, Any]]:
         """Get the generation prompt."""
         return self._prompt
 
-    @property
-    def input(self) -> Optional[str]:
-        """Get the generation input."""
-        return self._input
-
-    @property
-    def output(self) -> Optional[str]:
-        """Get the generation output."""
-        return self._output
-
     @property
     def input_tokens(self) -> Optional[int]:
         """Get the generation input tokens."""
@@ -91,7 +80,7 @@ def options(self, options: Dict[str, Any]):
         """Set the generation options."""
         self._options = options
 
-    def start(self, input: Optional[str] = None) -> 'Generation':
+    def start(self, input: Optional['Input'] = None) -> 'Generation':
         """
         Start the generation with an optional input.
 
@@ -107,23 +96,30 @@ def start(self, input: Optional[str] = None) -> 'Generation':
         super().start()
         return self
 
-    def end(self, output: Optional[Union[str, Dict[str, Any]]] = None) -> 'Generation':
+    def end(self,
+            output: Optional[Output] = None,
+            input_tokens: Optional[int] = None,
+            output_tokens: Optional[int] = None,
+            cost: Optional[float] = None
+            ) -> 'Generation':
         """
         End the generation with an optional output or update parameters.
 
         Args:
-            output (Optional[Union[str, Dict[str, Any]]]): The output of the generation
-                or a dictionary of parameters to update.
+            output (Optional[Output]): The output of the generation as string, array or a dictionary.
+            input_tokens (Optional[int]): Optional number of tokens used for the input.
+            output_tokens (Optional[int]): Optional number of tokens used for the output.
+            cost (Optional[float]): Cost of the generation.
 
         Returns:
             Generation: The generation instance.
         """
         super().end()
 
-        if isinstance(output, dict):
-            self.update(output)
-        elif isinstance(output, str):
-            self._output = output
+        self._output = output
+        self._input_tokens = input_tokens
+        self._output_tokens = output_tokens
+        self._cost = cost
 
         # If this is a single generation, end the trace as well
         if self._options and self._options.get("type") == "single":

basalt/objects/log.py

@@ -1,5 +1,6 @@
-from typing import Dict, Optional, Any, cast
-from ..ressources.monitor.log_types import LogParams
+from typing import Dict, Optional, cast, Any
+from ..ressources.monitor.base_log_types import BaseLogParamsWithType, LogType
+from ..ressources.monitor.log_types import LogParams, Input, Output
 from ..ressources.monitor.generation_types import GenerationParams
 from .base_log import BaseLog
 from .generation import Generation
@@ -9,21 +10,33 @@ class Log(BaseLog):
     Class representing a log in the monitoring system.
     """
     def __init__(self, params: LogParams):
-        super().__init__(params)
+        base_log_params = BaseLogParamsWithType(
+            name=params.get("name"),
+            ideal_output=params.get("ideal_output"),
+            start_time=params.get("start_time"),
+            end_time=params.get("end_time"),
+            metadata=params.get("metadata"),
+            parent=params.get("parent"),
+            trace=params.get("trace"),
+            evaluators=params.get("evaluators"),
+            type=LogType(params.get("type")),
+        )
+
+        super().__init__(base_log_params)
         self._input = params.get("input")
         self._output = None
 
     @property
-    def input(self) -> Optional[str]:
+    def input(self) -> Optional['Input']:
         """Get the log input."""
         return self._input
 
     @property
-    def output(self) -> Optional[str]:
+    def output(self) -> Optional['Output']:
         """Get the log output."""
         return self._output
 
-    def start(self, input: Optional[str] = None) -> 'Log':
+    def start(self, input: Optional['Input'] = None) -> 'Log':
         """
         Start the log with an optional input.
 
@@ -39,7 +52,7 @@ def start(self, input: Optional[str] = None) -> 'Log':
         super().start()
         return self
 
-    def end(self, output: Optional[str] = None) -> 'Log':
+    def end(self, output: Optional['Output'] = None) -> 'Log':
         """
         End the log with an optional output.