From 8be6ab65bec27b738993a8eb5edd89a99c40fa24 Mon Sep 17 00:00:00 2001 From: Abram Date: Sun, 10 Dec 2023 14:24:58 +0100 Subject: [PATCH] Feat - introduce func response type --- agenta-cli/agenta/sdk/types.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/agenta-cli/agenta/sdk/types.py b/agenta-cli/agenta/sdk/types.py index 8c22032bf8..a18680ecb9 100644 --- a/agenta-cli/agenta/sdk/types.py +++ b/agenta-cli/agenta/sdk/types.py @@ -1,5 +1,5 @@ import json -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional from pydantic import BaseModel, Extra, HttpUrl @@ -10,6 +10,19 @@ def __init__(self, file_name: str, file_path: str): self.file_path = file_path +class FuncTokenUsage(BaseModel): + completion_tokens: str + prompt_tokens: str + total_tokens: str + + +class FuncResponse(BaseModel): + message: str + usage: Optional[FuncTokenUsage] + cost: Optional[str] + latency: str + + class DictInput(dict): def __new__(cls, default_keys=None): instance = super().__new__(cls, default_keys)