From f1fbca3a80a35c3993071c387246681e013a3930 Mon Sep 17 00:00:00 2001 From: liblaf <30631553+liblaf@users.noreply.github.com> Date: Fri, 29 Nov 2024 15:14:41 +0800 Subject: [PATCH] refactor: improve tag extraction and type annotations - Ensures tag extraction is case-insensitive - Adds type annotations for better code clarity and maintainability --- src/llm_cli/utils/_extract_between_tags.py | 9 ++++++--- src/llm_cli/utils/_get_prompt.py | 5 +++-- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/src/llm_cli/utils/_extract_between_tags.py b/src/llm_cli/utils/_extract_between_tags.py index c1289e4..9195290 100644 --- a/src/llm_cli/utils/_extract_between_tags.py +++ b/src/llm_cli/utils/_extract_between_tags.py @@ -1,11 +1,14 @@ def extract_between_tags(content: str | None, tag: str = "Answer") -> str: if content is None: return "" - start: int = content.find("<" + tag + ">") + content_lower: str = content.lower() + opening_tag: str = f"<{tag}>".lower() + start: int = content_lower.find(opening_tag) if start >= 0: - start += len(tag) + 2 + start += len(opening_tag) content = content[start:] - end: int = content.find("") + closing_tag: str = f"".lower() + end: int = content_lower.find(closing_tag) if end >= 0: content = content[:end] return content.strip() diff --git a/src/llm_cli/utils/_get_prompt.py b/src/llm_cli/utils/_get_prompt.py index c0c4f00..f3b1bac 100644 --- a/src/llm_cli/utils/_get_prompt.py +++ b/src/llm_cli/utils/_get_prompt.py @@ -1,7 +1,8 @@ import importlib.resources +from importlib.resources.abc import Traversable def get_prompt(name: str) -> str: - prompts_dir = importlib.resources.files("llm_cli.assets.prompts") - fpath = prompts_dir / f"{name}.md" + prompts_dir: Traversable = importlib.resources.files("llm_cli.assets.prompts") + fpath: Traversable = prompts_dir / f"{name}.md" return fpath.read_text()