Feat ollama

.github/workflows/cicd.yml

@@ -24,6 +24,7 @@ jobs:
           - tests/test_mysql.py
           - tests/test_mongodb.py
           - tests/test_mssql.py
+          - tests/test_ollama.py
           - tests/test_neo4j.py
           - tests/test_redis.py
           - tests/test_opensearch.py
@@ -54,7 +55,7 @@ jobs:
       - name: Install Python dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install -e ".[test,graph,rag]"
+          python -m pip install -e ".[test,rag,graph]"
 
       - name: Register Jupyter kernel
         run: |

README.md

@@ -8,14 +8,14 @@
 
 `py-dockerdb` gives you easy Docker database setup in Python for PostgreSQL, MySQL, MongoDB, Microsoft SQL Server, Redis, and Neo4j. It is built for people who teach, demo, and prototype with notebooks or scripts and need repeatable local databases in minutes. Instead of writing Docker commands and per-engine setup code, you use one API to create, start, connect, and clean up containers.
 
-Switch from PostgreSQL to MongoDB and back without changing a line of connection code. This makes side-by-side database comparison a first-class workflow - useful for MVPs where the right engine isn't decided yet, and for RAG experiments where you want to test one storage backend, then swap it out without rewriting environment glue.
+Switch from PostgreSQL to MongoDB and back without changing a line of connection code. This makes side-by-side database comparison a first-class workflow - useful for MVPs where the right engine isn't decided yet, and for RAG and GraphRAG experiments where you want to test one storage backend, then swap it out without rewriting environment glue.
 
 If you teach SQL or data workflows, it removes the environment-setup section from your slides entirely: every student runs the same two lines and gets a working database.
 
 ## When to use this
 
 - **Teaching a SQL workshop or notebook tutorial** - every learner starts from the same environment with no per-machine Docker setup required.
-- **Comparing databases for an MVP** - run PostgreSQL, MySQL, MongoDB, MSSQL, and Redis under the same Python interface and switch engines without rewriting connection code.
+- **Comparing databases for an MVP** - run PostgreSQL, MySQL, MongoDB, MSSQL, Redis, and Neo4j under the same Python interface and switch engines without rewriting connection code.
 - **Building a local RAG prototype** - spin up a backing store, test your retrieval pipeline, then swap from PostgreSQL to MongoDB in one config change without touching orchestration code.
 - **GraphRAG with Neo4j** - provision a local knowledge graph for multi-hop retrieval pipelines. The `Neo4jDB.connection` property returns a `neo4j.Driver` that plugs directly into LlamaIndex's `Neo4jGraphStore` and LangChain's `Neo4jGraph`.
 
@@ -35,9 +35,10 @@ If you teach SQL or data workflows, it removes the environment-setup section fro
 - Database drivers (installed automatically with package dependencies):
   - `psycopg2-binary` for PostgreSQL
   - `mysql-connector-python` for MySQL
-  - `pymongo` for MongoDB
-  - `pyodbc` for MSSQL
-  - `redis` for Redis
+- `pymongo` for MongoDB
+- `pyodbc` for MSSQL
+- `redis` for Redis
+- `neo4j` for Neo4j
 
 ## Installation
 
@@ -51,7 +52,7 @@ pip install "py-dockerdb[graph]"
 
 ## Usage
 
-The API is consistent across all five engines: define a config, call `create_db()`, run your workload, then tear down with `delete_db()`.
+The API is consistent across all engines: define a config, call `create_db()`, run your workload, then tear down with `delete_db()`.
 
 ### PostgreSQL example
 
@@ -88,22 +89,32 @@ conn.close()
 db_manager.delete_db(running_ok=True)
 ```
 
-### Neo4j example (GraphRAG)
+### Neo4j / GraphRAG example
 
 ```python
+import uuid
+from pathlib import Path
 from docker_db.dbs.neo4j_db import Neo4jConfig, Neo4jDB
 
-config = Neo4jConfig(password="demopassword", project_name="demo")
+container_name = f"demo-neo4j-{uuid.uuid4().hex[:8]}"
+temp_dir = Path("tmp")
+temp_dir.mkdir(exist_ok=True)
+
+config = Neo4jConfig(
+    password="demopass",
+    project_name="demo",
+    container_name=container_name,
+    workdir=temp_dir.absolute(),
+)
+
 db_manager = Neo4jDB(config)
 db_manager.create_db()
 
-driver = db_manager.connection          # neo4j.Driver via Bolt
-# Neo4j Browser UI: http://localhost:7474
-
-with driver.session(database=config.database) as session:
-    session.run("CREATE (a:Person {name: 'Alice'})-[:KNOWS]->(b:Person {name: 'Bob'})")
-    result = session.run("MATCH (a)-[:KNOWS]->(b) RETURN b.name AS name")
-    print(result.single()["name"])      # Bob
+driver = db_manager.connection
+with driver.session() as session:
+    session.run("CREATE (n:Person {name: 'Alice'})")
+    result = session.run("MATCH (n:Person) RETURN n.name AS name")
+    print(result.single()["name"])  # Alice
 
 driver.close()
 db_manager.delete_db(running_ok=True)
@@ -118,7 +129,7 @@ Full runnable notebooks for each engine are in the [`usage/`](./usage/) director
 - [MongoDB](./usage/mongo_example.ipynb)
 - [MSSQL](./usage/mssql_example.ipynb)
 - [Redis](./usage/redis_example.ipynb)
-- [Neo4j + GraphRAG](./usage/neo4j_example.ipynb)
+- [Neo4j / GraphRAG](./usage/neo4j_example.ipynb)
 - [Container lifecycle and management](./usage/db_management_example.ipynb)
 
 ### Seeding data for demos and workshops
@@ -140,11 +151,10 @@ conn_string = db_manager.connection_string(sql_magic=True)
 
 ## Roadmap
 
-- [x] PostgreSQL + `pgvector` (vector store for RAG)
-- [x] Neo4j (GraphRAG, knowledge graph retrieval)
-- [ ] Redis Stack (semantic LLM cache, vector search)
-- [ ] Ollama (local LLM inference)
-- [ ] Elasticsearch / OpenSearch (hybrid BM25 + vector search)
+Docker-friendly vector and graph backends that can be run locally without paid licenses:
+
+- [x] PostgreSQL + `pgvector`
+- [x] Neo4j (knowledge graphs, GraphRAG)
 - [ ] Qdrant
 - [ ] Weaviate
 
@@ -165,6 +175,7 @@ python -m pytest -vv -s tests/test_postgres.py
 python -m pytest -vv -s tests/test_mysql.py
 python -m pytest -vv -s tests/test_mongodb.py
 python -m pytest -vv -s tests/test_mssql.py
+python -m pytest -vv -s tests/test_redis.py
 python -m pytest -vv -s tests/test_neo4j.py
 python -m pytest -vv -s tests/test_notebooks.py
 ```

docker_db/__init__.py

@@ -11,6 +11,8 @@
     RedisDB,
     MySQLConfig,
     MySQLDB,
+    OllamaConfig,
+    OllamaDB,
     Neo4jConfig,
     Neo4jDB,
     PostgresConfig,
@@ -27,6 +29,8 @@
     "PostgresConfig",
     "MySQLDB",
     "MySQLConfig",
+    "OllamaDB",
+    "OllamaConfig",
     "MSSQLDB",
     "MSSQLConfig",
     "Neo4jDB",

docker_db/dbs/__init__.py

@@ -4,6 +4,7 @@
 from docker_db.dbs.postgres_db import PostgresDB, PostgresConfig
 from docker_db.dbs.mysql_db import MySQLDB, MySQLConfig
 from docker_db.dbs.mssql_db import MSSQLDB, MSSQLConfig
+from docker_db.dbs.ollama_db import OllamaDB, OllamaConfig
 from docker_db.dbs.neo4j_db import Neo4jDB, Neo4jConfig
 from docker_db.dbs.redis_db import RedisDB, RedisConfig
 from docker_db.dbs.opensearch_db import OpenSearchDB, OpenSearchConfig
@@ -17,6 +18,8 @@
     "MySQLConfig",
     "MSSQLDB",
     "MSSQLConfig",
+    "OllamaDB",
+    "OllamaConfig",
     "Neo4jDB",
     "Neo4jConfig",
     "RedisDB",

docker_db/dbs/ollama_db.py

@@ -0,0 +1,193 @@
+"""
+Ollama container management module.
+
+This module provides classes to manage Ollama containers using Docker.
+It enables developers to easily create, configure, start, stop, and delete Ollama
+containers for development and testing purposes.
+"""
+import time
+
+import docker
+import requests
+from docker.models.containers import Container
+from pydantic import Field
+from requests.exceptions import RequestException
+
+from docker_db.docker import ContainerConfig, ContainerManager
+
+
+class OllamaConfig(ContainerConfig):
+    """Configuration for an Ollama container."""
+
+    database: str = Field(
+        default="ollama",
+        description="Compatibility field required by the base container manager API.",
+    )
+    model: str | None = Field(
+        default=None,
+        description="Optional model name to pull after container startup.",
+    )
+    pull_model_on_create: bool = Field(
+        default=False,
+        description="If True and model is set, pull the model in create_db().",
+    )
+    port: int = Field(default=11434, description="Port on which Ollama listens")
+    env_vars: dict = Field(
+        default_factory=dict,
+        description="A dictionary of environment variables to set in the container",
+    )
+    _type: str = "ollama"
+
+
+class OllamaDB(ContainerManager):
+    """Manages lifecycle of an Ollama container via Docker SDK."""
+
+    def __init__(self, config: OllamaConfig):
+        self.config: OllamaConfig = config
+        assert self._is_docker_running()
+        self.client = docker.from_env()
+
+    @property
+    def connection(self):
+        """
+        Establish a new HTTP session to the Ollama API.
+
+        Returns
+        -------
+        requests.Session
+            A new HTTP session object.
+        """
+        return requests.Session()
+
+    @property
+    def base_url(self) -> str:
+        """Return Ollama API base URL."""
+        return f"http://{self.config.host}:{self.config.port}"
+
+    def connection_string(self, db_name: str = None, sql_magic: bool = False) -> str:
+        """
+        Get Ollama base URL.
+
+        Parameters
+        ----------
+        db_name : str, optional
+            Unused placeholder for API compatibility with other managers.
+        sql_magic : bool, optional
+            Unused placeholder for API compatibility with SQL-oriented managers.
+
+        Returns
+        -------
+        str
+            Base URL of the Ollama API endpoint.
+        """
+        return self.base_url
+
+    def _get_environment_vars(self):
+        default_env_vars = {}
+        default_env_vars.update(self.config.env_vars)
+        return default_env_vars
+
+    def _get_volume_mounts(self):
+        return [
+            docker.types.Mount(
+                target="/root/.ollama",
+                source=str(self.config.volume_path),
+                type="bind",
+            )
+        ]
+
+    def _get_port_mappings(self):
+        return {"11434/tcp": self.config.port}
+
+    def _get_healthcheck(self):
+        return {
+            "Test": ["CMD", "ollama", "--version"],
+            "Interval": 5000000000,  # 5s
+            "Timeout": 3000000000,  # 3s
+            "Retries": 10,
+        }
+
+    def _wait_for_api_ready(self) -> bool:
+        for _ in range(self.config.retries):
+            try:
+                response = requests.get(
+                    f"{self.base_url}/api/tags",
+                    timeout=5,
+                )
+                if response.status_code == 200:
+                    return True
+            except RequestException:
+                pass
+            time.sleep(self.config.delay)
+        return False
+
+    def _create_db(
+        self,
+        db_name: str | None = None,
+        container: Container | None = None,
+    ):
+        """Validate Ollama API availability and optionally pull a model."""
+        container = container or self.client.containers.get(self.config.container_name)
+        container.reload()
+        if not container.attrs.get("State", {}).get("Running", False):
+            raise RuntimeError(f"Container {container.name} is not running.")
+
+        if not self._wait_for_api_ready():
+            raise RuntimeError("Ollama API did not become ready in time.")
+
+        if self.config.pull_model_on_create and self.config.model:
+            self.pull_model(self.config.model)
+
+        self.database_created = True
+        return container
+
+    def _wait_for_db(self, container: Container | None = None) -> bool:
+        """Wait until Ollama API is accepting connections and ready."""
+        try:
+            container = container or self.client.containers.get(self.config.container_name)
+            for _ in range(self.config.retries):
+                container.reload()
+                state = container.attrs.get("State", {})
+                if state.get("Running", False):
+                    break
+                time.sleep(self.config.delay)
+        except (docker.errors.NotFound, docker.errors.APIError):
+            pass
+
+        return self._wait_for_api_ready()
+
+    def pull_model(self, model: str):
+        """
+        Pull a model via Ollama API.
+
+        Parameters
+        ----------
+        model : str
+            Model name to download into the local Ollama store.
+
+        Returns
+        -------
+        dict
+            JSON response returned by Ollama pull endpoint.
+        """
+        response = requests.post(
+            f"{self.base_url}/api/pull",
+            json={"model": model, "stream": False},
+            timeout=600,
+        )
+        if response.status_code != 200:
+            raise RuntimeError(f"Failed to pull model '{model}': {response.text}")
+        return response.json()
+
+    def list_models(self) -> list[dict]:
+        """List locally available models."""
+        response = requests.get(f"{self.base_url}/api/tags", timeout=10)
+        if response.status_code != 200:
+            raise RuntimeError(f"Failed to list models: {response.text}")
+        payload = response.json()
+        return payload.get("models", [])
+
+    def test_connection(self):
+        """Ensure Ollama API is reachable."""
+        if not self._wait_for_api_ready():
+            raise ConnectionError("Ollama API is not reachable.")

docker_db/docker/model.py

@@ -12,9 +12,10 @@
     "mssql": "ms",
     "mongodb": "mg",
     "cassandra": "cs",
-    "neo4j": "n4j",
+    "ollama": "ol",
     "redis": "rd",
     "opensearch": "os",
+    "neo4j": "n4j",
 }
 
 DEFAULT_IMAGE_MAP = {
@@ -24,9 +25,10 @@
     "mssql": "mcr.microsoft.com/mssql/server:2022-latest",
     "mongodb": "mongo:6",
     "cassandra": "cassandra:4",
-    "neo4j": "neo4j:5",
+    "ollama": "ollama/ollama:latest",
     "redis": "redis:7",
     "opensearch": "opensearchproject/opensearch:2.13.0",
+    "neo4j": "neo4j:5",
 }