Implements #32

Author: johann-petrak

llms_wrapper/cost_logger.py

@@ -7,6 +7,39 @@
 from typing import Any, Optional
 from pathlib import Path
 import json
+import socket
+import getpass
+import datetime
+
+
+def get_username() -> str:
+    """
+    Get the current username in a portable way.
+
+    Returns:
+        Username as a string
+    """
+    try:
+        # getpass.getuser() works on Windows, Linux, and macOS
+        # It checks environment variables in order: LOGNAME, USER, LNAME, USERNAME
+        return getpass.getuser()
+    except Exception:
+        # Fallback to environment variables
+        return os.environ.get('USER') or os.environ.get('USERNAME') or 'unknown'
+
+
+def get_hostname() -> str:
+    """
+    Get the current hostname in a portable way.
+
+    Returns:
+        Hostname as a string
+    """
+    try:
+        # socket.gethostname() works on Windows, Linux, and macOS
+        return socket.gethostname()
+    except Exception:
+        return 'unknown'
 
 
 class Log2Sqlite:
@@ -39,6 +72,13 @@ def __init__(self, db_path: str, **defaults):
             Exception: If database initialization fails
         """
         self.db_path = Path(db_path)
+
+        if defaults.get("user") is None:
+            defaults["user"] = get_username()
+        if defaults.get("hostname") is None:
+            defaults["hostname"] = get_hostname()
+
+
         self.defaults = defaults
 
         # Validate that defaults only contain known fields
@@ -120,6 +160,10 @@ def log(self, row: dict):
         if invalid_fields:
             raise Exception(f"Invalid fields in row: {invalid_fields}")
 
+        # the datetime is always set fixed here!
+        row["datetime"] = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+
+
         # Merge defaults with provided row (row takes precedence)
         merged_row = {**self.defaults, **row}
 
@@ -245,9 +289,89 @@ def import_file(self, file: str):
                 raise Exception(f"Failed to import from {file}: {e}") from e
             raise
 
+    def rows(self, model=None, modelalias=None, hostname=None, user=None,
+             project=None, task=None, note=None, apikey_name=None,
+             date_from=None, date_to=None) -> list[dict]:
+        """
+        Get all rows matching specified criteria.
+
+        Args:
+            model: Filter by model name
+            modelalias: Filter by model alias
+            hostname: Filter by hostname
+            user: Filter by user
+            project: Filter by project
+            task: Filter by task
+            note: Filter by note
+            apikey_name: Filter by API key name
+            date_from: Filter by datetime >= this value (inclusive)
+            date_to: Filter by datetime <= this value (inclusive)
+
+        Returns:
+            List of dictionaries, each containing a matching row's data
+            (excluding the auto-increment id field)
+
+        Raises:
+            Exception: If query fails
+        """
+        conditions = []
+        values = []
+
+        # Add field filters
+        field_filters = {
+            'model': model,
+            'modelalias': modelalias,
+            'hostname': hostname,
+            'user': user,
+            'project': project,
+            'task': task,
+            'note': note,
+            'apikey_name': apikey_name
+        }
+
+        for field, value in field_filters.items():
+            if value is not None:
+                conditions.append(f'{field} = ?')
+                values.append(value)
+
+        # Add date range filters
+        if date_from is not None:
+            conditions.append('datetime >= ?')
+            values.append(date_from)
+
+        if date_to is not None:
+            conditions.append('datetime <= ?')
+            values.append(date_to)
+
+        # Build WHERE clause
+        where = ''
+        if conditions:
+            where = 'WHERE ' + ' AND '.join(conditions)
+
+        try:
+            conn = sqlite3.connect(self.db_path, timeout=5.0)
+            try:
+                conn.row_factory = sqlite3.Row  # Access columns by name
+                cursor = conn.execute(
+                    f'SELECT * FROM logs {where} ORDER BY datetime DESC, id DESC',
+                    values
+                )
+
+                # Convert rows to list of dicts, excluding id field
+                result = []
+                for row in cursor:
+                    row_dict = {key: row[key] for key in row.keys() if key != 'id'}
+                    result.append(row_dict)
+
+                return result
+            finally:
+                conn.close()
+        except Exception as e:
+            raise Exception(f"Failed to get rows: {e}") from e
+
     def get(self, model=None, modelalias=None, hostname=None, user=None,
             project=None, task=None, note=None, apikey_name=None,
-            date_from=None, date_to=None) -> tuple[float, int, int]:
+            date_from=None, date_to=None) -> tuple[float, int, int, int]:
         """
         Get aggregated cost and token sums for rows matching specified criteria.
 
@@ -264,7 +388,7 @@ def get(self, model=None, modelalias=None, hostname=None, user=None,
             date_to: Filter by datetime <= this value (inclusive)
 
         Returns:
-            Tuple of (total_cost, total_input_tokens, total_output_tokens)
+            Tuple of (total_cost, total_input_tokens, total_output_tokens, row_count)
 
         Raises:
             Exception: If query fails
@@ -310,12 +434,13 @@ def get(self, model=None, modelalias=None, hostname=None, user=None,
                     f'''SELECT 
                         COALESCE(SUM(cost), 0.0),
                         COALESCE(SUM(input_tokens), 0),
-                        COALESCE(SUM(output_tokens), 0)
+                        COALESCE(SUM(output_tokens), 0),
+                        COUNT(*)
                     FROM logs {where}''',
                     values
                 )
                 result = cursor.fetchone()
-                return (float(result[0]), int(result[1]), int(result[2]))
+                return (float(result[0]), int(result[1]), int(result[2]), int(result[3]))
             finally:
                 conn.close()
         except Exception as e:

llms_wrapper/llms.py

@@ -687,7 +687,8 @@ def query(
             tools: a list of tool dictionaries, each dictionary describing a tool. 
                 See https://docs.litellm.ai/docs/completion/function_call for the format. 
                 However, this can be created using the `make_tooling` function.
-            return_cost: whether or not LLM invocation costs should get returned
+            return_cost: whether or not LLM invocation costs should get returned. Gets automatically enabled if
+                cost logging is enabled.
             return_response: whether or not the complete reponse should get returned
             debug: if True, emits debug messages to aid development and debugging
             litellm_debug: if True, litellm debug logging is enabled, if False, disabled, if None, use debug setting
@@ -712,6 +713,8 @@ def cleaned_args(args: dict):
             return args
         if self.debug:
             debug = True
+        if self.cost_logger:
+            return_cost = True
         if litellm_debug is None and debug or litellm_debug:
             #  litellm.set_verbose = True    ## deprecated!
             os.environ['LITELLM_LOG'] = 'DEBUG'
@@ -863,7 +866,10 @@ def cost_callback(kwargs, completion_response, start_time, end_time):
                     ret["cost"] = callback_data.get("cost")
                     ret["n_prompt_tokens"] = callback_data.get("prompt_tokens")
                     ret["n_completion_tokens"] = callback_data.get("completion_tokens")
-                    self.cost_logger.log(dict(cost=ret["cost"], input_tokens=ret["input_tokens"], output_tokens=ret["output_tokens"]))
+                    self.cost_logger.log(
+                        dict(
+                            model=llm["llm"], modelalias=llm["alias"],
+                            cost=ret["cost"], input_tokens=ret["input_tokens"], output_tokens=ret["output_tokens"]))
                     return ret
                 except Exception as e:
                     tb = traceback.extract_tb(e.__traceback__)
@@ -937,7 +943,10 @@ def chunk_generator(model_generator, retobj):
                 ret["n_completion_tokens"] = usage.completion_tokens
                 ret["n_prompt_tokens"] = usage.prompt_tokens
                 ret["n_total_tokens"] = usage.total_tokens
-                self.cost_logger.log(dict(cost=ret["cost"], input_tokens=ret["n_prompt_tokens"], output_tokens=ret["n_completion_tokens"]))
+                self.cost_logger.log(
+                    dict(
+                        model=llm["llm"], modelalias=llm["alias"],
+                        cost=ret["cost"], input_tokens=ret["n_prompt_tokens"], output_tokens=ret["n_completion_tokens"]))
                 # add the cost and tokens from the recursive call info, if available
                 if recursive_call_info.get("cost") is not None:
                     ret["cost"] += recursive_call_info["cost"]