agentic-kit-common 0.0.9__py3-none-any.whl → 0.0.10__py3-none-any.whl

agentic_kit_common/vector/embedding/embedding.py

@@ -1,6 +1,7 @@
 import abc
 
 from langchain_community.embeddings import XinferenceEmbeddings
+from langchain_openai import OpenAIEmbeddings
 
 
 class EmbeddingFactoryBase:
@@ -26,14 +27,28 @@ class XinferenceEmbeddingFactory(EmbeddingFactoryBase):
         return embedding
 
 
+class VllmEmbeddingFactory(EmbeddingFactoryBase):
+    provider: str = 'vllm'
+
+    @classmethod
+    def create_embedding(cls, base_url: str, model_uid: str, api_key: str = 'api_key', dims: int = 1024):
+        embedding = OpenAIEmbeddings(
+            openai_api_key=api_key,  # vLLM 不校验 key
+            openai_api_base=base_url,
+            model=model_uid,
+        )
+        return embedding
+
+
 _FACTORY_LIST = [
-    XinferenceEmbeddingFactory
+    XinferenceEmbeddingFactory,
+    VllmEmbeddingFactory,
 ]
 
 
 class EmbeddingFactory:
     @classmethod
-    def create_embedding(cls, base_url: str, model_uid: str, api_key: str = '', dims: int = 1024, provider: str = 'xinference'):
+    def create_embedding(cls, base_url: str, model_uid: str, api_key: str = 'api_key', 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)

agentic_kit_common/vector/manager/milvus_manager.py

@@ -4,7 +4,7 @@ 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 pymilvus import db, 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
@@ -18,14 +18,6 @@ class MilvusManager:
 
     collection_name: str
 
-    # search_field: str
-    #
-    # query_fields: List[str]
-    #
-    # output_fields: List[str]
-    #
-    # fields: List[FieldSchema]
-    #
     model_cls: BaseModel
 
     @classmethod
@@ -33,9 +25,9 @@ class MilvusManager:
         cls,
         embed_model: Embeddings,
         vector_store_uri: str,
-        # database_name: str,
-        # collection_name: str,
         # model_cls: BaseModel,
+        database_name: str = None,
+        collection_name: str = None,
         search_field: str = None,
         query_fields: List[str] = None,
         output_fields: List[str] = None,
@@ -62,8 +54,8 @@ class MilvusManager:
         return cls(
             embed_model=embed_model,
             vector_store_uri=vector_store_uri,
-            # database_name=database_name,
-            # collection_name=collection_name,
+            database_name=database_name,
+            collection_name=collection_name,
             search_field=search_field,
             query_fields=query_fields,
             output_fields=output_fields,
@@ -78,8 +70,8 @@ class MilvusManager:
         self,
         embed_model: Embeddings,
         vector_store_uri: str,
-        # database_name: str,
-        # collection_name: str,
+        database_name: str,
+        collection_name: str,
         search_field: str,
         query_fields: List[str],
         output_fields: List[str],
@@ -89,8 +81,15 @@ class MilvusManager:
         index_params_auto: Dict,
         search_params: Dict,
     ):
-        # self.database_name = database_name
-        # self.collection_name = collection_name
+        if database_name is None:
+            self.database_name = self.__class__.database_name
+        else:
+            self.database_name = database_name
+        if collection_name is None:
+            self.collection_name = self.__class__.collection_name
+        else:
+            self.collection_name = collection_name
+
         self.embed_model = embed_model
         self.vector_store_uri = vector_store_uri
         self.search_field = search_field
@@ -106,6 +105,16 @@ class MilvusManager:
         self.do_init()
 
     def do_init(self):
+        # 1. 先连到系统默认库，才能操作 database 级别
+        res = connections.connect(alias="default", uri=self.vector_store_uri)  # 不指定 db_name 即连 default
+        # 2. 验证连接成功
+        if not connections.has_connection("default"):
+            raise RuntimeError("Milvus 连接失败")
+        # 3. 判断并创建目标库
+        if self.database_name not in db.list_database():
+            db.create_database(self.database_name)  # pymilvus ≥2.2.9
+            print(f"database {self.database_name} created")
+
         res = connections.connect(
             alias="default",
             db_name=self.database_name,
@@ -201,13 +210,15 @@ class MilvusManager:
     def insert(self, items: list[dict]):
         uids = []
         vectors = []
+        owner_uids = []
         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'])
+            owner_uids.append(item['owner_uid'])
             vectors.append(self.embed_model.embed_query(item['content']))
 
-        insert_fields = [uids, vectors]
+        insert_fields = [uids, owner_uids, vectors]
         insert_result = self.collection.insert(insert_fields)
         self.collection.flush()
         return insert_result

agentic_kit_common/vector/schema/milvus_schema.py

@@ -3,11 +3,12 @@ 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="owner_uid", dtype=DataType.VARCHAR, max_length=64),
     FieldSchema(name="vectors", dtype=DataType.FLOAT_VECTOR, dim=1024),
 ]
-default_output_fields = ['pk', 'uid']
+default_output_fields = ['pk', 'uid', 'owner_uid']
 default_search_field = "vectors"
-default_query_fields = ['uid']
+default_query_fields = ['uid', 'owner_uid']
 
 default_index_params_vector = {
     'metric_type': 'COSINE',  # IP 或 COSINE

{agentic_kit_common-0.0.9.dist-info → agentic_kit_common-0.0.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentic-kit-common
-Version: 0.0.9
+Version: 0.0.10
 Summary: Common utilities and tools for agentic kit ecosystem
 Home-page: 
 Author: manson

{agentic_kit_common-0.0.9.dist-info → agentic_kit_common-0.0.10.dist-info}/RECORD

@@ -10,23 +10,23 @@ agentic_kit_common/orm/schema.py,sha256=ukdVP71NE_JO5_HOe_FApJjBAKUsDluoElVn5vNS
 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/embedding/embedding.py,sha256=ZEvl2w2-PrAph60u2wJItj4lb68XxK805tRzOp49MdI,1861
 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/manager/milvus_manager.py,sha256=gDwWlkrnIUFqMoDNsvqnPAKTyAUnyxD-fh7jy93LcQY,8568
 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/vector/schema/milvus_schema.py,sha256=4kIUtIIE2vcKP1s1Ghn_eikuzDb1lsSeAaTwrsNzUlQ,784
 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_embedding.py,sha256=xcfHDHGL2_tpXp_VaLnNDAWvKdkc1K1BjiYoV5WxtFY,900
 test/test_minio.py,sha256=TOkX8A2pPkjrwAIH88xBZmFpDc1ZgTk1QS6mtubOZ-Y,2308
 test/test_orm.py,sha256=8fGCU7BWaD5sDbg0fgYN0Saf_hi7t-q8svHCHLQDceo,1620
-test/test_vector.py,sha256=4sJ7bXt5iRhc6j1Ig0bZ1j0huPU-clnsnT7O0-uFbkw,1311
-agentic_kit_common-0.0.9.dist-info/METADATA,sha256=2QyqNG2MQrLYyQ6TDuPXkyBnPdf5La0g4-rloQ4jAZY,7496
-agentic_kit_common-0.0.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-agentic_kit_common-0.0.9.dist-info/top_level.txt,sha256=nEKDlp84vqKSVWssGcxyuIsTqWLhMo45xqMs2GK4Dgg,24
-agentic_kit_common-0.0.9.dist-info/RECORD,,
+test/test_vector.py,sha256=Z3Bwvrw0XGBWXHYXk9pkF1cjnUZSZk0gmLsgQIIaZuY,1717
+agentic_kit_common-0.0.10.dist-info/METADATA,sha256=KIp7HPRyzygUBrXu8H9B4xu8NOFiD6_B5PTaYVN20Fc,7497
+agentic_kit_common-0.0.10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+agentic_kit_common-0.0.10.dist-info/top_level.txt,sha256=nEKDlp84vqKSVWssGcxyuIsTqWLhMo45xqMs2GK4Dgg,24
+agentic_kit_common-0.0.10.dist-info/RECORD,,

test/test_embedding.py

@@ -6,13 +6,22 @@ project_root = Path(__file__).parent.parent
 sys.path.insert(0, str(project_root))
 
 from agentic_kit_common.vector.embedding import EmbeddingFactory
+from langchain_openai import OpenAIEmbeddings
 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)
+        # embedding = EmbeddingFactory.create_embedding(base_url=global_settings.embedding_base_url, model_uid=global_settings.embedding_model_uid)
+
+        embedding = EmbeddingFactory.create_embedding(
+            base_url=global_settings.embedding_base_url,
+            model_uid=global_settings.embedding_model_uid,
+            provider='vllm'
+        )
+
         text = '你好'
+        # text = 'hello, world'
         text_vector = embedding.embed_query(text=text)
         print(text_vector)
 

test/test_vector.py

@@ -17,11 +17,11 @@ 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)
+embedding = EmbeddingFactory.create_embedding(base_url=global_settings.embedding_base_url, model_uid=global_settings.embedding_model_uid, provider='vllm')
 
 
 class McpServerManager(MilvusManager):
-    database_name: str = 'skill_center'
+    database_name: str = 'nl2sql'
 
     collection_name: str = 'mcp_server'
 
@@ -31,8 +31,16 @@ class McpServerManager(MilvusManager):
 class MyTestCase(unittest.TestCase):
     def test_vector(self):
         manager = McpServerManager.create(embed_model=embedding, vector_store_uri=global_settings.milvus_url)
+        # manager = MilvusManager.create(embed_model=embedding, vector_store_uri=global_settings.milvus_url, database_name='nl2sql', collection_name='mcp_server2')
+        # manager.insert([{
+        #     'uid': 'xxxx',
+        #     'content': 'hello',
+        #     'owner_uid': 'xxxx2'
+        # }])
 
-        manager.query()
+        owner_uid = 'xxxx2'
+        res = manager.query(expr=f'owner_uid == "{owner_uid}"')
+        print(res)
 
         # manager.search('你好')
 
@@ -43,9 +51,9 @@ class MyTestCase(unittest.TestCase):
         #     'content': 'hello',
         # }])
 
-        manager.delete_by_uid(uid='2050d9a9-b125-42b9-80ce-bd9eee89d9eb')
+        # manager.delete_by_uid(uid='2050d9a9-b125-42b9-80ce-bd9eee89d9eb')
 
-        manager.query()
+        # manager.query()
 
 
 if __name__ == '__main__':