Skip to content

Commit 7d80fc4

Browse files
authored
refine components retrieval and rewrite (infiniflow#2818)
### What problem does this PR solve? ### Type of change - [x] Performance Improvement
1 parent a20b820 commit 7d80fc4

File tree

2 files changed

+14
-7
lines changed

2 files changed

+14
-7
lines changed

agent/component/retrieval.py

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -50,12 +50,15 @@ class Retrieval(ComponentBase, ABC):
5050
component_name = "Retrieval"
5151

5252
def _run(self, history, **kwargs):
53-
query = []
54-
for role, cnt in history[::-1][:self._param.message_history_window_size]:
55-
if role != "user":continue
56-
query.append(cnt)
57-
# query = "\n".join(query)
58-
query = query[0]
53+
# query = []
54+
# for role, cnt in history[::-1][:self._param.message_history_window_size]:
55+
# if role != "user":continue
56+
# query.append(cnt)
57+
# # query = "\n".join(query)
58+
# query = query[0]
59+
query = self.get_input()
60+
query = str(query["content"][0]) if "content" in query else ""
61+
5962
kbs = KnowledgebaseService.get_by_ids(self._param.kb_ids)
6063
if not kbs:
6164
raise ValueError("Can't find knowledgebases by {}".format(self._param.kb_ids))

agent/component/rewrite.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,11 @@ def _run(self, history, **kwargs):
9191
raise Exception("Sorry! Nothing relevant found.")
9292
self._loop += 1
9393

94-
conv = self._canvas.get_history(4)
94+
hist = self._canvas.get_history(4)
95+
conv = []
96+
for m in hist:
97+
if m["role"] not in ["user", "assistant"]: continue
98+
conv.append("{}: {}".format(m["role"].upper(), m["content"]))
9599
conv = "\n".join(conv)
96100

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

0 commit comments

Comments
 (0)