feat: editable rag template

backend/apps/rag/main.py

@@ -62,6 +62,7 @@ from config import (
     CHROMA_CLIENT,
     CHUNK_SIZE,
     CHUNK_OVERLAP,
+    RAG_TEMPLATE,
 )
 from constants import ERROR_MESSAGES
 
@@ -73,6 +74,8 @@ app = FastAPI()
 
 app.state.CHUNK_SIZE = CHUNK_SIZE
 app.state.CHUNK_OVERLAP = CHUNK_OVERLAP
+app.state.RAG_TEMPLATE = RAG_TEMPLATE
+
 
 origins = ["*"]
 
@@ -154,6 +157,25 @@ async def update_chunk_params(
     }
 
 
+@app.get("/template")
+async def get_rag_template(user=Depends(get_current_user)):
+    return {
+        "status": True,
+        "template": app.state.RAG_TEMPLATE,
+    }
+
+
+class RAGTemplateForm(BaseModel):
+    template: str
+
+
+@app.post("/template/update")
+async def update_rag_template(form_data: RAGTemplateForm, user=Depends(get_admin_user)):
+    # TODO: check template requirements
+    app.state.RAG_TEMPLATE = form_data.template
+    return {"status": True, "template": app.state.RAG_TEMPLATE}
+
+
 class QueryDocForm(BaseModel):
     collection_name: str
     query: str

backend/config.py

@@ -144,6 +144,21 @@ CHROMA_CLIENT = chromadb.PersistentClient(
 CHUNK_SIZE = 1500
 CHUNK_OVERLAP = 100
 
+
+RAG_TEMPLATE = """Use the following context as your learned knowledge, inside <context></context> XML tags.
+<context>
+    [context]
+</context>
+
+When answer to user:
+- If you don't know, just say that you don't know.
+- If you don't know when you are not sure, ask for clarification.
+Avoid mentioning that you obtained the information from the context.
+And answer according to the language of the user's question.
+        
+Given the context information, answer the query.
+Query: [query]"""
+
 ####################################
 # Transcribe
 ####################################