From 4a8d81ce160dcdef2bc6f3a292c4ba8b70f42842 Mon Sep 17 00:00:00 2001 From: Si Changjun Date: Thu, 24 Aug 2023 15:54:52 +0800 Subject: [PATCH 01/29] knowleadge.drawio --- doc/未命名绘图.drawio | 10 ++++++++++ 1 file changed, 10 insertions(+) create mode 100644 doc/未命名绘图.drawio diff --git a/doc/未命名绘图.drawio b/doc/未命名绘图.drawio new file mode 100644 index 0000000..f8db6cf --- /dev/null +++ b/doc/未命名绘图.drawio @@ -0,0 +1,10 @@ + + + + + + + + + + From 076eeaf977d015ef57c095d13cef374f2ecbe149 Mon Sep 17 00:00:00 2001 From: Si Changjun Date: Thu, 24 Aug 2023 15:54:52 +0800 Subject: [PATCH 02/29] knowleadge.drawio --- doc/未命名绘图.drawio | 10 ++++++++++ 1 file changed, 10 insertions(+) create mode 100644 doc/未命名绘图.drawio diff --git a/doc/未命名绘图.drawio b/doc/未命名绘图.drawio new file mode 100644 index 0000000..f8db6cf --- /dev/null +++ b/doc/未命名绘图.drawio @@ -0,0 +1,10 @@ + + + + + + + + + + From 67677ba40fcd11300c3a86cdb47e1682e8f35d30 Mon Sep 17 00:00:00 2001 From: tsukasa Date: Tue, 29 Aug 2023 16:45:05 +0800 Subject: [PATCH 03/29] objected knowledge base design --- doc/mvp/knowleadge.drawio | 1478 +++++++++++++++++++++++++++++++++++++ doc/mvp/knowleage.md | 20 + doc/未命名绘图.drawio | 10 - 3 files changed, 1498 insertions(+), 10 deletions(-) create mode 100644 doc/mvp/knowleadge.drawio create mode 100644 doc/mvp/knowleage.md delete mode 100644 doc/未命名绘图.drawio diff --git a/doc/mvp/knowleadge.drawio b/doc/mvp/knowleadge.drawio new file mode 100644 index 0000000..bdb5534 --- /dev/null +++ b/doc/mvp/knowleadge.drawio @@ -0,0 +1,1478 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/doc/mvp/knowleage.md b/doc/mvp/knowleage.md new file mode 100644 index 0000000..7b48cac --- /dev/null +++ b/doc/mvp/knowleage.md @@ -0,0 +1,20 @@ +# Objected knowleadge base, a specialized implemention for emails +## Vectorized Knowledge +Large language models are trained on general corpora and without fine-tuning on user-specific data, they struggle to utilize user-related context effectively. + +Users accumulate a vast amount of content that reflects their personality during their regular internet usage. This includes personal photos, tweets, Facebook posts, emails, etc. While it's possible to include all this content in the prompt during each interaction with the large language model, this approach is costly and can easily reach the token limit. + +A common solution is to generate feature vectors from this content using word embedding techniques and store them in a vector database. During an interaction, the vector that is most relevant to the prompt is retrieved from the database, merged with the prompt, and then passed to the large language model. + +We refer to this vectorized content as "knowledge". + +## Objected knownleadge base +In a personal AI system, to build a user's own knowledge base, we first need to implement various spider programs to crawl and retrieve all user-related data. Modern web content is typically rich text, including text, images, videos, hyperlinks, etc. Organizing this rich text in a tree-like structure similar to HTML is necessary, hence the need to introduce an object structure to represent this content. + +Different parts of this content cannot be vectorized using the same embedding model. For instance, text and images, as well as the content of an image and its EXIF information, need separate embeddings. This means that in the vector database, the same content may have multiple vector values, and a row can represent a whole content item or just a part of it. + +We need a comprehensive object structure to represent the hierarchy and relationships of content, as well as to implement the indexing and storage of objects. In the Minimum Viable Product (MVP) version, we'll implement a specialized solution for email content. In future versions, we can generalize this to handle other types of content, such as Facebook posts, tweets, etc. + +## Agent with knowleadge base +At the same time, we also need to explore the paradigm of using the knowledge base in Agents and workflows, so that the agent can better complete tasks in interaction with users through the context provided by the knowledge base. + diff --git a/doc/未命名绘图.drawio b/doc/未命名绘图.drawio deleted file mode 100644 index f8db6cf..0000000 --- a/doc/未命名绘图.drawio +++ /dev/null @@ -1,10 +0,0 @@ - - - - - - - - - - From a7960861d12fd34411ce18af1b58d77e7630d826 Mon Sep 17 00:00:00 2001 From: tsukasa Date: Wed, 30 Aug 2023 17:16:37 +0800 Subject: [PATCH 04/29] objected knowledge base implement for email code structure --- doc/mvp/workflow.drawio | 6 +- doc/未命名绘图.drawio | 10 ---- src/aios_kernel/knowledge/__init__.py | 2 + src/aios_kernel/knowledge/knowledge_base.py | 10 ++++ src/aios_kernel/knowledge/object.py | 65 +++++++++++++++++++++ src/aios_kernel/knowledge/object_storage.py | 33 +++++++++++ src/aios_kernel/knowledge/vector_base.py | 14 +++++ src/service/email_spider/converter.py | 24 ++++++++ src/service/email_spider/main.py | 12 ++++ src/service/email_spider/spider.py | 17 ++++++ 10 files changed, 180 insertions(+), 13 deletions(-) delete mode 100644 doc/未命名绘图.drawio create mode 100644 src/aios_kernel/knowledge/__init__.py create mode 100644 src/aios_kernel/knowledge/knowledge_base.py create mode 100644 src/aios_kernel/knowledge/object.py create mode 100644 src/aios_kernel/knowledge/object_storage.py create mode 100644 src/aios_kernel/knowledge/vector_base.py create mode 100644 src/service/email_spider/converter.py create mode 100644 src/service/email_spider/main.py create mode 100644 src/service/email_spider/spider.py diff --git a/doc/mvp/workflow.drawio b/doc/mvp/workflow.drawio index e49374a..c6659c1 100644 --- a/doc/mvp/workflow.drawio +++ b/doc/mvp/workflow.drawio @@ -1,6 +1,6 @@ - + @@ -123,7 +123,7 @@ - + @@ -220,7 +220,7 @@ - + diff --git a/doc/未命名绘图.drawio b/doc/未命名绘图.drawio deleted file mode 100644 index f8db6cf..0000000 --- a/doc/未命名绘图.drawio +++ /dev/null @@ -1,10 +0,0 @@ - - - - - - - - - - diff --git a/src/aios_kernel/knowledge/__init__.py b/src/aios_kernel/knowledge/__init__.py new file mode 100644 index 0000000..13bb147 --- /dev/null +++ b/src/aios_kernel/knowledge/__init__.py @@ -0,0 +1,2 @@ +from .object import EmailObject, ImageObject, TextChunkObject, ObjectID +from .knowledge_base import KnowledgeBase diff --git a/src/aios_kernel/knowledge/knowledge_base.py b/src/aios_kernel/knowledge/knowledge_base.py new file mode 100644 index 0000000..59e4f25 --- /dev/null +++ b/src/aios_kernel/knowledge/knowledge_base.py @@ -0,0 +1,10 @@ +# define a knowledge base class +from src.aios_kernel.agent import AgentPrompt +from .object import KnowledgeObject + +class KnowledgeBase: + async def insert(self, object: KnowledgeObject): + pass + + async def query(self, prompt: AgentPrompt) -> AgentPrompt: + pass \ No newline at end of file diff --git a/src/aios_kernel/knowledge/object.py b/src/aios_kernel/knowledge/object.py new file mode 100644 index 0000000..320cd8f --- /dev/null +++ b/src/aios_kernel/knowledge/object.py @@ -0,0 +1,65 @@ + +# define a object type enum +from abc import ABC + + +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): + 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): + 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 = [] + self.images = [] + + diff --git a/src/aios_kernel/knowledge/object_storage.py b/src/aios_kernel/knowledge/object_storage.py new file mode 100644 index 0000000..1a6f85e --- /dev/null +++ b/src/aios_kernel/knowledge/object_storage.py @@ -0,0 +1,33 @@ +# 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/aios_kernel/knowledge/vector_base.py b/src/aios_kernel/knowledge/vector_base.py new file mode 100644 index 0000000..b2cf442 --- /dev/null +++ b/src/aios_kernel/knowledge/vector_base.py @@ -0,0 +1,14 @@ +# import the ObjectID class +from .object import ObjectID + +# define a vector base class +class VectorBase: + def __init__(self, db_url, model_name) -> None: + self.db_url = db_url + self.model_name = model_name + + async def insert(self, vector: [float], id: ObjectID): + pass + + async def query(self, vector: [float], top_k: int) -> [ObjectID]: + pass \ No newline at end of file diff --git a/src/service/email_spider/converter.py b/src/service/email_spider/converter.py new file mode 100644 index 0000000..193b220 --- /dev/null +++ b/src/service/email_spider/converter.py @@ -0,0 +1,24 @@ +from aios_kernel.knowledge import KnowledgeBase, EmailObject + +# define a email converter class + +class EmailConverter: + # define init method + def __init__(self, local_dir, knowledge_base: KnowledgeBase) -> None: + pass + + async def run(self): + # convert the email to knowledge object + for email_dir in self._next(): + # convert the email to knowledge object + knowledge_object = self._convert(email_dir) + # insert the knowledge object to knowledge base + await self.knowledge_base.insert(knowledge_object) + + def _next(self) -> str: + pass + + def _convert(self, email_dir) -> EmailObject: + pass + + \ No newline at end of file diff --git a/src/service/email_spider/main.py b/src/service/email_spider/main.py new file mode 100644 index 0000000..cb75abd --- /dev/null +++ b/src/service/email_spider/main.py @@ -0,0 +1,12 @@ +import asyncio +from .spider import EmailSpider, EmailConverter + + +if __name__ == "__main__": + spider = EmailSpider("smtp.163.com","user","pwd","./email") + asyncio.run(spider.run()) + + converter = EmailConverter("./email",KnowledgeBase()) + asyncio.run(converter.run()) + + diff --git a/src/service/email_spider/spider.py b/src/service/email_spider/spider.py new file mode 100644 index 0000000..0b5a212 --- /dev/null +++ b/src/service/email_spider/spider.py @@ -0,0 +1,17 @@ +# define a email spider class + +class EmailSpider: + def __init__(self, address, account, pwd, local_dir) -> None: + pass + + async def run(self): + # spide the email from the email server + for email_link in self._next(): + # save the email to local directory + self._save(email_link) + + def _next(self): + pass + + def _save(self, email_link) -> str: + pass \ No newline at end of file From fc0ce3b381f3884d1a1064a36a694f913de3faff Mon Sep 17 00:00:00 2001 From: Si Changjun Date: Thu, 24 Aug 2023 15:54:52 +0800 Subject: [PATCH 05/29] knowleadge.drawio --- doc/未命名绘图.drawio | 10 ++++++++++ 1 file changed, 10 insertions(+) create mode 100644 doc/未命名绘图.drawio diff --git a/doc/未命名绘图.drawio b/doc/未命名绘图.drawio new file mode 100644 index 0000000..f8db6cf --- /dev/null +++ b/doc/未命名绘图.drawio @@ -0,0 +1,10 @@ + + + + + + + + + + From aa05a3a003dd423ce4d6256f1d7cce494fdf4de7 Mon Sep 17 00:00:00 2001 From: tsukasa Date: Tue, 29 Aug 2023 16:45:05 +0800 Subject: [PATCH 06/29] objected knowledge base design --- doc/mvp/knowleadge.drawio | 1478 +++++++++++++++++++++++++++++++++++++ doc/mvp/knowleage.md | 20 + doc/未命名绘图.drawio | 10 - 3 files changed, 1498 insertions(+), 10 deletions(-) create mode 100644 doc/mvp/knowleadge.drawio create mode 100644 doc/mvp/knowleage.md delete mode 100644 doc/未命名绘图.drawio diff --git a/doc/mvp/knowleadge.drawio b/doc/mvp/knowleadge.drawio new file mode 100644 index 0000000..bdb5534 --- /dev/null +++ b/doc/mvp/knowleadge.drawio @@ -0,0 +1,1478 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/doc/mvp/knowleage.md b/doc/mvp/knowleage.md new file mode 100644 index 0000000..7b48cac --- /dev/null +++ b/doc/mvp/knowleage.md @@ -0,0 +1,20 @@ +# Objected knowleadge base, a specialized implemention for emails +## Vectorized Knowledge +Large language models are trained on general corpora and without fine-tuning on user-specific data, they struggle to utilize user-related context effectively. + +Users accumulate a vast amount of content that reflects their personality during their regular internet usage. This includes personal photos, tweets, Facebook posts, emails, etc. While it's possible to include all this content in the prompt during each interaction with the large language model, this approach is costly and can easily reach the token limit. + +A common solution is to generate feature vectors from this content using word embedding techniques and store them in a vector database. During an interaction, the vector that is most relevant to the prompt is retrieved from the database, merged with the prompt, and then passed to the large language model. + +We refer to this vectorized content as "knowledge". + +## Objected knownleadge base +In a personal AI system, to build a user's own knowledge base, we first need to implement various spider programs to crawl and retrieve all user-related data. Modern web content is typically rich text, including text, images, videos, hyperlinks, etc. Organizing this rich text in a tree-like structure similar to HTML is necessary, hence the need to introduce an object structure to represent this content. + +Different parts of this content cannot be vectorized using the same embedding model. For instance, text and images, as well as the content of an image and its EXIF information, need separate embeddings. This means that in the vector database, the same content may have multiple vector values, and a row can represent a whole content item or just a part of it. + +We need a comprehensive object structure to represent the hierarchy and relationships of content, as well as to implement the indexing and storage of objects. In the Minimum Viable Product (MVP) version, we'll implement a specialized solution for email content. In future versions, we can generalize this to handle other types of content, such as Facebook posts, tweets, etc. + +## Agent with knowleadge base +At the same time, we also need to explore the paradigm of using the knowledge base in Agents and workflows, so that the agent can better complete tasks in interaction with users through the context provided by the knowledge base. + diff --git a/doc/未命名绘图.drawio b/doc/未命名绘图.drawio deleted file mode 100644 index f8db6cf..0000000 --- a/doc/未命名绘图.drawio +++ /dev/null @@ -1,10 +0,0 @@ - - - - - - - - - - From 39eb96f9e330f5e7556a88f41624b525e6ecd29b Mon Sep 17 00:00:00 2001 From: tsukasa Date: Wed, 30 Aug 2023 17:16:37 +0800 Subject: [PATCH 07/29] objected knowledge base implement for email code structure --- doc/mvp/workflow.drawio | 6 +- doc/未命名绘图.drawio | 10 ---- src/aios_kernel/knowledge/__init__.py | 2 + src/aios_kernel/knowledge/knowledge_base.py | 10 ++++ src/aios_kernel/knowledge/object.py | 65 +++++++++++++++++++++ src/aios_kernel/knowledge/object_storage.py | 33 +++++++++++ src/aios_kernel/knowledge/vector_base.py | 14 +++++ src/service/email_spider/converter.py | 24 ++++++++ src/service/email_spider/main.py | 12 ++++ src/service/email_spider/spider.py | 17 ++++++ 10 files changed, 180 insertions(+), 13 deletions(-) delete mode 100644 doc/未命名绘图.drawio create mode 100644 src/aios_kernel/knowledge/__init__.py create mode 100644 src/aios_kernel/knowledge/knowledge_base.py create mode 100644 src/aios_kernel/knowledge/object.py create mode 100644 src/aios_kernel/knowledge/object_storage.py create mode 100644 src/aios_kernel/knowledge/vector_base.py create mode 100644 src/service/email_spider/converter.py create mode 100644 src/service/email_spider/main.py create mode 100644 src/service/email_spider/spider.py diff --git a/doc/mvp/workflow.drawio b/doc/mvp/workflow.drawio index 2b3e239..c6659c1 100644 --- a/doc/mvp/workflow.drawio +++ b/doc/mvp/workflow.drawio @@ -1,6 +1,6 @@ - + @@ -123,7 +123,7 @@ - + @@ -220,7 +220,7 @@ - + diff --git a/doc/未命名绘图.drawio b/doc/未命名绘图.drawio deleted file mode 100644 index f8db6cf..0000000 --- a/doc/未命名绘图.drawio +++ /dev/null @@ -1,10 +0,0 @@ - - - - - - - - - - diff --git a/src/aios_kernel/knowledge/__init__.py b/src/aios_kernel/knowledge/__init__.py new file mode 100644 index 0000000..13bb147 --- /dev/null +++ b/src/aios_kernel/knowledge/__init__.py @@ -0,0 +1,2 @@ +from .object import EmailObject, ImageObject, TextChunkObject, ObjectID +from .knowledge_base import KnowledgeBase diff --git a/src/aios_kernel/knowledge/knowledge_base.py b/src/aios_kernel/knowledge/knowledge_base.py new file mode 100644 index 0000000..59e4f25 --- /dev/null +++ b/src/aios_kernel/knowledge/knowledge_base.py @@ -0,0 +1,10 @@ +# define a knowledge base class +from src.aios_kernel.agent import AgentPrompt +from .object import KnowledgeObject + +class KnowledgeBase: + async def insert(self, object: KnowledgeObject): + pass + + async def query(self, prompt: AgentPrompt) -> AgentPrompt: + pass \ No newline at end of file diff --git a/src/aios_kernel/knowledge/object.py b/src/aios_kernel/knowledge/object.py new file mode 100644 index 0000000..320cd8f --- /dev/null +++ b/src/aios_kernel/knowledge/object.py @@ -0,0 +1,65 @@ + +# define a object type enum +from abc import ABC + + +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): + 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): + 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 = [] + self.images = [] + + diff --git a/src/aios_kernel/knowledge/object_storage.py b/src/aios_kernel/knowledge/object_storage.py new file mode 100644 index 0000000..1a6f85e --- /dev/null +++ b/src/aios_kernel/knowledge/object_storage.py @@ -0,0 +1,33 @@ +# 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/aios_kernel/knowledge/vector_base.py b/src/aios_kernel/knowledge/vector_base.py new file mode 100644 index 0000000..b2cf442 --- /dev/null +++ b/src/aios_kernel/knowledge/vector_base.py @@ -0,0 +1,14 @@ +# import the ObjectID class +from .object import ObjectID + +# define a vector base class +class VectorBase: + def __init__(self, db_url, model_name) -> None: + self.db_url = db_url + self.model_name = model_name + + async def insert(self, vector: [float], id: ObjectID): + pass + + async def query(self, vector: [float], top_k: int) -> [ObjectID]: + pass \ No newline at end of file diff --git a/src/service/email_spider/converter.py b/src/service/email_spider/converter.py new file mode 100644 index 0000000..193b220 --- /dev/null +++ b/src/service/email_spider/converter.py @@ -0,0 +1,24 @@ +from aios_kernel.knowledge import KnowledgeBase, EmailObject + +# define a email converter class + +class EmailConverter: + # define init method + def __init__(self, local_dir, knowledge_base: KnowledgeBase) -> None: + pass + + async def run(self): + # convert the email to knowledge object + for email_dir in self._next(): + # convert the email to knowledge object + knowledge_object = self._convert(email_dir) + # insert the knowledge object to knowledge base + await self.knowledge_base.insert(knowledge_object) + + def _next(self) -> str: + pass + + def _convert(self, email_dir) -> EmailObject: + pass + + \ No newline at end of file diff --git a/src/service/email_spider/main.py b/src/service/email_spider/main.py new file mode 100644 index 0000000..cb75abd --- /dev/null +++ b/src/service/email_spider/main.py @@ -0,0 +1,12 @@ +import asyncio +from .spider import EmailSpider, EmailConverter + + +if __name__ == "__main__": + spider = EmailSpider("smtp.163.com","user","pwd","./email") + asyncio.run(spider.run()) + + converter = EmailConverter("./email",KnowledgeBase()) + asyncio.run(converter.run()) + + diff --git a/src/service/email_spider/spider.py b/src/service/email_spider/spider.py new file mode 100644 index 0000000..0b5a212 --- /dev/null +++ b/src/service/email_spider/spider.py @@ -0,0 +1,17 @@ +# define a email spider class + +class EmailSpider: + def __init__(self, address, account, pwd, local_dir) -> None: + pass + + async def run(self): + # spide the email from the email server + for email_link in self._next(): + # save the email to local directory + self._save(email_link) + + def _next(self): + pass + + def _save(self, email_link) -> str: + pass \ No newline at end of file From febe28eecc3dc9e4e7c339956dfb87597d47ef85 Mon Sep 17 00:00:00 2001 From: tsukasa Date: Thu, 31 Aug 2023 15:45:02 +0800 Subject: [PATCH 08/29] Add embeding compute task --- doc/mvp/workflow.drawio | 2 +- src/aios_kernel/compute_task.py | 8 +++++ src/aios_kernel/open_ai_node.py | 54 +++++++++++++++++++-------------- 3 files changed, 40 insertions(+), 24 deletions(-) diff --git a/doc/mvp/workflow.drawio b/doc/mvp/workflow.drawio index c6659c1..e011651 100644 --- a/doc/mvp/workflow.drawio +++ b/doc/mvp/workflow.drawio @@ -123,7 +123,7 @@ - + diff --git a/src/aios_kernel/compute_task.py b/src/aios_kernel/compute_task.py index f2253e0..ea1d7dc 100644 --- a/src/aios_kernel/compute_task.py +++ b/src/aios_kernel/compute_task.py @@ -38,6 +38,14 @@ class ComputeTask: self.params["model_name"] = "gpt-4-0613" self.params["max_token_size"] = max_token_size + def set_embeding_params(self, model_name, input, callchain_id = None): + self.task_type = "embeding" + self.create_time = time.time() + self.task_id = uuid.uuid4().hex + self.callchain_id = callchain_id + self.params["model_name"] = model_name + self.params["input"] = input + def display(self) -> str: return f"ComputeTask: {self.task_id} {self.task_type} {self.state}" diff --git a/src/aios_kernel/open_ai_node.py b/src/aios_kernel/open_ai_node.py index 7f5c8e6..be0d184 100644 --- a/src/aios_kernel/open_ai_node.py +++ b/src/aios_kernel/open_ai_node.py @@ -47,31 +47,39 @@ class OpenAI_ComputeNode(ComputeNode): def _run_task(self,task:ComputeTask): task.state = ComputeTaskState.RUNNING - mode_name = task.params["model_name"] - # max_token_size = task.params["max_token_size"] - prompts = task.params["prompts"] + # switch tsak type + if task.task_type == "llm_completion": + mode_name = task.params["model_name"] + # max_token_size = task.params["max_token_size"] + prompts = task.params["prompts"] - logger.info(f"call openai {mode_name} prompts: {prompts}") - resp = openai.ChatCompletion.create(model=mode_name, - messages=prompts, - max_tokens=4000, - temperature=1.2) - logger.info(f"openai response: {resp}") - - status_code = resp["choices"][0]["finish_reason"] - if status_code != "stop": - task.state = ComputeTaskState.ERROR - task.error_str =f"The status code was {status_code}." - return None - - result = ComputeTaskResult() - result.set_from_task(task) - result.worker_id = self.node_id - result.result_str = resp["choices"][0]["message"]["content"] - result.result = resp["choices"][0]["message"] - - return result + mode_name = task.params["model_name"] + # max_token_size = task.params["max_token_size"] + prompts = task.params["prompts"] + logger.info(f"call openai {mode_name} prompts: {prompts}") + resp = openai.ChatCompletion.create(model=mode_name, + messages=prompts, + max_tokens=4000, + temperature=1.2) + logger.info(f"openai response: {resp}") + + status_code = resp["choices"][0]["finish_reason"] + if status_code != "stop": + task.state = ComputeTaskState.ERROR + task.error_str =f"The status code was {status_code}." + return None + + result = ComputeTaskResult() + result.set_from_task(task) + result.worker_id = self.node_id + result.result_str = resp["choices"][0]["message"]["content"] + result.result = resp["choices"][0]["message"] + + return result + if task.task_type == "embeding": + pass + def start(self): async def _run_task_loop(): while True: From 5bd9775e4f31e32258b262878f846ec9a2360caf Mon Sep 17 00:00:00 2001 From: tsukasa Date: Thu, 31 Aug 2023 16:32:20 +0800 Subject: [PATCH 09/29] Add openai text embedding compute task --- src/aios_kernel/compute_kernel.py | 30 +++++++++ src/aios_kernel/compute_node.py | 2 +- src/aios_kernel/compute_task.py | 9 ++- src/aios_kernel/knowledge/knowledge_base.py | 4 +- src/aios_kernel/open_ai_node.py | 70 +++++++++++++-------- 5 files changed, 84 insertions(+), 31 deletions(-) diff --git a/src/aios_kernel/compute_kernel.py b/src/aios_kernel/compute_kernel.py index 1d05afd..5829bbc 100644 --- a/src/aios_kernel/compute_kernel.py +++ b/src/aios_kernel/compute_kernel.py @@ -118,4 +118,34 @@ class ComputeKernel: return task_req.result.result_str return "error!" + + def text_embedding(self,input:str,model_name:Optional[str] = None): + task_req = ComputeTask() + task_req.set_text_embeding_params(input,model_name) + self.run(task_req) + return task_req + + async def do_text_embedding(self,input:str,model_name:Optional[str] = None) -> [float] + task_req = self.text_embedding(input,model_name) + async def check_timer(): + check_times = 0 + while True: + if task_req.state == ComputeTaskState.DONE: + break + + if task_req.state == ComputeTaskState.ERROR: + break + + if check_times >= 20: + task_req.state = ComputeTaskState.ERROR + break + + await asyncio.sleep(0.5) + check_times += 1 + + await asyncio.create_task(check_timer()) + if task_req.state == ComputeTaskState.DONE: + return task_req.result.result + + return "error!" \ No newline at end of file diff --git a/src/aios_kernel/compute_node.py b/src/aios_kernel/compute_node.py index 6fe0059..ef8f933 100644 --- a/src/aios_kernel/compute_node.py +++ b/src/aios_kernel/compute_node.py @@ -28,7 +28,7 @@ class ComputeNode(ABC): pass @abstractmethod - def is_support(self,task_type:str) -> bool: + def is_support(self, task: ComputeTask) -> bool: pass @abstractmethod diff --git a/src/aios_kernel/compute_task.py b/src/aios_kernel/compute_task.py index ea1d7dc..35a1afd 100644 --- a/src/aios_kernel/compute_task.py +++ b/src/aios_kernel/compute_task.py @@ -38,12 +38,15 @@ class ComputeTask: self.params["model_name"] = "gpt-4-0613" self.params["max_token_size"] = max_token_size - def set_embeding_params(self, model_name, input, callchain_id = None): - self.task_type = "embeding" + def set_text_embeding_params(self, input, model_name=None, callchain_id = None): + self.task_type = "text_embedding" self.create_time = time.time() self.task_id = uuid.uuid4().hex self.callchain_id = callchain_id - self.params["model_name"] = model_name + if model_name is not None: + self.params["model_name"] = model_name + else: + self.params["model_name"] = "text-embedding-ada-002" self.params["input"] = input def display(self) -> str: diff --git a/src/aios_kernel/knowledge/knowledge_base.py b/src/aios_kernel/knowledge/knowledge_base.py index 59e4f25..a576280 100644 --- a/src/aios_kernel/knowledge/knowledge_base.py +++ b/src/aios_kernel/knowledge/knowledge_base.py @@ -7,4 +7,6 @@ class KnowledgeBase: pass async def query(self, prompt: AgentPrompt) -> AgentPrompt: - pass \ No newline at end of file + pass + + \ No newline at end of file diff --git a/src/aios_kernel/open_ai_node.py b/src/aios_kernel/open_ai_node.py index be0d184..27177f0 100644 --- a/src/aios_kernel/open_ai_node.py +++ b/src/aios_kernel/open_ai_node.py @@ -53,32 +53,45 @@ class OpenAI_ComputeNode(ComputeNode): # max_token_size = task.params["max_token_size"] prompts = task.params["prompts"] - mode_name = task.params["model_name"] + mode_name = task.params["model_name"] # max_token_size = task.params["max_token_size"] - prompts = task.params["prompts"] + prompts = task.params["prompts"] - logger.info(f"call openai {mode_name} prompts: {prompts}") - resp = openai.ChatCompletion.create(model=mode_name, - messages=prompts, - max_tokens=4000, - temperature=1.2) - logger.info(f"openai response: {resp}") - - status_code = resp["choices"][0]["finish_reason"] - if status_code != "stop": - task.state = ComputeTaskState.ERROR - task.error_str =f"The status code was {status_code}." - return None - - result = ComputeTaskResult() - result.set_from_task(task) - result.worker_id = self.node_id - result.result_str = resp["choices"][0]["message"]["content"] - result.result = resp["choices"][0]["message"] - - return result - if task.task_type == "embeding": - pass + logger.info(f"call openai {mode_name} prompts: {prompts}") + resp = openai.ChatCompletion.create(model=mode_name, + messages=prompts, + max_tokens=4000, + temperature=1.2) + logger.info(f"openai response: {resp}") + + status_code = resp["choices"][0]["finish_reason"] + if status_code != "stop": + task.state = ComputeTaskState.ERROR + task.error_str =f"The status code was {status_code}." + return None + + result = ComputeTaskResult() + result.set_from_task(task) + result.worker_id = self.node_id + result.result_str = resp["choices"][0]["message"]["content"] + result.result = resp["choices"][0]["message"] + + return result + if task.task_type == "text_embedding": + model_name = task.params["model_name"] + input = task.params["input"] + logger.info(f"call openai {model_name} input: {input}") + + resp = openai.Embeding.create(model=model_name, + input=input) + logger.info(f"openai response: {resp}") + + result = ComputeTaskResult() + result.set_from_task(task) + result.worker_id = self.node_id + result.result = resp["data"][0]["embedding"] + + return result def start(self): async def _run_task_loop(): @@ -104,8 +117,13 @@ class OpenAI_ComputeNode(ComputeNode): pass - def is_support(self,task_type:str) -> bool: - return True + def is_support(self, task: ComputeTask) -> bool: + if task.task_type == "llm_completion": + return True + if task.task_type == "text_embedding": + if task.params["model_name"] == "text-embedding-ada-002": + return True + return False def is_local(self) -> bool: From cff34c6aba53cde3f9aaa235e9ba6dde5d81b1f9 Mon Sep 17 00:00:00 2001 From: tsukasa Date: Thu, 31 Aug 2023 20:13:41 +0800 Subject: [PATCH 10/29] knowledge base basic flow --- src/aios_kernel/compute_kernel.py | 2 +- src/aios_kernel/knowledge/knowledge_base.py | 46 +++++++++++++++++++-- src/aios_kernel/knowledge/object.py | 8 ++-- 3 files changed, 47 insertions(+), 9 deletions(-) diff --git a/src/aios_kernel/compute_kernel.py b/src/aios_kernel/compute_kernel.py index 5829bbc..f626578 100644 --- a/src/aios_kernel/compute_kernel.py +++ b/src/aios_kernel/compute_kernel.py @@ -125,7 +125,7 @@ class ComputeKernel: self.run(task_req) return task_req - async def do_text_embedding(self,input:str,model_name:Optional[str] = None) -> [float] + async def do_text_embedding(self,input:str,model_name:Optional[str] = None) -> [float]: task_req = self.text_embedding(input,model_name) async def check_timer(): check_times = 0 diff --git a/src/aios_kernel/knowledge/knowledge_base.py b/src/aios_kernel/knowledge/knowledge_base.py index a576280..bb06401 100644 --- a/src/aios_kernel/knowledge/knowledge_base.py +++ b/src/aios_kernel/knowledge/knowledge_base.py @@ -1,12 +1,50 @@ # define a knowledge base class -from src.aios_kernel.agent import AgentPrompt -from .object import KnowledgeObject +from ..agent import AgentPrompt +from ..compute_kernel import ComputeKernel +from .object import KnowledgeObject, ObjectType, EmailObject, TextChunkObject, ImageObject +from .object_storage import ObjectStorage +from .vector_base import VectorBase class KnowledgeBase: + def __init__(self) -> None: + self.object_store = ObjectStorage() + self.vector_base = VectorBase() + self.compute_kernel = ComputeKernel() + async def insert(self, object: KnowledgeObject): - pass + if object.object_type == ObjectType.Email: + email: EmailObject = object + for text_id in email.text: + [text, _] = self.object_store.get(text_id) + text: TextChunkObject = text + vector = await self.compute_kernel.do_text_embedding(text.text) + self.vector_base.insert(vector, text_id) + + for image_id in email.images: + [image, _] = self.object_store.get(image_id) + image: ImageObject = image + vector = await self.compute_kernel.do_text_embedding(image.meta) + self.vector_base.insert(vector, image_id) + + vector = await self.compute_kernel.do_text_embedding(email.meta) + self.vector_base.insert(vector, email.get_id()) + else: + pass async def query(self, prompt: AgentPrompt) -> AgentPrompt: - pass + for msg in prompt.messages: + if msg.role == "user": + vector = await self.compute_kernel.do_text_embedding(msg.content) + object_ids = self.vector_base.query(vector, 10) + for object_id in object_ids: + if object_id.object_type == ObjectType.Email: + + [object, email] = self.object_store.get(object_id) + if object.object_type == ObjectType.Email: + email: EmailObject = object + prompt.append(AgentPrompt()) + prompt + + \ No newline at end of file diff --git a/src/aios_kernel/knowledge/object.py b/src/aios_kernel/knowledge/object.py index 320cd8f..815d4e2 100644 --- a/src/aios_kernel/knowledge/object.py +++ b/src/aios_kernel/knowledge/object.py @@ -21,7 +21,7 @@ class ObjectID: # pylint: disable=too-few-public-methods # define a object class class KnowledgeObject(ABC): # pylint: disable=too-few-public-methods - def __init__(self, object_type): + def __init__(self, object_type: ObjectType): self.object_type = object_type @abstractmethod @@ -41,7 +41,7 @@ class KnowledgeObject(ABC): # pylint: disable=too-few-public-methods # define a text chunk class class TextChunkObject(KnowledgeObject): # pylint: disable=too-few-public-methods - def __init__(self, text): + def __init__(self, text: str): super().__init__(ObjectType.TextChunk) self.text = text @@ -59,7 +59,7 @@ class EmailObject(KnowledgeObject): # pylint: disable=too-few-public-methods def __init__(self, meta): super().__init__(ObjectType.Email) self.meta = meta - self.text = [] - self.images = [] + self.text = [ObjectID] + self.images = [ObjectID] From 188595beb623c81350ecd07de9fbfd01603b005a Mon Sep 17 00:00:00 2001 From: liyaxing Date: Tue, 5 Sep 2023 20:42:22 +0800 Subject: [PATCH 11/29] Add chromadb based vector store impl and test case --- .../{knowledge => }/knowledge_base.py | 20 ++++---- src/{aios_kernel => }/knowledge/__init__.py | 2 +- src/{aios_kernel => }/knowledge/object.py | 4 +- .../knowledge/object_storage.py | 0 src/knowledge/vector/__init__.py | 2 + src/knowledge/vector/chroma_store.py | 49 +++++++++++++++++++ .../vector}/vector_base.py | 5 +- test/test_vector_storage.py | 29 +++++++++++ 8 files changed, 96 insertions(+), 15 deletions(-) rename src/aios_kernel/{knowledge => }/knowledge_base.py (74%) rename src/{aios_kernel => }/knowledge/__init__.py (63%) rename src/{aios_kernel => }/knowledge/object.py (96%) rename src/{aios_kernel => }/knowledge/object_storage.py (100%) create mode 100644 src/knowledge/vector/__init__.py create mode 100644 src/knowledge/vector/chroma_store.py rename src/{aios_kernel/knowledge => knowledge/vector}/vector_base.py (76%) create mode 100644 test/test_vector_storage.py diff --git a/src/aios_kernel/knowledge/knowledge_base.py b/src/aios_kernel/knowledge_base.py similarity index 74% rename from src/aios_kernel/knowledge/knowledge_base.py rename to src/aios_kernel/knowledge_base.py index bb06401..a882ba7 100644 --- a/src/aios_kernel/knowledge/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -1,9 +1,8 @@ # define a knowledge base class -from ..agent import AgentPrompt -from ..compute_kernel import ComputeKernel -from .object import KnowledgeObject, ObjectType, EmailObject, TextChunkObject, ImageObject -from .object_storage import ObjectStorage -from .vector_base import VectorBase +from . import AgentPrompt, ComputeKernel +from ..knowledge.object import KnowledgeObject, ObjectType, EmailObject, TextChunkObject, ImageObject +from ..knowledge.object_storage import ObjectStorage +from ..knowledge.vector.vector_base import VectorBase class KnowledgeBase: def __init__(self) -> None: @@ -38,12 +37,11 @@ class KnowledgeBase: object_ids = self.vector_base.query(vector, 10) for object_id in object_ids: if object_id.object_type == ObjectType.Email: - - [object, email] = self.object_store.get(object_id) - if object.object_type == ObjectType.Email: - email: EmailObject = object - prompt.append(AgentPrompt()) - prompt + [object, email] = self.object_store.get(object_id) + if object.object_type == ObjectType.Email: + email: EmailObject = object + prompt.append(AgentPrompt()) + prompt diff --git a/src/aios_kernel/knowledge/__init__.py b/src/knowledge/__init__.py similarity index 63% rename from src/aios_kernel/knowledge/__init__.py rename to src/knowledge/__init__.py index 13bb147..c879518 100644 --- a/src/aios_kernel/knowledge/__init__.py +++ b/src/knowledge/__init__.py @@ -1,2 +1,2 @@ from .object import EmailObject, ImageObject, TextChunkObject, ObjectID -from .knowledge_base import KnowledgeBase +from .vector import * \ No newline at end of file diff --git a/src/aios_kernel/knowledge/object.py b/src/knowledge/object.py similarity index 96% rename from src/aios_kernel/knowledge/object.py rename to src/knowledge/object.py index 815d4e2..f7a127e 100644 --- a/src/aios_kernel/knowledge/object.py +++ b/src/knowledge/object.py @@ -1,7 +1,7 @@ # define a object type enum -from abc import ABC - +from abc import ABC, abstractmethod +from enum import Enum class ObjectType(Enum): TextChunk = 1 diff --git a/src/aios_kernel/knowledge/object_storage.py b/src/knowledge/object_storage.py similarity index 100% rename from src/aios_kernel/knowledge/object_storage.py rename to src/knowledge/object_storage.py diff --git a/src/knowledge/vector/__init__.py b/src/knowledge/vector/__init__.py new file mode 100644 index 0000000..b3f0d78 --- /dev/null +++ b/src/knowledge/vector/__init__.py @@ -0,0 +1,2 @@ +from .vector_base import VectorBase +from .chroma_store import ChromaVectorStore diff --git a/src/knowledge/vector/chroma_store.py b/src/knowledge/vector/chroma_store.py new file mode 100644 index 0000000..626cd4f --- /dev/null +++ b/src/knowledge/vector/chroma_store.py @@ -0,0 +1,49 @@ +from .vector_base import VectorBase +from ..object import ObjectID +import chromadb +import logging +import os + + +class ChromaVectorStore(VectorBase): + def __init__(self, db_url, model_name: str) -> None: + super().__init__(db_url, model_name) + + logging.info( + "will init chroma vector store, db={}, model={}".format(db_url, model_name) + ) + + directory = os.path.join( + os.path.dirname(__file__), "../../../rootfs/data/vector" + ) + logging.info("will use vector store: {}".format(directory)) + + client = chromadb.PersistentClient( + path=directory, settings=chromadb.Settings(anonymized_telemetry=False) + ) + # client = chromadb.Client() + + collection_name = "coll_{}".format(model_name) + logging.info("will init chroma colletion: %s", collection_name) + + collection = client.get_or_create_collection(collection_name) + self.collection = collection + + async def insert(self, vector: [float], id: ObjectID): + self.collection.add( + embeddings=vector, + ids=id, + ) + + async def query(self, vector: [float], top_k: int) -> [ObjectID]: + ret = self.collection.query( + query_embeddings=vector, + n_results=top_k, + ) + + return ret["ids"] + + async def delete(self, id: ObjectID): + self.collection.delete( + ids=id, + ) diff --git a/src/aios_kernel/knowledge/vector_base.py b/src/knowledge/vector/vector_base.py similarity index 76% rename from src/aios_kernel/knowledge/vector_base.py rename to src/knowledge/vector/vector_base.py index b2cf442..354c21e 100644 --- a/src/aios_kernel/knowledge/vector_base.py +++ b/src/knowledge/vector/vector_base.py @@ -1,5 +1,5 @@ # import the ObjectID class -from .object import ObjectID +from ..object import ObjectID # define a vector base class class VectorBase: @@ -11,4 +11,7 @@ class VectorBase: pass async def query(self, vector: [float], top_k: int) -> [ObjectID]: + pass + + async def delete(self, id: ObjectID): pass \ No newline at end of file diff --git a/test/test_vector_storage.py b/test/test_vector_storage.py new file mode 100644 index 0000000..7357d65 --- /dev/null +++ b/test/test_vector_storage.py @@ -0,0 +1,29 @@ +import sys +import os + +dir_path = os.path.dirname(os.path.realpath(__file__)) +print(dir_path) + +sys.path.append("{}/../src/".format(dir_path)) +print(sys.path) + +from knowledge import ChromaVectorStore + + +import asyncio +import unittest + + +async def test_vector(): + storage = ChromaVectorStore("", "test") + await storage.insert([1, 2, 3], "test") + ids = await storage.query([1, 2, 3], 10) + print(ids) + +class TestVectorSTorage(unittest.TestCase): + def test_run(self): + asyncio.run(test_vector()) + + +if __name__ == "__main__": + unittest.main() From 1e5b9afd0f2fdd36bee255c75be1c45c8847e83f Mon Sep 17 00:00:00 2001 From: liyaxing Date: Thu, 7 Sep 2023 20:19:18 +0800 Subject: [PATCH 12/29] Add chunk relate impl and test case --- src/knowledge/__init__.py | 5 ++- src/knowledge/data/__init__.py | 2 + src/knowledge/data/chunk.py | 43 +++++++++++++++++++ src/knowledge/data/tracker.py | 77 ++++++++++++++++++++++++++++++++++ test/test_chunk.py | 36 ++++++++++++++++ 5 files changed, 161 insertions(+), 2 deletions(-) create mode 100644 src/knowledge/data/__init__.py create mode 100644 src/knowledge/data/chunk.py create mode 100644 src/knowledge/data/tracker.py create mode 100644 test/test_chunk.py diff --git a/src/knowledge/__init__.py b/src/knowledge/__init__.py index c879518..6edb83e 100644 --- a/src/knowledge/__init__.py +++ b/src/knowledge/__init__.py @@ -1,2 +1,3 @@ -from .object import EmailObject, ImageObject, TextChunkObject, ObjectID -from .vector import * \ No newline at end of file +from .object import * +from .vector import * +from .data import * \ No newline at end of file diff --git a/src/knowledge/data/__init__.py b/src/knowledge/data/__init__.py new file mode 100644 index 0000000..81c672a --- /dev/null +++ b/src/knowledge/data/__init__.py @@ -0,0 +1,2 @@ +from .chunk import ChunkID, PositionType, PositionFileRange +from .tracker import ChunkTracker \ No newline at end of file diff --git a/src/knowledge/data/chunk.py b/src/knowledge/data/chunk.py new file mode 100644 index 0000000..ff355bd --- /dev/null +++ b/src/knowledge/data/chunk.py @@ -0,0 +1,43 @@ +from enum import Enum +from ..object import ObjectID + +ChunkID = ObjectID + +class PositionType(Enum): + Unknown = 1 + Device = 2 + File = 3 + FileRange = 4 + ChunkManager = 5 + + +class PositionFileRange: + def __init__(self, path: str, range_begin: int, range_end: int): + self.path = path + self.range_begin = range_begin + self.range_end = range_end + + def encode(self): + return f"{self.range_begin}:{self.range_end}:{self.path}" + + @staticmethod + def decode(value: str): + parts = value.split(":") + if len(parts) < 3: + raise ValueError("Invalid input string") + + try: + range_begin = int(parts[0]) + range_end = int(parts[1]) + except ValueError as e: + raise ValueError("Invalid range_begin or range_end string") from e + + path = ":".join(parts[2:]) + return PositionFileRange(path, range_begin, range_end) + + def __str__(self): + return self.encode() + + @staticmethod + def from_string(value: str): + return PositionFileRange.decode(value) diff --git a/src/knowledge/data/tracker.py b/src/knowledge/data/tracker.py new file mode 100644 index 0000000..efe4ede --- /dev/null +++ b/src/knowledge/data/tracker.py @@ -0,0 +1,77 @@ +import sqlite3 +import time +import logging +import os +from .chunk import ChunkID, PositionType, PositionFileRange + + +class ChunkTracker: + _instance = None + + def __new__(cls): + if cls._instance is None: + cls._instance = super().__new__(cls) + directory = os.path.join( + os.path.dirname(__file__), "../../../rootfs/data/chunk/" + ) + + if not os.path.exists(directory): + os.mkdir(directory) + file = os.path.join(directory, "chunk_tracker.db") + logging.info(f"will init chunk tracker, db={file}") + cls._instance.__singleton_init__(file) + + return cls._instance + + def __singleton_init__(self, db_path: str): + logging.info(f"will init chunk tracker, db={db_path}") + + self.conn = sqlite3.connect(db_path) + self.cursor = self.conn.cursor() + self.cursor.execute( + """ + CREATE TABLE IF NOT EXISTS chunks ( + id TEXT NOT NULL, + pos TEXT NOT NULL, + pos_type TINYINT NOT NULL, + insert_time UNSIGNED BIG INT NOT NULL, + update_time UNSIGNED BIG INT NOT NULL, + flags INTEGER DEFAULT 0, + PRIMARY KEY(id, pos, pos_type) + ) + """ + ) + self.conn.commit() + + def add_position(self, chunk_id: ChunkID, position: str, position_type: PositionType): + logging.debug(f"add chunk position: {chunk_id}, {position}, {position_type}") + + insert_time = update_time = int(time.time()) + self.cursor.execute( + """ + INSERT OR REPLACE INTO chunks (id, pos, pos_type, insert_time, update_time) + VALUES (?, ?, ?, ?, ?) + """, + (str(chunk_id), position, position_type.value, insert_time, update_time,), + ) + self.conn.commit() + + def remove_position(self, chunk_id: ChunkID): + logging.info(f"remove chunk position: {chunk_id}") + + self.cursor.execute( + """ + DELETE FROM chunks WHERE id = ? + """, + (str(chunk_id),), + ) + self.conn.commit() + + def get_position(self, chunk_id: ChunkID) -> (str, PositionType): + self.cursor.execute( + """ + SELECT pos, pos_type FROM chunks WHERE id = ? + """, + (str(chunk_id),), + ) + return self.cursor.fetchone() diff --git a/test/test_chunk.py b/test/test_chunk.py new file mode 100644 index 0000000..40a7aeb --- /dev/null +++ b/test/test_chunk.py @@ -0,0 +1,36 @@ +import sys +import os + +dir_path = os.path.dirname(os.path.realpath(__file__)) +print(dir_path) + +sys.path.append("{}/../src/".format(dir_path)) +print(sys.path) + +from knowledge import ChunkTracker, ChunkID, HashValue, PositionType + + +import asyncio +import unittest + + +class TestChunk(unittest.TestCase): + def test_chunk_tracker(self): + tracker = ChunkTracker() + + hash = HashValue.hash_data("1234567890".encode("utf-8")); + cid = ChunkID.new_chunk_id(hash) + print(cid) + + tracker.add_position(cid, "/tmp/1", PositionType.File) + ret = tracker.get_position(cid) + print(ret[0]) + + tracker.remove_position(cid) + ret = tracker.get_position(cid) + self.assertEqual(ret, None) + + + +if __name__ == "__main__": + unittest.main() From e8e3812d287eeae5b4275883b0927928932280bd Mon Sep 17 00:00:00 2001 From: liyaxing Date: Fri, 8 Sep 2023 15:52:05 +0800 Subject: [PATCH 13/29] Add chunk store and chunk list helper --- src/knowledge/data/__init__.py | 4 +- src/knowledge/data/chunk.py | 2 +- src/knowledge/data/chunk_store.py | 26 +++++++++ src/knowledge/data/gen.py | 94 +++++++++++++++++++++++++++++++ src/knowledge/data/tracker.py | 40 ++++++------- test/test_chunk.py | 26 ++++++++- 6 files changed, 164 insertions(+), 28 deletions(-) create mode 100644 src/knowledge/data/chunk_store.py create mode 100644 src/knowledge/data/gen.py diff --git a/src/knowledge/data/__init__.py b/src/knowledge/data/__init__.py index 81c672a..f8e69aa 100644 --- a/src/knowledge/data/__init__.py +++ b/src/knowledge/data/__init__.py @@ -1,2 +1,4 @@ from .chunk import ChunkID, PositionType, PositionFileRange -from .tracker import ChunkTracker \ No newline at end of file +from .tracker import ChunkTracker +from .chunk_store import ChunkStore +from .gen import ChunkListGenerator \ No newline at end of file diff --git a/src/knowledge/data/chunk.py b/src/knowledge/data/chunk.py index ff355bd..5267dc3 100644 --- a/src/knowledge/data/chunk.py +++ b/src/knowledge/data/chunk.py @@ -8,7 +8,7 @@ class PositionType(Enum): Device = 2 File = 3 FileRange = 4 - ChunkManager = 5 + ChunkStore = 5 class PositionFileRange: diff --git a/src/knowledge/data/chunk_store.py b/src/knowledge/data/chunk_store.py new file mode 100644 index 0000000..d06eb38 --- /dev/null +++ b/src/knowledge/data/chunk_store.py @@ -0,0 +1,26 @@ +import os +import logging +from ..object import FileBlobStorage +from .chunk import ChunkID + + +class ChunkStore: + def __init__(self, root_dir: str): + logging.info(f"will init chunk store, root_dir={root_dir}") + + if not os.path.exists(root_dir): + os.makedirs(root_dir) + + self.root = root_dir + self.blob = FileBlobStorage(root_dir) + + def put_chunk(self, chunk_id: ChunkID, contents: bytes): + self.blob.put(chunk_id, contents) + + def get_chunk(self, chunk_id: ChunkID) -> bytes: + return self.blob.get(chunk_id) + + def delete_chunk(self, chunk_id: ChunkID): + self.blob.delete(chunk_id) + + \ No newline at end of file diff --git a/src/knowledge/data/gen.py b/src/knowledge/data/gen.py new file mode 100644 index 0000000..a6d9425 --- /dev/null +++ b/src/knowledge/data/gen.py @@ -0,0 +1,94 @@ +import os +import hashlib +import re +from typing import Tuple, List +from .chunk_store import ChunkStore +from .chunk import ChunkID, PositionFileRange, PositionType +from ..object import HashValue +from .tracker import ChunkTracker + + +class ChunkListGenerator: + def __init__(self, chunk_store: ChunkStore, chunk_tracker: ChunkTracker): + self.chunk_store = chunk_store + self.chunk_tracker = chunk_tracker + + def create_chunk_list_from_file( + self, file_path: str, chunk_size: int, restore: bool + ) -> Tuple[List[ChunkID], HashValue]: + assert ( + chunk_size % (1024 * 1024) == 0 + ), "chunk size should be an integral multiple of 1MB" + chunk_list = [] + hash_obj = hashlib.sha256() + + with open(file_path, "rb") as file: + while True: + chunk = file.read(chunk_size) + if not chunk: + break + chunk_id = ChunkID.hash_data(chunk) + chunk_list.append(chunk_id) + + hash_obj.update(chunk) + + if restore: + self.chunk_tracker.add_position( + chunk_id, file_path, PositionType.ChunkStore + ) + self.chunk_store.put_chunk(chunk_id, chunk) + else: + file_range = PositionFileRange( + file_path, file.tell() - chunk_size, chunk_size + ) + self.chunk_tracker.add_position( + chunk_id, file_range, PositionType.FileRange + ) + + file_hash = HashValue(hash_obj.digest()) + print(f"calc file hash: {file_path}, {file_hash}") + return chunk_list, file_hash + + def create_chunk_list_from_text( + self, text: str, chunk_max_words: int, separator_chars: str = ".," + ) -> Tuple[List[ChunkID], HashValue]: + text_list = self._split_text_list(text, chunk_max_words, separator_chars) + chunk_list = [] + hash_obj = hashlib.sha256() + + for text in text_list: + chunk_bytes = text.encode("utf-8") + hash_obj.update(chunk_bytes) + + chunk_id = ChunkID.hash_data(chunk_bytes) + chunk_list.append(chunk_id) + self.chunk_tracker.add_position( + chunk_id, "", PositionType.ChunkStore + ) + self.chunk_store.put_chunk(chunk_id, chunk_bytes) + + hash = HashValue(hash_obj.digest()) + return 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 + diff --git a/src/knowledge/data/tracker.py b/src/knowledge/data/tracker.py index efe4ede..2905d11 100644 --- a/src/knowledge/data/tracker.py +++ b/src/knowledge/data/tracker.py @@ -2,31 +2,17 @@ import sqlite3 import time import logging import os -from .chunk import ChunkID, PositionType, PositionFileRange +from .chunk import ChunkID, PositionType, PositionFileRange class ChunkTracker: - _instance = None - - def __new__(cls): - if cls._instance is None: - cls._instance = super().__new__(cls) - directory = os.path.join( - os.path.dirname(__file__), "../../../rootfs/data/chunk/" - ) - - if not os.path.exists(directory): - os.mkdir(directory) - file = os.path.join(directory, "chunk_tracker.db") - logging.info(f"will init chunk tracker, db={file}") - cls._instance.__singleton_init__(file) - - return cls._instance + def __init__(self, root_dir: str): + if not os.path.exists(root_dir): + os.makedirs(root_dir) + file = os.path.join(root_dir, "chunk_tracker.db") + logging.info(f"will init chunk tracker, db={file}") - def __singleton_init__(self, db_path: str): - logging.info(f"will init chunk tracker, db={db_path}") - - self.conn = sqlite3.connect(db_path) + self.conn = sqlite3.connect(file) self.cursor = self.conn.cursor() self.cursor.execute( """ @@ -43,7 +29,9 @@ class ChunkTracker: ) self.conn.commit() - def add_position(self, chunk_id: ChunkID, position: str, position_type: PositionType): + def add_position( + self, chunk_id: ChunkID, position: str, position_type: PositionType + ): logging.debug(f"add chunk position: {chunk_id}, {position}, {position_type}") insert_time = update_time = int(time.time()) @@ -52,7 +40,13 @@ class ChunkTracker: INSERT OR REPLACE INTO chunks (id, pos, pos_type, insert_time, update_time) VALUES (?, ?, ?, ?, ?) """, - (str(chunk_id), position, position_type.value, insert_time, update_time,), + ( + str(chunk_id), + position, + position_type.value, + insert_time, + update_time, + ), ) self.conn.commit() diff --git a/test/test_chunk.py b/test/test_chunk.py index 40a7aeb..b2f2e7e 100644 --- a/test/test_chunk.py +++ b/test/test_chunk.py @@ -7,16 +7,27 @@ print(dir_path) sys.path.append("{}/../src/".format(dir_path)) print(sys.path) -from knowledge import ChunkTracker, ChunkID, HashValue, PositionType + +import logging +import sys + +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) +from knowledge import ChunkTracker, ChunkID, HashValue, PositionType, KnowledgeStore, ChunkListGenerator import asyncio import unittest class TestChunk(unittest.TestCase): def test_chunk_tracker(self): - tracker = ChunkTracker() + tracker = KnowledgeStore().get_chunk_tracker() hash = HashValue.hash_data("1234567890".encode("utf-8")); cid = ChunkID.new_chunk_id(hash) @@ -29,8 +40,17 @@ class TestChunk(unittest.TestCase): tracker.remove_position(cid) ret = tracker.get_position(cid) self.assertEqual(ret, None) - + + def test_chunk(self): + gen = ChunkListGenerator(KnowledgeStore().get_chunk_store(), KnowledgeStore().get_chunk_tracker()) + gen.create_chunk_list_from_file("H:/test", 1024 * 1024, True) + + # Read the file + text_file = "H:/test.txt" + with open(text_file, 'r', encoding = "utf-8") as file: + text = file.read() + gen.create_chunk_list_from_text(text, 1024) if __name__ == "__main__": unittest.main() From 36476eb5b4078e8fd4f5edfdd4aee90887b515ef Mon Sep 17 00:00:00 2001 From: liyaxing Date: Mon, 11 Sep 2023 00:46:56 +0800 Subject: [PATCH 14/29] Add object and chunk base impl and test cases --- src/aios_kernel/knowledge_base.py | 2 +- src/knowledge/__init__.py | 3 +- src/knowledge/data/__init__.py | 3 +- src/knowledge/data/chunk_list.py | 14 +++++ src/knowledge/data/{gen.py => writer.py} | 26 ++++---- src/knowledge/object.py | 65 ------------------- src/knowledge/object/__init__.py | 5 ++ src/knowledge/object/blob.py | 54 ++++++++++++++++ src/knowledge/object/hash.py | 42 +++++++++++++ src/knowledge/object/object.py | 79 ++++++++++++++++++++++++ src/knowledge/object/object_id.py | 53 ++++++++++++++++ src/knowledge/object/object_store.py | 24 +++++++ src/knowledge/object_storage.py | 33 ---------- src/knowledge/store.py | 46 ++++++++++++++ test/test_chunk.py | 30 ++++++--- test/test_object.py | 31 ++++++++++ 16 files changed, 385 insertions(+), 125 deletions(-) create mode 100644 src/knowledge/data/chunk_list.py rename src/knowledge/data/{gen.py => writer.py} (89%) delete mode 100644 src/knowledge/object.py create mode 100644 src/knowledge/object/__init__.py create mode 100644 src/knowledge/object/blob.py create mode 100644 src/knowledge/object/hash.py create mode 100644 src/knowledge/object/object.py create mode 100644 src/knowledge/object/object_id.py create mode 100644 src/knowledge/object/object_store.py delete mode 100644 src/knowledge/object_storage.py create mode 100644 src/knowledge/store.py create mode 100644 test/test_object.py diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index a882ba7..4f4ba4f 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -1,7 +1,7 @@ # define a knowledge base class from . import AgentPrompt, ComputeKernel from ..knowledge.object import KnowledgeObject, ObjectType, EmailObject, TextChunkObject, ImageObject -from ..knowledge.object_storage import ObjectStorage +from ..knowledge.store import ObjectStorage from ..knowledge.vector.vector_base import VectorBase class KnowledgeBase: diff --git a/src/knowledge/__init__.py b/src/knowledge/__init__.py index 6edb83e..e2cdeb2 100644 --- a/src/knowledge/__init__.py +++ b/src/knowledge/__init__.py @@ -1,3 +1,4 @@ from .object import * from .vector import * -from .data import * \ No newline at end of file +from .data import * +from .store import KnowledgeStore \ No newline at end of file diff --git a/src/knowledge/data/__init__.py b/src/knowledge/data/__init__.py index f8e69aa..2b391c9 100644 --- a/src/knowledge/data/__init__.py +++ b/src/knowledge/data/__init__.py @@ -1,4 +1,5 @@ from .chunk import ChunkID, PositionType, PositionFileRange from .tracker import ChunkTracker from .chunk_store import ChunkStore -from .gen import ChunkListGenerator \ No newline at end of file +from .writer import ChunkListWriter +from .chunk_list import ChunkList diff --git a/src/knowledge/data/chunk_list.py b/src/knowledge/data/chunk_list.py new file mode 100644 index 0000000..969b03f --- /dev/null +++ b/src/knowledge/data/chunk_list.py @@ -0,0 +1,14 @@ +from ..object import HashValue +from .chunk import ChunkID +from typing import List + +class ChunkList: + def __init__(self, chunk_list: List[ChunkID], hash: HashValue): + self.chunk_list = chunk_list + self.hash = hash + + def __str__(self): + return self.hash.to_base58() + + def __repr__(self): + return f"chunk_list: {self.chunk_list}, hash: {self.hash}" \ No newline at end of file diff --git a/src/knowledge/data/gen.py b/src/knowledge/data/writer.py similarity index 89% rename from src/knowledge/data/gen.py rename to src/knowledge/data/writer.py index a6d9425..7f79447 100644 --- a/src/knowledge/data/gen.py +++ b/src/knowledge/data/writer.py @@ -6,16 +6,16 @@ from .chunk_store import ChunkStore from .chunk import ChunkID, PositionFileRange, PositionType from ..object import HashValue from .tracker import ChunkTracker +from .chunk_list import ChunkList - -class ChunkListGenerator: +class ChunkListWriter: def __init__(self, chunk_store: ChunkStore, chunk_tracker: ChunkTracker): self.chunk_store = chunk_store self.chunk_tracker = chunk_tracker def create_chunk_list_from_file( self, file_path: str, chunk_size: int, restore: bool - ) -> Tuple[List[ChunkID], HashValue]: + ) -> ChunkList: assert ( chunk_size % (1024 * 1024) == 0 ), "chunk size should be an integral multiple of 1MB" @@ -47,7 +47,8 @@ class ChunkListGenerator: file_hash = HashValue(hash_obj.digest()) print(f"calc file hash: {file_path}, {file_hash}") - return chunk_list, file_hash + + return ChunkList(chunk_list, file_hash) def create_chunk_list_from_text( self, text: str, chunk_max_words: int, separator_chars: str = ".," @@ -55,20 +56,18 @@ class ChunkListGenerator: text_list = self._split_text_list(text, chunk_max_words, separator_chars) chunk_list = [] hash_obj = hashlib.sha256() - + for text in text_list: - chunk_bytes = text.encode("utf-8") + chunk_bytes = text.encode("utf-8") hash_obj.update(chunk_bytes) - + chunk_id = ChunkID.hash_data(chunk_bytes) chunk_list.append(chunk_id) - self.chunk_tracker.add_position( - chunk_id, "", PositionType.ChunkStore - ) + self.chunk_tracker.add_position(chunk_id, "", PositionType.ChunkStore) self.chunk_store.put_chunk(chunk_id, chunk_bytes) - + hash = HashValue(hash_obj.digest()) - return chunk_list, hash + return ChunkList(chunk_list, hash) @staticmethod def _split_text_list( @@ -90,5 +89,4 @@ class ChunkListGenerator: word_count = 1 if chunk: chunk_list.append(" ".join(chunk)) - return chunk_list - + return chunk_list \ 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/__init__.py b/src/knowledge/object/__init__.py new file mode 100644 index 0000000..ef1db5c --- /dev/null +++ b/src/knowledge/object/__init__.py @@ -0,0 +1,5 @@ +from .object import EmailObject, ImageObject, TextChunkObject, ObjectID +from .blob import FileBlobStorage +from .hash import HashValue, hash_data +from .object_store import ObjectStore +from .object_id import ObjectID, ObjectType diff --git a/src/knowledge/object/blob.py b/src/knowledge/object/blob.py new file mode 100644 index 0000000..b005e62 --- /dev/null +++ b/src/knowledge/object/blob.py @@ -0,0 +1,54 @@ +import os +import shutil +from .object import ObjectID + + +class FileBlobStorage: + def __init__(self, root): + self.root = root + + def get_full_path(self, object_id: ObjectID, auto_create: bool = True): + if os.name == "nt": # Windows + hash_str = object_id.to_base36() + len = 3 + else: + hash_str = str(object_id) + len = 2 + + tmp, first = hash_str[:-len], hash_str[-len:] + second = tmp[-len:] + + if os.name == "nt": # Windows + if second in ["con", "aux", "nul", "prn"]: + second = tmp[-(len + 1) :] + if first in ["con", "aux", "nul", "prn"]: + first = f"{first}_" + + path = os.path.join(self.root, first, second) + if auto_create and not os.path.exists(path): + os.makedirs(path) + + path = os.path.join(path, hash_str) + + return path + + def write_sync(self, path: str, contents: bytes): + with open(path, "wb") as f: + f.write(contents) + + def put(self, object_id: ObjectID, contents: bytes): + full_path = self.get_full_path(object_id) + self.write_sync(full_path, contents) + + def get(self, object_id: ObjectID) -> bytes: + full_path = self.get_full_path(object_id) + 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) + + def exists(self, object_id: ObjectID) -> bool: + full_path = self.get_full_path(object_id) + return os.path.exists(full_path) diff --git a/src/knowledge/object/hash.py b/src/knowledge/object/hash.py new file mode 100644 index 0000000..5063814 --- /dev/null +++ b/src/knowledge/object/hash.py @@ -0,0 +1,42 @@ +import hashlib +import base58 +import base36 + +class HashValue: + def __init__(self, value: bytes): + assert len(value) == 32, "HashValue must be 32 bytes long" + self.value = value + + def __str__(self) -> str: + return self.to_base58() + + @staticmethod + def hash_data(data): + return hash_data(data) + + def to_base58(self): + return base58.b58encode(self.value).decode() + + @staticmethod + def from_base58(s): + return HashValue(base58.b58decode(s)) + + def to_base36(self): + # Convert the bytes to int before encoding + num = int.from_bytes(self.value, 'big') + return base36.dumps(num) + + @staticmethod + def from_base36(s): + # Decode to int and then convert to bytes + num = base36.loads(s) + return HashValue(num.to_bytes((num.bit_length() + 7) // 8, 'big')) + + +HASH_VALUE_LEN = 32 + + +def hash_data(data: bytes): + sha256 = hashlib.sha256() + sha256.update(data) + return HashValue(sha256.digest()) diff --git a/src/knowledge/object/object.py b/src/knowledge/object/object.py new file mode 100644 index 0000000..0eff535 --- /dev/null +++ b/src/knowledge/object/object.py @@ -0,0 +1,79 @@ + +# define a object type enum +from abc import ABC, abstractmethod +from enum import Enum +from .object_id import ObjectID +import hashlib +import json +import pickle + +class KnowledgeObject(ABC): + def __init__(self, object_type: int, desc: dict = {}, body: dict = {}): + self.desc = desc + self.body = body + self.object_type = object_type + + def get_object_type(self): + return self.object_type + + def object_id(self) -> ObjectID: + return self.calculate_id() + + def set_desc_with_key_value(self, key, value): + self.desc[key] = value + + def get_desc_with_key(self, key): + return self.desc.get(key) + + def get_desc(self) -> dict: + return self.desc + + def set_body_with_key_value(self, key, value): + self.body[key] = value + + def get_body_with_key(self, key): + return self.body.get(key) + + def get_body(self) -> dict: + return self.body + + def calculate_id(self): + # Convert the object_type and desc to string and compute the SHA256 hash + data = json.dumps({"object_type": self.object_type, "desc": self.desc}) + sha256 = hashlib.sha256() + sha256.update(data.encode()) + return ObjectID(sha256.digest()) + + def encode(self) -> bytes: + return pickle.dumps(self) + + @staticmethod + def decode(data: bytes): + return pickle.loads(data) + + + +# 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/object_id.py b/src/knowledge/object/object_id.py new file mode 100644 index 0000000..27c0063 --- /dev/null +++ b/src/knowledge/object/object_id.py @@ -0,0 +1,53 @@ +# define a object type enum +from abc import ABC, abstractmethod +from enum import Enum +from .hash import HashValue +import base58 +import base36 + + +class ObjectType(Enum): + Chunk = 7 + TextChunk = 100 + Image = 101 + Email = 102 + + +# define a object ID class to identify a object +class ObjectID: # pylint: disable=too-few-public-methods + def __init__(self, object_type: ObjectType, value: bytes): + assert len(value) == 32, "ObjectID must be 32 bytes long" + self.object_type = object_type + self.value = value + + def __str__(self): + return self.to_base58() + + def get_object_type(self): + return self.object_type + + def to_base58(self): + return base58.b58encode(self.value).decode() + + @staticmethod + def from_base58(s): + return ObjectID(base58.b58decode(s)) + + def to_base36(self): + # Convert the bytes to int before encoding + num = int.from_bytes(self.value, "big") + return base36.dumps(num) + + @staticmethod + def from_base36(s): + # Decode to int and then convert to bytes + num = base36.loads(s) + return ObjectID(num.to_bytes((num.bit_length() + 7) // 8, "big")) + + @staticmethod + def new_chunk_id(chunk_hash: HashValue): + return ObjectID(ObjectType.Chunk, chunk_hash.value) + + @staticmethod + def hash_data(data: bytes): + return ObjectID.new_chunk_id(HashValue.hash_data(data)) \ No newline at end of file diff --git a/src/knowledge/object/object_store.py b/src/knowledge/object/object_store.py new file mode 100644 index 0000000..d903b8f --- /dev/null +++ b/src/knowledge/object/object_store.py @@ -0,0 +1,24 @@ +import os +import logging +from .blob import FileBlobStorage +from .object_id import ObjectID + + +class ObjectStore: + def __init__(self, root_dir: str): + logging.info(f"will init object blob store, root_dir={root_dir}") + + blob_dir = os.path.join(root_dir, "blob") + if not os.path.exists(blob_dir): + logging.info(f"will create blob dir: {blob_dir}") + os.makedirs(blob_dir) + self.blob = FileBlobStorage(blob_dir) + + def put_object(self, object_id: ObjectID, contents: bytes): + self.blob.put(object_id, contents) + + def get_object(self, object_id: ObjectID) -> bytes: + return self.blob.get(object_id) + + def delete_object(self, object_id: ObjectID): + self.blob.delete(object_id) 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 new file mode 100644 index 0000000..a4a0af6 --- /dev/null +++ b/src/knowledge/store.py @@ -0,0 +1,46 @@ +import os +from .object import ObjectStore +from .data import ChunkStore, ChunkTracker +import logging + + +# KnowledgeStore class, 聚合ChunkStore，ChunkTracker， ObjectStore，并且是一个全局单例，可以方便的使用这三个内置的store示例 +class KnowledgeStore: + _instance = None + + def __new__(cls): + if cls._instance is None: + cls._instance = super().__new__(cls) + directory = os.path.join( + os.path.dirname(__file__), "../../rootfs/data/" + ) + directory = os.path.normpath(directory) + print(directory) + + if not os.path.exists(directory): + os.makedirs(directory) + + cls._instance.__singleton_init__(directory) + + return cls._instance + + def __singleton_init__(self, root_dir: str): + logging.info(f"will init knowledge store, root_dir={root_dir}") + + self.root = root_dir + + object_store_dir = os.path.join(root_dir, "object") + self.object_store = ObjectStore(object_store_dir) + + chunk_store_dir = os.path.join(root_dir, "chunk") + self.chunk_store = ChunkStore(chunk_store_dir) + self.chunk_tracker = ChunkTracker(chunk_store_dir) + + def get_object_store(self) -> ObjectStore: + return self.object_store + + def get_chunk_store(self) -> ChunkStore: + return self.chunk_store + + def get_chunk_tracker(self) -> ChunkTracker: + return self.chunk_tracker diff --git a/test/test_chunk.py b/test/test_chunk.py index b2f2e7e..f1f496b 100644 --- a/test/test_chunk.py +++ b/test/test_chunk.py @@ -15,12 +15,19 @@ 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') +formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") handler.setFormatter(formatter) root.addHandler(handler) -from knowledge import ChunkTracker, ChunkID, HashValue, PositionType, KnowledgeStore, ChunkListGenerator +from knowledge import ( + ChunkTracker, + ChunkID, + HashValue, + PositionType, + KnowledgeStore, + ChunkListWriter, +) import asyncio import unittest @@ -28,29 +35,32 @@ import unittest class TestChunk(unittest.TestCase): def test_chunk_tracker(self): tracker = KnowledgeStore().get_chunk_tracker() - - hash = HashValue.hash_data("1234567890".encode("utf-8")); + + hash = HashValue.hash_data("1234567890".encode("utf-8")) cid = ChunkID.new_chunk_id(hash) print(cid) - + tracker.add_position(cid, "/tmp/1", PositionType.File) ret = tracker.get_position(cid) print(ret[0]) - + tracker.remove_position(cid) ret = tracker.get_position(cid) self.assertEqual(ret, None) - + def test_chunk(self): - gen = ChunkListGenerator(KnowledgeStore().get_chunk_store(), KnowledgeStore().get_chunk_tracker()) + gen = ChunkListWriter( + KnowledgeStore().get_chunk_store(), KnowledgeStore().get_chunk_tracker() + ) gen.create_chunk_list_from_file("H:/test", 1024 * 1024, True) - + # Read the file text_file = "H:/test.txt" - with open(text_file, 'r', encoding = "utf-8") as file: + with open(text_file, "r", encoding="utf-8") as file: text = file.read() gen.create_chunk_list_from_text(text, 1024) + if __name__ == "__main__": unittest.main() diff --git a/test/test_object.py b/test/test_object.py new file mode 100644 index 0000000..6a603fd --- /dev/null +++ b/test/test_object.py @@ -0,0 +1,31 @@ +import sys +import os + +dir_path = os.path.dirname(os.path.realpath(__file__)) +print(dir_path) + +sys.path.append("{}/../src/".format(dir_path)) +print(sys.path) + +from knowledge import ObjectID, HashValue + + +import asyncio +import unittest + + +class TestVectorSTorage(unittest.TestCase): + def test_object(self): + data = HashValue.hash_data("1233".encode("utf-8")); + print(data.to_base58()) + print(data.to_base36()) + + data2 = HashValue.from_base58(data.to_base58()) + self.assertEqual(data.to_base36(), data2.to_base36()) + + data2 = HashValue.from_base36(data.to_base36()) + self.assertEqual(data.to_base58(), data2.to_base58()) + + +if __name__ == "__main__": + unittest.main() From 2761883ae81a2d8bb3048baef665c85b22b70039 Mon Sep 17 00:00:00 2001 From: liyaxing Date: Mon, 11 Sep 2023 02:05:24 +0800 Subject: [PATCH 15/29] Add some core objects impl --- src/knowledge/__init__.py | 3 +- src/knowledge/core_object/__init__.py | 4 + src/knowledge/core_object/document_object.py | 57 ++++++++++++ src/knowledge/core_object/email_object.py | 47 ++++++++++ src/knowledge/core_object/image_object.py | 89 +++++++++++++++++++ src/knowledge/core_object/rich_text_object.py | 80 +++++++++++++++++ src/knowledge/core_object/video_object.py | 84 +++++++++++++++++ src/knowledge/object/__init__.py | 2 +- src/knowledge/object/object.py | 31 +------ src/knowledge/object/object_id.py | 5 +- src/knowledge/store.py | 2 +- 11 files changed, 372 insertions(+), 32 deletions(-) create mode 100644 src/knowledge/core_object/__init__.py create mode 100644 src/knowledge/core_object/document_object.py create mode 100644 src/knowledge/core_object/email_object.py create mode 100644 src/knowledge/core_object/image_object.py create mode 100644 src/knowledge/core_object/rich_text_object.py create mode 100644 src/knowledge/core_object/video_object.py diff --git a/src/knowledge/__init__.py b/src/knowledge/__init__.py index e2cdeb2..680c632 100644 --- a/src/knowledge/__init__.py +++ b/src/knowledge/__init__.py @@ -1,4 +1,5 @@ from .object import * from .vector import * from .data import * -from .store import KnowledgeStore \ No newline at end of file +from .store import KnowledgeStore +from .core_object import * \ No newline at end of file diff --git a/src/knowledge/core_object/__init__.py b/src/knowledge/core_object/__init__.py new file mode 100644 index 0000000..065de74 --- /dev/null +++ b/src/knowledge/core_object/__init__.py @@ -0,0 +1,4 @@ +from .document_object import DocumentObject, DocumentObjectBuilder +from .image_object import ImageObject, ImageObjectBuilder +from .video_object import VideoObject, VideoObjectBuilder +from .email_object import EmailObject, EmailObjectBuilder \ No newline at end of file diff --git a/src/knowledge/core_object/document_object.py b/src/knowledge/core_object/document_object.py new file mode 100644 index 0000000..dd9bb8b --- /dev/null +++ b/src/knowledge/core_object/document_object.py @@ -0,0 +1,57 @@ +from ..object import KnowledgeObject +from ..data import ChunkList, ChunkListWriter +from ..object import ObjectType + + +# desc +# meta +# hash: "file-hash", +# tags: {} +# body +# chunk_list: [chunk_id, chunk_id, ...] + + +class DocumentObject(KnowledgeObject): + def __init__(self, meta: dict, tags: dict, chunk_list: ChunkList): + desc = dict() + body = dict() + desc["meta"] = meta + desc["tags"] = tags + desc["hash"] = chunk_list.hash.to_base58() + body["chunk_list"] = chunk_list.chunk_list + + super().__init__(ObjectType.Document, desc, body) + + def get_meta(self): + return self.desc["meta"] + + def get_tags(self): + return self.desc["tags"] + + def get_hash(self): + return self.desc["hash"] + + def get_chunk_list(self): + return self.body["chunk_list"] + + +class DocumentObjectBuilder: + def __init__(self, meta: dict, tags: dict, text: str): + self.meta = meta + self.tags = tags + self.text = text + + def set_meta(self, meta: dict): + self.meta = meta + return self + + def set_text(self, text: str): + self.text = text + return self + + def build(self) -> DocumentObject: + chunk_list = ChunkListWriter.create_chunk_list_from_text( + self.text, + 1024 * 4, + ) + return DocumentObject(self.meta, self.tags, chunk_list) diff --git a/src/knowledge/core_object/email_object.py b/src/knowledge/core_object/email_object.py new file mode 100644 index 0000000..fe5e0d4 --- /dev/null +++ b/src/knowledge/core_object/email_object.py @@ -0,0 +1,47 @@ +from .rich_text_object import RichTextObject, RichTextObjectBuilder +from ..object import ObjectID, ObjectType, KnowledgeObject + + +class EmailObject(KnowledgeObject): + def __init__(self, meta: dict, tags: dict, rich_text: RichTextObject): + desc = dict() + body = dict() + desc["meta"] = meta + desc["tags"] = tags + + # FIXME rich text content store in desc or body? which one is better? + body["content"] = rich_text + + super().__init__(ObjectType.Email, desc, body) + + def get_meta(self): + return self.desc["meta"] + + def get_tags(self): + return self.desc["tags"] + + def get_rich_text(self): + return self.body["rich_text"] + + +class EmailObjectBuilder: + def __init__(self, meta: dict, tags: dict, folder: str): + self.meta = meta + self.tags = tags + self.folder = folder + + def set_meta(self, meta: dict): + self.meta = meta + return self + + def set_tags(self, tags: dict): + self.tags = tags + return self + + def set_folder(self, folder: str): + self.folder = folder + return self + + def build(self) -> EmailObject: + content = RichTextObjectBuilder(self.folder).build() + return EmailObject(self.meta, self.tags, content) diff --git a/src/knowledge/core_object/image_object.py b/src/knowledge/core_object/image_object.py new file mode 100644 index 0000000..d6a0cde --- /dev/null +++ b/src/knowledge/core_object/image_object.py @@ -0,0 +1,89 @@ +from ..object import KnowledgeObject +from ..data import ChunkList, ChunkListWriter +from ..object import ObjectType + + +# desc +# meta +# tags +# hash: "file-hash", +# exif: {} +# body +# chunk_list: [chunk_id, chunk_id, ...] + + +class ImageObject(KnowledgeObject): + def __init__(self, meta: dict, tags: dict, exif: dict, chunk_list: ChunkList): + desc = dict() + body = dict() + desc["meta"] = meta + desc["exif"] = exif + desc + desc["hash"] = chunk_list.hash.to_base58() + body["chunk_list"] = chunk_list.chunk_list + + super().__init__(ObjectType.Image, desc, body) + + def get_meta(self): + return self.desc["meta"] + + def get_exif(self): + return self.desc["exif"] + + def get_tags(self): + return self.desc["tags"] + + def get_hash(self): + return self.desc["hash"] + + def get_chunk_list(self): + return self.body["chunk_list"] + + +from PIL import Image +from PIL.ExifTags import TAGS + + +def get_exif_data(image_path: str): + with Image.open(image_path) as image: + exif_data = image._getexif() + + if exif_data is not None: + return { + TAGS.get(key): exif_data[key] + for key in exif_data.keys() + if key in TAGS and isinstance(exif_data[key], (bytes, str)) + } + else: + return {} + + +class ImageObjectBuilder: + def __init__(self, meta: dict, tags: dict, image_file: str): + self.meta = meta + self.tags = tags + self.image_file = image_file + self.restore_file = False + + def set_meta(self, meta: dict): + self.meta = meta + return self + + def set_tags(self, tags: dict): + self.tags = tags + return self + + def set_image_file(self, image_file: str): + self.image_file = image_file + return self + + def set_restore_file(self, restore_file: bool): + self.restore_file = restore_file + return self + + def build(self) -> ImageObject: + chunk_list = ChunkListWriter.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) diff --git a/src/knowledge/core_object/rich_text_object.py b/src/knowledge/core_object/rich_text_object.py new file mode 100644 index 0000000..ac3fcbe --- /dev/null +++ b/src/knowledge/core_object/rich_text_object.py @@ -0,0 +1,80 @@ +from knowledge.object.object_id import ObjectType +from ..object import KnowledgeObject +from ..data import ChunkList, ChunkListWriter +from ..object import ObjectType +from .video_object import VideoObjectBuilder, VideoObject +from .image_object import ImageObjectBuilder, ImageObject +from .document_object import DocumentObjectBuilder, DocumentObject + +class RichTextObject(KnowledgeObject): + def __init__(self, images: dict = {}, videos: dict = {}, documents: dict = {}, rich_texts: dict = {}}): + desc = dict() + desc.images = dict() + desc.videos = dict() + desc.documents = dict() + desc.rich_texts = dict() + + super().__init__(ObjectType.RichText, desc) + + + def add_image_with_key(self, key, image_object: ImageObject): + assert self.desc.images[key] == None + self.desc.images[key] = image_object + + def add_image(self, image_object: ImageObject): + self.desc.images[image_object.object_id()] = image_object + + def get_image_with_key(self, key) -> ImageObject: + return self.desc.images[key] + + def get_images(self) -> dict: + return self.desc.images + + def add_video_with_key(self, key, video_object: VideoObject): + assert self.desc.videos[key] == None + self.desc.videos[key] = video_object + + def add_video(self, video_object: VideoObject): + self.desc.videos[video_object.object_id()] = video_object + + def get_video_with_key(self, key) -> VideoObject: + return self.desc.videos[key] + + def get_videos(self) -> dict: + return self.desc.videos + + + def add_document_with_key(self, key, document_object: DocumentObject): + assert self.desc.documents[key] == None + self.desc.documents[key] = document_object + + def add_document(self, document_object: DocumentObject): + self.desc.documents[document_object.object_id()] = document_object + + def get_document_with_key(self, key) -> DocumentObject: + return self.desc.documents[key] + + def get_documents(self) -> dict: + return self.desc.documents + + def add_rich_text_with_key(self, key, rich_text_object): + assert self.desc.rich_texts[key] == None + self.desc.rich_texts[key] = rich_text_object + + def add_rich_text(self, rich_text_object): + self.desc.rich_texts[rich_text_object.object_id()] = rich_text_object + + def get_rich_text_with_key(self, key): + return self.desc.rich_texts[key] + + def get_rich_texts(self) -> dict: + return self.desc.rich_texts + + +class RichTextObjectBuilder: + def __init__(self, folder: str): + self.folder = folder + + def build(self) -> RichTextObject: + # TODO + return RichTextObject() \ No newline at end of file diff --git a/src/knowledge/core_object/video_object.py b/src/knowledge/core_object/video_object.py new file mode 100644 index 0000000..ab4c262 --- /dev/null +++ b/src/knowledge/core_object/video_object.py @@ -0,0 +1,84 @@ +from ..object import KnowledgeObject +from ..data import ChunkList, ChunkListWriter +from ..object import ObjectType + + +# desc +# meta +# tags +# hash: "file-hash", +# info: {} +# body +# chunk_list: [chunk_id, chunk_id, ...] + + +class VideoObject(KnowledgeObject): + def __init__(self, meta: dict, tags: dict, info: dict, chunk_list: ChunkList): + desc = dict() + body = dict() + desc["meta"] = meta + desc["tags"] = tags + desc["info"] = info + desc["hash"] = chunk_list.hash.to_base58() + body["chunk_list"] = chunk_list.chunk_list + + super().__init__(ObjectType.Video, desc, body) + + def get_meta(self): + return self.desc["meta"] + + def get_tags(self): + return self.desc["tags"] + + def get_info(self): + return self.desc["info"] + + def get_hash(self): + return self.desc["hash"] + + def get_chunk_list(self): + return self.body["chunk_list"] + + +from moviepy.editor import VideoFileClip + + +def get_video_info(video_path: str) -> dict: + clip = VideoFileClip(video_path) + return { + "duration": clip.duration, # Duration in seconds + "fps": clip.fps, # Frames per second + "nframes": clip.reader.nframes, # Total number of frames + "size": clip.size, # Size of the frames (width, height) + } + + +class VideoObjectBuilder: + def __init__(self, meta: dict, tags: dict, video_file: str): + self.meta = meta + self.tags = tags + self.video_file = video_file + self.restore_file = False + + def set_meta(self, meta: dict): + self.meta = meta + return self + + def set_tags(self, tags: dict): + self.tags = tags + return self + + def set_video_file(self, video_file: str): + self.video_file = video_file + return self + + def set_restore_file(self, restore_file: bool): + self.restore_file = restore_file + return self + + def build(self) -> VideoObject: + chunk_list = ChunkListWriter.create_chunk_list_from_file( + self.video_file, 1024 * 1024 * 4, self.restore_file + ) + info = get_video_info(self.video_file) + return VideoObject(self.meta, self.tags, info, chunk_list) diff --git a/src/knowledge/object/__init__.py b/src/knowledge/object/__init__.py index ef1db5c..7886d7a 100644 --- a/src/knowledge/object/__init__.py +++ b/src/knowledge/object/__init__.py @@ -1,4 +1,4 @@ -from .object import EmailObject, ImageObject, TextChunkObject, ObjectID +from .object import KnowledgeObject from .blob import FileBlobStorage from .hash import HashValue, hash_data from .object_store import ObjectStore diff --git a/src/knowledge/object/object.py b/src/knowledge/object/object.py index 0eff535..b3841ed 100644 --- a/src/knowledge/object/object.py +++ b/src/knowledge/object/object.py @@ -2,18 +2,18 @@ # define a object type enum from abc import ABC, abstractmethod from enum import Enum -from .object_id import ObjectID +from .object_id import ObjectID, ObjectType import hashlib import json import pickle class KnowledgeObject(ABC): - def __init__(self, object_type: int, desc: dict = {}, body: dict = {}): + def __init__(self, object_type: ObjectType, desc: dict = {}, body: dict = {}): self.desc = desc self.body = body self.object_type = object_type - def get_object_type(self): + def get_object_type(self) -> ObjectType: return self.object_type def object_id(self) -> ObjectID: @@ -50,30 +50,5 @@ class KnowledgeObject(ABC): @staticmethod def decode(data: bytes): return pickle.loads(data) - - - -# 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/object_id.py b/src/knowledge/object/object_id.py index 27c0063..3c3ac24 100644 --- a/src/knowledge/object/object_id.py +++ b/src/knowledge/object/object_id.py @@ -10,7 +10,10 @@ class ObjectType(Enum): Chunk = 7 TextChunk = 100 Image = 101 - Email = 102 + Video = 102 + Document = 103 + RichText = 104 + Email = 105 # define a object ID class to identify a object diff --git a/src/knowledge/store.py b/src/knowledge/store.py index a4a0af6..ef697fa 100644 --- a/src/knowledge/store.py +++ b/src/knowledge/store.py @@ -4,7 +4,7 @@ from .data import ChunkStore, ChunkTracker import logging -# KnowledgeStore class, 聚合ChunkStore，ChunkTracker， ObjectStore，并且是一个全局单例，可以方便的使用这三个内置的store示例 +# 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: _instance = None From f3cc3f60435880d1e451e8bb6d31780e10f9a245 Mon Sep 17 00:00:00 2001 From: liyaxing Date: Mon, 11 Sep 2023 02:07:01 +0800 Subject: [PATCH 16/29] Update requirements --- .gitignore | 2 ++ src/requirements.txt | 5 +++++ 2 files changed, 7 insertions(+) create mode 100644 src/requirements.txt diff --git a/.gitignore b/.gitignore index 4ef429f..948ea70 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,4 @@ .vscode/ *.pyc +rootfs/data +*.log diff --git a/src/requirements.txt b/src/requirements.txt new file mode 100644 index 0000000..811e8e1 --- /dev/null +++ b/src/requirements.txt @@ -0,0 +1,5 @@ +chromadb==0.4 +openai==0.28 +toml==0.10 +Pillow==10.0 +moviepy==1.0 \ No newline at end of file From 58c0d74ca131c8c505aed52dc73102b4fa166370 Mon Sep 17 00:00:00 2001 From: liyaxing Date: Mon, 11 Sep 2023 17:30:06 +0800 Subject: [PATCH 17/29] Add chunk reader impl --- src/knowledge/data/__init__.py | 1 + src/knowledge/data/chunk_store.py | 4 +- src/knowledge/data/reader.py | 85 +++++++++++++++++++++++++++++++ src/knowledge/data/tracker.py | 6 +-- src/requirements.txt | 4 +- 5 files changed, 95 insertions(+), 5 deletions(-) create mode 100644 src/knowledge/data/reader.py diff --git a/src/knowledge/data/__init__.py b/src/knowledge/data/__init__.py index 2b391c9..3f470f0 100644 --- a/src/knowledge/data/__init__.py +++ b/src/knowledge/data/__init__.py @@ -3,3 +3,4 @@ from .tracker import ChunkTracker from .chunk_store import ChunkStore from .writer import ChunkListWriter from .chunk_list import ChunkList +from .reader import ChunkReader, Chunk \ No newline at end of file diff --git a/src/knowledge/data/chunk_store.py b/src/knowledge/data/chunk_store.py index d06eb38..666b027 100644 --- a/src/knowledge/data/chunk_store.py +++ b/src/knowledge/data/chunk_store.py @@ -22,5 +22,7 @@ class ChunkStore: def delete_chunk(self, chunk_id: ChunkID): self.blob.delete(chunk_id) - + + def get_chunk_file_path(self, chunk_id: ChunkID) -> str: + return self.blob.get_full_path(chunk_id, False) \ No newline at end of file diff --git a/src/knowledge/data/reader.py b/src/knowledge/data/reader.py new file mode 100644 index 0000000..c4cbb96 --- /dev/null +++ b/src/knowledge/data/reader.py @@ -0,0 +1,85 @@ +from .chunk import ChunkID, PositionType, PositionFileRange +from .chunk_store import ChunkStore +from .tracker import ChunkTracker +from ..object import HashValue +import logging +from typing import List +import hashlib + +class Chunk: + def __init__(self, file_path: str, range_start: int, size: int = -1): + self.file_path = file_path + self.range_start = range_start + self.size = size + + def read(self): + with open(self.file_path, 'rb') as f: + f.seek(self.range_start) + return f.read(self.size) + + + +class ChunkReader: + def __init__(self, chunk_store: ChunkStore, chunk_tracker: ChunkTracker): + self.chunk_store = chunk_store + self.chunk_tracker = chunk_tracker + + def get_chunk(self, chunk_id: ChunkID) -> Chunk: + positions = self.chunk_tracker.get_position(chunk_id) + if positions is None: + logging.warning(f"chunk not found: {chunk_id}") + return None + + if len(positions) == 0: + logging.warning(f"chunk not found: {chunk_id}") + return None + + for pos in positions: + [position, position_type] = pos + logging.info(f"chunk position: {chunk_id}, {position}, {position_type}") + if position_type == PositionType.ChunkStore: + file_path = self.chunk_store.get_chunk_file_path(chunk_id) + return Chunk(file_path, 0, -1) + elif position_type == PositionType.File: + return Chunk(position, 0, -1) + elif position_type == PositionType.FileRange: + file_range = PositionFileRange.decode(position) + return Chunk(file_range.path, file_range.range_begin, file_range.range_end - file_range.range_begin) + else: + raise ValueError(f"invalid position type: {position_type}") + + logging.error(f"chunk not found: {chunk_id}") + return None + + 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]): + 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() + + def read_text_chunk_list(self, chunk_ids: List[ChunkID]): + for chunk_id in chunk_ids: + chunk = self.get_chunk(chunk_id) + if chunk is None: + raise ValueError(f"text chunk not found: {chunk_id}") + + yield chunk.read().decode("utf-8") + + def calc_file_hash(self, file_path: str) -> HashValue: + hash_obj = hashlib.sha256() + with open(file_path, "rb") as file: + while True: + chunk = file.read(1024 * 1024) + if not chunk: + break + hash_obj.update(chunk) + return HashValue(hash_obj.digest()) + + def calc_text_hash(self, text: str) -> HashValue: + hash_obj = hashlib.sha256() + hash_obj.update(text.encode("utf-8")) \ No newline at end of file diff --git a/src/knowledge/data/tracker.py b/src/knowledge/data/tracker.py index 2905d11..57daa38 100644 --- a/src/knowledge/data/tracker.py +++ b/src/knowledge/data/tracker.py @@ -3,7 +3,7 @@ import time import logging import os from .chunk import ChunkID, PositionType, PositionFileRange - +from typing import List class ChunkTracker: def __init__(self, root_dir: str): @@ -61,11 +61,11 @@ class ChunkTracker: ) self.conn.commit() - def get_position(self, chunk_id: ChunkID) -> (str, PositionType): + def get_position(self, chunk_id: ChunkID) -> List[(str, PositionType)]: self.cursor.execute( """ SELECT pos, pos_type FROM chunks WHERE id = ? """, (str(chunk_id),), ) - return self.cursor.fetchone() + return self.cursor.fetchmany() diff --git a/src/requirements.txt b/src/requirements.txt index 811e8e1..1b90399 100644 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -2,4 +2,6 @@ chromadb==0.4 openai==0.28 toml==0.10 Pillow==10.0 -moviepy==1.0 \ No newline at end of file +moviepy==1.0 +base58==2.1 +base36==0.1 \ No newline at end of file From ddc103af38c79c562a94e895babcb6a357df3b48 Mon Sep 17 00:00:00 2001 From: liyaxing Date: Mon, 11 Sep 2023 20:32:52 +0800 Subject: [PATCH 18/29] Add EmailObjectBuilder impl and fix some bugs --- src/knowledge/core_object/document_object.py | 4 +- src/knowledge/core_object/email_object.py | 102 ++++++++++++++++-- src/knowledge/core_object/image_object.py | 4 +- src/knowledge/core_object/rich_text_object.py | 50 ++++----- src/knowledge/core_object/video_object.py | 4 +- src/knowledge/data/tracker.py | 4 +- src/knowledge/data/writer.py | 2 +- src/knowledge/store.py | 11 +- test/test_object.py | 14 ++- 9 files changed, 150 insertions(+), 45 deletions(-) diff --git a/src/knowledge/core_object/document_object.py b/src/knowledge/core_object/document_object.py index dd9bb8b..6d6270e 100644 --- a/src/knowledge/core_object/document_object.py +++ b/src/knowledge/core_object/document_object.py @@ -1,7 +1,7 @@ from ..object import KnowledgeObject from ..data import ChunkList, ChunkListWriter from ..object import ObjectType - +from .. import KnowledgeStore # desc # meta @@ -50,7 +50,7 @@ class DocumentObjectBuilder: return self def build(self) -> DocumentObject: - chunk_list = ChunkListWriter.create_chunk_list_from_text( + chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_text( self.text, 1024 * 4, ) diff --git a/src/knowledge/core_object/email_object.py b/src/knowledge/core_object/email_object.py index fe5e0d4..8b8efb5 100644 --- a/src/knowledge/core_object/email_object.py +++ b/src/knowledge/core_object/email_object.py @@ -1,5 +1,12 @@ from .rich_text_object import RichTextObject, RichTextObjectBuilder from ..object import ObjectID, ObjectType, KnowledgeObject +from .document_object import DocumentObjectBuilder +from .image_object import ImageObjectBuilder +from .video_object import VideoObjectBuilder +from .rich_text_object import RichTextObjectBuilder +import os +import json +import logging class EmailObject(KnowledgeObject): @@ -24,16 +31,35 @@ class EmailObject(KnowledgeObject): return self.body["rich_text"] +""" +EmailObject folder structure: +. +├── email.txt +└── meta.json + ├── image + │ ├── image1.jpg + │ ├── image2.jpg + │ └── ... + ├── video + │ ├── video1.mp4 + │ ├── video2.mv + │ └── ... + └── audio + ├── audio1.m4a + ├── audio2.flac + └── ... +EmailObjectBuilder will read the target folder and build the EmailObject +Store meta.json to meta in EmailObject +Store email.txt to DocumentObject and RichTextObject in EmailObject +Store very image file in image folder to ImageObject and RichTextObject in EmailObject, etc +""" + + class EmailObjectBuilder: - def __init__(self, meta: dict, tags: dict, folder: str): - self.meta = meta + def __init__(self, tags: dict, folder: str): self.tags = tags self.folder = folder - def set_meta(self, meta: dict): - self.meta = meta - return self - def set_tags(self, tags: dict): self.tags = tags return self @@ -43,5 +69,65 @@ class EmailObjectBuilder: return self def build(self) -> EmailObject: - content = RichTextObjectBuilder(self.folder).build() - return EmailObject(self.meta, self.tags, content) + # Read meta.json + meta = {} + meta_file = os.path.join(self.folder, "meta.json") + if os.path.exists(meta_file): + logging.info(f"Will read meta.json {meta_file}") + with open(meta_file, "r", encoding="utf-8") as f: + meta = json.load(f) + else: + logging.info(f"Meta file missing! {meta_file}") + + # Read email.txt + documents = {} + content_file = os.path.join(self.folder, "email.txt") + if os.path.exists(content_file): + logging.info(f"Will read email.txt {content_file}") + + try: + with open(content_file, "r", encoding="utf-8") as f: + text = f.read() + + document = DocumentObjectBuilder({}, {}, text).build() + documents = {"email.txt": document} + except Exception as e: + logging.error(f"Failed to read email.txt {content_file} {e}") + else: + logging.info(f"Content file missing! {content_file}") + + # Process image files + images = {} + image_dir = os.path.join(self.folder, "image") + if os.path.exists(image_dir): + for image_file in os.listdir(image_dir): + image_path = os.path.join(image_dir, image_file) + logging.info(f"Will read image file {image_path}") + + try: + image = ImageObjectBuilder({}, {}, image_path).build() + images[image_file] = image + except Exception as e: + logging.error(f"Failed to read image file {image_path} {e}") + continue + + # Process video files + videos = {} + video_dir = os.path.join(self.folder, "video") + if os.path.exists(video_dir): + for video_file in os.listdir(video_dir): + video_path = os.path.join(video_dir, video_file) + logging.info(f"Will read video file {video_path}") + + try: + video = VideoObjectBuilder({}, {}, video_path).build() + videos[video_file] = video + except Exception as e: + logging.error(f"Failed to read video file {video_path} {e}") + continue + + # Create RichTextObject + rich_text = RichTextObject(images, videos, documents) + + # Create EmailObject + return EmailObject(meta, {}, rich_text) diff --git a/src/knowledge/core_object/image_object.py b/src/knowledge/core_object/image_object.py index d6a0cde..83ab434 100644 --- a/src/knowledge/core_object/image_object.py +++ b/src/knowledge/core_object/image_object.py @@ -1,7 +1,7 @@ from ..object import KnowledgeObject from ..data import ChunkList, ChunkListWriter from ..object import ObjectType - +from .. import KnowledgeStore # desc # meta @@ -82,7 +82,7 @@ class ImageObjectBuilder: return self def build(self) -> ImageObject: - chunk_list = ChunkListWriter.create_chunk_list_from_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) diff --git a/src/knowledge/core_object/rich_text_object.py b/src/knowledge/core_object/rich_text_object.py index ac3fcbe..7791689 100644 --- a/src/knowledge/core_object/rich_text_object.py +++ b/src/knowledge/core_object/rich_text_object.py @@ -7,68 +7,68 @@ from .image_object import ImageObjectBuilder, ImageObject from .document_object import DocumentObjectBuilder, DocumentObject class RichTextObject(KnowledgeObject): - def __init__(self, images: dict = {}, videos: dict = {}, documents: dict = {}, rich_texts: dict = {}}): + def __init__(self, images: dict = {}, videos: dict = {}, documents: dict = {}, rich_texts: dict = {}): desc = dict() - desc.images = dict() - desc.videos = dict() - desc.documents = dict() - desc.rich_texts = dict() + desc["images"] = images + desc["videos"] = videos + desc["documents"] = documents + desc["rich_texts"] = rich_texts super().__init__(ObjectType.RichText, desc) def add_image_with_key(self, key, image_object: ImageObject): - assert self.desc.images[key] == None - self.desc.images[key] = image_object + assert self.desc["images"][key] == None + self.desc["images"][key] = image_object def add_image(self, image_object: ImageObject): - self.desc.images[image_object.object_id()] = image_object + self.desc["images"][image_object.object_id()] = image_object def get_image_with_key(self, key) -> ImageObject: - return self.desc.images[key] + return self.desc["images"][key] def get_images(self) -> dict: - return self.desc.images + return self.desc["images"] def add_video_with_key(self, key, video_object: VideoObject): - assert self.desc.videos[key] == None - self.desc.videos[key] = video_object + assert self.desc["videos"][key] == None + self.desc["videos"][key] = video_object def add_video(self, video_object: VideoObject): - self.desc.videos[video_object.object_id()] = video_object + self.desc["videos"][video_object.object_id()] = video_object def get_video_with_key(self, key) -> VideoObject: - return self.desc.videos[key] + return self.desc["videos"][key] def get_videos(self) -> dict: - return self.desc.videos + return self.desc["videos"] def add_document_with_key(self, key, document_object: DocumentObject): - assert self.desc.documents[key] == None - self.desc.documents[key] = document_object + assert self.desc["documents"][key] == None + self.desc["documents"][key] = document_object def add_document(self, document_object: DocumentObject): - self.desc.documents[document_object.object_id()] = document_object + self.desc["documents"][document_object.object_id()] = document_object def get_document_with_key(self, key) -> DocumentObject: - return self.desc.documents[key] + return self.desc["documents"][key] def get_documents(self) -> dict: - return self.desc.documents + return self.desc["documents"] def add_rich_text_with_key(self, key, rich_text_object): - assert self.desc.rich_texts[key] == None - self.desc.rich_texts[key] = rich_text_object + assert self.desc["rich_texts"][key] == None + self.desc["rich_texts"][key] = rich_text_object def add_rich_text(self, rich_text_object): - self.desc.rich_texts[rich_text_object.object_id()] = rich_text_object + self.desc["rich_texts"][rich_text_object.object_id()] = rich_text_object def get_rich_text_with_key(self, key): - return self.desc.rich_texts[key] + return self.desc["rich_texts"][key] def get_rich_texts(self) -> dict: - return self.desc.rich_texts + return self.desc["rich_texts"] class RichTextObjectBuilder: diff --git a/src/knowledge/core_object/video_object.py b/src/knowledge/core_object/video_object.py index ab4c262..56e9829 100644 --- a/src/knowledge/core_object/video_object.py +++ b/src/knowledge/core_object/video_object.py @@ -1,7 +1,7 @@ from ..object import KnowledgeObject from ..data import ChunkList, ChunkListWriter from ..object import ObjectType - +from .. import KnowledgeStore # desc # meta @@ -77,7 +77,7 @@ class VideoObjectBuilder: return self def build(self) -> VideoObject: - chunk_list = ChunkListWriter.create_chunk_list_from_file( + chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_file( self.video_file, 1024 * 1024 * 4, self.restore_file ) info = get_video_info(self.video_file) diff --git a/src/knowledge/data/tracker.py b/src/knowledge/data/tracker.py index 57daa38..66024c9 100644 --- a/src/knowledge/data/tracker.py +++ b/src/knowledge/data/tracker.py @@ -3,7 +3,7 @@ import time import logging import os from .chunk import ChunkID, PositionType, PositionFileRange -from typing import List +from typing import List, Tuple class ChunkTracker: def __init__(self, root_dir: str): @@ -61,7 +61,7 @@ class ChunkTracker: ) self.conn.commit() - def get_position(self, chunk_id: ChunkID) -> List[(str, PositionType)]: + def get_position(self, chunk_id: ChunkID) -> List[Tuple[str, PositionType]]: self.cursor.execute( """ SELECT pos, pos_type FROM chunks WHERE id = ? diff --git a/src/knowledge/data/writer.py b/src/knowledge/data/writer.py index 7f79447..7d38e9d 100644 --- a/src/knowledge/data/writer.py +++ b/src/knowledge/data/writer.py @@ -42,7 +42,7 @@ class ChunkListWriter: file_path, file.tell() - chunk_size, chunk_size ) self.chunk_tracker.add_position( - chunk_id, file_range, PositionType.FileRange + chunk_id, str(file_range), PositionType.FileRange ) file_hash = HashValue(hash_obj.digest()) diff --git a/src/knowledge/store.py b/src/knowledge/store.py index ef697fa..dc6b17a 100644 --- a/src/knowledge/store.py +++ b/src/knowledge/store.py @@ -1,6 +1,6 @@ import os from .object import ObjectStore -from .data import ChunkStore, ChunkTracker +from .data import ChunkStore, ChunkTracker, ChunkListWriter, ChunkReader import logging @@ -35,6 +35,9 @@ class KnowledgeStore: chunk_store_dir = os.path.join(root_dir, "chunk") self.chunk_store = ChunkStore(chunk_store_dir) self.chunk_tracker = ChunkTracker(chunk_store_dir) + self.chunk_list_writer = ChunkListWriter(self.chunk_store, self.chunk_tracker) + self.chunk_reader = ChunkReader(self.chunk_store, self.chunk_tracker) + def get_object_store(self) -> ObjectStore: return self.object_store @@ -44,3 +47,9 @@ class KnowledgeStore: def get_chunk_tracker(self) -> ChunkTracker: return self.chunk_tracker + + def get_chunk_list_writer(self) -> ChunkListWriter: + return self.chunk_list_writer + + def get_chunk_reader(self) -> ChunkReader: + return self.chunk_reader diff --git a/test/test_object.py b/test/test_object.py index 6a603fd..1e1bc56 100644 --- a/test/test_object.py +++ b/test/test_object.py @@ -1,5 +1,6 @@ import sys import os +import logging dir_path = os.path.dirname(os.path.realpath(__file__)) print(dir_path) @@ -7,9 +8,16 @@ print(dir_path) sys.path.append("{}/../src/".format(dir_path)) print(sys.path) -from knowledge import ObjectID, HashValue +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) +from knowledge import ObjectID, HashValue, EmailObjectBuilder import asyncio import unittest @@ -26,6 +34,8 @@ class TestVectorSTorage(unittest.TestCase): data2 = HashValue.from_base36(data.to_base36()) self.assertEqual(data.to_base58(), data2.to_base58()) - + email_folder = "F:\\system\\Downloads\\8081ffdb80925f5bff9c6ab9c4756c7d" + email_object = EmailObjectBuilder({}, email_folder).build() + if __name__ == "__main__": unittest.main() From 8ac6065332dc34de4da67191d6be1dae5ae15826 Mon Sep 17 00:00:00 2001 From: tsukasa Date: Tue, 12 Sep 2023 15:28:59 +0800 Subject: [PATCH 19/29] embedding email object --- src/aios_kernel/__init__.py | 3 +- src/aios_kernel/knowledge_base.py | 113 ++++++++++++++++------ src/knowledge/core_object/__init__.py | 1 + src/knowledge/core_object/image_object.py | 2 +- src/knowledge/object/object_id.py | 1 - src/knowledge/store.py | 7 ++ src/knowledge/vector/chroma_store.py | 6 +- src/knowledge/vector/vector_base.py | 3 +- test/test_knowledge_base.py | 57 +++++++++++ test/test_vector_storage.py | 4 +- 10 files changed, 155 insertions(+), 42 deletions(-) create mode 100644 test/test_knowledge_base.py diff --git a/src/aios_kernel/__init__.py b/src/aios_kernel/__init__.py index bd1434b..9817b76 100644 --- a/src/aios_kernel/__init__.py +++ b/src/aios_kernel/__init__.py @@ -4,7 +4,8 @@ from .chatsession import AIChatSession from .agent import AIAgent,AIAgentTemplete,AgentPrompt from .compute_kernel import ComputeKernel,ComputeTask from .compute_node import ComputeNode,LocalComputeNode -from .open_ai_node import OpenAI_ComputeNode +from .open_ai_node import OpenAI_ComputeNode +from .knowledge_base import KnowledgeBase from .role import AIRole,AIRoleGroup from .workflow import Workflow from .bus import AIBus \ No newline at end of file diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index 4f4ba4f..52b2943 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -1,47 +1,96 @@ # define a knowledge base class +import json from . import AgentPrompt, ComputeKernel -from ..knowledge.object import KnowledgeObject, ObjectType, EmailObject, TextChunkObject, ImageObject -from ..knowledge.store import ObjectStorage -from ..knowledge.vector.vector_base import VectorBase +from ..knowledge import * + class KnowledgeBase: - def __init__(self) -> None: - self.object_store = ObjectStorage() - self.vector_base = VectorBase() + _instance = None + + def __new__(cls): + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance.__singleton_init__() + + return cls._instance + + def __singleton_init__(self) -> None: + self.store = KnowledgeStore() self.compute_kernel = ComputeKernel() - async def insert(self, object: KnowledgeObject): - if object.object_type == ObjectType.Email: - email: EmailObject = object - for text_id in email.text: - [text, _] = self.object_store.get(text_id) - text: TextChunkObject = text - vector = await self.compute_kernel.do_text_embedding(text.text) - self.vector_base.insert(vector, text_id) - - for image_id in email.images: - [image, _] = self.object_store.get(image_id) - image: ImageObject = image - vector = await self.compute_kernel.do_text_embedding(image.meta) - self.vector_base.insert(vector, image_id) - - vector = await self.compute_kernel.do_text_embedding(email.meta) - self.vector_base.insert(vector, email.get_id()) + async def __embedding_document(self, document: DocumentObject): + for chunk_id in document.get_chunk_list(): + chunk = self.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 self.compute_kernel.do_text_embedding(text) + self.store.get_vector_store("default").insert(vector, chunk_id) + + async def __embedding_image(self, image: ImageObject): + desc = {} + if not image.get_meta(): + desc["meta"] = image.get_meta() + if not image.get_exif(): + desc["exif"] = image.get_exif() + if not image.get_tags(): + desc["tags"] = image.get_tags() + vector = await self.compute_kernel.do_text_embedding(json.dumps(desc)) + self.store.get_vector_store("default").insert(vector, image.calculate_id()) + + async def __embedding_vedio(self, vedio: VideoObject): + desc = {} + if not vedio.get_meta(): + desc["meta"] = vedio.get_meta() + if not vedio.get_info(): + desc["info"] = vedio.get_info() + if not vedio.get_tags(): + desc["tags"] = vedio.get_tags() + vector = await self.compute_kernel.do_text_embedding(json.dumps(desc)) + self.store.get_vector_store("default").insert(vector, vedio.calculate_id()) + + async def __embedding_rich_text(self, rich_text: RichTextObject): + for document in rich_text.get_documents().values(): + await self.__embedding_document(document) + for image in rich_text.get_images().values(): + await self.__embedding_image(image) + for vedio in rich_text.get_videos().values(): + await self.__embedding_vedio(vedio) + for rich_text in rich_text.get_rich_texts().values(): + 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())) + self.store.get_vector_store("default").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_vedio(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 query(self, prompt: AgentPrompt) -> AgentPrompt: + async def query(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 = self.vector_base.query(vector, 10) - for object_id in object_ids: - if object_id.object_type == ObjectType.Email: - [object, email] = self.object_store.get(object_id) - if object.object_type == ObjectType.Email: - email: EmailObject = object - prompt.append(AgentPrompt()) - prompt + object_ids = self.store.get_vector_store("default").query(vector, 10) + results.append(object_ids) + return results + + + + diff --git a/src/knowledge/core_object/__init__.py b/src/knowledge/core_object/__init__.py index 065de74..dee13e2 100644 --- a/src/knowledge/core_object/__init__.py +++ b/src/knowledge/core_object/__init__.py @@ -1,4 +1,5 @@ from .document_object import DocumentObject, DocumentObjectBuilder from .image_object import ImageObject, ImageObjectBuilder from .video_object import VideoObject, VideoObjectBuilder +from .rich_text_object import RichTextObject, RichTextObjectBuilder from .email_object import EmailObject, EmailObjectBuilder \ No newline at end of file diff --git a/src/knowledge/core_object/image_object.py b/src/knowledge/core_object/image_object.py index 83ab434..30b0e8a 100644 --- a/src/knowledge/core_object/image_object.py +++ b/src/knowledge/core_object/image_object.py @@ -18,7 +18,7 @@ class ImageObject(KnowledgeObject): body = dict() desc["meta"] = meta desc["exif"] = exif - desc + desc["tags"] = tags desc["hash"] = chunk_list.hash.to_base58() body["chunk_list"] = chunk_list.chunk_list diff --git a/src/knowledge/object/object_id.py b/src/knowledge/object/object_id.py index 3c3ac24..6ce77ae 100644 --- a/src/knowledge/object/object_id.py +++ b/src/knowledge/object/object_id.py @@ -8,7 +8,6 @@ import base36 class ObjectType(Enum): Chunk = 7 - TextChunk = 100 Image = 101 Video = 102 Document = 103 diff --git a/src/knowledge/store.py b/src/knowledge/store.py index dc6b17a..c0cc6b5 100644 --- a/src/knowledge/store.py +++ b/src/knowledge/store.py @@ -1,6 +1,7 @@ import os from .object import ObjectStore from .data import ChunkStore, ChunkTracker, ChunkListWriter, ChunkReader +from .vector import ChromaVectorStore, VectorBase import logging @@ -37,6 +38,7 @@ class KnowledgeStore: self.chunk_tracker = ChunkTracker(chunk_store_dir) self.chunk_list_writer = ChunkListWriter(self.chunk_store, self.chunk_tracker) self.chunk_reader = ChunkReader(self.chunk_store, self.chunk_tracker) + self.vector_store = {} def get_object_store(self) -> ObjectStore: @@ -53,3 +55,8 @@ class KnowledgeStore: def get_chunk_reader(self) -> ChunkReader: return self.chunk_reader + + 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) + return self.vector_store[model_name] diff --git a/src/knowledge/vector/chroma_store.py b/src/knowledge/vector/chroma_store.py index 626cd4f..0056375 100644 --- a/src/knowledge/vector/chroma_store.py +++ b/src/knowledge/vector/chroma_store.py @@ -6,11 +6,11 @@ import os class ChromaVectorStore(VectorBase): - def __init__(self, db_url, model_name: str) -> None: - super().__init__(db_url, model_name) + def __init__(self, model_name: str) -> None: + super().__init__(model_name) logging.info( - "will init chroma vector store, db={}, model={}".format(db_url, model_name) + "will init chroma vector store, model={}".format(model_name) ) directory = os.path.join( diff --git a/src/knowledge/vector/vector_base.py b/src/knowledge/vector/vector_base.py index 354c21e..83276ed 100644 --- a/src/knowledge/vector/vector_base.py +++ b/src/knowledge/vector/vector_base.py @@ -3,8 +3,7 @@ from ..object import ObjectID # define a vector base class class VectorBase: - def __init__(self, db_url, model_name) -> None: - self.db_url = db_url + def __init__(self, model_name) -> None: self.model_name = model_name async def insert(self, vector: [float], id: ObjectID): diff --git a/test/test_knowledge_base.py b/test/test_knowledge_base.py new file mode 100644 index 0000000..efed1bd --- /dev/null +++ b/test/test_knowledge_base.py @@ -0,0 +1,57 @@ +import sys +import os +import logging + +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) + + +from knowledge import ObjectID, HashValue, EmailObjectBuilder +from aios_kernel import KnowledgeBase, AgentPrompt +import asyncio +import unittest + +async def test_embedding_email(): + data = HashValue.hash_data("1233".encode("utf-8")); + print(data.to_base58()) + print(data.to_base36()) + + data2 = HashValue.from_base58(data.to_base58()) + self.assertEqual(data.to_base36(), data2.to_base36()) + + data2 = HashValue.from_base36(data.to_base36()) + self.assertEqual(data.to_base58(), data2.to_base58()) + + email_folder = "F:\\system\\Downloads\\8081ffdb80925f5bff9c6ab9c4756c7d" + email_object = EmailObjectBuilder({}, email_folder).build() + + await KnowledgeBase().do_embedding(email_object) + + +async def test_query_email(): + msg_prompt = AgentPrompt() + msg_prompt.messages = [{"role":"user","content":"abcdef"}] + + KnowledgeBase().query(msg_prompt) + +class TestVectorSTorage(unittest.TestCase): + def test_embedding(self): + asyncio.run(test_embedding_email()) + + def test_query(self): + asyncio.run(test_query_email()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/test_vector_storage.py b/test/test_vector_storage.py index 7357d65..6023ab1 100644 --- a/test/test_vector_storage.py +++ b/test/test_vector_storage.py @@ -14,8 +14,8 @@ import asyncio import unittest -async def test_vector(): - storage = ChromaVectorStore("", "test") +async def test_embedding_email(): + storage = ChromaVectorStore("test") await storage.insert([1, 2, 3], "test") ids = await storage.query([1, 2, 3], 10) print(ids) From 6e5beb4cdffb42c85eeb720ab0035b7c102ac58f Mon Sep 17 00:00:00 2001 From: tsukasa Date: Wed, 13 Sep 2023 17:55:27 +0800 Subject: [PATCH 20/29] save relationship of objects --- src/aios_kernel/knowledge_base.py | 107 +++++++++++++++++++++- src/knowledge/core_object/email_object.py | 2 +- src/knowledge/object/__init__.py | 1 + src/knowledge/object/relation.py | 16 ++++ src/knowledge/store.py | 10 +- test/test_knowledge_base.py | 2 +- test/test_vector_storage.py | 4 +- 7 files changed, 131 insertions(+), 11 deletions(-) create mode 100644 src/knowledge/object/relation.py diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index 52b2943..ce689e6 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -1,5 +1,6 @@ # define a knowledge base class import json +import pickle from . import AgentPrompt, ComputeKernel from ..knowledge import * @@ -39,7 +40,7 @@ class KnowledgeBase: vector = await self.compute_kernel.do_text_embedding(json.dumps(desc)) self.store.get_vector_store("default").insert(vector, image.calculate_id()) - async def __embedding_vedio(self, vedio: VideoObject): + async def __embedding_video(self, vedio: VideoObject): desc = {} if not vedio.get_meta(): desc["meta"] = vedio.get_meta() @@ -56,7 +57,7 @@ class KnowledgeBase: for image in rich_text.get_images().values(): await self.__embedding_image(image) for vedio in rich_text.get_videos().values(): - await self.__embedding_vedio(vedio) + await self.__embedding_video(vedio) for rich_text in rich_text.get_rich_texts().values(): await self.__embedding_rich_text(rich_text) @@ -65,13 +66,14 @@ class KnowledgeBase: self.store.get_vector_store("default").insert(vector, email.calculate_id()) await self.__embedding_rich_text(email.get_rich_text()) - async def do_embedding(self, object: KnowledgeObject): + + 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_vedio(object) + 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: @@ -79,7 +81,80 @@ class KnowledgeBase: else: pass - async def query(self, prompt: AgentPrompt) -> [ObjectID]: + def __save_document(self, document: DocumentObject): + doc_id = document.calculate_id() + self.store.get_object_store().put_object(doc_id, document.encode()) + for chunk_id in document.get_chunk_list(): + self.store.get_relation_store().add_relation(chunk_id, doc_id) + + def __save_image(self, image: ImageObject): + image_id = image.calculate_id() + self.store.get_object_store().put_object(image_id, image.encode()) + + def __save_video(self, video: VideoObject): + video_id = video.calculate_id() + self.store.get_object_store().put_object(video_id, video.encode()) + + def __save_rich_text(self, rich_text: RichTextObject): + rich_text_id = rich_text.calculate_id() + rich_text_enc = dict() + rich_text_enc["desc"] = rich_text.desc + rich_text_enc["body"] = {"documents": {}, "images": {}, "videos": {}, "rich_texts": {}} + for key, document in rich_text.get_documents().items(): + self.__save_document(document) + doc_id = document.calculate_id() + self.store.get_relation_store().add_relation(doc_id, rich_text_id) + rich_text_enc["body"]["documents"][key] = doc_id + for key, image in rich_text.get_images().items(): + self.__save_image(image) + image_id = image.calculate_id() + self.store.get_relation_store().add_relation(image_id, rich_text_id) + rich_text_enc["body"]["images"][key] = image_id + for key, video in rich_text.get_videos().items(): + self.__save_video(video) + video_id = video.calculate_id() + self.store.get_relation_store().add_relation(video_id, rich_text_id) + rich_text_enc["body"]["videos"][key] = video_id + for key, rich_text in rich_text.get_rich_texts().items(): + self.__save_rich_text(rich_text) + rich_text_id = rich_text.calculate_id() + self.store.get_relation_store().add_relation(rich_text_id, rich_text_id) + rich_text_enc["body"]["rich_texts"][key] = rich_text_id + + + self.store.get_object_store().put_object(rich_text_id, pickle.dumps(rich_text_enc)) + + def __save_email(self, email: EmailObject): + email_id = email.calculate_id() + email_enc = dict() + email_enc["desc"] = email.desc + email_enc["body"] = {"content": None} + self.__save_rich_text(email.get_rich_text()) + rich_text_id = email.get_rich_text().calculate_id() + self.store.get_relation_store().add_relation(rich_text_id, email_id) + email_enc["body"]["content"] = rich_text_id + self.store.get_object_store().put_object(email_id, pickle.dumps(email_enc)) + + + def __save_object(self, object: KnowledgeObject): + if object.get_object_type() == ObjectType.Document: + self.__save_document(object) + if object.get_object_type() == ObjectType.Image: + self.__save_image(object) + if object.get_object_type() == ObjectType.Video: + self.__save_video(object) + if object.get_object_type() == ObjectType.RichText: + self.__save_rich_text(object) + if object.get_object_type() == ObjectType.Email: + self.__save_email(object) + else: + pass + + async def insert_object(self, object: KnowledgeObject): + self.__save_object(object) + self.__do_embedding(object) + + async def query_objects(self, prompt: AgentPrompt) -> [ObjectID]: results = [] for msg in prompt.messages: if msg.role == "user": @@ -89,6 +164,28 @@ class KnowledgeBase: return results + async def __prompt_from_objects(self, object_ids: [ObjectID]) -> AgentPrompt: + prompt = AgentPrompt() + for object_id in object_ids: + object = self.store.get_object_reader().get_object(object_id) + if object is None: + raise ValueError(f"object not found: {object_id}") + if object.get_object_type() == ObjectType.Document: + document = object + prompt.messages.append({"role": "agent", "content": document.get_body()}) + if object.get_object_type() == ObjectType.Image: + image = object + prompt.messages.append({"role": "agent", "content": json.dumps(image.get_desc())}) + if object.get_object_type() == ObjectType.Video: + video = object + prompt.messages.append({"role": "agent", "content": json.dumps(video.get_desc())}) + if object.get_object_type() == ObjectType.RichText: + rich_text = object + prompt.messages.append({"role": "agent", "content": json.dumps(rich_text.get_desc())}) + if object.get_object_type() == ObjectType.Email: + email = object + prompt.messages.append({"role": "agent", "content": json.dumps(email.get_desc())}) + return prompt diff --git a/src/knowledge/core_object/email_object.py b/src/knowledge/core_object/email_object.py index 8b8efb5..f563052 100644 --- a/src/knowledge/core_object/email_object.py +++ b/src/knowledge/core_object/email_object.py @@ -28,7 +28,7 @@ class EmailObject(KnowledgeObject): return self.desc["tags"] def get_rich_text(self): - return self.body["rich_text"] + return self.body["content"] """ diff --git a/src/knowledge/object/__init__.py b/src/knowledge/object/__init__.py index 7886d7a..291715f 100644 --- a/src/knowledge/object/__init__.py +++ b/src/knowledge/object/__init__.py @@ -1,5 +1,6 @@ from .object import KnowledgeObject from .blob import FileBlobStorage from .hash import HashValue, hash_data +from .relation import ObjectRelationStore from .object_store import ObjectStore from .object_id import ObjectID, ObjectType diff --git a/src/knowledge/object/relation.py b/src/knowledge/object/relation.py new file mode 100644 index 0000000..1b56e48 --- /dev/null +++ b/src/knowledge/object/relation.py @@ -0,0 +1,16 @@ +# define a relation store class +from .object_id import ObjectID + + +class ObjectRelationStore: + def __init__(self, root_dir: str): + pass + + def add_relation(self, object: ObjectID, parent: ObjectID): + pass + + def get_related_objects(self, object: ObjectID) -> [ObjectID]: + pass + + def delete_relation(self, object_id: ObjectID) -> [ObjectID]: + pass \ No newline at end of file diff --git a/src/knowledge/store.py b/src/knowledge/store.py index c0cc6b5..03e48d0 100644 --- a/src/knowledge/store.py +++ b/src/knowledge/store.py @@ -1,5 +1,6 @@ import os -from .object import ObjectStore + +from .object import ObjectStore, ObjectRelationStore from .data import ChunkStore, ChunkTracker, ChunkListWriter, ChunkReader from .vector import ChromaVectorStore, VectorBase import logging @@ -30,6 +31,9 @@ class KnowledgeStore: self.root = root_dir + relation_store_dir = os.path.join(root_dir, "relation") + self.relation_store = ObjectRelationStore(relation_store_dir) + object_store_dir = os.path.join(root_dir, "object") self.object_store = ObjectStore(object_store_dir) @@ -39,7 +43,9 @@ class KnowledgeStore: self.chunk_list_writer = ChunkListWriter(self.chunk_store, self.chunk_tracker) self.chunk_reader = ChunkReader(self.chunk_store, self.chunk_tracker) self.vector_store = {} - + + def get_relation_store(self) -> ObjectRelationStore: + return self.relation_store def get_object_store(self) -> ObjectStore: return self.object_store diff --git a/test/test_knowledge_base.py b/test/test_knowledge_base.py index efed1bd..ebda00c 100644 --- a/test/test_knowledge_base.py +++ b/test/test_knowledge_base.py @@ -45,7 +45,7 @@ async def test_query_email(): KnowledgeBase().query(msg_prompt) -class TestVectorSTorage(unittest.TestCase): +class TestKnowledgeBase(unittest.TestCase): def test_embedding(self): asyncio.run(test_embedding_email()) diff --git a/test/test_vector_storage.py b/test/test_vector_storage.py index 6023ab1..ff4dc72 100644 --- a/test/test_vector_storage.py +++ b/test/test_vector_storage.py @@ -14,13 +14,13 @@ import asyncio import unittest -async def test_embedding_email(): +async def test_vector(): storage = ChromaVectorStore("test") await storage.insert([1, 2, 3], "test") ids = await storage.query([1, 2, 3], 10) print(ids) -class TestVectorSTorage(unittest.TestCase): +class TestVectorStorage(unittest.TestCase): def test_run(self): asyncio.run(test_vector()) From aaf0f152289196606eb05593bbaf8caffe5670c8 Mon Sep 17 00:00:00 2001 From: tsukasa Date: Thu, 14 Sep 2023 15:44:58 +0800 Subject: [PATCH 21/29] put email object with nested object not object id --- src/aios_kernel/knowledge_base.py | 43 +++++++++++++++++-------------- 1 file changed, 24 insertions(+), 19 deletions(-) diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index ce689e6..ebbfca8 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -31,22 +31,22 @@ class KnowledgeBase: async def __embedding_image(self, image: ImageObject): desc = {} - if not image.get_meta(): + if not not image.get_meta(): desc["meta"] = image.get_meta() - if not image.get_exif(): + if not not image.get_exif(): desc["exif"] = image.get_exif() - if not image.get_tags(): + if not not image.get_tags(): desc["tags"] = image.get_tags() vector = await self.compute_kernel.do_text_embedding(json.dumps(desc)) self.store.get_vector_store("default").insert(vector, image.calculate_id()) async def __embedding_video(self, vedio: VideoObject): desc = {} - if not vedio.get_meta(): + if not not vedio.get_meta(): desc["meta"] = vedio.get_meta() - if not vedio.get_info(): + if not not vedio.get_info(): desc["info"] = vedio.get_info() - if not vedio.get_tags(): + if not not vedio.get_tags(): desc["tags"] = vedio.get_tags() vector = await self.compute_kernel.do_text_embedding(json.dumps(desc)) self.store.get_vector_store("default").insert(vector, vedio.calculate_id()) @@ -97,43 +97,43 @@ class KnowledgeBase: def __save_rich_text(self, rich_text: RichTextObject): rich_text_id = rich_text.calculate_id() - rich_text_enc = dict() - rich_text_enc["desc"] = rich_text.desc - rich_text_enc["body"] = {"documents": {}, "images": {}, "videos": {}, "rich_texts": {}} + # rich_text_enc = dict() + # rich_text_enc["desc"] = rich_text.desc + # rich_text_enc["body"] = {"documents": {}, "images": {}, "videos": {}, "rich_texts": {}} for key, document in rich_text.get_documents().items(): self.__save_document(document) doc_id = document.calculate_id() self.store.get_relation_store().add_relation(doc_id, rich_text_id) - rich_text_enc["body"]["documents"][key] = doc_id + # rich_text_enc["body"]["documents"][key] = doc_id for key, image in rich_text.get_images().items(): self.__save_image(image) image_id = image.calculate_id() self.store.get_relation_store().add_relation(image_id, rich_text_id) - rich_text_enc["body"]["images"][key] = image_id + # rich_text_enc["body"]["images"][key] = image_id for key, video in rich_text.get_videos().items(): self.__save_video(video) video_id = video.calculate_id() self.store.get_relation_store().add_relation(video_id, rich_text_id) - rich_text_enc["body"]["videos"][key] = video_id + # rich_text_enc["body"]["videos"][key] = video_id for key, rich_text in rich_text.get_rich_texts().items(): self.__save_rich_text(rich_text) rich_text_id = rich_text.calculate_id() self.store.get_relation_store().add_relation(rich_text_id, rich_text_id) - rich_text_enc["body"]["rich_texts"][key] = rich_text_id + # rich_text_enc["body"]["rich_texts"][key] = rich_text_id - self.store.get_object_store().put_object(rich_text_id, pickle.dumps(rich_text_enc)) + self.store.get_object_store().put_object(rich_text_id, rich_text.encode()) def __save_email(self, email: EmailObject): email_id = email.calculate_id() - email_enc = dict() - email_enc["desc"] = email.desc - email_enc["body"] = {"content": None} + # email_enc = dict() + # email_enc["desc"] = email.desc + # email_enc["body"] = {"content": None} self.__save_rich_text(email.get_rich_text()) rich_text_id = email.get_rich_text().calculate_id() self.store.get_relation_store().add_relation(rich_text_id, email_id) - email_enc["body"]["content"] = rich_text_id - self.store.get_object_store().put_object(email_id, pickle.dumps(email_enc)) + # email_enc["body"]["content"] = rich_text_id + self.store.get_object_store().put_object(email_id, email.encode()) def __save_object(self, object: KnowledgeObject): @@ -163,6 +163,11 @@ class KnowledgeBase: results.append(object_ids) return results + async def __load_object(self, object_id: ObjectID) -> KnowledgeObject: + if object_id.get_object_type() == ObjectType.Document: + return DocumentObject.decode(self.store.get_object_store().get_object(object_id)) + object = self.store.get_object_store().get_object(object_id) + async def __prompt_from_objects(self, object_ids: [ObjectID]) -> AgentPrompt: prompt = AgentPrompt() From 9cbf29e0b0b379dcf53ac71362cf099ce9e5f873 Mon Sep 17 00:00:00 2001 From: tsukasa Date: Thu, 14 Sep 2023 16:03:56 +0800 Subject: [PATCH 22/29] implement object relation store with sqlite --- src/aios_kernel/knowledge_base.py | 1 - src/knowledge/object/relation.py | 56 ++++++++++++++++++++++++++++--- 2 files changed, 52 insertions(+), 5 deletions(-) diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index ebbfca8..18d78a3 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -1,6 +1,5 @@ # define a knowledge base class import json -import pickle from . import AgentPrompt, ComputeKernel from ..knowledge import * diff --git a/src/knowledge/object/relation.py b/src/knowledge/object/relation.py index 1b56e48..f6eca2e 100644 --- a/src/knowledge/object/relation.py +++ b/src/knowledge/object/relation.py @@ -1,16 +1,64 @@ # define a relation store class +import sqlite3 +import os +import logging from .object_id import ObjectID class ObjectRelationStore: def __init__(self, root_dir: str): - pass + if not os.path.exists(root_dir): + os.makedirs(root_dir) + file = os.path.join(root_dir, "relationships.db") + logging.info(f"will init chunk tracker, db={file}") + + self.conn = sqlite3.connect(file) + self.cursor = self.conn.cursor() + self.cursor.execute( + """ + CREATE TABLE IF NOT EXISTS relationships ( + id TEXT NOT NULL, + parent TEXT NOT NULL, + PRIMARY KEY(id, parent) + ) + """ + ) + self.conn.commit() def add_relation(self, object: ObjectID, parent: ObjectID): - pass + self.cursor.execute( + """ + INSERT INTO relationships (id, parent) + VALUES (?, ?, ?, ?, ?) + """, + ( + str(object), + str(parent), + ), + ) + self.conn.commit() def get_related_objects(self, object: ObjectID) -> [ObjectID]: - pass + parents = [] + cur = object + while True: + self.cursor.execute( + """ + SELECT id, parent FROM relationships WHERE id = ? + """, + (str(cur),), + ) + parent = self.cursor.fetchone() + if parent is None: + break + parents.append(parent) + cur = parent def delete_relation(self, object_id: ObjectID) -> [ObjectID]: - pass \ No newline at end of file + self.cursor.execute( + """ + DELETE FROM relationships WHERE id = ? + """, + (str(object_id),), + ) + self.conn.commit() \ No newline at end of file From a0cb5480df61bcba76466d8b39f868e1a527194e Mon Sep 17 00:00:00 2001 From: tsukasa Date: Thu, 14 Sep 2023 16:42:14 +0800 Subject: [PATCH 23/29] prompt from knowledge object --- src/aios_kernel/knowledge_base.py | 80 ++++++++++++++++++++++--------- 1 file changed, 58 insertions(+), 22 deletions(-) diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index 18d78a3..cabf005 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -165,32 +165,68 @@ class KnowledgeBase: async def __load_object(self, object_id: ObjectID) -> KnowledgeObject: if object_id.get_object_type() == ObjectType.Document: return DocumentObject.decode(self.store.get_object_store().get_object(object_id)) - object = self.store.get_object_store().get_object(object_id) + if object_id.get_object_type() == ObjectType.Image: + return ImageObject.decode(self.store.get_object_store().get_object(object_id)) + if object_id.get_object_type() == ObjectType.Video: + return VideoObject.decode(self.store.get_object_store().get_object(object_id)) + if object_id.get_object_type() == ObjectType.RichText: + return RichTextObject.decode(self.store.get_object_store().get_object(object_id)) + if object_id.get_object_type() == ObjectType.Email: + return EmailObject.decode(self.store.get_object_store().get_object(object_id)) + else: + pass - async def __prompt_from_objects(self, object_ids: [ObjectID]) -> AgentPrompt: - prompt = AgentPrompt() + async def prompt_from_objects(self, object_ids: [ObjectID]) -> AgentPrompt: + results = dict() for object_id in object_ids: - object = self.store.get_object_reader().get_object(object_id) - if object is None: - raise ValueError(f"object not found: {object_id}") - if object.get_object_type() == ObjectType.Document: - document = object - prompt.messages.append({"role": "agent", "content": document.get_body()}) - if object.get_object_type() == ObjectType.Image: - image = object - prompt.messages.append({"role": "agent", "content": json.dumps(image.get_desc())}) - if object.get_object_type() == ObjectType.Video: - video = object - prompt.messages.append({"role": "agent", "content": json.dumps(video.get_desc())}) - if object.get_object_type() == ObjectType.RichText: - rich_text = object - prompt.messages.append({"role": "agent", "content": json.dumps(rich_text.get_desc())}) - if object.get_object_type() == ObjectType.Email: - email = object - prompt.messages.append({"role": "agent", "content": json.dumps(email.get_desc())}) - return prompt + parents = self.store.get_relation_store().get_related_objects(object_id) + # last parent is the root object + root_object_id = parents[-1] + if results[root_object_id] is None: + results[root_object_id] = [root_object_id, object_id] + else: + results[root_object_id].append(object_id) + content = "I found the following contents described with json format:\n" + result_desc = [] + for result in results.values(): + # first element in result is the root object + root_object = await self.__load_object(result[0]) + if root_object.get_object_type() == ObjectType.Email: + email = await self.__load_object(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": self.store.get_chunk_reader().get_chunk(object_id).read().decode("utf-8")}) + if object_id.get_object_type() == ObjectType.Image: + image = await self.__load_object(object_id) + desc = image.get_desc() + desc["type"] = "image" + upper_list.append(desc) + if object_id.get_object_type() == ObjectType.Video: + video = await self.__load_object(object_id) + desc = video.get_desc() + desc["type"] = "video" + upper_list.append(desc) + else: + pass + content += json.dumps(result_desc) + content += ".\n" + + prompt = AgentPrompt() + prompt.add_message("knowledge", content) + + return prompt + From 937385f83cf080912099d5888e7646e41b3b217e Mon Sep 17 00:00:00 2001 From: liyaxing Date: Thu, 14 Sep 2023 20:09:20 +0800 Subject: [PATCH 24/29] Improve ObjectID with remove object_type and __eq__ method --- src/knowledge/object/object_id.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/knowledge/object/object_id.py b/src/knowledge/object/object_id.py index 6ce77ae..8278183 100644 --- a/src/knowledge/object/object_id.py +++ b/src/knowledge/object/object_id.py @@ -17,16 +17,12 @@ class ObjectType(Enum): # define a object ID class to identify a object class ObjectID: # pylint: disable=too-few-public-methods - def __init__(self, object_type: ObjectType, value: bytes): + def __init__(self, value: bytes): assert len(value) == 32, "ObjectID must be 32 bytes long" - self.object_type = object_type self.value = value def __str__(self): return self.to_base58() - - def get_object_type(self): - return self.object_type def to_base58(self): return base58.b58encode(self.value).decode() @@ -48,8 +44,11 @@ class ObjectID: # pylint: disable=too-few-public-methods @staticmethod def new_chunk_id(chunk_hash: HashValue): - return ObjectID(ObjectType.Chunk, chunk_hash.value) + return ObjectID(chunk_hash.value) @staticmethod def hash_data(data: bytes): - return ObjectID.new_chunk_id(HashValue.hash_data(data)) \ No newline at end of file + return ObjectID.new_chunk_id(HashValue.hash_data(data)) + + def __eq__(self, other) -> bool: + return self.value == other.value \ No newline at end of file From e1bba92ec2c708ed0108448dfaf94f8bae8f7326 Mon Sep 17 00:00:00 2001 From: liyaxing Date: Thu, 14 Sep 2023 20:10:07 +0800 Subject: [PATCH 25/29] Add object relation store impl and test case --- src/knowledge/object/relation.py | 102 +++++++++++++++++++++---------- test/test_object.py | 41 +++++++++++-- 2 files changed, 107 insertions(+), 36 deletions(-) diff --git a/src/knowledge/object/relation.py b/src/knowledge/object/relation.py index f6eca2e..34c0911 100644 --- a/src/knowledge/object/relation.py +++ b/src/knowledge/object/relation.py @@ -1,64 +1,104 @@ # define a relation store class -import sqlite3 -import os -import logging from .object_id import ObjectID +import sqlite3 +from typing import List, Tuple, Optional +import logging +import os +from enum import Enum + + +class ObjectRelationType(Enum): + Parent = 1 class ObjectRelationStore: def __init__(self, root_dir: str): if not os.path.exists(root_dir): os.makedirs(root_dir) - file = os.path.join(root_dir, "relationships.db") - logging.info(f"will init chunk tracker, db={file}") + file = os.path.join(root_dir, "relation.db") + logging.info(f"will init object relation store, db={file}") self.conn = sqlite3.connect(file) self.cursor = self.conn.cursor() self.cursor.execute( """ - CREATE TABLE IF NOT EXISTS relationships ( - id TEXT NOT NULL, - parent TEXT NOT NULL, - PRIMARY KEY(id, parent) + CREATE TABLE IF NOT EXISTS relations ( + object_id TEXT, + assoc_id TEXT, + relation_type TEXT, + PRIMARY KEY (object_id, assoc_id, relation_type) ) """ ) - self.conn.commit() - def add_relation(self, object: ObjectID, parent: ObjectID): + def add_relation( + self, + object_id: ObjectID, + assoc_id: ObjectID, + relation_type: ObjectRelationType = ObjectRelationType.Parent, + ): + if relation_type == None: + relation_type = ObjectRelationType.Parent + self.cursor.execute( """ - INSERT INTO relationships (id, parent) - VALUES (?, ?, ?, ?, ?) + INSERT OR IGNORE INTO relations (object_id, assoc_id, relation_type) + VALUES (?, ?, ?) """, - ( - str(object), - str(parent), - ), + (str(object_id), str(assoc_id), relation_type.value), ) self.conn.commit() - def get_related_objects(self, object: ObjectID) -> [ObjectID]: - parents = [] - cur = object - while True: + def get_related_objects( + self, object_id: ObjectID, relation_type: Optional[ObjectRelationType] = None + ) -> List[ObjectID]: + if relation_type: self.cursor.execute( """ - SELECT id, parent FROM relationships WHERE id = ? + SELECT assoc_id FROM relations WHERE object_id = ? AND relation_type = ? """, - (str(cur),), + (str(object_id), relation_type.value), ) - parent = self.cursor.fetchone() - if parent is None: - break - parents.append(parent) - cur = parent + else: + self.cursor.execute( + """ + SELECT assoc_id FROM relations WHERE object_id = ? + """, + (str(object_id),), + ) + return [ObjectID.from_base58(row[0]) for row in self.cursor.fetchall()] - def delete_relation(self, object_id: ObjectID) -> [ObjectID]: + def get_related_root_objects( + self, object_id: ObjectID, relation_type: Optional[ObjectRelationType] = None + ) -> List[ObjectID]: + root_objects = [] + related_objects = self.get_related_objects(object_id, relation_type) + history = [] + history.append(object_id) + + while related_objects: + for obj in related_objects: + next_related_objects = self.get_related_objects(obj, relation_type) + if not next_related_objects: + if obj not in root_objects: + root_objects.append(obj) + else: + for related_object in next_related_objects: + if obj not in history: + related_objects.append(related_object) + else: + logging.warning( + f"loop detected: {obj} <-> {related_object}" + ) + related_objects = next_related_objects + + return root_objects + + def delete_relation(self, object_id: ObjectID): self.cursor.execute( """ - DELETE FROM relationships WHERE id = ? + DELETE FROM relations WHERE object_id = ? """, (str(object_id),), ) - self.conn.commit() \ No newline at end of file + self.conn.commit() diff --git a/test/test_object.py b/test/test_object.py index 1e1bc56..0f5620e 100644 --- a/test/test_object.py +++ b/test/test_object.py @@ -17,25 +17,56 @@ handler.setFormatter(formatter) root.addHandler(handler) -from knowledge import ObjectID, HashValue, EmailObjectBuilder +from knowledge import ( + ObjectID, + HashValue, + EmailObjectBuilder, + ObjectRelationStore, + KnowledgeStore, +) import asyncio import unittest class TestVectorSTorage(unittest.TestCase): def test_object(self): - data = HashValue.hash_data("1233".encode("utf-8")); + data = HashValue.hash_data("1233".encode("utf-8")) print(data.to_base58()) print(data.to_base36()) - + data2 = HashValue.from_base58(data.to_base58()) self.assertEqual(data.to_base36(), data2.to_base36()) - + data2 = HashValue.from_base36(data.to_base36()) self.assertEqual(data.to_base58(), data2.to_base58()) email_folder = "F:\\system\\Downloads\\8081ffdb80925f5bff9c6ab9c4756c7d" email_object = EmailObjectBuilder({}, email_folder).build() - + + def test_relation(self): + obj1 = ObjectID.hash_data("12345".encode("utf-8")) + obj2 = ObjectID.hash_data("67890".encode("utf-8")) + obj3 = ObjectID.hash_data("abcde".encode("utf-8")) + obj4 = ObjectID.hash_data("fghij".encode("utf-8")) + print(obj1.to_base58(), obj2.to_base58(), obj3.to_base58()) + relation_store = KnowledgeStore().get_relation_store() + relation_store.add_relation(obj1, obj2) + relation_store.add_relation(obj1, obj2) + relation_store.add_relation(obj2, obj3) + + relation_store.add_relation(obj1, obj3) + relation_store.add_relation(obj1, obj4) + + objs = relation_store.get_related_objects(obj2) + self.assertEqual(len(objs), 1) + self.assertEqual(objs[0], obj3) + + objs = relation_store.get_related_root_objects(obj1) + self.assertEqual(len(objs), 2) + self.assertEqual(obj3 in objs, True) + self.assertEqual(obj4 in objs, True) + # self.assertCountEqual(objs, [obj3, obj4]) + + if __name__ == "__main__": unittest.main() From 6ea83729a570711a36682d1b0a86976206963f2f Mon Sep 17 00:00:00 2001 From: liyaxing Date: Thu, 14 Sep 2023 20:57:53 +0800 Subject: [PATCH 26/29] Refactoring the EmailObject build process --- src/knowledge/core_object/document_object.py | 13 +++++-- src/knowledge/core_object/email_object.py | 41 ++++++++++++++++---- src/knowledge/data/writer.py | 2 +- 3 files changed, 45 insertions(+), 11 deletions(-) diff --git a/src/knowledge/core_object/document_object.py b/src/knowledge/core_object/document_object.py index 6d6270e..02056bc 100644 --- a/src/knowledge/core_object/document_object.py +++ b/src/knowledge/core_object/document_object.py @@ -1,4 +1,4 @@ -from ..object import KnowledgeObject +from ..object import KnowledgeObject, ObjectRelationStore from ..data import ChunkList, ChunkListWriter from ..object import ObjectType from .. import KnowledgeStore @@ -49,9 +49,16 @@ class DocumentObjectBuilder: self.text = text return self - def build(self) -> DocumentObject: + def build(self, relation_store: ObjectRelationStore) -> DocumentObject: chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_text( self.text, 1024 * 4, ) - return DocumentObject(self.meta, self.tags, chunk_list) + doc = DocumentObject(self.meta, self.tags, chunk_list) + doc_id = doc.calculate_id() + + # Add relation to store + for chunk_id in chunk_list.chunk_list: + relation_store.add_relation(chunk_id, doc_id) + + return doc diff --git a/src/knowledge/core_object/email_object.py b/src/knowledge/core_object/email_object.py index f563052..c86c3a2 100644 --- a/src/knowledge/core_object/email_object.py +++ b/src/knowledge/core_object/email_object.py @@ -1,9 +1,9 @@ +from .. import KnowledgeStore from .rich_text_object import RichTextObject, RichTextObjectBuilder from ..object import ObjectID, ObjectType, KnowledgeObject from .document_object import DocumentObjectBuilder from .image_object import ImageObjectBuilder from .video_object import VideoObjectBuilder -from .rich_text_object import RichTextObjectBuilder import os import json import logging @@ -69,6 +69,11 @@ class EmailObjectBuilder: return self def build(self) -> EmailObject: + + # Just get the object store and relation store from global KnowledgeStore + store = KnowledgeStore().get_object_store() + relation = KnowledgeStore().get_relation_store() + # Read meta.json meta = {} meta_file = os.path.join(self.folder, "meta.json") @@ -89,8 +94,10 @@ class EmailObjectBuilder: with open(content_file, "r", encoding="utf-8") as f: text = f.read() - document = DocumentObjectBuilder({}, {}, text).build() - documents = {"email.txt": document} + document = DocumentObjectBuilder({}, {}, text).build(relation_store=relation) + document_id = document.calculate_id() + store.put_object(document_id, document.encode()) + documents = {"email.txt": document_id} except Exception as e: logging.error(f"Failed to read email.txt {content_file} {e}") else: @@ -106,7 +113,9 @@ class EmailObjectBuilder: try: image = ImageObjectBuilder({}, {}, image_path).build() - images[image_file] = image + image_id = image.calculate_id() + store.put_object(image_id, image.encode()) + images[image_file] = image_id except Exception as e: logging.error(f"Failed to read image file {image_path} {e}") continue @@ -121,13 +130,31 @@ class EmailObjectBuilder: try: video = VideoObjectBuilder({}, {}, video_path).build() - videos[video_file] = video + video_id = video.calculate_id() + store.put_object(video_id, video.encode()) + videos[video_file] = video_id except Exception as e: logging.error(f"Failed to read video file {video_path} {e}") continue # Create RichTextObject rich_text = RichTextObject(images, videos, documents) - + rich_text_id = rich_text.calculate_id() + + # build relations with rich_text + for image_id in images.values(): + relation.add_relation(image_id, rich_text_id) + for video_id in videos.values(): + relation.add_relation(video_id, rich_text_id) + for document_id in documents.values(): + relation.add_relation(document_id, rich_text_id) + # Create EmailObject - return EmailObject(meta, {}, rich_text) + email_object = EmailObject(meta, {}, rich_text) + email_object_id = email_object.calculate_id() + store.put_object(email_object_id, email_object.encode()) + + # build relations with email_object + relation.add_relation(rich_text_id, email_object_id) + + return email_object diff --git a/src/knowledge/data/writer.py b/src/knowledge/data/writer.py index 7d38e9d..599c51b 100644 --- a/src/knowledge/data/writer.py +++ b/src/knowledge/data/writer.py @@ -52,7 +52,7 @@ class ChunkListWriter: def create_chunk_list_from_text( self, text: str, chunk_max_words: int, separator_chars: str = ".," - ) -> Tuple[List[ChunkID], HashValue]: + ) -> ChunkList: text_list = self._split_text_list(text, chunk_max_words, separator_chars) chunk_list = [] hash_obj = hashlib.sha256() From b59af5d4a7cc0acf61793fe1d7de5fdefd75229b Mon Sep 17 00:00:00 2001 From: liyaxing Date: Fri, 15 Sep 2023 14:01:26 +0800 Subject: [PATCH 27/29] Fix object encode --- src/knowledge/data/chunk.py | 4 ++-- src/knowledge/object/object.py | 24 +++++++++++++++++------- src/knowledge/object/object_id.py | 4 ++-- src/knowledge/object/relation.py | 4 ++-- test/test_object.py | 15 +++++++++++++++ 5 files changed, 38 insertions(+), 13 deletions(-) diff --git a/src/knowledge/data/chunk.py b/src/knowledge/data/chunk.py index 5267dc3..76f2058 100644 --- a/src/knowledge/data/chunk.py +++ b/src/knowledge/data/chunk.py @@ -1,9 +1,9 @@ -from enum import Enum +from enum import IntEnum from ..object import ObjectID ChunkID = ObjectID -class PositionType(Enum): +class PositionType(IntEnum): Unknown = 1 Device = 2 File = 3 diff --git a/src/knowledge/object/object.py b/src/knowledge/object/object.py index b3841ed..490555b 100644 --- a/src/knowledge/object/object.py +++ b/src/knowledge/object/object.py @@ -1,4 +1,3 @@ - # define a object type enum from abc import ABC, abstractmethod from enum import Enum @@ -6,6 +5,16 @@ from .object_id import ObjectID, ObjectType import hashlib import json import pickle +from typing import Any + + +class ObjectEnhancedJSONEncoder(json.JSONEncoder): + def default(self, o: Any) -> Any: + if isinstance(o, ObjectID): + return o.to_base58() + + return super().default(o) + class KnowledgeObject(ABC): def __init__(self, object_type: ObjectType, desc: dict = {}, body: dict = {}): @@ -18,7 +27,7 @@ class KnowledgeObject(ABC): def object_id(self) -> ObjectID: return self.calculate_id() - + def set_desc_with_key_value(self, key, value): self.desc[key] = value @@ -27,7 +36,7 @@ class KnowledgeObject(ABC): def get_desc(self) -> dict: return self.desc - + def set_body_with_key_value(self, key, value): self.body[key] = value @@ -36,10 +45,13 @@ class KnowledgeObject(ABC): def get_body(self) -> dict: return self.body - + def calculate_id(self): # Convert the object_type and desc to string and compute the SHA256 hash - data = json.dumps({"object_type": self.object_type, "desc": self.desc}) + data = json.dumps( + {"object_type": self.object_type, "desc": self.desc}, + cls=ObjectEnhancedJSONEncoder, + ) sha256 = hashlib.sha256() sha256.update(data.encode()) return ObjectID(sha256.digest()) @@ -50,5 +62,3 @@ class KnowledgeObject(ABC): @staticmethod def decode(data: bytes): return pickle.loads(data) - - diff --git a/src/knowledge/object/object_id.py b/src/knowledge/object/object_id.py index 8278183..42f313a 100644 --- a/src/knowledge/object/object_id.py +++ b/src/knowledge/object/object_id.py @@ -1,12 +1,12 @@ # define a object type enum from abc import ABC, abstractmethod -from enum import Enum +from enum import IntEnum from .hash import HashValue import base58 import base36 -class ObjectType(Enum): +class ObjectType(IntEnum): Chunk = 7 Image = 101 Video = 102 diff --git a/src/knowledge/object/relation.py b/src/knowledge/object/relation.py index 34c0911..f273279 100644 --- a/src/knowledge/object/relation.py +++ b/src/knowledge/object/relation.py @@ -4,10 +4,10 @@ import sqlite3 from typing import List, Tuple, Optional import logging import os -from enum import Enum +from enum import IntEnum -class ObjectRelationType(Enum): +class ObjectRelationType(IntEnum): Parent = 1 diff --git a/test/test_object.py b/test/test_object.py index 0f5620e..655e2a2 100644 --- a/test/test_object.py +++ b/test/test_object.py @@ -23,6 +23,7 @@ from knowledge import ( EmailObjectBuilder, ObjectRelationStore, KnowledgeStore, + EmailObject, ) import asyncio import unittest @@ -42,6 +43,20 @@ class TestVectorSTorage(unittest.TestCase): email_folder = "F:\\system\\Downloads\\8081ffdb80925f5bff9c6ab9c4756c7d" email_object = EmailObjectBuilder({}, email_folder).build() + + id = email_object.calculate_id() + print(f"got email object: {id.to_base58()}") + + # test encode & decode + ret = email_object.encode() + obj = EmailObject.decode(ret) + id2 = obj.calculate_id() + print(f"got email object: {id2.to_base58()}") + self.assertEqual(id.to_base58(), id2.to_base58()) + + ret2 = obj.encode() + self.assertEqual(ret, ret2) + def test_relation(self): obj1 = ObjectID.hash_data("12345".encode("utf-8")) From 805dc88de28d841ac1472803100e36a11ae2ae9c Mon Sep 17 00:00:00 2001 From: tsukasa Date: Mon, 18 Sep 2023 09:51:51 +0800 Subject: [PATCH 28/29] test prompt from knowledge object --- src/aios_kernel/compute_kernel.py | 2 +- src/aios_kernel/compute_task.py | 2 +- src/aios_kernel/knowledge_base.py | 151 ++++++++++--------- src/aios_kernel/open_ai_node.py | 2 +- src/knowledge/core_object/document_object.py | 1 + test/test_knowledge_base.py | 39 ++--- 6 files changed, 98 insertions(+), 99 deletions(-) diff --git a/src/aios_kernel/compute_kernel.py b/src/aios_kernel/compute_kernel.py index f626578..b1da61b 100644 --- a/src/aios_kernel/compute_kernel.py +++ b/src/aios_kernel/compute_kernel.py @@ -121,7 +121,7 @@ class ComputeKernel: def text_embedding(self,input:str,model_name:Optional[str] = None): task_req = ComputeTask() - task_req.set_text_embeding_params(input,model_name) + task_req.set_text_embedding_params(input,model_name) self.run(task_req) return task_req diff --git a/src/aios_kernel/compute_task.py b/src/aios_kernel/compute_task.py index 35a1afd..6916a2a 100644 --- a/src/aios_kernel/compute_task.py +++ b/src/aios_kernel/compute_task.py @@ -38,7 +38,7 @@ class ComputeTask: self.params["model_name"] = "gpt-4-0613" self.params["max_token_size"] = max_token_size - def set_text_embeding_params(self, input, model_name=None, callchain_id = None): + def set_text_embedding_params(self, input, model_name=None, callchain_id = None): self.task_type = "text_embedding" self.create_time = time.time() self.task_id = uuid.uuid4().hex diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index cabf005..dbb96e6 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -1,7 +1,7 @@ # define a knowledge base class import json from . import AgentPrompt, ComputeKernel -from ..knowledge import * +from knowledge import * class KnowledgeBase: @@ -62,7 +62,7 @@ class KnowledgeBase: async def __embedding_email(self, email: EmailObject): vector = await self.compute_kernel.do_text_embedding(json.dumps(email.get_desc())) - self.store.get_vector_store("default").insert(vector, email.calculate_id()) + await self.store.get_vector_store("default").insert(vector, email.calculate_id()) await self.__embedding_rich_text(email.get_rich_text()) @@ -80,89 +80,94 @@ class KnowledgeBase: else: pass - def __save_document(self, document: DocumentObject): - doc_id = document.calculate_id() - self.store.get_object_store().put_object(doc_id, document.encode()) - for chunk_id in document.get_chunk_list(): - self.store.get_relation_store().add_relation(chunk_id, doc_id) + # def __save_document(self, document: DocumentObject): + # doc_id = document.calculate_id() + # self.store.get_object_store().put_object(doc_id, document.encode()) + # for chunk_id in document.get_chunk_list(): + # self.store.get_relation_store().add_relation(chunk_id, doc_id) - def __save_image(self, image: ImageObject): - image_id = image.calculate_id() - self.store.get_object_store().put_object(image_id, image.encode()) + # def __save_image(self, image: ImageObject): + # image_id = image.calculate_id() + # self.store.get_object_store().put_object(image_id, image.encode()) - def __save_video(self, video: VideoObject): - video_id = video.calculate_id() - self.store.get_object_store().put_object(video_id, video.encode()) + # def __save_video(self, video: VideoObject): + # video_id = video.calculate_id() + # self.store.get_object_store().put_object(video_id, video.encode()) - def __save_rich_text(self, rich_text: RichTextObject): - rich_text_id = rich_text.calculate_id() - # rich_text_enc = dict() - # rich_text_enc["desc"] = rich_text.desc - # rich_text_enc["body"] = {"documents": {}, "images": {}, "videos": {}, "rich_texts": {}} - for key, document in rich_text.get_documents().items(): - self.__save_document(document) - doc_id = document.calculate_id() - self.store.get_relation_store().add_relation(doc_id, rich_text_id) - # rich_text_enc["body"]["documents"][key] = doc_id - for key, image in rich_text.get_images().items(): - self.__save_image(image) - image_id = image.calculate_id() - self.store.get_relation_store().add_relation(image_id, rich_text_id) - # rich_text_enc["body"]["images"][key] = image_id - for key, video in rich_text.get_videos().items(): - self.__save_video(video) - video_id = video.calculate_id() - self.store.get_relation_store().add_relation(video_id, rich_text_id) - # rich_text_enc["body"]["videos"][key] = video_id - for key, rich_text in rich_text.get_rich_texts().items(): - self.__save_rich_text(rich_text) - rich_text_id = rich_text.calculate_id() - self.store.get_relation_store().add_relation(rich_text_id, rich_text_id) - # rich_text_enc["body"]["rich_texts"][key] = rich_text_id + # def __save_rich_text(self, rich_text: RichTextObject): + # rich_text_id = rich_text.calculate_id() + # # rich_text_enc = dict() + # # rich_text_enc["desc"] = rich_text.desc + # # rich_text_enc["body"] = {"documents": {}, "images": {}, "videos": {}, "rich_texts": {}} + # for key, document in rich_text.get_documents().items(): + # self.__save_document(document) + # doc_id = document.calculate_id() + # self.store.get_relation_store().add_relation(doc_id, rich_text_id) + # # rich_text_enc["body"]["documents"][key] = doc_id + # for key, image in rich_text.get_images().items(): + # self.__save_image(image) + # image_id = image.calculate_id() + # self.store.get_relation_store().add_relation(image_id, rich_text_id) + # # rich_text_enc["body"]["images"][key] = image_id + # for key, video in rich_text.get_videos().items(): + # self.__save_video(video) + # video_id = video.calculate_id() + # self.store.get_relation_store().add_relation(video_id, rich_text_id) + # # rich_text_enc["body"]["videos"][key] = video_id + # for key, rich_text in rich_text.get_rich_texts().items(): + # self.__save_rich_text(rich_text) + # rich_text_id = rich_text.calculate_id() + # self.store.get_relation_store().add_relation(rich_text_id, rich_text_id) + # # rich_text_enc["body"]["rich_texts"][key] = rich_text_id - self.store.get_object_store().put_object(rich_text_id, rich_text.encode()) + # self.store.get_object_store().put_object(rich_text_id, rich_text.encode()) - def __save_email(self, email: EmailObject): - email_id = email.calculate_id() - # email_enc = dict() - # email_enc["desc"] = email.desc - # email_enc["body"] = {"content": None} - self.__save_rich_text(email.get_rich_text()) - rich_text_id = email.get_rich_text().calculate_id() - self.store.get_relation_store().add_relation(rich_text_id, email_id) - # email_enc["body"]["content"] = rich_text_id - self.store.get_object_store().put_object(email_id, email.encode()) + # def __save_email(self, email: EmailObject): + # email_id = email.calculate_id() + # # email_enc = dict() + # # email_enc["desc"] = email.desc + # # email_enc["body"] = {"content": None} + # self.__save_rich_text(email.get_rich_text()) + # rich_text_id = email.get_rich_text().calculate_id() + # self.store.get_relation_store().add_relation(rich_text_id, email_id) + # # email_enc["body"]["content"] = rich_text_id + # self.store.get_object_store().put_object(email_id, email.encode()) - def __save_object(self, object: KnowledgeObject): - if object.get_object_type() == ObjectType.Document: - self.__save_document(object) - if object.get_object_type() == ObjectType.Image: - self.__save_image(object) - if object.get_object_type() == ObjectType.Video: - self.__save_video(object) - if object.get_object_type() == ObjectType.RichText: - self.__save_rich_text(object) - if object.get_object_type() == ObjectType.Email: - self.__save_email(object) - else: - pass + # def __save_object(self, object: KnowledgeObject): + # if object.get_object_type() == ObjectType.Document: + # self.__save_document(object) + # if object.get_object_type() == ObjectType.Image: + # self.__save_image(object) + # if object.get_object_type() == ObjectType.Video: + # self.__save_video(object) + # if object.get_object_type() == ObjectType.RichText: + # self.__save_rich_text(object) + # if object.get_object_type() == ObjectType.Email: + # self.__save_email(object) + # else: + # pass async def insert_object(self, object: KnowledgeObject): - self.__save_object(object) - self.__do_embedding(object) + # self.__save_object(object) + await self.__do_embedding(object) + + async def query_prompt(self, prompt: AgentPrompt): + objects = await self.query_objects(prompt) + knowledge_prompt = self.prompt_from_objects(objects) + prompt.append(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 = self.store.get_vector_store("default").query(vector, 10) + object_ids = await self.store.get_vector_store("default").query(vector, 10) results.append(object_ids) return results - async def __load_object(self, object_id: ObjectID) -> KnowledgeObject: + def __load_object(self, object_id: ObjectID) -> KnowledgeObject: if object_id.get_object_type() == ObjectType.Document: return DocumentObject.decode(self.store.get_object_store().get_object(object_id)) if object_id.get_object_type() == ObjectType.Image: @@ -177,10 +182,10 @@ class KnowledgeBase: pass - async def prompt_from_objects(self, object_ids: [ObjectID]) -> AgentPrompt: + def prompt_from_objects(self, object_ids: [ObjectID]) -> AgentPrompt: results = dict() for object_id in object_ids: - parents = self.store.get_relation_store().get_related_objects(object_id) + parents = self.store.get_relation_store().get_related_root_objects(object_id) # last parent is the root object root_object_id = parents[-1] if results[root_object_id] is None: @@ -192,9 +197,9 @@ class KnowledgeBase: result_desc = [] for result in results.values(): # first element in result is the root object - root_object = await self.__load_object(result[0]) - if root_object.get_object_type() == ObjectType.Email: - email = await self.__load_object(object_id) + root_object_id = result[0] + if root_object_id.get_object_type() == ObjectType.Email: + email = self.__load_object(root_object_id) desc = email.get_desc() desc["type"] = "email" desc["contents"] = [] @@ -208,12 +213,12 @@ 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 = await self.__load_object(object_id) + image = self.__load_object(object_id) desc = image.get_desc() desc["type"] = "image" upper_list.append(desc) if object_id.get_object_type() == ObjectType.Video: - video = await self.__load_object(object_id) + video = self.__load_object(object_id) desc = video.get_desc() desc["type"] = "video" upper_list.append(desc) diff --git a/src/aios_kernel/open_ai_node.py b/src/aios_kernel/open_ai_node.py index 27177f0..a37706e 100644 --- a/src/aios_kernel/open_ai_node.py +++ b/src/aios_kernel/open_ai_node.py @@ -82,7 +82,7 @@ class OpenAI_ComputeNode(ComputeNode): input = task.params["input"] logger.info(f"call openai {model_name} input: {input}") - resp = openai.Embeding.create(model=model_name, + resp = openai.Embedding.create(model=model_name, input=input) logger.info(f"openai response: {resp}") diff --git a/src/knowledge/core_object/document_object.py b/src/knowledge/core_object/document_object.py index 02056bc..ab02b8d 100644 --- a/src/knowledge/core_object/document_object.py +++ b/src/knowledge/core_object/document_object.py @@ -53,6 +53,7 @@ class DocumentObjectBuilder: chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_text( self.text, 1024 * 4, + "." ) doc = DocumentObject(self.meta, self.tags, chunk_list) doc_id = doc.calculate_id() diff --git a/test/test_knowledge_base.py b/test/test_knowledge_base.py index ebda00c..5c9d516 100644 --- a/test/test_knowledge_base.py +++ b/test/test_knowledge_base.py @@ -18,39 +18,32 @@ root.addHandler(handler) from knowledge import ObjectID, HashValue, EmailObjectBuilder -from aios_kernel import KnowledgeBase, AgentPrompt +from aios_kernel import KnowledgeBase, AgentPrompt, OpenAI_ComputeNode, ComputeKernel import asyncio import unittest -async def test_embedding_email(): - data = HashValue.hash_data("1233".encode("utf-8")); - print(data.to_base58()) - print(data.to_base36()) - - data2 = HashValue.from_base58(data.to_base58()) - self.assertEqual(data.to_base36(), data2.to_base36()) - - data2 = HashValue.from_base36(data.to_base36()) - self.assertEqual(data.to_base58(), data2.to_base58()) +async def test_embedding_email(test): + open_ai_node = OpenAI_ComputeNode() + open_ai_node.start() + ComputeKernel().add_compute_node(open_ai_node) + + email_folder = os.path.join(dir_path, "../rootfs/data/email/") + print("explore emails in folder ", email_folder) + for root, dirs, files in os.walk(email_folder): + for dir in dirs: + email_object = EmailObjectBuilder({}, os.path.join(root, dir)).build() + await KnowledgeBase().insert_object(email_object) - email_folder = "F:\\system\\Downloads\\8081ffdb80925f5bff9c6ab9c4756c7d" - email_object = EmailObjectBuilder({}, email_folder).build() - - await KnowledgeBase().do_embedding(email_object) - - -async def test_query_email(): msg_prompt = AgentPrompt() msg_prompt.messages = [{"role":"user","content":"abcdef"}] - KnowledgeBase().query(msg_prompt) + await KnowledgeBase().query_prompt(msg_prompt) + + class TestKnowledgeBase(unittest.TestCase): def test_embedding(self): - asyncio.run(test_embedding_email()) - - def test_query(self): - asyncio.run(test_query_email()) + asyncio.run(test_embedding_email(self)) if __name__ == "__main__": From 441583b7f8658d5306663499b5b10afd1f3401cb Mon Sep 17 00:00:00 2001 From: tsukasa Date: Mon, 18 Sep 2023 11:28:21 +0800 Subject: [PATCH 29/29] test prompt from knowledge object --- src/aios_kernel/knowledge_base.py | 40 ++++++++++++-------- src/aios_kernel/open_ai_node.py | 16 ++++++++ src/knowledge/core_object/document_object.py | 2 +- src/knowledge/data/writer.py | 33 ++++++++-------- src/knowledge/object/object.py | 3 +- src/knowledge/object/object_id.py | 6 ++- src/knowledge/vector/chroma_store.py | 9 +++-- 7 files changed, 71 insertions(+), 38 deletions(-) diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index dbb96e6..b229864 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -1,5 +1,6 @@ # define a knowledge base class import json +import logging from . import AgentPrompt, ComputeKernel from knowledge import * @@ -26,7 +27,7 @@ class KnowledgeBase: text = chunk.read().decode("utf-8") vector = await self.compute_kernel.do_text_embedding(text) - self.store.get_vector_store("default").insert(vector, chunk_id) + await self.store.get_vector_store("default").insert(vector, chunk_id) async def __embedding_image(self, image: ImageObject): desc = {} @@ -37,7 +38,7 @@ class KnowledgeBase: if not not image.get_tags(): desc["tags"] = image.get_tags() vector = await self.compute_kernel.do_text_embedding(json.dumps(desc)) - self.store.get_vector_store("default").insert(vector, image.calculate_id()) + await self.store.get_vector_store("default").insert(vector, image.calculate_id()) async def __embedding_video(self, vedio: VideoObject): desc = {} @@ -48,16 +49,20 @@ class KnowledgeBase: if not not vedio.get_tags(): desc["tags"] = vedio.get_tags() vector = await self.compute_kernel.do_text_embedding(json.dumps(desc)) - self.store.get_vector_store("default").insert(vector, vedio.calculate_id()) + await self.store.get_vector_store("default").insert(vector, vedio.calculate_id()) async def __embedding_rich_text(self, rich_text: RichTextObject): - for document in rich_text.get_documents().values(): + for document_id in rich_text.get_documents().values(): + document = DocumentObject.decode(self.store.get_object_store().get_object(document_id)) await self.__embedding_document(document) - for image in rich_text.get_images().values(): + for image_id in rich_text.get_images().values(): + image = ImageObject.decode(self.store.get_object_store().get_object(image_id)) await self.__embedding_image(image) - for vedio in rich_text.get_videos().values(): - await self.__embedding_video(vedio) - for rich_text in rich_text.get_rich_texts().values(): + for video_id in rich_text.get_videos().values(): + video = VideoObject.decode(self.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.store.get_object_store().get_object(rich_text_id)) await self.__embedding_rich_text(rich_text) async def __embedding_email(self, email: EmailObject): @@ -154,17 +159,19 @@ class KnowledgeBase: 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()}") prompt.append(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) + 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.append(object_ids) + results.extend(object_ids) return results def __load_object(self, object_id: ObjectID) -> KnowledgeObject: @@ -187,11 +194,12 @@ class KnowledgeBase: for object_id in object_ids: parents = self.store.get_relation_store().get_related_root_objects(object_id) # last parent is the root object - root_object_id = parents[-1] - if results[root_object_id] is None: - results[root_object_id] = [root_object_id, object_id] + 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[root_object_id].append(object_id) + results[str(root_object_id)] = [root_object_id, object_id] content = "I found the following contents described with json format:\n" result_desc = [] @@ -228,7 +236,7 @@ class KnowledgeBase: content += ".\n" prompt = AgentPrompt() - prompt.add_message("knowledge", content) + prompt.messages.append({"role": "knowledge", "content": content}) return prompt diff --git a/src/aios_kernel/open_ai_node.py b/src/aios_kernel/open_ai_node.py index a37706e..1c39c96 100644 --- a/src/aios_kernel/open_ai_node.py +++ b/src/aios_kernel/open_ai_node.py @@ -84,6 +84,22 @@ class OpenAI_ComputeNode(ComputeNode): resp = openai.Embedding.create(model=model_name, input=input) + + # resp = { + # "object": "list", + # "data": [ + # { + # "object": "embedding", + # "index": 0, + # "embedding": [ + # -0.00930514745414257, + # 0.00765434792265296, + # -0.007167573552578688, + # -0.012373941019177437, + # -0.04884673282504082 + # ]}] + # } + logger.info(f"openai response: {resp}") result = ComputeTaskResult() diff --git a/src/knowledge/core_object/document_object.py b/src/knowledge/core_object/document_object.py index ab02b8d..7210f16 100644 --- a/src/knowledge/core_object/document_object.py +++ b/src/knowledge/core_object/document_object.py @@ -53,7 +53,7 @@ class DocumentObjectBuilder: chunk_list = KnowledgeStore().get_chunk_list_writer().create_chunk_list_from_text( self.text, 1024 * 4, - "." + ".?!\n" ) doc = DocumentObject(self.meta, self.tags, chunk_list) doc_id = doc.calculate_id() diff --git a/src/knowledge/data/writer.py b/src/knowledge/data/writer.py index 599c51b..3ac0ed1 100644 --- a/src/knowledge/data/writer.py +++ b/src/knowledge/data/writer.py @@ -74,19 +74,20 @@ class ChunkListWriter: 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 \ No newline at end of file + # 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 diff --git a/src/knowledge/object/object.py b/src/knowledge/object/object.py index 490555b..a63e30a 100644 --- a/src/knowledge/object/object.py +++ b/src/knowledge/object/object.py @@ -54,7 +54,8 @@ class KnowledgeObject(ABC): ) sha256 = hashlib.sha256() sha256.update(data.encode()) - return ObjectID(sha256.digest()) + hash_bytes = sha256.digest() + return ObjectID(bytes([self.object_type]) + hash_bytes[1:]) def encode(self) -> bytes: return pickle.dumps(self) diff --git a/src/knowledge/object/object_id.py b/src/knowledge/object/object_id.py index 42f313a..46955f4 100644 --- a/src/knowledge/object/object_id.py +++ b/src/knowledge/object/object_id.py @@ -44,7 +44,11 @@ class ObjectID: # pylint: disable=too-few-public-methods @staticmethod def new_chunk_id(chunk_hash: HashValue): - return ObjectID(chunk_hash.value) + assert len(chunk_hash.value) == 32, "ObjectID must be 32 bytes long" + return ObjectID(bytes([ObjectType.Chunk]) + chunk_hash.value[1:]) + + def get_object_type(self) -> ObjectType: + return ObjectType(self.value[0]) @staticmethod def hash_data(data: bytes): diff --git a/src/knowledge/vector/chroma_store.py b/src/knowledge/vector/chroma_store.py index 0056375..03006ae 100644 --- a/src/knowledge/vector/chroma_store.py +++ b/src/knowledge/vector/chroma_store.py @@ -30,9 +30,10 @@ class ChromaVectorStore(VectorBase): self.collection = collection async def insert(self, vector: [float], id: ObjectID): + logging.info(f"will insert vector: {vector} id: {str(id)}") self.collection.add( embeddings=vector, - ids=id, + ids=str(id), ) async def query(self, vector: [float], top_k: int) -> [ObjectID]: @@ -40,8 +41,10 @@ class ChromaVectorStore(VectorBase): query_embeddings=vector, n_results=top_k, ) - - return ret["ids"] + logging.info(f"query result {ret}") + if len(ret['ids']) == 0: + return [] + return list(map(ObjectID.from_base58, ret["ids"][0])) async def delete(self, id: ObjectID): self.collection.delete(