feat: SQL model for collections

app/collections/router.py

@@ -0,0 +1,31 @@
+from typing import Annotated, List
+from fastapi import APIRouter, Path
+from sqlmodel import Session, select
+
+from app.common.schema import Collection, DbDep
+
+router = APIRouter(tags=["collections"], prefix="/collections")
+
+@router.put("/")
+async def add(db: DbDep, collection: Collection) -> Collection:
+    """Create a collection."""
+    with Session(db) as session:
+        session.add(collection)
+        session.commit()
+        session.refresh(collection)
+        return collection
+
+@router.get("/")
+async def list(db: DbDep) -> List[Collection]:
+    """List collections."""
+    with Session(db) as session:
+        collections = session.exec(select(Collection)).all()
+        return collections
+
+PathCollectionId = Annotated[int, Path(..., description="The collection ID.")]
+
+@router.get("/{id}")
+async def get(id: PathCollectionId, db: DbDep) -> Collection:
+    """Get a specific collection."""
+    with Session(db) as session:
+        return session.get(Collection, id)

app/common/schema.py

@@ -0,0 +1,27 @@
+from typing import Annotated, Optional
+from fastapi import Depends, Request
+from sqlalchemy import Engine, UniqueConstraint
+from sqlmodel import Field, SQLModel
+
+class Collection(SQLModel, table=True):
+    id: Optional[int] = Field(default=None, primary_key=True)
+
+    # TODO: We may want this to be unique per-tenant rather than globally unique names.
+    name: str = Field(index=True, unique=True)
+
+class Document(SQLModel, table=True):
+    """Schema for documents in the SQL DB.
+    """
+    __table_args__ = (UniqueConstraint("collection_id", "url"),
+                      UniqueConstraint("collection_id", "doc_id"))
+
+    id: Optional[int] = Field(default=None, primary_key=True)
+    collection_id: int = Field(foreign_key="collection.id")
+
+    url: str = Field(index=True)
+    doc_id: Optional[str] = Field(default = None)
+
+def _db(request: Request) -> Engine:
+    return request.state.engine
+
+DbDep = Annotated[Engine, Depends(_db)]

app/config.py

@@ -19,6 +19,9 @@ class Config:
         env_file = ".env"
         env_file_encoding = "utf-8"
 
+    DB: str = "sqlite:///database.db?check_same_thread=false"
+    """The database to connect to."""
+
     ENVIRONMENT: Environment = Environment.PRODUCTION
     """The environment the application is running in."""
 
@@ -114,7 +117,7 @@ def custom_generate_unique_id_function(route: APIRoute) -> str:
     from a variety of sources -- documents, web pages, audio, etc.
     """,
     "servers": [
-        {"url": "http://127.0.0.1:8000", "description": "Local server"},
+        {"url": "http://localhost:8000", "description": "Local server"},
     ],
     "openapi_tags": [
         {

app/documents/router.py

@@ -16,12 +16,7 @@ async def add(
     store: StoreDep,
     url: Annotated[str, Body(..., description="The URL of the document to add.")],
 ):
-    """Add a document to the unstructured collection.
-
-    Parameters:
-    - collection: The ID of the collection to add to.
-    - document: The URL of the document to add.
-    """
+    """Add a document."""
 
     # Load the content.
     logger.debug("Loading content from {}", url)

app/main.py

@@ -1,23 +1,31 @@
 import contextlib
-from typing import AsyncIterator, TypedDict
+from typing import Annotated, AsyncIterator, TypedDict
 
-from fastapi import FastAPI
+from fastapi import Depends, FastAPI, Request
 from fastapi.routing import APIRoute
 from llama_index import StorageContext
+from sqlalchemy import Engine
+from sqlmodel import SQLModel, create_engine
 
-from app.config import app_configs
+from app.config import app_configs, settings
 from app.ingest.store import Store
 from app.routes import api_router
 
 
 class State(TypedDict):
-    storage_context: StorageContext
-
+    store: Store
+    db: Engine
 
 @contextlib.asynccontextmanager
 async def lifespan(_app: FastAPI) -> AsyncIterator[State]:
     """Function creating instances used during the lifespan of the service."""
-    state = {"store": Store()}
+    engine = create_engine(settings.DB, echo=True)
+    SQLModel.metadata.create_all(engine)
+
+    state = {
+        "store": Store(),
+        "engine": engine,
+    }
 
     yield state
 

app/routes.py

@@ -1,9 +1,11 @@
 from fastapi import APIRouter
 
 from app.chunks.router import router as chunks_router
+from app.collections.router import router as collections_router
 from app.documents.router import router as documents_router
 
 api_router = APIRouter(prefix="/api")
 
+api_router.include_router(collections_router)
 api_router.include_router(documents_router)
 api_router.include_router(chunks_router)

pyproject.toml

@@ -18,6 +18,7 @@ loguru = "^0.7.2"
 redis = "^5.0.1"
 accelerate = "^0.26.1"
 safetensors = "^0.4.1"
+sqlmodel = "^0.0.14"
 
 [tool.poetry.group.dev.dependencies]
 ruff = "^0.1.11"