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

agentic_kit_common/orm/execution.py

@@ -0,0 +1,53 @@
+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 []

agentic_kit_common/orm/multi_session.py

@@ -0,0 +1,245 @@
+import logging
+from contextlib import contextmanager
+from threading import Lock
+from typing import Dict, Optional, Any, List
+
+from sqlalchemy import create_engine, text
+from sqlalchemy.exc import OperationalError
+from sqlalchemy.orm import sessionmaker, scoped_session
+
+from .schema import DatabaseEngineModel, _DEFAULT_CONFIG
+
+# 配置日志
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class DatabaseEngineManager:
+    """数据库引擎管理器"""
+
+    def __init__(self, init_engines: Optional[List[DatabaseEngineModel]] = None):
+        self._engines: Dict[str, Any] = {}
+        self._sessions: Dict[str, Any] = {}
+        self._engine_configs: Dict[str, Dict] = {}
+        self._default_engine_name = ""
+        self._lock = Lock()
+
+        # 初始化默认配置
+        self._default_config = _DEFAULT_CONFIG.model_dump()
+
+        # 初始化引擎
+        if init_engines:
+            self._init_engines(init_engines=init_engines)
+
+    def _init_engines(self, init_engines: List[DatabaseEngineModel]):
+        for engine in init_engines:
+            self.add_engine(engine)
+
+    def _create_engine(self, database_uri: str, database_name: str, config: Optional[Dict] = None) -> Any:
+        """创建数据库引擎"""
+        if config is None:
+            config = self._default_config.copy()
+
+        # 构建完整的数据库URL
+        if database_uri.endswith('/'):
+            full_url = f"{database_uri}{database_name}"
+        else:
+            full_url = f"{database_uri}/{database_name}"
+
+        try:
+            engine = create_engine(
+                url=full_url,
+                echo=config.get("echo", False),
+                pool_size=config.get("pool_size", 10),
+                max_overflow=config.get("max_overflow", 5),
+                pool_pre_ping=config.get("pool_pre_ping", True),
+                pool_recycle=config.get("pool_recycle", 1800),
+            )
+            logger.info(f"成功创建引擎: {database_name} -> {database_uri}")
+            return engine
+        except Exception as e:
+            logger.error(f"创建引擎失败: {database_name}, 错误: {e}")
+            raise
+
+    def add_engine(self, engine_info: DatabaseEngineModel) -> bool:
+        """添加数据库引擎"""
+        with self._lock:
+            engine_name = engine_info.engine_name
+            database_uri = engine_info.database_uri
+            database_name = engine_info.database_name
+            config = engine_info.config.model_dump().copy()
+
+            if engine_name in self._engines:
+                logger.warning(f"引擎 '{engine_name}' 已存在")
+                return False
+
+            try:
+                # 创建引擎
+                engine = self._create_engine(database_uri, database_name, config)
+
+                # 创建会话工厂
+                session_factory = scoped_session(
+                    sessionmaker(
+                        bind=engine,
+                        expire_on_commit=False,
+                        autocommit=False,
+                        autoflush=False
+                    )
+                )
+
+                # 存储配置和实例
+                self._engines[engine_name] = engine
+                self._sessions[engine_name] = session_factory
+                self._engine_configs[engine_name] = {
+                    "database_uri": database_uri,
+                    "database_name": database_name,
+                    "config": config or self._default_config.copy()
+                }
+
+                logger.info(f"成功添加引擎: {engine_name}")
+                return True
+
+            except Exception as e:
+                logger.error(f"添加引擎失败: {engine_name}, 错误: {e}")
+                return False
+
+    def remove_engine(self, engine_name: str) -> bool:
+        """移除数据库引擎"""
+        with self._lock:
+            if engine_name not in self._engines:
+                logger.warning(f"引擎 '{engine_name}' 不存在")
+                return False
+
+            try:
+                # 关闭所有会话
+                if engine_name in self._sessions:
+                    self._sessions[engine_name].close_all()
+                    del self._sessions[engine_name]
+
+                # 处置引擎
+                if engine_name in self._engines:
+                    self._engines[engine_name].dispose()
+                    del self._engines[engine_name]
+
+                # 移除配置
+                if engine_name in self._engine_configs:
+                    del self._engine_configs[engine_name]
+
+                # 如果移除的是默认引擎
+                if engine_name == self._default_engine_name and self._engines:
+                    self._default_engine_name = ''
+
+                logger.info(f"成功移除引擎: {engine_name}")
+                return True
+
+            except Exception as e:
+                logger.error(f"移除引擎失败: {engine_name}, 错误: {e}")
+                return False
+
+    def get_engine(self, engine_name: Optional[str] = None) -> Any:
+        """获取数据库引擎"""
+        if engine_name is None:
+            engine_name = self._default_engine_name
+
+        if engine_name not in self._engines:
+            raise ValueError(f"数据库引擎 '{engine_name}' 不存在")
+
+        return self._engines[engine_name]
+
+    def get_session_factory(self, engine_name: Optional[str] = None) -> Any:
+        """获取会话工厂"""
+        if engine_name is None:
+            engine_name = self._default_engine_name
+
+        if engine_name not in self._sessions:
+            raise ValueError(f"数据库引擎 '{engine_name}' 不存在")
+
+        return self._sessions[engine_name]
+
+    @contextmanager
+    def get_db_session(self,
+                       engine_name: Optional[str] = None,
+                       auto_commit_by_exit: bool = False,
+                       auto_close: bool = True):
+        """获取数据库会话（上下文管理器）"""
+        if engine_name is None:
+            engine_name = self._default_engine_name
+
+        if engine_name not in self._sessions:
+            raise ValueError(f"数据库引擎 '{engine_name}' 不存在")
+
+        session_factory = self._sessions[engine_name]
+        session = session_factory()
+
+        try:
+            yield session
+            if auto_commit_by_exit:
+                session.commit()
+        except OperationalError as e:
+            logger.warning(f"数据库连接异常，尝试重新连接: {e}")
+            session.rollback()
+            # 重新创建会话
+            session.close()
+            session = session_factory()
+            yield session
+            if auto_commit_by_exit:
+                session.commit()
+        except Exception as e:
+            session.rollback()
+            logger.error(f"数据库操作异常: {e}")
+            raise e
+        finally:
+            if auto_close:
+                session.close()
+
+    def set_default_engine(self, engine_name: str) -> bool:
+        """设置默认引擎"""
+        with self._lock:
+            if engine_name not in self._engines:
+                logger.warning(f"引擎 '{engine_name}' 不存在")
+                return False
+
+            self._default_engine_name = engine_name
+            logger.info(f"设置默认引擎为: {engine_name}")
+            return True
+
+    def list_engines(self) -> List[str]:
+        """列出所有引擎名称"""
+        return list(self._engines.keys())
+
+    def get_engine_info(self, engine_name: str) -> Optional[Dict]:
+        """获取引擎信息"""
+        if engine_name not in self._engine_configs:
+            return None
+
+        info = self._engine_configs[engine_name].copy()
+        info["is_default"] = (engine_name == self._default_engine_name)
+        return info
+
+    def health_check(self, engine_name: Optional[str] = None) -> Dict[str, bool]:
+        """健康检查"""
+        results = {}
+
+        if engine_name:
+            engines_to_check = [engine_name] if engine_name in self._engines else []
+        else:
+            engines_to_check = list(self._engines.keys())
+
+        for name in engines_to_check:
+            try:
+                with self.get_db_session(name, auto_close=True) as session:
+                    session.execute(text("SELECT 1"))
+                results[name] = True
+                logger.debug(f"健康检查通过: {name}")
+            except Exception as e:
+                results[name] = False
+                logger.error(f"健康检查失败: {name}, 错误: {e}")
+
+        return results
+
+
+# 初始化数据库引擎管理器
+def initialize_database_engine_manager(init_engines: Optional[List[DatabaseEngineModel]] = None):
+    """初始化默认数据库"""
+    manager = DatabaseEngineManager(init_engines=init_engines)
+    return manager

agentic_kit_common/orm/schema.py

@@ -0,0 +1,19 @@
+from pydantic import BaseModel, Field
+
+
+class DatabaseEngineConfigModel(BaseModel):
+    echo: bool = Field(False, description="数据库的 SQL 语句 + 参数 打印到标准输出")
+    pool_pre_ping: bool = Field(True, description="每次从池里拿连接前，先发一句做“心跳”")
+    pool_size: int = Field(10, description="连接池里 长期保持的“永久”连接 数量")
+    max_overflow: int = Field(5, description="当 pool_size 用光后，最多还能再新建多少条“临时”连接")
+    pool_recycle: int = Field(1800, description="一条连接被 复用多久之后强制回收（关闭并新建）")
+
+
+_DEFAULT_CONFIG = DatabaseEngineConfigModel()
+
+
+class DatabaseEngineModel(BaseModel):
+    engine_name: str = Field(..., description="数据库引擎名字")
+    database_uri: str = Field(..., description="数据库uir地址")
+    database_name: str = Field(..., description="数据库名称")
+    config: DatabaseEngineConfigModel = Field(_DEFAULT_CONFIG, description="配置信息")

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.8.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.8
+Version: 0.0.10
 Summary: Common utilities and tools for agentic kit ecosystem
 Home-page: 
 Author: manson

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

@@ -3,26 +3,30 @@ agentic_kit_common/minio/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZ
 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/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_vector.py,sha256=4sJ7bXt5iRhc6j1Ig0bZ1j0huPU-clnsnT7O0-uFbkw,1311
-agentic_kit_common-0.0.8.dist-info/METADATA,sha256=mLi0procHOdh1DiCFCstlsD1RM-wAMWEKFpHTRMbLxo,7496
-agentic_kit_common-0.0.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-agentic_kit_common-0.0.8.dist-info/top_level.txt,sha256=nEKDlp84vqKSVWssGcxyuIsTqWLhMo45xqMs2GK4Dgg,24
-agentic_kit_common-0.0.8.dist-info/RECORD,,
+test/test_orm.py,sha256=8fGCU7BWaD5sDbg0fgYN0Saf_hi7t-q8svHCHLQDceo,1620
+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_orm.py

@@ -0,0 +1,47 @@
+import sys
+import unittest
+from pathlib import Path
+
+from agentic_kit_common.orm.execution import session_sql_execute
+from agentic_kit_common.orm.multi_session import initialize_database_engine_manager
+from agentic_kit_common.orm.schema import DatabaseEngineModel
+
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root))
+
+
+class MyTestCase(unittest.TestCase):
+    def test_multi_session(self):
+        engine = DatabaseEngineModel(
+            engine_name='default',
+            database_name='czailab_llm',
+            database_uri='mysql+mysqlconnector://ailab_dev:Qwert!%40%234@45.120.102.236'
+        )
+        # manager = initialize_database_engine_manager(init_engines=[engine])
+
+        manager = initialize_database_engine_manager()
+        manager.add_engine(engine_info=engine)
+
+        print(manager.get_engine('default'))
+        print(manager.get_engine_info('default'))
+        print(manager.list_engines())
+
+        # session = manager.get_db_session(engine_name='default')
+        with manager.get_db_session(engine_name='default') as session:
+            print(session)
+            print(type(session))
+
+            result = session_sql_execute(db_session=session, sql_text='select * from tenant; select * from rerank_model')
+            print(result)
+            print(type(result))
+            for item in result:
+                print(item)
+                print(item.keys())
+                print(item.values())
+                print(item.items())
+                for _item in item.items():
+                    print(_item)
+
+
+if __name__ == '__main__':
+    unittest.main()

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__':