Add chromadb based vector store impl and test case
This commit is contained in:
@@ -1,2 +0,0 @@
|
|||||||
from .object import EmailObject, ImageObject, TextChunkObject, ObjectID
|
|
||||||
from .knowledge_base import KnowledgeBase
|
|
||||||
@@ -1,46 +0,0 @@
|
|||||||
# 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
|
|
||||||
|
|
||||||
class KnowledgeBase:
|
|
||||||
def __init__(self) -> None:
|
|
||||||
self.object_store = ObjectStorage()
|
|
||||||
self.vector_base = VectorBase()
|
|
||||||
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())
|
|
||||||
else:
|
|
||||||
pass
|
|
||||||
|
|
||||||
async def query(self, prompt: AgentPrompt) -> AgentPrompt:
|
|
||||||
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
|
|
||||||
@@ -1,14 +0,0 @@
|
|||||||
# 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
|
|
||||||
@@ -1,7 +1,7 @@
|
|||||||
|
|
||||||
# define a object type enum
|
# define a object type enum
|
||||||
from abc import ABC
|
from abc import ABC, abstractmethod
|
||||||
|
from enum import Enum
|
||||||
|
|
||||||
class ObjectType(Enum):
|
class ObjectType(Enum):
|
||||||
TextChunk = 1
|
TextChunk = 1
|
||||||
Reference in New Issue
Block a user