Skip to content

Commit

Permalink
feat: implement knowledge update functionality for bots
Browse files Browse the repository at this point in the history
- Added a new hook `useKnowledgeUpdate` to manage knowledge updates.
- Integrated knowledge update feature in `BotCard` component with loading state.
- Created `updateKnowledge` API in `BotsController` to handle knowledge updates.
- Updated FastAPI router to include endpoint for knowledge updates, validating bot existence.
- Refactored retrieval logic to check and update knowledge based on document changes.
  • Loading branch information
sedationh committed Dec 1, 2024
1 parent bc76d17 commit 4f69245
Show file tree
Hide file tree
Showing 6 changed files with 130 additions and 11 deletions.
43 changes: 34 additions & 9 deletions client/app/factory/list/components/BotCard.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import ErrorBadgeIcon from '@/public/icons/ErrorBadgeIcon';
import CheckBadgeIcon from '@/public/icons/CheckBadgeIcon';
import LoadingIcon from '@/public/icons/LoadingIcon';
import { RagTask } from '@/app/services/BotsController';
import { useKnowledgeUpdate } from '@/app/hooks/useKnowledgeUpdate';

declare type Bot = Tables<'bots'>;

Expand All @@ -34,6 +35,8 @@ const BotCard = (props: { bot: Bot }) => {
const router = useRouter();
const { deleteBot, isLoading, isSuccess } = useBotDelete();
const { data: taskInfo } = useGetBotRagTask(bot.id, true, false);
const { mutate: updateKnowledge, isPending: isUpdating } =
useKnowledgeUpdate();

useEffect(() => {
if (isSuccess) {
Expand Down Expand Up @@ -69,6 +72,19 @@ const BotCard = (props: { bot: Bot }) => {
);
};

const handleUpdateKnowledge = () => {
updateKnowledge(
{
bot_id: bot.id,
},
{
onSuccess: () => {
// TODO
},
},
);
};

return (
<>
<Card
Expand Down Expand Up @@ -124,21 +140,30 @@ const BotCard = (props: { bot: Bot }) => {
placement="top"
content={I18N.components.BotCard.gengXinZhiShiKu}
classNames={{
base: [
// arrow color
'before:bg-[#3F3F46] dark:before:bg-white',
],
base: ['before:bg-[#3F3F46] dark:before:bg-white'],
content: [
'py-2 px-4 rounded-lg shadow-xl text-white',
'bg-[#3F3F46]',
],
}}
>
<Image
src="../images/refresh.svg"
alt={I18N.components.BotCard.gengXinZhiShi}
className="z-10 cursor-pointer"
/>
{isUpdating ? (
<div className="w-6 h-6 flex items-center justify-center">
<span className="animate-spinner-ease-spin">
<LoadingIcon />
</span>
</div>
) : (
<Image
src="../images/refresh.svg"
onClick={(e) => {
e.stopPropagation();
handleUpdateKnowledge();
}}
alt={I18N.components.BotCard.gengXinZhiShi}
className="z-10 cursor-pointer"
/>
)}
</Tooltip>
</div>
</div>
Expand Down
9 changes: 9 additions & 0 deletions client/app/hooks/useKnowledgeUpdate.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
import { useMutation } from '@tanstack/react-query';
import { updateKnowledge } from '../services/BotsController';

export function useKnowledgeUpdate() {
return useMutation({
mutationKey: ['updateKnowledge'],
mutationFn: updateKnowledge
});
}
7 changes: 7 additions & 0 deletions client/app/services/BotsController.ts
Original file line number Diff line number Diff line change
Expand Up @@ -135,3 +135,10 @@ export async function bindBotToRepo(repsConfigs: BindBotToRepoConfig[]) {
});
return response.data;
}

// Add knowledge update API
export async function updateKnowledge(config: {
bot_id: string;
}) {
return axios.post(`${apiDomain}/api/rag/update_knowledge`, config);
}
2 changes: 1 addition & 1 deletion petercat_utils/rag_helper/git_doc_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,7 @@ def handle_tree_node(self):
)

def handle_blob_node(self):
retrieval.add_knowledge_by_doc(
retrieval.check_and_update_knowledge(
RAGGitDocConfig(
repo_name=self.repo_name,
file_path=self.path,
Expand Down
37 changes: 37 additions & 0 deletions petercat_utils/rag_helper/retrieval.py
Original file line number Diff line number Diff line change
Expand Up @@ -197,3 +197,40 @@ def get_chunk_list(repo_name: str, page_size: int, page_number: int):
)
total_count = len(count_response.data)
return {"rows": query.data, "total": total_count}


def check_and_update_knowledge(config: RAGGitDocConfig):
# 初始化 GitHub loader 获取最新的文件信息
loader = init_github_file_loader(config)
latest_sha = loader.file_sha

# 获取当前存储的文档
client = get_client()
existing_docs = (
client.table(TABLE_NAME)
.select("id, file_sha")
.eq("repo_name", config.repo_name)
.eq("file_path", config.file_path)
.execute()
)

if not existing_docs.data:
# 如果不存在文档,直接添加
return add_knowledge_by_doc(config)

# 检查 SHA 是否变化
current_sha = existing_docs.data[0]["file_sha"]

if current_sha == latest_sha:
return False

# SHA 不同,需要更新
# 1. 删除旧文档
client.table(TABLE_NAME)\
.delete()\
.eq("repo_name", config.repo_name)\
.eq("file_path", config.file_path)\
.execute()

# 2. 添加新文档
return add_knowledge_by_doc(config)
43 changes: 42 additions & 1 deletion server/rag/router.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import json
from typing import Optional
from typing import Annotated, Optional
from pydantic import BaseModel

from auth.get_user_info import get_user_id
from fastapi import APIRouter, Depends
from petercat_utils.db.client.supabase import get_client

Expand Down Expand Up @@ -116,3 +118,42 @@ def get_rag_task(repo_name: str):
return response
except Exception as e:
return json.dumps({"success": False, "message": str(e)})

class UpdateKnowledgeRequest(BaseModel):
bot_id: str

@router.post("/rag/update_knowledge", dependencies=[Depends(verify_rate_limit)])
def update_knowledge(request: UpdateKnowledgeRequest, user_id: Annotated[str | None, Depends(get_user_id)] = None):
try:
# Get config from database using bot_id
supabase = get_client()
response = (
supabase.table("bots")
.select("*")
.eq("id", request.bot_id)
.eq("uid", user_id)
.single()
.execute()
)

if not response.data:
return json.dumps({
"success": False,
"message": f"Bot with id {request.bot_id} not found"
})

bot_config = RAGGitDocConfig(**response.data)
result = retrieval.check_and_update_knowledge(bot_config)

if result:
return json.dumps({
"success": True,
"message": "Knowledge updated successfully!"
})
else:
return json.dumps({
"success": False,
"message": "Knowledge not updated!"
})
except Exception as e:
return json.dumps({"success": False, "message": str(e)})

0 comments on commit 4f69245

Please sign in to comment.