Merge pull request #188 from presenton/refactor/replace-redis-with-sqlite

refactor/replace redis with sqlite
This commit is contained in:
Saurav Niraula 2025-08-04 21:58:01 +05:45 committed by GitHub
commit 18aee8fc78
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 96 additions and 40 deletions

View file

@ -6,4 +6,6 @@ servers/fastapi/debug
servers/fastapi/.venv
servers/nextjs/node_modules
servers/nextjs/.next
servers/nextjs/.next
container.db

3
.gitignore vendored
View file

@ -14,4 +14,5 @@ debug
my-doc.txt
generated_models
nltk
chroma
chroma
container.db

View file

@ -3,8 +3,7 @@ FROM python:3.11-slim-bookworm
# Install Node.js and npm
RUN apt-get update && apt-get install -y \
nginx \
curl \
redis-server
curl
# Install Node.js 20 using NodeSource repository
RUN curl -fsSL https://deb.nodesource.com/setup_20.x | bash - && \
@ -25,7 +24,7 @@ RUN curl -fsSL https://ollama.com/install.sh | sh
# Install dependencies for FastAPI
RUN pip install aiohttp aiomysql aiosqlite asyncpg fastapi[standard] \
pathvalidate pdfplumber nltk chromadb sqlmodel redis \
pathvalidate pdfplumber nltk chromadb sqlmodel \
anthropic google-genai openai fastmcp
RUN pip install docling --extra-index-url https://download.pytorch.org/whl/cpu

View file

@ -3,8 +3,7 @@ FROM python:3.11-slim-bookworm
# Install Node.js and npm
RUN apt-get update && apt-get install -y \
nginx \
curl \
redis-server
curl
# Install Node.js 20 using NodeSource repository
@ -27,7 +26,7 @@ RUN curl -fsSL http://ollama.com/install.sh | sh
# Install dependencies for FastAPI
RUN pip install aiohttp aiomysql aiosqlite asyncpg fastapi[standard] \
pathvalidate pdfplumber nltk chromadb sqlmodel redis \
pathvalidate pdfplumber nltk chromadb sqlmodel \
anthropic google-genai openai fastmcp
RUN pip install docling --extra-index-url https://download.pytorch.org/whl/cpu

View file

@ -1,9 +1,11 @@
import json
from datetime import datetime
from fastapi import HTTPException
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from models.ollama_model_status import OllamaModelStatus
from services import REDIS_SERVICE
from models.sql.ollama_pull_status import OllamaPullStatus
from services.database import get_container_db_async_session
from utils.ollama import pull_ollama_model
@ -15,6 +17,8 @@ async def pull_ollama_model_background_task(model: str):
)
log_event_count = 0
session = await get_container_db_async_session().__anext__()
try:
async for event in pull_ollama_model(model):
log_event_count += 1
@ -30,18 +34,13 @@ async def pull_ollama_model_background_task(model: str):
if "status" in event:
saved_model_status.status = event["status"]
REDIS_SERVICE.set(
f"ollama_models/{model}",
json.dumps(saved_model_status.model_dump(mode="json")),
)
await upsert_ollama_pull_status(session, model, saved_model_status)
except Exception as e:
saved_model_status.status = "error"
saved_model_status.done = True
REDIS_SERVICE.set(
f"ollama_models/{model}",
json.dumps(saved_model_status.model_dump(mode="json")),
)
await upsert_ollama_pull_status(session, model, saved_model_status)
await session.close()
raise HTTPException(
status_code=500,
detail=f"Failed to pull model: {e}",
@ -51,9 +50,27 @@ async def pull_ollama_model_background_task(model: str):
saved_model_status.status = "pulled"
saved_model_status.downloaded = saved_model_status.size
REDIS_SERVICE.set(
f"ollama_models/{model}",
json.dumps(saved_model_status.model_dump(mode="json")),
)
await upsert_ollama_pull_status(session, model, saved_model_status)
await session.close()
return saved_model_status
async def upsert_ollama_pull_status(
session: AsyncSession, model: str, model_status: OllamaModelStatus
):
stmt = select(OllamaPullStatus).where(OllamaPullStatus.id == model)
result = await session.execute(stmt)
existing_record = result.scalar_one_or_none()
if existing_record:
existing_record.status = model_status.model_dump(mode="json")
existing_record.last_updated = datetime.now()
else:
new_record = OllamaPullStatus(
id=model,
status=model_status.model_dump(mode="json"),
last_updated=datetime.now(),
)
session.add(new_record)
await session.commit()
await session.flush()

View file

@ -1,12 +1,15 @@
from datetime import datetime, timedelta
import json
from typing import List
from fastapi import APIRouter, BackgroundTasks, HTTPException
from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException
from sqlalchemy.ext.asyncio import AsyncSession
from api.v1.ppt.background_tasks import pull_ollama_model_background_task
from constants.supported_ollama_models import SUPPORTED_OLLAMA_MODELS
from models.ollama_model_metadata import OllamaModelMetadata
from models.ollama_model_status import OllamaModelStatus
from services import REDIS_SERVICE
from models.sql.ollama_pull_status import OllamaPullStatus
from services.database import get_container_db_async_session
from utils.ollama import list_pulled_ollama_models
OLLAMA_ROUTER = APIRouter(prefix="/ollama", tags=["Ollama"])
@ -23,7 +26,11 @@ async def get_available_models():
@OLLAMA_ROUTER.get("/model/pull", response_model=OllamaModelStatus)
async def pull_model(model: str, background_tasks: BackgroundTasks):
async def pull_model(
model: str,
background_tasks: BackgroundTasks,
session: AsyncSession = Depends(get_container_db_async_session),
):
if model not in SUPPORTED_OLLAMA_MODELS:
raise HTTPException(
@ -46,21 +53,27 @@ async def pull_model(model: str, background_tasks: BackgroundTasks):
detail=f"Failed to check pulled models: {e}",
)
saved_model_status = REDIS_SERVICE.get(f"ollama_models/{model}")
saved_pull_status = None
saved_model_status = None
try:
saved_pull_status = await session.get(OllamaPullStatus, model)
saved_model_status = saved_pull_status.status
except Exception as e:
pass
# If the model is being pulled, return the model
if saved_model_status:
saved_model_status_json = json.loads(saved_model_status)
# If the model is being pulled, return the model
# ? If the model status is pulled in redis but was not found while listing pulled models,
# ? it means the model was deleted and we need to pull it again
if (
saved_model_status_json["status"] == "error"
or saved_model_status_json["status"] == "pulled"
saved_model_status["status"] == "error"
or saved_model_status["status"] == "pulled"
or saved_pull_status.last_updated < (datetime.now() - timedelta(seconds=10))
):
REDIS_SERVICE.delete(f"ollama_models/{model}")
await session.delete(saved_pull_status)
else:
return saved_model_status_json
return saved_model_status
# If the model is not being pulled, pull the model
background_tasks.add_task(pull_ollama_model_background_task, model)

View file

@ -2,7 +2,7 @@ from datetime import datetime
from typing import Optional
from sqlalchemy import JSON, Column, DateTime
from sqlmodel import SQLModel, Field
from sqlmodel import Field, SQLModel
from utils.randomizers import get_random_uuid

View file

@ -1,4 +1,4 @@
from sqlmodel import SQLModel, Field, Column, JSON
from sqlmodel import Field, Column, JSON, SQLModel
from utils.randomizers import get_random_uuid

View file

@ -0,0 +1,8 @@
from datetime import datetime
from sqlmodel import Field, Column, JSON, SQLModel, DateTime
class OllamaPullStatus(SQLModel, table=True):
id: str = Field(primary_key=True)
last_updated: datetime = Field(sa_column=Column(DateTime, default=datetime.now))
status: dict = Field(sa_column=Column(JSON))

View file

@ -1,7 +1,7 @@
from datetime import datetime
from typing import List, Optional
from sqlalchemy import JSON, Column, DateTime
from sqlmodel import SQLModel, Field
from sqlmodel import Field, SQLModel
from models.presentation_layout import PresentationLayoutModel
from models.presentation_outline_model import PresentationOutlineModel

View file

@ -1,5 +1,5 @@
from typing import Optional
from sqlmodel import SQLModel, Field, Column, JSON
from sqlmodel import Field, Column, JSON, SQLModel
from utils.randomizers import get_random_uuid

View file

@ -1,6 +1,4 @@
from services.redis_service import RedisService
from services.temp_file_service import TempFileService
TEMP_FILE_SERVICE = TempFileService()
REDIS_SERVICE = RedisService()

View file

@ -37,6 +37,25 @@ async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
yield session
# Container DB (Lives inside the container)
container_db_url = "sqlite+aiosqlite:////app/container.db"
container_db_engine: AsyncEngine = create_async_engine(
container_db_url, connect_args={"check_same_thread": False}
)
container_db_async_session_maker = async_sessionmaker(
container_db_engine, expire_on_commit=False
)
async def get_container_db_async_session() -> AsyncGenerator[AsyncSession, None]:
async with container_db_async_session_maker() as session:
yield session
# Create Database and Tables
async def create_db_and_tables():
async with sql_engine.begin() as conn:
await conn.run_sync(SQLModel.metadata.create_all)
async with container_db_engine.begin() as conn:
await conn.run_sync(SQLModel.metadata.create_all)

View file

@ -5,7 +5,7 @@ import nltk
from models.document_chunk import DocumentChunk
try:
nltk.data.find("tokenizers/punkt")
nltk.data.find("tokenizers/punkt", paths=["./nltk"])
except LookupError:
nltk.download("punkt", download_dir="./nltk")