Modules

Derivative based on the original work here: https://github.com/thehyve/omop-cdm/blob/main/src/omop_cdm/regular/cdm600/tables.py Modifications made to this file: - Removed support for schema. - Added new tables

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at *
https://www.apache.org/licenses/LICENSE-2.0 *
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Derivative based on the original work here: https://github.com/thehyve/omop-cdm/blob/main/src/omop_cdm/regular/cdm54/tables.py Modifications made to this file: - Removed support for schema. - Added new tables

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at *
https://www.apache.org/licenses/LICENSE-2.0 *
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Engine and session factory for OMOP CDM databases.

This module provides an asynchronous SQLAlchemy engine factory with helpers to create/init CDM schemas and obtain async sessions across supported backends (SQLite, MySQL, PostgreSQL).

`CdmEngineFactory` ¶

Bases: object

Factory to create async SQLAlchemy engines and sessions for OMOP CDM.

Supports SQLite (default), MySQL, and PostgreSQL. Exposes convenience properties for the configured engine and async session maker.

Parameters:

Name	Description	Default
`db`	Database type: "sqlite", "mysql", or "pgsql".	`'sqlite'`
`host`	Database host (ignored for SQLite).	`'localhost'`
`port`	Database port (ignored for SQLite).	`5432`
`user`	Database user (ignored for SQLite).	`'root'`
`pw`	Database password (ignored for SQLite).	`'pass'`
`name`	Database name or SQLite filename.	`'cdm.sqlite'`
`schema`	PostgreSQL schema to use for CDM.	`''`

Source code in src/pyomop/engine_factory.py

class CdmEngineFactory(object):
    """Factory to create async SQLAlchemy engines and sessions for OMOP CDM.

    Supports SQLite (default), MySQL, and PostgreSQL. Exposes convenience
    properties for the configured engine and async session maker.

    Args:
        db: Database type: "sqlite", "mysql", or "pgsql".
        host: Database host (ignored for SQLite).
        port: Database port (ignored for SQLite).
        user: Database user (ignored for SQLite).
        pw: Database password (ignored for SQLite).
        name: Database name or SQLite filename.
        schema: PostgreSQL schema to use for CDM.
    """

    def __init__(
        self,
        db="sqlite",
        host="localhost",
        port=5432,
        user="root",
        pw="pass",
        name="cdm.sqlite",
        schema="",
    ):
        self._db = db
        self._name = name
        self._host = host
        self._port = port
        self._user = user
        self._pw = pw
        self._schema = schema
        self._engine = None
        self._base = None

    async def init_models(self, metadata):
        """Drop and re-create all tables from provided metadata.

        This is mainly used for tests and quick local setups.

        Args:
            metadata: SQLAlchemy ``MetaData`` containing table definitions.

        Raises:
            ValueError: If the engine has not been initialized.
        """
        if self._engine is None:
            raise ValueError("Database engine is not initialized.")
        async with self._engine.begin() as conn:
            await conn.run_sync(metadata.drop_all)
            await conn.run_sync(metadata.create_all)

    @property
    def db(self):
        """Return the configured database type (sqlite/mysql/pgsql)."""
        return self._db

    @property
    def host(self):
        """Return the configured database host (if applicable)."""
        return self._host

    @property
    def port(self):
        """Return the configured database port (if applicable)."""
        return self._port

    @property
    def name(self):
        """Return the configured database name or SQLite filename."""
        return self._name

    @property
    def user(self):
        """Return the configured database user (if applicable)."""
        return self._user

    @property
    def pw(self):
        """Return the configured database password (if applicable)."""
        return self._pw

    @property
    def schema(self):
        """Return the configured schema (PostgreSQL)."""
        return self._schema

    @property
    def base(self):
        """Return automapped classes when an engine exists, otherwise None."""
        if self.engine is not None:  # Not self_engine
            Base = automap_base()
            Base.prepare(self.engine, reflect=True)
            return Base.classes
        return None

    @property
    def engine(self):
        """Create or return the async engine for the configured backend.

        Returns:
            Async engine instance bound to the configured database.
        """
        if self._db == "sqlite":
            # Schemas are not supported for SQLite; warn if a non-default schema was provided
            if self._schema and self._schema not in ("", "public"):
                logger.warning(
                    "Schema is not supported for SQLite; ignoring schema='%s'",
                    self._schema,
                )
            self._engine = create_async_engine("sqlite+aiosqlite:///" + self._name)
        elif self._db == "mysql":
            # Schemas are not supported for MySQL in the same way as PostgreSQL; warn and ignore
            if self._schema and self._schema not in ("", "public"):
                logger.warning(
                    "Schema is not supported for MySQL; ignoring schema='%s'",
                    self._schema,
                )
            mysql_url = "mysql://{}:{}@{}:{}/{}"
            mysql_url = mysql_url.format(
                self._user, self._pw, self._host, self._port, self._name
            )
            self._engine = create_async_engine(
                mysql_url, isolation_level="READ UNCOMMITTED"
            )
        elif self._db == "pgsql":
            pgsql_url = "postgresql+asyncpg://{}:{}@{}:{}/{}"
            pgsql_url = pgsql_url.format(
                self._user, self._pw, self._host, self._port, self._name
            )
            connect_args = {}
            # If a schema is provided, set the PostgreSQL search_path so that all
            # operations (reflection, DDL/DML) use this schema by default.
            if self._schema and self._schema != "":
                connect_args = {"server_settings": {"search_path": self._schema}}
            self._engine = create_async_engine(pgsql_url, connect_args=connect_args)
        else:
            # Unknown DB type; create no engine and warn
            logger.warning("Unknown database type '%s'—no engine created.", self._db)
            return None
        return self._engine

    @property
    def session(self):
        """Return an async_sessionmaker for creating AsyncSession objects."""
        if self._engine is not None:
            async_session = async_sessionmaker(
                self._engine, expire_on_commit=False, class_=AsyncSession
            )
            return async_session
        return None

    @property
    def async_session(self):
        """Alias for session to maintain backward compatibility."""
        if self._engine is not None:
            async_session = async_sessionmaker(
                self._engine, expire_on_commit=False, class_=AsyncSession
            )
            return async_session
        return None

    @db.setter
    def db(self, value):
        self._db = value

    @name.setter
    def name(self, value):
        self._name = value

    @port.setter
    def port(self, value):
        self._port = value

    @host.setter
    def host(self, value):
        self._host = value

    @user.setter
    def user(self, value):
        self._user = value

    @pw.setter
    def pw(self, value):
        self._pw = value

    @schema.setter
    def schema(self, value):
        self._schema = value

`async_session` `property` ¶

Alias for session to maintain backward compatibility.

`base` `property` ¶

Return automapped classes when an engine exists, otherwise None.

`db` `property` `writable` ¶

Return the configured database type (sqlite/mysql/pgsql).

`engine` `property` ¶

Create or return the async engine for the configured backend.

Returns:

Type	Description
	Async engine instance bound to the configured database.

`host` `property` `writable` ¶

Return the configured database host (if applicable).

`name` `property` `writable` ¶

Return the configured database name or SQLite filename.

`port` `property` `writable` ¶

Return the configured database port (if applicable).

`pw` `property` `writable` ¶

Return the configured database password (if applicable).

`schema` `property` `writable` ¶

Return the configured schema (PostgreSQL).

`session` `property` ¶

Return an async_sessionmaker for creating AsyncSession objects.

`user` `property` `writable` ¶

Return the configured database user (if applicable).

`init_models(metadata)` `async` ¶

Drop and re-create all tables from provided metadata.

This is mainly used for tests and quick local setups.

Parameters:

Name	Type	Description	Default
`metadata`		SQLAlchemy `MetaData` containing table definitions.	required

Raises:

Type	Description
`ValueError`	If the engine has not been initialized.

Source code in src/pyomop/engine_factory.py

async def init_models(self, metadata):
    """Drop and re-create all tables from provided metadata.

    This is mainly used for tests and quick local setups.

    Args:
        metadata: SQLAlchemy ``MetaData`` containing table definitions.

    Raises:
        ValueError: If the engine has not been initialized.
    """
    if self._engine is None:
        raise ValueError("Database engine is not initialized.")
    async with self._engine.begin() as conn:
        await conn.run_sync(metadata.drop_all)
        await conn.run_sync(metadata.create_all)

LLM-oriented SQLDatabase wrapper for OMOP CDM.

This module provides a thin wrapper around llama_index.core.SQLDatabase that is aware of the OMOP CDM metadata reflected from this package's SQLAlchemy models. It enables LLM-powered query components to reason about available tables, columns, and foreign keys using the OMOP metadata directly.

This file is import-safe even when the optional LLM extras are not installed; in that case, attempting to instantiate CDMDatabase will raise a clear ImportError directing you to install pyomop[llm].

`CDMDatabase` ¶

Bases: SQLDatabase

OMOP-aware SQLDatabase for LLM query engines.

This class adapts llama-index's SQLDatabase to use the OMOP CDM SQLAlchemy metadata bundled with this package, making it easy to expose concise schema information to LLM components.

Parameters:

Name	Type	Description	Default
`engine`	`Engine`	SQLAlchemy `Engine` connected to the OMOP database.	required
`schema`	`Optional[str]`	Optional database schema name.	`None`
`ignore_tables`	`Optional[List[str]]`	Tables to hide from the LLM context.	`None`
`include_tables`	`Optional[List[str]]`	Explicit subset of tables to expose.	`None`
`sample_rows_in_table_info`	`int`	Kept for API parity (unused here).	`3`
`indexes_in_table_info`	`bool`	Kept for API parity (unused here).	`False`
`custom_table_info`	`Optional[dict]`	Optional overrides for table descriptions.	`None`
`view_support`	`bool`	Whether to reflect views as well (unused here).	`False`
`max_string_length`	`int`	Max length of generated descriptions.	`300`
`version`	`str`	OMOP CDM version label ("cdm54" or "cdm6").	`'cdm54'`

Source code in src/pyomop/llm_engine.py

class CDMDatabase(SQLDatabase):
    """OMOP-aware SQLDatabase for LLM query engines.

    This class adapts llama-index's ``SQLDatabase`` to use the OMOP CDM
    SQLAlchemy metadata bundled with this package, making it easy to expose
    concise schema information to LLM components.

    Args:
        engine: SQLAlchemy ``Engine`` connected to the OMOP database.
        schema: Optional database schema name.
        ignore_tables: Tables to hide from the LLM context.
        include_tables: Explicit subset of tables to expose.
        sample_rows_in_table_info: Kept for API parity (unused here).
        indexes_in_table_info: Kept for API parity (unused here).
        custom_table_info: Optional overrides for table descriptions.
        view_support: Whether to reflect views as well (unused here).
        max_string_length: Max length of generated descriptions.
        version: OMOP CDM version label ("cdm54" or "cdm6").
    """

    def __init__(
        self,
        engine: Engine,
        schema: Optional[str] = None,
        ignore_tables: Optional[List[str]] = None,
        include_tables: Optional[List[str]] = None,
        sample_rows_in_table_info: int = 3,
        indexes_in_table_info: bool = False,
        custom_table_info: Optional[dict] = None,
        view_support: bool = False,
        max_string_length: int = 300,
        version: str = "cdm54",
    ) -> None:
        if not _LLM_AVAILABLE:  # pragma: no cover - import-safe guard
            raise ImportError("Install 'pyomop[llm]' to use LLM features.")

        # Basic configuration
        self._engine = engine
        self._schema = schema

        if include_tables and ignore_tables:
            raise ValueError("Cannot specify both include_tables and ignore_tables")

        # Load OMOP metadata for the chosen version
        if version == "cdm6":
            from .cdm6 import Base
        else:
            from .cdm54 import Base
        metadata: MetaData = Base.metadata

        # All known tables (no view reflection here)
        self._all_tables = set(metadata.tables.keys())

        # Validate include/ignore lists
        self._include_tables = set(include_tables) if include_tables else set()
        if self._include_tables:
            missing = self._include_tables - self._all_tables
            if missing:
                raise ValueError(f"include_tables {missing} not found in OMOP metadata")

        self._ignore_tables = set(ignore_tables) if ignore_tables else set()
        if self._ignore_tables:
            missing = self._ignore_tables - self._all_tables
            if missing:
                raise ValueError(f"ignore_tables {missing} not found in OMOP metadata")

        if self._include_tables:
            usable = set(self._include_tables)
        elif self._ignore_tables:
            usable = self._all_tables - self._ignore_tables
        else:
            usable = set(self._all_tables)
        self._usable_tables = usable

        if not isinstance(sample_rows_in_table_info, int):
            raise TypeError("sample_rows_in_table_info must be an integer")
        self._sample_rows_in_table_info = sample_rows_in_table_info
        self._indexes_in_table_info = indexes_in_table_info

        # Optional custom descriptions
        self._custom_table_info = custom_table_info
        if self._custom_table_info is not None:
            if not isinstance(self._custom_table_info, dict):
                raise TypeError(
                    "custom_table_info must be a dict of {table_name: description}"
                )
            self._custom_table_info = {
                t: info
                for t, info in self._custom_table_info.items()
                if t in self._all_tables
            }

        self._max_string_length = max_string_length
        self._metadata = metadata

        # Initialize parent so llama-index internals are configured too.
        # llama-index expects a synchronous SQLAlchemy Engine. If an AsyncEngine
        # is provided, create a synchronous engine from the same URL.
        parent_engine: Engine
        if AsyncEngine is not None and isinstance(self._engine, AsyncEngine):
            url_str = str(self._engine.url)
            # Convert common async driver URLs to sync variants.
            # This is conservative and primarily targets SQLite used in tests.
            url_str = (
                url_str.replace("+aiosqlite", "")
                .replace("+asyncpg", "")
                .replace("+psycopg_async", "+psycopg2")
            )
            parent_engine = create_engine(url_str)
        else:
            parent_engine = self._engine  # type: ignore[assignment]

        super().__init__(
            engine=parent_engine,
            schema=schema,
            include_tables=sorted(self._usable_tables) if self._usable_tables else None,
        )

    # --- llama-index compatibility helpers (use OMOP metadata directly) ---
    def get_table_columns(self, table_name: str) -> List[str]:
        """Return list of column names for a table.

        This uses the OMOP SQLAlchemy ``MetaData`` instead of DB inspector.
        """

        return [col.name for col in self._metadata.tables[table_name].columns]

    def get_single_table_info(self, table_name: str) -> str:
        """Return a concise description of columns and foreign keys for a table.

        The format matches what llama-index expects when building table context.
        """

        template = "Table '{table_name}' has columns: {columns}, and foreign keys: {foreign_keys}."
        columns: List[str] = []
        foreign_keys: List[str] = []
        for column in self._metadata.tables[table_name].columns:
            columns.append(f"{column.name} ({column.type!s})")
            for fk in column.foreign_keys:
                foreign_keys.append(
                    f"{column.name} -> {fk.column.table.name}.{fk.column.name}"
                )
        column_str = ", ".join(columns)
        fk_str = ", ".join(foreign_keys)
        return template.format(
            table_name=table_name, columns=column_str, foreign_keys=fk_str
        )

    def usable_tables(self) -> List[str]:
        """Return the sorted list of tables exposed to the LLM.

        This respects include/ignore settings passed at initialization.
        """

        return sorted(self._usable_tables)

    # Backwards/compat helper name used in some code paths
    def get_usable_table_names(self) -> List[str]:  # pragma: no cover - thin wrapper
        return self.usable_tables()

`get_single_table_info(table_name)` ¶

Return a concise description of columns and foreign keys for a table.

The format matches what llama-index expects when building table context.

Source code in src/pyomop/llm_engine.py

def get_single_table_info(self, table_name: str) -> str:
    """Return a concise description of columns and foreign keys for a table.

    The format matches what llama-index expects when building table context.
    """

    template = "Table '{table_name}' has columns: {columns}, and foreign keys: {foreign_keys}."
    columns: List[str] = []
    foreign_keys: List[str] = []
    for column in self._metadata.tables[table_name].columns:
        columns.append(f"{column.name} ({column.type!s})")
        for fk in column.foreign_keys:
            foreign_keys.append(
                f"{column.name} -> {fk.column.table.name}.{fk.column.name}"
            )
    column_str = ", ".join(columns)
    fk_str = ", ".join(foreign_keys)
    return template.format(
        table_name=table_name, columns=column_str, foreign_keys=fk_str
    )

`get_table_columns(table_name)` ¶

Return list of column names for a table.

This uses the OMOP SQLAlchemy MetaData instead of DB inspector.

Source code in src/pyomop/llm_engine.py

def get_table_columns(self, table_name: str) -> List[str]:
    """Return list of column names for a table.

    This uses the OMOP SQLAlchemy ``MetaData`` instead of DB inspector.
    """

    return [col.name for col in self._metadata.tables[table_name].columns]

`usable_tables()` ¶

Return the sorted list of tables exposed to the LLM.

This respects include/ignore settings passed at initialization.

Source code in src/pyomop/llm_engine.py

def usable_tables(self) -> List[str]:
    """Return the sorted list of tables exposed to the LLM.

    This respects include/ignore settings passed at initialization.
    """

    return sorted(self._usable_tables)

LLM query utilities over the OMOP CDM schema.

This module wires llama-index components to an OMOP-aware CDMDatabase so you can build semantic and SQL-first query engines that know about your CDM tables. All LLM-related imports are optional and performed lazily at runtime.

`CdmLLMQuery` ¶

Helper that prepares an LLM-backed SQL query engine for OMOP.

It constructs an object index of selected CDM tables and exposes a retriever-backed query engine that can generate SQL or run SQL-only queries depending on configuration.

Parameters:

Name	Type	Description	Default
`sql_database`	`CDMDatabase`	A `CDMDatabase` instance connected to the OMOP DB.	required
`llm`	`Any`	Optional LLM implementation to plug into llama-index settings.	`None`
`similarity_top_k`	`int`	Top-k tables to retrieve for each query.	`1`
`embed_model`	`str`	HuggingFace embedding model name.	`'BAAI/bge-small-en-v1.5'`
`**kwargs`	`Any`	Reserved for future expansion.	`{}`

Source code in src/pyomop/llm_query.py

class CdmLLMQuery:
    """Helper that prepares an LLM-backed SQL query engine for OMOP.

    It constructs an object index of selected CDM tables and exposes a
    retriever-backed query engine that can generate SQL or run SQL-only queries
    depending on configuration.

    Args:
        sql_database: A ``CDMDatabase`` instance connected to the OMOP DB.
        llm: Optional LLM implementation to plug into llama-index settings.
        similarity_top_k: Top-k tables to retrieve for each query.
        embed_model: HuggingFace embedding model name.
        **kwargs: Reserved for future expansion.
    """

    def __init__(
        self,
        sql_database: CDMDatabase,
        llm: Any = None,  # FIXME: type
        similarity_top_k: int = 1,
        embed_model: str = "BAAI/bge-small-en-v1.5",
        **kwargs: Any,
    ):
        # Lazy import optional dependencies so the package imports without them
        try:
            sql_query_mod = importlib.import_module(
                "llama_index.core.indices.struct_store.sql_query"
            )
            objects_mod = importlib.import_module("llama_index.core.objects")
            core_mod = importlib.import_module("llama_index.core")
            hf_mod = importlib.import_module("langchain_huggingface")

            SQLTableRetrieverQueryEngine = getattr(
                sql_query_mod, "SQLTableRetrieverQueryEngine"
            )
            SQLTableNodeMapping = getattr(objects_mod, "SQLTableNodeMapping")
            ObjectIndex = getattr(objects_mod, "ObjectIndex")
            SQLTableSchema = getattr(objects_mod, "SQLTableSchema")
            VectorStoreIndex = getattr(core_mod, "VectorStoreIndex")
            Settings = getattr(core_mod, "Settings")
            HuggingFaceEmbeddings = getattr(hf_mod, "HuggingFaceEmbeddings")
        except Exception as e:  # pragma: no cover
            raise ImportError("Install 'pyomop[llm]' to use LLM query features.") from e
        self._sql_database = sql_database
        self._similarity_top_k = similarity_top_k
        self._embed_model = HuggingFaceEmbeddings(model_name=embed_model)
        self._llm = llm
        Settings.llm = llm
        Settings.embed_model = self._embed_model
        self._table_node_mapping = SQLTableNodeMapping(sql_database)
        usable_tables = []
        if hasattr(sql_database, "usable_tables"):
            usable_tables = list(sql_database.usable_tables())  # type: ignore[attr-defined]
        elif hasattr(sql_database, "get_usable_table_names"):
            usable_tables = list(sql_database.get_usable_table_names())  # type: ignore[attr-defined]
        self._table_schema_objs = [
            SQLTableSchema(table_name=t) for t in sorted(set(usable_tables))
        ]

        self._object_index = ObjectIndex.from_objects(
            self._table_schema_objs,
            self._table_node_mapping,
            VectorStoreIndex,  # type: ignore
        )

        self._query_engine = SQLTableRetrieverQueryEngine(
            self._sql_database,
            self._object_index.as_retriever(similarity_top_k=1),
            sql_only=True,
        )

    @property
    def table_node_mapping(self) -> Any:
        """Mapping between tables and nodes used by the object index."""
        return self._table_node_mapping

    @property
    def table_schema_objs(self) -> list[Any]:
        """List of table schema objects indexed for retrieval."""
        return self._table_schema_objs

    @property
    def object_index(self) -> Any:
        """The underlying llama-index object index used for retrieval."""
        return self._object_index

    @property
    def query_engine(self) -> Any:
        """A retriever-backed SQL query engine over the CDM tables."""
        return self._query_engine

`object_index` `property` ¶

The underlying llama-index object index used for retrieval.

`query_engine` `property` ¶

A retriever-backed SQL query engine over the CDM tables.

`table_node_mapping` `property` ¶

Mapping between tables and nodes used by the object index.

`table_schema_objs` `property` ¶

List of table schema objects indexed for retrieval.

CSV-to-OMOP loader.

This module implements a flexible CSV loader that can populate multiple OMOP CDM tables according to a JSON mapping file. It also performs helpful cleanup operations like foreign key normalization, birthdate backfilling, gender mapping, and concept code lookups.

`CdmCsvLoader` ¶

Load a single CSV into multiple OMOP CDM tables using a JSON mapping file.

Mapping file format (JSON):

{ "csv_key": "patient_id", # optional, CSV column that contains the patient/person identifier "tables": [ { "name": "cohort", # target table name as in the database "filters": [ # optional row filters applied to CSV before mapping {"column": "resourceType", "equals": "Encounter"} ], "columns": { # mapping of target_table_column -> value "cohort_definition_id": {"const": 1}, # constant value "subject_id": "patient_id", # copy from CSV column "cohort_start_date": "period.start", # copy from CSV column "cohort_end_date": "period.end" # copy from CSV column } } ] }

Notes

Constants are provided via {"const": value}.
If a required column is missing from mapping, it's left as None (DB default or nullable required).
Primary keys that are Integer types will autoincrement where supported (SQLite/PostgreSQL typical behavior).
Dates/times are converted to proper Python types where possible based on reflected column types.