Merge pull request #69 from photosssa/MVP

Add image insert/query in knowledge base
This commit is contained in:
Liu Zhicong
2023-09-28 09:17:30 -07:00
committed by GitHub
8 changed files with 103 additions and 45 deletions
+4 -3
View File
@@ -1,10 +1,10 @@
instance_id = "Mia"
fullname = "Mia"
llm_model_name = "gpt-3.5-turbo-16k-0613"
llm_model_name = "gpt-4"
max_token_size = 16000
#enable_function =["add_event"]
#enable_kb = "true"
enable_timestamp = "true"
enable_timestamp = "false"
owner_prompt = "我是你的主人{name}"
contact_prompt = "我是你的朋友{name}"
owner_env = "knowledge"
@@ -18,6 +18,7 @@ content = """
你在收到我的信息后,按如下规则处理
1. 在第一次接受到一条信息时,优先尝试用合适的关键字查询去查询知识库。
2. 如果信息中包含一段知识库的查询结果,尝试用查询结果处理,如果还是不能处理,尝试递增index继续查询。
3. 如果知识库返回不了结果了,请尽力返回
3. 如果要返回知识库结果条目,在消息开头附上他的json字符串
4. 如果知识库返回不了结果了,请尽力返回。
"""
-12
View File
@@ -358,13 +358,6 @@ class AIAgent:
old_content = msg.get("content")
msg["content"] = old_content.format_map(self.owner_env)
async def _get_knowlege_prompt(self,input_msg:AgentPrompt) -> AgentPrompt:
if self.enable_kb is False:
return None
from .knowledge_base import KnowledgeBase
return await KnowledgeBase().query_prompt(input_msg)
async def _process_msg(self,msg:AgentMsg) -> AgentMsg:
from .compute_kernel import ComputeKernel
from .bus import AIBus
@@ -393,11 +386,6 @@ class AIAgent:
history_prmpt,history_token_len = await self._get_prompt_from_session(chatsession,system_prompt_len + function_token_len,input_len)
prompt.append(history_prmpt) # chat context
kb_prompt = await self._get_knowlege_prompt(msg_prompt)
prompt.append(kb_prompt)
prompt.append(msg_prompt)
logger.debug(f"Agent {self.agent_id} do llm token static system:{system_prompt_len},function:{function_token_len},history:{history_token_len},input:{input_len}, totoal prompt:{system_prompt_len + function_token_len + history_token_len} ")
task_result:ComputeTaskResult = await ComputeKernel.get_instance().do_llm_completion(prompt,self.llm_model_name,self.max_token_size,inner_functions)
if task_result.result_code != ComputeTaskResultCode.OK:
+16
View File
@@ -5,6 +5,7 @@ import logging
import asyncio
from asyncio import Queue
from knowledge import ObjectID
from .agent import AgentPrompt
from .compute_node import ComputeNode
from .compute_task import ComputeTask, ComputeTaskState, ComputeTaskResult, ComputeTaskType,ComputeTaskResultCode
@@ -152,6 +153,21 @@ class ComputeKernel:
return "error!"
def image_embedding(self,input:ObjectID,model_name:Optional[str] = None):
task_req = ComputeTask()
task_req.set_image_embedding_params(input,model_name)
self.run(task_req)
return task_req
async def do_image_embedding(self,input:ObjectID,model_name:Optional[str] = None) -> [float]:
task_req = self.image_embedding(input,model_name)
task_result = await self._send_task(task_req)
if task_req.state == ComputeTaskState.DONE:
return task_result.result_str
return "error!"
async def do_text_to_speech(self,
input:str,
language_code:Optional[str] = None,
+54 -18
View File
@@ -23,6 +23,7 @@ class KnowledgeBase:
self.store = KnowledgeStore()
self.compute_kernel = ComputeKernel.get_instance()
self._default_text_model = "all-MiniLM-L6-v2"
self._default_image_model = "clip-ViT-B-32"
async def __embedding_document(self, document: DocumentObject):
for chunk_id in document.get_chunk_list():
@@ -35,15 +36,16 @@ class KnowledgeBase:
await self.store.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)
await self.store.get_vector_store(self._default_text_model).insert(vector, image.calculate_id())
# 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 self.compute_kernel.do_image_embedding(image.calculate_id(), self._default_image_model)
await self.store.get_vector_store(self._default_image_model).insert(vector, image.calculate_id())
async def __embedding_video(self, vedio: VideoObject):
desc = {}
@@ -163,9 +165,16 @@ class KnowledgeBase:
self.store.get_object_store().put_object(object.calculate_id(), object.encode())
await self.__do_embedding(object)
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)
async def query_objects(self, tokens: str, types: list[str], topk: int) -> [ObjectID]:
texts = []
if "text" in types:
vector = await self.compute_kernel.do_text_embedding(tokens, self._default_text_model)
texts = await self.store.get_vector_store(self._default_text_model).query(vector, topk)
images = []
if "image" in types:
vector = await self.compute_kernel.do_text_embedding(tokens, self._default_image_model)
images = await self.store.get_vector_store(self._default_image_model).query(vector, topk)
return texts + images
def __load_object(self, object_id: ObjectID) -> KnowledgeObject:
if object_id.get_object_type() == ObjectType.Document:
@@ -213,8 +222,9 @@ class KnowledgeBase:
if object_id.get_object_type() == ObjectType.Chunk:
upper_list.append({"type": "text", "content": self.store.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 = image.get_desc()
# 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:
@@ -229,13 +239,39 @@ class KnowledgeBase:
return content
def parse_object_in_message(self, message: str) -> KnowledgeObject:
# get message's first line
lines = message.split("\n")
if len(lines) > 0:
message = lines[0]
try:
desc = json.loads(message)
object_id = desc["object_id"]
except:
return None
if object_id is not None:
return self.__load_object(ObjectID(object_id))
def bytes_from_object(self, object: KnowledgeObject) -> bytes:
if object.get_object_type() == ObjectType.Image:
image_object = object
return self.store.get_chunk_reader().read_chunk_list_to_single_bytes(image_object.get_chunk_list())
class KnowledgeEnvironment(Environment):
def __init__(self, env_id: str) -> None:
super().__init__(env_id)
query_param = {
"tokens": "tokens to query",
"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",
@@ -243,10 +279,10 @@ class KnowledgeEnvironment(Environment):
self._query,
query_param))
async def _query(self, tokens: str, index: int=0):
object_ids = await KnowledgeBase().query_objects(tokens, 4)
async def _query(self, tokens: str, types: list[str] = ["text"], index: int=0):
object_ids = await KnowledgeBase().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 + KnowledgeBase().tokens_from_objects(object_ids[index:index + 1])
return content + KnowledgeBase().tokens_from_objects(object_ids[index:index+1])
+6
View File
@@ -267,6 +267,7 @@ class KnowledgeEmailSource:
class KnowledgeDirSource:
def __init__(self, config):
self.config = config
config["path"] = os.path.abspath(config["path"])
self.config["type"] = "dir"
@classmethod
@@ -342,6 +343,11 @@ class KnowledgePipline:
self.add_email_source(KnowledgeEmailSource(source_config))
if source_config['type'] == 'dir':
self.add_dir_source(KnowledgeDirSource(source_config))
user_data_dir = AIStorage.get_instance().get_myai_dir()
default_dir = os.path.abspath(f"{user_data_dir}/data")
if not os.path.exists(default_dir):
os.makedirs(default_dir)
self.add_dir_source(KnowledgeDirSource({"path": default_dir}))
return True
+2 -2
View File
@@ -146,7 +146,7 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
return None
file_size = image_obj.get_file_size()
print(f"got image object: {source.to_base58()}, size: {file_size}")
# print(f"got image object: {source.to_base58()}, size: {file_size}")
image_data = (
KnowledgeStore()
@@ -207,7 +207,7 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode):
"error": {"code": -1, "message": "load image failed"},
}
sentence_embeddings = self.model.encode(img)
sentence_embeddings = self.model.encode(img, show_progress_bar=False).tolist()
# logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}")
return {
+18 -7
View File
@@ -9,6 +9,10 @@ from telegram import Bot
from telegram.ext import Updater
from telegram.error import Forbidden, NetworkError
from knowledge.object.object_id import ObjectType
from .knowledge_base import KnowledgeBase
from .tunnel import AgentTunnel
from .storage import AIStorage
from .contact_manager import ContactManager,Contact,FamilyMember
@@ -171,13 +175,20 @@ class TelegramTunnel(AgentTunnel):
else:
if resp_msg.body_mime is None:
if resp_msg.body is not None:
pos = resp_msg.body.find("audio file")
if pos != -1:
audio_file = resp_msg.body[pos+11:].strip()
if audio_file.startswith("\""):
audio_file = audio_file[1:-1]
await update.message.reply_voice(audio_file)
return
knowledge_object = KnowledgeBase().parse_object_in_message(resp_msg.body)
if knowledge_object is not None:
if knowledge_object.get_object_type() == ObjectType.Image:
image = KnowledgeBase().bytes_from_object(knowledge_object)
await update.message.reply_photo(image)
return
else:
pos = resp_msg.body.find("audio file")
if pos != -1:
audio_file = resp_msg.body[pos+11:].strip()
if audio_file.startswith("\""):
audio_file = audio_file[1:-1]
await update.message.reply_voice(audio_file)
return
await update.message.reply_text(resp_msg.body)
else:
if resp_msg.body_mime.startswith("image"):
+3 -3
View File
@@ -319,7 +319,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 add dir\n"
"/knowledge journal [$topn]\n")])
if len(args) < 1:
return show_text
@@ -585,7 +585,7 @@ def print_welcome_screen():
\033[1;94m\tGive your Agent a Telegram account :\033[0m /connect $agent_name
\033[1;94m\tAdd personal files to the AI Knowledge Base. \033[0m
\t\t1) Copy your file to ~/myai/data
\t\t2) /knowlege add dir
\t\t2) /knowledge add dir
\033[1;94m\tSearch your knowledge base :\033[0m /open Mia
\033[1;94m\tCheck the progress of AI reading personal data :\033[0m /knowledge journal
\033[1;94m\tOpen AI Bash (For Developer Only):\033[0m /open ai_bash
@@ -665,7 +665,7 @@ async def main():
'/history $num $offset',
'/connect $target',
'/contact $name',
'/knowledge add email | dir',
'/knowledge add dir',
'/knowledge journal [$topn]',
'/set_config $key',
'/enable $feature',