infomankit 0.3.23__py3-none-any.whl

infoman/service/infrastructure/db_relation/mysql.py

@@ -0,0 +1,52 @@
+# database/backends.py
+"""
+数据库后端实现
+"""
+
+from typing import Dict, Any
+from infoman.config import DatabaseInstanceConfig
+
+
+class MySQLBackend:
+    """MySQL 后端"""
+
+    @staticmethod
+    def get_engine() -> str:
+        return "tortoise.backends.mysql"
+
+    @staticmethod
+    def get_credentials(config: DatabaseInstanceConfig) -> Dict[str, Any]:
+        """生成 MySQL 连接凭证"""
+        credentials = {
+            "host": config.host,
+            "port": config.port,
+            "user": config.user,
+            "password": config.password,
+            "database": config.database,
+            "charset": config.charset,
+            # 连接池
+            "minsize": config.pool_min_size,
+            "maxsize": config.pool_max_size,
+            "pool_recycle": config.pool_recycle,
+            "echo": config.echo,
+            # 超时
+            "connect_timeout": config.connect_timeout,
+            # 其他
+            "autocommit": True,
+            "init_command": "SET sql_mode='STRICT_TRANS_TABLES'",
+        }
+
+        # SSL
+        if config.ssl_enabled:
+            ssl_config = {}
+            if config.ssl_ca:
+                ssl_config["ca"] = config.ssl_ca
+            if config.ssl_cert:
+                ssl_config["cert"] = config.ssl_cert
+            if config.ssl_key:
+                ssl_config["key"] = config.ssl_key
+
+            if ssl_config:
+                credentials["ssl"] = ssl_config
+
+        return credentials

infoman/service/infrastructure/db_relation/pgsql.py

@@ -0,0 +1,54 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+
+"""
+# Time       ：2025/12/22 21:10
+# Author     ：Maxwell
+# Description：
+"""
+from typing import Dict, Any
+from infoman.config import DatabaseInstanceConfig
+
+
+class PostgreSQLBackend:
+    """PostgreSQL 后端"""
+
+    @staticmethod
+    def get_engine() -> str:
+        return "tortoise.backends.asyncpg"
+
+    @staticmethod
+    def get_credentials(config: DatabaseInstanceConfig) -> Dict[str, Any]:
+        """生成 PostgreSQL 连接凭证"""
+        credentials = {
+            "host": config.host,
+            "port": config.port,
+            "user": config.user,
+            "password": config.password,
+            "database": config.database,
+            # 连接池（注意：asyncpg 参数名不同）
+            "min_size": config.pool_min_size,
+            "max_size": config.pool_max_size,
+            "max_queries": 50000,
+            "max_inactive_connection_lifetime": config.pool_recycle,
+            # 超时
+            "timeout": config.connect_timeout,
+            "command_timeout": config.command_timeout,
+            # PostgreSQL 特有
+            "schema": config.schema_name,
+            "server_settings": {
+                "application_name": "infoman",
+                "jit": "off",
+            },
+        }
+
+        # SSL
+        if config.ssl_enabled:
+            import ssl
+
+            ssl_context = ssl.create_default_context()
+            if config.ssl_ca:
+                ssl_context.load_verify_locations(cafile=config.ssl_ca)
+            credentials["ssl"] = ssl_context
+
+        return credentials

infoman/service/infrastructure/db_relation/sqllite.py

@@ -0,0 +1,25 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+
+"""
+# Time       ：2025/12/22 21:11
+# Author     ：Maxwell
+# Description：
+"""
+from typing import Dict, Any
+from infoman.config import DatabaseConfig
+
+
+class SQLiteBackend:
+    """SQLite 后端"""
+
+    @staticmethod
+    def get_engine() -> str:
+        return "tortoise.backends.sqlite"
+
+    @staticmethod
+    def get_credentials(config: DatabaseConfig) -> Dict[str, Any]:
+        """生成 SQLite 连接凭证"""
+        return {
+            "file_path": f"{config.database}.db",
+        }

infoman/service/infrastructure/db_vector/__init__.py

@@ -0,0 +1,40 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+
+"""
+向量数据库基础设施
+
+支持:
+- Qdrant: 高性能向量搜索引擎
+- Milvus: 待实现
+
+Usage:
+    # 方式 1: 通过 lifespan 自动初始化
+    from infoman.service.app import application
+    # VectorDBManager 会在应用启动时自动初始化
+
+    # 方式 2: 手动使用
+    from infoman.service.infrastructure.db_vector import VectorDBManager
+
+    manager = VectorDBManager()
+    await manager.startup()
+
+    # 使用 Qdrant
+    client = manager.qdrant.client
+    await client.create_collection(...)
+
+    # 便捷方法
+    await manager.create_collection("my_collection", vector_size=768)
+    results = await manager.search("my_collection", query_vector=[...])
+
+    # 关闭
+    await manager.shutdown()
+"""
+
+from .manager import VectorDBManager
+from .qdrant import QdrantBackend
+
+__all__ = [
+    "VectorDBManager",
+    "QdrantBackend",
+]

infoman/service/infrastructure/db_vector/manager.py

@@ -0,0 +1,201 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+
+"""
+向量数据库管理器（支持延迟导入）
+
+支持：
+- Qdrant
+- Milvus (待实现)
+- 统一的管理接口
+"""
+
+from typing import Optional, Dict, Any, TYPE_CHECKING
+from fastapi import FastAPI
+from loguru import logger
+
+from infoman.config import settings
+
+if TYPE_CHECKING:
+    from infoman.service.infrastructure.db_vector.qdrant import QdrantBackend
+
+
+class VectorDBManager:
+    """向量数据库管理器"""
+
+    def __init__(self):
+        self.qdrant: Optional[Any] = None
+        self.initialized = False
+
+    @property
+    def is_available(self) -> bool:
+        """是否有可用的向量数据库"""
+        return (
+            (self.qdrant and self.qdrant.is_available) or
+            False  # 其他数据库
+        )
+
+    @property
+    def client(self):
+        """
+        获取默认客户端
+
+        优先级：Qdrant > Milvus
+        """
+        if self.qdrant and self.qdrant.is_available:
+            return self.qdrant.client
+
+        return None
+
+    async def startup(self, app: Optional[FastAPI] = None) -> bool:
+        if self.initialized:
+            logger.warning("⚠️ VectorDBManager 已初始化，跳过重复初始化")
+            return True
+
+        success_count = 0
+
+        # ========== 启动 Qdrant ==========
+        if settings.qdrant_configured:
+            # 延迟导入 QdrantBackend
+            try:
+                from infoman.service.infrastructure.db_vector.qdrant import QdrantBackend
+            except ImportError as e:
+                logger.error(f"❌ Qdrant 依赖未安装: {e}")
+                logger.error("请运行: pip install infomankit[vector]")
+                return False
+
+            self.qdrant = QdrantBackend(settings)
+
+            if await self.qdrant.startup():
+                success_count += 1
+
+                if app:
+                    app.state.qdrant_client = self.qdrant.client
+                    logger.debug("✅ Qdrant 客户端已挂载到 app.state")
+
+            # ========== 初始化完成 ==========
+            if success_count > 0:
+                self.initialized = True
+                logger.success(f"✅ 向量数据库初始化完成（{success_count} 个）")
+                return True
+        else:
+            logger.info("⏭️ 向量数据库未配置，跳过初始化")
+        return False
+
+    async def shutdown(self):
+        if not self.initialized:
+            return
+
+        logger.info("⏹️ 关闭向量数据库连接...")
+
+        if self.qdrant:
+            await self.qdrant.shutdown()
+
+        self.initialized = False
+        logger.success("✅ 所有向量数据库已关闭")
+
+    async def health_check(self) -> Dict[str, Any]:
+        """
+        健康检查
+
+        Returns:
+            {
+                "qdrant": {"status": "healthy", ...},
+                "milvus": {"status": "not_configured", ...},
+            }
+        """
+        health = {}
+
+        # Qdrant
+        if self.qdrant:
+            health["qdrant"] = await self.qdrant.health_check()
+        else:
+            health["qdrant"] = {
+                "status": "not_configured",
+                "name": "qdrant",
+                "details": {"enabled": False}
+            }
+
+        # Milvus
+        # if self.milvus:
+        #     health["milvus"] = await self.milvus.health_check()
+
+        return health
+
+    async def get_stats(self) -> Dict[str, Any]:
+        """
+        获取统计信息
+
+        Returns:
+            {
+                "qdrant": {...},
+                "milvus": {...},
+            }
+        """
+        stats = {}
+
+        if self.qdrant and self.qdrant.is_available:
+            stats["qdrant"] = await self.qdrant.get_stats()
+
+        # if self.milvus and self.milvus.is_available:
+        #     stats["milvus"] = await self.milvus.get_stats()
+
+        return stats
+
+    # ========== 便捷方法 ==========
+
+    async def create_collection(
+        self,
+        collection_name: str,
+        vector_size: int,
+        backend: str = "qdrant",
+        **kwargs
+    ) -> bool:
+        """
+        创建集合
+
+        Args:
+            collection_name: 集合名称
+            vector_size: 向量维度
+            backend: 使用的后端（qdrant/milvus）
+            **kwargs: 其他参数
+
+        Returns:
+            是否创建成功
+        """
+        if backend == "qdrant" and self.qdrant:
+            return await self.qdrant.create_collection(
+                collection_name, vector_size, **kwargs
+            )
+
+        logger.error(f"后端 {backend} 不可用")
+        return False
+
+    async def search(
+        self,
+        collection_name: str,
+        query_vector: list,
+        limit: int = 10,
+        backend: str = "qdrant",
+        **kwargs
+    ):
+        """
+        向量搜索
+
+        Args:
+            collection_name: 集合名称
+            query_vector: 查询向量
+            limit: 返回结果数量
+            backend: 使用的后端
+            **kwargs: 其他参数
+
+        Returns:
+            搜索结果
+        """
+        if backend == "qdrant" and self.qdrant:
+            return await self.qdrant.search(
+                collection_name, query_vector, limit, **kwargs
+            )
+
+        logger.error(f"后端 {backend} 不可用")
+        return []

infoman/service/infrastructure/db_vector/qdrant.py

@@ -0,0 +1,322 @@
+# !/usr/bin/env python
+# -*-coding:utf-8 -*-
+
+"""
+Qdrant 向量数据库后端
+
+功能：
+- 异步连接管理
+- 健康检查
+- 集合管理
+- 向量操作
+"""
+
+from typing import Optional, Dict, Any, List
+from qdrant_client import AsyncQdrantClient
+from qdrant_client.models import Distance, VectorParams, PointStruct
+from loguru import logger
+
+from infoman.config import VectorDBConfig
+from infoman.service.infrastructure.base import (
+    BaseInfrastructureComponent,
+    ComponentType,
+    ComponentStatus,
+)
+
+
+class QdrantBackend(BaseInfrastructureComponent):
+    """Qdrant 向量数据库后端"""
+
+    def __init__(self, config: VectorDBConfig):
+        super().__init__(
+            component_type=ComponentType.OTHER,  # 向量数据库可以归类为 OTHER
+            name="qdrant",
+            enabled=config.qdrant_configured,
+        )
+        self.config = config
+        self._client: Optional[AsyncQdrantClient] = None
+
+    async def startup(self) -> bool:
+        """启动 Qdrant 连接"""
+        if not self.enabled:
+            logger.info("⏭️ Qdrant 未启用，跳过初始化")
+            self._set_status(ComponentStatus.NOT_CONFIGURED)
+            return False
+
+        try:
+            self._set_status(ComponentStatus.INITIALIZING)
+            logger.info(f"🔌 正在连接 Qdrant [{self.config.QDRANT_HOST}]...")
+
+            # 创建 Qdrant 客户端
+            self._client = AsyncQdrantClient(
+                host=self.config.QDRANT_HOST,
+                port=self.config.QDRANT_HTTP_PORT,
+                grpc_port=self.config.QDRANT_GRPC_PORT,
+                api_key=self.config.QDRANT_API_KEY,
+                timeout=self.config.QDRANT_TIMEOUT,
+                prefer_grpc=True,  # 优先使用 gRPC（性能更好）
+            )
+
+            # 测试连接
+            await self._test_connection()
+
+            self._set_status(ComponentStatus.HEALTHY)
+            self._log_startup(
+                True,
+                f"{self.config.QDRANT_HOST}:{self.config.QDRANT_HTTP_PORT}"
+            )
+
+            return True
+
+        except Exception as e:
+            self._set_status(ComponentStatus.UNHEALTHY)
+            self._log_startup(False, str(e))
+            return False
+
+    async def shutdown(self):
+        """关闭 Qdrant 连接"""
+        if not self._client:
+            return
+
+        try:
+            logger.info("⏹️ 正在关闭 Qdrant 连接...")
+
+            # Qdrant 客户端会自动清理连接
+            # 如果需要显式关闭，可以添加逻辑
+            self._client = None
+
+            self._set_status(ComponentStatus.STOPPED)
+            self._log_shutdown(True)
+
+        except Exception as e:
+            self._log_shutdown(False, str(e))
+
+    async def health_check(self) -> Dict[str, Any]:
+        """健康检查"""
+        if not self.enabled:
+            return {
+                "status": "not_configured",
+                "component_type": self.component_type.value,
+                "name": self.name,
+                "details": {"enabled": False}
+            }
+
+        if not self._client:
+            return {
+                "status": "unhealthy",
+                "component_type": self.component_type.value,
+                "name": self.name,
+                "details": {"error": "客户端未初始化"}
+            }
+
+        try:
+            # 尝试获取集合列表
+            collections = await self._client.get_collections()
+
+            return {
+                "status": "healthy",
+                "component_type": self.component_type.value,
+                "name": self.name,
+                "details": {
+                    "host": self.config.QDRANT_HOST,
+                    "http_port": self.config.QDRANT_HTTP_PORT,
+                    "grpc_port": self.config.QDRANT_GRPC_PORT,
+                    "collections_count": len(collections.collections),
+                }
+            }
+
+        except Exception as e:
+            logger.error(f"Qdrant 健康检查失败: {e}")
+            self._set_status(ComponentStatus.UNHEALTHY)
+
+            return {
+                "status": "unhealthy",
+                "component_type": self.component_type.value,
+                "name": self.name,
+                "details": {"error": str(e)}
+            }
+
+    async def _test_connection(self):
+        """测试连接"""
+        try:
+            # 尝试获取集合列表
+            await self._client.get_collections()
+            logger.success("✅ Qdrant 连接测试成功")
+
+        except Exception as e:
+            logger.error(f"❌ Qdrant 连接测试失败: {e}")
+            raise
+
+    # ========== 便捷方法 ==========
+
+    async def create_collection(
+        self,
+        collection_name: str,
+        vector_size: int,
+        distance: Distance = Distance.COSINE,
+        **kwargs
+    ) -> bool:
+        """
+        创建集合
+
+        Args:
+            collection_name: 集合名称
+            vector_size: 向量维度
+            distance: 距离度量（COSINE/EUCLID/DOT）
+            **kwargs: 其他参数
+
+        Returns:
+            是否创建成功
+        """
+        if not self._client:
+            logger.error("Qdrant 客户端未初始化")
+            return False
+
+        try:
+            # 检查集合是否存在
+            collections = await self._client.get_collections()
+            collection_names = [c.name for c in collections.collections]
+
+            if collection_name in collection_names:
+                logger.warning(f"集合 {collection_name} 已存在")
+                return True
+
+            # 创建集合
+            await self._client.create_collection(
+                collection_name=collection_name,
+                vectors_config=VectorParams(
+                    size=vector_size,
+                    distance=distance,
+                ),
+                **kwargs
+            )
+
+            logger.success(f"✅ 创建集合成功: {collection_name}")
+            return True
+
+        except Exception as e:
+            logger.error(f"创建集合失败: {e}")
+            return False
+
+    async def delete_collection(self, collection_name: str) -> bool:
+        """删除集合"""
+        if not self._client:
+            return False
+
+        try:
+            await self._client.delete_collection(collection_name)
+            logger.success(f"✅ 删除集合成功: {collection_name}")
+            return True
+
+        except Exception as e:
+            logger.error(f"删除集合失败: {e}")
+            return False
+
+    async def upsert_points(
+        self,
+        collection_name: str,
+        points: List[PointStruct],
+    ) -> bool:
+        """插入或更新向量点"""
+        if not self._client:
+            return False
+
+        try:
+            await self._client.upsert(
+                collection_name=collection_name,
+                points=points,
+            )
+            logger.debug(f"✅ 插入/更新 {len(points)} 个点到 {collection_name}")
+            return True
+
+        except Exception as e:
+            logger.error(f"插入向量失败: {e}")
+            return False
+
+    async def search(
+        self,
+        collection_name: str,
+        query_vector: List[float],
+        limit: int = 10,
+        score_threshold: Optional[float] = None,
+        **kwargs
+    ) -> List[Dict[str, Any]]:
+        """
+        向量搜索
+
+        Args:
+            collection_name: 集合名称
+            query_vector: 查询向量
+            limit: 返回结果数量
+            score_threshold: 相似度阈值
+            **kwargs: 其他搜索参数
+
+        Returns:
+            搜索结果列表
+        """
+        if not self._client:
+            return []
+
+        try:
+            results = await self._client.search(
+                collection_name=collection_name,
+                query_vector=query_vector,
+                limit=limit,
+                score_threshold=score_threshold,
+                **kwargs
+            )
+
+            return [
+                {
+                    "id": r.id,
+                    "score": r.score,
+                    "payload": r.payload,
+                }
+                for r in results
+            ]
+
+        except Exception as e:
+            logger.error(f"向量搜索失败: {e}")
+            return []
+
+    async def get_collection_info(self, collection_name: str) -> Optional[Dict[str, Any]]:
+        """获取集合信息"""
+        if not self._client:
+            return None
+
+        try:
+            info = await self._client.get_collection(collection_name)
+            return {
+                "name": collection_name,
+                "vectors_count": info.vectors_count,
+                "points_count": info.points_count,
+                "status": info.status,
+            }
+
+        except Exception as e:
+            logger.error(f"获取集合信息失败: {e}")
+            return None
+
+    async def get_stats(self) -> Dict[str, Any]:
+        """获取统计信息"""
+        if not self._client or not self.is_available:
+            return {}
+
+        try:
+            collections = await self._client.get_collections()
+
+            stats = {
+                "collections_count": len(collections.collections),
+                "collections": []
+            }
+
+            for collection in collections.collections:
+                info = await self.get_collection_info(collection.name)
+                if info:
+                    stats["collections"].append(info)
+
+            return stats
+
+        except Exception as e:
+            logger.error(f"获取统计信息失败: {e}")
+            return {}