Spaces:
Paused
Paused
| # | |
| # Copyright 2024 The InfiniFlow Authors. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # | |
| from abc import ABC | |
| from api.db import LLMType | |
| from api.db.services.llm_service import LLMBundle | |
| from agent.component import GenerateParam, Generate | |
| from rag.utils import num_tokens_from_string, encoder | |
| class RelevantParam(GenerateParam): | |
| """ | |
| Define the Relevant component parameters. | |
| """ | |
| def __init__(self): | |
| super().__init__() | |
| self.prompt = "" | |
| self.yes = "" | |
| self.no = "" | |
| def check(self): | |
| super().check() | |
| self.check_empty(self.yes, "[Relevant] 'Yes'") | |
| self.check_empty(self.no, "[Relevant] 'No'") | |
| def get_prompt(self): | |
| self.prompt = """ | |
| You are a grader assessing relevance of a retrieved document to a user question. | |
| It does not need to be a stringent test. The goal is to filter out erroneous retrievals. | |
| If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. | |
| Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question. | |
| No other words needed except 'yes' or 'no'. | |
| """ | |
| return self.prompt | |
| class Relevant(Generate, ABC): | |
| component_name = "Relevant" | |
| def _run(self, history, **kwargs): | |
| q = "" | |
| for r, c in self._canvas.history[::-1]: | |
| if r == "user": | |
| q = c | |
| break | |
| ans = self.get_input() | |
| ans = " - ".join(ans["content"]) if "content" in ans else "" | |
| if not ans: | |
| return Relevant.be_output(self._param.no) | |
| ans = "Documents: \n" + ans | |
| ans = f"Question: {q}\n" + ans | |
| chat_mdl = LLMBundle(self._canvas.get_tenant_id(), LLMType.CHAT, self._param.llm_id) | |
| if num_tokens_from_string(ans) >= chat_mdl.max_length - 4: | |
| ans = encoder.decode(encoder.encode(ans)[:chat_mdl.max_length - 4]) | |
| ans = chat_mdl.chat(self._param.get_prompt(), [{"role": "user", "content": ans}], | |
| self._param.gen_conf()) | |
| print(ans, ":::::::::::::::::::::::::::::::::") | |
| if ans.lower().find("yes") >= 0: | |
| return Relevant.be_output(self._param.yes) | |
| if ans.lower().find("no") >= 0: | |
| return Relevant.be_output(self._param.no) | |
| assert False, f"Relevant component got: {ans}" | |