diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index 306836f..86ee98d 100644 --- a/.github/workflows/python-tests.yml +++ b/.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/ diff --git a/basalt/endpoints/monitor/send_trace.py b/basalt/endpoints/monitor/send_trace.py index fe58696..e09bc3b 100644 --- a/basalt/endpoints/monitor/send_trace.py +++ b/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", {}) \ No newline at end of file + + return None, response.get("trace", {}) diff --git a/basalt/objects/base_log.py b/basalt/objects/base_log.py index 19be5e6..73f6b90 100644 --- a/basalt/objects/base_log.py +++ b/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.""" diff --git a/basalt/objects/generation.py b/basalt/objects/generation.py index b6e5ccb..704c1f4 100644 --- a/basalt/objects/generation.py +++ b/basalt/objects/generation.py @@ -1,9 +1,8 @@ -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): @@ -11,23 +10,23 @@ 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": diff --git a/basalt/objects/log.py b/basalt/objects/log.py index 9e6a512..b700274 100644 --- a/basalt/objects/log.py +++ b/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. diff --git a/basalt/objects/trace.py b/basalt/objects/trace.py index 85ad6a3..7b108f9 100644 --- a/basalt/objects/trace.py +++ b/basalt/objects/trace.py @@ -1,13 +1,13 @@ from datetime import datetime from typing import Dict, Optional, Any, List -from ..ressources.monitor.trace_types import TraceParams, User, Organization +from ..ressources.monitor.trace_types import TraceParams, User, Organization, Input, Output, IdealOutput from .base_log import BaseLog from .generation import Generation from .log import Log from ..utils.flusher import Flusher from .experiment import Experiment -from ..ressources.monitor.evaluator_types import Evaluator +from ..ressources.monitor.evaluator_types import Evaluator, EvaluationConfig from ..ressources.monitor.generation_types import GenerationParams from ..ressources.monitor.log_types import LogParams from ..utils.protocols import ILogger @@ -53,17 +53,17 @@ def name(self) -> Optional[str]: return self._name @property - def input(self) -> Optional[str]: + def input(self) -> Optional[Input]: """Get the trace input.""" return self._input @property - def output(self) -> Optional[str]: + def output(self) -> Optional[Output]: """Get the trace output.""" return self._output @property - def ideal_output(self) -> Optional[str]: + def ideal_output(self) -> Optional[IdealOutput]: """Get the trace ideal_output.""" return self._ideal_output @@ -122,12 +122,12 @@ def evaluators(self) -> Optional[List[Dict[str, Any]]]: """Get the evaluators.""" return self._evaluators - def start(self, input: Optional[str] = None) -> 'Trace': + def start(self, input: Optional[Input] = None) -> 'Trace': """ Start the trace with an optional input. Args: - input (Optional[str]): The input to the trace. + input (Optional[Input]): The input to the trace. Can be a string, an array or a dictionary. Returns: Trace: The trace instance. @@ -171,7 +171,7 @@ def set_metadata(self, metadata: Dict[str, Any]) -> 'Trace': self._metadata = metadata return self - def set_evaluation_config(self, config: Dict[str, Any]) -> 'Trace': + def set_evaluation_config(self, config: EvaluationConfig) -> 'Trace': """ Set the evaluation configuration for the trace. @@ -213,12 +213,12 @@ def add_evaluator(self, evaluator: Evaluator) -> 'Trace': self._evaluators.append(evaluator) return self - def update(self, params: Dict[str, Any]) -> 'Trace': + def update(self, params: TraceParams) -> 'Trace': """ Update the trace. Args: - params (Dict[str, Any]): Parameters to update. + params (TraceParams): Parameters to update. Returns: Trace: The trace instance. @@ -237,7 +237,7 @@ def update(self, params: Dict[str, Any]) -> 'Trace': self._name = params.get("name", self._name) self._evaluators = params.get("evaluators", self._evaluators) - self._evaluation_config = params.get("evaluationConfig", self._evaluation_config) + self._evaluation_config = params.get("evaluation_config", self._evaluation_config) return self @@ -294,7 +294,7 @@ def create_log(self, params: LogParams) -> 'BaseLog': return log - async def end(self, output: Optional[str] = None) -> 'Trace': + async def end(self, output: Optional[Output] = None) -> 'Trace': """ End the trace with an optional output. @@ -314,7 +314,7 @@ async def end(self, output: Optional[str] = None) -> 'Trace': return self - def end_sync(self, output: Optional[str] = None) -> 'Trace': + def end_sync(self, output: Optional[Output] = None) -> 'Trace': """ End the trace with an optional output synchronously. diff --git a/basalt/ressources/monitor/base_log_types.py b/basalt/ressources/monitor/base_log_types.py index eb28b78..7bb69ca 100644 --- a/basalt/ressources/monitor/base_log_types.py +++ b/basalt/ressources/monitor/base_log_types.py @@ -33,6 +33,10 @@ class LogType(str, Enum): # Type alias for log type strings - use this in TypedDict parameters LogTypeStr = Literal['span', 'generation', 'function', 'tool', 'retrieval', 'event'] +Input = str | List[Any] | Dict[str, Any] +Output = str | List[Any] | Dict[str, Any] +IdealOutput = str | List[Any] | Dict[str, Any] + class _BaseLogParamsRequired(TypedDict): """Required fields for BaseLogParams.""" name: str @@ -54,13 +58,19 @@ class BaseLogParams(_BaseLogParamsRequired, TypedDict, total=False): Every log must be associated with a trace. evaluators: The evaluators to attach to the log. """ - ideal_output: Optional[str] + input: Optional[Input] + output: Optional[Output] + ideal_output: Optional[IdealOutput] start_time: Optional[Union[datetime, str]] end_time: Optional[Union[datetime, str]] metadata: Optional[Dict[str, Any]] parent: Optional['Log'] trace: Optional['Trace'] - evaluators: Optional[List[Evaluator]] + evaluators: List[Evaluator] + +class BaseLogParamsWithType(BaseLogParams, TypedDict, total=False): + """Base parameters for creating a log entry.""" + type: LogType @dataclass class BaseLog: @@ -84,6 +94,9 @@ class BaseLog: Every log must be associated with a trace. evaluators: List of evaluators attached to the log. """ + input: Optional[Input] + output: Optional[Output] + ideal_output: Optional[IdealOutput] name: str type: LogType id: str = field(default_factory=lambda: str(f'log-{uuid4().hex[:8]}')) @@ -94,7 +107,7 @@ class BaseLog: trace: Optional['Trace'] = None evaluators: List[Evaluator] = field(default_factory=list) - def start(self: SelfType) -> SelfType: + def start(self: SelfType, input: Optional[Input] = None) -> SelfType: """Marks the log as started and sets the start time if not already set. Returns: @@ -135,7 +148,7 @@ def update(self: SelfType, params: Dict[str, Any]) -> SelfType: """ ... - def end(self: SelfType) -> SelfType: + def end(self: SelfType, **kwargs) -> SelfType: """Marks the log as ended. Returns: diff --git a/basalt/ressources/monitor/generation_types.py b/basalt/ressources/monitor/generation_types.py index caf6b5b..325ed38 100644 --- a/basalt/ressources/monitor/generation_types.py +++ b/basalt/ressources/monitor/generation_types.py @@ -1,7 +1,7 @@ from typing import Dict, Optional, Union, Any, TypedDict from dataclasses import dataclass, field -from .base_log_types import BaseLog, BaseLogParams, LogType +from .base_log_types import BaseLog, BaseLogParams, LogType, Input, Output class _PromptReferenceRequired(TypedDict): """Required fields for PromptReference.""" @@ -58,8 +58,6 @@ class GenerationParams(BaseLogParams, total=False): ``` """ prompt: Optional[PromptReference] - input: Optional[str] - output: Optional[str] variables: Optional[Dict[str, Any]] options: Optional[Dict[str, Any]] input_tokens: Optional[int] @@ -110,15 +108,13 @@ class Generation(BaseLog): ``` """ prompt: Optional[PromptReference] = None - input: Optional[str] = None - output: Optional[str] = None variables: Optional[Dict[str, Any]] = None type: LogType = field(default=LogType.GENERATION) input_tokens: Optional[int] = None output_tokens: Optional[int] = None cost: Optional[float] = None - def start(self, input: Optional[str] = None) -> 'Generation': + def start(self, input: Optional[Input] = None) -> 'Generation': """Marks the generation as started and sets the input if provided. Args: @@ -138,12 +134,20 @@ def start(self, input: Optional[str] = None) -> 'Generation': """ ... - 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': """Marks the generation as ended and sets the output if provided. Args: output (Optional[Union[str, Dict[str, Any]]]): Optional output data from the model. Can be either a string or a dictionary containing output parameters. + 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 for method chaining. @@ -157,11 +161,11 @@ def end(self, output: Optional[Union[str, Dict[str, Any]]] = None) -> 'Generatio generation.end("The capital of France is Paris.") # End a generation with output params - generation.end({ - "output": "The capital of France is Paris.", - "input_tokens": 10, - "output_tokens": 10, - "cost": 0.01 + generation.end( + output="The capital of France is Paris.", + input_tokens=10, + output_tokens=10, + cost=0.01 }) ``` """ diff --git a/basalt/ressources/monitor/log_types.py b/basalt/ressources/monitor/log_types.py index 9e0c4bc..2cf5256 100644 --- a/basalt/ressources/monitor/log_types.py +++ b/basalt/ressources/monitor/log_types.py @@ -1,7 +1,7 @@ from typing import Optional, TYPE_CHECKING, TypedDict from dataclasses import dataclass -from .base_log_types import BaseLog, BaseLogParams, LogType, LogTypeStr +from .base_log_types import BaseLog, BaseLogParams, LogTypeStr, Input, Output if TYPE_CHECKING: from .generation_types import Generation, GenerationParams @@ -22,8 +22,7 @@ class LogParams(BaseLogParams, _LogParamsRequired, total=False): input: Optional input data for this operation. output: Optional output data generated by the operation. """ - input: Optional[str] - output: Optional[str] + ... class UpdateLogParams(LogParams, total=False): """Parameters for updating a log.""" @@ -72,11 +71,7 @@ class Log(BaseLog): log.end('Processed user data') ``` """ - input: Optional[str] = None - output: Optional[str] = None - ideal_output: Optional[str] = None - - def start(self, input: Optional[str] = None) -> 'Log': + def start(self, input: Optional[Input] = None) -> 'Log': """Marks the log as started and sets the input if provided. Args: @@ -96,7 +91,7 @@ def start(self, input: Optional[str] = None) -> 'Log': """ ... - def end(self, output: Optional[str] = None) -> 'Log': + def end(self, output: Optional[Output] = None) -> 'Log': """Marks the log as ended and sets the output if provided. Args: diff --git a/basalt/ressources/monitor/trace_types.py b/basalt/ressources/monitor/trace_types.py index c397079..6d8cb3a 100644 --- a/basalt/ressources/monitor/trace_types.py +++ b/basalt/ressources/monitor/trace_types.py @@ -22,19 +22,23 @@ class Organization(TypedDict): id: str name: str +Input = str | List[Any] | Dict[str, Any] +Output = str | List[Any] | Dict[str, Any] +IdealOutput = str | List[Any] | Dict[str, Any] + class TraceParams(TypedDict, total=False): """Parameters for creating or updating a trace.""" name: Optional[str] - input: Optional[str] - output: Optional[str] - ideal_output: Optional[str] + input: Optional[Input] + output: Optional[Output] + ideal_output: Optional[IdealOutput] start_time: Optional[datetime] end_time: Optional[datetime] user: Optional[User] organization: Optional['Organization'] metadata: Optional[Dict[str, Any]] experiment: Optional[Experiment] - evaluators: Optional[List[Evaluator]] + evaluators: List[Evaluator] evaluation_config: Optional[EvaluationConfig] @@ -76,20 +80,20 @@ class Trace: ``` """ name: Optional[str] - input: Optional[str] - output: Optional[str] - ideal_output: Optional[str] + input: Optional[Input] + output: Optional[Output] + ideal_output: Optional[IdealOutput] start_time: datetime end_time: Optional[datetime] user: Optional[User] organization: Optional[Organization] metadata: Optional[Dict[str, Any]] experiment: Optional['Experiment'] - evaluators: Optional[List[Evaluator]] + evaluators: List[Evaluator] evaluation_config: Optional[EvaluationConfig] logs: List['BaseLog'] = field(default_factory=list) - def start(self, input: Optional[str] = None) -> 'Trace': + def start(self, input: Optional[Input] = None) -> 'Trace': """Marks the trace as started and sets the input if provided. Args: @@ -109,7 +113,7 @@ def start(self, input: Optional[str] = None) -> 'Trace': """ ... - def set_ideal_output(self, ideal_output: str) -> 'Trace': + def set_ideal_output(self, ideal_output: IdealOutput) -> 'Trace': """Sets the ideal output for the trace.""" ... @@ -295,7 +299,7 @@ def create_log(self, params: 'LogParams') -> 'Log': """ ... - async def end(self, output: Optional[str] = None) -> 'Trace': + async def end(self, output: Optional[Output] = None) -> 'Trace': """Marks the trace as ended and sets the output if provided. Args: @@ -309,13 +313,16 @@ async def end(self, output: Optional[str] = None) -> 'Trace': # End a trace without output trace.end() - # End a trace with output + # End a trace with text output trace.end('The capital of France is Paris.') + + # End a trace with json output + trace.end({ "response": "The capital of France is Paris." }) ``` """ ... - def end_sync(self, output: Optional[str] = None) -> 'Trace': + def end_sync(self, output: Optional[Output] = None) -> 'Trace': """Marks the trace as ended and sets the output if provided. Args: @@ -327,10 +334,12 @@ def end_sync(self, output: Optional[str] = None) -> 'Trace': Example: ```python # End a trace without output - trace.end() + trace.end_sync() # End a trace with output - trace.end('The capital of France is Paris.') + trace.end_sync('The capital of France is Paris.') + + trace.end_sync({ "response": 'The capital of France is Paris.' }) ``` """ ... diff --git a/basalt/utils/flusher.py b/basalt/utils/flusher.py index fcd44a6..9a304ff 100644 --- a/basalt/utils/flusher.py +++ b/basalt/utils/flusher.py @@ -80,9 +80,11 @@ def _log_to_dict(log: Any) -> Dict[str, Any]: # Add generation-specific fields if it's a generation if log.type == "generation" and hasattr(log, "prompt"): base_dict.update({ + "input_tokens": log.input_tokens, + "output_tokens": log.output_tokens, + "cost": log.cost, "prompt": log.prompt, "variables": log.variables if log.variables else [], # Ensure variables is always a list - "options": log.options if hasattr(log, "options") else None }) return base_dict diff --git a/setup.py b/setup.py index ec70b2b..025ca8a 100644 --- a/setup.py +++ b/setup.py @@ -30,5 +30,5 @@ def get_version(): "aiohttp>=3.8.0", "jinja2>=3.1.0", ], - python_requires=">=3.8" + python_requires=">=3.10" )