From 57cfb9ee76a64ab3948901085cc99bf74fe535b9 Mon Sep 17 00:00:00 2001 From: yingying Date: Wed, 21 Aug 2024 09:07:27 +0800 Subject: [PATCH] fix: fix qa_chat --- server/agent/qa_chat.py | 9 +++--- server/event_handler/discussion.py | 2 +- server/event_handler/issue.py | 33 ++++++++++------------ server/routers/chat.py | 44 ++++++++++++++++++++++++------ 4 files changed, 56 insertions(+), 32 deletions(-) diff --git a/server/agent/qa_chat.py b/server/agent/qa_chat.py index 1f9d2c2e..c55ba53b 100644 --- a/server/agent/qa_chat.py +++ b/server/agent/qa_chat.py @@ -1,4 +1,4 @@ -from typing import AsyncIterator +from typing import AsyncIterator, Optional from agent.base import AgentBuilder from prompts.bot_template import generate_prompt_by_repo_name from petercat_utils import get_client @@ -6,7 +6,7 @@ from tools import issue, sourcecode, knowledge, git_info -def get_tools(bot_id, token): +def get_tools(bot_id: str, token: Optional[str]): issue_tools = issue.factory(access_token=token) return { "search_knowledge": knowledge.factory(bot_id=bot_id), @@ -49,8 +49,9 @@ def agent_stream_chat(input_data: ChatData, user_token: str) -> AsyncIterator[st return agent.run_stream_chat(input_data) -def agent_chat(input_data: ChatData) -> AsyncIterator[str]: +def agent_chat(input_data: ChatData, user_token: Optional[str]) -> AsyncIterator[str]: agent = AgentBuilder( - prompt=init_prompt(input_data), tools=get_tools(input_data.bot_id) + prompt=init_prompt(input_data), + tools=get_tools(input_data.bot_id, token=user_token), ) return agent.run_chat(input_data) diff --git a/server/event_handler/discussion.py b/server/event_handler/discussion.py index 2820c017..f5061d63 100644 --- a/server/event_handler/discussion.py +++ b/server/event_handler/discussion.py @@ -88,7 +88,7 @@ async def handle_discussion_event(self, action: str): discussion_number = discussion["number"] message = Message(role="user", content=[text_block]) - analysis_result = await agent_chat(ChatData(messages=[message])) + analysis_result = await agent_chat(ChatData(messages=[message]), None) discussion_id = await self.get_discussion_id(owner, repo, discussion_number) await self.create_discussion_comment(discussion_id, analysis_result["output"]) diff --git a/server/event_handler/issue.py b/server/event_handler/issue.py index a04f2f0a..84d6b431 100644 --- a/server/event_handler/issue.py +++ b/server/event_handler/issue.py @@ -1,4 +1,3 @@ - from typing import Any from github import Github, Auth from github import GithubException @@ -6,7 +5,8 @@ from petercat_utils.data_class import ChatData, Message, TextContentBlock -class IssueEventHandler(): + +class IssueEventHandler: event: Any auth: Auth.AppAuth g: Github @@ -18,23 +18,20 @@ def __init__(self, payload: Any, auth: Auth.AppAuth) -> None: async def execute(self): try: - print('actions:', self.event['action']) - if self.event['action'] == "opened": - repo_name = self.event['repository']["full_name"] - issue_number = self.event["issue"]["number"] - repo = self.g.get_repo(repo_name) - issue = repo.get_issue(number=issue_number) - issue_content = f"{issue.title}: {issue.body}" - text_block = TextContentBlock( - type="text", - text=issue_content - ) - issue_content = issue.body - message = Message(role="user", content=[text_block]) - analysis_result = await agent_chat(ChatData(messages=[message])) - issue.create_comment(analysis_result['output']) + print("actions:", self.event["action"]) + if self.event["action"] == "opened": + repo_name = self.event["repository"]["full_name"] + issue_number = self.event["issue"]["number"] + repo = self.g.get_repo(repo_name) + issue = repo.get_issue(number=issue_number) + issue_content = f"{issue.title}: {issue.body}" + text_block = TextContentBlock(type="text", text=issue_content) + issue_content = issue.body + message = Message(role="user", content=[text_block]) + analysis_result = await agent_chat(ChatData(messages=[message]), None) + issue.create_comment(analysis_result["output"]) - return {"success": True} + return {"success": True} except GithubException as e: print(f"处理 GitHub 请求时出错：{e}") return {"success": False, "error": str(e)} diff --git a/server/routers/chat.py b/server/routers/chat.py index 16ccb957..0baaf94b 100644 --- a/server/routers/chat.py +++ b/server/routers/chat.py @@ -14,26 +14,52 @@ responses={404: {"description": "Not found"}}, ) + async def generate_auth_failed_stream(): message = "Auth failed, please login first\n\n" yield message -@router.post("/stream_qa", response_class=StreamingResponse, dependencies=[Depends(verify_rate_limit)]) -def run_qa_chat(input_data: ChatData, user_access_token: Annotated[str | None, Depends(get_user_access_token)] = None): - print(f"run_qa_chat: input_data={input_data}, user_access_token={user_access_token}") - result = qa_chat.agent_stream_chat(input_data=input_data, user_token=user_access_token) + +@router.post( + "/stream_qa", + response_class=StreamingResponse, + dependencies=[Depends(verify_rate_limit)], +) +def run_qa_chat( + input_data: ChatData, + user_access_token: Annotated[str | None, Depends(get_user_access_token)] = None, +): + print( + f"run_qa_chat: input_data={input_data}, user_access_token={user_access_token}" + ) + result = qa_chat.agent_stream_chat( + input_data=input_data, user_token=user_access_token + ) return StreamingResponse(result, media_type="text/event-stream") @router.post("/qa", dependencies=[Depends(verify_rate_limit)]) -async def run_issue_helper(input_data: ChatData): - result = await qa_chat.agent_chat(input_data) +async def run_issue_helper( + input_data: ChatData, + user_access_token: Annotated[str | None, Depends(get_user_access_token)] = None, +): + result = await qa_chat.agent_chat(input_data, user_access_token) return result -@router.post("/stream_builder", response_class=StreamingResponse, dependencies=[Depends(verify_rate_limit)]) -def run_bot_builder(input_data: ChatData, bot_id: Optional[str] = None, user_id: Annotated[str | None, Depends(get_user_id)] = None): +@router.post( + "/stream_builder", + response_class=StreamingResponse, + dependencies=[Depends(verify_rate_limit)], +) +def run_bot_builder( + input_data: ChatData, + bot_id: Optional[str] = None, + user_id: Annotated[str | None, Depends(get_user_id)] = None, +): if not user_id: - return StreamingResponse(generate_auth_failed_stream(), media_type="text/event-stream") + return StreamingResponse( + generate_auth_failed_stream(), media_type="text/event-stream" + ) result = bot_builder.agent_stream_chat(input_data, user_id, bot_id) return StreamingResponse(result, media_type="text/event-stream")