diff --git a/Makefile b/Makefile index 5aaad2e2a..9f57f63f1 100644 --- a/Makefile +++ b/Makefile @@ -39,8 +39,8 @@ test-ingest: poetry run pytest ingest/tests --cov=ingest -v --cov-report=term-missing --cov-fail-under=40 test-django: - docker-compose up -d db - docker-compose run django-app poetry run pytest django_app/tests/ --ds redbox_app.settings -v --cov=redbox_app.redbox_core --cov-fail-under 10 + docker compose up -d --wait db + docker compose run django-app poetry run pytest django_app/tests/ --ds redbox_app.settings -v --cov=redbox_app.redbox_core --cov-fail-under 10 lint: poetry run ruff check . @@ -61,10 +61,10 @@ checktypes: # poetry run mypy legacy_app --follow-imports skip --ignore-missing-imports check-migrations: - docker-compose build django-app - docker-compose run django-app poetry run python django_app/manage.py migrate - docker-compose run django-app poetry run python django_app/manage.py makemigrations --check + docker compose build django-app + docker compose run django-app poetry run python django_app/manage.py migrate + docker compose run django-app poetry run python django_app/manage.py makemigrations --check reset-db: - docker-compose down db --volumes - docker-compose up -d db + docker compose down db --volumes + docker compose up -d db diff --git a/core_api/Dockerfile b/core_api/Dockerfile index aa25ed770..ae809b212 100644 --- a/core_api/Dockerfile +++ b/core_api/Dockerfile @@ -6,7 +6,7 @@ RUN pip install poetry ADD redbox/ /app/redbox ADD pyproject.toml poetry.lock /app/ WORKDIR /app/ -RUN poetry install --no-root --no-ansi --with api --without ai,ingest,dev,worker +RUN poetry install --no-root --no-ansi --with api --without ai,ingest,dev,worker,pytest-django,streamlit-app WORKDIR /app diff --git a/django_app/redbox_app/redbox_core/views.py b/django_app/redbox_app/redbox_core/views.py index c58ec5780..0fec2c039 100644 --- a/django_app/redbox_app/redbox_core/views.py +++ b/django_app/redbox_app/redbox_core/views.py @@ -13,8 +13,7 @@ def homepage_view(request): def documents_view(request): - - # Testing with dummy data for now + # Testing with dummy data for now if not File.objects.exists(): File.objects.create(name="Document 1", path="#download1", processing_status=ProcessingStatusEnum.complete) File.objects.create(name="Document 2", path="#download2", processing_status=ProcessingStatusEnum.parsing) @@ -54,4 +53,4 @@ def remove_doc_view(request, doc_id: str): request, template_name="remove-doc.html", context={"request": request, "doc_id": doc_id, "doc_name": doc_name}, - ) \ No newline at end of file + ) diff --git a/django_app/redbox_app/urls.py b/django_app/redbox_app/urls.py index 08893e035..cfb14891a 100644 --- a/django_app/redbox_app/urls.py +++ b/django_app/redbox_app/urls.py @@ -15,7 +15,7 @@ path("accounts/", include("allauth.urls")), path("documents/", views.documents_view, name="documents"), path("upload/", views.upload_view, name="upload"), - path("remove-doc/", views.remove_doc_view, name="remove_doc") + path("remove-doc/", views.remove_doc_view, name="remove_doc"), ] urlpatterns = info_urlpatterns + other_urlpatterns diff --git a/docker-compose.yml b/docker-compose.yml index dcfc7024d..9e858b75f 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -40,7 +40,7 @@ services: networks: - redbox-app-network streamlit-app: - image: redbox-app:latest + image: redbox-streamlit-app:latest build: context: . dockerfile: ./streamlit_app/Dockerfile diff --git a/embed/Dockerfile b/embed/Dockerfile index 1d07dcd7a..f00d6d8c6 100644 --- a/embed/Dockerfile +++ b/embed/Dockerfile @@ -9,7 +9,7 @@ ADD pyproject.toml poetry.lock /app/ ADD ./download_embedder.py /app/ ADD ./model_db.py /app/ WORKDIR /app/ -RUN poetry install --no-root --no-ansi --with worker,api --without ai,ingest,dev +RUN poetry install --no-root --no-ansi --with worker,api --without ai,ingest,dev,pytest-django,streamlit-app # Add the rest of the files ADD ./embed/src/app.py /app/app.py diff --git a/ingest/Dockerfile b/ingest/Dockerfile index 36b9e8eb4..bdc99aa8e 100644 --- a/ingest/Dockerfile +++ b/ingest/Dockerfile @@ -10,7 +10,7 @@ ADD download_embedder.py /app/ ADD model_db.py /app/ WORKDIR /app/ -RUN poetry install --no-root --no-ansi --with worker,ingest --without ai,dev,api,django-app,pytest-django +RUN poetry install --no-root --no-ansi --with worker,ingest --without ai,dev,api,django-app,pytest-django,streamlit-app # Download the model RUN poetry run download-model diff --git a/pyproject.toml b/pyproject.toml index 3d412496d..0e2236ab6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,6 +63,8 @@ boto3-stubs = "^1.34.67" streamlit = "^1.32.2" lxml = "^5.1.0" loguru = "^0.7.2" +streamlit-feedback = "^0.1.3" +html2markdown = "^0.1.7" [tool.poetry.group.ai.dependencies] anthropic = "^0.21.1" @@ -70,12 +72,6 @@ litellm = "^1.32.7" openai = "^1.14.2" -[tool.poetry.group.streamlitapp.dependencies] -cognitojwt = "^1.4.1" -pandas = "^2.2.1" -html2markdown = "^0.1.7" -streamlit-feedback = "^0.1.3" - [tool.poetry.group.django-app.dependencies] python = "^3.11.2" django = "^4.2.8" diff --git a/redbox/Dockerfile b/redbox/Dockerfile index d3bab0db8..def1b2ebe 100644 --- a/redbox/Dockerfile +++ b/redbox/Dockerfile @@ -5,7 +5,7 @@ RUN apt-get update RUN pip install poetry ADD ../pyproject.toml ../poetry.lock ./ -RUN poetry install --no-root --no-ansi --with worker,api --without ai,streamlit-app,ingest,dev +RUN poetry install --no-root --no-ansi --with worker,api --without ai,streamlit-app,ingest,dev,pytest-django ADD ./redbox /app/redbox ADD ./Makefile /app/ diff --git a/streamlit_app/Dockerfile b/streamlit_app/Dockerfile index d1ca6a7cf..d0ed711a4 100644 --- a/streamlit_app/Dockerfile +++ b/streamlit_app/Dockerfile @@ -1,25 +1,41 @@ -FROM python:3.11 +FROM python:3.11-buster as builder + -RUN apt-get update -RUN apt-get install -y libgl-dev libmagic-dev inetutils-ping RUN pip install poetry +ENV POETRY_NO_INTERACTION=1 \ + POETRY_VIRTUALENVS_IN_PROJECT=1 \ + POETRY_VIRTUALENVS_CREATE=1 \ + POETRY_CACHE_DIR=/tmp/poetry_cache + WORKDIR /app/ ADD pyproject.toml poetry.lock ./ # Streamlit app needs a lot of dependencies + +RUN --mount=type=cache,target=$POETRY_CACHE_DIR poetry install --no-root --no-ansi --with streamlit-app,ai,ingest,embed --without dev,worker,api,pytest-django --no-root + +FROM python:3.11-slim-buster as runtime + +RUN apt-get update +RUN apt-get install -y libgl-dev libmagic-dev inetutils-ping + + +ENV VIRTUAL_ENV=/app/.venv \ + PATH="/app/.venv/bin:$PATH" + +COPY --from=builder ${VIRTUAL_ENV} ${VIRTUAL_ENV} + +WORKDIR /app/ + ADD ./download_embedder.py /app/ ADD ./model_db.py /app/ -RUN poetry install --no-root --no-ansi --with streamlit-app,ai,ingest --without dev,worker,api +RUN python download_embedder.py ADD redbox/ /app/redbox ADD streamlit_app/ /app -# Download the model - -RUN poetry run download-model - EXPOSE 8501 -ENTRYPOINT ["poetry", "run", "streamlit", "run", "--server.address", "0.0.0.0", "--server.port", "8501", "Welcome.py" ] +ENTRYPOINT ["streamlit", "run", "--server.address", "0.0.0.0", "--server.port", "8501", "Welcome.py" ] diff --git a/streamlit_app/pages/5_Persona_Chat.py b/streamlit_app/pages/5_Persona_Chat.py index 5a887d933..de3362262 100644 --- a/streamlit_app/pages/5_Persona_Chat.py +++ b/streamlit_app/pages/5_Persona_Chat.py @@ -23,10 +23,12 @@ # Model selector + def change_selected_model(): load_llm_handler(ENV, update=True) st.write(st.session_state.llm) + persona_select = st.sidebar.selectbox( "What is your role?", options=st.session_state.available_personas, @@ -72,6 +74,7 @@ def get_files_by_uuid(file_uuids): files = st.session_state.storage_handler.read_items(file_uuids, "File") return files + def render_citation_response(response): cited_chunks = [ ( diff --git a/streamlit_app/utils.py b/streamlit_app/utils.py index 031a1b951..a6182d8e8 100644 --- a/streamlit_app/utils.py +++ b/streamlit_app/utils.py @@ -7,7 +7,6 @@ from typing import Optional import boto3 -import cognitojwt import dotenv import html2markdown import pandas as pd @@ -18,17 +17,13 @@ from langchain.callbacks.base import BaseCallbackHandler from langchain.chains.base import Chain from langchain.schema.output import LLMResult -from langchain.vectorstores.elasticsearch import ( - ApproxRetrievalStrategy, - ElasticsearchStore, -) +from langchain.vectorstores.elasticsearch import ApproxRetrievalStrategy, ElasticsearchStore from langchain_community.chat_models import ChatLiteLLM from langchain_community.embeddings import SentenceTransformerEmbeddings from loguru import logger from lxml.html.clean import Cleaner -from sentence_transformers import SentenceTransformer -from streamlit.web.server.websocket_headers import _get_websocket_headers +from model_db import SentenceTransformerDB from redbox.llm.llm_base import LLMHandler from redbox.models.feedback import Feedback from redbox.models.file import File @@ -134,22 +129,18 @@ def init_session_state() -> dict: "Foreign Policy Experts", ] + if "model_db" not in st.session_state: + st.session_state.model_db = SentenceTransformerDB() + st.session_state.model_db.init_from_disk() + if "embedding_model" not in st.session_state: available_models = [] - models = {} - for dirpath, _, filenames in os.walk("models"): - # Check if the current directory contains a file named "config.json" - if "pytorch_model.bin" in filenames: - # If it does, print the path to the directory - available_models.append(dirpath) + for model_name in st.session_state.model_db: + available_models.append(model_name) - for model_path in available_models: - model_name = model_path.split("/")[-3] - model = model_name.split("--")[-1] - models[model] = SentenceTransformer(model_path) + default_model = available_models[0] - model_name = ENV.get("EMBEDDING_MODEL", "all-mpnet-base-v2") - st.session_state.embedding_model = models[model_name] + st.session_state.embedding_model = st.session_state.model_db[default_model] if "BUCKET_NAME" not in st.session_state: st.session_state.BUCKET_NAME = f"redbox-storage-{st.session_state.user_uuid}" @@ -160,9 +151,7 @@ def init_session_state() -> dict: # The bucket does not exist or you have no access. if err.response["Error"]["Code"] == "404": print("The bucket does not exist.") - st.session_state.s3_client.create_bucket( - Bucket=st.session_state.BUCKET_NAME - ) + st.session_state.s3_client.create_bucket(Bucket=st.session_state.BUCKET_NAME) print("Bucket created successfully.") else: raise err @@ -178,9 +167,7 @@ def init_session_state() -> dict: ], basic_auth=(ENV["ELASTIC_USER"], ENV["ELASTIC_PASSWORD"]), ) - st.session_state.storage_handler = ElasticsearchStorageHandler( - es_client=es, root_index="redbox-data" - ) + st.session_state.storage_handler = ElasticsearchStorageHandler(es_client=es, root_index="redbox-data") if st.session_state.user_uuid == "dev": st.sidebar.info("**DEV MODE**") @@ -245,9 +232,7 @@ def init_session_state() -> dict: return ENV -def get_link_html( - page: str, text: str, query_dict: Optional[dict] = None, target: str = "_self" -) -> str: +def get_link_html(page: str, text: str, query_dict: Optional[dict] = None, target: str = "_self") -> str: """Returns a link in HTML format Args: @@ -403,16 +388,12 @@ def st_render(self, file: File) -> None: """ render_method = self.render_methods[file.type] - stream = st.session_state.s3_client.get_object( - Bucket=st.session_state.BUCKET_NAME, Key=file.name - ) + stream = st.session_state.s3_client.get_object(Bucket=st.session_state.BUCKET_NAME, Key=file.name) file_bytes = stream["Body"].read() render_method(file, file_bytes) def _render_pdf(self, file: File, page_number: Optional[int] = None) -> None: - stream = st.session_state.s3_client.get_object( - Bucket=st.session_state.BUCKET_NAME, Key=file.name - ) + stream = st.session_state.s3_client.get_object(Bucket=st.session_state.BUCKET_NAME, Key=file.name) base64_pdf = base64.b64encode(stream["Body"].read()).decode("utf-8") if page_number is not None: @@ -444,9 +425,7 @@ def _render_csv(self, file: File, file_bytes: bytes) -> None: st.dataframe(df, use_container_width=True) def _render_eml(self, file: File, file_bytes: bytes) -> None: - st.markdown( - self.cleaner.clean_html(file_bytes.decode("utf-8")), unsafe_allow_html=True - ) + st.markdown(self.cleaner.clean_html(file_bytes.decode("utf-8")), unsafe_allow_html=True) def _render_html(self, file: File, file_bytes: bytes) -> None: markdown_html = html2markdown.convert(file_bytes.decode("utf-8"))