change /knowledge commands in shell
This commit is contained in:
@@ -0,0 +1,68 @@
|
||||
import os
|
||||
import aiofiles
|
||||
import chardet
|
||||
import logging
|
||||
import string
|
||||
from knowledge import ImageObjectBuilder, DocumentObjectBuilder, KnowledgePipelineEnvironment, KnowledgePipelineJournal
|
||||
from aios_kernel.storage import AIStorage
|
||||
|
||||
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)
|
||||
@@ -0,0 +1,102 @@
|
||||
# define a knowledge base class
|
||||
import json
|
||||
import string
|
||||
from aios_kernel import ComputeKernel, AIStorage
|
||||
from knowledge import *
|
||||
|
||||
|
||||
class EmbeddingParser:
|
||||
def __init__(self, env: KnowledgePipelineEnvironment, config: dict):
|
||||
self._default_text_model = "all-MiniLM-L6-v2"
|
||||
self._default_image_model = "clip-ViT-B-32"
|
||||
|
||||
path = string.Template(config["path"]).substitute(myai_dir=AIStorage.get_instance().get_myai_dir())
|
||||
if not os.path.exists(path):
|
||||
os.makedirs(path)
|
||||
config["path"] = path
|
||||
|
||||
self.env = env
|
||||
self.config = config
|
||||
|
||||
def get_path(self) -> str:
|
||||
return self.config["path"]
|
||||
|
||||
def __get_vector_store(self, model_name: str) -> ChromaVectorStore:
|
||||
return ChromaVectorStore(self.get_path(), model_name)
|
||||
|
||||
async def __embedding_document(self, document: DocumentObject):
|
||||
for chunk_id in document.get_chunk_list():
|
||||
chunk = self.env.get_knowledge_store().get_chunk_reader().get_chunk(chunk_id)
|
||||
if chunk is None:
|
||||
raise ValueError(f"text chunk not found: {chunk_id}")
|
||||
|
||||
text = chunk.read().decode("utf-8")
|
||||
vector = await ComputeKernel.get_instance().do_text_embedding(text, self._default_text_model)
|
||||
if vector:
|
||||
await self.__get_vector_store(self._default_text_model).insert(vector, chunk_id)
|
||||
|
||||
async def __embedding_image(self, image: ImageObject):
|
||||
# desc = {}
|
||||
# if not not image.get_meta():
|
||||
# desc["meta"] = image.get_meta()
|
||||
# if not not image.get_exif():
|
||||
# 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), self._default_text_model)
|
||||
vector = await ComputeKernel.get_instance().do_image_embedding(image.calculate_id(), self._default_image_model)
|
||||
if vector:
|
||||
await self.__get_vector_store(self._default_image_model).insert(vector, image.calculate_id())
|
||||
|
||||
async def __embedding_video(self, vedio: VideoObject):
|
||||
desc = {}
|
||||
if not not vedio.get_meta():
|
||||
desc["meta"] = vedio.get_meta()
|
||||
if not not vedio.get_info():
|
||||
desc["info"] = vedio.get_info()
|
||||
if not not vedio.get_tags():
|
||||
desc["tags"] = vedio.get_tags()
|
||||
vector = await ComputeKernel.get_instance().do_text_embedding(json.dumps(desc), self._default_text_model)
|
||||
await self.__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():
|
||||
document = DocumentObject.decode(self.env.get_knowledge_store().get_object_store().get_object(document_id))
|
||||
await self.__embedding_document(document)
|
||||
for image_id in rich_text.get_images().values():
|
||||
image = ImageObject.decode(self.env.get_knowledge_store().get_object_store().get_object(image_id))
|
||||
await self.__embedding_image(image)
|
||||
for video_id in rich_text.get_videos().values():
|
||||
video = VideoObject.decode(self.env.get_knowledge_store().get_object_store().get_object(video_id))
|
||||
await self.__embedding_video(video)
|
||||
for rich_text_id in rich_text.get_rich_texts().values():
|
||||
rich_text = RichTextObject.decode(self.env.get_knowledge_store().get_object_store().get_object(rich_text_id))
|
||||
await self.__embedding_rich_text(rich_text)
|
||||
|
||||
async def __embedding_email(self, email: EmailObject):
|
||||
vector = await ComputeKernel.get_instance().do_text_embedding(json.dumps(email.get_desc()), self._default_text_model)
|
||||
await self.__get_vector_store(self._default_text_model).insert(vector, email.calculate_id())
|
||||
await self.__embedding_rich_text(email.get_rich_text())
|
||||
|
||||
|
||||
async def __do_embedding(self, object: KnowledgeObject):
|
||||
if object.get_object_type() == ObjectType.Document:
|
||||
await self.__embedding_document(object)
|
||||
if object.get_object_type() == ObjectType.Image:
|
||||
await self.__embedding_image(object)
|
||||
if object.get_object_type() == ObjectType.Video:
|
||||
await self.__embedding_video(object)
|
||||
if object.get_object_type() == ObjectType.RichText:
|
||||
await self.__embedding_rich_text(object)
|
||||
if object.get_object_type() == ObjectType.Email:
|
||||
await self.__embedding_email(object)
|
||||
else:
|
||||
pass
|
||||
|
||||
async def parse(self, object: ObjectID) -> str:
|
||||
obj = self.env.get_knowledge_store().load_object(object)
|
||||
await self.__do_embedding(obj)
|
||||
return "insert into vector store"
|
||||
|
||||
def init(env: KnowledgePipelineEnvironment, params: dict) -> EmbeddingParser:
|
||||
return EmbeddingParser(env, params)
|
||||
+3
-3
@@ -1,6 +1,6 @@
|
||||
name = "LocalEmbedding"
|
||||
input.module = "local_dir"
|
||||
name = "Mia"
|
||||
input.module = "input.py"
|
||||
input.params.path = "${myai_dir}/data"
|
||||
parser.module = "embedding"
|
||||
parser.module = "parser.py"
|
||||
parser.params.path = "${myai_dir}/knowledge/indices/embedding"
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
import os
|
||||
import logging
|
||||
import json
|
||||
from aios_kernel import *
|
||||
from knowledge import *
|
||||
|
||||
class KnowledgeEnvironment(Environment):
|
||||
def __init__(self, env_id: str) -> None:
|
||||
super().__init__(env_id)
|
||||
self.path = os.path.join(AIStorage.get_instance().get_myai_dir(), "knowledge/indices/embedding")
|
||||
self._default_text_model = "all-MiniLM-L6-v2"
|
||||
self._default_image_model = "clip-ViT-B-32"
|
||||
|
||||
query_param = {
|
||||
"tokens": "key words to query",
|
||||
"types": "prefered knowledge types, one or more of [text, image]",
|
||||
"index": "index of query result"
|
||||
}
|
||||
self.add_ai_function(SimpleAIFunction("query_knowledge",
|
||||
"vector query content from local knowledge base",
|
||||
self._query,
|
||||
query_param))
|
||||
|
||||
def __get_vector_store(self, model_name: str) -> ChromaVectorStore:
|
||||
return ChromaVectorStore(self.path, model_name)
|
||||
|
||||
async def query_objects(self, tokens: str, types: list[str], topk: int) -> [ObjectID]:
|
||||
texts = []
|
||||
if "text" in types:
|
||||
vector = await ComputeKernel.get_instance().do_text_embedding(tokens, self._default_text_model)
|
||||
texts = await self.__get_vector_store(self._default_text_model).query(vector, topk)
|
||||
images = []
|
||||
if "image" in types:
|
||||
vector = await ComputeKernel.get_instance().do_text_embedding(tokens, self._default_image_model)
|
||||
images = await self.__get_vector_store(self._default_image_model).query(vector, topk)
|
||||
return texts + images
|
||||
|
||||
|
||||
def tokens_from_objects(self, object_ids: [ObjectID]) -> list[str]:
|
||||
results = dict()
|
||||
for object_id in object_ids:
|
||||
parents = KnowledgeStore().get_relation_store().get_related_root_objects(object_id)
|
||||
# last parent is the root object
|
||||
root_object_id = parents[0] if parents else object_id
|
||||
logging.info(f"object_id: {str(object_id)} root_object_id: {str(root_object_id)}")
|
||||
if str(root_object_id) in results:
|
||||
results[str(root_object_id)].append(object_id)
|
||||
else:
|
||||
results[str(root_object_id)] = [root_object_id, object_id]
|
||||
content = ""
|
||||
result_desc = []
|
||||
for result in results.values():
|
||||
# first element in result is the root object
|
||||
root_object_id = result[0]
|
||||
if root_object_id.get_object_type() == ObjectType.Email:
|
||||
email = KnowledgeStore().load_object(root_object_id)
|
||||
desc = email.get_desc()
|
||||
desc["type"] = "email"
|
||||
desc["contents"] = []
|
||||
result_desc.append(desc)
|
||||
upper_list = desc["contents"]
|
||||
result = result[1:]
|
||||
else:
|
||||
upper_list = result_desc
|
||||
|
||||
for object_id in result:
|
||||
if object_id.get_object_type() == ObjectType.Chunk:
|
||||
upper_list.append({"type": "text", "content": KnowledgeStore().get_chunk_reader().get_chunk(object_id).read().decode("utf-8")})
|
||||
if object_id.get_object_type() == ObjectType.Image:
|
||||
# image = self.load_object(object_id)
|
||||
desc = dict()
|
||||
desc["id"] = str(object_id)
|
||||
desc["type"] = "image"
|
||||
upper_list.append(desc)
|
||||
if object_id.get_object_type() == ObjectType.Video:
|
||||
video = KnowledgeStore().load_object(object_id)
|
||||
desc = video.get_desc()
|
||||
desc["type"] = "video"
|
||||
upper_list.append(desc)
|
||||
else:
|
||||
pass
|
||||
content += json.dumps(result_desc)
|
||||
content += ".\n"
|
||||
|
||||
return content
|
||||
|
||||
async def _query(self, tokens: str, types: list[str] = ["text"], index: str=0):
|
||||
index = int(index)
|
||||
object_ids = await self.query_objects(tokens, types, 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 + self.tokens_from_objects(object_ids[index:index+1])
|
||||
|
||||
def init() -> KnowledgeEnvironment:
|
||||
return KnowledgeEnvironment("embedding")
|
||||
@@ -1,3 +1,3 @@
|
||||
pipelines = [
|
||||
"local_embedding"
|
||||
"Mia"
|
||||
]
|
||||
Reference in New Issue
Block a user