Bläddra i källkod

feat: editable chunk params

Timothy J. Baek 1 år sedan
förälder
incheckning
ccf08fb91e

+ 37 - 2
backend/apps/rag/main.py

@@ -71,6 +71,9 @@ from constants import ERROR_MESSAGES
 
 app = FastAPI()
 
+app.state.CHUNK_SIZE = CHUNK_SIZE
+app.state.CHUNK_OVERLAP = CHUNK_OVERLAP
+
 origins = ["*"]
 
 app.add_middleware(
@@ -92,7 +95,7 @@ class StoreWebForm(CollectionNameForm):
 
 def store_data_in_vector_db(data, collection_name) -> bool:
     text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP
+        chunk_size=app.state.CHUNK_SIZE, chunk_overlap=app.state.CHUNK_OVERLAP
     )
     docs = text_splitter.split_documents(data)
 
@@ -116,7 +119,39 @@ def store_data_in_vector_db(data, collection_name) -> bool:
 
 @app.get("/")
 async def get_status():
-    return {"status": True}
+    return {
+        "status": True,
+        "chunk_size": app.state.CHUNK_SIZE,
+        "chunk_overlap": app.state.CHUNK_OVERLAP,
+    }
+
+
+@app.get("/chunk")
+async def get_chunk_params(user=Depends(get_admin_user)):
+    return {
+        "status": True,
+        "chunk_size": app.state.CHUNK_SIZE,
+        "chunk_overlap": app.state.CHUNK_OVERLAP,
+    }
+
+
+class ChunkParamUpdateForm(BaseModel):
+    chunk_size: int
+    chunk_overlap: int
+
+
+@app.post("/chunk/update")
+async def update_chunk_params(
+    form_data: ChunkParamUpdateForm, user=Depends(get_admin_user)
+):
+    app.state.CHUNK_SIZE = form_data.chunk_size
+    app.state.CHUNK_OVERLAP = form_data.chunk_overlap
+
+    return {
+        "status": True,
+        "chunk_size": app.state.CHUNK_SIZE,
+        "chunk_overlap": app.state.CHUNK_OVERLAP,
+    }
 
 
 class QueryDocForm(BaseModel):

+ 58 - 0
src/lib/apis/rag/index.ts

@@ -1,5 +1,63 @@
 import { RAG_API_BASE_URL } from '$lib/constants';
 
+export const getChunkParams = async (token: string) => {
+	let error = null;
+
+	const res = await fetch(`${RAG_API_BASE_URL}/chunk`, {
+		method: 'GET',
+		headers: {
+			'Content-Type': 'application/json',
+			Authorization: `Bearer ${token}`
+		}
+	})
+		.then(async (res) => {
+			if (!res.ok) throw await res.json();
+			return res.json();
+		})
+		.catch((err) => {
+			console.log(err);
+			error = err.detail;
+			return null;
+		});
+
+	if (error) {
+		throw error;
+	}
+
+	return res;
+};
+
+export const updateChunkParams = async (token: string, size: number, overlap: number) => {
+	let error = null;
+
+	const res = await fetch(`${RAG_API_BASE_URL}/chunk/update`, {
+		method: 'POST',
+		headers: {
+			'Content-Type': 'application/json',
+			Authorization: `Bearer ${token}`
+		},
+		body: JSON.stringify({
+			chunk_size: size,
+			chunk_overlap: overlap
+		})
+	})
+		.then(async (res) => {
+			if (!res.ok) throw await res.json();
+			return res.json();
+		})
+		.catch((err) => {
+			console.log(err);
+			error = err.detail;
+			return null;
+		});
+
+	if (error) {
+		throw error;
+	}
+
+	return res;
+};
+
 export const uploadDocToVectorDB = async (token: string, collection_name: string, file: File) => {
 	const data = new FormData();
 	data.append('file', file);

+ 57 - 5
src/lib/components/documents/Settings/General.svelte

@@ -1,6 +1,6 @@
 <script lang="ts">
 	import { getDocs } from '$lib/apis/documents';
-	import { scanDocs } from '$lib/apis/rag';
+	import { getChunkParams, scanDocs, updateChunkParams } from '$lib/apis/rag';
 	import { documents } from '$lib/stores';
 	import { onMount } from 'svelte';
 	import toast from 'svelte-french-toast';
@@ -9,6 +9,9 @@
 
 	let loading = false;
 
+	let chunkSize = 0;
+	let chunkOverlap = 0;
+
 	const scanHandler = async () => {
 		loading = true;
 		const res = await scanDocs(localStorage.token);
@@ -20,13 +23,24 @@
 		}
 	};
 
-	onMount(async () => {});
+	const submitHandler = async () => {
+		const res = await updateChunkParams(localStorage.token, chunkSize, chunkOverlap);
+	};
+
+	onMount(async () => {
+		const res = await getChunkParams(localStorage.token);
+
+		if (res) {
+			chunkSize = res.chunk_size;
+			chunkOverlap = res.chunk_overlap;
+		}
+	});
 </script>
 
 <form
 	class="flex flex-col h-full justify-between space-y-3 text-sm"
 	on:submit|preventDefault={() => {
-		// console.log('submit');
+		submitHandler();
 		saveHandler();
 	}}
 >
@@ -93,14 +107,52 @@
 				</button>
 			</div>
 		</div>
+
+		<hr class=" dark:border-gray-700" />
+
+		<div class=" ">
+			<div class=" text-sm font-medium">Chunk Params</div>
+
+			<div class=" flex">
+				<div class="  flex w-full justify-between">
+					<div class="self-center text-xs font-medium min-w-fit">Chunk Size</div>
+
+					<div class="self-center p-3">
+						<input
+							class=" w-full rounded py-1.5 px-4 text-sm dark:text-gray-300 dark:bg-gray-800 outline-none border border-gray-100 dark:border-gray-600"
+							type="number"
+							placeholder="Enter Chunk Size"
+							bind:value={chunkSize}
+							autocomplete="off"
+							min="0"
+						/>
+					</div>
+				</div>
+
+				<div class="flex w-full">
+					<div class=" self-center text-xs font-medium min-w-fit">Chunk Overlap</div>
+
+					<div class="self-center p-3">
+						<input
+							class="w-full rounded py-1.5 px-4 text-sm dark:text-gray-300 dark:bg-gray-800 outline-none border border-gray-100 dark:border-gray-600"
+							type="number"
+							placeholder="Enter Chunk Overlap"
+							bind:value={chunkOverlap}
+							autocomplete="off"
+							min="0"
+						/>
+					</div>
+				</div>
+			</div>
+		</div>
 	</div>
 
-	<!-- <div class="flex justify-end pt-3 text-sm font-medium">
+	<div class="flex justify-end pt-3 text-sm font-medium">
 		<button
 			class=" px-4 py-2 bg-emerald-600 hover:bg-emerald-700 text-gray-100 transition rounded"
 			type="submit"
 		>
 			Save
 		</button>
-	</div> -->
+	</div>
 </form>