diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index d2debd8..b528192 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -20,6 +20,7 @@ class KnowledgeBase: def __singleton_init__(self) -> None: self.store = KnowledgeStore() self.compute_kernel = ComputeKernel.get_instance() + self._default_text_model = "all-MiniLM-L6-v2" async def __embedding_document(self, document: DocumentObject): for chunk_id in document.get_chunk_list(): @@ -28,8 +29,8 @@ class KnowledgeBase: raise ValueError(f"text chunk not found: {chunk_id}") text = chunk.read().decode("utf-8") - vector = await self.compute_kernel.do_text_embedding(text) - await self.store.get_vector_store("default").insert(vector, chunk_id) + vector = await self.compute_kernel.do_text_embedding(text, self._default_text_model) + await self.store.get_vector_store(self._default_text_model).insert(vector, chunk_id) async def __embedding_image(self, image: ImageObject): desc = {} @@ -39,8 +40,8 @@ class KnowledgeBase: desc["exif"] = image.get_exif() if not not image.get_tags(): desc["tags"] = image.get_tags() - vector = await self.compute_kernel.do_text_embedding(json.dumps(desc)) - await self.store.get_vector_store("default").insert(vector, image.calculate_id()) + vector = await self.compute_kernel.do_text_embedding(json.dumps(desc), self._default_text_model) + await self.store.get_vector_store(self._default_text_model).insert(vector, image.calculate_id()) async def __embedding_video(self, vedio: VideoObject): desc = {} @@ -50,8 +51,8 @@ class KnowledgeBase: desc["info"] = vedio.get_info() if not not vedio.get_tags(): desc["tags"] = vedio.get_tags() - vector = await self.compute_kernel.do_text_embedding(json.dumps(desc)) - await self.store.get_vector_store("default").insert(vector, vedio.calculate_id()) + vector = await self.compute_kernel.do_text_embedding(json.dumps(desc), self._default_text_model) + await self.store.get_vector_store(self._default_text_model).insert(vector, vedio.calculate_id()) async def __embedding_rich_text(self, rich_text: RichTextObject): for document_id in rich_text.get_documents().values(): @@ -68,8 +69,8 @@ class KnowledgeBase: await self.__embedding_rich_text(rich_text) async def __embedding_email(self, email: EmailObject): - vector = await self.compute_kernel.do_text_embedding(json.dumps(email.get_desc())) - await self.store.get_vector_store("default").insert(vector, email.calculate_id()) + vector = await self.compute_kernel.do_text_embedding(json.dumps(email.get_desc()), self._default_text_model) + await self.store.get_vector_store(self._default_text_model).insert(vector, email.calculate_id()) await self.__embedding_rich_text(email.get_rich_text()) @@ -172,8 +173,8 @@ class KnowledgeBase: results = [] for msg in prompt.messages: if msg["role"] == "user": - vector = await self.compute_kernel.do_text_embedding(msg["content"]) - object_ids = await self.store.get_vector_store("default").query(vector, 10) + vector = await self.compute_kernel.do_text_embedding(msg["content"], self._default_text_model) + object_ids = await self.store.get_vector_store(self._default_text_model).query(vector, 10) results.extend(object_ids) return results diff --git a/src/aios_kernel/local_st_compute_node.py b/src/aios_kernel/local_st_compute_node.py index e73b9f9..bc9f4af 100644 --- a/src/aios_kernel/local_st_compute_node.py +++ b/src/aios_kernel/local_st_compute_node.py @@ -89,7 +89,7 @@ class LocalSentenceTransformer_ComputeNode(Queue_ComputeNode): def is_support(self, task: ComputeTask) -> bool: return task.task_type == ComputeTaskType.TEXT_EMBEDDING and ( - not task.params["model_name"] or task.params["model_name"] == "llama" + not task.params["model_name"] or task.params["model_name"] == "all-MiniLM-L6-v2" ) def is_local(self) -> bool: diff --git a/src/knowledge/data/writer.py b/src/knowledge/data/writer.py index 0a87ac1..0381bc6 100644 --- a/src/knowledge/data/writer.py +++ b/src/knowledge/data/writer.py @@ -3,14 +3,14 @@ import hashlib import re import tiktoken import logging -from typing import Tuple, List +from typing import Callable, Iterable, Optional, Tuple, List from .chunk_store import ChunkStore from .chunk import ChunkID, PositionFileRange, PositionType from ..object import HashValue from .tracker import ChunkTracker from .chunk_list import ChunkList -def _join_docs(self, docs: List[str], separator: str) -> Optional[str]: +def _join_docs(docs: List[str], separator: str) -> Optional[str]: text = separator.join(docs) text = text.strip() if text == "": @@ -19,7 +19,6 @@ def _join_docs(self, docs: List[str], separator: str) -> Optional[str]: return text def _merge_splits( - self, splits: Iterable[str], separator: str, chunk_size: int,