agentic-kit-common 0.0.4__tar.gz → 0.0.6__tar.gz

{agentic_kit_common-0.0.4 → agentic_kit_common-0.0.6}/PKG-INFO

@@ -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.6
+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.4 → agentic_kit_common-0.0.6}/agentic_kit_common/orm/manager.py

@@ -14,8 +14,7 @@ class BaseOrmManager(object):
     def get_by_id(cls, obj_id, close_session_after_curd=False):
         with get_db_session() as db:
             query = db.query(cls._model_cls) \
-                .filter(cls._model_cls.id == obj_id) \
-                .filter(cls._model_cls.active == True)
+                .filter(cls._model_cls.id == obj_id)
 
             if hasattr(cls._model_cls, 'active'):
                 query = query.filter(cls._model_cls.active == 1)
@@ -124,9 +123,10 @@ class BaseOrmManager(object):
     @classmethod
     def soft_delete_obj(cls, obj):
         """软删除obj信息"""
-        cls.update_obj(obj, **{
-            'active': False
-        })
+        if hasattr(cls._model_cls, 'active'):
+            cls.update_obj(obj, **{
+                'active': False
+            })
         return obj
 
     @classmethod

agentic_kit_common-0.0.6/agentic_kit_common/vector/embedding/__init__.py

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

agentic_kit_common-0.0.6/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-0.0.6/agentic_kit_common/vector/manager/__init__.py

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

agentic_kit_common-0.0.6/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-0.0.6/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-0.0.6/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-0.0.6/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 → agentic_kit_common-0.0.6}/agentic_kit_common.egg-info/PKG-INFO

@@ -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.6
+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.4 → agentic_kit_common-0.0.6}/agentic_kit_common.egg-info/SOURCES.txt

@@ -12,9 +12,20 @@ agentic_kit_common/orm/__init__.py
 agentic_kit_common/orm/base.py
 agentic_kit_common/orm/manager.py
 agentic_kit_common/orm/session.py
+agentic_kit_common/vector/__init__.py
+agentic_kit_common/vector/embedding/__init__.py
+agentic_kit_common/vector/embedding/embedding.py
+agentic_kit_common/vector/manager/__init__.py
+agentic_kit_common/vector/manager/milvus_manager.py
+agentic_kit_common/vector/schema/__init__.py
+agentic_kit_common/vector/schema/base.py
+agentic_kit_common/vector/schema/milvus_schema.py
 agentic_kit_common/web/__init__.py
 agentic_kit_common/web/http/__init__.py
 agentic_kit_common/web/http/response.py
 test/__init__.py
 test/config.py
-test/test_minio.py
+test/settings.py
+test/test_embedding.py
+test/test_minio.py
+test/test_vector.py

{agentic_kit_common-0.0.4 → agentic_kit_common-0.0.6}/agentic_kit_common.egg-info/requires.txt

@@ -8,3 +8,5 @@ langchain_community
 langchain_experimental
 mysql-connector-python
 sqlalchemy
+pymilvus
+xinference_client

{agentic_kit_common-0.0.4 → agentic_kit_common-0.0.6}/setup.py

@@ -4,10 +4,10 @@ from setuptools import setup, find_packages
 
 setup(
     name='agentic-kit-common',
-    version="0.0.4",
+    version="0.0.6",
     author="manson",
     author_email="manson.li3307@gmail.com",
-    description='Common utilities and tools for agentic kit ecosystem, including MinIO manager and other shared components',
+    description='Common utilities and tools for agentic kit ecosystem',
     long_description=open('README.md', encoding='utf-8').read(),
     long_description_content_type='text/markdown',
     url='',
@@ -41,6 +41,12 @@ setup(
         # sqlalchemy
         "mysql-connector-python",
         "sqlalchemy",
+
+        # milvus
+        "pymilvus",
+
+        # xinference
+        "xinference_client",
     ],
     keywords=['AI', 'LLM', 'Agent', 'MinIO', 'Storage', 'Common', 'Utilities'],
     include_package_data=True,

agentic_kit_common-0.0.6/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)

agentic_kit_common-0.0.6/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()

agentic_kit_common-0.0.6/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()