perf: update ingestion, add todo for retriever

rag_chatbot/core/engine/__init__.py

@@ -1,5 +1,6 @@
 from .engine import LocalChatEngine
-
+from .retriever import LocalRetriever
 __all__ = [
     "LocalChatEngine",
+    "LocalRetriever"
 ]

rag_chatbot/core/engine/retriever.py

@@ -19,6 +19,7 @@
 
 class TwoStageRetriever:
     def __init__(self) -> None:
+        # TODO
         pass
 
 
@@ -43,7 +44,7 @@ def _get_two_stage_retriever(
             similarity_top_k=self._setting.retriever.similarity_top_k,
             embed_model=Settings.embed_model,
             verbose=True
-        )
+        )  # TODO
 
     def _get_fusion_retriever(
         self,

rag_chatbot/core/ingestion/ingestion.py

@@ -3,7 +3,7 @@
 from llama_index.core.schema import BaseNode
 from llama_index.core.node_parser import SentenceSplitter
 from dotenv import load_dotenv
-from typing import List
+from typing import Any, List
 from tqdm import tqdm
 from ...setting import RAGSettings
 
@@ -19,6 +19,8 @@ def __init__(self, setting: RAGSettings | None = None) -> None:
     def store_nodes(
         self,
         input_files: list[str],
+        embed_nodes: bool = True,
+        embed_model: Any | None = None
     ) -> List[BaseNode]:
         splitter = SentenceSplitter.from_defaults(
             chunk_size=self._setting.ingestion.chunk_size,
@@ -32,6 +34,8 @@ def store_nodes(
         ]
         return_nodes = []
         self._ingested_file = []
+        if embed_nodes:
+            Settings.embed_model = embed_model or Settings.embed_model
         for input_file in tqdm(input_files, desc="Ingesting data"):
             file_name = input_file.strip().split('/')[-1]
             self._ingested_file.append(file_name)
@@ -50,7 +54,8 @@ def store_nodes(
                     doc.excluded_llm_metadata_keys = excluded_keys
 
                 nodes = splitter(document)
-                nodes = Settings.embed_model(nodes)
+                if embed_nodes:
+                    nodes = Settings.embed_model(nodes)
                 self._node_store[file_name] = nodes
                 return_nodes.extend(nodes)
 

rag_chatbot/setting/setting.py

@@ -3,6 +3,9 @@
 
 
 class OllamaSettings(BaseModel):
+    llm: str = Field(
+        default="llama3:8b-instruct-q8_0", description="LLM model"
+    )
     keep_alive: str = Field(
         default="1h", description="Keep alive time for the server"
     )