Upload folder using huggingface_hub

README.md

@@ -1,12 +1,6 @@
 ---
-title: Discord Bot
-emoji: 😻
-colorFrom: gray
-colorTo: blue
+title: discord-bot
+app_file: __main__.py
 sdk: gradio
 sdk_version: 4.33.0
-app_file: app.py
-pinned: false
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

__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)

__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)

_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}"],
+)

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),
+                )

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)

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"]

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,
+                )

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()

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()

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"

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)

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"

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
+    )

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

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

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)