refactor: remove the dependency on bot_id in the task scheduling process

petercat_utils/data_class.py

@@ -24,6 +24,7 @@ class ImageURLContentBlock(BaseModel):
     image_url: ImageURL
     type: Literal["image_url"]
 
+
 class ImageRawURLContentBlock(BaseModel):
     image_url: str
     type: Literal["image_url"]
@@ -71,7 +72,7 @@ class GitDocConfig(BaseModel):
 
 
 class RAGGitDocConfig(GitDocConfig):
-    bot_id: str
+    bot_id: Optional[str] = ""
 
 
 class GitIssueConfig(BaseModel):

petercat_utils/rag_helper/git_doc_task.py

@@ -48,7 +48,6 @@ def add_rag_git_doc_task(config: RAGGitDocConfig, extra=None):
         sha=sha,
         repo_name=config.repo_name,
         node_type=extra["node_type"],
-        bot_id=config.bot_id,
         path=config.file_path,
     )
     res = doc_task.save()
@@ -64,15 +63,13 @@ def __init__(
         commit_id,
         node_type: GitDocTaskNodeType,
         sha,
-        bot_id,
         path,
         repo_name,
         status=TaskStatus.NOT_STARTED,
         from_id=None,
         id=None,
     ):
         super().__init__(
-            bot_id=bot_id,
             type=TaskType.GIT_DOC,
             from_id=from_id,
             id=id,
@@ -127,7 +124,6 @@ def handle_tree_node(self):
                     repo_name=record["repo_name"],
                     node_type=record["node_type"],
                     path=record["path"],
-                    bot_id=self.bot_id,
                 )
                 doc_task.send()
 
@@ -137,7 +133,6 @@ def handle_tree_node(self):
                 {
                     "metadata": {
                         "tree": list(map(lambda item: item.raw_data, tree_data.tree)),
-                        "bot_id": self.bot_id,
                     },
                     "status": TaskStatus.COMPLETED.value,
                 }
@@ -152,7 +147,6 @@ def handle_blob_node(self):
                 repo_name=self.repo_name,
                 file_path=self.path,
                 commit_id=self.commit_id,
-                bot_id=self.bot_id,
             )
         )
         return self.update_status(TaskStatus.COMPLETED)

petercat_utils/rag_helper/git_task.py

@@ -21,7 +21,6 @@ def __init__(
         self,
         type,
         repo_name,
-        bot_id,
         status=TaskStatus.NOT_STARTED,
         from_id=None,
         id=None,
@@ -31,7 +30,6 @@ def __init__(
         self.from_id = from_id
         self.status = status
         self.repo_name = repo_name
-        self.bot_id = bot_id
 
     @staticmethod
     def get_table_name(type: TaskType):

petercat_utils/rag_helper/retrieval.py

@@ -164,7 +164,6 @@ def reload_knowledge(config: RAGGitDocConfig):
         commit_id=loader.commit_id,
         file_sha=loader.file_sha,
         file_path=config.file_path,
-        bot_id=config.bot_id,
     )
     return store
 

petercat_utils/rag_helper/task.py

@@ -57,7 +57,6 @@ def get_task_by_id(task_id):
 def get_task(
     task_type: TaskType,
     task_id: str,
-    bot_id: Optional[str],
 ) -> GitTask:
     supabase = get_client()
     response = (
@@ -75,7 +74,6 @@ def get_task(
                 sha=data["sha"],
                 repo_name=data["repo_name"],
                 node_type=data["node_type"],
-                bot_id=bot_id,
                 path=data["path"],
                 status=data["status"],
                 from_id=data["from_task_id"],
@@ -92,8 +90,8 @@ def get_task(
             )
 
 
-def trigger_task(task_type: TaskType, bot_id: Optional[str], task_id: Optional[str]):
-    task = get_task(task_type, bot_id, task_id) if task_id else get_oldest_task()
+def trigger_task(task_type: TaskType, task_id: Optional[str]):
+    task = get_task(task_type, task_id) if task_id else get_oldest_task()
     if task is None:
         return task
     return task.handle()

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "petercat_utils"
-version = "0.1.38"
+version = "0.1.39"
 description = ""
 authors = ["raoha.rh <[email protected]>"]
 readme = "README.md"

server/rag/router.py

@@ -4,7 +4,11 @@
 from fastapi import APIRouter, Depends
 from petercat_utils.db.client.supabase import get_client
 
-from petercat_utils.data_class import RAGGitDocConfig, RAGGitIssueConfig, TaskType
+from petercat_utils.data_class import (
+    GitDocConfig,
+    RAGGitIssueConfig,
+    TaskType,
+)
 from petercat_utils.rag_helper import (
     retrieval,
     task,
@@ -24,7 +28,7 @@
 
 
 @router.post("/rag/add_knowledge_by_doc", dependencies=[Depends(verify_rate_limit)])
-def add_knowledge_by_doc(config: RAGGitDocConfig):
+def add_knowledge_by_doc(config: GitDocConfig):
     try:
         result = retrieval.add_knowledge_by_doc(config)
         if result:
@@ -64,7 +68,7 @@ def search_knowledge(query: str, repo_name: str, filter: dict = {}):
 
 
 @router.post("/rag/add_git_doc_task", dependencies=[Depends(verify_rate_limit)])
-def add_git_doc_task(config: RAGGitDocConfig):
+def add_git_doc_task(config: GitDocConfig):
     try:
         data = git_doc_task.add_rag_git_doc_task(config)
         return data
@@ -82,9 +86,9 @@ def add_git_issue_task(config: RAGGitIssueConfig):
 
 
 @router.post("/rag/trigger_task", dependencies=[Depends(verify_rate_limit)])
-def trigger_task(task_type: TaskType, bot_id: str, task_id: Optional[str] = None):
+def trigger_task(task_type: TaskType, task_id: Optional[str] = None):
     try:
-        task.trigger_task(task_type, task_id, bot_id)
+        task.trigger_task(task_type, task_id)
     except Exception as e:
         return json.dumps({"success": False, "message": str(e)})