diff --git a/petercat_utils/data_class.py b/petercat_utils/data_class.py index 4426baec..27b03985 100644 --- a/petercat_utils/data_class.py +++ b/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): diff --git a/petercat_utils/rag_helper/git_doc_task.py b/petercat_utils/rag_helper/git_doc_task.py index 9e6caa18..74d16e0d 100644 --- a/petercat_utils/rag_helper/git_doc_task.py +++ b/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,7 +63,6 @@ def __init__( commit_id, node_type: GitDocTaskNodeType, sha, - bot_id, path, repo_name, status=TaskStatus.NOT_STARTED, @@ -72,7 +70,6 @@ def __init__( 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) diff --git a/petercat_utils/rag_helper/git_task.py b/petercat_utils/rag_helper/git_task.py index cd22fa88..61d6cf58 100644 --- a/petercat_utils/rag_helper/git_task.py +++ b/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): diff --git a/petercat_utils/rag_helper/retrieval.py b/petercat_utils/rag_helper/retrieval.py index a9762f09..5a3207dd 100644 --- a/petercat_utils/rag_helper/retrieval.py +++ b/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 diff --git a/petercat_utils/rag_helper/task.py b/petercat_utils/rag_helper/task.py index 98c3990a..144632ab 100644 --- a/petercat_utils/rag_helper/task.py +++ b/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() diff --git a/pyproject.toml b/pyproject.toml index 51413e33..ef116755 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "petercat_utils" -version = "0.1.38" +version = "0.1.39" description = "" authors = ["raoha.rh "] readme = "README.md" diff --git a/server/rag/router.py b/server/rag/router.py index 65efe4ee..5fbb69ca 100644 --- a/server/rag/router.py +++ b/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)})