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

agentic_kit_common/log/logger.py

@@ -0,0 +1,16 @@
+import logging
+
+
+def setup_logger(tag, level):
+    logger = logging.getLogger(tag)
+    logger.setLevel(level)
+
+    if not logger.handlers:
+        handler = logging.StreamHandler()
+        formatter = logging.Formatter(
+            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+        )
+        handler.setFormatter(formatter)
+        logger.addHandler(handler)
+
+    return logger

agentic_kit_common/orm/execution.py

@@ -1,53 +1,90 @@
-from typing import List, Union
+from typing import List, Union, Optional, Set
 
-import sqlparse
+import sqlglot
 from sqlalchemy import text
 from sqlalchemy.orm import Session
+from sqlglot import expressions
 
 
-def is_readonly_sql(sql: str) -> bool:
+def get_operation_type(statement) -> str:
+    """获取SQL操作类型"""
+    return type(statement).__name__.upper()
+
+
+def get_sql_operation_info(sql: str) -> Optional[dict]:
     """
-    只允许：
-      SELECT / WITH / VALUES / EXPLAIN / DESCRIBE / SHOW
-    拒绝：
-      INSERT/UPDATE/DELETE/CREATE/ALTER/DROP/TRUNCATE/LOAD/REPLACE/LOCK/UNLOCK/GRANT/REVOKE/EXECUTE/CALL
+    获取SQL操作信息（用于调试和日志记录）
+
+    Returns:
+        dict: 包含操作类型、涉及表等信息
     """
-    sql_clean = sqlparse.format(sql.strip(), strip_comments=True)
-    tokens = [t.normalized for t in sqlparse.parse(sql_clean)[0].flatten()
-              if t.ttype in (sqlparse.tokens.Keyword, sqlparse.tokens.DML, sqlparse.tokens.DDL)]
+    try:
+        statement = sqlglot.parse_one(sql, read="mysql")
+        operation_type = get_operation_type(statement)
+
+        tables = []
+        for table in statement.find_all(expressions.Table):
+            table_name = table.name
+            if table_name.startswith('`') and table_name.endswith('`'):
+                table_name = table_name[1:-1]
+            tables.append(table_name)
+
+        return {
+            'operation': operation_type,
+            'tables': tables,
+            'has_wildcard': bool(statement.find(expressions.Star)) if isinstance(statement, expressions.Select) else False
+        }
+    except Exception as e:
+        print(f"获取SQL操作信息失败: {e}")
+        return None
+
+
+def is_readonly_expression(node: expressions.Expression, allowed_operations: Optional[Set[str]] = None, enable_wildcard_check: bool = True):
+    """递归检查表达式树中是否出现写操作节点"""
+    # 写操作黑名单节点类型
+    write_types = {
+        expressions.Insert,
+        expressions.Update,
+        expressions.Delete,
+        expressions.Create,
+        expressions.Alter,
+        expressions.Drop,
+        expressions.Replace,
+        expressions.Merge
+    }
+    for descendant in node.walk():
+        if type(descendant) in write_types:
+            return False, type(descendant)
+
+    if allowed_operations:
+        op_type = get_operation_type(node)
+        # 检查操作类型
+        if op_type not in allowed_operations:
+            return False, op_type
+
+    if enable_wildcard_check and isinstance(node, expressions.Select) and node.find(expressions.Star):
+        return False, expressions.Star
 
-    forbidden = {"INSERT", "UPDATE", "DELETE", "CREATE", "ALTER", "DROP", "TRUNCATE",
-                 "LOAD", "REPLACE", "LOCK", "UNLOCK", "GRANT", "REVOKE", "EXEC", "CALL"}
-    allowed_root = {"SELECT", "WITH", "EXPLAIN", "DESCRIBE", "SHOW", "VALUES"}
+    return True, None
 
-    root = tokens[0].upper() if tokens else ""
-    if root not in allowed_root:
-        return False
-    if any(k in forbidden for k in tokens):
-        return False
-    return True
 
+def session_sql_execute(db_session: Session, sql_text: Union[str, List], format_result: bool = True):
+    def __do_execute(_sql: str):
+        _result = db_session.execute(text(f"{sql_text}"))
+        if format_result:
+            columns = list(_result.keys())
+            rows = [dict(zip(columns, r)) for r in _result.fetchall()]
+            return rows
+        else:
+            return _result
 
-def session_sql_execute(db_session: Session, sql_text: Union[str, List], query_only: bool = True):
     if isinstance(sql_text, str):
-        # 处理单条SQL语句
-        if query_only:
-            if not is_readonly_sql(sql_text):
-                return []
-        result = db_session.execute(text(f"{sql_text}"))
-        columns = list(result.keys())
-        rows = [dict(zip(columns, r)) for r in result.fetchall()]
-        return rows
+        return __do_execute(sql_text)
     elif isinstance(sql_text, list):
         results = []
         for sub_sql_text in sql_text:
-            if query_only:
-                if not is_readonly_sql(sub_sql_text):
-                    continue
-            result = db_session.execute(text(f"{sub_sql_text}"))
-            columns = list(result.keys())
-            rows = [dict(zip(columns, r)) for r in result.fetchall()]
-            results.append(rows)
+            result = __do_execute(sub_sql_text)
+            results.append(result)
             return results
     else:
         return []

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.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentic-kit-common
-Version: 0.0.9
+Version: 0.0.11
 Summary: Common utilities and tools for agentic kit ecosystem
 Home-page: 
 Author: manson
@@ -26,6 +26,7 @@ Requires-Dist: langchain_community
 Requires-Dist: langchain_experimental
 Requires-Dist: mysql-connector-python
 Requires-Dist: sqlalchemy
+Requires-Dist: sqlglot
 Requires-Dist: pymilvus
 Requires-Dist: xinference_client
 Dynamic: author

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

@@ -1,32 +1,34 @@
 agentic_kit_common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+agentic_kit_common/log/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+agentic_kit_common/log/logger.py,sha256=EH8pOW6KVpb-c4RdK2sohpHgfs3_hBcaAletRs1O23k,391
 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/execution.py,sha256=sCEcInFGFZ9ZiIs0eHKbQPtN4Z02v7vGdFMyVEWygcU,1951
+agentic_kit_common/orm/execution.py,sha256=beyRJqVGY5nOqMDpvLcfwimUzP6eSz8ayNHLwh1hGOo,2858
 agentic_kit_common/orm/manager.py,sha256=lWgFk5fUu_9m6yN_fskWHSYGaW30ty--KZKj8AuvIh0,4852
 agentic_kit_common/orm/multi_session.py,sha256=CAEnOLl0I-r77JAknKQ2sERF40OkHD7BOrYVYoHuQV4,8814
 agentic_kit_common/orm/schema.py,sha256=ukdVP71NE_JO5_HOe_FApJjBAKUsDluoElVn5vNSGJQ,1023
 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.11.dist-info/METADATA,sha256=i2EbkVNKrTVCcfwE9ky0Sfr64Lg8A97GsPOQzDCgKwM,7520
+agentic_kit_common-0.0.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+agentic_kit_common-0.0.11.dist-info/top_level.txt,sha256=nEKDlp84vqKSVWssGcxyuIsTqWLhMo45xqMs2GK4Dgg,24
+agentic_kit_common-0.0.11.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__':