Upload folder using huggingface_hub

.flake8

@@ -0,0 +1,6 @@
+[flake8]
+ignore = E203, E501, W503
+max-line-length = 88
+max-complexity = 10
+per-file-ignores =
+    */__init__.py: F401, F403

.gitignore

@@ -0,0 +1,180 @@
+# Created by https://www.toptal.com/developers/gitignore/api/python,pycharm
+# Edit at https://www.toptal.com/developers/gitignore?templates=python,pycharm
+
+### PyCharm ###
+# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider
+# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
+
+.idea/
+
+# AWS User-
+# CMake
+cmake-build-*/
+
+# File-based project format
+*.iws
+
+# IntelliJ
+out/
+
+# mpeltonen/sbt-idea plugin
+.idea_modules/
+
+# JIRA plugin
+atlassian-ide-plugin.xml
+
+# Crashlytics plugin (for Android Studio and IntelliJ)
+com_crashlytics_export_strings.xml
+crashlytics.properties
+crashlytics-build.properties
+fabric.properties
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+poetry.lock
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# End of https://www.toptal.com/developers/gitignore/api/python,pycharm
+
+coverage_report/
+
+local_data/
+models/
+.DS_Store

.isort.cfg

@@ -0,0 +1,2 @@
+[settings]
+profile = black

.pre-commit-config.yaml

@@ -0,0 +1,14 @@
+repos:
+  - repo: https://github.com/psf/black
+    rev: 23.3.0
+    hooks:
+      - id: black
+  - repo: https://github.com/PyCQA/flake8
+    rev: 6.0.0
+    hooks:
+      - id: flake8
+  - repo: https://github.com/PyCQA/isort
+    rev: 5.12.0
+    hooks:
+      - id: isort
+

README.md

@@ -1,6 +1,62 @@
 ---
 title: discord-bot
-app_file: __main__.py
+app_file: app/__main__.py
 sdk: gradio
 sdk_version: 4.33.0
 ---
+# Capstone Project
+
+## Requirements
+
+* [Docker](https://www.docker.com/).
+* [Docker Compose](https://docs.docker.com/compose/install/).
+* [Poetry](https://python-poetry.org/) for Python package and environment management.
+
+
+## Installation
+
+### Set up virtual environment
+
+```shell
+python -m venv venv
+. venv/bin/activate
+```
+
+### Install dependencies
+
+- Install poetry: https://python-poetry.org/
+- Install dependencies
+
+```shell
+poetry install
+```
+
+#### Local Requirements
+```shell
+poetry install --with local
+```
+
+Download embedding and(or) LLM models
+```shell
+bash prestart.sh
+```
+
+### Install `pre-commit` hooks
+
+```shell
+pre-commit install
+```
+
+## Running
+```shell
+docker compose up -d
+poetry run python -m app
+```
+
+## URLs
+### Development URLs
+#### Gradio UI
+http://localhost:8000
+
+#### API Documentation
+http://localhost:8000/docs

app/README.md

@@ -0,0 +1,6 @@
+---
+title: discord-bot
+app_file: __main__.py
+sdk: gradio
+sdk_version: 4.33.0
+---

app/__init__.py

@@ -0,0 +1,9 @@
+import logging
+
+ROOT_LOG_LEVEL = "INFO"
+
+PRETTY_LOG_FORMAT = (
+    "%(asctime)s.%(msecs)03d [%(levelname)-8s] %(name)+25s - %(message)s"
+)
+logging.basicConfig(level=ROOT_LOG_LEVEL, format=PRETTY_LOG_FORMAT, datefmt="%H:%M:%S")
+logging.captureWarnings(True)

app/__main__.py

@@ -0,0 +1,7 @@
+import uvicorn
+
+from app._config import settings
+from app.main import app
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=settings.PORT)

app/_config.py

@@ -0,0 +1,62 @@
+import os
+from typing import Literal, Optional
+
+from pydantic import Field
+from pydantic_settings import BaseSettings
+
+
+class Settings(BaseSettings):
+    ENVIRONMENT: str
+    PORT: int = 8000
+    VECTOR_DATABASE: Literal["weaviate"] = "weaviate"
+
+    OPENAI_API_KEY: Optional[str] = None
+    OPENAI_MODEL: str = "gpt-3.5-turbo"
+
+    WEAVIATE_CLIENT_URL: str = "http://localhost:8080"
+
+    LLM_MODE: Literal["openai", "mock", "local"] = "mock"
+    EMBEDDING_MODE: Literal["openai", "mock", "local"] = "mock"
+
+    LOCAL_DATA_FOLDER: str = "local_data/test"
+
+    DEFAULT_QUERY_SYSTEM_PROMPT: str = "You can only answer questions about the provided context. If you know the answer but it is not based in the provided context, don't provide the answer, just state the answer is not in the context provided."
+
+    LOCAL_HF_EMBEDDING_MODEL_NAME: str = "BAAI/bge-small-en-v1.5"
+
+    LOCAL_HF_LLM_REPO_ID: str = "TheBloke/Llama-2-7B-Chat-GGUF"
+    LOCAL_HF_LLM_MODEL_FILE: str = "llama-2-7b-chat.Q4_K_M.gguf"
+
+    # LLM config
+    LLM_TEMPERATURE: float = Field(
+        default=0.1, description="The temperature to use for sampling."
+    )
+    LLM_MAX_NEW_TOKENS: int = Field(
+        default=256,
+        description="The maximum number of tokens to generate.",
+    )
+    LLM_CONTEXT_WINDOW: int = Field(
+        default=3900,
+        description="The maximum number of context tokens for the model.",
+    )
+
+    # UI
+    IS_UI_ENABLED: bool = True
+    UI_PATH: str = "/"
+
+    # Rerank
+    IS_RERANK_ENABLED: bool = True
+    RERANK_TOP_N: int = 3
+    RERANK_MODEL_NAME: str = "cross-encoder/ms-marco-MiniLM-L-2-v2"
+
+    class Config:
+        case_sensitive = True
+        env_file_encoding = "utf-8"
+
+
+environment = os.environ.get("ENVIRONMENT", "local")
+settings = Settings(
+    ENVIRONMENT=environment,
+    # ".env.{environment}" takes priority over ".env"
+    _env_file=[".env", f".env.{environment}"],
+)

app/components/embedding/component.py

@@ -0,0 +1,38 @@
+import logging
+
+from llama_index import MockEmbedding
+from llama_index.embeddings.base import BaseEmbedding
+
+from app._config import settings
+from app.enums import EmbeddingMode
+from app.paths import models_cache_path
+
+logger = logging.getLogger(__name__)
+
+MOCK_EMBEDDING_DIM = 1536
+
+
+class EmbeddingComponent:
+    embedding_model: BaseEmbedding
+
+    def __init__(self) -> None:
+        embedding_mode = settings.EMBEDDING_MODE
+        logger.info("Initializing the embedding model in mode=%s", embedding_mode)
+        match embedding_mode:
+            case EmbeddingMode.OPENAI:
+                from llama_index import OpenAIEmbedding
+
+                self.embedding_model = OpenAIEmbedding(api_key=settings.OPENAI_API_KEY)
+
+            case EmbeddingMode.MOCK:
+                # Not a random number, is the dimensionality used by
+                # the default embedding model
+                self.embedding_model = MockEmbedding(MOCK_EMBEDDING_DIM)
+
+            case EmbeddingMode.LOCAL:
+                from llama_index.embeddings import HuggingFaceEmbedding
+
+                self.embedding_model = HuggingFaceEmbedding(
+                    model_name=settings.LOCAL_HF_EMBEDDING_MODEL_NAME,
+                    cache_folder=str(models_cache_path),
+                )

app/components/ingest/component.py

@@ -0,0 +1,143 @@
+import abc
+import logging
+import threading
+from pathlib import Path
+from typing import Any
+
+from llama_index import (
+    Document,
+    ServiceContext,
+    StorageContext,
+    VectorStoreIndex,
+    load_index_from_storage,
+)
+from llama_index.data_structs import IndexDict
+from llama_index.indices.base import BaseIndex
+
+from app.components.ingest.helpers import IngestionHelper
+from app.paths import local_data_path
+
+logger = logging.getLogger(__name__)
+
+
+class BaseIngestComponent(abc.ABC):
+    def __init__(
+        self,
+        storage_context: StorageContext,
+        service_context: ServiceContext,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        logger.debug(f"Initializing base ingest component type={type(self).__name__}")
+        self.storage_context = storage_context
+        self.service_context = service_context
+
+    @abc.abstractmethod
+    def ingest(self, file_name: str, file_data: Path) -> list[Document]:
+        pass
+
+    @abc.abstractmethod
+    def bulk_ingest(self, files: list[tuple[str, Path]]) -> list[Document]:
+        pass
+
+    @abc.abstractmethod
+    def delete(self, doc_id: str) -> None:
+        pass
+
+
+class BaseIngestComponentWithIndex(BaseIngestComponent, abc.ABC):
+    def __init__(
+        self,
+        storage_context: StorageContext,
+        service_context: ServiceContext,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(storage_context, service_context, *args, **kwargs)
+
+        self.show_progress = True
+        self._index_thread_lock = (
+            threading.Lock()
+        )  # Thread lock! Not Multiprocessing lock
+        self._index = self._initialize_index()
+
+    def _initialize_index(self) -> BaseIndex[IndexDict]:
+        """Initialize the index from the storage context."""
+        try:
+            # Load the index with store_nodes_override=True to be able to delete them
+            index = load_index_from_storage(
+                storage_context=self.storage_context,
+                service_context=self.service_context,
+                store_nodes_override=True,  # Force store nodes in index and document stores
+                show_progress=self.show_progress,
+            )
+        except ValueError:
+            # There are no index in the storage context, creating a new one
+            logger.info("Creating a new vector store index")
+            index = VectorStoreIndex.from_documents(
+                [],
+                storage_context=self.storage_context,
+                service_context=self.service_context,
+                store_nodes_override=True,  # Force store nodes in index and document stores
+                show_progress=self.show_progress,
+            )
+            index.storage_context.persist(persist_dir=local_data_path)
+        return index
+
+    def _save_index(self) -> None:
+        self._index.storage_context.persist(persist_dir=local_data_path)
+
+    def delete(self, doc_id: str) -> None:
+        with self._index_thread_lock:
+            # Delete the document from the index
+            self._index.delete_ref_doc(doc_id, delete_from_docstore=True)
+
+            # Save the index
+            self._save_index()
+
+
+class SimpleIngestComponent(BaseIngestComponentWithIndex):
+    def __init__(
+        self,
+        storage_context: StorageContext,
+        service_context: ServiceContext,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(storage_context, service_context, *args, **kwargs)
+
+    def ingest(self, file_name: str, file_data: Path) -> list[Document]:
+        logger.info("Ingesting file_name=%s", file_name)
+        documents = IngestionHelper.transform_file_into_documents(file_name, file_data)
+        logger.info(
+            "Transformed file=%s into count=%s documents", file_name, len(documents)
+        )
+        logger.debug("Saving the documents in the index and doc store")
+        return self._save_docs(documents)
+
+    def bulk_ingest(self, files: list[tuple[str, Path]]) -> list[Document]:
+        saved_documents = []
+        for file_name, file_data in files:
+            documents = IngestionHelper.transform_file_into_documents(
+                file_name, file_data
+            )
+            saved_documents.extend(self._save_docs(documents))
+        return saved_documents
+
+    def _save_docs(self, documents: list[Document]) -> list[Document]:
+        logger.debug("Transforming count=%s documents into nodes", len(documents))
+        with self._index_thread_lock:
+            for document in documents:
+                self._index.insert(document, show_progress=True)
+            logger.debug("Persisting the index and nodes")
+            # persist the index and nodes
+            self._save_index()
+            logger.debug("Persisted the index and nodes")
+        return documents
+
+
+def get_ingestion_component(
+    storage_context: StorageContext,
+    service_context: ServiceContext,
+) -> BaseIngestComponent:
+    return SimpleIngestComponent(storage_context, service_context)

app/components/ingest/helpers.py

@@ -0,0 +1,61 @@
+import logging
+from pathlib import Path
+
+from llama_index import Document
+from llama_index.readers import JSONReader, StringIterableReader
+from llama_index.readers.file.base import DEFAULT_FILE_READER_CLS
+
+logger = logging.getLogger(__name__)
+
+# Patching the default file reader to support other file types
+FILE_READER_CLS = DEFAULT_FILE_READER_CLS.copy()
+FILE_READER_CLS.update(
+    {
+        ".json": JSONReader,
+    }
+)
+
+
+class IngestionHelper:
+    """Helper class to transform a file into a list of documents.
+
+    This class should be used to transform a file into a list of documents.
+    These methods are thread-safe (and multiprocessing-safe).
+    """
+
+    @staticmethod
+    def transform_file_into_documents(
+        file_name: str, file_data: Path
+    ) -> list[Document]:
+        documents = IngestionHelper._load_file_to_documents(file_name, file_data)
+        for document in documents:
+            document.metadata["file_name"] = file_name
+        IngestionHelper._exclude_metadata(documents)
+        return documents
+
+    @staticmethod
+    def _load_file_to_documents(file_name: str, file_data: Path) -> list[Document]:
+        logger.debug("Transforming file_name=%s into documents", file_name)
+        extension = Path(file_name).suffix
+        reader_cls = FILE_READER_CLS.get(extension)
+        if reader_cls is None:
+            logger.debug(
+                "No reader found for extension=%s, using default string reader",
+                extension,
+            )
+            # Read as a plain text
+            string_reader = StringIterableReader()
+            return string_reader.load_data([file_data.read_text()])
+
+        logger.debug("Specific reader found for extension=%s", extension)
+        return reader_cls().load_data(file_data)
+
+    @staticmethod
+    def _exclude_metadata(documents: list[Document]) -> None:
+        logger.debug("Excluding metadata from count=%s documents", len(documents))
+        for document in documents:
+            document.metadata["doc_id"] = document.doc_id
+            # We don't want the Embeddings search to receive this metadata
+            document.excluded_embed_metadata_keys = ["doc_id"]
+            # We don't want the LLM to receive these metadata in the context
+            document.excluded_llm_metadata_keys = ["file_name", "doc_id", "page_label"]

app/components/llm/component.py

@@ -0,0 +1,50 @@
+import logging
+
+from llama_index.llms import LLM, MockLLM
+
+from app._config import settings
+from app.enums import LLMMode
+from app.paths import models_path
+
+logger = logging.getLogger(__name__)
+
+
+class LLMComponent:
+    llm: LLM
+
+    def __init__(self) -> None:
+        llm_mode = settings.LLM_MODE
+        logger.info(f"Initializing the LLM in mode={llm_mode}")
+        match settings.LLM_MODE:
+            case LLMMode.OPENAI:
+                from llama_index.llms import OpenAI
+
+                self.llm = OpenAI(
+                    api_key=settings.OPENAI_API_KEY,
+                    model=settings.OPENAI_MODEL,
+                )
+            case LLMMode.MOCK:
+                self.llm = MockLLM()
+
+            case LLMMode.LOCAL:
+                from llama_index.llms import LlamaCPP
+                from llama_index.llms.llama_utils import (
+                    completion_to_prompt,
+                    messages_to_prompt,
+                )
+
+                self.llm = LlamaCPP(
+                    model_path=str(models_path / settings.LOCAL_HF_LLM_MODEL_FILE),
+                    temperature=settings.LLM_TEMPERATURE,
+                    max_new_tokens=settings.LLM_MAX_NEW_TOKENS,
+                    context_window=settings.LLM_CONTEXT_WINDOW,
+                    generate_kwargs={},
+                    # set to at least 1 to use GPU
+                    # set to -1 for all gpu
+                    # set to 0 for cpu
+                    model_kwargs={"n_gpu_layers": 0},
+                    # transform inputs into Llama2 format
+                    messages_to_prompt=messages_to_prompt,
+                    completion_to_prompt=completion_to_prompt,
+                    verbose=True,
+                )

app/components/node_store/component.py

@@ -0,0 +1,31 @@
+import logging
+
+from llama_index.storage.docstore import BaseDocumentStore, SimpleDocumentStore
+from llama_index.storage.index_store import SimpleIndexStore
+from llama_index.storage.index_store.types import BaseIndexStore
+
+from app.paths import local_data_path
+
+logger = logging.getLogger(__name__)
+
+
+class NodeStoreComponent:
+    index_store: BaseIndexStore
+    doc_store: BaseDocumentStore
+
+    def __init__(self) -> None:
+        try:
+            self.index_store = SimpleIndexStore.from_persist_dir(
+                persist_dir=str(local_data_path)
+            )
+        except FileNotFoundError:
+            logger.debug("Local index store not found, creating a new one")
+            self.index_store = SimpleIndexStore()
+
+        try:
+            self.doc_store = SimpleDocumentStore.from_persist_dir(
+                persist_dir=str(local_data_path)
+            )
+        except FileNotFoundError:
+            logger.debug("Local document store not found, creating a new one")
+            self.doc_store = SimpleDocumentStore()

app/components/vector_store/component.py

@@ -0,0 +1,51 @@
+import logging
+import typing
+
+from llama_index import VectorStoreIndex
+from llama_index.indices.vector_store import VectorIndexRetriever
+from llama_index.vector_stores.types import VectorStore
+
+from app._config import settings
+from app.enums import WEAVIATE_INDEX_NAME, VectorDatabase
+
+logger = logging.getLogger(__name__)
+
+
+class VectorStoreComponent:
+    vector_store: VectorStore
+
+    def __init__(self) -> None:
+        match settings.VECTOR_DATABASE:
+            case VectorDatabase.WEAVIATE:
+                import weaviate
+                from llama_index.vector_stores import WeaviateVectorStore
+
+                client = weaviate.Client(settings.WEAVIATE_CLIENT_URL)
+                self.vector_store = typing.cast(
+                    VectorStore,
+                    WeaviateVectorStore(
+                        weaviate_client=client, index_name=WEAVIATE_INDEX_NAME
+                    ),
+                )
+            case _:
+                # Should be unreachable
+                # The settings validator should have caught this
+                raise ValueError(
+                    f"Vectorstore database {settings.VECTOR_DATABASE} not supported"
+                )
+
+    @staticmethod
+    def get_retriever(
+        index: VectorStoreIndex,
+        doc_ids: list[str] | None = None,
+        similarity_top_k: int = 2,
+    ) -> VectorIndexRetriever:
+        return VectorIndexRetriever(
+            index=index,
+            similarity_top_k=similarity_top_k,
+            doc_ids=doc_ids,
+        )
+
+    def close(self) -> None:
+        if hasattr(self.vector_store.client, "close"):
+            self.vector_store.client.close()

app/enums.py

@@ -0,0 +1,39 @@
+from enum import Enum, auto, unique
+from pathlib import Path
+
+PROJECT_ROOT_PATH: Path = Path(__file__).parents[1]
+
+
+@unique
+class BaseEnum(str, Enum):
+    @staticmethod
+    def _generate_next_value_(name: str, *_):
+        """
+        Automatically generate values for enum.
+        Enum values are lower-cased enum member names.
+        """
+        return name.lower()
+
+    @classmethod
+    def get_values(cls) -> list[str]:
+        # noinspection PyUnresolvedReferences
+        return [m.value for m in cls]
+
+
+class LLMMode(BaseEnum):
+    MOCK = auto()
+    OPENAI = auto()
+    LOCAL = auto()
+
+
+class EmbeddingMode(BaseEnum):
+    MOCK = auto()
+    OPENAI = auto()
+    LOCAL = auto()
+
+
+class VectorDatabase(BaseEnum):
+    WEAVIATE = auto()
+
+
+WEAVIATE_INDEX_NAME = "LlamaIndex"

app/main.py

@@ -0,0 +1,38 @@
+import logging
+
+from fastapi import FastAPI
+
+from app._config import settings
+from app.components.embedding.component import EmbeddingComponent
+from app.components.llm.component import LLMComponent
+from app.components.node_store.component import NodeStoreComponent
+from app.components.vector_store.component import VectorStoreComponent
+from app.server.chat.router import chat_router
+from app.server.chat.service import ChatService
+from app.server.embedding.router import embedding_router
+from app.server.ingest.service import IngestService
+
+logger = logging.getLogger(__name__)
+
+app = FastAPI()
+app.include_router(chat_router)
+app.include_router(embedding_router)
+
+if settings.IS_UI_ENABLED:
+    logger.debug("Importing the UI module")
+    from app.ui.ui import PrivateGptUi
+
+    llm_component = LLMComponent()
+    vector_store_component = VectorStoreComponent()
+    embedding_component = EmbeddingComponent()
+    node_store_component = NodeStoreComponent()
+
+    ingest_service = IngestService(
+        llm_component, vector_store_component, embedding_component, node_store_component
+    )
+    chat_service = ChatService(
+        llm_component, vector_store_component, embedding_component, node_store_component
+    )
+
+    ui = PrivateGptUi(ingest_service, chat_service)
+    ui.mount_in_app(app, settings.UI_PATH)

app/paths.py

@@ -0,0 +1,15 @@
+from pathlib import Path
+
+from app._config import settings
+from app.enums import PROJECT_ROOT_PATH
+
+
+def _absolute_or_from_project_root(path: str) -> Path:
+    if path.startswith("/"):
+        return Path(path)
+    return PROJECT_ROOT_PATH / path
+
+
+local_data_path: Path = _absolute_or_from_project_root(settings.LOCAL_DATA_FOLDER)
+models_path: Path = PROJECT_ROOT_PATH / "models"
+models_cache_path: Path = models_path / "cache"

app/server/chat/router.py

@@ -0,0 +1,70 @@
+from fastapi import APIRouter
+from llama_index.llms import ChatMessage, MessageRole
+from pydantic import BaseModel
+
+from app.components.embedding.component import EmbeddingComponent
+from app.components.llm.component import LLMComponent
+from app.components.node_store.component import NodeStoreComponent
+from app.components.vector_store.component import VectorStoreComponent
+from app.server.chat.service import ChatService
+from app.server.chat.utils import OpenAICompletion, OpenAIMessage, to_openai_response
+
+chat_router = APIRouter()
+
+
+class ChatBody(BaseModel):
+    messages: list[OpenAIMessage]
+    include_sources: bool = True
+
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "messages": [
+                        {
+                            "role": "system",
+                            "content": "You are a rapper. Always answer with a rap.",
+                        },
+                        {
+                            "role": "user",
+                            "content": "How do you fry an egg?",
+                        },
+                    ],
+                    "include_sources": True,
+                }
+            ]
+        }
+    }
+
+
+@chat_router.post(
+    "/chat",
+    response_model=None,
+    responses={200: {"model": OpenAICompletion}},
+    tags=["Contextual Completions"],
+)
+def chat_completion(body: ChatBody) -> OpenAICompletion:
+    """Given a list of messages comprising a conversation, return a response.
+
+    Optionally include an initial `role: system` message to influence the way
+    the LLM answers.
+
+    When using `'include_sources': true`, the API will return the source Chunks used
+    to create the response, which come from the context provided.
+    """
+    llm_component = LLMComponent()
+    vector_store_component = VectorStoreComponent()
+    embedding_component = EmbeddingComponent()
+    node_store_component = NodeStoreComponent()
+
+    chat_service = ChatService(
+        llm_component, vector_store_component, embedding_component, node_store_component
+    )
+    all_messages = [
+        ChatMessage(content=m.content, role=MessageRole(m.role)) for m in body.messages
+    ]
+
+    completion = chat_service.chat(messages=all_messages)
+    return to_openai_response(
+        completion.response, completion.sources if body.include_sources else None
+    )

app/server/chat/schemas.py

@@ -0,0 +1,45 @@
+from typing import Literal
+
+from llama_index.schema import NodeWithScore
+from pydantic import BaseModel, Field
+
+from app.server.ingest.schemas import IngestedDoc
+
+
+class Chunk(BaseModel):
+    object: Literal["context.chunk"]
+    score: float = Field(examples=[0.023])
+    document: IngestedDoc
+    text: str = Field(examples=["Outbound sales increased 20%, driven by new leads."])
+    previous_texts: list[str] | None = Field(
+        default=None,
+        examples=[["SALES REPORT 2023", "Inbound didn't show major changes."]],
+    )
+    next_texts: list[str] | None = Field(
+        default=None,
+        examples=[
+            [
+                "New leads came from Google Ads campaign.",
+                "The campaign was run by the Marketing Department",
+            ]
+        ],
+    )
+
+    @classmethod
+    def from_node(cls: type["Chunk"], node: NodeWithScore) -> "Chunk":
+        doc_id = node.node.ref_doc_id if node.node.ref_doc_id is not None else "-"
+        return cls(
+            object="context.chunk",
+            score=node.score or 0.0,
+            document=IngestedDoc(
+                object="ingest.document",
+                doc_id=doc_id,
+                doc_metadata=node.metadata,
+            ),
+            text=node.get_content(),
+        )
+
+
+class Completion(BaseModel):
+    response: str
+    sources: list[Chunk] | None = None

app/server/chat/service.py

@@ -0,0 +1,122 @@
+from dataclasses import dataclass
+
+from llama_index import ServiceContext, StorageContext, VectorStoreIndex
+from llama_index.chat_engine import ContextChatEngine
+from llama_index.chat_engine.types import BaseChatEngine
+from llama_index.core.postprocessor import SentenceTransformerRerank
+from llama_index.indices.postprocessor import MetadataReplacementPostProcessor
+from llama_index.llms import ChatMessage, MessageRole
+
+from app._config import settings
+from app.components.embedding.component import EmbeddingComponent
+from app.components.llm.component import LLMComponent
+from app.components.node_store.component import NodeStoreComponent
+from app.components.vector_store.component import VectorStoreComponent
+from app.server.chat.schemas import Chunk, Completion
+
+
+@dataclass
+class ChatEngineInput:
+    system_message: ChatMessage | None = None
+    last_message: ChatMessage | None = None
+    chat_history: list[ChatMessage] | None = None
+
+    @classmethod
+    def from_messages(cls, messages: list[ChatMessage]) -> "ChatEngineInput":
+        # Detect if there is a system message, extract the last message and chat history
+        system_message = (
+            messages[0]
+            if len(messages) > 0 and messages[0].role == MessageRole.SYSTEM
+            else None
+        )
+        last_message = (
+            messages[-1]
+            if len(messages) > 0 and messages[-1].role == MessageRole.USER
+            else None
+        )
+        # Remove from messages list the system message and last message,
+        # if they exist. The rest is the chat history.
+        if system_message:
+            messages.pop(0)
+        if last_message:
+            messages.pop(-1)
+        chat_history = messages if len(messages) > 0 else None
+
+        return cls(
+            system_message=system_message,
+            last_message=last_message,
+            chat_history=chat_history,
+        )
+
+
+class ChatService:
+    def __init__(
+        self,
+        llm_component: LLMComponent,
+        vector_store_component: VectorStoreComponent,
+        embedding_component: EmbeddingComponent,
+        node_store_component: NodeStoreComponent,
+    ) -> None:
+        self.llm_service = llm_component
+        self.vector_store_component = vector_store_component
+        self.storage_context = StorageContext.from_defaults(
+            vector_store=vector_store_component.vector_store,
+            docstore=node_store_component.doc_store,
+            index_store=node_store_component.index_store,
+        )
+        self.service_context = ServiceContext.from_defaults(
+            llm=llm_component.llm, embed_model=embedding_component.embedding_model
+        )
+        self.index = VectorStoreIndex.from_vector_store(
+            vector_store_component.vector_store,
+            storage_context=self.storage_context,
+            service_context=self.service_context,
+            show_progress=True,
+        )
+
+    def _chat_engine(self, system_prompt: str | None = None) -> BaseChatEngine:
+        vector_index_retriever = self.vector_store_component.get_retriever(
+            index=self.index
+        )
+
+        node_postprocessors = [
+            MetadataReplacementPostProcessor(target_metadata_key="window")
+        ]
+        if settings.IS_RERANK_ENABLED:
+            rerank = SentenceTransformerRerank(
+                top_n=settings.RERANK_TOP_N, model=settings.RERANK_MODEL_NAME
+            )
+            node_postprocessors.append(rerank)
+
+        return ContextChatEngine.from_defaults(
+            system_prompt=system_prompt,
+            retriever=vector_index_retriever,
+            service_context=self.service_context,
+            node_postprocessors=node_postprocessors,
+        )
+
+    def chat(self, messages: list[ChatMessage]):
+        chat_engine_input = ChatEngineInput.from_messages(messages)
+        last_message = (
+            chat_engine_input.last_message.content
+            if chat_engine_input.last_message
+            else None
+        )
+        system_prompt = (
+            chat_engine_input.system_message.content
+            if chat_engine_input.system_message
+            else None
+        )
+        chat_history = (
+            chat_engine_input.chat_history if chat_engine_input.chat_history else None
+        )
+
+        chat_engine = self._chat_engine(system_prompt=system_prompt)
+        wrapped_response = chat_engine.chat(
+            message=last_message if last_message is not None else "",
+            chat_history=chat_history,
+        )
+
+        sources = [Chunk.from_node(node) for node in wrapped_response.source_nodes]
+        completion = Completion(response=wrapped_response.response, sources=sources)
+        return completion

app/server/chat/utils.py

@@ -0,0 +1,68 @@
+import time
+import uuid
+from typing import Literal
+
+from llama_index.llms import ChatResponse
+from pydantic import BaseModel, Field
+
+from app.server.chat.schemas import Chunk
+
+
+class OpenAIMessage(BaseModel):
+    """Inference result, with the source of the message.
+
+    Role could be the assistant or system
+    (providing a default response, not AI generated).
+    """
+
+    role: Literal["assistant", "system", "user"] = Field(default="user")
+    content: str | None
+
+
+class OpenAIChoice(BaseModel):
+    """Response from AI."""
+
+    finish_reason: str | None = Field(examples=["stop"])
+    message: OpenAIMessage | None = None
+    sources: list[Chunk] | None = None
+    index: int = 0
+
+
+class OpenAICompletion(BaseModel):
+    """Clone of OpenAI Completion model.
+
+    For more information see: https://platform.openai.com/docs/api-reference/chat/object
+    """
+
+    id: str
+    object: Literal["completion", "completion.chunk"] = Field(default="completion")
+    created: int = Field(..., examples=[1623340000])
+    model: Literal["llm-agriculture"]
+    choices: list[OpenAIChoice]
+
+    @classmethod
+    def from_text(
+        cls,
+        text: str | None,
+        finish_reason: str | None = None,
+        sources: list[Chunk] | None = None,
+    ) -> "OpenAICompletion":
+        return OpenAICompletion(
+            id=str(uuid.uuid4()),
+            object="completion",
+            created=int(time.time()),
+            model="llm-agriculture",
+            choices=[
+                OpenAIChoice(
+                    message=OpenAIMessage(role="assistant", content=text),
+                    finish_reason=finish_reason,
+                    sources=sources,
+                )
+            ],
+        )
+
+
+def to_openai_response(
+    response: str | ChatResponse, sources: list[Chunk] | None = None
+) -> OpenAICompletion:
+    return OpenAICompletion.from_text(response, finish_reason="stop", sources=sources)

app/server/embedding/router.py

@@ -0,0 +1,18 @@
+from fastapi import APIRouter
+
+from app.components.embedding.component import EmbeddingComponent
+from app.server.embedding.schemas import EmbeddingsBody, EmbeddingsResponse
+from app.server.embedding.service import EmbeddingsService
+
+embedding_router = APIRouter()
+
+
+@embedding_router.post("/embedding", tags=["Embeddings"])
+def generate_embeddings(body: EmbeddingsBody) -> EmbeddingsResponse:
+    embedding_component = EmbeddingComponent()
+    service = EmbeddingsService(embedding_component)
+    input_texts = body.input if isinstance(body.input, list) else [body.input]
+    embeddings = service.embed_texts(input_texts)
+    return EmbeddingsResponse(
+        object="list", model=service.embedding_model.model_name, data=embeddings
+    )

app/server/embedding/schemas.py

@@ -0,0 +1,19 @@
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class Embedding(BaseModel):
+    index: int
+    object: Literal["embedding"]
+    embedding: list[float] = Field(examples=[[0.1, -0.2]])
+
+
+class EmbeddingsBody(BaseModel):
+    input: str | list[str]
+
+
+class EmbeddingsResponse(BaseModel):
+    object: Literal["list"]
+    model: str
+    data: list[Embedding]

app/server/embedding/service.py

@@ -0,0 +1,18 @@
+from app.components.embedding.component import EmbeddingComponent
+from app.server.embedding.schemas import Embedding
+
+
+class EmbeddingsService:
+    def __init__(self, embedding_component: EmbeddingComponent) -> None:
+        self.embedding_model = embedding_component.embedding_model
+
+    def embed_texts(self, texts: list[str]) -> list[Embedding]:
+        texts_embeddings = self.embedding_model.get_text_embedding_batch(texts)
+        return [
+            Embedding(
+                index=texts_embeddings.index(embedding),
+                object="embedding",
+                embedding=embedding,
+            )
+            for embedding in texts_embeddings
+        ]

app/server/ingest/schemas.py

@@ -0,0 +1,32 @@
+from typing import Any, Literal
+
+from llama_index import Document
+from pydantic import BaseModel, Field
+
+
+class IngestedDoc(BaseModel):
+    object: Literal["ingest.document"]
+    doc_id: str = Field(examples=["c202d5e6-7b69-4869-81cc-dd574ee8ee11"])
+    doc_metadata: dict[str, Any] | None = Field(
+        examples=[
+            {
+                "page_label": "2",
+                "file_name": "agriculture.pdf",
+            }
+        ]
+    )
+
+    @staticmethod
+    def curate_metadata(metadata: dict[str, Any]) -> dict[str, Any]:
+        """Remove unwanted metadata keys."""
+        for key in ["doc_id", "window", "original_text"]:
+            metadata.pop(key, None)
+        return metadata
+
+    @staticmethod
+    def from_document(document: Document) -> "IngestedDoc":
+        return IngestedDoc(
+            object="ingest.document",
+            doc_id=document.doc_id,
+            doc_metadata=IngestedDoc.curate_metadata(document.metadata),
+        )

app/server/ingest/service.py

@@ -0,0 +1,123 @@
+import logging
+import tempfile
+from pathlib import Path
+from typing import AnyStr, BinaryIO
+
+from llama_index import ServiceContext, StorageContext
+from llama_index.node_parser import SentenceWindowNodeParser
+
+from app.components.embedding.component import EmbeddingComponent
+from app.components.ingest.component import get_ingestion_component
+from app.components.llm.component import LLMComponent
+from app.components.node_store.component import NodeStoreComponent
+from app.components.vector_store.component import VectorStoreComponent
+from app.server.ingest.schemas import IngestedDoc
+
+logger = logging.getLogger(__name__)
+
+
+class IngestService:
+    def __init__(
+        self,
+        llm_component: LLMComponent,
+        vector_store_component: VectorStoreComponent,
+        embedding_component: EmbeddingComponent,
+        node_store_component: NodeStoreComponent,
+    ) -> None:
+        self.llm_service = llm_component
+        self.storage_context = StorageContext.from_defaults(
+            vector_store=vector_store_component.vector_store,
+            docstore=node_store_component.doc_store,
+            index_store=node_store_component.index_store,
+        )
+        node_parser = SentenceWindowNodeParser.from_defaults()
+        self.ingest_service_context = ServiceContext.from_defaults(
+            llm=self.llm_service.llm,
+            embed_model=embedding_component.embedding_model,
+            node_parser=node_parser,
+            # Embeddings done early in the pipeline of node transformations, right
+            # after the node parsing
+            transformations=[node_parser, embedding_component.embedding_model],
+        )
+
+        self.ingest_component = get_ingestion_component(
+            self.storage_context, self.ingest_service_context
+        )
+
+    def _ingest_data(self, file_name: str, file_data: AnyStr) -> list[IngestedDoc]:
+        logger.debug(f"Got file data of size={len(file_data)} to ingest")
+        # llama-index mainly supports reading from files, so
+        # we have to create a tmp file to read for it to work
+        # delete=False to avoid a Windows 11 permission error.
+        with tempfile.NamedTemporaryFile(delete=False) as tmp:
+            try:
+                path_to_tmp = Path(tmp.name)
+                if isinstance(file_data, bytes):
+                    path_to_tmp.write_bytes(file_data)
+                else:
+                    path_to_tmp.write_text(str(file_data))
+                return self.ingest_file(file_name, path_to_tmp)
+            finally:
+                tmp.close()
+                path_to_tmp.unlink()
+
+    def ingest_file(self, file_name: str, file_data: Path) -> list[IngestedDoc]:
+        logger.info(f"Ingesting file_name={file_name}")
+        documents = self.ingest_component.ingest(file_name, file_data)
+        logger.info(f"Finished ingestion file_name={file_name}")
+        return [IngestedDoc.from_document(document) for document in documents]
+
+    def ingest_text(self, file_name: str, text: str) -> list[IngestedDoc]:
+        logger.debug(f"Ingesting text data with file_name={file_name}")
+        return self._ingest_data(file_name, text)
+
+    def ingest_bin_data(
+        self, file_name: str, raw_file_data: BinaryIO
+    ) -> list[IngestedDoc]:
+        logger.debug(f"Ingesting binary data with file_name={file_name}")
+        file_data = raw_file_data.read()
+        return self._ingest_data(file_name, file_data)
+
+    def bulk_ingest(self, files: list[tuple[str, Path]]) -> list[IngestedDoc]:
+        logger.info(f"Ingesting file_names={[f[0] for f in files]}")
+        documents = self.ingest_component.bulk_ingest(files)
+        logger.info(f"Finished ingestion file_name={[f[0] for f in files]}")
+        return [IngestedDoc.from_document(document) for document in documents]
+
+    def list_ingested(self) -> list[IngestedDoc]:
+        ingested_docs = []
+        try:
+            docstore = self.storage_context.docstore
+            ingested_docs_ids: set[str] = set()
+
+            for node in docstore.docs.values():
+                if node.ref_doc_id is not None:
+                    ingested_docs_ids.add(node.ref_doc_id)
+
+            for doc_id in ingested_docs_ids:
+                ref_doc_info = docstore.get_ref_doc_info(ref_doc_id=doc_id)
+                doc_metadata = None
+                if ref_doc_info is not None and ref_doc_info.metadata is not None:
+                    doc_metadata = IngestedDoc.curate_metadata(ref_doc_info.metadata)
+                ingested_docs.append(
+                    IngestedDoc(
+                        object="ingest.document",
+                        doc_id=doc_id,
+                        doc_metadata=doc_metadata,
+                    )
+                )
+        except ValueError:
+            logger.warning("Got an exception when getting list of docs", exc_info=True)
+            pass
+        logger.debug(f"Found count={len(ingested_docs)} ingested documents")
+        return ingested_docs
+
+    def delete(self, doc_id: str) -> None:
+        """Delete an ingested document.
+
+        :raises ValueError: if the document does not exist
+        """
+        logger.info(
+            "Deleting the ingested document=%s in the doc and index store", doc_id
+        )
+        self.ingest_component.delete(doc_id)

app/ui/__init__.py

@@ -0,0 +1 @@
+"""Gradio based UI."""

app/ui/schemas.py

@@ -0,0 +1,27 @@
+from pydantic import BaseModel
+
+from app.server.chat.schemas import Chunk
+
+
+class Source(BaseModel):
+    file: str
+    page: str
+    text: str
+
+    class Config:
+        frozen = True
+
+    @staticmethod
+    def curate_sources(sources: list[Chunk]) -> set["Source"]:
+        curated_sources = set()
+
+        for chunk in sources:
+            doc_metadata = chunk.document.doc_metadata
+
+            file_name = doc_metadata.get("file_name", "-") if doc_metadata else "-"
+            page_label = doc_metadata.get("page_label", "-") if doc_metadata else "-"
+
+            source = Source(file=file_name, page=page_label, text=chunk.text)
+            curated_sources.add(source)
+
+        return curated_sources

app/ui/ui.py

@@ -0,0 +1,228 @@
+"""This file should be imported only and only if you want to run the UI locally."""
+import itertools
+import logging
+from pathlib import Path
+from typing import Any
+
+import gradio as gr
+from fastapi import FastAPI
+from gradio.themes.utils.colors import slate
+from llama_index.llms import ChatMessage, MessageRole
+
+from app._config import settings
+from app.components.embedding.component import EmbeddingComponent
+from app.components.llm.component import LLMComponent
+from app.components.node_store.component import NodeStoreComponent
+from app.components.vector_store.component import VectorStoreComponent
+from app.enums import PROJECT_ROOT_PATH
+from app.server.chat.service import ChatService
+from app.server.ingest.service import IngestService
+from app.ui.schemas import Source
+
+logger = logging.getLogger(__name__)
+
+THIS_DIRECTORY_RELATIVE = Path(__file__).parent.relative_to(PROJECT_ROOT_PATH)
+AVATAR_BOT = THIS_DIRECTORY_RELATIVE / "dodge_ava.jpg"
+
+UI_TAB_TITLE = "Agriculture Chatbot"
+
+SOURCES_SEPARATOR = "\n\n Sources: \n"
+
+
+class PrivateGptUi:
+    def __init__(
+        self,
+        ingest_service: IngestService,
+        chat_service: ChatService,
+    ) -> None:
+        self._ingest_service = ingest_service
+        self._chat_service = chat_service
+
+        # Cache the UI blocks
+        self._ui_block = None
+
+        # Initialize system prompt
+        self._system_prompt = self._get_default_system_prompt()
+
+    def _chat(self, message: str, history: list[list[str]], *_: Any) -> Any:
+        def build_history() -> list[ChatMessage]:
+            history_messages: list[ChatMessage] = list(
+                itertools.chain(
+                    *[
+                        [
+                            ChatMessage(content=interaction[0], role=MessageRole.USER),
+                            ChatMessage(
+                                # Remove from history content the Sources information
+                                content=interaction[1].split(SOURCES_SEPARATOR)[0],
+                                role=MessageRole.ASSISTANT,
+                            ),
+                        ]
+                        for interaction in history
+                    ]
+                )
+            )
+
+            # max 20 messages to try to avoid context overflow
+            return history_messages[:20]
+
+        new_message = ChatMessage(content=message, role=MessageRole.USER)
+        all_messages = [*build_history(), new_message]
+        # If a system prompt is set, add it as a system message
+        if self._system_prompt:
+            all_messages.insert(
+                0,
+                ChatMessage(
+                    content=self._system_prompt,
+                    role=MessageRole.SYSTEM,
+                ),
+            )
+
+        completion = self._chat_service.chat(messages=all_messages)
+        full_response = completion.response
+
+        if completion.sources:
+            full_response += SOURCES_SEPARATOR
+            curated_sources = Source.curate_sources(completion.sources)
+            sources_text = "\n\n\n".join(
+                f"{index}. {source.file} (page {source.page})"
+                for index, source in enumerate(curated_sources, start=1)
+            )
+            full_response += sources_text
+
+        return full_response
+
+    # On initialization this function set the system prompt
+    # to the default prompt based on settings.
+    @staticmethod
+    def _get_default_system_prompt() -> str:
+        return settings.DEFAULT_QUERY_SYSTEM_PROMPT
+
+    def _set_system_prompt(self, system_prompt_input: str) -> None:
+        logger.info(f"Setting system prompt to: {system_prompt_input}")
+        self._system_prompt = system_prompt_input
+
+    def _list_ingested_files(self) -> list[list[str]]:
+        files = set()
+        for ingested_document in self._ingest_service.list_ingested():
+            if ingested_document.doc_metadata is None:
+                # Skipping documents without metadata
+                continue
+            file_name = ingested_document.doc_metadata.get(
+                "file_name", "[FILE NAME MISSING]"
+            )
+            files.add(file_name)
+        return [[row] for row in files]
+
+    def _upload_file(self, files: list[str]) -> None:
+        logger.debug("Loading count=%s files", len(files))
+        paths = [Path(file) for file in files]
+        self._ingest_service.bulk_ingest([(str(path.name), path) for path in paths])
+
+    def _build_ui_blocks(self) -> gr.Blocks:
+        logger.debug("Creating the UI blocks")
+        with gr.Blocks(
+            title=UI_TAB_TITLE,
+            theme=gr.themes.Soft(primary_hue=slate),
+            css=".logo { "
+            "display:flex;"
+            "height: 80px;"
+            "border-radius: 8px;"
+            "align-content: center;"
+            "justify-content: center;"
+            "align-items: center;"
+            "}"
+            ".logo img { height: 25% }"
+            ".contain { display: flex !important; flex-direction: column !important; }"
+            "#component-0, #component-3, #component-10, #component-8  { height: 100% !important; }"
+            "#chatbot { flex-grow: 1 !important; overflow: auto !important;}"
+            "#col { height: calc(100vh - 112px - 16px) !important; }",
+        ) as blocks:
+            with gr.Row():
+                gr.HTML(f"<div class='logo'/><h1>{UI_TAB_TITLE}</h1></div")
+
+            with gr.Row(equal_height=False):
+                with gr.Column(scale=3):
+                    upload_button = gr.components.UploadButton(
+                        "Upload File(s)",
+                        type="filepath",
+                        file_count="multiple",
+                        size="sm",
+                    )
+                    ingested_dataset = gr.List(
+                        self._list_ingested_files,
+                        headers=["File name"],
+                        label="Ingested Files",
+                        interactive=False,
+                        render=False,  # Rendered under the button
+                    )
+                    upload_button.upload(
+                        self._upload_file,
+                        inputs=upload_button,
+                        outputs=ingested_dataset,
+                    )
+                    ingested_dataset.change(
+                        self._list_ingested_files,
+                        outputs=ingested_dataset,
+                    )
+                    ingested_dataset.render()
+                    system_prompt_input = gr.Textbox(
+                        placeholder=self._system_prompt,
+                        label="System Prompt",
+                        lines=2,
+                        interactive=True,
+                        render=False,
+                    )
+
+                    # On blur, set system prompt to use in queries
+                    system_prompt_input.blur(
+                        self._set_system_prompt,
+                        inputs=system_prompt_input,
+                    )
+
+                with gr.Column(scale=7, elem_id="col"):
+                    _ = gr.ChatInterface(
+                        self._chat,
+                        chatbot=gr.Chatbot(
+                            label=f"LLM: {settings.LLM_MODE}",
+                            show_copy_button=True,
+                            elem_id="chatbot",
+                            render=False,
+                            avatar_images=(
+                                None,
+                                AVATAR_BOT,
+                            ),
+                        ),
+                        additional_inputs=[upload_button, system_prompt_input],
+                    )
+        return blocks
+
+    def get_ui_blocks(self) -> gr.Blocks:
+        if self._ui_block is None:
+            self._ui_block = self._build_ui_blocks()
+        return self._ui_block
+
+    def mount_in_app(self, app: FastAPI, path: str) -> None:
+        blocks = self.get_ui_blocks()
+        blocks.queue()
+        logger.info("Mounting the gradio UI, at path=%s", path)
+        gr.mount_gradio_app(app, blocks, path=path)
+
+
+if __name__ == "__main__":
+    llm_component = LLMComponent()
+    vector_store_component = VectorStoreComponent()
+    embedding_component = EmbeddingComponent()
+    node_store_component = NodeStoreComponent()
+
+    ingest_service = IngestService(
+        llm_component, vector_store_component, embedding_component, node_store_component
+    )
+    chat_service = ChatService(
+        llm_component, vector_store_component, embedding_component, node_store_component
+    )
+
+    ui = PrivateGptUi(ingest_service, chat_service)
+
+    _blocks = ui.get_ui_blocks()
+    _blocks.queue()
+    _blocks.launch(debug=False, show_api=False)

docker-compose.yml

@@ -0,0 +1,26 @@
+version: '3.4'
+services:
+  weaviate:
+    command:
+    - --host
+    - 0.0.0.0
+    - --port
+    - '8080'
+    - --scheme
+    - http
+    image: semitechnologies/weaviate:1.23.0
+    ports:
+    - 8080:8080
+    - 50051:50051
+    volumes:
+    - weaviate_data:/var/lib/weaviate
+    restart: on-failure:0
+    environment:
+      QUERY_DEFAULTS_LIMIT: 25
+      AUTHENTICATION_ANONYMOUS_ACCESS_ENABLED: 'true'
+      PERSISTENCE_DATA_PATH: '/var/lib/weaviate'
+      DEFAULT_VECTORIZER_MODULE: 'none'
+      CLUSTER_HOSTNAME: 'node1'
+
+volumes:
+  weaviate_data:

prestart.sh

@@ -0,0 +1,5 @@
+#! /usr/bin/env bash
+
+# Download Embedding model and LLM model
+python setup.py
+

pyproject.toml

@@ -0,0 +1,31 @@
+[tool.poetry]
+name = "capstone"
+version = "0.1.0"
+description = ""
+authors = ["PhucVu <[email protected]>"]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = "^3.10"
+llama-index = "^0.9.22"
+weaviate-client = "^3.26.0"
+pydantic-settings = "^2.1.0"
+fastapi = "^0.108.0"
+uvicorn = "^0.25.0"
+pydantic = "^2.5.3"
+gradio = "^4.12.0"
+
+# reranker
+torch = {version="^2.3.0", optional=true}
+sentence-transformers = {version="^2.7.0", optional=true}
+
+[tool.poetry.group.local]
+optional = true
+[tool.poetry.group.local.dependencies]
+transformers = "^4.36.2"
+torch = "^2.1.2"
+llama-cpp-python = "^0.2.29"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

setup.py

@@ -0,0 +1,31 @@
+import logging
+import os
+
+from huggingface_hub import hf_hub_download, snapshot_download
+
+from app._config import settings
+from app.paths import models_cache_path, models_path
+
+os.makedirs(models_path, exist_ok=True)
+
+logger = logging.getLogger(__name__)
+
+# Download embedding model
+embedding_path = models_path / "embedding"
+logger.info(f"Downloading embedding: {settings.LOCAL_HF_EMBEDDING_MODEL_NAME}")
+snapshot_download(
+    repo_id=settings.LOCAL_HF_EMBEDDING_MODEL_NAME,
+    cache_dir=models_cache_path,
+    local_dir=embedding_path,
+)
+logger.info("Embedding model downloaded")
+
+# Download LLM model
+logger.info(f"Downloading LLM: {settings.LOCAL_HF_LLM_MODEL_FILE}")
+hf_hub_download(
+    repo_id=settings.LOCAL_HF_LLM_REPO_ID,
+    filename=settings.LOCAL_HF_LLM_MODEL_FILE,
+    cache_dir=models_cache_path,
+    local_dir=models_path,
+)
+logger.info("LLM model downloaded")