from flowllm.core.context import C
from flowllm.core.op import BaseAsyncOp
from flowllm.core.schema import VectorNode
from flowllm.main import FlowLLMApp
import asyncio


@C.register_op()
class RAGSearchOp(BaseAsyncOp):
    """RAG 搜索 Op，展示 embedding、vectorstore 和并发能力"""

    async def async_execute(self):
        """执行 RAG 搜索逻辑"""
        # 1. 读取输入
        query = self.context.get("query", "")
        workspace_id = self.context.get("workspace_id", "default")
        top_k = self.context.get("top_k", 3)

        assert query, "query 不能为空"

        # 2. 使用 Embedding 生成查询向量（一行代码示意 emb）
        query_vector = await self.embedding_model.get_embeddings_async(query)

        # 3. 使用 VectorStore 进行语义搜索（一行代码示意 vectorstore）
        nodes = await self.vector_store.async_search(query=query, workspace_id=workspace_id, top_k=top_k)

        # 4. 并发处理多个查询（两行代码示意并发）
        # 提交多个并发任务
        for node in nodes:
            self.submit_async_task(self._process_node, node)
        # 等待所有任务完成
        results = await self.join_async_task()

        # 5. 写回输出
        self.context.response.nodes = nodes
        self.context.response.results = results
        return results

    async def _process_node(self, node: VectorNode):
        """并发处理单个节点的辅助方法"""
        # 模拟一些异步处理，比如进一步分析节点内容
        await asyncio.sleep(0.1)  # 模拟异步操作
        return {"node_id": node.unique_id, "content": node.content[:50]}


# 使用示例
async def main():
    async with FlowLLMApp() as app:
        # 基本使用
        result = await RAGSearchOp().async_call(
            query="什么是机器学习？",
            workspace_id="knowledge_base",
            top_k=5
        )
        print(f"搜索结果：{result}")

        # 使用自定义 embedding_model 和 vector_store
        result2 = await RAGSearchOp(
            embedding_model="openai_compatible",
            vector_store="memory"
        ).async_call(
            query="深度学习",
            workspace_id="ai_knowledge"
        )
        print(f"搜索结果：{result2}")


if __name__ == "__main__":
    asyncio.run(main())