LM Proxy DB Connector

A minimalistic SQLAlchemy-based database connector for LM-Proxy.

Features

• 📊 Database connection management with SQLAlchemy
• 📝 Includes a component for logging of LLM requests and responses to various databases
• 🔄 Support for SQLite, PostgreSQL, MySQL, and other SQLAlchemy-supported databases
• 🛡️ Thread-safe implementation

Installation

pip install lm-proxy-db-connector

Quick Start

Using with LM-Proxy

Add the DB component to your LM-Proxy configuration:

# config.toml
[components.db]
dsn = "postgresql+psycopg2://user:password@localhost:5432/mydb"
class = "lm_proxy_db_connector.Component"

# Add database logging
[[loggers]]
class = "lm_proxy.loggers.BaseLogger"
[loggers.log_writer]
class = "lm_proxy_db_connector.logging.DBLogWriter"
table_name = "llm_logs"

Using with YAML config

# config.yml
components:
  db:
    class: "lm_proxy_db_connector.init_db"
    db_url: "postgresql+psycopg2://user:password@localhost:5432/mydb"

loggers:
  - class: "lm_proxy_db_connector.logging.DBLogger"
    table_name: "llm_logs"

Supported Database URLs

The connector uses SQLAlchemy's URL format:

• SQLite: sqlite:///path/to/database.db or sqlite:///:memory: (in-memory)
• PostgreSQL: postgresql+psycopg2://user:password@localhost:5432/dbname
• MySQL: mysql+pymysql://user:password@localhost:3306/dbname

Database Session Usage

from lm_proxy_db_connector import db_session
# DB initialization is handled by the LM Proxy component

# Use a session
with db_session() as session:
    result = session.execute("SELECT * FROM users")
    # Session is automatically committed or rolled back

Database Logger

from lm_proxy_db_connector.logging import DBLogger

# Create a logger that logs to a database table
logger = DBLogger(
    table_name="llm_logs",
    schema="public",  # Optional
    # Define column structure and mapping
    columns={
        "id": {"type": "string", "primary_key": True, "length": 36},
        "request": {"type": "json", "src": "request.messages"},
        "response": {"type": "text"},
        "created_at": {"type": "datetime", "default": "now"},
        # Completion tokens extracted from response.usage
        "completion_tokens": {"type": "integer", "src": "response.usage.completion_tokens"}
    }
)

Development

Setup Development Environment

# Clone the repository
git clone https://github.com/Nayjest/lm-proxy-db-connector.git
cd lm-proxy-db-connector

# Install in development mode
pip install -e .

Running Tests

# Start test databases (PostgreSQL and MySQL)
docker compose up -d

# Run tests
pytest

Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

License

This project is licensed under the MIT License - see the LICENSE file for details.