feat: editable chunk params

This commit is contained in:
Timothy J. Baek 2024-02-17 22:29:52 -08:00
parent a32ab5afbd
commit ccf08fb91e
3 changed files with 152 additions and 7 deletions

View file

@ -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):

View file

@ -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);

View file

@ -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 justify-end pt-3 text-sm font-medium">
<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">
<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>