diff --git a/rootfs/agents/Mia/agent.toml b/rootfs/agents/Mia/agent.toml index 5a3bac9..15590c8 100644 --- a/rootfs/agents/Mia/agent.toml +++ b/rootfs/agents/Mia/agent.toml @@ -2,34 +2,22 @@ instance_id = "Mia" fullname = "Mia" llm_model_name = "gpt-3.5-turbo-16k-0613" max_token_size = 16000 -enable_kb = "true" -#enable_timestamp = "true" -#owner_prompt = "我是你的主人{name}" -contact_prompt = "我是你的主人朋友{name}，请回答主人允许回答的问题" -guest_prompt = "我是你的的主人的客人{name}，请不要回答我的任何问题" -owner_env = "calender" +#enable_function =["add_event"] +#enable_kb = "true" +enable_timestamp = "true" +owner_prompt = "我是你的主人{name}" +contact_prompt = "我是你的朋友{name}" +owner_env = "knowledge" [[prompt]] role = "system" content = """ -你叫Jarvis,是我的超级私人助理。 -你领导一个团队为我服务，团队的成员有： -Tracy,私人英语老师 -David,私人画家 +你叫Mia,你可以访问我的个人知识库。 *** -你看到的信息里有的有时会带上时间标签，这是为了让你更好的理解时间。你回复的信息不用创建这个时间标签。 你在收到我的信息后，按如下规则处理 -1. 如果你认为团队里有人更适合处理该信息，用下面方法转发消息给他们处理 -``` -##/send_msg 成员名字 -消息内容 -``` -2.你可以访问我的Calender,查看我的日程安排。如果你在处理信息的过程中需要修改我的日程安排，请直接用合适的方法修改。 -3.不符合上述规则的信息，请尽力处理。 +1. 在第一次接受到一条信息时，优先尝试用合适的关键字查询去查询知识库。 +2. 如果信息中包含一段知识库的查询结果，尝试用查询结果处理，如果还是不能处理，尝试递增index继续查询。 +3. 如果知识库返回不了结果了，请尽力返回。 """ - -#3.你可以访问我的个人信息库，当你处理我的信息时，如果需要用到我的个人信息，请先用合适的方法进行查询，然后再基于查询的结果进行进一步处理后再将结果发给我。 -#4.你能根据我的需要对系统进行配置，但在修改任何配置前，请先和我确认。 - diff --git a/src/aios_kernel/__init__.py b/src/aios_kernel/__init__.py index c4d6362..debd524 100644 --- a/src/aios_kernel/__init__.py +++ b/src/aios_kernel/__init__.py @@ -5,7 +5,7 @@ from .agent import AIAgent,AIAgentTemplete,AgentPrompt from .compute_kernel import ComputeKernel,ComputeTask,ComputeTaskResult,ComputeTaskState,ComputeTaskType from .compute_node import ComputeNode,LocalComputeNode from .open_ai_node import OpenAI_ComputeNode -from .knowledge_base import KnowledgeBase +from .knowledge_base import KnowledgeBase, KnowledgeEnvironment from .knowledge_pipeline import KnowledgeEmailSource, KnowledgeDirSource, KnowledgePipline from .role import AIRole,AIRoleGroup from .workflow import Workflow @@ -23,6 +23,7 @@ from .text_to_speech_function import TextToSpeechFunction from .workspace_env import WorkspaceEnvironment from .local_stability_node import Local_Stability_ComputeNode from .stability_node import Stability_ComputeNode +from .local_st_compute_node import LocalSentenceTransformer_Text_ComputeNode,LocalSentenceTransformer_Image_ComputeNode AIOS_Version = "0.5.1, build 2023-9-27" diff --git a/src/aios_kernel/compute_kernel.py b/src/aios_kernel/compute_kernel.py index 843a951..a243374 100644 --- a/src/aios_kernel/compute_kernel.py +++ b/src/aios_kernel/compute_kernel.py @@ -148,7 +148,7 @@ class ComputeKernel: task_result = await self._send_task(task_req) if task_req.state == ComputeTaskState.DONE: - return task_result.result + return task_result.result_str return "error!" diff --git a/src/aios_kernel/compute_task.py b/src/aios_kernel/compute_task.py index 12fd5cd..6ab9dd2 100644 --- a/src/aios_kernel/compute_task.py +++ b/src/aios_kernel/compute_task.py @@ -2,6 +2,8 @@ from enum import Enum import uuid import time +from typing import Union +from knowledge import ObjectID class ComputeTaskResultCode(Enum): OK = 0 @@ -25,6 +27,7 @@ class ComputeTaskType(Enum): VOICE_2_TEXT = "voice_2_text" TEXT_2_VOICE = "text_2_voice" TEXT_EMBEDDING ="text_embedding" + IMAGE_EMBEDDING ="image_embedding" class ComputeTask: @@ -60,7 +63,7 @@ class ComputeTask: if inner_functions is not None: self.params["inner_functions"] = inner_functions - def set_text_embedding_params(self, input, model_name=None, callchain_id = None): + def set_text_embedding_params(self, input: str, model_name=None, callchain_id = None): self.task_type = ComputeTaskType.TEXT_EMBEDDING self.create_time = time.time() self.task_id = uuid.uuid4().hex @@ -70,6 +73,17 @@ 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() + self.task_id = uuid.uuid4().hex + self.callchain_id = callchain_id + if model_name is not None: + self.params["model_name"] = model_name + 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 diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index d2debd8..8e4de5a 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -4,6 +4,8 @@ import logging from .agent import AgentPrompt from .compute_kernel import ComputeKernel from .storage import AIStorage +from .environment import Environment +from .ai_function import SimpleAIFunction from knowledge import * @@ -20,6 +22,7 @@ class KnowledgeBase: def __singleton_init__(self) -> None: self.store = KnowledgeStore() self.compute_kernel = ComputeKernel.get_instance() + self._default_text_model = "all-MiniLM-L6-v2" async def __embedding_document(self, document: DocumentObject): for chunk_id in document.get_chunk_list(): @@ -28,8 +31,8 @@ class KnowledgeBase: raise ValueError(f"text chunk not found: {chunk_id}") text = chunk.read().decode("utf-8") - vector = await self.compute_kernel.do_text_embedding(text) - await self.store.get_vector_store("default").insert(vector, chunk_id) + vector = await self.compute_kernel.do_text_embedding(text, self._default_text_model) + await self.store.get_vector_store(self._default_text_model).insert(vector, chunk_id) async def __embedding_image(self, image: ImageObject): desc = {} @@ -39,8 +42,8 @@ class KnowledgeBase: desc["exif"] = image.get_exif() if not not image.get_tags(): desc["tags"] = image.get_tags() - vector = await self.compute_kernel.do_text_embedding(json.dumps(desc)) - await self.store.get_vector_store("default").insert(vector, image.calculate_id()) + vector = await self.compute_kernel.do_text_embedding(json.dumps(desc), self._default_text_model) + await self.store.get_vector_store(self._default_text_model).insert(vector, image.calculate_id()) async def __embedding_video(self, vedio: VideoObject): desc = {} @@ -50,8 +53,8 @@ class KnowledgeBase: desc["info"] = vedio.get_info() if not not vedio.get_tags(): desc["tags"] = vedio.get_tags() - vector = await self.compute_kernel.do_text_embedding(json.dumps(desc)) - await self.store.get_vector_store("default").insert(vector, vedio.calculate_id()) + vector = await self.compute_kernel.do_text_embedding(json.dumps(desc), self._default_text_model) + await self.store.get_vector_store(self._default_text_model).insert(vector, vedio.calculate_id()) async def __embedding_rich_text(self, rich_text: RichTextObject): for document_id in rich_text.get_documents().values(): @@ -68,8 +71,8 @@ class KnowledgeBase: await self.__embedding_rich_text(rich_text) async def __embedding_email(self, email: EmailObject): - vector = await self.compute_kernel.do_text_embedding(json.dumps(email.get_desc())) - await self.store.get_vector_store("default").insert(vector, email.calculate_id()) + vector = await self.compute_kernel.do_text_embedding(json.dumps(email.get_desc()), self._default_text_model) + await self.store.get_vector_store(self._default_text_model).insert(vector, email.calculate_id()) await self.__embedding_rich_text(email.get_rich_text()) @@ -159,23 +162,10 @@ class KnowledgeBase: async def insert_object(self, object: KnowledgeObject): self.store.get_object_store().put_object(object.calculate_id(), object.encode()) await self.__do_embedding(object) - - async def query_prompt(self, prompt: AgentPrompt): - logging.info(f"query_prompt: {prompt}") - objects = await self.query_objects(prompt) - knowledge_prompt = self.prompt_from_objects(objects) - logging.info(f"prompt_from_objects result: {knowledge_prompt.as_str()}") - - return knowledge_prompt - - async def query_objects(self, prompt: AgentPrompt) -> [ObjectID]: - results = [] - for msg in prompt.messages: - if msg["role"] == "user": - vector = await self.compute_kernel.do_text_embedding(msg["content"]) - object_ids = await self.store.get_vector_store("default").query(vector, 10) - results.extend(object_ids) - return results + + async def query_objects(self, tokens: str, topk: int) -> [ObjectID]: + vector = await self.compute_kernel.do_text_embedding(tokens, self._default_text_model) + return await self.store.get_vector_store(self._default_text_model).query(vector, topk) def __load_object(self, object_id: ObjectID) -> KnowledgeObject: if object_id.get_object_type() == ObjectType.Document: @@ -192,7 +182,7 @@ class KnowledgeBase: pass - def prompt_from_objects(self, object_ids: [ObjectID]) -> AgentPrompt: + def tokens_from_objects(self, object_ids: [ObjectID]) -> list[str]: results = dict() for object_id in object_ids: parents = self.store.get_relation_store().get_related_root_objects(object_id) @@ -203,8 +193,7 @@ class KnowledgeBase: results[str(root_object_id)].append(object_id) else: results[str(root_object_id)] = [root_object_id, object_id] - - content = "*** I have provided the following known information for your reference with json format:\n" + content = "" result_desc = [] for result in results.values(): # first element in result is the root object @@ -236,12 +225,28 @@ class KnowledgeBase: else: pass content += json.dumps(result_desc) - content += ".\n" + content += ".\n" - prompt = AgentPrompt() - prompt.messages.append({"role": "user", "content": content}) - - return prompt + return content - \ No newline at end of file +class KnowledgeEnvironment(Environment): + def __init__(self, env_id: str) -> None: + super().__init__(env_id) + + query_param = { + "tokens": "tokens to query", + "index": "index of query result" + } + self.add_ai_function(SimpleAIFunction("query_knowledge", + "vector query content from local knowledge base", + self._query, + query_param)) + + async def _query(self, tokens: str, index: int=0): + object_ids = await KnowledgeBase().query_objects(tokens, 4) + if len(object_ids) <= index: + return "*** I have no more information for your reference.\n" + else: + content = "*** I have provided the following known information for your reference with json format:\n" + return content + KnowledgeBase().tokens_from_objects(object_ids[index:index + 1]) \ No newline at end of file diff --git a/src/aios_kernel/knowledge_pipeline.py b/src/aios_kernel/knowledge_pipeline.py index 24dc0c8..9b6b111 100644 --- a/src/aios_kernel/knowledge_pipeline.py +++ b/src/aios_kernel/knowledge_pipeline.py @@ -32,7 +32,7 @@ import requests import os import toml from .storage import AIStorage, UserConfigItem -from .knowledge_base import KnowledgeBase, ImageObjectBuilder, ObjectID, ObjectType, DocumentObjectBuilder +from .knowledge_base import KnowledgeBase, ImageObjectBuilder, ObjectID, ObjectType, DocumentObjectBuilder, EmailObjectBuilder, EmailObject class KnowledgeJournal: def __init__(self, source_type: str, source_id: str, item_id: str, object_id: str, timestamp=None): @@ -53,7 +53,10 @@ class KnowledgeJournal: pass return f"Add {object_type} from {os.path.join(self.source_id, self.item_id)}" if self.source_type == "email": - pass + object_id = ObjectID.from_base58(self.object_id) + email = EmailObject.decode(KnowledgeBase().store.get_object_store().get_object(object_id)) + meta = email.get_meta() + return f'Add email from {os.path.join(self.source_id)} subject {meta["subject"]}' # init sqlite3 client @@ -108,7 +111,7 @@ class KnowledgeEmailSource: self.config["type"] = "email" def id(self): - "::".join([self.config["imap_server"], self.config["address"]]) + return self.config["address"] @classmethod def user_config_items(cls): @@ -121,13 +124,15 @@ class KnowledgeEmailSource: @classmethod def local_root(cls): user_data_dir = AIStorage.get_instance().get_myai_dir() - return os.path.abspath(f"{user_data_dir}/email") + return os.path.abspath(f"{user_data_dir}/knowledge/email") async def run_once(self): # read config from toml file # and read from config config.local.toml if exists (config.local.toml is ignored by git) + logging.debug(f"knowledge email source {self.id()} run once") + filter = "ALL" self.client = self.email_client() - await self.read_emails() + await self.read_emails(imap_keyword=filter) def email_client(self) -> imaplib.IMAP4_SSL: logging.info(f"read email config from {self.config.get('imap_server')}") @@ -139,26 +144,37 @@ class KnowledgeEmailSource: return client async def read_emails(self, folder: str = 'INBOX', imap_keyword: str = "UNSEEN"): + journal_client = KnowledgeJournalClient() + latest_journal = journal_client.latest_journal(self.id()) + latest_uid = 0 if latest_journal is None else int(latest_journal.item_id) self.client.select(folder) _, data = self.client.uid('search', None, imap_keyword) # get email uid list email_list = data[0].split() logging.info(f"got {len(email_list)} emails") - email_list.reverse() + journal_client = KnowledgeJournalClient() for uid in email_list: - if self.check_email_saved(uid): - logging.info(f"email uid {uid} already saved") - else: - self.read_and_save_email(uid) - logging.info(f"email uid {uid} saved") + _uid = int.from_bytes(uid) + if _uid > latest_uid: + email_dir = self.check_email_saved(uid) + if email_dir is not None: + logging.info(f"email uid {uid} already saved") + else: + email_dir = self.read_and_save_email(uid) + logging.info(f"email uid {uid} saved") + email_object = EmailObjectBuilder({}, email_dir).build() + await KnowledgeBase().insert_object(email_object) + journal_client.insert(KnowledgeJournal("email", self.id(), str(int.from_bytes(uid)), str(email_object.calculate_id()))) - def read_and_save_email(self, uid: str): + + def read_and_save_email(self, uid: str) -> str: message_parts = "(BODY.PEEK[])" _, email_data = self.client.uid('fetch', uid, message_parts) mail = mailparser.parse_from_bytes(email_data[0][1]) logging.info(f"got email subject [{mail.subject}]") self.save_email(mail) + return self.get_local_dir_name(mail) def get_local_dir_name(self, mail: mailparser.MailParser) -> str: dir = f"{self.local_root()}/{self.config.get('address')}" @@ -166,7 +182,7 @@ class KnowledgeEmailSource: name = hashlib.md5(name.encode('utf-8')).hexdigest() return f"{dir}/{name}" - def check_email_saved(self, uid: str): + def check_email_saved(self, uid: str) -> str: message_parts = "(BODY[HEADER])" _, email_data = self.client.uid('fetch', uid, message_parts) mail = mailparser.parse_from_bytes(email_data[0][1]) @@ -175,8 +191,8 @@ class KnowledgeEmailSource: logging.info(f"check email saved {dir}") file = f"{dir}/email.txt" if os.path.exists(file): - return False - return False + return dir + return None # save email attachment(images) def save_email_attachment(self, mail: mailparser.MailParser, email_dir: str): @@ -205,12 +221,16 @@ class KnowledgeEmailSource: img_urls = [img['src'] for img in img_tags if 'src' in img.attrs] logging.info(f'Found {len(img_urls)} images in email body') + name_count = 0 + if not os.path.exists(email_dir): os.makedirs(email_dir) for img_url in img_urls: # keep the original image filename(last of url) - img_filename = os.path.join(email_dir, img_url.split('/')[-1]) + ext = img_url.split('/')[-1].split('.')[-1] + img_filename = os.path.join(email_dir, f"{name_count}.{ext}") + name_count += 1 # download image response = requests.get(img_url, stream=True) if response.status_code == 200: @@ -230,7 +250,8 @@ class KnowledgeEmailSource: logging.info(f"save email to {email_dir}") if not os.path.exists(email_dir): os.makedirs(email_dir) - with open(f"{email_dir}/email.txt", "w") as f: + with open(f"{email_dir}/email.txt", "w", encoding='utf-8') as f: + # soup = BeautifulSoup(mail.body, 'html.parser') f.write(mail.body) with open(f"{email_dir}/meta.json", "w", encoding='utf-8') as f: mail_dict = json.loads(mail.mail_json) diff --git a/src/aios_kernel/local_st_compute_node.py b/src/aios_kernel/local_st_compute_node.py new file mode 100644 index 0000000..cfef1e0 --- /dev/null +++ b/src/aios_kernel/local_st_compute_node.py @@ -0,0 +1,246 @@ +import logging +import requests +from typing import Optional, List +from pydantic import BaseModel +from typing import Union +from PIL import Image +import io +from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskType +from .queue_compute_node import Queue_ComputeNode +from knowledge import ObjectID + +logger = logging.getLogger(__name__) + + + +class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode): + # For valid pretrained models, see https://www.sbert.net/docs/pretrained_models.html + def __init__(self, model_name: str = "all-MiniLM-L6-v2"): + super().__init__() + + self.node_id = "local_sentence_transformer_text_embedding_node" + self.model_name = model_name + self.model = None + + def initial(self) -> bool: + logger.info( + f"LocalSentenceTransformer_Text_ComputeNode init, model_name: {self.model_name}" + ) + + assert self.model_name is not None + assert self.model is None + try: + from sentence_transformers import SentenceTransformer + + self.model = SentenceTransformer(self.model_name) + except Exception as err: + logger.error(f"load model {self.model} failed: {err}") + return False + self.start() + return True + + async def execute_task( + self, task: ComputeTask + ) -> { + "task_type": str, + "content": str, + "message": str, + "state": ComputeTaskState, + "error": { + "code": int, + "message": str, + }, + }: + try: + # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task: {task}") + if task.task_type == ComputeTaskType.TEXT_EMBEDDING: + input = task.params["input"] + logger.debug( + f"LocalSentenceTransformer_Text_ComputeNode task input: {input}" + ) + sentence_embeddings = self.model.encode(input, show_progress_bar=False).tolist() + # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}") + return { + "state": ComputeTaskState.DONE, + "content": sentence_embeddings, + "message": None, + } + else: + return { + "state": ComputeTaskState.ERROR, + "error": {"code": -1, "message": "unsupport embedding task type"}, + } + except Exception as err: + import traceback + + logger.error(f"{traceback.format_exc()}, error: {err}") + + return { + "state": ComputeTaskState.ERROR, + "error": {"code": -1, "message": "unknown exception: " + str(err)}, + } + + def display(self) -> str: + return f"LocalSentenceTransformer_Text_ComputeNode: {self.node_id}, {self.model_name}" + + def get_capacity(self): + pass + + def is_support(self, task: ComputeTask) -> bool: + return task.task_type == ComputeTaskType.TEXT_EMBEDDING and task.params["model_name"] == "all-MiniLM-L6-v2" + + def is_local(self) -> bool: + return True + + +class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode): + # For valid pretrained models, see https://www.sbert.net/docs/pretrained_models.html + def __init__( + self, + model_name: str = "clip-ViT-B-32", + multi_model_name: str = "clip-ViT-B-32-multilingual-v1", + ): + super().__init__() + + self.node_id = "local_sentence_transformer_image_embedding_node" + self.model_name = model_name + self.multi_model_name = multi_model_name + self.model = None + self.multi_model = None + + def initial(self) -> bool: + logger.info( + f"LocalSentenceTransformer_Image_ComputeNode init, model_name: {self.model_name} {self.multi_model_name}" + ) + + assert self.model_name is not None + assert self.multi_model_name is not None + assert self.model is None + assert self.multi_model is None + + try: + from sentence_transformers import SentenceTransformer + + self.model = SentenceTransformer(self.model_name) + self.multi_model = SentenceTransformer(self.multi_model_name) + except Exception as err: + logger.error(f"load model {self.model} failed: {err}") + return False + self.start() + return True + + def _load_image(self, source: Union[ObjectID, bytes]) -> Optional[Image]: + image_data = None + if isinstance(source, ObjectID): + from knowledge import KnowledgeStore, ImageObject + + buf = KnowledgeStore().get_object_store().get_object(source) + if buf is None: + logger.error(f"load image object but not found! {source}") + return None + + try: + image_obj = ImageObject.decode(buf) + except Exception as err: + logger.error(f"decode ImageObject from buffer failed: {source}, {err}") + return None + + file_size = image_obj.get_file_size() + print(f"got image object: {source.to_base58()}, size: {file_size}") + + image_data = ( + KnowledgeStore() + .get_chunk_reader() + .read_chunk_list_to_single_bytes(image_obj.get_chunk_list()) + ) + + elif isinstance(source, bytes): + image_data = source + else: + logger.error(f"unsupport image source type: {type(source)}, {source}") + return None + + try: + img = Image.open(io.BytesIO(image_data)) + + return img + except Exception as err: + logger.error(f"load image from buffer failed: {source}, {err}") + return None + + async def execute_task( + self, task: ComputeTask + ) -> { + "task_type": str, + "content": str, + "message": str, + "state": ComputeTaskState, + "error": { + "code": int, + "message": str, + }, + }: + try: + # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task: {task}") + if task.task_type == ComputeTaskType.TEXT_EMBEDDING: + input = task.params["input"] + logger.debug( + f"LocalSentenceTransformer_Text_ComputeNode task text input: {input}" + ) + sentence_embeddings = self.multi_model.encode(input, show_progress_bar=False).tolist() + # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}") + return { + "state": ComputeTaskState.DONE, + "content": sentence_embeddings, + "message": None, + } + elif task.task_type == ComputeTaskType.IMAGE_EMBEDDING: + input = task.params["input"] + logger.debug( + f"LocalSentenceTransformer_Image_ComputeNode task image input: {input}" + ) + + img = self._load_image(input) + if img is None: + return { + "state": ComputeTaskState.ERROR, + "error": {"code": -1, "message": "load image failed"}, + } + + sentence_embeddings = self.model.encode(img) + + # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}") + return { + "state": ComputeTaskState.DONE, + "content": sentence_embeddings, + "message": None, + } + else: + return { + "state": ComputeTaskState.ERROR, + "error": {"code": -1, "message": "unsupport embedding task type"}, + } + except Exception as err: + import traceback + + logger.error(f"{traceback.format_exc()}, error: {err}") + + return { + "state": ComputeTaskState.ERROR, + "error": {"code": -1, "message": "unknown exception: " + str(err)}, + } + + def display(self) -> str: + return f"LocalSentenceTransformer_Image_ComputeNode: {self.node_id}, {self.model_name}" + + def get_capacity(self): + pass + + def is_support(self, task: ComputeTask) -> bool: + return ( + (task.task_type == ComputeTaskType.TEXT_EMBEDDING and task.params["model_name"] == "clip-ViT-B-32") + or task.task_type == ComputeTaskType.IMAGE_EMBEDDING + ) + + def is_local(self) -> bool: + return True diff --git a/src/aios_kernel/open_ai_node.py b/src/aios_kernel/open_ai_node.py index e16d76d..6716c2b 100644 --- a/src/aios_kernel/open_ai_node.py +++ b/src/aios_kernel/open_ai_node.py @@ -98,7 +98,7 @@ class OpenAI_ComputeNode(ComputeNode): task.state = ComputeTaskState.DONE result.result_code = ComputeTaskResultCode.OK result.worker_id = self.node_id - result.result = resp["data"][0]["embedding"] + result.result_str = resp["data"][0]["embedding"] return result case ComputeTaskType.LLM_COMPLETION: diff --git a/src/aios_kernel/queue_compute_node.py b/src/aios_kernel/queue_compute_node.py index 1475672..6d97446 100644 --- a/src/aios_kernel/queue_compute_node.py +++ b/src/aios_kernel/queue_compute_node.py @@ -4,7 +4,7 @@ from asyncio import Queue import logging from abc import abstractmethod -from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType +from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskResultCode, ComputeTaskState, ComputeTaskType from .compute_node import ComputeNode logger = logging.getLogger(__name__) @@ -13,6 +13,7 @@ class Queue_ComputeNode(ComputeNode): def __init__(self): super().__init__() self.task_queue = Queue() + self.is_start = False @abstractmethod async def execute_task(self, task: ComputeTask) -> { @@ -39,28 +40,32 @@ class Queue_ComputeNode(ComputeNode): resp = await self.execute_task(task) result = ComputeTaskResult() - result.set_from_task(task) - - task.state = resp["state"] - - if task.state == ComputeTaskState.ERROR: - task.error_str = resp["error"]["message"] - result.worker_id = self.node_id - result.result_str = resp["content"] - result.result_message = resp["message"] + task.state = resp["state"] + + if task.state == ComputeTaskState.ERROR: + result.result_code = ComputeTaskResultCode.ERROR + task.error_str = resp["error"]["message"] + else: + result.result_code = ComputeTaskResultCode.OK + result.result_str = resp["content"] + result.result_message = resp["message"] + + result.set_from_task(task) return result - async def start(self): + def start(self): + if self.is_start is True: + return + self.is_start = True + async def _run_task_loop(): while True: task = await self.task_queue.get() - logger.info(f"{self.display()} get task: {task.display()}") - result = await self._run_task(task) - if result is not None: - task.result = result + logger.info(f"openai_node get task: {task.display()}") + await self._run_task(task) asyncio.create_task(_run_task_loop()) diff --git a/src/knowledge/core_object/document_object.py b/src/knowledge/core_object/document_object.py index 672b4a3..b69f59f 100644 --- a/src/knowledge/core_object/document_object.py +++ b/src/knowledge/core_object/document_object.py @@ -50,11 +50,7 @@ class DocumentObjectBuilder: return self def build(self) -> DocumentObject: - chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_text( - self.text, - 1024 * 4, - ".?!\n" - ) + chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_text(self.text) doc = DocumentObject(self.meta, self.tags, chunk_list) doc_id = doc.calculate_id() diff --git a/src/knowledge/core_object/email_object.py b/src/knowledge/core_object/email_object.py index 37531ff..108c9cc 100644 --- a/src/knowledge/core_object/email_object.py +++ b/src/knowledge/core_object/email_object.py @@ -21,13 +21,13 @@ class EmailObject(KnowledgeObject): super().__init__(ObjectType.Email, desc, body) - def get_meta(self): + def get_meta(self) -> dict: return self.desc["meta"] - def get_tags(self): + def get_tags(self) -> dict: return self.desc["tags"] - def get_rich_text(self): + def get_rich_text(self) -> RichTextObject: return self.body["content"] diff --git a/src/knowledge/core_object/image_object.py b/src/knowledge/core_object/image_object.py index 2c02ad6..340bcc2 100644 --- a/src/knowledge/core_object/image_object.py +++ b/src/knowledge/core_object/image_object.py @@ -2,6 +2,7 @@ from ..object import KnowledgeObject from ..data import ChunkList, ChunkListWriter from ..object import ObjectType from .. import KnowledgeStore +import os # desc # meta @@ -13,30 +14,34 @@ from .. import KnowledgeStore class ImageObject(KnowledgeObject): - def __init__(self, meta: dict, tags: dict, exif: dict, chunk_list: ChunkList): + def __init__(self, meta: dict, tags: dict, exif: dict, file_size: int, chunk_list: ChunkList): desc = dict() body = dict() desc["meta"] = meta desc["exif"] = exif desc["tags"] = tags desc["hash"] = chunk_list.hash.to_base58() + desc["file_size"] = file_size body["chunk_list"] = chunk_list.chunk_list super().__init__(ObjectType.Image, desc, body) - def get_meta(self): + def get_meta(self) -> dict: return self.desc["meta"] - def get_exif(self): + def get_exif(self) -> dict: return self.desc["exif"] - def get_tags(self): + def get_tags(self) -> dict: return self.desc["tags"] - def get_hash(self): + def get_hash(self) -> str: return self.desc["hash"] - def get_chunk_list(self): + def get_file_size(self) -> int: + return self.desc["file_size"] + + def get_chunk_list(self) -> ChunkList: return self.body["chunk_list"] @@ -82,8 +87,10 @@ class ImageObjectBuilder: return self def build(self) -> ImageObject: + + file_size = os.path.getsize(self.image_file) chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_file( self.image_file, 1024 * 1024 * 4, self.restore_file ) exif = get_exif_data(self.image_file) - return ImageObject(self.meta, self.tags, exif, chunk_list) + return ImageObject(self.meta, self.tags, exif, file_size, chunk_list) diff --git a/src/knowledge/data/reader.py b/src/knowledge/data/reader.py index c4cbb96..ede2648 100644 --- a/src/knowledge/data/reader.py +++ b/src/knowledge/data/reader.py @@ -12,7 +12,7 @@ class Chunk: self.range_start = range_start self.size = size - def read(self): + def read(self) -> bytes: with open(self.file_path, 'rb') as f: f.seek(self.range_start) return f.read(self.size) @@ -26,6 +26,8 @@ class ChunkReader: def get_chunk(self, chunk_id: ChunkID) -> Chunk: positions = self.chunk_tracker.get_position(chunk_id) + logging.info(f"chunk positions: {chunk_id}, {positions}") + if positions is None: logging.warning(f"chunk not found: {chunk_id}") return None @@ -54,15 +56,23 @@ class ChunkReader: def get_chunk_list(self, chunk_list: List[ChunkID]) -> List[Chunk]: return [self.get_chunk(chunk_id) for chunk_id in chunk_list] - def read_chunk_list(self, chunk_ids: List[ChunkID]): + def read_chunk_list(self, chunk_ids: List[ChunkID]) -> bytes: for chunk_id in chunk_ids: chunk = self.get_chunk(chunk_id) if chunk is None: raise ValueError(f"chunk not found: {chunk_id}") - yield from chunk.read() + yield chunk.read() - def read_text_chunk_list(self, chunk_ids: List[ChunkID]): + def read_chunk_list_to_single_bytes(self, chunk_ids: List[ChunkID]) -> bytes: + chunks = [] + for chunk in self.read_chunk_list(chunk_ids): + chunks.append(chunk) + + image_data = b''.join(chunks) + return image_data + + def read_text_chunk_list(self, chunk_ids: List[ChunkID]) -> str: for chunk_id in chunk_ids: chunk = self.get_chunk(chunk_id) if chunk is None: diff --git a/src/knowledge/data/writer.py b/src/knowledge/data/writer.py index ad6c520..0381bc6 100644 --- a/src/knowledge/data/writer.py +++ b/src/knowledge/data/writer.py @@ -1,13 +1,139 @@ import os import hashlib import re -from typing import Tuple, List +import tiktoken +import logging +from typing import Callable, Iterable, Optional, Tuple, List from .chunk_store import ChunkStore from .chunk import ChunkID, PositionFileRange, PositionType from ..object import HashValue from .tracker import ChunkTracker from .chunk_list import ChunkList +def _join_docs(docs: List[str], separator: str) -> Optional[str]: + text = separator.join(docs) + text = text.strip() + if text == "": + return None + else: + return text + +def _merge_splits( + splits: Iterable[str], + separator: str, + chunk_size: int, + chunk_overlap: int, + length_function: Callable[[str], int] + ) -> List[str]: + # We now want to combine these smaller pieces into medium size + # chunks to send to the LLM. + separator_len = length_function(separator) + + docs = [] + current_doc: List[str] = [] + total = 0 + for d in splits: + _len = length_function(d) + if ( + total + _len + (separator_len if len(current_doc) > 0 else 0) + > chunk_size + ): + if total > chunk_size: + logging.warning( + f"Created a chunk of size {total}, " + f"which is longer than the specified {self._chunk_size}" + ) + if len(current_doc) > 0: + doc = _join_docs(current_doc, separator) + if doc is not None: + docs.append(doc) + # Keep on popping if: + # - we have a larger chunk than in the chunk overlap + # - or if we still have any chunks and the length is long + while total > chunk_overlap or ( + total + _len + (separator_len if len(current_doc) > 0 else 0) + > chunk_size + and total > 0 + ): + total -= length_function(current_doc[0]) + ( + separator_len if len(current_doc) > 1 else 0 + ) + current_doc = current_doc[1:] + current_doc.append(d) + total += _len + (separator_len if len(current_doc) > 1 else 0) + doc = _join_docs(current_doc, separator) + if doc is not None: + docs.append(doc) + return docs + + +def _split_text_with_regex( + text: str, separator: str, keep_separator: bool +) -> List[str]: + # Now that we have the separator, split the text + if separator: + if keep_separator: + # The parentheses in the pattern keep the delimiters in the result. + _splits = re.split(f"({separator})", text) + splits = [_splits[i] + _splits[i + 1] for i in range(1, len(_splits), 2)] + if len(_splits) % 2 == 0: + splits += _splits[-1:] + splits = [_splits[0]] + splits + else: + splits = re.split(separator, text) + else: + splits = list(text) + return [s for s in splits if s != ""] + + +def _split_text( + text: str, + separators: List[str], + chunk_size: int, + chunk_overlap: int, + length_function: Callable[[str], int] + ) -> List[str]: + + """Split incoming text and return chunks.""" + final_chunks = [] + # Get appropriate separator to use + separator = separators[-1] + new_separators = [] + for i, _s in enumerate(separators): + _separator = re.escape(_s) + if _s == "": + separator = _s + break + if re.search(_separator, text): + separator = _s + new_separators = separators[i + 1 :] + break + + keep_separator = True + _separator = re.escape(separator) + splits = _split_text_with_regex(text, _separator, keep_separator) + + # Now go merging things, recursively splitting longer texts. + _good_splits = [] + _separator = "" if keep_separator else separator + for s in splits: + if length_function(s) < chunk_size: + _good_splits.append(s) + else: + if _good_splits: + merged_text = _merge_splits(_good_splits, _separator, chunk_size, chunk_overlap, length_function) + final_chunks.extend(merged_text) + _good_splits = [] + if not new_separators: + final_chunks.append(s) + else: + other_info = _split_text(s, new_separators, chunk_size, chunk_overlap, length_function) + final_chunks.extend(other_info) + if _good_splits: + merged_text = _merge_splits(_good_splits, _separator, chunk_size, chunk_overlap, length_function) + final_chunks.extend(merged_text) + return final_chunks + class ChunkListWriter: def __init__(self, chunk_store: ChunkStore, chunk_tracker: ChunkTracker): self.chunk_store = chunk_store @@ -27,6 +153,8 @@ class ChunkListWriter: chunk = file.read(chunk_size) if not chunk: break + + chunk_len = len(chunk) chunk_id = ChunkID.hash_data(chunk) chunk_list.append(chunk_id) @@ -38,8 +166,9 @@ class ChunkListWriter: ) self.chunk_store.put_chunk(chunk_id, chunk) else: + pos = file.tell() file_range = PositionFileRange( - file_path, file.tell() - chunk_size, chunk_size + file_path, pos - chunk_len, pos ) self.chunk_tracker.add_position( chunk_id, str(file_range), PositionType.FileRange @@ -51,9 +180,24 @@ class ChunkListWriter: return ChunkList(chunk_list, file_hash) def create_chunk_list_from_text( - self, text: str, chunk_max_words: int, separator_chars: str = ".," + self, + text: str, + chunk_size: int = 4000, + chunk_overlap: int = 200, + separators: str = ["\n\n", "\n", " ", ""] ) -> ChunkList: - text_list = self._split_text_list(text, chunk_max_words, separator_chars) + enc = tiktoken.encoding_for_model("gpt-3.5-turbo") + + def length_function(text: str) -> int: + return len( + enc.encode( + text, + allowed_special=set(), + disallowed_special="all", + ) + ) + + text_list = _split_text(text, separators, chunk_size, chunk_overlap, length_function) chunk_list = [] hash_obj = hashlib.sha256() @@ -67,27 +211,4 @@ class ChunkListWriter: self.chunk_store.put_chunk(chunk_id, chunk_bytes) hash = HashValue(hash_obj.digest()) - return ChunkList(chunk_list, hash) - - @staticmethod - def _split_text_list( - text: str, chunk_max_words: int, separator_chars: str = ".," - ) -> List[str]: - sentences = re.split(f"[{separator_chars}]", text) - # chunk_list = [] - # chunk = [] - # word_count = 0 - # for sentence in sentences: - # words = sentence.split() - # for word in words: - # if word_count < chunk_max_words: - # chunk.append(word) - # word_count += 1 - # else: - # chunk_list.append(" ".join(chunk)) - # chunk = [word] - # word_count = 1 - # if chunk: - # chunk_list.append(" ".join(chunk)) - # return chunk_list - return sentences \ No newline at end of file + return ChunkList(chunk_list, hash) \ No newline at end of file diff --git a/src/knowledge/object.py b/src/knowledge/object.py deleted file mode 100644 index f7a127e..0000000 --- a/src/knowledge/object.py +++ /dev/null @@ -1,65 +0,0 @@ - -# define a object type enum -from abc import ABC, abstractmethod -from enum import Enum - -class ObjectType(Enum): - TextChunk = 1 - Image = 2 - Email = 101 - - -# define a object ID class to identify a object -class ObjectID: # pylint: disable=too-few-public-methods - def __init__(self, object_type, digist): - self.object_type = object_type - self.digist = digist - - def __str__(self): - return f"{self.object_type.name}:{self.digist}" - - -# define a object class -class KnowledgeObject(ABC): # pylint: disable=too-few-public-methods - def __init__(self, object_type: ObjectType): - self.object_type = object_type - - @abstractmethod - def get_id(self) -> ObjectID: - pass - - # define a to binary method to convert object to binary - @abstractmethod - def to_binary(self) -> bytes: - pass - - # define a from binary method to convert binary to object - @abstractmethod - def from_binary(self, binary: bytes): - pass - - -# define a text chunk class -class TextChunkObject(KnowledgeObject): # pylint: disable=too-few-public-methods - def __init__(self, text: str): - super().__init__(ObjectType.TextChunk) - self.text = text - - -# define a image class -class ImageObject(KnowledgeObject): # pylint: disable=too-few-public-methods - def __init__(self, meta, path): - super().__init__(ObjectType.Image) - self.meta = meta - self.path = path - - -# define a email class -class EmailObject(KnowledgeObject): # pylint: disable=too-few-public-methods - def __init__(self, meta): - super().__init__(ObjectType.Email) - self.meta = meta - self.text = [ObjectID] - self.images = [ObjectID] - - diff --git a/src/knowledge/object/blob.py b/src/knowledge/object/blob.py index b005e62..08fb0bf 100644 --- a/src/knowledge/object/blob.py +++ b/src/knowledge/object/blob.py @@ -1,6 +1,8 @@ import os import shutil from .object import ObjectID +import logging +logger = logging.getLogger(__name__) class FileBlobStorage: @@ -38,16 +40,23 @@ class FileBlobStorage: def put(self, object_id: ObjectID, contents: bytes): full_path = self.get_full_path(object_id) + if os.path.exists(full_path): + logger.warning(f"will replace object: {object_id}") + self.write_sync(full_path, contents) def get(self, object_id: ObjectID) -> bytes: full_path = self.get_full_path(object_id) + if not os.path.exists(full_path): + return None + with open(full_path, "rb") as f: return f.read() def delete(self, object_id: ObjectID): full_path = self.get_full_path(object_id) - os.remove(full_path) + if os.path.exists(full_path): + os.remove(full_path) def exists(self, object_id: ObjectID) -> bool: full_path = self.get_full_path(object_id) diff --git a/src/knowledge/object/object.py b/src/knowledge/object/object.py index a63e30a..ac6f2af 100644 --- a/src/knowledge/object/object.py +++ b/src/knowledge/object/object.py @@ -1,6 +1,8 @@ # define a object type enum +from __future__ import annotations from abc import ABC, abstractmethod from enum import Enum + from .object_id import ObjectID, ObjectType import hashlib import json @@ -61,5 +63,5 @@ class KnowledgeObject(ABC): return pickle.dumps(self) @staticmethod - def decode(data: bytes): + def decode(data: bytes) -> "ImageObject": return pickle.loads(data) diff --git a/src/knowledge/object_storage.py b/src/knowledge/object_storage.py deleted file mode 100644 index 1a6f85e..0000000 --- a/src/knowledge/object_storage.py +++ /dev/null @@ -1,33 +0,0 @@ -# import RDB LargeBinary -from sqlalchemy import Column, String, LargeBinary, create_engine, sessionmaker, pickle -from .object import KnowledgeObject - -# implement object storage with RDB -# define object storage table -class ObjectStorageTable(Base): - __tablename__ = 'object_storage' - id = Column(String, primary_key=True) - parent = Column(String, nullable=True) - object = Column(LargeBinary, nullable=False) - - def __init__(self, id, parent, object): # pylint: disable=redefined-builtin - self.id = id - self.parent = parent - self.object = object - -# define object storage class -class ObjectStorage: - async def __init__(self, db_url): - self.engine = create_engine(db_url) - self.session = sessionmaker(bind=self.engine)() # pylint: disable=not-callable - - async def get(self, id) -> [KnowledgeObject, KnowledgeObject]: - obj = self.session.query(ObjectStorageTable).filter(ObjectStorageTable.id == id).first() - if obj is None: - return None - return pickle.loads(obj.object) - - # define insert method - async def insert(self, object, parent): # pylint: disable=redefined-builtin - obj = ObjectStorageTable(id, parent, pickle.dumps(object)) - \ No newline at end of file diff --git a/src/knowledge/store.py b/src/knowledge/store.py index adcc85a..d8b4bbb 100644 --- a/src/knowledge/store.py +++ b/src/knowledge/store.py @@ -4,7 +4,7 @@ from .object import ObjectStore, ObjectRelationStore from .data import ChunkStore, ChunkTracker, ChunkListWriter, ChunkReader from .vector import ChromaVectorStore, VectorBase import logging -import aios_kernel + # KnowledgeStore class, which aggregates ChunkStore, ChunkTracker, and ObjectStore, and is a global singleton that makes it easy to use these three built-in store examples class KnowledgeStore: @@ -13,6 +13,8 @@ class KnowledgeStore: def __new__(cls): if cls._instance is None: cls._instance = super().__new__(cls) + + import aios_kernel knowledge_dir = aios_kernel.storage.AIStorage().get_myai_dir() / "knowledge" if not os.path.exists(knowledge_dir): @@ -60,5 +62,5 @@ class KnowledgeStore: def get_vector_store(self, model_name: str) -> VectorBase: if model_name not in self.vector_store: - self.vector_store[model_name] = ChromaVectorStore(model_name) + self.vector_store[model_name] = ChromaVectorStore(self.root, model_name) return self.vector_store[model_name] diff --git a/src/requirements.txt b/src/requirements.txt index 7bd02d6..c6128ee 100644 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -80,10 +80,61 @@ toml>=0.10.0 protobuf grpcio grpcio-status +h11==0.14.0 +httpcore==0.17.3 +httptools==0.6.0 +httpx==0.24.1 +huggingface-hub==0.16.4 +humanfriendly==10.0 +idna==3.4 +imageio==2.31.3 +imageio-ffmpeg==0.4.8 +importlib-resources==6.0.1 +mail-parser==3.15.0 +monotonic==1.6 +moviepy==1.0.0 +mpmath==1.3.0 +multidict==6.0.4 +numpy==1.25.2 +onnxruntime==1.15.1 +openai==0.28.0 +overrides==7.4.0 +packaging==23.1 +pandas==2.1.0 +Pillow==10.0.0 +posthog==3.0.2 +proglog==0.1.10 +prompt-toolkit==3.0.39 +proto-plus==1.22.3 +protobuf +pulsar-client==3.3.0 +pyasn1==0.5.0 +pyasn1-modules==0.3.0 +pydantic==1.10.12 +PyPika==0.48.9 +pyreadline3==3.4.1 +python-dateutil==2.8.2 +python-dotenv==1.0.0 +python-telegram-bot==20.5 +pytz==2023.3.post1 +PyYAML==6.0.1 +requests==2.31.0 +rsa==4.9 +simplejson==3.19.1 +six==1.16.0 +sniffio==1.3.0 +soupsieve==2.5 +starlette==0.27.0 +sympy==1.12 +telegram==0.0.1 +tokenizers==0.14.0 +toml==0.10.0 pysocks chardet pydub aiosqlite python-telegram-bot pydub -stability_sdk \ No newline at end of file +stability_sdk +sentence-transformers==2.2.2 +tiktoken diff --git a/src/service/aios_shell/aios_shell.py b/src/service/aios_shell/aios_shell.py index 094793d..a3aef85 100644 --- a/src/service/aios_shell/aios_shell.py +++ b/src/service/aios_shell/aios_shell.py @@ -24,8 +24,7 @@ directory = os.path.dirname(__file__) sys.path.append(directory + '/../../') -from aios_kernel import AIOS_Version,AgentMsgType,UserConfigItem,AIStorage,Workflow,AIAgent,AgentMsg,AgentMsgStatus,ComputeKernel,OpenAI_ComputeNode,AIBus,AIChatSession,AgentTunnel,TelegramTunnel,CalenderEnvironment,Environment,EmailTunnel,LocalLlama_ComputeNode,Local_Stability_ComputeNode,Stability_ComputeNode,PaintEnvironment -from aios_kernel import ContactManager,Contact + import proxy from aios_kernel import * @@ -114,6 +113,10 @@ class AIOS_Shell: cm.add_family_member(self.username,owenr) + knowledge_env = KnowledgeEnvironment("knowledge") + Environment.set_env_by_id("knowledge",knowledge_env) + + cal_env = CalenderEnvironment("calender") await cal_env.start() Environment.set_env_by_id("calender",cal_env) @@ -137,6 +140,7 @@ class AIOS_Shell: return False ComputeKernel.get_instance().add_compute_node(open_ai_node) + if await AIStorage.get_instance().is_feature_enable("llama"): llama_ai_node = LocalLlama_ComputeNode() if await llama_ai_node.initial() is True: @@ -146,6 +150,7 @@ class AIOS_Shell: logger.error("llama node initial failed!") await AIStorage.get_instance().set_feature_init_result("llama",False) + if await AIStorage.get_instance().is_feature_enable("aigc"): try: google_text_to_speech_node = GoogleTextToSpeechNode.get_instance() @@ -160,13 +165,20 @@ class AIOS_Shell: # logger.error("stability api node initial failed!") # ComputeKernel.get_instance().add_compute_node(stability_api_node) - local_sd_node = Local_Stability_ComputeNode.get_instance() - if await local_sd_node.initial() is True: - ComputeKernel.get_instance().add_compute_node(local_sd_node) - else: - logger.error("local stability node initial failed!") - await AIStorage.get_instance.set_feature_init_result("aigc",False) - + + + local_st_text_compute_node = LocalSentenceTransformer_Text_ComputeNode() + if local_st_text_compute_node.initial() is not True: + logger.error("local sentence transformer text embedding node initial failed!") + else: + ComputeKernel.get_instance().add_compute_node(local_st_text_compute_node) + + local_st_image_compute_node = LocalSentenceTransformer_Image_ComputeNode() + if local_st_image_compute_node.initial() is not True: + logger.error("local sentence transformer image embedding node initial failed!") + else: + ComputeKernel.get_instance().add_compute_node(local_st_image_compute_node) + await ComputeKernel.get_instance().start() @@ -308,8 +320,7 @@ class AIOS_Shell: async def handle_knowledge_commands(self, args): show_text = FormattedText([("class:title", "sub command not support!\n" "/knowledge add email | dir\n" - "/knowledge journal [$topn]\n" - "/knowledge query $query\n")]) + "/knowledge journal [$topn]\n")]) if len(args) < 1: return show_text sub_cmd = args[0] @@ -346,13 +357,6 @@ class AIOS_Shell: topn = 10 if len(args) == 1 else int(args[1]) journals = [str(journal) for journal in KnowledgePipline.get_instance().get_latest_journals(topn)] print_formatted_text("\r\n".join(journals)) - if sub_cmd == "query": - if len(args) < 2: - return show_text - prompt = AgentPrompt() - prompt.messages.append({"role": "user", "content":" ".join(args[1:])}) - result = await KnowledgeBase().query_prompt(prompt) - print_formatted_text(result.as_str()) async def call_func(self,func_name, args): match func_name: @@ -662,8 +666,7 @@ async def main(): '/connect $target', '/contact $name', '/knowledge add email | dir', - '/knowledge journal [$topn]', - '/knowledge query $query' + '/knowledge journal [$topn]', '/set_config $key', '/enable $feature', '/disable $feature', diff --git a/test/test_chunk.py b/test/test_chunk.py index f1f496b..0771f03 100644 --- a/test/test_chunk.py +++ b/test/test_chunk.py @@ -59,7 +59,7 @@ class TestChunk(unittest.TestCase): with open(text_file, "r", encoding="utf-8") as file: text = file.read() - gen.create_chunk_list_from_text(text, 1024) + gen.create_chunk_list_from_text(text) if __name__ == "__main__": diff --git a/test/test_embedding.py b/test/test_embedding.py new file mode 100644 index 0000000..e6fe4d7 --- /dev/null +++ b/test/test_embedding.py @@ -0,0 +1,98 @@ +import sys +import os +import logging +from sentence_transformers import SentenceTransformer, util + + +dir_path = os.path.dirname(os.path.realpath(__file__)) +print(dir_path) + +sys.path.append("{}/../src/".format(dir_path)) +print(sys.path) + +root = logging.getLogger() +root.setLevel(logging.DEBUG) +handler = logging.StreamHandler(sys.stdout) +handler.setLevel(logging.DEBUG) +formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") +handler.setFormatter(formatter) +root.addHandler(handler) + + +def test_st(): + image_model = SentenceTransformer('clip-ViT-B-32-multilingual-v1') + model = SentenceTransformer("all-MiniLM-L6-v2") + + # Our sentences we like to encode + sentences = [ + "This framework generates embeddings for each input sentence", + "Sentences are passed as a list of string.", + "The quick brown fox jumps over the lazy dog.", + ] + + # Sentences are encoded by calling model.encode() + sentence_embeddings = model.encode(sentences) + + # Print the embeddings + for sentence, embedding in zip(sentences, sentence_embeddings): + print("Sentence:", sentence) + print("Embedding:", embedding) + print("") + + # Single list of sentences + + """ + sentences = [ + "The cat sits outside", + "A man is playing guitar", + "I love pasta", + "The new movie is awesome", + "The cat plays in the garden", + "A woman watches TV", + "The new movie is so great", + "Do you like pizza?", + ] + """ + sentences = [ + "猫坐在外面", + "狗坐在上面", + "狗坐在里面", + "一个男人在弹吉他", + "我爱意大利面", + "新电影太精彩了", + "猫在花园里玩耍", + "一个女人在看电视", + "新电影太棒了", + "你喜欢披萨吗？", + ] + + # Compute embeddings + #embeddings = model.encode(sentences, convert_to_tensor=True) + embeddings = model.encode(sentences) + print("embeddings as follows: ") + print(embeddings) + + + # Compute cosine-similarities for each sentence with each other sentence + cosine_scores = util.cos_sim(embeddings, embeddings) + + # Find the pairs with the highest cosine similarity scores + pairs = [] + for i in range(len(cosine_scores) - 1): + for j in range(i + 1, len(cosine_scores)): + pairs.append({"index": [i, j], "score": cosine_scores[i][j]}) + + # Sort scores in decreasing order + pairs = sorted(pairs, key=lambda x: x["score"], reverse=True) + + for pair in pairs[0:10]: + i, j = pair["index"] + print( + "{} \t\t {} \t\t Score: {:.4f}".format( + sentences[i], sentences[j], pair["score"] + ) + ) + + +if __name__ == "__main__": + test_st() diff --git a/test/test_object.py b/test/test_object.py index 655e2a2..e464aa0 100644 --- a/test/test_object.py +++ b/test/test_object.py @@ -24,13 +24,15 @@ from knowledge import ( ObjectRelationStore, KnowledgeStore, EmailObject, + ImageObject, ) +from aios_kernel import LocalSentenceTransformer_Image_ComputeNode, ComputeTask import asyncio import unittest -class TestVectorSTorage(unittest.TestCase): - def test_object(self): +class TestVectorSTorage(unittest.IsolatedAsyncioTestCase): + async def test_object(self): data = HashValue.hash_data("1233".encode("utf-8")) print(data.to_base58()) print(data.to_base36()) @@ -57,6 +59,41 @@ class TestVectorSTorage(unittest.TestCase): ret2 = obj.encode() self.assertEqual(ret, ret2) + images = email_object.get_rich_text().get_images() + image_keys = list(images.keys()) + print("got image list: ", image_keys) + + image_id = images[image_keys[1]] + print(f"got image object: {image_keys[1]} {image_id.to_base58()}") + + node = LocalSentenceTransformer_Image_ComputeNode(); + ret = node.initial() + self.assertEqual(ret, True) + + task = ComputeTask() + task.set_image_embedding_params(image_id) + ret = await node.execute_task(task) + print(ret) + ''' + buf = KnowledgeStore().get_object_store().get_object(image_id) + image_obj= ImageObject.decode(buf) + file_size = image_obj.get_file_size() + print(f"got image object: {image_id.to_base58()}, size: {file_size}") + + + image_data = KnowledgeStore().get_chunk_reader().read_chunk_list_to_single_bytes(image_obj.get_chunk_list()) + self.assertEqual(file_size, len(image_data)) + + from PIL import Image + import io + image = Image.open(io.BytesIO(image_data)) + image.show() + + from sentence_transformers import SentenceTransformer + #model = SentenceTransformer('clip-ViT-B-32-multilingual-v1') + model = SentenceTransformer('clip-ViT-B-32') + model.encode(image, convert_to_tensor=True) + ''' def test_relation(self): obj1 = ObjectID.hash_data("12345".encode("utf-8"))