diff --git a/petercat_utils/data_class.py b/petercat_utils/data_class.py index 27b03985..8b0b7a99 100644 --- a/petercat_utils/data_class.py +++ b/petercat_utils/data_class.py @@ -50,6 +50,7 @@ class ChatData(BaseModel): llm: Optional[str] = "openai" prompt: Optional[str] = None bot_id: Optional[str] = None + repo_name: Optional[str] = None class ExecuteMessage(BaseModel): diff --git a/server/agent/bot/__init__.py b/server/agent/bot/__init__.py index 06d1923c..2d958a05 100644 --- a/server/agent/bot/__init__.py +++ b/server/agent/bot/__init__.py @@ -2,11 +2,11 @@ from agent.llm import LLM, LLMTokenLike - class Bot: _bot: BotModel _llm: LLM _llm_token: LLMTokenLike + def __init__(self, bot: BotModel, llm_token: LLMTokenLike): self._bot = bot self._llm_token = llm_token @@ -15,7 +15,7 @@ def __init__(self, bot: BotModel, llm_token: LLMTokenLike): @property def id(self): return self._bot.id - + @property def prompt(self): return self._bot.prompt @@ -30,4 +30,8 @@ def llm(self): @property def token_id(self): - return self._bot.token_id \ No newline at end of file + return self._bot.token_id + + @property + def repo_name(self): + return self._bot.repo_name diff --git a/server/agent/qa_chat.py b/server/agent/qa_chat.py index f6b9ddf4..c9b5caac 100644 --- a/server/agent/qa_chat.py +++ b/server/agent/qa_chat.py @@ -16,7 +16,7 @@ def get_tools(bot: Bot, auth_token: Optional[Auth.Token]): return { "check_login": login_tools["check_login"], - "search_knowledge": knowledge.factory(repo_name=bot.repo_name), + "search_knowledge": knowledge.factory(bot_id=bot.id), "create_issue": issue_tools["create_issue"], "get_issues": issue_tools["get_issues"], "get_file_content": pull_request_tools["get_file_content"], diff --git a/server/agent/tools/knowledge.py b/server/agent/tools/knowledge.py index 8adf7433..a920ccd3 100644 --- a/server/agent/tools/knowledge.py +++ b/server/agent/tools/knowledge.py @@ -1,9 +1,10 @@ +from agent.bot.get_bot import get_bot_by_id from langchain.tools import tool from petercat_utils import retrieval -def factory(repo_name: str): - repo_name = repo_name +def factory(bot_id: str): + bot_id = bot_id @tool(parse_docstring=True) def search_knowledge( @@ -15,6 +16,8 @@ def search_knowledge( query: The user's question. """ try: + bot = get_bot_by_id(bot_id) + repo_name = bot.repo_name if bot.repo_name else "" return retrieval.search_knowledge(query, repo_name) except Exception as e: print(f"An error occurred: {e}") diff --git a/server/bot/router.py b/server/bot/router.py index b7fd1adc..a165a6c1 100644 --- a/server/bot/router.py +++ b/server/bot/router.py @@ -77,7 +77,7 @@ def get_bot_detail( data = ( supabase.table("bots") .select( - "id, created_at, updated_at, avatar, description, name, starters, public, hello_message" + "id, created_at, updated_at, avatar, description, name, starters, public, hello_message, repo_name" ) .eq("id", id) .execute() diff --git a/server/core/models/bot.py b/server/core/models/bot.py index 230cc55d..297edce6 100644 --- a/server/core/models/bot.py +++ b/server/core/models/bot.py @@ -6,6 +6,7 @@ class BotModel(BaseModel): id: str uid: str + repo_name: str avatar: Optional[str] = "" description: Optional[str] prompt: Optional[str] = "" @@ -15,7 +16,6 @@ class BotModel(BaseModel): token_id: Optional[str] = "" created_at: datetime = datetime.now() domain_whitelist: Optional[list[str]] = [] - repo_name: str = "" class RepoBindBotConfigVO(BaseModel): diff --git a/server/github_app/router.py b/server/github_app/router.py index bf5dbd07..bae7dc94 100644 --- a/server/github_app/router.py +++ b/server/github_app/router.py @@ -106,7 +106,6 @@ async def github_app_webhook( if "installation" not in payload: return {"success": False, "message": "Invalid Webhook request"} - print(f"payload: {payload}") installation_id = payload["installation"]["id"] try: auth = Auth.AppAuth( diff --git a/subscriber/requirements.txt b/subscriber/requirements.txt index ad3699c1..23c995c0 100644 --- a/subscriber/requirements.txt +++ b/subscriber/requirements.txt @@ -1 +1 @@ -petercat_utils>=0.1.36 +petercat_utils>=0.1.39