agentic-kit-common 0.0.4__py3-none-any.whl → 0.0.5__py3-none-any.whl

agentic_kit_common/vector/embedding/__init__.py

@@ -0,0 +1 @@
+from .embedding import EmbeddingFactory

agentic_kit_common/vector/embedding/embedding.py

@@ -0,0 +1,40 @@
+import abc
+
+from langchain_community.embeddings import XinferenceEmbeddings
+
+
+class EmbeddingFactoryBase:
+    provider: str
+
+    @classmethod
+    @abc.abstractmethod
+    def create_embedding(cls, base_url: str, model_uid: str, api_key: str = '', dims: int = 1024):
+        raise NotImplemented
+
+
+class XinferenceEmbeddingFactory(EmbeddingFactoryBase):
+    provider: str = 'xinference'
+
+    @classmethod
+    def create_embedding(cls, base_url: str, model_uid: str, api_key: str = '', dims: int = 1024):
+        # assert embedding_server_url is not None and embedding_server_url != ''
+        # assert embedding_model_uid is not None and embedding_model_uid != ''
+        embedding = XinferenceEmbeddings(
+            server_url=base_url,
+            model_uid=model_uid,  # replace model_uid with the model UID return from launching the model
+        )
+        return embedding
+
+
+_FACTORY_LIST = [
+    XinferenceEmbeddingFactory
+]
+
+
+class EmbeddingFactory:
+    @classmethod
+    def create_embedding(cls, base_url: str, model_uid: str, api_key: str = '', dims: int = 1024, provider: str = 'xinference'):
+        for factory in _FACTORY_LIST:
+            if factory.provider == provider:
+                return factory.create_embedding(base_url=base_url, model_uid=model_uid, api_key=api_key, dims=dims)
+        raise Exception(f'不支持的embedding provider配置: [{provider}]')

agentic_kit_common/vector/manager/__init__.py

@@ -0,0 +1 @@
+from .milvus_manager import MilvusManager

agentic_kit_common/vector/manager/milvus_manager.py

@@ -0,0 +1,226 @@
+import logging
+from typing import List, Dict
+
+from langchain_core.embeddings import Embeddings
+from pydantic import BaseModel
+from pymilvus import connections, Collection, CollectionSchema, FieldSchema
+from pymilvus.orm import utility
+
+from ..schema import default_search_params, default_index_params_auto, default_index_params_vector, default_fields, \
+    default_output_fields, default_search_field, default_query_fields
+
+logger = logging.getLogger(__name__)
+logger.setLevel(level=logging.DEBUG)
+
+
+class MilvusManager:
+    database_name: str
+
+    collection_name: str
+
+    # search_field: str
+    #
+    # query_fields: List[str]
+    #
+    # output_fields: List[str]
+    #
+    # fields: List[FieldSchema]
+    #
+    model_cls: BaseModel
+
+    @classmethod
+    def create(
+        cls,
+        embed_model: Embeddings,
+        vector_store_uri: str,
+        # database_name: str,
+        # collection_name: str,
+        # model_cls: BaseModel,
+        search_field: str = None,
+        query_fields: List[str] = None,
+        output_fields: List[str] = None,
+        fields: List[FieldSchema] = None,
+        index_params_vector: Dict = None,
+        index_params_auto: Dict = None,
+        search_params: Dict = None,
+    ):
+        if fields is None:
+            fields = default_fields
+        if output_fields is None:
+            output_fields = default_output_fields
+        if search_field is None:
+            search_field = default_search_field
+        if query_fields is None:
+            query_fields = default_query_fields
+        if index_params_vector is None:
+            index_params_vector = default_index_params_vector
+        if index_params_auto is None:
+            index_params_auto = default_index_params_auto
+        if search_params is None:
+            search_params = default_search_params
+
+        return cls(
+            embed_model=embed_model,
+            vector_store_uri=vector_store_uri,
+            # database_name=database_name,
+            # collection_name=collection_name,
+            search_field=search_field,
+            query_fields=query_fields,
+            output_fields=output_fields,
+            fields=fields,
+            # model_cls=model_cls,
+            index_params_vector=index_params_vector,
+            index_params_auto=index_params_auto,
+            search_params=search_params,
+        )
+
+    def __init__(
+        self,
+        embed_model: Embeddings,
+        vector_store_uri: str,
+        # database_name: str,
+        # collection_name: str,
+        search_field: str,
+        query_fields: List[str],
+        output_fields: List[str],
+        fields: List[FieldSchema],
+        # model_cls: BaseModel,
+        index_params_vector: Dict,
+        index_params_auto: Dict,
+        search_params: Dict,
+    ):
+        # self.database_name = database_name
+        # self.collection_name = collection_name
+        self.embed_model = embed_model
+        self.vector_store_uri = vector_store_uri
+        self.search_field = search_field
+        self.query_fields = query_fields
+        self.output_fields = output_fields
+        self.fields = fields
+        # self.model_cls = model_cls
+        self.index_params_vector = index_params_vector
+        self.index_params_auto = index_params_auto
+        self.search_params = search_params
+        self.collection = None
+
+        self.do_init()
+
+    def do_init(self):
+        res = connections.connect(
+            alias="default",
+            db_name=self.database_name,
+            uri=self.vector_store_uri
+        )
+        assert res is None  # no error
+
+        if utility.has_collection(self.collection_name):
+            self.collection = Collection(self.collection_name)
+        else:
+            tool_schema = CollectionSchema(self.fields, self.database_name)
+            self.collection = Collection(self.collection_name, tool_schema)
+            self.collection.create_index(self.search_field, self.index_params_vector)
+            for field in self.query_fields:
+                self.collection.create_index(field, self.index_params_auto)
+        self.collection.load()
+
+    def delete_by_uid(self, uid: str):
+        expr = f"uid=='{uid}'"
+        return self.delete(expr=expr)
+
+    def delete(self, expr: str):
+        delete_result = self.collection.delete(expr=expr)
+        self.collection.flush()
+        return delete_result
+
+    def search(self, query: str, expr: str = None, top_k: int = 5, score_threshold: float = 0, search_params: dict = None):
+        """vector搜索"""
+        if search_params is None:
+            search_params = self.search_params
+        qv = self.embed_model.embed_query(query)
+        results = self.collection.search(
+            data=[qv],  # 查询向量
+            anns_field=self.search_field,  # 指定向量字段名称
+            param=search_params,  # 搜索参数
+            limit=top_k,  # 返回最相似的top-k个结果
+            expr=expr,
+            output_fields=self.output_fields  # 返回的字段
+        )
+        logger.debug(f'search results: {results}')
+        return self._parse_result_to_schema(results=results, score_threshold=score_threshold)
+
+    def query(self, expr: str = '', is_page: bool = False, page_size: int = 20, page: int = 1):
+        """表达式查询，支持分页"""
+
+        # note: 如果没有表达式，必须分页，否则数据太多
+        if expr is None or expr == '':
+            is_page = True
+
+        if is_page:
+            _total = self.collection.query(
+                expr=expr,
+                output_fields=["count(*)"]
+            )
+            total = _total[0]['count(*)']
+
+            results = self.collection.query(
+                expr=expr,
+                limit=page_size,
+                offset=(page - 1) * page_size,
+                output_fields=self.output_fields  # 返回的字段
+            )
+            _result = []
+            for result in results:
+                _result.append(self.model_cls(**result))
+            logger.debug(f'query paged results: {results}')
+
+            data = {
+                'total': total,
+                'page': page,
+                'page_size': page_size if page_size < total else total,
+                'data': _result
+            }
+            return data
+        else:
+            results = self.collection.query(
+                expr=expr,
+                output_fields=self.output_fields  # 返回的字段
+            )
+            _result = []
+            for result in results:
+                _result.append(self.model_cls(**result))
+            logger.debug(f'query not paged results: {results}')
+            return _result
+
+    def clear(self):
+        if self.collection:
+            self.collection.flush()
+            self.collection.release()
+            self.collection.drop()
+            self.do_init()
+
+    def insert(self, items: list[dict]):
+        uids = []
+        vectors = []
+        for item in items:
+            if 'uid' not in item or 'content' not in item:
+                raise Exception(f'insert items must include [uid, content], illegal for {item}')
+            uids.append(item['uid'])
+            vectors.append(self.embed_model.embed_query(item['content']))
+
+        insert_fields = [uids, vectors]
+        insert_result = self.collection.insert(insert_fields)
+        self.collection.flush()
+        return insert_result
+
+    def _parse_result_to_schema(self, results, score_threshold: float = 0):
+        """过滤结果"""
+        _results = []
+        for hits in results:
+            for hit in hits:
+                if score_threshold != 0:
+                    if hit.score >= score_threshold:
+                        _results.append(self.model_cls(score=hit.score, **hit.fields))
+                else:
+                    _results.append(self.model_cls(score=hit.score, **hit.fields))
+
+        return _results

agentic_kit_common/vector/schema/__init__.py

@@ -0,0 +1,3 @@
+from .base import VectorBaseModel
+from .milvus_schema import default_fields, default_output_fields, default_query_fields, default_search_field, \
+    default_index_params_auto, default_index_params_vector, default_search_params

agentic_kit_common/vector/schema/base.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel
+
+
+class VectorBaseModel(BaseModel):
+    uid: str
+
+    score: float = 0

agentic_kit_common/vector/schema/milvus_schema.py

@@ -0,0 +1,25 @@
+from pymilvus import FieldSchema, DataType
+
+default_fields = [
+    FieldSchema(name="pk", dtype=DataType.INT64, is_primary=True, auto_id=True),
+    FieldSchema(name="uid", dtype=DataType.VARCHAR, max_length=36),
+    FieldSchema(name="vectors", dtype=DataType.FLOAT_VECTOR, dim=1024),
+]
+default_output_fields = ['pk', 'uid']
+default_search_field = "vectors"
+default_query_fields = ['uid']
+
+default_index_params_vector = {
+    'metric_type': 'COSINE',  # IP 或 COSINE
+    'index_type': 'FLAT',  # FLAT
+    "params": {"nlist": 128},
+}
+
+default_index_params_auto = {
+    'index_type': 'AUTOINDEX'
+}
+
+default_search_params = {
+    'metric_type': 'COSINE',
+    "params": {"nprobe": 10},
+}

{agentic_kit_common-0.0.4.dist-info → agentic_kit_common-0.0.5.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: agentic-kit-common
-Version: 0.0.4
-Summary: Common utilities and tools for agentic kit ecosystem, including MinIO manager and other shared components
+Version: 0.0.5
+Summary: Common utilities and tools for agentic kit ecosystem
 Home-page: 
 Author: manson
 Author-email: manson.li3307@gmail.com
@@ -26,6 +26,8 @@ Requires-Dist: langchain_community
 Requires-Dist: langchain_experimental
 Requires-Dist: mysql-connector-python
 Requires-Dist: sqlalchemy
+Requires-Dist: pymilvus
+Requires-Dist: xinference_client
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier

agentic_kit_common-0.0.5.dist-info/RECORD

@@ -0,0 +1,28 @@
+agentic_kit_common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+agentic_kit_common/minio/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+agentic_kit_common/minio/minio_manager.py,sha256=PzoUWn0YqXTHx1UClbkwLkIUmv5O4aSDc7eVf08qOzs,21494
+agentic_kit_common/orm/__init__.py,sha256=wqY81g3P7FftFvLK5SaxFJzNNxrQwtmcb4RCXOSAZa8,71
+agentic_kit_common/orm/base.py,sha256=QIura_i2nIY2XeA3-KkO2loLNbEAoJK2qx0hu_8nhYU,2277
+agentic_kit_common/orm/manager.py,sha256=HU8dHNb76uTo2KMOXdNG3nn3uea7vPYIFPO61SQoNIU,4648
+agentic_kit_common/orm/session.py,sha256=LX4ZUKJNXdcQ0KqRt5L0pJX-QG-tDyKXEDijcPUkGD0,2716
+agentic_kit_common/vector/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+agentic_kit_common/vector/embedding/__init__.py,sha256=mp--cfCDzTn5hNxzcIAU4m0g4F0d0rrZErFVWsGDvx4,40
+agentic_kit_common/vector/embedding/embedding.py,sha256=k1hHNkBcucu6lSe2hYSrCal-ML26Eo6Bfe8cVkXSww8,1380
+agentic_kit_common/vector/manager/__init__.py,sha256=w2uAmKGRx9Nv3QySXIAgxQGlYR549AQvbBuBFsQttKI,42
+agentic_kit_common/vector/manager/milvus_manager.py,sha256=HaoKKsyMOfmA3jX45XEVhmK6SprjTmwULOyhCqqAIGg,7840
+agentic_kit_common/vector/schema/__init__.py,sha256=F8WCi4ybnujFW-qHcDI9KOqzw7i8r07dmBUO2R2_Iu0,228
+agentic_kit_common/vector/schema/base.py,sha256=oo5OSFCX9UeQuWskiebqctF1dyXISD7_czzoKYoHnKg,102
+agentic_kit_common/vector/schema/milvus_schema.py,sha256=V8igxe6_FmZ46jicNdKxBbf0vl3bwTC4AyRf9gqB_VA,684
+agentic_kit_common/web/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+agentic_kit_common/web/http/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+agentic_kit_common/web/http/response.py,sha256=uxd_MRHsFfQ0pUeHITDQD8tuOY6fo6IJ0MO6uL62AlM,1347
+test/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+test/config.py,sha256=VMcvfWnuWZKr_p0rC57i90jTq-XVsRuVzfEIe1C7Drg,3804
+test/settings.py,sha256=vU1dqUvGzshi6MG7JbGfW4jKfUsfAObgUfky-LMjUNs,344
+test/test_embedding.py,sha256=78M6t9tt6-oCfdQZfOartCIkUFZMvogcu5UKJh1lJb8,609
+test/test_minio.py,sha256=TOkX8A2pPkjrwAIH88xBZmFpDc1ZgTk1QS6mtubOZ-Y,2308
+test/test_vector.py,sha256=4sJ7bXt5iRhc6j1Ig0bZ1j0huPU-clnsnT7O0-uFbkw,1311
+agentic_kit_common-0.0.5.dist-info/METADATA,sha256=YBhN8h9vA8H2kYHP0SyGh0Vgj-f1PiCjQLmVhoYCRG4,7496
+agentic_kit_common-0.0.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+agentic_kit_common-0.0.5.dist-info/top_level.txt,sha256=nEKDlp84vqKSVWssGcxyuIsTqWLhMo45xqMs2GK4Dgg,24
+agentic_kit_common-0.0.5.dist-info/RECORD,,

test/settings.py

@@ -0,0 +1,18 @@
+from dotenv import load_dotenv
+from pydantic_settings import BaseSettings
+
+load_dotenv()
+
+class GlobalSettings(BaseSettings):
+    # embedding
+    embedding_base_url: str
+    embedding_model_uid: str
+    embedding_dims: int
+    embedding_api_key: str
+
+    # milvus
+    milvus_url: str
+
+
+global_settings = GlobalSettings()
+print(global_settings)

test/test_embedding.py

@@ -0,0 +1,21 @@
+import sys
+import unittest
+from pathlib import Path
+
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root))
+
+from agentic_kit_common.vector.embedding import EmbeddingFactory
+from .settings import global_settings
+
+
+class MyTestCase(unittest.TestCase):
+    def test_embedding(self):
+        embedding = EmbeddingFactory.create_embedding(base_url=global_settings.embedding_base_url, model_uid=global_settings.embedding_model_uid)
+        text = '你好'
+        text_vector = embedding.embed_query(text=text)
+        print(text_vector)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/test_vector.py

@@ -0,0 +1,52 @@
+import logging
+import sys
+import unittest
+import uuid
+from pathlib import Path
+
+from pydantic import BaseModel
+
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root))
+
+from agentic_kit_common.vector.embedding import EmbeddingFactory
+from agentic_kit_common.vector.manager import MilvusManager
+from agentic_kit_common.vector.schema import VectorBaseModel
+from .settings import global_settings
+
+
+logging.basicConfig(level=logging.DEBUG)
+
+embedding = EmbeddingFactory.create_embedding(base_url=global_settings.embedding_base_url, model_uid=global_settings.embedding_model_uid)
+
+
+class McpServerManager(MilvusManager):
+    database_name: str = 'skill_center'
+
+    collection_name: str = 'mcp_server'
+
+    model_cls: BaseModel = VectorBaseModel
+
+
+class MyTestCase(unittest.TestCase):
+    def test_vector(self):
+        manager = McpServerManager.create(embed_model=embedding, vector_store_uri=global_settings.milvus_url)
+
+        manager.query()
+
+        # manager.search('你好')
+
+        # uid = str(uuid.uuid4())
+        # print(uid)
+        # manager.insert([{
+        #     'uid': uid,
+        #     'content': 'hello',
+        # }])
+
+        manager.delete_by_uid(uid='2050d9a9-b125-42b9-80ce-bd9eee89d9eb')
+
+        manager.query()
+
+
+if __name__ == '__main__':
+    unittest.main()

agentic_kit_common-0.0.4.dist-info/RECORD

@@ -1,17 +0,0 @@
-agentic_kit_common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-agentic_kit_common/minio/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-agentic_kit_common/minio/minio_manager.py,sha256=PzoUWn0YqXTHx1UClbkwLkIUmv5O4aSDc7eVf08qOzs,21494
-agentic_kit_common/orm/__init__.py,sha256=wqY81g3P7FftFvLK5SaxFJzNNxrQwtmcb4RCXOSAZa8,71
-agentic_kit_common/orm/base.py,sha256=QIura_i2nIY2XeA3-KkO2loLNbEAoJK2qx0hu_8nhYU,2277
-agentic_kit_common/orm/manager.py,sha256=HU8dHNb76uTo2KMOXdNG3nn3uea7vPYIFPO61SQoNIU,4648
-agentic_kit_common/orm/session.py,sha256=LX4ZUKJNXdcQ0KqRt5L0pJX-QG-tDyKXEDijcPUkGD0,2716
-agentic_kit_common/web/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-agentic_kit_common/web/http/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-agentic_kit_common/web/http/response.py,sha256=uxd_MRHsFfQ0pUeHITDQD8tuOY6fo6IJ0MO6uL62AlM,1347
-test/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-test/config.py,sha256=VMcvfWnuWZKr_p0rC57i90jTq-XVsRuVzfEIe1C7Drg,3804
-test/test_minio.py,sha256=TOkX8A2pPkjrwAIH88xBZmFpDc1ZgTk1QS6mtubOZ-Y,2308
-agentic_kit_common-0.0.4.dist-info/METADATA,sha256=MXCYS7a4zdlactafV82PK_TcW933GqrDOTaspAK_g4Q,7492
-agentic_kit_common-0.0.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-agentic_kit_common-0.0.4.dist-info/top_level.txt,sha256=nEKDlp84vqKSVWssGcxyuIsTqWLhMo45xqMs2GK4Dgg,24
-agentic_kit_common-0.0.4.dist-info/RECORD,,