Fixes Docker setup and sessions

Author: vmesel

.devcontainer/Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.13-slim
+FROM python:3.12-slim
 
 ENV PIP_DEFAULT_TIMEOUT=100
 ENV PIP_DISABLE_PIP_VERSION_CHECK=on

Dockerfile

@@ -12,27 +12,24 @@ ENV PYTHONFAULTHANDLER=1
 ENV PYTHONHASHSEED=random
 ENV PYTHONUNBUFFERED=1
 
-USER talkd
-WORKDIR /app
+RUN apt update -y && \
+    apt install -y gcc libpq-dev && \
+    apt-get autoremove -y && \
+    apt-get clean -y
+
+RUN pip install poetry
 
+WORKDIR /app
 COPY poetry.lock pyproject.toml README.md /app/
 COPY pytest.ini /app/src/
 
-USER root
-RUN apt update -y && apt upgrade -y && apt install -y gcc libpq-dev
-RUN pip install -U pip poetry
-
 RUN  poetry config virtualenvs.create false && \
     poetry install --only main
 
-USER talkd
-
 COPY /static /app/static
 COPY /etc /app/etc
 COPY /src /app/src
 
-USER root
-
 RUN chmod +x /app/etc/run.sh
 RUN chmod +x /app/etc/run-tests.sh
 

Dockerfile.test

@@ -0,0 +1,38 @@
+FROM python:3.12.3-slim
+LABEL org.opencontainers.image.source https://github.com/talkdai/dialog
+LABEL org.opencontainers.image.licenses MIT
+
+RUN useradd --user-group --system --create-home --no-log-init talkd
+RUN export PATH="/home/talkd/.local/bin:$PATH"
+
+ENV PIP_DEFAULT_TIMEOUT=100
+ENV PIP_DISABLE_PIP_VERSION_CHECK=on
+ENV PIP_NO_CACHE_DIR=on
+ENV PYTHONFAULTHANDLER=1
+ENV PYTHONHASHSEED=random
+ENV PYTHONUNBUFFERED=1
+
+RUN apt update -y && \
+    apt install -y gcc libpq-dev && \
+    apt-get autoremove -y && \
+    apt-get clean -y
+
+RUN pip install poetry
+
+WORKDIR /app
+COPY poetry.lock pyproject.toml README.md /app/
+COPY pytest.ini /app/src/
+
+RUN  poetry config virtualenvs.create false && \
+    poetry install
+
+COPY /static /app/static
+COPY /etc /app/etc
+COPY /src /app/src
+
+RUN chmod +x /app/etc/run.sh
+RUN chmod +x /app/etc/run-tests.sh
+
+WORKDIR /app/src
+
+CMD [ "/app/etc/run.sh" ]

Makefile

@@ -1,9 +1,6 @@
 lint:
 	poetry run black .
 
-db-up:
-	docker compose up db
-
 run:
 	poetry run uvicorn --app-dir src main:app --reload --host 0.0.0.0 --port 8000 --lifespan on --env-file .env
 
@@ -15,4 +12,13 @@ test-build:
 	docker compose -f docker-compose.test.yml run --rm --build dialog
 
 test:
-	docker compose -f docker-compose.test.yml run --rm dialog
+	docker compose -f docker-compose.test.yml run --rm dialog
+
+run-build:
+	docker compose up --build
+
+run-db:
+	docker compose up db
+
+run-ui:
+	docker compose -f docker-compose-open-webui.yml up

docker-compose.dev-container.yml

@@ -1,17 +1,17 @@
-version: '3.3'
+version: "3.3"
 services:
   db:
     image: pgvector/pgvector:pg15
     restart: always
     ports:
-      - '5432:5432'
+      - "5432:5432"
     environment:
       POSTGRES_USER: talkdai
       POSTGRES_PASSWORD: talkdai
       POSTGRES_DB: talkdai
     volumes:
-       - ./data/db:/var/lib/postgresql/data
-       - ./psql/db-ext-vector.sql:/docker-entrypoint-initdb.d/db-ext-vector.sql
+      - ./data/db:/var/lib/postgresql/data
+      - ./psql/db-ext-vector.sql:/docker-entrypoint-initdb.d/db-ext-vector.sql
     healthcheck:
       test: ["CMD", "pg_isready", "-d", "talkdai", "-U", "talkdai"]
       interval: 10s
@@ -23,13 +23,14 @@ services:
       dockerfile: Dockerfile
     stdin_open: true
     tty: true
+    entrypoint: ["/bin/bash"]
     volumes:
       - ./:/app
       - ./data/:/app/data/
       - ./.empty:/app/data/db
       - ./static:/app/static
     ports:
-      - '8000:8000'
+      - "8000:8000"
     depends_on:
       db:
         condition: service_healthy
@@ -42,7 +43,7 @@ services:
   openwebui:
     image: ghcr.io/open-webui/open-webui:main
     ports:
-      - '3000:8080'
+      - "3000:8080"
     environment:
       - OPENAI_API_KEYS=FAKE-KEY;
       - OPENAI_API_BASE_URLS=http://dialog:8000/openai;
@@ -56,4 +57,4 @@ services:
         condition: service_started
 
 volumes:
-  open-webui:
+  open-webui:

etc/run-tests.sh

@@ -1,7 +1,5 @@
 #!/bin/bash
 set -e
 
-cd /app/src/
-poetry install --only dev
 poetry run alembic upgrade head
-exec poetry run pytest --cov=. --cov-config=.coveragerc . -vvv
+exec poetry run pytest --cov=. --cov-config=.coveragerc . -vvv --timeout 10

etc/run.sh

@@ -1,9 +1,9 @@
 #!/bin/bash
 
-alembic upgrade head
+poetry run alembic upgrade head
 [[ -z "${DIALOG_LOADCSV_CLEARDB}" ]] || CLEARDB_COMMAND=--cleardb
 [[ -z "${DIALOG_LOADCSV_EMBED_COLUMNS}" ]] || EMBED_COLUMNS="--embed-columns ${DIALOG_LOADCSV_EMBED_COLUMNS}"
-python load_csv.py --path ${DIALOG_DATA_PATH} ${CLEARDB_COMMAND} ${EMBED_COLUMNS}
+poetry run python load_csv.py --path ${DIALOG_DATA_PATH} ${CLEARDB_COMMAND} ${EMBED_COLUMNS}
 
 /app/etc/install-plugins.sh
 
@@ -13,7 +13,7 @@ if  [ -n "${TEST}" ]; then
 fi
 
 if [ -n "${DEBUG}" ]; then
-    exec uvicorn main:app --host 0.0.0.0 --port ${PORT} --reload
+    exec poetry run uvicorn main:app --host 0.0.0.0 --port ${PORT} --reload
 else
-    exec uvicorn main:app --host 0.0.0.0 --port ${PORT}
+    exec poetry run uvicorn main:app --host 0.0.0.0 --port ${PORT}
 fi

poetry.lock

@@ -37,6 +37,7 @@ responses = "^0.25.0"
 pytest-mock = "^3.12.0"
 pytest-alembic = "^0.11.0"
 pytest-cov = "^4.1.0"
+pytest-timeout = "^2.3.1"
 
 [build-system]
 requires = ["poetry-core"]

pyproject.toml

@@ -25,14 +25,14 @@ def __init__(self, *args, **kwargs):
 
     @property
     def memory(self) -> BaseChatMemory:
-        if self.session_id:
-            return generate_memory_instance(
-                session_id=self.session_id,
+        # Returns the memory instance, if a session_id is not provided, a new session_id is generated
+        # and the memory instance is created.
+        return generate_memory_instance(
+                session_id=self.session_id if self.session_id else str(uuid4()),
                 parent_session_id=self.parent_session_id,
                 dbsession=self.dbsession,
                 database_url=Settings().DATABASE_URL
             )
-        return None
 
     def generate_prompt(self, text):
         self.relevant_contents = get_most_relevant_contents_from_message(

src/dialog/llm/agents/default.py

@@ -9,7 +9,7 @@
 from dialog_lib.embeddings.generate import generate_embeddings
 from dialog.llm.embeddings import EMBEDDINGS_LLM
 from dialog_lib.db.models import CompanyContent
-from dialog.db import get_session
+from dialog.db import session_scope
 from dialog.settings import Settings
 
 import logging
@@ -21,7 +21,6 @@
 
 logger = logging.getLogger("make_embeddings")
 
-session = next(get_session())
 NECESSARY_COLS = ["category", "subcategory", "question", "content"]
 PK_METADATA_COLS = ["category", "subcategory", "question"]
 
@@ -33,7 +32,7 @@ def _get_csv_cols(path: str) -> List[str]:
         return reader.fieldnames
 
 
-def retrieve_docs_from_vectordb() -> List[Document]:
+def retrieve_docs_from_vectordb(session) -> List[Document]:
     """Retrieve all documents from the vector store."""
     company_contents: List[CompanyContent] = session.query(CompanyContent).all()
     return [
@@ -97,7 +96,7 @@ def load_csv_with_metadata(
 
 
 def load_csv_and_generate_embeddings(
-    path, cleardb=False, embed_columns: Optional[list[str]] = None
+    path, session, cleardb=False, embed_columns: Optional[list[str]] = None
 ):
     """
     Load the knowledge base CSV, get their embeddings and store them into the vector store.
@@ -121,7 +120,7 @@ def load_csv_and_generate_embeddings(
         session.commit()
 
     # Get existing docs
-    docs_in_db: List[Document] = retrieve_docs_from_vectordb()
+    docs_in_db: List[Document] = retrieve_docs_from_vectordb(session)
     logging.info(f"Existing docs: {len(docs_in_db)}")
     existing_pks: List[str] = [
         get_document_pk(doc, PK_METADATA_COLS) for doc in docs_in_db
@@ -160,6 +159,7 @@ def load_csv_and_generate_embeddings(
     parser.add_argument("--embed-columns", default="content")
     args = parser.parse_args()
 
-    load_csv_and_generate_embeddings(
-        args.path, args.cleardb, args.embed_columns.split(",")
-    )
+    with session_scope() as session:
+        load_csv_and_generate_embeddings(
+            args.path, session, args.cleardb, args.embed_columns.split(",")
+        )

src/load_csv.py

@@ -30,9 +30,11 @@ def dbsession(mocker):
 
     with Session() as session:
         yield session
+        session.rollback()
 
     Base.metadata.drop_all(bind=engine)
 
+
 @pytest.fixture
 def client(dbsession):
     def get_session_override():