Merge branch 'MVP' into mvp-dev

This commit is contained in:
Si Changjun
2023-12-06 09:39:52 +08:00
committed by GitHub
7 changed files with 213 additions and 23 deletions
+1 -1
View File
@@ -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)
+152 -8
View File
@@ -1,16 +1,29 @@
import copy
import os
from typing import List
import aiofiles
import chardet
import logging
import string
from aios import AIStorage,ImageObjectBuilder, DocumentObjectBuilder, KnowledgePipelineEnvironment, KnowledgePipelineJournal
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):
@@ -18,16 +31,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)
@@ -41,7 +54,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)
@@ -61,7 +74,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)
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