gyt/opendan
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #109 from wugren/MVP

Browse Source 
Add pdf、docx parser and image、video to text function for knowledge
This commit is contained in:
Liu Zhicong
2023-12-05 10:56:07 -08:00
committed by GitHub
parent 66e407add5 75d5c0066b
commit 0425c80ffd
7 changed files with 214 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rootfs/agents/TextSummary/agent.py
+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)
rootfs/knowledge_pipelines/Mia/input.py
+146 -2
View File
@@ -1,10 +1,23 @@
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):
@@ -66,3 +79,134 @@ class KnowledgeDirSource:
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
src/aios/agent/agent.py
+19 -1
View File
@@ -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
src/aios/knowledge/object/object_id.py
+3
View File
@@ -67,3 +67,6 @@ class ObjectID: # pylint: disable=too-few-public-methods
def __eq__(self, other) -> bool:
return self.value == other.value
def __hash__(self):
return hash(self.value)
src/aios/proto/agent_msg.py
+7
View File
@@ -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
src/aios/proto/compute_task.py
+1 -1
View File
@@ -127,7 +127,7 @@ class ComputeTaskResult:
self.callchain_id: str = None
self.worker_id: str = None
self.error_str : str = None
self.result_code: int = 0
self.result_code: int = ComputeTaskResultCode.OK
self.result_str: str = None # easy to use,can read from result
self.result : dict = {}
src/component/tg_tunnel.py
+19 -1
View File
@@ -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