cognitor 0.0.0__tar.gz

cognitor-0.0.0/PKG-INFO

@@ -0,0 +1,112 @@
+Metadata-Version: 2.4
+Name: cognitor
+Version: 0.0.0
+Summary: Python SDK to extract relevant metrics from Small Language Model inference calls.
+Author-email: Riccardo <riccardo@tanaos.com>
+Project-URL: Homepage, https://github.com/riccardo/cognitor-py
+Project-URL: Bug Tracker, https://github.com/riccardo/cognitor-py/issues
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: transformers
+Requires-Dist: psutil
+Requires-Dist: torch
+Requires-Dist: pydantic
+Requires-Dist: psycopg2-binary
+Requires-Dist: sqlalchemy
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: black; extra == "dev"
+Requires-Dist: isort; extra == "dev"
+Requires-Dist: build; extra == "dev"
+Requires-Dist: twine; extra == "dev"
+Requires-Dist: python-dotenv>=1.0.1; extra == "dev"
+
+# cognitor-py
+
+`cognitor-py` is a Python SDK that wraps `transformers` inference calls to extract useful metadata and performance metrics.
+
+## Features
+
+- **Model Information**: Automatically captures the model name.
+- **Performance Metrics**: Tracks CPU and RAM usage during inference.
+- **GPU Monitoring**: Captures peak GPU memory usage (if CUDA is available).
+- **Token Counting**: Calculates input and output token counts for common pipeline tasks.
+- **Latency Tracking**: Measures inference duration.
+- **Error Handling**: Captures and reports errors during inference.
+- **Flexible Logging Targets**: Automatically saves all inference logs to either a local PostgreSQL database or a local file (JSON lines).
+- **Graceful Error Handling**: Ensures the program continues to run even if the database is unreachable.
+
+## Installation
+
+```bash
+pip install cognitor-py
+```
+
+## Usage
+
+### Using the Inference Monitor
+
+```python
+from transformers import pipeline, AutoTokenizer
+from cognitor import Cognitor
+
+# Initialize your model and tokenizer
+model_name = "gpt2"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+pipe = pipeline("text-generation", model=model_name, tokenizer=tokenizer)
+
+# Initialize Cognitor with PostgreSQL configuration (default)
+# Or use log_type="file" and log_path="logs.jsonl" for file logging
+cognitor = Cognitor(
+    model_name=model_name,
+    tokenizer=tokenizer,
+    log_type="database", # or "file"
+    host="localhost",
+    port=5432,
+    user="postgres",
+    password="postgres",
+    dbname="cognitor"
+)
+
+# Run inference within the monitor context
+with cognitor.monitor() as m:
+    input_text = "Once upon a time,"
+    # Use track() to capture only the inference duration
+    with m.track():
+        output = pipe(input_text, max_length=50)
+    m.capture(input_data=input_text, output=output)
+
+# The metadata is now available via the cognitor instance
+metadata = cognitor.get_last_metadata()
+print(output)
+print(metadata)
+```
+
+### Metadata Structure
+
+The extracted metadata follows this structure:
+
+```python
+{
+    "model_name": "gpt2",
+    "timestamp": "2026-04-01T14:34:14+0200",
+    "input_tokens": 5,
+    "output_tokens": 45,
+    "cpu_percent": 12.5,
+    "ram_usage_percent": 1.2,
+    "gpu_usage_percent": 5.5, # Optional
+    "duration": 0.45, # Inference-only duration
+    "input": "Once upon a time,",
+    "output": [...],
+    "error": None
+}
+```
+
+## License
+
+MIT

cognitor-0.0.0/README.md

@@ -0,0 +1,84 @@
+# cognitor-py
+
+`cognitor-py` is a Python SDK that wraps `transformers` inference calls to extract useful metadata and performance metrics.
+
+## Features
+
+- **Model Information**: Automatically captures the model name.
+- **Performance Metrics**: Tracks CPU and RAM usage during inference.
+- **GPU Monitoring**: Captures peak GPU memory usage (if CUDA is available).
+- **Token Counting**: Calculates input and output token counts for common pipeline tasks.
+- **Latency Tracking**: Measures inference duration.
+- **Error Handling**: Captures and reports errors during inference.
+- **Flexible Logging Targets**: Automatically saves all inference logs to either a local PostgreSQL database or a local file (JSON lines).
+- **Graceful Error Handling**: Ensures the program continues to run even if the database is unreachable.
+
+## Installation
+
+```bash
+pip install cognitor-py
+```
+
+## Usage
+
+### Using the Inference Monitor
+
+```python
+from transformers import pipeline, AutoTokenizer
+from cognitor import Cognitor
+
+# Initialize your model and tokenizer
+model_name = "gpt2"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+pipe = pipeline("text-generation", model=model_name, tokenizer=tokenizer)
+
+# Initialize Cognitor with PostgreSQL configuration (default)
+# Or use log_type="file" and log_path="logs.jsonl" for file logging
+cognitor = Cognitor(
+    model_name=model_name,
+    tokenizer=tokenizer,
+    log_type="database", # or "file"
+    host="localhost",
+    port=5432,
+    user="postgres",
+    password="postgres",
+    dbname="cognitor"
+)
+
+# Run inference within the monitor context
+with cognitor.monitor() as m:
+    input_text = "Once upon a time,"
+    # Use track() to capture only the inference duration
+    with m.track():
+        output = pipe(input_text, max_length=50)
+    m.capture(input_data=input_text, output=output)
+
+# The metadata is now available via the cognitor instance
+metadata = cognitor.get_last_metadata()
+print(output)
+print(metadata)
+```
+
+### Metadata Structure
+
+The extracted metadata follows this structure:
+
+```python
+{
+    "model_name": "gpt2",
+    "timestamp": "2026-04-01T14:34:14+0200",
+    "input_tokens": 5,
+    "output_tokens": 45,
+    "cpu_percent": 12.5,
+    "ram_usage_percent": 1.2,
+    "gpu_usage_percent": 5.5, # Optional
+    "duration": 0.45, # Inference-only duration
+    "input": "Once upon a time,",
+    "output": [...],
+    "error": None
+}
+```
+
+## License
+
+MIT

cognitor-0.0.0/cognitor/__init__.py

@@ -0,0 +1,3 @@
+from .monitor import Cognitor, InferenceMetadata
+
+__all__ = ["Cognitor", "InferenceMetadata"]

cognitor-0.0.0/cognitor/database.py

@@ -0,0 +1,58 @@
+import json
+import logging
+from typing import Any, Dict, Optional
+from sqlalchemy import create_engine, Column, Integer, String, Float, Text
+from sqlalchemy.orm import sessionmaker, declarative_base
+
+logger = logging.getLogger(__name__)
+
+Base = declarative_base()
+
+class InferenceLog(Base):
+    __tablename__ = 'inference_logs'
+
+    id = Column(Integer, primary_key=True, autoincrement=True)
+    model_name = Column(String)
+    timestamp = Column(String)
+    input_tokens = Column(Integer)
+    output_tokens = Column(Integer)
+    cpu_percent = Column(Float)
+    ram_usage_percent = Column(Float)
+    gpu_usage_percent = Column(Float, nullable=True)
+    duration = Column(Float)
+    input = Column(Text)
+    output = Column(Text)
+    error = Column(Text, nullable=True)
+    extra = Column(Text)
+
+class DatabaseManager:
+    """
+    Manages the database connection and schema initialization using SQLAlchemy.
+    """
+    def __init__(self, host: str = "localhost", port: int = 5432, user: str = "postgres", password: str = "postgres", dbname: str = "cognitor") -> None:
+        """
+        Initializes the DatabaseManager with connection parameters.
+        Args:
+            host (str): The database host address.
+            port (int): The database port number.
+            user (str): The database username.
+            password (str): The database password.
+            dbname (str): The name of the database.
+        """
+        self.db_url: str = f"postgresql://{user}:{password}@{host}:{port}/{dbname}"
+        self.engine: Any = create_engine(self.db_url)
+        self.Session: Any = sessionmaker(bind=self.engine)
+        self._initialized: bool = False
+
+    def init_db(self) -> None:
+        """
+        Initializes the database schema by creating all defined tables.
+        Exits gracefully if the database is unreachable.
+        """
+        try:
+            Base.metadata.create_all(self.engine)
+            self._initialized = True
+        except Exception as e:
+            print(f"CRITICAL: Database unreachable at {self.engine.url}. Error: {e}")
+            print("Exiting gracefully...")
+            exit(0)

cognitor-0.0.0/cognitor/monitor.py

@@ -0,0 +1,238 @@
+import time
+from typing import Any, Dict, Optional, Union, List
+from pydantic import BaseModel
+from .utils import get_resource_usage, Timer
+from .database import DatabaseManager
+from .service import LoggingService
+import torch
+
+class InferenceMetadata(BaseModel):
+    model_name: str
+    timestamp: str
+    input_tokens: int
+    output_tokens: int
+    cpu_percent: float
+    ram_usage_percent: float
+    gpu_usage_percent: Optional[float] = None
+    duration: float
+    input: Any
+    output: Any
+    error: Optional[str] = None
+    extra: Dict[str, Any] = {}
+
+class InferenceMonitor:
+    """
+    A context manager for monitoring a single inference call, capturing metrics and metadata.
+    """
+    def __init__(self, cognitor: 'Cognitor') -> None:
+        """
+        Initializes the InferenceMonitor with a Cognitor instance.
+        Args:
+            cognitor (Cognitor): The parent Cognitor instance.
+        """
+        self.cognitor: 'Cognitor' = cognitor
+        self.input: Any = None
+        self.output: Any = None
+        self.error: Optional[str] = None
+        self.start_usage: Optional[Dict[str, float]] = None
+        self.full_timer: Timer = Timer()
+        self.inference_timer: Timer = Timer()
+        self.metadata: Optional[InferenceMetadata] = None
+
+    def __enter__(self) -> 'InferenceMonitor':
+        """
+        Enters the monitoring context, recording initial resource usage and starting the timer.
+        Returns:
+            InferenceMonitor: The monitor instance.
+        """
+        self.start_usage = get_resource_usage()
+        self.full_timer.__enter__()
+        return self
+
+    def track(self) -> Timer:
+        """
+        Returns a context manager to track the specific inference duration.
+        Returns:
+            Timer: The inference timer context manager.
+        """
+        return self.inference_timer
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """
+        Exits the monitoring context, generating metadata and logging it to the configured target.
+        Args:
+            exc_type: The exception type if an error occurred.
+            exc_val: The exception value if an error occurred.
+            exc_tb: The traceback if an error occurred.
+        """
+        self.full_timer.__exit__(exc_type, exc_val, exc_tb)
+        if exc_type:
+            self.error = str(exc_val)
+        
+        self.metadata = self.cognitor._generate_metadata(
+            input_data=self.input,
+            output=self.output,
+            error=self.error,
+            full_timer=self.full_timer,
+            inference_timer=self.inference_timer,
+            start_usage=self.start_usage if self.start_usage else {}
+        )
+        self.cognitor._last_metadata = self.metadata
+        
+        # Log to database via LoggingService
+        if self.cognitor.logging_service:
+            self.cognitor.logging_service.log_inference(self.metadata.model_dump())
+
+    def capture(self, output: Any, input_data: Any = None) -> None:
+        """
+        Captures the output and optionally the input of the inference.
+        Args:
+            output (Any): The inference output data.
+            input_data (Any, optional): The inference input data. Defaults to None.
+        """
+        if input_data is not None:
+            self.input = input_data
+        self.output = output
+
+class Cognitor:
+    """
+    The main SDK class for monitoring SLM inference calls.
+    """
+    def __init__(self, model_name: str, tokenizer: Any = None, log_type: str = "database", log_path: Optional[str] = None, host: str = "localhost", port: int = 5432, user: str = "postgres", password: str = "postgres", dbname: str = "cognitor") -> None:
+        """
+        Initializes the Cognitor instance with model, tokenizer, and logging configuration.
+        Args:
+            model_name (str): The name of the model being monitored.
+            tokenizer (Any, optional): The tokenizer for token counting. Defaults to None.
+            log_type (str): The logging target ('database' or 'file'). Defaults to "database".
+            log_path (Optional[str]): The path for file logging. Defaults to None.
+            host (str): The database host. Defaults to "localhost".
+            port (int): The database port. Defaults to 5432.
+            user (str): The database user. Defaults to "postgres".
+            password (str): The database password. Defaults to "postgres".
+            dbname (str): The database name. Defaults to "cognitor".
+        """
+        self.model_name: str = model_name
+        self.tokenizer: Any = tokenizer
+        self._last_metadata: Optional[InferenceMetadata] = None
+        
+        self.db_manager: Optional[DatabaseManager] = None
+        self.logging_service: Optional[LoggingService] = None
+
+        if log_type == "database":
+            self.db_manager = DatabaseManager(host=host, port=port, user=user, password=password, dbname=dbname)
+            self.db_manager.init_db()
+            self.logging_service = LoggingService(db_manager=self.db_manager)
+        elif log_type == "file":
+            if not log_path:
+                log_path = f"{model_name}_logs.jsonl"
+            self.logging_service = LoggingService(log_file=log_path)
+        else:
+            raise ValueError(f"Invalid log_type: {log_type}. Must be 'database' or 'file'.")
+
+    def monitor(self) -> InferenceMonitor:
+        """
+        Returns a context manager to monitor an inference call.
+        Returns:
+            InferenceMonitor: The inference monitor context manager.
+        """
+        return InferenceMonitor(self)
+
+    def get_last_metadata(self) -> Optional[Dict[str, Any]]:
+        """
+        Returns the metadata from the last monitored inference.
+        Returns:
+            Optional[Dict[str, Any]]: The metadata dictionary, or None if no inference has been monitored.
+        """
+        return self._last_metadata.model_dump() if self._last_metadata else None
+
+    def monitor_inference(self, func: Any, *args: Any, **kwargs: Any) -> Dict[str, Any]:
+        """
+        Monitors a single inference call using the context manager.
+        Args:
+            func (Callable): The inference function to monitor.
+            *args: Positional arguments for the inference function.
+            **kwargs: Keyword arguments for the inference function.
+        Returns:
+            Dict[str, Any]: A dictionary containing the 'output' and 'metadata'.
+        """
+        input_data: Any = args[0] if args else kwargs
+        with self.monitor() as m:
+            try:
+                with m.track():
+                    output = func(*args, **kwargs)
+                m.capture(input_data=input_data, output=output)
+                return {"output": output, "metadata": m.metadata.model_dump() if m.metadata else {}}
+            except Exception as e:
+                raise e
+
+    def _generate_metadata(self, input_data: Any, output: Any, error: Optional[str], full_timer: Timer, inference_timer: Timer, start_usage: Dict[str, float]) -> InferenceMetadata:
+        """
+        Internal method to generate InferenceMetadata from collected metrics.
+        Args:
+            input_data (Any): The inference input data.
+            output (Any): The inference output data.
+            error (Optional[str]): The error message if an error occurred.
+            full_timer (Timer): The timer for the full context duration.
+            inference_timer (Timer): The timer for the specific inference duration.
+            start_usage (Dict[str, float]): The initial resource usage metrics.
+        Returns:
+            InferenceMetadata: The generated metadata object.
+        """
+        end_usage: Dict[str, float] = get_resource_usage()
+        
+        # Extract token counts if tokenizer is available
+        input_tokens: int = 0
+        output_tokens: int = 0
+        
+        if self.tokenizer:
+            if isinstance(input_data, str):
+                input_tokens = len(self.tokenizer.encode(input_data))
+            elif isinstance(input_data, list):
+                if all(isinstance(i, str) for i in input_data):
+                    input_tokens = sum(len(self.tokenizer.encode(i)) for i in input_data)
+                elif all(isinstance(i, dict) for i in input_data):
+                    # Handle chat-like inputs
+                    for msg in input_data:
+                        if isinstance(msg, dict) and "content" in msg:
+                            input_tokens += len(self.tokenizer.encode(msg["content"]))
+
+            if output:
+                if isinstance(output, list):
+                    for item in output:
+                        if isinstance(item, dict):
+                            if "generated_text" in item:
+                                output_tokens += len(self.tokenizer.encode(item["generated_text"]))
+                            elif "summary_text" in item:
+                                output_tokens += len(self.tokenizer.encode(item["summary_text"]))
+                            elif "translation_text" in item:
+                                output_tokens += len(self.tokenizer.encode(item["translation_text"]))
+                        elif isinstance(item, str):
+                            output_tokens += len(self.tokenizer.encode(item))
+                elif isinstance(output, dict):
+                    if "generated_text" in output:
+                        output_tokens = len(self.tokenizer.encode(output["generated_text"]))
+                elif isinstance(output, str):
+                    output_tokens = len(self.tokenizer.encode(output))
+
+        gpu_usage_percent: Optional[float] = None
+        if torch.cuda.is_available():
+            device = torch.cuda.current_device()
+            total_mem = torch.cuda.get_device_properties(device).total_memory
+            peak_mem = torch.cuda.max_memory_allocated(device)
+            gpu_usage_percent = (peak_mem / total_mem) * 100
+            torch.cuda.reset_peak_memory_stats(device)
+
+        return InferenceMetadata(
+            model_name=self.model_name,
+            timestamp=full_timer.timestamp,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cpu_percent=end_usage["cpu_percent"],
+            ram_usage_percent=end_usage["ram_usage_percent"],
+            gpu_usage_percent=gpu_usage_percent,
+            duration=inference_timer.duration,
+            input=input_data,
+            output=output,
+            error=error
+        )

cognitor-0.0.0/cognitor/service.py

@@ -0,0 +1,80 @@
+import json
+import logging
+from typing import Any, Dict, Optional
+from .database import DatabaseManager, InferenceLog
+
+logger = logging.getLogger(__name__)
+
+class LoggingService:
+    """
+    Service responsible for logging inference metadata to either a database or a local file.
+    """
+    def __init__(self, db_manager: Optional[DatabaseManager] = None, log_file: Optional[str] = None) -> None:
+        """
+        Initializes the LoggingService with a DatabaseManager and/or a log file path.
+        Args:
+            db_manager (Optional[DatabaseManager]): The database manager instance for DB logging.
+            log_file (Optional[str]): The path to the local file for file-based logging.
+        """
+        self.db_manager: Optional[DatabaseManager] = db_manager
+        self.log_file: Optional[str] = log_file
+
+    def log_inference(self, metadata: Dict[str, Any]) -> None:
+        """
+        Saves inference metadata to the configured target (database or file).
+        Args:
+            metadata (Dict[str, Any]): The inference metadata to be logged.
+        """
+        if self.db_manager and self.db_manager._initialized:
+            self._log_to_db(metadata)
+        
+        if self.log_file:
+            self._log_to_file(metadata)
+
+    def _log_to_db(self, metadata: Dict[str, Any]) -> None:
+        """
+        Internal method to save inference metadata to the database using SQLAlchemy.
+        Args:
+            metadata (Dict[str, Any]): The inference metadata to be logged.
+        """
+        if not self.db_manager:
+            return
+
+        session: Any = self.db_manager.Session()
+        try:
+            log_entry = InferenceLog(
+                model_name=metadata.get("model_name"),
+                timestamp=metadata.get("timestamp"),
+                input_tokens=metadata.get("input_tokens"),
+                output_tokens=metadata.get("output_tokens"),
+                cpu_percent=metadata.get("cpu_percent"),
+                ram_usage_percent=metadata.get("ram_usage_percent"),
+                gpu_usage_percent=metadata.get("gpu_usage_percent"),
+                duration=metadata.get("duration"),
+                input=json.dumps(metadata.get("input")),
+                output=json.dumps(metadata.get("output")),
+                error=metadata.get("error"),
+                extra=json.dumps(metadata.get("extra", {}))
+            )
+            session.add(log_entry)
+            session.commit()
+        except Exception as e:
+            session.rollback()
+            print(f"WARNING: Failed to log to database: {e}")
+        finally:
+            session.close()
+
+    def _log_to_file(self, metadata: Dict[str, Any]) -> None:
+        """
+        Internal method to save inference metadata to a local file in JSON lines format.
+        Args:
+            metadata (Dict[str, Any]): The inference metadata to be logged.
+        """
+        if not self.log_file:
+            return
+
+        try:
+            with open(self.log_file, "a") as f:
+                f.write(json.dumps(metadata) + "\n")
+        except Exception as e:
+            print(f"WARNING: Failed to log to file {self.log_file}: {e}")

cognitor-0.0.0/cognitor/utils.py

@@ -0,0 +1,51 @@
+import psutil
+import time
+import os
+from datetime import datetime, timezone
+from typing import Dict, Any
+
+def get_resource_usage() -> Dict[str, float]:
+    """
+    Returns current CPU and RAM usage percentages.
+    Returns:
+        Dict[str, float]: A dictionary containing 'cpu_percent' and 'ram_usage_percent'.
+    """
+    process = psutil.Process(os.getpid())
+    return {
+        "cpu_percent": psutil.cpu_percent(interval=None),
+        "ram_usage_percent": process.memory_percent()
+    }
+
+class Timer:
+    """
+    A context manager to measure duration and record a UTC timestamp.
+    """
+    def __init__(self) -> None:
+        """
+        Initializes the Timer instance.
+        """
+        self.start: float = 0.0
+        self.end: float = 0.0
+        self.duration: float = 0.0
+        self.timestamp: str = ""
+
+    def __enter__(self) -> 'Timer':
+        """
+        Starts the timer and records the current UTC timestamp.
+        Returns:
+            Timer: The timer instance.
+        """
+        self.start = time.perf_counter()
+        self.timestamp = datetime.now(timezone.utc).isoformat()
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """
+        Stops the timer and calculates the duration.
+        Args:
+            exc_type: The exception type if an error occurred.
+            exc_val: The exception value if an error occurred.
+            exc_tb: The traceback if an error occurred.
+        """
+        self.end = time.perf_counter()
+        self.duration = self.end - self.start

cognitor-0.0.0/cognitor.egg-info/PKG-INFO

@@ -0,0 +1,112 @@
+Metadata-Version: 2.4
+Name: cognitor
+Version: 0.0.0
+Summary: Python SDK to extract relevant metrics from Small Language Model inference calls.
+Author-email: Riccardo <riccardo@tanaos.com>
+Project-URL: Homepage, https://github.com/riccardo/cognitor-py
+Project-URL: Bug Tracker, https://github.com/riccardo/cognitor-py/issues
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: transformers
+Requires-Dist: psutil
+Requires-Dist: torch
+Requires-Dist: pydantic
+Requires-Dist: psycopg2-binary
+Requires-Dist: sqlalchemy
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: black; extra == "dev"
+Requires-Dist: isort; extra == "dev"
+Requires-Dist: build; extra == "dev"
+Requires-Dist: twine; extra == "dev"
+Requires-Dist: python-dotenv>=1.0.1; extra == "dev"
+
+# cognitor-py
+
+`cognitor-py` is a Python SDK that wraps `transformers` inference calls to extract useful metadata and performance metrics.
+
+## Features
+
+- **Model Information**: Automatically captures the model name.
+- **Performance Metrics**: Tracks CPU and RAM usage during inference.
+- **GPU Monitoring**: Captures peak GPU memory usage (if CUDA is available).
+- **Token Counting**: Calculates input and output token counts for common pipeline tasks.
+- **Latency Tracking**: Measures inference duration.
+- **Error Handling**: Captures and reports errors during inference.
+- **Flexible Logging Targets**: Automatically saves all inference logs to either a local PostgreSQL database or a local file (JSON lines).
+- **Graceful Error Handling**: Ensures the program continues to run even if the database is unreachable.
+
+## Installation
+
+```bash
+pip install cognitor-py
+```
+
+## Usage
+
+### Using the Inference Monitor
+
+```python
+from transformers import pipeline, AutoTokenizer
+from cognitor import Cognitor
+
+# Initialize your model and tokenizer
+model_name = "gpt2"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+pipe = pipeline("text-generation", model=model_name, tokenizer=tokenizer)
+
+# Initialize Cognitor with PostgreSQL configuration (default)
+# Or use log_type="file" and log_path="logs.jsonl" for file logging
+cognitor = Cognitor(
+    model_name=model_name,
+    tokenizer=tokenizer,
+    log_type="database", # or "file"
+    host="localhost",
+    port=5432,
+    user="postgres",
+    password="postgres",
+    dbname="cognitor"
+)
+
+# Run inference within the monitor context
+with cognitor.monitor() as m:
+    input_text = "Once upon a time,"
+    # Use track() to capture only the inference duration
+    with m.track():
+        output = pipe(input_text, max_length=50)
+    m.capture(input_data=input_text, output=output)
+
+# The metadata is now available via the cognitor instance
+metadata = cognitor.get_last_metadata()
+print(output)
+print(metadata)
+```
+
+### Metadata Structure
+
+The extracted metadata follows this structure:
+
+```python
+{
+    "model_name": "gpt2",
+    "timestamp": "2026-04-01T14:34:14+0200",
+    "input_tokens": 5,
+    "output_tokens": 45,
+    "cpu_percent": 12.5,
+    "ram_usage_percent": 1.2,
+    "gpu_usage_percent": 5.5, # Optional
+    "duration": 0.45, # Inference-only duration
+    "input": "Once upon a time,",
+    "output": [...],
+    "error": None
+}
+```
+
+## License
+
+MIT

cognitor-0.0.0/cognitor.egg-info/SOURCES.txt

@@ -0,0 +1,14 @@
+README.md
+pyproject.toml
+cognitor/__init__.py
+cognitor/database.py
+cognitor/monitor.py
+cognitor/service.py
+cognitor/utils.py
+cognitor.egg-info/PKG-INFO
+cognitor.egg-info/SOURCES.txt
+cognitor.egg-info/dependency_links.txt
+cognitor.egg-info/requires.txt
+cognitor.egg-info/top_level.txt
+tests/test_integration.py
+tests/test_monitor.py

cognitor-0.0.0/cognitor.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

cognitor-0.0.0/cognitor.egg-info/requires.txt

@@ -0,0 +1,14 @@
+transformers
+psutil
+torch
+pydantic
+psycopg2-binary
+sqlalchemy
+
+[dev]
+pytest
+black
+isort
+build
+twine
+python-dotenv>=1.0.1

cognitor-0.0.0/cognitor.egg-info/top_level.txt

@@ -0,0 +1 @@
+cognitor

cognitor-0.0.0/pyproject.toml

@@ -0,0 +1,45 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "cognitor"
+dynamic = ["version"]
+description = "Python SDK to extract relevant metrics from Small Language Model inference calls."
+readme = "README.md"
+requires-python = ">=3.8"
+authors = [
+    { name = "Riccardo", email = "riccardo@tanaos.com" },
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Programming Language :: Python :: 3.10",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+    "transformers",
+    "psutil",
+    "torch",
+    "pydantic",
+    "psycopg2-binary",
+    "sqlalchemy",
+]
+
+[project.urls]
+"Homepage" = "https://github.com/riccardo/cognitor-py"
+"Bug Tracker" = "https://github.com/riccardo/cognitor-py/issues"
+
+[tool.setuptools_scm]
+version_file = "cognitor/__version__.py"
+
+[project.optional-dependencies]
+dev = [
+    "pytest",
+    "black",
+    "isort",
+    "build",
+    "twine",
+    "python-dotenv>=1.0.1",
+]

cognitor-0.0.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

cognitor-0.0.0/tests/test_integration.py

@@ -0,0 +1,46 @@
+import pytest
+from transformers import pipeline, AutoTokenizer
+from cognitor import Cognitor
+from unittest.mock import patch
+
+@patch("cognitor.monitor.DatabaseManager")
+def test_text_classification_pipeline(mock_db_manager, tmp_path):
+    model_name = "distilbert-base-uncased-finetuned-sst-2-english"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    # Use 'text-classification' which is the standard name for sentiment-analysis
+    pipe = pipeline("text-classification", model=model_name, tokenizer=tokenizer)
+    
+    cognitor = Cognitor(model_name=model_name, tokenizer=tokenizer)
+    
+    input_text = "I love this library!"
+    with cognitor.monitor() as m:
+        with m.track():
+            output = pipe(input_text)
+        m.capture(input_data=input_text, output=output)
+    
+    assert output is not None
+    metadata = cognitor.get_last_metadata()
+    assert metadata["model_name"] == model_name
+    assert metadata["input_tokens"] > 0
+    assert metadata["duration"] > 0
+
+@patch("cognitor.monitor.DatabaseManager")
+def test_text_generation_pipeline(mock_db_manager, tmp_path):
+    model_name = "gpt2"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    pipe = pipeline("text-generation", model=model_name, tokenizer=tokenizer)
+    
+    cognitor = Cognitor(model_name=model_name, tokenizer=tokenizer)
+    
+    input_text = "Once upon a time,"
+    with cognitor.monitor() as m:
+        with m.track():
+            output = pipe(input_text, max_length=10, do_sample=False)
+        m.capture(input_data=input_text, output=output)
+    
+    assert output is not None
+    metadata = cognitor.get_last_metadata()
+    assert metadata["model_name"] == model_name
+    assert metadata["input_tokens"] > 0
+    assert metadata["output_tokens"] > 0
+    assert metadata["duration"] > 0

cognitor-0.0.0/tests/test_monitor.py

@@ -0,0 +1,70 @@
+import pytest
+from cognitor import Cognitor
+from unittest.mock import MagicMock, patch
+
+@patch("cognitor.monitor.DatabaseManager")
+def test_metrics_collection(mock_db_manager, tmp_path):
+    cognitor = Cognitor(model_name="test-model")
+    
+    with cognitor.monitor() as m:
+        input_val = 5
+        with m.track():
+            output = input_val * 2
+        m.capture(input_data=input_val, output=output)
+    
+    assert output == 10
+    metadata = cognitor.get_last_metadata()
+    assert metadata["model_name"] == "test-model"
+    assert "cpu_percent" in metadata
+    assert "ram_usage_percent" in metadata
+    assert "duration" in metadata
+    assert metadata["input"] == 5
+
+@patch("cognitor.monitor.DatabaseManager")
+def test_token_counting(mock_db_manager, tmp_path):
+    mock_tokenizer = MagicMock()
+    mock_tokenizer.encode.side_effect = lambda x: [1] * len(x.split())
+    
+    cognitor = Cognitor(model_name="test-model", tokenizer=mock_tokenizer)
+    
+    input_text = "This is a test input"
+    with cognitor.monitor() as m:
+        with m.track():
+            output = {"generated_text": "This is a test output"}
+        m.capture(input_data=input_text, output=output)
+    
+    metadata = cognitor.get_last_metadata()
+    assert metadata["input_tokens"] == 5 # "This is a test input" -> 5 words
+    assert metadata["output_tokens"] == 5 # "This is a test output" -> 5 words
+
+@patch("cognitor.monitor.DatabaseManager")
+def test_error_handling(mock_db_manager, tmp_path):
+    cognitor = Cognitor(model_name="test-model")
+    
+    with pytest.raises(ValueError):
+        with cognitor.monitor() as m:
+            with m.track():
+                m.capture(input_data=5, output=None)
+                raise ValueError("Inference failed")
+    
+    metadata = cognitor.get_last_metadata()
+    assert metadata["error"] == "Inference failed"
+
+def test_file_logging(tmp_path):
+    log_file = tmp_path / "test_logs.jsonl"
+    cognitor = Cognitor(model_name="test-model", log_type="file", log_path=str(log_file))
+    
+    with cognitor.monitor() as m:
+        input_val = 5
+        with m.track():
+            output = input_val * 2
+        m.capture(input_data=input_val, output=output)
+    
+    assert log_file.exists()
+    with open(log_file, "r") as f:
+        line = f.readline()
+        import json
+        data = json.loads(line)
+        assert data["model_name"] == "test-model"
+        assert data["input"] == 5
+        assert data["output"] == 10