| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| from abc import ABC |
| from api.db import LLMType |
| from api.db.services.llm_service import LLMBundle |
| from agent.component import GenerateParam, Generate |
| from agent.settings import DEBUG |
|
|
|
|
| class CategorizeParam(GenerateParam): |
|
|
| """ |
| Define the Categorize component parameters. |
| """ |
| def __init__(self): |
| super().__init__() |
| self.category_description = {} |
| self.prompt = "" |
|
|
| def check(self): |
| super().check() |
| self.check_empty(self.category_description, "[Categorize] Category examples") |
| for k, v in self.category_description.items(): |
| if not k: raise ValueError(f"[Categorize] Category name can not be empty!") |
| if not v.get("to"): raise ValueError(f"[Categorize] 'To' of category {k} can not be empty!") |
|
|
| def get_prompt(self): |
| cate_lines = [] |
| for c, desc in self.category_description.items(): |
| for l in desc.get("examples", "").split("\n"): |
| if not l: continue |
| cate_lines.append("Question: {}\tCategory: {}".format(l, c)) |
| descriptions = [] |
| for c, desc in self.category_description.items(): |
| if desc.get("description"): |
| descriptions.append( |
| "--------------------\nCategory: {}\nDescription: {}\n".format(c, desc["description"])) |
|
|
| self.prompt = """ |
| You're a text classifier. You need to categorize the user’s questions into {} categories, |
| namely: {} |
| Here's description of each category: |
| {} |
| |
| You could learn from the following examples: |
| {} |
| You could learn from the above examples. |
| Just mention the category names, no need for any additional words. |
| """.format( |
| len(self.category_description.keys()), |
| "/".join(list(self.category_description.keys())), |
| "\n".join(descriptions), |
| "- ".join(cate_lines) |
| ) |
| return self.prompt |
|
|
|
|
| class Categorize(Generate, ABC): |
| component_name = "Categorize" |
|
|
| def _run(self, history, **kwargs): |
| input = self.get_input() |
| input = "Question: " + ("; ".join(input["content"]) if "content" in input else "") + "Category: " |
| chat_mdl = LLMBundle(self._canvas.get_tenant_id(), LLMType.CHAT, self._param.llm_id) |
| ans = chat_mdl.chat(self._param.get_prompt(), [{"role": "user", "content": input}], |
| self._param.gen_conf()) |
| if DEBUG: print(ans, ":::::::::::::::::::::::::::::::::", input) |
| for c in self._param.category_description.keys(): |
| if ans.lower().find(c.lower()) >= 0: |
| return Categorize.be_output(self._param.category_description[c]["to"]) |
|
|
| return Categorize.be_output(self._param.category_description.items()[-1][1]["to"]) |
|
|
|
|
|
|
