Skip to content

Commit

Permalink
refactor: remove the dependency on bot_id in the task scheduling proc…
Browse files Browse the repository at this point in the history
…ess (#441)

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

* chore: fix typo
  • Loading branch information
xingwanying authored Oct 16, 2024
1 parent 1772e73 commit 0f08935
Show file tree
Hide file tree
Showing 7 changed files with 12 additions and 18 deletions.
3 changes: 2 additions & 1 deletion petercat_utils/data_class.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ class ImageURLContentBlock(BaseModel):
image_url: ImageURL
type: Literal["image_url"]


class ImageRawURLContentBlock(BaseModel):
image_url: str
type: Literal["image_url"]
Expand Down Expand Up @@ -71,7 +72,7 @@ class GitDocConfig(BaseModel):


class RAGGitDocConfig(GitDocConfig):
bot_id: str
bot_id: Optional[str] = ""


class GitIssueConfig(BaseModel):
Expand Down
6 changes: 0 additions & 6 deletions petercat_utils/rag_helper/git_doc_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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,
Expand Down Expand Up @@ -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()

Expand All @@ -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,
}
Expand All @@ -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)
Expand Down
2 changes: 0 additions & 2 deletions petercat_utils/rag_helper/git_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ def __init__(
self,
type,
repo_name,
bot_id,
status=TaskStatus.NOT_STARTED,
from_id=None,
id=None,
Expand All @@ -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):
Expand Down
1 change: 0 additions & 1 deletion petercat_utils/rag_helper/retrieval.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
6 changes: 2 additions & 4 deletions petercat_utils/rag_helper/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = (
Expand All @@ -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"],
Expand All @@ -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()
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
10 changes: 7 additions & 3 deletions server/rag/router.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
RAGGitDocConfig,
RAGGitIssueConfig,
TaskType,
)
from petercat_utils.rag_helper import (
retrieval,
task,
Expand Down Expand Up @@ -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)})

Expand Down

0 comments on commit 0f08935

Please sign in to comment.