modaic 0.1.0__py3-none-any.whl

modaic/databases/__init__.py

@@ -0,0 +1,35 @@
+from .graph_database import GraphDatabase, MemgraphConfig, Neo4jConfig
+from .sql_database import SQLDatabase, SQLiteBackend
+from .vector_database.vector_database import (
+    CollectionConfig,
+    IndexConfig,
+    IndexType,
+    Metric,
+    SearchResult,
+    SupportsHybridSearch,
+    VDBExtensions,
+    VectorDatabase,
+    VectorDBBackend,
+    VectorType,
+)
+from .vector_database.vendors.milvus import MilvusBackend
+
+__all__ = [
+    "CollectionConfig",
+    "SQLDatabase",
+    "SQLiteBackend",
+    "VectorDatabase",
+    "MilvusBackend",
+    "SearchResult",
+    "VectorDBBackend",
+    "IndexConfig",
+    "IndexType",
+    "Metric",
+    "SupportsHybridSearch",
+    "VDBExtensions",
+    "VectorDBBackend",
+    "VectorType",
+    "GraphDatabase",
+    "MemgraphConfig",
+    "Neo4jConfig",
+]

modaic/databases/graph_database.py

@@ -0,0 +1,269 @@
+import os
+from dataclasses import asdict, dataclass
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Protocol,
+    Type,
+)
+
+from dotenv import load_dotenv
+
+from ..context.base import Context, Relation
+from ..observability import Trackable, track_modaic_obj
+
+load_dotenv()
+
+
+if TYPE_CHECKING:
+    import gqlalchemy
+
+
+class GraphDBConfig(Protocol):
+    _client_class: ClassVar[Type["gqlalchemy.DatabaseClient"]]
+    # CAVEAT: checking for this attribute is currently the most reliable way to ascertain that something is a dataclass
+    __dataclass_fields__: ClassVar[Dict[str, Any]]
+
+
+class GraphDatabase(Trackable):
+    """
+    A database that stores context objects and relationships between them in a graph database.
+    """
+
+    def __init__(self, config: GraphDBConfig, **kwargs):
+        Trackable.__init__(self, **kwargs)
+        self.config = config
+        if getattr(self.config, "_client_class", None) is None:
+            raise ImportError("gqlalchemy is not installed. Install 'gqlalchemy' to use GraphDatabase backends.")
+        self._client = self.config._client_class(**asdict(self.config))
+
+    @track_modaic_obj
+    def execute_and_fetch(self, query: str) -> List[Dict[str, Any]]:
+        return self._client.execute_and_fetch(query)
+
+    def execute(
+        self,
+        query: str,
+        parameters: Dict[str, Any] = None,
+        connection: Optional["gqlalchemy.Connection"] = None,
+    ) -> None:
+        self._client.execute(query, parameters or {}, connection)
+
+    def create_index(self, index: "gqlalchemy.Index") -> None:
+        self._client.create_index(index)
+
+    def drop_index(self, index: "gqlalchemy.Index") -> None:
+        self._client.drop_index(index)
+
+    def get_indexes(self) -> List["gqlalchemy.Index"]:
+        return self._client.get_indexes()
+
+    def ensure_indexes(self, indexes: List["gqlalchemy.Index"]) -> None:
+        self._client.ensure_indexes(indexes)
+
+    def drop_indexes(self) -> None:
+        self._client.drop_indexes()
+
+    def create_constraint(self, constraint: "gqlalchemy.Constraint") -> None:
+        self._client.create_constraint(constraint)
+
+    def drop_constraint(self, constraint: "gqlalchemy.Constraint") -> None:
+        self._client.drop_constraint(constraint)
+
+    def get_constraints(self) -> List["gqlalchemy.Constraint"]:
+        return self._client.get_constraints()
+
+    def get_exists_constraints(self) -> List["gqlalchemy.Constraint"]:
+        return self._client.get_exists_constraints()
+
+    def get_unique_constraints(self) -> List["gqlalchemy.Constraint"]:
+        return self._client.get_unique_constraints()
+
+    def ensure_constraints(self, constraints: List["gqlalchemy.Constraint"]) -> None:
+        self._client.ensure_constraints(constraints)
+
+    def drop_database(self) -> None:
+        self._client.drop_database()
+
+    def new_connection(self) -> "gqlalchemy.Connection":
+        return self._client.new_connection()
+
+    def get_variable_assume_one(self, query_result: Iterator[Dict[str, Any]], variable_name: str) -> Any:
+        return self._client.get_variable_assume_one(query_result, variable_name)
+
+    def create_node(self, node: Context) -> Optional[Context]:
+        """
+        <Danger>This method is not thread safe. We are actively working on a solution to make it thread safe.</Danger>
+        """
+        node = node.to_gqlalchemy()
+        created_node = self._client.create_node(node)
+        if created_node is not None:
+            return Context.from_gqlalchemy(created_node)
+
+    def save_node(self, node: Context) -> "gqlalchemy.Node":
+        """
+        <Danger>This method is not thread safe. We are actively working on a solution to make it thread safe.</Danger>
+        """
+        node = node.to_gqlalchemy(self)
+        result = self._client.save_node(node)
+        return result
+
+    def save_nodes(self, nodes: List[Context]) -> None:
+        """
+        <Danger>This method is not thread safe. We are actively working on a solution to make it thread safe.</Danger>
+        """
+        nodes = [node.to_gqlalchemy(self) for node in nodes]
+        self._client.save_nodes(nodes)
+
+    def save_node_with_id(self, node: Context) -> Optional["gqlalchemy.Node"]:
+        """
+        <Danger>This method is not thread safe. We are actively working on a solution to make it thread safe.</Danger>
+        """
+        node = node.to_gqlalchemy(self)
+        result = self._client.save_node_with_id(node)
+        return result
+
+    def load_node(self, node: Context) -> Optional["gqlalchemy.Node"]:
+        """
+        <Danger>This method is not thread safe. We are actively working on a solution to make it thread safe.</Danger>
+        """
+        node = node.to_gqlalchemy(self)
+        result = self._client.load_node(node)
+        return result
+
+    def load_node_with_all_properties(self, node: Context) -> Optional["gqlalchemy.Node"]:
+        """
+        <Danger>This method is not thread safe. We are actively working on a solution to make it thread safe.</Danger>
+        """
+        node = node.to_gqlalchemy(self)
+        result = self._client.load_node_with_all_properties(node)
+        return result
+
+    def load_node_with_id(self, node: Context) -> Optional["gqlalchemy.Node"]:
+        """
+        <Danger>This method is not thread safe. We are actively working on a solution to make it thread safe.</Danger>
+        """
+        node = node.to_gqlalchemy(self)
+        result = self._client.load_node_with_id(node)
+        return result
+
+    def load_relationship(self, relationship: Relation) -> Optional["gqlalchemy.Relationship"]:
+        """
+        <Danger>This method is not thread safe. We are actively working on a solution to make it thread safe.</Danger>
+        """
+        relationship = relationship.to_gqlalchemy(self)
+        result = self._client.load_relationship(relationship)
+        return result
+
+    def load_relationship_with_id(self, relationship: Relation) -> Optional["gqlalchemy.Relationship"]:
+        """
+        <Danger>This method is not thread safe. We are actively working on a solution to make it thread safe.</Danger>
+        """
+        relationship = relationship.to_gqlalchemy(self)
+        result = self._client.load_relationship_with_id(relationship)
+        return result
+
+    def load_relationship_with_start_node_id_and_end_node_id(
+        self, relationship: Relation
+    ) -> Optional["gqlalchemy.Relationship"]:
+        """
+        <Danger>This method is not thread safe. We are actively working on a solution to make it thread safe.</Danger>
+        """
+        relationship = relationship.to_gqlalchemy(self)
+        result = self._client.load_relationship_with_start_node_id_and_end_node_id(relationship)
+        return result
+
+    def save_relationship(self, relationship: Relation) -> Optional["gqlalchemy.Relationship"]:
+        """
+        <Danger>This method is not thread safe. We are actively working on a solution to make it thread safe.</Danger>
+        """
+        relationship = relationship.to_gqlalchemy(self)
+
+        result = self._client.save_relationship(relationship)
+        return result
+
+    def save_relationships(self, relationships: List[Relation]) -> None:
+        """
+        <Danger>This method is not thread safe. We are actively working on a solution to make it thread safe.</Danger>
+        """
+        relationships = [relationship.to_gqlalchemy(self) for relationship in relationships]
+        self._client.save_relationships(relationships)
+
+    def save_relationship_with_id(self, relationship: Relation) -> Optional["gqlalchemy.Relationship"]:
+        """
+        <Danger>This method is not thread safe. We are actively working on a solution to make it thread safe.</Danger>
+        """
+        relationship = relationship.to_gqlalchemy(self)
+        result = self._client.save_relationship_with_id(relationship)
+        return result
+
+    def create_relationship(self, relationship: Relation) -> Optional["gqlalchemy.Relationship"]:
+        """
+        <Danger>This method is not thread safe. We are actively working on a solution to make it thread safe.</Danger>
+        """
+        relationship = relationship.to_gqlalchemy(self)
+        result = self._client.create_relationship(relationship)
+        return result
+
+
+NEO4J_HOST = os.getenv("NEO4J_HOST", "localhost")
+NEO4J_PORT = int(os.getenv("NEO4J_PORT", "7687"))
+NEO4J_USERNAME = os.getenv("NEO4J_USERNAME", "neo4j")
+NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD", "test")
+NEO4J_ENCRYPTED = os.getenv("NEO4J_ENCRYPT", "false").lower() == "true"
+NEO4J_CLIENT_NAME = os.getenv("NEO4J_CLIENT_NAME", "neo4j")
+
+
+def load_neo4j() -> Optional[Type[Any]]:
+    try:
+        import gqlalchemy
+    except ImportError:
+        return None
+    return gqlalchemy.Neo4j
+
+
+def load_memgraph() -> Optional[Type[Any]]:
+    try:
+        import gqlalchemy
+    except ImportError:
+        return None
+    return gqlalchemy.Memgraph
+
+
+@dataclass
+class Neo4jConfig:
+    host: str = NEO4J_HOST
+    port: int = NEO4J_PORT
+    username: str = NEO4J_USERNAME
+    password: str = NEO4J_PASSWORD
+    encrypted: bool = (NEO4J_ENCRYPTED,)
+    client_name: str = (NEO4J_CLIENT_NAME,)
+
+    _client_class: ClassVar[Optional[Type[Any]]] = load_neo4j()
+
+
+MG_HOST = os.getenv("MG_HOST", "127.0.0.1")
+MG_PORT = int(os.getenv("MG_PORT", "7687"))
+MG_USERNAME = os.getenv("MG_USERNAME", "")
+MG_PASSWORD = os.getenv("MG_PASSWORD", "")
+MG_ENCRYPTED = os.getenv("MG_ENCRYPT", "false").lower() == "true"
+MG_CLIENT_NAME = os.getenv("MG_CLIENT_NAME", "GQLAlchemy")
+MG_LAZY = os.getenv("MG_LAZY", "false").lower() == "true"
+
+
+@dataclass
+class MemgraphConfig:
+    host: str = MG_HOST
+    port: int = MG_PORT
+    username: str = MG_USERNAME
+    password: str = MG_PASSWORD
+    encrypted: bool = MG_ENCRYPTED
+    client_name: str = MG_CLIENT_NAME
+    lazy: bool = MG_LAZY
+
+    _client_class: ClassVar[Optional[Type[Any]]] = load_memgraph()

modaic/databases/sql_database.py

@@ -0,0 +1,355 @@
+import json
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Any, Callable, Iterable, List, Literal, Optional, Tuple
+from urllib.parse import urlencode
+
+import pandas as pd
+from sqlalchemy import Column, CursorResult, MetaData, String, Text, create_engine, inspect, text
+from sqlalchemy import Table as SQLTable
+from sqlalchemy.dialects import sqlite
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy.sql.compiler import IdentifierPreparer
+from tqdm import tqdm
+
+from ..context.table import BaseTable, Table, TableFile
+from ..storage import FileStore
+
+
+@dataclass
+class SQLDatabaseBackend:
+    """
+    Base class for SQL database backends.
+    Each subclass must implement the `url` property.
+    """
+
+    @property
+    def url(self) -> str:
+        raise NotImplementedError("Subclasses must implement this method")
+
+
+@dataclass
+class SQLServerBackend(SQLDatabaseBackend):
+    """
+    Backend configuration for a SQL served over a port or remote connection. (MySQL, PostgreSQL, etc.)
+
+    Args:
+        user: The username to connect to the database.
+        password: The password to connect to the database.
+        host: The host of the database.
+        database: The name of the database.
+        port: The port of the database.
+    """
+
+    user: str
+    password: str
+    host: str
+    database: str
+    port: Optional[str] = None
+    dialect: str = "mysql"
+    driver: Optional[str] = None
+    query_params: Optional[dict] = None
+
+    @property
+    def url(self) -> str:
+        port = f":{self.port}" if self.port else ""
+        driver = f"+{self.driver}" if self.driver else ""
+        query = f"?{urlencode(self.query_params)}" if self.query_params else ""
+        return f"{self.dialect}{driver}://{self.user}:{self.password}@{self.host}{port}/{self.database}{query}"
+
+
+@dataclass
+class SQLiteBackend(SQLDatabaseBackend):
+    """
+    Backend configuration for a SQLite database.
+
+    Args:
+        db_path: Path to the SQLite database file.
+        in_memory: Whether to create an in-memory SQLite database.
+        query_params: Query parameters to pass to the database.
+    """
+
+    db_path: Optional[str] = None
+    in_memory: bool = False
+    query_params: Optional[dict] = None
+
+    @property
+    def url(self) -> str:
+        base = "sqlite:///:memory:" if self.in_memory else f"sqlite:///{self.db_path}"
+        query = f"?{urlencode(self.query_params)}" if self.query_params else ""
+        return f"{base}{query}"
+
+
+class SQLDatabase:
+    def __init__(
+        self,
+        backend: SQLDatabaseBackend | str,
+        engine_kwargs: dict = None,  # TODO: This may not be a smart idea, may want to enforce specific kwargs
+        session_kwargs: dict = None,  # TODO: This may not be a smart idea, may want to enforce specific kwargs
+    ):
+        self.url = backend.url if isinstance(backend, SQLDatabaseBackend) else backend
+        self.engine = create_engine(self.url, **(engine_kwargs or {}))
+        self.metadata = MetaData()
+        self.session = sessionmaker(bind=self.engine, **(session_kwargs or {}))
+        self.inspector = inspect(self.engine)
+        self.preparer = IdentifierPreparer(sqlite.dialect())
+
+        # Create metadata table to store table metadata
+        self.metadata_table = SQLTable(
+            "metadata",
+            self.metadata,
+            Column("table_name", String(255), primary_key=True),
+            Column("metadata_json", Text),
+        )
+        self.metadata.create_all(self.engine)
+        self.connection = None
+        self._in_transaction = False
+
+    def add_table(
+        self,
+        table: BaseTable,
+        if_exists: Literal["fail", "replace", "append"] = "replace",
+        schema: str = None,
+    ):
+        # TODO: support batch inserting for large dataframes
+        with self.connect() as connection:
+            # Use the connection for to_sql to respect transaction context
+            table._df.to_sql(table.name, connection, if_exists=if_exists, index=False)
+
+            # Remove existing metadata for this table if it exists
+            connection.execute(self.metadata_table.delete().where(self.metadata_table.c.table_name == table.name))
+
+            # Insert new metadata
+            connection.execute(
+                self.metadata_table.insert().values(
+                    table_name=table.name,
+                    metadata_json=json.dumps(table.metadata),
+                )
+            )
+            if self._should_commit():
+                connection.commit()
+
+    def add_tables(
+        self,
+        tables: Iterable[BaseTable],
+        if_exists: Literal["fail", "replace", "append"] = "replace",
+        schema: str = None,
+    ):
+        for table in tables:
+            self.add_table(table, if_exists, schema)
+
+    def drop_table(self, name: str, must_exist: bool = False):
+        """
+        Drop a table from the database and remove its metadata.
+
+        Args:
+            name: The name of the table to drop
+        """
+        if_exists = "IF EXISTS" if not must_exist else ""
+        safe_name = self.preparer.quote(name)
+        with self.connect() as connection:
+            command = text(f"DROP TABLE {if_exists} {safe_name}")
+            connection.execute(command)
+            # Also remove metadata for this table
+            connection.execute(self.metadata_table.delete().where(self.metadata_table.c.table_name == name))
+            if self._should_commit():
+                connection.commit()
+
+    def drop_tables(self, names: Iterable[str], must_exist: bool = False):
+        for name in names:
+            self.drop_table(name, must_exist)
+
+    def list_tables(self) -> List[str]:
+        """
+        List all tables currently in the database.
+
+        Returns:
+            List of table names in the database.
+        """
+        # Refresh the inspector to ensure we get current table list
+        self.inspector = inspect(self.engine)
+        return self.inspector.get_table_names()
+
+    def get_table(self, name: str) -> BaseTable:
+        df = pd.read_sql_table(name, self.engine)
+
+        return Table(df=df, name=name, metadata=self.get_table_metadata(name))
+
+    def get_table_schema(self, name: str) -> List[Column]:
+        """
+        Return column schema for a given table.
+
+        Args:
+            name: The name of the table to get schema for
+
+        Returns:
+            Column schema information for the table.
+        """
+        return self.inspector.get_columns(name)
+
+    def get_table_metadata(self, name: str) -> dict:
+        """
+        Get metadata for a specific table.
+
+        Args:
+            name: The name of the table to get metadata for
+
+        Returns:
+            Dictionary containing the table's metadata, or empty dict if not found.
+        """
+        with self.connect() as connection:
+            result = connection.execute(
+                self.metadata_table.select().where(self.metadata_table.c.table_name == name)
+            ).fetchone()
+
+        if result:
+            return json.loads(result.metadata_json)
+        return {}
+
+    def query(self, query: str) -> CursorResult:
+        with self.connect() as connection:
+            result = connection.execute(text(query))
+        return result
+
+    def fetchall(self, query: str) -> List[Tuple]:
+        result = self.query(query)
+        return result.fetchall()
+
+    def fetchone(self, query: str) -> Tuple:
+        result = self.query(query)
+        return result.fetchone()
+
+    @classmethod
+    def from_file_store(
+        cls,
+        file_store: FileStore,
+        backend: SQLDatabaseBackend,
+        folder: Optional[str] = None,
+        table_created_hook: Optional[Callable[[TableFile], Any]] = None,
+    ) -> "SQLDatabase":
+        # TODO: support batch inserting and parallel processing
+        """
+        Initializes a new SQLDatabase from a file store.
+
+        Args:
+            file_store: File store containing files to load
+            backend: SQL database backend
+            folder: Folder in the file store to load
+
+        Returns:
+            New SQLDatabase instance loaded with data from the file store.
+        """
+        # TODO: make sure the loaded sql database is empty if not raise error and tell user to use __init__ for an already existing database
+        instance = cls(backend)
+        instance.add_file_store(file_store, folder, table_created_hook)
+        return instance
+
+    def add_file_store(
+        self,
+        file_store: FileStore,
+        folder: Optional[str] = None,
+        table_created_hook: Optional[Callable[[TableFile], Any]] = None,
+    ):
+        with self.begin():
+            for key, _ in tqdm(file_store.items(folder), desc="Uploading files to SQL database"):
+                table = TableFile.from_file_store(key, file_store)
+                self.add_table(table, if_exists="fail")
+                if table_created_hook:
+                    table_created_hook(table)
+
+    @contextmanager
+    def connect(self):
+        """
+        Context manager for database connections.
+        Reuses existing connection if available, otherwise creates a temporary one.
+        """
+        connection_existed = self.connection is not None
+        if not connection_existed:
+            self.connection = self.engine.connect()
+
+        try:
+            yield self.connection
+        finally:
+            # Only close if we created the connection for this operation
+            if not connection_existed:
+                self.close()
+
+    def open_persistent_connection(self):
+        """
+        Opens a persistent connection that will be reused across operations.
+        Call close() to close the persistent connection.
+        """
+        if self.connection is None:
+            self.connection = self.engine.connect()
+
+    def close(self):
+        """
+        Closes the current connection if one exists.
+        """
+        if self.connection:
+            self.connection.close()
+            self.connection = None
+
+    def _should_commit(self) -> bool:
+        """
+        Returns True if operations should commit immediately.
+        Returns False if we're within an explicit transaction context.
+        """
+        return not self._in_transaction
+
+    @contextmanager
+    def begin(self):
+        """
+        Context manager for database transactions using existing connection.
+        Requires an active connection. Commits on success, rolls back on exception.
+
+        Raises:
+            RuntimeError: If no active connection exists
+        """
+        if self.connection is None:
+            raise RuntimeError("No active connection. Use connect_and_begin() or open a connection first.")
+
+        transaction = self.connection.begin()
+        old_in_transaction = self._in_transaction
+        self._in_transaction = True
+
+        try:
+            yield self.connection
+            transaction.commit()
+        except Exception:
+            transaction.rollback()
+            raise
+        finally:
+            self._in_transaction = old_in_transaction
+
+    @contextmanager
+    def connect_and_begin(self):
+        """
+        Context manager that establishes a connection and starts a transaction.
+        Reuses existing connection if available, otherwise creates a temporary one.
+        Commits on success, rolls back on exception.
+        """
+        connection_existed = self.connection is not None
+        if not connection_existed:
+            self.connection = self.engine.connect()
+
+        transaction = self.connection.begin()
+        old_in_transaction = self._in_transaction
+        self._in_transaction = True
+
+        try:
+            yield self.connection
+            transaction.commit()
+        except Exception:
+            transaction.rollback()
+            raise
+        finally:
+            self._in_transaction = old_in_transaction
+            # Only close if we created the connection for this operation
+            if not connection_existed:
+                self.close()
+
+
+class MultiTenantSQLDatabase:
+    def __init__(self):
+        raise NotImplementedError("Not implemented")

modaic/databases/vector_database/__init__.py

@@ -0,0 +1,12 @@
+from .vector_database import IndexConfig, IndexType, Metric, SupportsHybridSearch, VectorDatabase, VectorType
+from .vendors.milvus import MilvusBackend
+
+__all__ = [
+    "VectorDatabase",
+    "SupportsHybridSearch",
+    "MilvusBackend",
+    "IndexConfig",
+    "IndexType",
+    "VectorType",
+    "Metric",
+]