From 5f451107c9d3473f2e4df5395c72e6448eaf05b4 Mon Sep 17 00:00:00 2001 From: wugren Date: Tue, 5 Dec 2023 18:08:32 +0800 Subject: [PATCH 1/2] =?UTF-8?q?Add=20pdf=E3=80=81docx=20parser=20and=20ima?= =?UTF-8?q?ge=E3=80=81video=20to=20text=20function=20for=20knowledge?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- rootfs/agents/TextSummary/agent.py | 2 +- rootfs/knowledge_pipelines/Mia/input.py | 162 ++++++++++++++++++++++-- src/aios/knowledge/object/object_id.py | 17 ++- 3 files changed, 164 insertions(+), 17 deletions(-) diff --git a/rootfs/agents/TextSummary/agent.py b/rootfs/agents/TextSummary/agent.py index 537f520..831b786 100644 --- a/rootfs/agents/TextSummary/agent.py +++ b/rootfs/agents/TextSummary/agent.py @@ -20,7 +20,7 @@ class TextSummaryAgent(CustomAIAgent): chunks = split_text(msg.body, separators=["\n\n", "\n"], chunk_size=4000, chunk_overlap=200, length_function=len) prompt = AgentPrompt() - prompt.system_message = "Your job is to generate a summary based on the input." + prompt.system_message = {"role":"system","content":"Your job is to generate a summary based on the input."} if len(chunks) == 1: prompt.append(AgentPrompt(chunks[0])) resp = await self.do_llm_complection(prompt) diff --git a/rootfs/knowledge_pipelines/Mia/input.py b/rootfs/knowledge_pipelines/Mia/input.py index e9a54f4..4166bef 100644 --- a/rootfs/knowledge_pipelines/Mia/input.py +++ b/rootfs/knowledge_pipelines/Mia/input.py @@ -1,17 +1,30 @@ +import copy import os +from typing import List + import aiofiles import chardet import logging import string -from knowledge import ImageObjectBuilder, DocumentObjectBuilder, KnowledgePipelineEnvironment, KnowledgePipelineJournal -from aios_kernel.storage import AIStorage + +import docx2txt +from PyPDF2 import PdfReader + +from aios import KnowledgePipelineEnvironment, ImageObjectBuilder, DocumentObjectBuilder, KnowledgeStore, RichTextObject +from aios.agent.agent_base import AgentPrompt +from aios.frame.compute_kernel import ComputeKernel +from aios.knowledge.data.writer import split_text +from aios.proto.compute_task import ComputeTaskResult, ComputeTaskResultCode +from aios.storage.storage import AIStorage +from aios.utils import video_utils, image_utils + class KnowledgeDirSource: def __init__(self, env: KnowledgePipelineEnvironment, config): self.env = env path = string.Template(config["path"]).substitute(myai_dir=AIStorage.get_instance().get_myai_dir()) config["path"] = path - self.config = config + self.config = config # @classmethod # def user_config_items(cls): @@ -19,16 +32,16 @@ class KnowledgeDirSource: def path(self): return self.config["path"] - + @staticmethod async def read_txt_file(file_path:str)->str: cur_encode = "utf-8" async with aiofiles.open(file_path,'rb') as f: cur_encode = chardet.detect(await f.read())['encoding'] - + async with aiofiles.open(file_path,'r',encoding=cur_encode) as f: return await f.read() - + async def next(self): while True: journals = self.env.journal.latest_journals(1) @@ -42,7 +55,7 @@ class KnowledgeDirSource: if os.path.getmtime(self.path()) <= from_time: yield (None, None) continue - + file_pathes = sorted(os.listdir(self.path()), key=lambda x: os.path.getctime(os.path.join(self.path(), x))) for rel_path in file_pathes: file_path = os.path.join(self.path(), rel_path) @@ -62,7 +75,138 @@ class KnowledgeDirSource: await self.env.get_knowledge_store().insert_object(document) yield (document.calculate_id(), file_path) yield (None, None) - + def init(env: KnowledgePipelineEnvironment, params: dict) -> KnowledgeDirSource: - return KnowledgeDirSource(env, params) \ No newline at end of file + return KnowledgeDirSource(env, params) + + +async def image_to_text(images: List[str]) -> str: + msg_prompt = AgentPrompt() + image_prompt = "What's in this image?" + content = [{"type": "text", "text": image_prompt}] + content.extend([{"type": "image_url", "image_url": {"url": image_utils.to_base64(image)}} for image in images]) + msg_prompt.messages = [{"role": "user", "content": content}] + + resp: ComputeTaskResult = await (ComputeKernel.get_instance() + .do_llm_completion(prompt=msg_prompt, + resp_mode="text", + mode_name="gpt-4-vision-preview", + max_token=4000, + inner_functions=None, + timeout=None)) + if resp.result_code != ComputeTaskResultCode.OK: + raise Exception(f"image_to_text error: {resp.result_code} msg:{resp.error_str}") + return resp.result_str + + +async def video_to_text(video: str) -> str: + prompt = "These pictures are key frames extracted from the video. Please describe the content of the video based on these key frames." + frames = video_utils.extract_frames(video, (1024, 1024)) + msg_prompt = AgentPrompt() + content = [{"type": "text", "text": prompt}] + content.extend([{"type": "image_url", "image_url": {"url": frame}} for frame in frames]) + msg_prompt.messages = [{"role": "user", "content": content}] + resp: ComputeTaskResult = await (ComputeKernel.get_instance() + .do_llm_completion(prompt=msg_prompt, + resp_mode="text", + mode_name="gpt-4-vision-preview", + max_token=4000, + inner_functions=None, + timeout=None)) + if resp.result_code != ComputeTaskResultCode.OK: + raise Exception(f"video_to_text error: {resp.result_code} msg:{resp.error_str}") + return resp.result_str + + +async def summary_document(text: str, separators: List[str]=["\n\n", "\n"]) -> str: + chunks = split_text(text, separators=separators, chunk_size=4000, chunk_overlap=200, length_function=len) + + prompt = AgentPrompt() + prompt.system_message = {"role":"system","content":"Your job is to generate a summary based on the input."} + if len(chunks) == 1: + prompt.append(AgentPrompt(chunks[0])) + resp = await (ComputeKernel.get_instance() + .do_llm_completion(prompt=prompt, + resp_mode="text", + mode_name="gpt-4-1106-preview", + max_token=4000, + inner_functions=None, + timeout=None)) + if resp.result_code != ComputeTaskResultCode.OK: + raise Exception(f"summary_document error: {resp.result_code} msg:{resp.error_str}") + return resp.result_str + + segments = [] + for i, chunk in enumerate(chunks): + seg_prompt = copy.deepcopy(prompt) + seg_prompt.append(AgentPrompt(chunk)) + resp = await (ComputeKernel.get_instance() + .do_llm_completion(prompt=seg_prompt, + resp_mode="text", + mode_name="gpt-4-1106-preview", + max_token=4000, + inner_functions=None, + timeout=None)) + if resp.result_code != ComputeTaskResultCode.OK: + raise Exception(f"summary_document error: {resp.result_code} msg:{resp.error_str}") + segments.append(resp.result_str) + + segments_str = "\n".join(segments) + prompt.append(AgentPrompt(f"Please combine the summaries of the following paragraphs into one complete summary:\n{segments_str}")) + resp = await (ComputeKernel.get_instance() + .do_llm_completion(prompt=prompt, + resp_mode="text", + mode_name="gpt-4-1106-preview", + max_token=4000, + inner_functions=None, + timeout=None)) + if resp.result_code != ComputeTaskResultCode.OK: + raise Exception(f"summary_document error: {resp.result_code} msg:{resp.error_str}") + return resp.result_str + + + +def pdf_to_rich_text_object(pdf: str, store: KnowledgeStore) -> RichTextObject: + base_name = os.path.basename(pdf) + cache_path = os.path.join(AIStorage.get_instance().get_myai_dir(), "knowledge", "doc_cache", base_name) + if not os.path.exists(cache_path): + os.makedirs(cache_path) + + reader = PdfReader(pdf) + rich_text = RichTextObject() + page_texts = [] + image_count = 0 + for page in reader.pages: + text = page.extract_text() + page_texts.append(text) + for image in page.images: + image_path = os.path.join(cache_path, f"{image_count}_{image.name}") + with open(image_path, "wb") as f: + f.write(image.data) + image_object = ImageObjectBuilder({}, {}, image_path).build(store) + rich_text.add_image(image_object) + + document = DocumentObjectBuilder({}, {}, "".join(page_texts)).build(store) + rich_text.add_document(document) + + return rich_text + + +def doc_to_rich_text_object(doc: str, store: KnowledgeStore) -> RichTextObject: + base_name = os.path.basename(doc) + cache_path = os.path.join(AIStorage.get_instance().get_myai_dir(), "knowledge", "doc_cache", base_name) + if not os.path.exists(cache_path): + os.makedirs(cache_path) + text = docx2txt.process(doc, cache_path) + + rich_text = RichTextObject() + for image in os.listdir(cache_path): + image_path = os.path.join(cache_path, image) + image_object = ImageObjectBuilder({}, {}, image_path).build(store) + rich_text.add_image(image_object) + + document = DocumentObjectBuilder({}, {}, text).build(store) + rich_text.add_document(document) + + return rich_text diff --git a/src/aios/knowledge/object/object_id.py b/src/aios/knowledge/object/object_id.py index 6cef161..1aca474 100644 --- a/src/aios/knowledge/object/object_id.py +++ b/src/aios/knowledge/object/object_id.py @@ -17,10 +17,10 @@ class ObjectType(IntEnum): def is_user_def(self) -> bool: return self.value >= 200 - + def get_user_def_type_code(self): return (self.value - 200) if self.is_user_def() else None - + @classmethod def from_user_def_type_code(cls, value): return value + 200 @@ -34,7 +34,7 @@ class ObjectID: # pylint: disable=too-few-public-methods def __str__(self): return self.to_base58() - + def to_base58(self): return base58.b58encode(self.value).decode() @@ -57,13 +57,16 @@ class ObjectID: # pylint: disable=too-few-public-methods def new_chunk_id(chunk_hash: HashValue): assert len(chunk_hash.value) == 32, "ObjectID must be 32 bytes long" return ObjectID(bytes([ObjectType.Chunk]) + chunk_hash.value[1:]) - + def get_object_type(self) -> ObjectType: return ObjectType(self.value[0]) - + @staticmethod def hash_data(data: bytes): return ObjectID.new_chunk_id(HashValue.hash_data(data)) - + def __eq__(self, other) -> bool: - return self.value == other.value \ No newline at end of file + return self.value == other.value + + def __hash__(self): + return hash(self.value) From 75d5c0066b9d277aec5b6dfd1206de91135ba40d Mon Sep 17 00:00:00 2001 From: wugren Date: Tue, 5 Dec 2023 19:59:41 +0800 Subject: [PATCH 2/2] AgentMsg support audio type --- src/aios/agent/agent.py | 20 +++++++++++++++++++- src/aios/proto/agent_msg.py | 7 +++++++ src/aios/proto/compute_task.py | 10 +++++----- src/component/tg_tunnel.py | 20 +++++++++++++++++++- 4 files changed, 50 insertions(+), 7 deletions(-) diff --git a/src/aios/agent/agent.py b/src/aios/agent/agent.py index 2b97326..dc4588f 100644 --- a/src/aios/agent/agent.py +++ b/src/aios/agent/agent.py @@ -13,7 +13,6 @@ import copy import sys from ..proto.agent_msg import AgentMsg -from ..proto.compute_task import ComputeTaskResult,ComputeTaskResultCode from .agent_base import * from .chatsession import * @@ -28,6 +27,7 @@ from ..storage.storage import AIStorage from ..knowledge import * from ..utils import video_utils, image_utils +from ..proto.compute_task import ComputeTaskResult,ComputeTaskResultCode logger = logging.getLogger(__name__) @@ -455,6 +455,15 @@ class AIAgent(BaseAIAgent): content = [{"type": "text", "text": f"{msg.sender}:{video_prompt}"}] content.extend([{"type": "image_url", "image_url": {"url": frame}} for frame in frames]) msg_prompt.messages = [{"role": "user", "content": content}] + elif msg.is_audio_msg(): + audio_file = msg.body + resp = await ComputeKernel.get_instance().do_speech_to_text(audio_file, None, prompt=None, response_format="text") + if resp.result_code != ComputeTaskResultCode.OK: + error_resp = msg.create_error_resp(resp.error_str) + return error_resp + else: + msg.body = resp.result_str + msg_prompt.messages = [{"role":"user","content":f"{msg.sender}:{resp.result_str}"}] else: msg_prompt.messages = [{"role":"user","content":f"{msg.sender}:{msg.body}"}] session_topic = msg.target + "#" + msg.topic @@ -487,6 +496,15 @@ class AIAgent(BaseAIAgent): content = [{"type": "text", "text": video_prompt}] content.extend([{"type": "image_url", "image_url": {"url": frame}} for frame in frames]) msg_prompt.messages = [{"role": "user", "content": content}] + elif msg.is_audio_msg(): + audio_file = msg.body + resp = await (ComputeKernel.get_instance().do_speech_to_text(audio_file, None, prompt=None, response_format="text")) + if resp.result_code != ComputeTaskResultCode.OK: + error_resp = msg.create_error_resp(resp.error_str) + return error_resp + else: + msg.body = resp.result_str + msg_prompt.messages = [{"role":"user","content":resp.result_str}] else: msg_prompt.messages = [{"role":"user","content":msg.body}] session_topic = msg.get_sender() + "#" + msg.topic diff --git a/src/aios/proto/agent_msg.py b/src/aios/proto/agent_msg.py index c7209bf..21ec3b6 100644 --- a/src/aios/proto/agent_msg.py +++ b/src/aios/proto/agent_msg.py @@ -210,6 +210,13 @@ class AgentMsg: return True return False + def is_audio_msg(self) -> bool: + if self.body_mime is None: + return False + if self.body_mime.startswith("audio/"): + return True + return False + def get_msg_id(self) -> str: return self.msg_id diff --git a/src/aios/proto/compute_task.py b/src/aios/proto/compute_task.py index c226b51..3390465 100644 --- a/src/aios/proto/compute_task.py +++ b/src/aios/proto/compute_task.py @@ -75,7 +75,7 @@ class ComputeTask: else: self.params["model_name"] = "text-embedding-ada-002" self.params["input"] = input - + def set_image_embedding_params(self, input = Union[ObjectID, bytes], model_name=None, callchain_id = None): self.task_type = ComputeTaskType.IMAGE_EMBEDDING self.create_time = time.time() @@ -86,7 +86,7 @@ class ComputeTask: else: self.params["model_name"] = None self.params["input"] = input - + def set_text_2_image_params(self, prompt: str, model_name, negative_prompt="", callchain_id=None): self.task_type = ComputeTaskType.TEXT_2_IMAGE self.create_time = time.time() @@ -126,15 +126,15 @@ class ComputeTaskResult: self.task_id: str = None self.callchain_id: str = None self.worker_id: str = None - self.error_str : str = None - self.result_code: int = 0 + self.error_str : str = None + self.result_code: int = ComputeTaskResultCode.OK self.result_str: str = None # easy to use,can read from result self.result : dict = {} self.result_refers: dict = {} self.pading_data: bytearray = None - + def set_from_task(self, task: ComputeTask): self.task_id = task.task_id diff --git a/src/component/tg_tunnel.py b/src/component/tg_tunnel.py index fa0e5e5..4406591 100644 --- a/src/component/tg_tunnel.py +++ b/src/component/tg_tunnel.py @@ -63,7 +63,7 @@ class TelegramTunnel(AgentTunnel): for update in updates: next_update_id = update.update_id + 1 - if update.message and (update.message.text or (update.message.photo and len(update.message.photo) > 0) or update.message.video): + if update.message and (update.message.text or (update.message.photo and len(update.message.photo) > 0) or update.message.video or update.message.voice or update.message.audio): await self.on_message(bot,update) return next_update_id @@ -89,6 +89,10 @@ class TelegramTunnel(AgentTunnel): update_id = (await self.bot.get_updates())[0].update_id except IndexError: update_id = None + except Exception as e: + logger.error(f"tg_tunnel error:{e}") + logger.exception(e) + update_id = None #logger.info("listening for new messages...") while True: @@ -179,6 +183,20 @@ class TelegramTunnel(AgentTunnel): await video_file.download_to_drive(file_path) agent_msg.body = agent_msg.create_video_body(file_path, message.caption) agent_msg.body_mime = f"video/{ext}" + elif message.audio is not None: + audio_file = await message.audio.get_file() + ext = audio_file.file_path.rsplit(".")[-1] + file_path = os.path.join(self.get_cache_path(), audio_file.file_id + f".{ext}") + await audio_file.download_to_drive(file_path) + agent_msg.body = file_path + agent_msg.body_mime = f"audio/{ext}" + elif message.voice is not None: + audio_file = await message.voice.get_file() + ext = audio_file.file_path.rsplit(".")[-1] + file_path = os.path.join(self.get_cache_path(), audio_file.file_id + f".{ext}") + await audio_file.download_to_drive(file_path) + agent_msg.body = file_path + agent_msg.body_mime = f"audio/{ext}" agent_msg.create_time = time.time() messag_type = message.chat.type