Skip to content

Commit

Permalink
refine components retrieval and rewrite (#2818)
Browse files Browse the repository at this point in the history
### What problem does this PR solve?


### Type of change

- [x] Performance Improvement
  • Loading branch information
KevinHuSh authored Oct 12, 2024
1 parent a20b820 commit 7d80fc4
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 7 deletions.
15 changes: 9 additions & 6 deletions agent/component/retrieval.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,12 +50,15 @@ class Retrieval(ComponentBase, ABC):
component_name = "Retrieval"

def _run(self, history, **kwargs):
query = []
for role, cnt in history[::-1][:self._param.message_history_window_size]:
if role != "user":continue
query.append(cnt)
# query = "\n".join(query)
query = query[0]
# query = []
# for role, cnt in history[::-1][:self._param.message_history_window_size]:
# if role != "user":continue
# query.append(cnt)
# # query = "\n".join(query)
# query = query[0]
query = self.get_input()
query = str(query["content"][0]) if "content" in query else ""

kbs = KnowledgebaseService.get_by_ids(self._param.kb_ids)
if not kbs:
raise ValueError("Can't find knowledgebases by {}".format(self._param.kb_ids))
Expand Down
6 changes: 5 additions & 1 deletion agent/component/rewrite.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,11 @@ def _run(self, history, **kwargs):
raise Exception("Sorry! Nothing relevant found.")
self._loop += 1

conv = self._canvas.get_history(4)
hist = self._canvas.get_history(4)
conv = []
for m in hist:
if m["role"] not in ["user", "assistant"]: continue
conv.append("{}: {}".format(m["role"].upper(), m["content"]))
conv = "\n".join(conv)

chat_mdl = LLMBundle(self._canvas.get_tenant_id(), LLMType.CHAT, self._param.llm_id)
Expand Down

0 comments on commit 7d80fc4

Please sign in to comment.