agentic-kit-common 0.0.9__tar.gz → 0.0.11__tar.gz

{agentic_kit_common-0.0.9 → agentic_kit_common-0.0.11}/PKG-INFO

@@ -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.11/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-0.0.11/agentic_kit_common/orm/execution.py

@@ -0,0 +1,90 @@
+from typing import List, Union, Optional, Set
+
+import sqlglot
+from sqlalchemy import text
+from sqlalchemy.orm import Session
+from sqlglot import expressions
+
+
+def get_operation_type(statement) -> str:
+    """获取SQL操作类型"""
+    return type(statement).__name__.upper()
+
+
+def get_sql_operation_info(sql: str) -> Optional[dict]:
+    """
+    获取SQL操作信息（用于调试和日志记录）
+
+    Returns:
+        dict: 包含操作类型、涉及表等信息
+    """
+    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
+
+    return True, None
+
+
+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
+
+    if isinstance(sql_text, str):
+        return __do_execute(sql_text)
+    elif isinstance(sql_text, list):
+        results = []
+        for sub_sql_text in sql_text:
+            result = __do_execute(sub_sql_text)
+            results.append(result)
+            return results
+    else:
+        return []

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

@@ -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 → agentic_kit_common-0.0.11}/agentic_kit_common.egg-info/SOURCES.txt

@@ -6,6 +6,8 @@ agentic_kit_common.egg-info/SOURCES.txt
 agentic_kit_common.egg-info/dependency_links.txt
 agentic_kit_common.egg-info/requires.txt
 agentic_kit_common.egg-info/top_level.txt
+agentic_kit_common/log/__init__.py
+agentic_kit_common/log/logger.py
 agentic_kit_common/minio/__init__.py
 agentic_kit_common/minio/minio_manager.py
 agentic_kit_common/orm/__init__.py

{agentic_kit_common-0.0.9 → agentic_kit_common-0.0.11}/agentic_kit_common.egg-info/requires.txt

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

{agentic_kit_common-0.0.9 → agentic_kit_common-0.0.11}/setup.py

@@ -4,7 +4,7 @@ from setuptools import setup, find_packages
 
 setup(
     name='agentic-kit-common',
-    version="0.0.9",
+    version="0.0.11",
     author="manson",
     author_email="manson.li3307@gmail.com",
     description='Common utilities and tools for agentic kit ecosystem',
@@ -42,6 +42,9 @@ setup(
         "mysql-connector-python",
         "sqlalchemy",
 
+        # orm
+        "sqlglot",
+
         # milvus
         "pymilvus",
 

{agentic_kit_common-0.0.9 → agentic_kit_common-0.0.11}/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)
 

{agentic_kit_common-0.0.9 → agentic_kit_common-0.0.11}/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__':

agentic_kit_common-0.0.9/agentic_kit_common/orm/execution.py

@@ -1,53 +0,0 @@
-from typing import List, Union
-
-import sqlparse
-from sqlalchemy import text
-from sqlalchemy.orm import Session
-
-
-def is_readonly_sql(sql: str) -> bool:
-    """
-    只允许：
-      SELECT / WITH / VALUES / EXPLAIN / DESCRIBE / SHOW
-    拒绝：
-      INSERT/UPDATE/DELETE/CREATE/ALTER/DROP/TRUNCATE/LOAD/REPLACE/LOCK/UNLOCK/GRANT/REVOKE/EXECUTE/CALL
-    """
-    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)]
-
-    forbidden = {"INSERT", "UPDATE", "DELETE", "CREATE", "ALTER", "DROP", "TRUNCATE",
-                 "LOAD", "REPLACE", "LOCK", "UNLOCK", "GRANT", "REVOKE", "EXEC", "CALL"}
-    allowed_root = {"SELECT", "WITH", "EXPLAIN", "DESCRIBE", "SHOW", "VALUES"}
-
-    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], 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
-    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)
-            return results
-    else:
-        return []