import copy import os from typing import List import aiofiles import chardet import logging import string 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 # @classmethod # def user_config_items(cls): # return [("path", "local dir path")] 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) from_time = 0 if len(journals) == 1: latest_journal = journals[0] if latest_journal.is_finish(): yield None continue from_time = os.path.getctime(latest_journal.get_input()) 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) timestamp = os.path.getctime(file_path) if timestamp <= from_time: continue ext = os.path.splitext(file_path)[1].lower() if ext in ['.jpg', '.jpeg', '.png', '.gif', '.bmp']: logging.info(f"knowledge dir source found image file {file_path}") image = ImageObjectBuilder({}, {}, file_path).build(self.env.get_knowledge_store()) await self.env.get_knowledge_store().insert_object(image) yield (image.calculate_id(), file_path) if ext in ['.txt']: logging.info(f"knowledge dir source found text file {file_path}") text = await self.read_txt_file(file_path) document = DocumentObjectBuilder({}, {}, text).build(self.env.get_knowledge_store()) 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) 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