model-library 0.1.2__py3-none-any.whl → 0.1.3__py3-none-any.whl

model_library/base/base.py

@@ -218,6 +218,10 @@ class LLM(ABC):
         Join input with history
         Log, Time, and Retry
         """
+
+        # verbose on debug
+        verbose = self.logger.isEnabledFor(logging.DEBUG)
+
         # format str input
         if isinstance(input, str):
             input = [TextInput(text=input)]
@@ -226,11 +230,11 @@ class LLM(ABC):
         input = [*files, *images, *input]
 
         # format input info
-        item_info = f"--- input ({len(input)}): {get_pretty_input_types(input)}\n"
+        item_info = (
+            f"--- input ({len(input)}): {get_pretty_input_types(input, verbose)}\n"
+        )
         if history:
-            item_info += (
-                f"--- history({len(history)}): {get_pretty_input_types(history)}\n"
-            )
+            item_info += f"--- history({len(history)}): {get_pretty_input_types(history, verbose)}\n"
 
         # format tool info
         tool_results = [t for t in input if isinstance(t, ToolResult)]
@@ -251,7 +255,7 @@ class LLM(ABC):
 
         # unique logger for the query
         query_id = uuid.uuid4().hex[:14]
-        query_logger = logging.getLogger(f"{self.logger.name}<query={query_id}>")
+        query_logger = self.logger.getChild(f"query={query_id}")
 
         query_logger.info(
             "Query started:\n" + item_info + tool_info + f"--- kwargs: {short_kwargs}\n"
@@ -277,6 +281,7 @@ class LLM(ABC):
         output.metadata.cost = await self._calculate_cost(output.metadata)
 
         query_logger.info(f"Query completed: {repr(output)}")
+        query_logger.debug(output.model_dump(exclude={"history", "raw"}))
 
         return output
 
@@ -316,7 +321,7 @@ class LLM(ABC):
             )
 
         # costs for long context
-        total_in = metadata.in_tokens + (metadata.cache_read_tokens or 0)
+        total_in = metadata.total_input_tokens
         if costs.context and total_in > costs.context.threshold:
             input_cost, output_cost = costs.context.get_costs(
                 input_cost,

model_library/base/output.py

@@ -59,6 +59,33 @@ class QueryResultCost(BaseModel):
             )
         )
 
+    @computed_field
+    @property
+    def total_input(self) -> float:
+        return sum(
+            filter(
+                None,
+                [
+                    self.input,
+                    self.cache_read,
+                    self.cache_write,
+                ],
+            )
+        )
+
+    @computed_field
+    @property
+    def total_output(self) -> float:
+        return sum(
+            filter(
+                None,
+                [
+                    self.output,
+                    self.reasoning,
+                ],
+            )
+        )
+
     @override
     def __repr__(self):
         use_cents = self.total < 1
@@ -92,6 +119,33 @@ class QueryResultMetadata(BaseModel):
     def default_duration_seconds(self) -> float:
         return self.duration_seconds or 0
 
+    @computed_field
+    @property
+    def total_input_tokens(self) -> int:
+        return sum(
+            filter(
+                None,
+                [
+                    self.in_tokens,
+                    self.cache_read_tokens,
+                    self.cache_write_tokens,
+                ],
+            )
+        )
+
+    @computed_field
+    @property
+    def total_output_tokens(self) -> int:
+        return sum(
+            filter(
+                None,
+                [
+                    self.out_tokens,
+                    self.reasoning_tokens,
+                ],
+            )
+        )
+
     def __add__(self, other: "QueryResultMetadata") -> "QueryResultMetadata":
         return QueryResultMetadata(
             in_tokens=self.in_tokens + other.in_tokens,

model_library/base/utils.py

@@ -21,12 +21,13 @@ def sum_optional(a: int | None, b: int | None) -> int | None:
     return (a or 0) + (b or 0)
 
 
-def get_pretty_input_types(input: Sequence["InputItem"]) -> str:
+def get_pretty_input_types(input: Sequence["InputItem"], verbose: bool = False) -> str:
     # for logging
     def process_item(item: "InputItem"):
         match item:
             case TextInput():
-                return truncate_str(repr(item))
+                item_str = repr(item)
+                return item_str if verbose else truncate_str(item_str)
             case FileBase():  # FileInput
                 return repr(item)
             case ToolResult():

model_library/config/ai21labs_models.yaml

@@ -19,6 +19,7 @@ ai21labs-models:
       supports_temperature: true
     default_parameters:
       temperature: 0.4
+      max_output_tokens: 4096
 
   ai21labs/jamba-large-1.7:
     label: Jamba 1.7 Large