Refacto prompt

.env.example

@@ -30,7 +30,7 @@ RERANKER_ENABLED=true
 RERANKER_MODEL=Alibaba-NLP/gte-multilingual-reranker-base # or jinaai/jina-reranker-v2-base-multilingual
 
 # Prompts
-PROMPTS_DIR=../prompts/example3_en # you can change it to ../prompts/example3 for french prompts
+PROMPTS_DIR=../prompts/example1
 
 # Ray
 RAY_DEDUP_LOGS=0 # turns off ray log deduplication that appear across multiple processes

.hydra_config/config.yaml

@@ -60,16 +60,20 @@ verbose:
   level: DEBUG
 
 paths:
-  prompts_dir: ${oc.env:PROMPTS_DIR, ../prompts/example3}
+  prompts_dir: ${oc.env:PROMPTS_DIR, ../prompts/example1}
   data_dir: ${oc.env:DATA_DIR, ../data}
   db_dir: ${oc.env:DB_DIR, /app/db}
   log_dir: ${oc.env:LOG_DIR, /app/logs}
 
-prompt:
-  rag_sys_pmpt: rag_sys_prompt_template.txt
-  contextualizer_pmpt: contextualizer_pmpt.txt
-  chunk_contextualizer_pmpt: chunk_contextualizer_tmpl.txt
-  image_describer: image_captioning.txt
+prompts:
+  sys_prompt: sys_prompt_tmpl.txt
+  query_contextualizer: query_contextualizer_tmpl.txt
+  chunk_contextualizer: chunk_contextualizer_tmpl.txt
+  image_describer: image_captioning_tmpl.txt
+
+  # query templates for different retriever types
+  hyde: hyde.txt
+  multi_query: multi_query_pmpt_tmpl.txt
 
 loader:
   image_captioning: true

.hydra_config/retriever/hyde.yaml

@@ -3,5 +3,4 @@ defaults:
 
 type: hyde
 # Extra params
-combine: False
-prompt_tmpl: 'hyde.txt'
+combine: False

.hydra_config/retriever/multiQuery.yaml

@@ -2,5 +2,4 @@ defaults:
   - base
 
 type: multiQuery
-k_queries: 3
-prompt_tmpl: multi_query_prompt_template.txt
+k_queries: 3

charts/openrag-stack/values.yaml

@@ -191,8 +191,7 @@ env:
     RERANKER_MODEL_TYPE: "infinity"
 
     # Prompts
-    PROMPTS_DIR: "../prompts/example3"
-    COLBERT_LOAD_TORCH_EXTENSION_VERBOSE: "True"
+    PROMPTS_DIR: "../prompts/example1"
 
     # Loaders
     PDFLoader: "MarkerLoader"

docs/assets/env_linux_gpu.env

@@ -30,7 +30,7 @@ RERANKER_ENABLED=true
 RERANKER_MODEL=Alibaba-NLP/gte-multilingual-reranker-base # or jinaai/jina-reranker-v2-base-multilingual
 
 # Prompts
-PROMPTS_DIR=../prompts/example3_en # you can change it to ../prompts/example3 for french prompts
+PROMPTS_DIR=../prompts/example1
 
 # Ray
 RAY_DEDUP_LOGS=0 # turns off ray log deduplication that appear across multiple processes

docs/assets/env_ollama_cpu.env

@@ -41,7 +41,7 @@ RERANKER_TOP_K=5
 RERANKER_BASE_URL=
 
 # Prompts
-PROMPTS_DIR=../prompts/example3
+PROMPTS_DIR=../prompts/example1
 
 # Loaders
 PDFLoader=MarkerLoader

openrag/components/indexer/chunker/chunker.py

@@ -3,7 +3,8 @@
 from pathlib import Path
 from typing import Optional
 
-from components.utils import get_llm_semaphore, load_config, load_sys_template
+from components.prompts import CHUNK_CONTEXTUALIZER
+from components.utils import get_llm_semaphore, load_config
 from langchain_core.documents.base import Document
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
@@ -20,9 +21,6 @@
 
 logger = get_logger()
 config = load_config()
-prompt_paths = Path(config.paths.get("prompts_dir"))
-chunk_contextualizer_pmpt = config.prompt.get("chunk_contextualizer_pmpt")
-CHUNK_CONTEXTUALIZER = load_sys_template(prompt_paths / chunk_contextualizer_pmpt)
 
 
 class BaseChunker(ABC):

openrag/components/indexer/loaders/base.py

@@ -5,7 +5,8 @@
 from pathlib import Path
 from typing import Dict, Optional, Union
 
-from components.utils import get_vlm_semaphore, load_config, load_sys_template
+from components.prompts import IMAGE_DESCRIBER
+from components.utils import get_vlm_semaphore, load_config
 from langchain_core.messages import HumanMessage
 from langchain_openai import ChatOpenAI
 from PIL import Image
@@ -14,11 +15,6 @@
 logger = get_logger()
 config = load_config()
 
-# Load the image description prompt from the configuration
-prompts_dir = Path(config.paths.prompts_dir)
-img_desc_prompt_path = prompts_dir / config.prompt["image_describer"]
-IMAGE_DESCRIPTION_PROMPT = load_sys_template(img_desc_prompt_path)
-
 
 class BaseLoader(ABC):
     def __init__(self, **kwargs) -> None:
@@ -140,7 +136,7 @@ async def get_image_description(
                             "type": "image_url",
                             "image_url": {"url": image_url},
                         },
-                        {"type": "text", "text": IMAGE_DESCRIPTION_PROMPT},
+                        {"type": "text", "text": IMAGE_DESCRIBER},
                     ]
                 )
 

openrag/components/pipeline.py

@@ -1,7 +1,7 @@
 import copy
 from enum import Enum
-from pathlib import Path
 
+from components.prompts import QUERY_CONTEXTUALIZER_PROMPT, SYS_PROMPT_TMPLT
 from langchain_core.documents.base import Document
 from openai import AsyncOpenAI
 from utils.logger import get_logger
@@ -10,7 +10,7 @@
 from .map_reduce import RAGMapReduce
 from .reranker import Reranker
 from .retriever import ABCRetriever, RetrieverFactory
-from .utils import format_context, load_sys_template
+from .utils import format_context
 
 logger = get_logger()
 
@@ -70,17 +70,6 @@ def __init__(self, config, logger=None) -> None:
         # retriever pipeline
         self.retriever_pipeline = RetrieverPipeline(config=config, logger=self.logger)
 
-        self.prompts_dir = Path(config.paths.prompts_dir)
-        # contextualizer prompt
-        self.contextualizer_pmpt = load_sys_template(
-            self.prompts_dir / config.prompt["contextualizer_pmpt"]
-        )
-
-        # rag sys prompt
-        self.rag_sys_prompt: str = load_sys_template(
-            self.prompts_dir / config.prompt["rag_sys_pmpt"]
-        )
-
         self.rag_mode = config.rag["mode"]
         self.chat_history_depth = config.rag["chat_history_depth"]
 
@@ -117,7 +106,7 @@ async def generate_query(self, messages: list[dict]) -> str:
                 response = await self.contextualizer.chat.completions.create(
                     model=self.config.vlm["model"],
                     messages=[
-                        {"role": "system", "content": self.contextualizer_pmpt},
+                        {"role": "system", "content": QUERY_CONTEXTUALIZER_PROMPT},
                         {
                             "role": "user",
                             "content": f"Given the following chat, generate a query. \n{chat_history}\n",
@@ -171,7 +160,7 @@ async def _prepare_for_chat_completion(self, partition: list[str], payload: dict
             0,
             {
                 "role": "system",
-                "content": self.rag_sys_prompt.format(context=context),
+                "content": SYS_PROMPT_TMPLT.format(context=context),
             },
         )
         payload["messages"] = messages

openrag/components/prompts/__init__.py

@@ -0,0 +1 @@
+from .prompts import *

openrag/components/prompts/prompts.py

@@ -0,0 +1,38 @@
+from pathlib import Path
+
+from config import load_config
+
+config = load_config()
+
+prompts_dir: Path = config.paths.prompts_dir
+prompt_mapping: dict = config.prompts
+
+
+def load_prompt(
+    prompt_name: str,
+    prompts_dir: Path = prompts_dir,
+    prompt_mapping: dict = prompt_mapping,
+) -> tuple[str, str]:
+    file_name = prompt_mapping.get(prompt_name, None)
+    if not file_name:
+        raise ValueError(f"No associated file name found for prompt: `{prompt_name}`")
+
+    file_path = prompts_dir / file_name
+
+    if not file_path.exists():
+        raise FileNotFoundError(f"Prompt file not found: `{file_path}`")
+
+    with open(file_path, mode="r") as f:
+        sys_msg = f.read()
+        return sys_msg
+
+
+# Load prompts
+SYS_PROMPT_TMPLT = load_prompt("sys_prompt")
+QUERY_CONTEXTUALIZER_PROMPT = load_prompt("query_contextualizer")
+CHUNK_CONTEXTUALIZER = load_prompt("chunk_contextualizer")
+IMAGE_DESCRIBER = load_prompt("image_describer")
+
+# Retrievers prompts
+HYDE_PROMPT = load_prompt("hyde")
+MULTI_QUERY_PROMPT = load_prompt("multi_query")

openrag/components/retriever.py

@@ -2,15 +2,14 @@
 from abc import ABC, abstractmethod
 from pathlib import Path
 
+from components.prompts import HYDE_PROMPT, MULTI_QUERY_PROMPT
 from langchain_core.documents.base import Document
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_openai import ChatOpenAI
 from omegaconf import OmegaConf
 from utils.dependencies import get_vectordb
 
-from .utils import load_sys_template
-
 CRITERIAS = ["similarity"]
 
 
@@ -114,12 +113,8 @@ def __init__(
                 raise TypeError(f"`k_queries` should be of type {int}")
             self.k_queries = k_queries
 
-            pmpt_tmpl_path = extra_args.get("prompts_dir") / extra_args.get(
-                "prompt_tmpl"
-            )
-            multi_query_tmpl = load_sys_template(pmpt_tmpl_path)
             prompt: ChatPromptTemplate = ChatPromptTemplate.from_template(
-                multi_query_tmpl
+                MULTI_QUERY_PROMPT
             )
             self.generate_queries = (
                 prompt | llm | StrOutputParser() | (lambda x: x.split("[SEP]"))
@@ -159,11 +154,7 @@ def __init__(
             if not isinstance(llm, ChatOpenAI):
                 raise TypeError(f"`llm` should be of type {ChatOpenAI}")
 
-            pmpt_tmpl_path = extra_args.get("prompts_dir") / extra_args.get(
-                "prompt_tmpl"
-            )
-            hyde_template = load_sys_template(pmpt_tmpl_path)
-            prompt: ChatPromptTemplate = ChatPromptTemplate.from_template(hyde_template)
+            prompt: ChatPromptTemplate = ChatPromptTemplate.from_template(HYDE_PROMPT)
 
             self.generate_hyde = prompt | llm | StrOutputParser()
             self.combine = extra_args.get("combine", False)

openrag/components/utils.py

@@ -2,7 +2,6 @@
 import atexit
 import threading
 from abc import ABCMeta
-from pathlib import Path
 
 import ray
 from config import load_config
@@ -119,12 +118,6 @@ def cleanup(self):
         ray.get(self._actor.cleanup.remote())
 
 
-def load_sys_template(file_path: Path) -> tuple[str, str]:
-    with open(file_path, mode="r") as f:
-        sys_msg = f.read()
-        return sys_msg
-
-
 def format_context(docs: list[Document]) -> str:
     if not docs:
         return "No document found from the database"
@@ -158,4 +151,4 @@ def get_vlm_semaphore() -> DistributedSemaphore:
 
 
 get_llm_semaphore()
-get_vlm_semaphore()
+get_vlm_semaphore()