Agent/knowhub/vector_store.py

"""
Milvus Lite 存储封装

单一存储架构，存储完整知识数据 + 向量。
"""

from milvus import default_server
from pymilvus import (
    connections, Collection, FieldSchema,
    CollectionSchema, DataType, utility
)
from typing import List, Dict, Optional
import json
import time


class MilvusStore:
    def __init__(self, data_dir: str = "./milvus_data"):
        """
        初始化 Milvus Lite 存储

        Args:
            data_dir: 数据存储目录
        """
        # 启动内嵌服务器
        default_server.set_base_dir(data_dir)

        # 检查是否已经有 Milvus 实例在运行
        try:
            # 尝试连接到可能已存在的实例
            connections.connect(
                alias="default",
                host='127.0.0.1',
                port=default_server.listen_port,
                timeout=5
            )
            print(f"[Milvus] 连接到已存在的 Milvus 实例 (端口 {default_server.listen_port})")
        except Exception:
            # 没有运行的实例，启动新的
            print(f"[Milvus] 启动新的 Milvus Lite 实例...")
            try:
                default_server.start()
                print(f"[Milvus] Milvus Lite 启动成功 (端口 {default_server.listen_port})")
                # 启动后建立连接
                connections.connect(
                    alias="default",
                    host='127.0.0.1',
                    port=default_server.listen_port,
                    timeout=5
                )
                print(f"[Milvus] 已连接到新启动的实例")
            except Exception as e:
                print(f"[Milvus] 启动失败: {e}")
                # 尝试连接到可能已经在运行的实例
                try:
                    connections.connect(
                        alias="default",
                        host='127.0.0.1',
                        port=default_server.listen_port,
                        timeout=5
                    )
                    print(f"[Milvus] 连接到已存在的实例")
                except Exception as e2:
                    raise RuntimeError(f"无法启动或连接到 Milvus: {e}, {e2}")

        self._init_collection()

    def _init_collection(self):
        """初始化 collection"""
        collection_name = "knowledge"

        if utility.has_collection(collection_name):
            self.collection = Collection(collection_name)
        else:
            # 定义 schema
            fields = [
                FieldSchema(name="id", dtype=DataType.VARCHAR,
                           max_length=100, is_primary=True),
                FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR,
                           dim=1536),
                FieldSchema(name="message_id", dtype=DataType.VARCHAR,
                           max_length=100),
                FieldSchema(name="task", dtype=DataType.VARCHAR,
                           max_length=2000),
                FieldSchema(name="content", dtype=DataType.VARCHAR,
                           max_length=50000),
                FieldSchema(name="types", dtype=DataType.ARRAY,
                           element_type=DataType.VARCHAR, max_capacity=20, max_length=50),
                FieldSchema(name="tags", dtype=DataType.JSON),
                FieldSchema(name="tag_keys", dtype=DataType.ARRAY,
                           element_type=DataType.VARCHAR, max_capacity=50, max_length=100),
                FieldSchema(name="scopes", dtype=DataType.ARRAY,
                           element_type=DataType.VARCHAR, max_capacity=20, max_length=100),
                FieldSchema(name="owner", dtype=DataType.VARCHAR,
                           max_length=200),
                FieldSchema(name="resource_ids", dtype=DataType.ARRAY,
                           element_type=DataType.VARCHAR, max_capacity=50, max_length=200),
                FieldSchema(name="source", dtype=DataType.JSON),
                FieldSchema(name="eval", dtype=DataType.JSON),
                FieldSchema(name="created_at", dtype=DataType.INT64),
                FieldSchema(name="updated_at", dtype=DataType.INT64),
                FieldSchema(name="status", dtype=DataType.VARCHAR,
                           max_length=20, default_value="approved"),
                FieldSchema(name="relationships", dtype=DataType.VARCHAR,
                           max_length=10000, default_value="[]"),
            ]

            schema = CollectionSchema(fields, description="KnowHub Knowledge")
            self.collection = Collection(collection_name, schema)

            # 创建向量索引
            index_params = {
                "metric_type": "COSINE",
                "index_type": "HNSW",
                "params": {"M": 16, "efConstruction": 200}
            }
            self.collection.create_index("embedding", index_params)

            # 为 status 创建 Trie 标量索引（加速过滤）
            try:
                self.collection.create_index("status", {"index_type": "Trie"})
            except Exception:
                pass

        self.collection.load()

    def insert(self, knowledge: Dict):
        """
        插入单条知识

        Args:
            knowledge: 知识数据（包含 embedding）
        """
        self.collection.insert([knowledge])
        self.collection.flush()

    def insert_batch(self, knowledge_list: List[Dict]):
        """
        批量插入知识

        Args:
            knowledge_list: 知识列表
        """
        if not knowledge_list:
            return
        self.collection.insert(knowledge_list)
        self.collection.flush()

    def search(self,
               query_embedding: List[float],
               filters: Optional[str] = None,
               limit: int = 10) -> List[Dict]:
        """
        向量检索 + 标量过滤

        Args:
            query_embedding: 查询向量
            filters: 过滤表达式（如: 'owner == "agent"'）
            limit: 返回数量

        Returns:
            知识列表
        """
        search_params = {"metric_type": "COSINE", "params": {"ef": 100}}

        results = self.collection.search(
            data=[query_embedding],
            anns_field="embedding",
            param=search_params,
            limit=limit,
            expr=filters,
            output_fields=["id", "message_id", "task", "content", "types",
                          "tags", "tag_keys", "scopes", "owner", "resource_ids",
                          "source", "eval", "created_at", "updated_at",
                          "status", "relationships"]
        )

        if not results or not results[0]:
            return []

        # 返回实体字典，包含所有字段
        # 注意：时间戳需要转换为毫秒（JavaScript Date 需要）
        return [
            {
                "id": hit.entity.get("id"),
                "message_id": hit.entity.get("message_id"),
                "task": hit.entity.get("task"),
                "content": hit.entity.get("content"),
                "types": list(hit.entity.get("types")) if hit.entity.get("types") else [],
                "tags": hit.entity.get("tags"),
                "tag_keys": list(hit.entity.get("tag_keys")) if hit.entity.get("tag_keys") else [],
                "scopes": list(hit.entity.get("scopes")) if hit.entity.get("scopes") else [],
                "owner": hit.entity.get("owner"),
                "resource_ids": list(hit.entity.get("resource_ids")) if hit.entity.get("resource_ids") else [],
                "source": hit.entity.get("source"),
                "eval": hit.entity.get("eval"),
                "created_at": hit.entity.get("created_at") * 1000 if hit.entity.get("created_at") else None,
                "updated_at": hit.entity.get("updated_at") * 1000 if hit.entity.get("updated_at") else None,
                "status": hit.entity.get("status", "approved"),
                "relationships": json.loads(hit.entity.get("relationships") or "[]"),
            }
            for hit in results[0]
        ]

    def query(self, filters: str, limit: int = 100) -> List[Dict]:
        """
        纯标量查询（不使用向量）

        Args:
            filters: 过滤表达式
            limit: 返回数量

        Returns:
            知识列表
        """
        results = self.collection.query(
            expr=filters,
            output_fields=["id", "message_id", "task", "content", "types",
                          "tags", "tag_keys", "scopes", "owner", "resource_ids",
                          "source", "eval", "created_at", "updated_at",
                          "status", "relationships"],
            limit=limit
        )

        # 转换时间戳为毫秒，确保数组字段格式正确
        for r in results:
            if r.get("created_at"):
                r["created_at"] = r["created_at"] * 1000
            if r.get("updated_at"):
                r["updated_at"] = r["updated_at"] * 1000
            # 确保数组字段是列表格式
            if r.get("types") and not isinstance(r["types"], list):
                r["types"] = list(r["types"])
            if r.get("tag_keys") and not isinstance(r["tag_keys"], list):
                r["tag_keys"] = list(r["tag_keys"])
            if r.get("scopes") and not isinstance(r["scopes"], list):
                r["scopes"] = list(r["scopes"])
            if r.get("resource_ids") and not isinstance(r["resource_ids"], list):
                r["resource_ids"] = list(r["resource_ids"])
            # 兼容旧数据（无 status/relationships 字段）
            if "status" not in r:
                r["status"] = "approved"
            if "relationships" not in r or r["relationships"] is None:
                r["relationships"] = []
            else:
                r["relationships"] = json.loads(r["relationships"]) if isinstance(r["relationships"], str) else r["relationships"]

        return results

    def get_by_id(self, knowledge_id: str) -> Optional[Dict]:
        """
        根据 ID 获取知识

        Args:
            knowledge_id: 知识 ID

        Returns:
            知识数据，不存在返回 None
        """
        results = self.collection.query(
            expr=f'id == "{knowledge_id}"',
            output_fields=["id", "embedding", "message_id", "task", "content", "types",
                          "tags", "tag_keys", "scopes", "owner", "resource_ids",
                          "source", "eval", "created_at", "updated_at",
                          "status", "relationships"]
        )

        if not results:
            return None

        # 转换时间戳和数组字段
        r = results[0]
        if r.get("created_at"):
            r["created_at"] = r["created_at"] * 1000
        if r.get("updated_at"):
            r["updated_at"] = r["updated_at"] * 1000
        if r.get("types") and not isinstance(r["types"], list):
            r["types"] = list(r["types"])
        if r.get("tag_keys") and not isinstance(r["tag_keys"], list):
            r["tag_keys"] = list(r["tag_keys"])
        if r.get("scopes") and not isinstance(r["scopes"], list):
            r["scopes"] = list(r["scopes"])
        if r.get("resource_ids") and not isinstance(r["resource_ids"], list):
            r["resource_ids"] = list(r["resource_ids"])
        # 兼容旧数据
        if "status" not in r:
            r["status"] = "approved"
        if "relationships" not in r or r["relationships"] is None:
            r["relationships"] = []
        else:
            r["relationships"] = json.loads(r["relationships"]) if isinstance(r["relationships"], str) else r["relationships"]

        return r

    def update(self, knowledge_id: str, updates: Dict):
        """
        更新知识（先删除再插入）

        Args:
            knowledge_id: 知识 ID
            updates: 更新字段
        """
        # 1. 查询现有数据
        existing = self.get_by_id(knowledge_id)
        if not existing:
            raise ValueError(f"Knowledge not found: {knowledge_id}")

        # 2. 合并更新
        existing.update(updates)
        existing["updated_at"] = int(time.time())

        # 3. 还原 get_by_id 的展示层转换，确保存储格式正确
        # created_at 被 get_by_id 乘以 1000（毫秒），需还原为秒
        if existing.get("created_at") and existing["created_at"] > 1_000_000_000_000:
            existing["created_at"] = existing["created_at"] // 1000
        # relationships 被 get_by_id 反序列化为 list，需还原为 JSON 字符串
        if isinstance(existing.get("relationships"), list):
            existing["relationships"] = json.dumps(existing["relationships"])

        # 4. 删除旧数据
        self.delete(knowledge_id)

        # 5. 插入新数据
        self.insert(existing)

    def delete(self, knowledge_id: str):
        """
        删除知识

        Args:
            knowledge_id: 知识 ID
        """
        self.collection.delete(f'id == "{knowledge_id}"')
        self.collection.flush()

    def count(self) -> int:
        """返回知识总数"""
        return self.collection.num_entities

    def drop_collection(self):
        """删除 collection（危险操作）"""
        utility.drop_collection("knowledge")

    def migrate_schema(self) -> int:
        """
        将旧 collection（无 status/relationships 字段）迁移到新 schema。
        采用中转 collection 模式，Step 3 之前数据始终有两份副本。
        返回迁移的知识条数。
        """
        MIGRATION_NAME = "knowledge_migration"
        MAIN_NAME = "knowledge"

        # 如果中转 collection 已存在（上次迁移中途失败），先清理
        if utility.has_collection(MIGRATION_NAME):
            print(f"[Migrate] 检测到残留中转 collection，清理...")
            utility.drop_collection(MIGRATION_NAME)

        # Step 1: 创建中转 collection（新 schema）
        print(f"[Migrate] Step 1: 创建中转 collection {MIGRATION_NAME}...")
        fields = [
            FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=100, is_primary=True),
            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=1536),
            FieldSchema(name="message_id", dtype=DataType.VARCHAR, max_length=100),
            FieldSchema(name="task", dtype=DataType.VARCHAR, max_length=2000),
            FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=50000),
            FieldSchema(name="types", dtype=DataType.ARRAY, element_type=DataType.VARCHAR, max_capacity=20, max_length=50),
            FieldSchema(name="tags", dtype=DataType.JSON),
            FieldSchema(name="tag_keys", dtype=DataType.ARRAY, element_type=DataType.VARCHAR, max_capacity=50, max_length=100),
            FieldSchema(name="scopes", dtype=DataType.ARRAY, element_type=DataType.VARCHAR, max_capacity=20, max_length=100),
            FieldSchema(name="owner", dtype=DataType.VARCHAR, max_length=200),
            FieldSchema(name="resource_ids", dtype=DataType.ARRAY, element_type=DataType.VARCHAR, max_capacity=50, max_length=200),
            FieldSchema(name="source", dtype=DataType.JSON),
            FieldSchema(name="eval", dtype=DataType.JSON),
            FieldSchema(name="created_at", dtype=DataType.INT64),
            FieldSchema(name="updated_at", dtype=DataType.INT64),
            FieldSchema(name="status", dtype=DataType.VARCHAR, max_length=20, default_value="approved"),
            FieldSchema(name="relationships", dtype=DataType.VARCHAR, max_length=10000, default_value="[]"),
        ]
        schema = CollectionSchema(fields, description="KnowHub Knowledge")
        migration_col = Collection(MIGRATION_NAME, schema)
        migration_col.create_index("embedding", {"metric_type": "COSINE", "index_type": "HNSW", "params": {"M": 16, "efConstruction": 200}})
        try:
            migration_col.create_index("status", {"index_type": "Trie"})
        except Exception:
            pass
        migration_col.load()

        # Step 2: 从旧 collection 逐批读取，补字段，插入中转
        print(f"[Migrate] Step 2: 读取旧数据并插入中转 collection...")
        batch_size = 200
        offset = 0
        total = 0
        while True:
            batch = self.collection.query(
                expr='id != ""',
                output_fields=["id", "embedding", "message_id", "task", "content", "types",
                               "tags", "tag_keys", "scopes", "owner", "resource_ids",
                               "source", "eval", "created_at", "updated_at"],
                limit=batch_size,
                offset=offset
            )
            if not batch:
                break
            for item in batch:
                item["status"] = item.get("status", "approved")
                item["relationships"] = item.get("relationships") or []
                # 时间戳已是秒级（query 返回原始值，未乘 1000）
            migration_col.insert(batch)
            migration_col.flush()
            total += len(batch)
            offset += len(batch)
            print(f"[Migrate] 已迁移 {total} 条...")
            if len(batch) < batch_size:
                break

        # Step 3: drop 旧 collection
        print(f"[Migrate] Step 3: drop 旧 collection {MAIN_NAME}...")
        self.collection.release()
        utility.drop_collection(MAIN_NAME)

        # Step 4: 创建新 collection（同名，新 schema）
        print(f"[Migrate] Step 4: 创建新 collection {MAIN_NAME}...")
        new_col = Collection(MAIN_NAME, schema)
        new_col.create_index("embedding", {"metric_type": "COSINE", "index_type": "HNSW", "params": {"M": 16, "efConstruction": 200}})
        try:
            new_col.create_index("status", {"index_type": "Trie"})
        except Exception:
            pass
        new_col.load()

        # Step 5: 从中转 collection 读取，插入新 collection
        print(f"[Migrate] Step 5: 从中转 collection 回写到新 collection...")
        offset = 0
        while True:
            batch = migration_col.query(
                expr='id != ""',
                output_fields=["id", "embedding", "message_id", "task", "content", "types",
                               "tags", "tag_keys", "scopes", "owner", "resource_ids",
                               "source", "eval", "created_at", "updated_at",
                               "status", "relationships"],
                limit=batch_size,
                offset=offset
            )
            if not batch:
                break
            new_col.insert(batch)
            new_col.flush()
            offset += len(batch)
            if len(batch) < batch_size:
                break

        # Step 6: drop 中转 collection
        print(f"[Migrate] Step 6: drop 中转 collection {MIGRATION_NAME}...")
        migration_col.release()
        utility.drop_collection(MIGRATION_NAME)

        # Step 7: 更新 self.collection 引用
        print(f"[Migrate] Step 7: 更新 collection 引用...")
        self.collection = new_col

        print(f"[Migrate] 迁移完成，共迁移 {total} 条知识。")
        return total