diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index 9685f57..3c889dc 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -11,13 +11,13 @@ jobs: os: [ ubuntu-latest ] env: OS: ${{ matrix.os }} - PYTHON: '3.10' + PYTHON: '3.11' steps: - uses: actions/checkout@master - name: Setup Python uses: actions/setup-python@master with: - python-version: '3.10' + python-version: '3.11' - name: Update setuptools run: | pip install -U setuptools wheel @@ -35,4 +35,4 @@ jobs: echo -e "\e[41m [**FAIL**] Please install pre-commit and format your code first. \e[0m" exit 1 fi - echo -e "\e[46m ********************************Passed******************************** \e[0m" \ No newline at end of file + echo -e "\e[46m ********************************Passed******************************** \e[0m" diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index dc60503..2a872ab 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -22,11 +22,11 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6 with: - python-version: '3.10' + python-version: '3.11' - name: Install dependencies run: | python -m pip install --upgrade pip diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml new file mode 100644 index 0000000..972afe0 --- /dev/null +++ b/.github/workflows/unittest.yml @@ -0,0 +1,43 @@ +name: Tests FlowLLM + +on: + push: + branches: [main, master, dev, develop] + pull_request: + branches: [main, master, dev, develop] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + unit-tests: + name: Unit Tests - py${{ matrix.python-version }} + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.12", "3.13"] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: 'pip' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip setuptools wheel + pip install -e ".[dev,core]" + + - name: Run unit tests + run: | + pytest tests/unit \ + -v \ + --tb=long \ + -s \ + --log-cli-level=WARNING diff --git a/.gitignore b/.gitignore index 67404c9..a0618d9 100644 --- a/.gitignore +++ b/.gitignore @@ -1,25 +1,54 @@ -.vscode -.env* +# OS / editor .DS_Store -.idea +.idea/ +.vscode/ +*.code-workspace + +# Local environment +.env +.env.* +!.env.example +!example.env +.venv/ venv/ -.ipynb_checkpoints -.__pycache__ -__pycache__ -*.log -tmp* -temp* -private* +env/ +private*/ + +# Python caches / test artifacts +__pycache__/ +*.py[cod] +*$py.class +.ipynb_checkpoints/ +.pytest_cache/ +.ruff_cache/ +.mypy_cache/ +.coverage +coverage.xml +htmlcov/ + +# Packaging / build outputs +build/ dist/ -nohup* +*.egg-info/ + +# Logs / temporary files +*.log +nohup.out +nohup*.out log/ -**/logs/* +logs/ +runs/ +tmp*/ +temp*/ .trash/ -runs -logs -*.egg-info/* -build/* -**/cache/** -*.pdf -*.csv -site/* \ No newline at end of file + +# FlowLLM runtime data +.flowllm/ +vault/ +*.db +*.sqlite +*.sqlite3 + +# Documentation build outputs +docs/_build/ +site/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7067f85..51c1504 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -3,35 +3,25 @@ repos: rev: v6.0.0 hooks: - id: check-ast - exclude: ^old/ - id: check-yaml - exclude: ^old/ - id: check-xml - exclude: ^old/ - id: check-toml - exclude: ^old/ - id: check-json - exclude: ^old/ - id: detect-private-key - exclude: ^old/ - id: trailing-whitespace - exclude: ^old/ - repo: https://github.com/asottile/add-trailing-comma rev: v4.0.0 hooks: - id: add-trailing-comma - exclude: ^old/ - repo: https://github.com/psf/black rev: 25.9.0 hooks: - id: black - exclude: ^old/ args: [--line-length=120] - repo: https://github.com/PyCQA/flake8 rev: 7.3.0 hooks: - id: flake8 - exclude: ^old/ args: [ "--extend-ignore=E203", "--max-line-length=120" @@ -42,8 +32,7 @@ repos: - id: pylint exclude: (?x)( - ^old/ - | ^docs + ^docs | pb2\.py$ | grpc\.py$ | \.demo$ @@ -54,6 +43,7 @@ repos: --disable=W0511, --disable=W0718, --disable=W0122, + --disable=W1203, --disable=C0103, --disable=R0913, --disable=R0917, @@ -81,8 +71,9 @@ repos: --disable=C3001, --disable=R1702, --disable=R0912, + --max-statements=120, --max-line-length=120, - --max-statements=75, + --max-module-lines=1500, ] - repo: https://github.com/regebro/pyroma rev: "5.0" diff --git a/LICENSE b/LICENSE index f8b0edd..65c2c5c 100644 --- a/LICENSE +++ b/LICENSE @@ -186,7 +186,7 @@ same "printed page" as the copyright notice for easier identification within third-party archives. - Copyright 2024 FlowLLM + Copyright 2025 Alibaba Group Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. diff --git a/README.md b/README.md index 72ca67d..8c0e020 100644 --- a/README.md +++ b/README.md @@ -3,229 +3,250 @@

- FlowLLM: Simplifying LLM-based HTTP/MCP Service Development
- If you find it useful, please give us a ⭐ Star. Your support drives our continuous improvement. -

- -

- Python Version - PyPI Version - License + Python Version + PyPI Version + PyPI Downloads + GitHub commit activity + License + English + 简体中文 GitHub Stars + DeepWiki

- English | 简体中文 + FlowLLM: Build LLM applications with ease.

---- +FlowLLM is a configuration-driven LLM application framework that organizes workflows, service entrypoints, and long-lived components with **Service, Job, Step, and Component**. -## 📖 Introduction +## Core Features -FlowLLM encapsulates LLM, Embedding, and vector_store capabilities as HTTP/MCP services. It is suitable for AI assistants, RAG applications, and workflow services, and can be integrated into MCP-compatible client tools. - -### 🏗️ Architecture Overview +- **Configuration-driven**: Starts from `flowllm/config/default.yaml`, with config files and dot-notation overrides. +- **Unified path**: `CLI / Client -> Service -> Application -> Job -> Step -> Component`. +- **Flexible Jobs**: Supports sync, streaming, background, and scheduled tasks exposed through HTTP or MCP. +- **Pluggable components**: Extend Steps, Services, Clients, LLMs, Embeddings, Embedding Stores, and Agent Wrappers through registry `R`.

- FlowLLM Framework + FlowLLM Architecture

-### 🌟 Applications Based on FlowLLM - -| Project Name | Description | -|-----------------------------------------------|--------------------------------------| -| [ReMe](https://github.com/agentscope-ai/ReMe) | Memory management toolkit for agents | +## Quick Start -### 📢 Recent Updates +### Installation -| Date | Update Content | -|------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| 2025-11-15 | Added [File Tool Op](docs/zh/guide/file_tool_op_guide.md) feature with 13 file operation tools, supporting file reading, writing, editing, searching, directory operations, system command execution, and task management | -| 2025-11-14 | Added Token counting capability, supporting accurate calculation of token counts for messages and tools via `self.token_count()` method, with support for multiple backends (base, openai, hf). See configuration examples in [default.yaml](flowllm/config/default.yaml) | +FlowLLM requires Python 3.11+. -### 📚 Learning Resources +Install development dependencies from source: -Project developers will share their latest learning materials here. +```bash +git clone https://github.com/flowllm-ai/flowllm.git +cd flowllm +pip install -e ".[dev]" +``` -| Date | Title | Description | -|------------|--------------------------------------------------------------------------------------------------------|-----------------------------------------------------------------------------------| -| 2025-11-24 | [Mem-PAL: Memory-Augmented Personalized Assistant](./docs/zh/reading/20251124-mem-pal.md) | Mem-PAL: Memory-Augmented Personalized Assistant with Log-based Structured Memory | -| 2025-11-14 | [HaluMem Analysis](./docs/zh/reading/20251114-halumem.md) | HaluMem: Evaluating Hallucinations in Memory Systems of Agents Analysis | -| 2025-11-13 | [Gemini CLI Context Management Mechanism](./docs/zh/reading/20251113-gemini-cli-context-management.md) | Multi-layer Context Management Strategy for Gemini CLI | -| 2025-11-10 | [Context Management Guide](./docs/zh/reading/20251110-manus-context-report.md) | Context Management Guide | -| 2025-11-10 | [LangChain&Manus Video Materials](./docs/zh/reading/20251110-manus-context-raw.md) | LangChain & Manus Context Management Video | +Install the optional Claude Code wrapper when needed: -### ⭐ Core Features +```bash +pip install -e ".[claude-code]" +``` -- **Simple Op Development**: Inherit from `BaseOp` or `BaseAsyncOp` and implement your business logic. FlowLLM provides lazy-initialized LLM, Embedding models, and vector stores accessible via `self.llm`, `self.embedding_model`, and `self.vector_store`. It also offers prompt template management through `prompt_format()` and `get_prompt()` methods. Additionally, FlowLLM includes built-in token counting capabilities. Use `self.token_count()` to accurately calculate token counts for messages and tools, supporting multiple backends (base, openai, hf, etc.). +Install all optional dependencies: -- **Flexible Flow Orchestration**: Compose Ops into Flows via YAML configuration. `>>` denotes serial composition; `|` denotes parallel composition. For example, `SearchOp() >> (AnalyzeOp() | TranslateOp()) >> FormatOp()` builds complex workflows. Define input/output schemas and start the service with `flowllm config=your_config`. +```bash +pip install -e ".[full]" +``` -- **Automatic Service Generation**: FlowLLM automatically generates HTTP, MCP, and CMD services. The HTTP service provides RESTful APIs with synchronous JSON and HTTP Stream responses. The MCP service registers as Model Context Protocol tools for MCP-compatible clients. The CMD service executes a single Op in command-line mode for quick testing and debugging. +### Start the Service ---- +```bash +flowllm start +``` -## ⚡ Quick Start +The default service address is `127.0.0.1:2333`, and the default workspace is `.flowllm/`. Startup automatically creates: -### 📦 Step0 Installation +```text +.flowllm/ +├── metadata/ +└── session/ +``` -#### 📥 From PyPI +You can override configuration from the command line: ```bash -pip install flowllm +flowllm start service.port=8181 enable_logo=false +flowllm start workspace_dir=/tmp/flowllm-demo service.host=127.0.0.1 service.port=8181 ``` -#### 🔧 From Source +See the [Quick Start](docs/en/quick_start.md) for more startup and invocation examples. -```bash -git clone https://github.com/flowllm-ai/flowllm.git -cd flowllm -pip install -e . -``` +## Calling Jobs -For detailed installation and configuration, refer to the [Installation Guide](docs/zh/guide/installation.md). +After the service starts, CLI commands other than `start` call server-side Jobs with the same name through the client: -### ⚙️ Configuration +```bash +flowllm version +flowllm health_check +flowllm help +flowllm demo query="Hello FlowLLM" min_score=0.8 +flowllm add a=1 b=2 +``` -Create a `.env` file and configure your API keys. Copy from `example.env` and modify: +The HTTP entry point is `POST /`: ```bash -cp example.env .env +curl -s http://127.0.0.1:2333/add \ + -H 'Content-Type: application/json' \ + -d '{"a":1,"b":2}' ``` -Configure your API keys in the `.env` file: +Streaming Jobs return SSE: ```bash -FLOW_LLM_API_KEY=sk-xxxx -FLOW_LLM_BASE_URL=https://xxxx/v1 -FLOW_EMBEDDING_API_KEY=sk-xxxx -FLOW_EMBEDDING_BASE_URL=https://xxxx/v1 +flowllm stream_demo query="Hi" repeat=3 interval=0.05 + +curl -N http://127.0.0.1:2333/stream_demo \ + -H 'Content-Type: application/json' \ + -d '{"query":"Hi","repeat":3,"interval":0.05}' ``` -For detailed configuration, refer to the [Configuration Guide](docs/zh/guide/config_guide.md). +Built-in Jobs: -### 🛠️ Step1 Build Op +| Job | Backend | Description | +| --- | --- | --- | +| `version` | `base` | Returns the FlowLLM package version. | +| `health_check` | `base` | Returns a component health-check summary. | +| `help` | `base` | Lists registered Jobs and their metadata. | +| `demo` | `base` | Two-step echo example. | +| `add` | `base` | Adds two numbers. | +| `stream_demo` | `stream` | Streams the input text character by character. | -```python -from flowllm.core.context import C -from flowllm.core.op import BaseAsyncOp -from flowllm.core.schema import Message -from flowllm.core.enumeration import Role +## Configuration -@C.register_op() -class SimpleChatOp(BaseAsyncOp): - async def async_execute(self): - query = self.context.get("query", "") - messages = [Message(role=Role.USER, content=query)] +The default configuration is located at `flowllm/config/default.yaml`. Root configuration sections include application parameters, service, Jobs, and Components: - # Use token_count method to calculate token count - token_num = self.token_count(messages) - print(f"Input tokens: {token_num}") +```yaml +service: + backend: http + +jobs: + add: + backend: base + description: "add two numbers" + steps: + - backend: add_step +``` + +You can override model and embedding variables through `.env`: - response = await self.llm.achat(messages=messages) - self.context.response.answer = response.content.strip() +```bash +cat > .env <<'EOF' +LLM_BACKEND=openai +LLM_MODEL_NAME=qwen3.7-plus +LLM_API_KEY=your_api_key +LLM_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 + +EMBEDDING_BACKEND=openai +EMBEDDING_MODEL_NAME=text-embedding-v4 +EMBEDDING_API_KEY=your_api_key +EMBEDDING_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 +EOF ``` -For details, refer to the [Simple Op Guide](docs/zh/guide/async_op_minimal_guide.md), [LLM Op Guide](docs/zh/guide/async_op_llm_guide.md), and [Advanced Op Guide](docs/zh/guide/async_op_advance_guide.md) (including Embedding, VectorStore, and concurrent execution). +You can also specify a YAML or JSON configuration file: -### 📝 Step2 Configure Config +```bash +flowllm start config=/path/to/app.yaml +``` -The following example demonstrates building an MCP (Model Context Protocol) service. Create a configuration file `my_mcp_config.yaml`: +Configuration parsing supports `${VAR}`, `${VAR:-default}`, booleans, numbers, JSON lists and dictionaries, and automatic `null` conversion. -```yaml -backend: mcp +## Code Framework -mcp: - transport: sse - host: "0.0.0.0" - port: 8001 - -flow: - demo_mcp_flow: - flow_content: MockSearchOp() - description: "Search results for a given query." - input_schema: - query: - type: string - description: "User query" - required: true - -llm: - default: - backend: openai_compatible - model_name: qwen3-30b-a3b-instruct-2507 - params: - temperature: 0.6 - token_count: # Optional, configure token counting backend - model_name: Qwen/Qwen3-30B-A3B-Instruct-2507 - backend: hf # Supports base, openai, hf, etc. - params: - use_mirror: true -``` - -### 🚀 Step3 Start MCP Service +FlowLLM's core layering is: -```bash -flowllm \ - config=my_mcp_config \ - backend=mcp \ # Optional, overrides config - mcp.transport=sse \ # Optional, overrides config - mcp.port=8001 \ # Optional, overrides config - llm.default.model_name=qwen3-30b-a3b-thinking-2507 # Optional, overrides config +```text +CLI / Client -> Service -> Application -> Job -> Step -> Component ``` -After the service starts, refer to the [Client Guide](docs/zh/guide/client_guide.md) to use the service and obtain the tool_call required by the model. +| Layer | Role | +| --- | --- | +| `Application` | Loads configuration, creates the workspace, assembles Service, Components, and Jobs, and starts lifecycles in dependency order. | +| `Service` | Exposes servable Jobs as HTTP routes or MCP tools. | +| `Job` | Orchestrates externally callable capabilities or background processes by executing Steps in order. | +| `Step` | Atomic workflow operation that reads and writes `RuntimeContext`, `Response`, and streaming queues. | +| `Component` | Long-lived infrastructure such as LLMs, Embeddings, Embedding Stores, and Agent Wrappers. | +| `Registry` | Global registry `R`, which looks up implementations by `(component_type, backend_name)`. | + +Minimal example for adding a Step: + +```python +from flowllm.components import R +from flowllm.steps import BaseStep ---- -## 📚 Detailed Documentation +@R.register("reverse_step") +class ReverseStep(BaseStep): + async def execute(self): + text = self.context.get("text", "") + self.context.response.answer = text[::-1] + return self.context.response +``` -### 🚀 Getting Started -- [Installation Guide](docs/zh/guide/installation.md) -- [Configuration Guide](docs/zh/guide/config_guide.md) +Then declare the Job in configuration: -### 🔧 Op Development -- [Op Introduction](docs/zh/guide/op_introduction.md) -- [Simple Op Guide](docs/zh/guide/async_op_minimal_guide.md) -- [LLM Op Guide](docs/zh/guide/async_op_llm_guide.md) -- [Advanced Op Guide](docs/zh/guide/async_op_advance_guide.md) -- [Tool Op Guide](docs/zh/guide/async_tool_op_guide.md) -- [File Tool Op Guide](docs/zh/guide/file_tool_op_guide.md) -- [Vector Store Guide](docs/zh/guide/vector_store_guide.md) +```yaml +jobs: + reverse: + backend: base + description: "reverse text" + parameters: + type: object + properties: + text: + type: string + required: + - text + steps: + - backend: reverse_step +``` -### 🔀 Flow Orchestration -- [Flow Guide](docs/zh/guide/flow_guide.md) +After adding an implementation, make sure the module is imported in the package `__init__.py`; otherwise, the registration decorator will not run. See the [code framework](docs/en/framework.md) for details. -### 🌐 Service Usage -- [HTTP Service Guide](docs/zh/guide/http_service_guide.md) -- [HTTP Stream Guide](docs/zh/guide/http_stream_guide.md) -- [MCP Service Guide](docs/zh/guide/mcp_service_guide.md) -- [CMD Service Guide](docs/zh/guide/cmd_service_guide.md) -- [Client Guide](docs/zh/guide/client_guide.md) +## MCP Service ---- +When the service backend is set to `mcp`, FlowLLM exposes non-streaming Jobs with `enable_serve: true` as MCP tools. `StreamJob` is not exposed by MCP Service. -## 🤝 Contributing +```yaml +service: + backend: mcp + transport: sse + host: 127.0.0.1 + port: 2333 +``` -Contributions of all forms are welcome! For participation methods, refer to the [Contribution Guide](docs/zh/guide/contribution.md). +MCP transport supports `stdio`, `sse`, and `streamable-http`. ---- +## Documentation -## 📄 License +- [Quick Start](docs/en/quick_start.md) +- [Code Framework](docs/en/framework.md) +- [Contributing](docs/en/contributing.md) +- [FlowLLM Development Skill](skills/flowllm_dev/SKILL.md) -This project is licensed under the [Apache 2.0](LICENSE) license. +## Open Source and Contributing ---- +FlowLLM is licensed under Apache 2.0. Before contributing, read the [contribution guide](docs/en/contributing.md) and [development skill](skills/flowllm_dev/SKILL.md), then run: -## Star 历史 +```bash +pre-commit run --all-files +pytest +``` -[![Star History Chart](https://api.star-history.com/svg?repos=flowllm-ai/flowllm&type=Date)](https://www.star-history.com/#flowllm-ai/flowllm&Date) +## License ---- +This project is open source under the Apache License 2.0. See [LICENSE](./LICENSE) for details. -

- GitHub • - Documentation • - PyPI -

+## Star History +[![Star History Chart](https://api.star-history.com/svg?repos=flowllm-ai/flowllm&type=Date)](https://www.star-history.com/#flowllm-ai/flowllm&Date) diff --git a/README_ZH.md b/README_ZH.md index 2af68d6..605f0b3 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -3,229 +3,250 @@

- FlowLLM：让基于LLM的HTTP/MCP服务开发更简单
- 如果觉得有用，欢迎给个 ⭐ Star，您的支持是我们持续改进的动力 + Python Version + PyPI Version + PyPI Downloads + GitHub commit activity + License + English + 简体中文 + GitHub Stars + DeepWiki

- Python Version - PyPI Version - License - GitHub Stars + FlowLLM: 轻松构建 LLM 应用。

+FlowLLM 是一个配置驱动的 LLM 应用框架，用 **Service、Job、Step、Component** 组织工作流、服务入口和长期组件。 + +## 核心特性 + +- **配置驱动**：默认从 `flowllm/config/default.yaml` 启动，支持配置文件和 dot notation 覆盖。 +- **统一链路**：`CLI / Client -> Service -> Application -> Job -> Step -> Component`。 +- **多种 Job 形态**：支持同步、流式、后台和定时任务，并可通过 HTTP 或 MCP 暴露。 +- **可插拔组件**：通过注册表 `R` 扩展 Step、Service、Client、LLM、Embedding、Embedding Store 和 Agent Wrapper。 +

- English | 简体中文 + FlowLLM Architecture

---- +## 快速开始 -## 📖 简介 +### 安装 -FlowLLM 将 LLM/Embedding/vector_store 能力封装为 HTTP/MCP 服务，适用于 AI 对话助手、RAG 应用、工作流服务等场景，并可集成到支持 -MCP 的客户端工具中。 +FlowLLM 要求 Python 3.11+。 -### 🏗️ 架构概览 +从源码安装开发依赖： -

- FlowLLM Framework -

+```bash +git clone https://github.com/flowllm-ai/flowllm.git +cd flowllm +pip install -e ".[dev]" +``` -### 🌟 基于FlowLLM的应用 +如果需要 Claude Code wrapper，再安装可选依赖： -| 项目名 | 描述 | -|-----------------------------------------------|---------------| -| [ReMe](https://github.com/agentscope-ai/ReMe) | 面向智能体的记忆管理工具包 | +```bash +pip install -e ".[claude-code]" +``` -### 📢 最近更新 +需要完整可选依赖时： -| 日期 | 更新内容 | -|------------|---------------------------------------------------------------| -| 2025-11-15 | 新增 [File Tool Op](docs/zh/guide/file_tool_op_guide.md) 功能，提供 13 个文件操作工具，支持文件读取、写入、编辑、搜索、目录操作、系统命令执行和任务管理等功能 | -| 2025-11-14 | 新增 Token 计数能力，支持通过 `self.token_count()` 方法准确计算消息和工具的 token 数量，支持多种后端（base、openai、hf），配置示例参考 [default.yaml](flowllm/config/default.yaml) | +```bash +pip install -e ".[full]" +``` -### 📚 学习资料分享 +### 启动服务 -项目开发者会在这里分享最近的学习资料。 +```bash +flowllm start +``` -| 日期 | 标题 | 描述 | -|------------|-----------------------------------------------------------------------------------|-------------------------------------------------------------------| -| 2025-11-24 | [Mem-PAL：基于日志结构化记忆的个性化助手](./docs/zh/reading/20251124-mem-pal.md) | Mem-PAL: Memory-Augmented Personalized Assistant with Log-based Structured Memory 解读 | -| 2025-11-14 | [HaluMem解读](./docs/zh/reading/20251114-halumem.md) | HaluMem: Evaluating Hallucinations in Memory Systems of Agents 解读 | -| 2025-11-13 | [Gemini CLI 上下文管理机制](./docs/zh/reading/20251113-gemini-cli-context-management.md) | Gemini CLI 的多层上下文管理策略 | -| 2025-11-10 | [上下文管理指南](./docs/zh/reading/20251110-manus-context-report.md) | 上下文管理指南 | -| 2025-11-10 | [LangChain&Manus视频资料](./docs/zh/reading/20251110-manus-context-raw.md) | LangChain & Manus Context Management Video | +默认服务地址是 `127.0.0.1:2333`，默认 workspace 是 `.flowllm/`。启动时会自动创建： -### ⭐ 核心特性 +```text +.flowllm/ +├── metadata/ +└── session/ +``` -- **简单易用的 Op 开发**：继承 BaseOp 或 BaseAsyncOp 基类，实现业务逻辑即可。FlowLLM提供了延迟初始化的 LLM、Embedding 模型和向量库，开发者只需通过 `self.llm`、`self.embedding_model`、`self.vector_store` 即可轻松使用这些资源。同时FlowLLM提供了完整的 Prompt 模板管理能力，通过 `prompt_format()` 和 `get_prompt()` 方法进行格式化和使用。此外，FlowLLM 还内置了 Token 计数能力，通过 `self.token_count()` 方法可以准确计算消息和工具的 token 数量，支持多种后端（base、openai、hf 等）。 +可以用命令行覆盖配置： -- **灵活的 Flow 编排**：通过 YAML 配置文件将 Op 组合成 Flow，支持灵活的编排方式。`>>` 表示串行组合，`|` 表示并行组合，例如 `SearchOp() >> (AnalyzeOp() | TranslateOp()) >> FormatOp()` 可构建复杂的工作流。定义输入输出 Schema 后，使用 `flowllm config=your_config` 命令即可启动服务。 +```bash +flowllm start service.port=8181 enable_logo=false +flowllm start workspace_dir=/tmp/flowllm-demo service.host=127.0.0.1 service.port=8181 +``` -- **自动生成服务**：配置完成后，FlowLLM 会自动生成 HTTP、MCP 和 CMD 服务。HTTP 服务提供标准的 RESTFul API，支持同步 JSON 响应和 HTTP Stream 流式响应。MCP 服务会自动注册为 Model Context Protocol 工具，可集成到支持 MCP 的客户端中。CMD 服务支持命令行模式执行单个 Op，适合快速测试和调试。 +更多启动和调用示例见 [快速开始](docs/zh/quick_start.md)。 ---- +## 调用 Job -## ⚡ 快速开始 +启动服务后，CLI 的非 `start` 命令会通过 client 调用服务端同名 Job： -### 📦 Step0 安装 +```bash +flowllm version +flowllm health_check +flowllm help +flowllm demo query="Hello FlowLLM" min_score=0.8 +flowllm add a=1 b=2 +``` -#### 📥 From PyPI +HTTP 入口是 `POST /`： ```bash -pip install flowllm +curl -s http://127.0.0.1:2333/add \ + -H 'Content-Type: application/json' \ + -d '{"a":1,"b":2}' ``` -#### 🔧 From Source +流式 Job 返回 SSE： ```bash -git clone https://github.com/flowllm-ai/flowllm.git -cd flowllm -pip install -e . +flowllm stream_demo query="Hi" repeat=3 interval=0.05 + +curl -N http://127.0.0.1:2333/stream_demo \ + -H 'Content-Type: application/json' \ + -d '{"query":"Hi","repeat":3,"interval":0.05}' ``` -详细安装与配置方法请参考 [安装指南](docs/zh/guide/installation.md)。 +内置 Job： -### ⚙️ 配置 +| Job | backend | 说明 | +| --- | --- | --- | +| `version` | `base` | 返回 FlowLLM 包版本。 | +| `health_check` | `base` | 返回组件健康检查摘要。 | +| `help` | `base` | 列出已注册 Job 及其 metadata。 | +| `demo` | `base` | 两步 echo 示例。 | +| `add` | `base` | 两数相加示例。 | +| `stream_demo` | `stream` | 将输入按字符流式输出。 | -创建 `.env` 文件，配置 API Key。你可以从 `example.env` 复制并修改： +## 配置 -```bash -cp example.env .env +默认配置位于 `flowllm/config/default.yaml`。配置根节点包括应用参数、服务、Job 和 Component： + +```yaml +service: + backend: http + +jobs: + add: + backend: base + description: "add two numbers" + steps: + - backend: add_step ``` -然后在 `.env` 文件中配置你的 API Key： +可通过 `.env` 覆盖模型和 embedding 相关变量： ```bash -FLOW_LLM_API_KEY=sk-xxxx -FLOW_LLM_BASE_URL=https://xxxx/v1 -FLOW_EMBEDDING_API_KEY=sk-xxxx -FLOW_EMBEDDING_BASE_URL=https://xxxx/v1 +cat > .env <<'EOF' +LLM_BACKEND=openai +LLM_MODEL_NAME=qwen3.7-plus +LLM_API_KEY=your_api_key +LLM_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 + +EMBEDDING_BACKEND=openai +EMBEDDING_MODEL_NAME=text-embedding-v4 +EMBEDDING_API_KEY=your_api_key +EMBEDDING_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 +EOF ``` -详细配置说明请参考 [配置指南](docs/zh/guide/config_guide.md)。 +也可以指定 YAML 或 JSON 配置文件： -### 🛠️ Step1 构建Op +```bash +flowllm start config=/path/to/app.yaml +``` -```python -from flowllm.core.context import C -from flowllm.core.op import BaseAsyncOp -from flowllm.core.schema import Message -from flowllm.core.enumeration import Role +配置解析支持 `${VAR}`、`${VAR:-default}`、bool、数字、JSON list/dict 和 `null` 自动转换。 -@C.register_op() -class SimpleChatOp(BaseAsyncOp): - async def async_execute(self): - query = self.context.get("query", "") - messages = [Message(role=Role.USER, content=query)] +## 代码框架 - # 使用 token_count 方法计算 token 数量 - token_num = self.token_count(messages) - print(f"Input tokens: {token_num}") +FlowLLM 的核心分层如下： - response = await self.llm.achat(messages=messages) - self.context.response.answer = response.content.strip() +```text +CLI / Client -> Service -> Application -> Job -> Step -> Component ``` -详细内容请参考 [简单 Op 指南](docs/zh/guide/async_op_minimal_guide.md)、[LLM Op 指南](docs/zh/guide/async_op_llm_guide.md) 和 [高级 Op 指南](docs/zh/guide/async_op_advance_guide.md)（包含 Embedding、VectorStore 和并发执行等高级功能）。 - -### 📝 Step2 配置config +| 层级 | 作用 | +| --- | --- | +| `Application` | 加载配置，创建 workspace，装配 Service、Component 和 Job，并按依赖顺序启动生命周期。 | +| `Service` | 将可服务的 Job 暴露为 HTTP 路由或 MCP 工具。 | +| `Job` | 外部可调用能力或后台流程的编排单元，按顺序执行 Step。 | +| `Step` | 工作流原子操作，读写 `RuntimeContext`、`Response` 和流式队列。 | +| `Component` | 长期存在的基础设施，例如 LLM、Embedding、Embedding Store、Agent Wrapper。 | +| `Registry` | 全局注册表 `R`，按 `(component_type, backend_name)` 查找实现。 | -以下示例展示如何构建一个 MCP（Model Context Protocol）服务。创建配置文件 `my_mcp_config.yaml`： +新增 Step 的最小示例： -```yaml -backend: mcp +```python +from flowllm.components import R +from flowllm.steps import BaseStep -mcp: - transport: sse - host: "0.0.0.0" - port: 8001 - -flow: - demo_mcp_flow: - flow_content: MockSearchOp() - description: "Search results for a given query." - input_schema: - query: - type: string - description: "User query" - required: true - -llm: - default: - backend: openai_compatible - model_name: qwen3-30b-a3b-instruct-2507 - params: - temperature: 0.6 - token_count: # 可选，配置 token 计数后端 - model_name: Qwen/Qwen3-30B-A3B-Instruct-2507 - backend: hf # 支持 base、openai、hf 等 - params: - use_mirror: true -``` - -### 🚀 Step3 启动 MCP 服务 -```bash -flowllm \ - config=my_mcp_config \ - backend=mcp \ # 可选，覆盖config配置 - mcp.transport=sse \ # 可选，覆盖config配置 - mcp.port=8001 \ # 可选，覆盖config配置 - llm.default.model_name=qwen3-30b-a3b-thinking-2507 # 可选，覆盖config配置 +@R.register("reverse_step") +class ReverseStep(BaseStep): + async def execute(self): + text = self.context.get("text", "") + self.context.response.answer = text[::-1] + return self.context.response ``` -服务启动后可以参考[Client Guide](docs/zh/guide/client_guide.md)来使用服务，可以直接获取模型所需要的tool_call。 +然后在配置中声明 Job： ---- +```yaml +jobs: + reverse: + backend: base + description: "reverse text" + parameters: + type: object + properties: + text: + type: string + required: + - text + steps: + - backend: reverse_step +``` -## 📚 详细文档 +新增实现后要确保模块在包的 `__init__.py` 中被 import，否则注册装饰器不会执行。更多细节见 [代码框架](docs/zh/framework.md)。 -### 🚀 入门指南 -- [安装指南](docs/zh/guide/installation.md) -- [配置指南](docs/zh/guide/config_guide.md) +## MCP 服务 -### 🔧 Op 开发 -- [Op 介绍](docs/zh/guide/op_introduction.md) -- [简单 Op 指南](docs/zh/guide/async_op_minimal_guide.md) -- [LLM Op 指南](docs/zh/guide/async_op_llm_guide.md) -- [高级 Op 指南](docs/zh/guide/async_op_advance_guide.md) -- [Tool Op 指南](docs/zh/guide/async_tool_op_guide.md) -- [File Tool Op 指南](docs/zh/guide/file_tool_op_guide.md) -- [Vector Store 指南](docs/zh/guide/vector_store_guide.md) +将 service backend 改成 `mcp` 后，FlowLLM 会把非流式、`enable_serve: true` 的 Job 暴露为 MCP tool。`StreamJob` 不会被 MCP Service 暴露。 -### 🔀 Flow 编排 -- [Flow 指南](docs/zh/guide/flow_guide.md) +```yaml +service: + backend: mcp + transport: sse + host: 127.0.0.1 + port: 2333 +``` -### 🌐 服务使用 -- [HTTP 服务指南](docs/zh/guide/http_service_guide.md) -- [HTTP Stream 指南](docs/zh/guide/http_stream_guide.md) -- [MCP 服务指南](docs/zh/guide/mcp_service_guide.md) -- [CMD 服务指南](docs/zh/guide/cmd_service_guide.md) -- [客户端指南](docs/zh/guide/client_guide.md) +MCP transport 支持 `stdio`、`sse` 和 `streamable-http`。 ---- +## 文档 -## 🤝 参与贡献 +- [快速开始](docs/zh/quick_start.md) +- [代码框架](docs/zh/framework.md) +- [开源与贡献](docs/zh/contributing.md) +- [FlowLLM 开发 Skill](skills/flowllm_dev/SKILL.md) -欢迎各种形式的贡献！具体参与方式请参考 [贡献指南](docs/zh/guide/contribution.md)。 +## 开源与贡献 ---- +FlowLLM 使用 Apache 2.0 许可证。贡献前请阅读 [贡献指南](docs/zh/contributing.md) 和 [开发 Skill](skills/flowllm_dev/SKILL.md)，并尽量运行： -## 📄 许可证 +```bash +pre-commit run --all-files +pytest +``` -本项目采用 [Apache 2.0](LICENSE) 许可证。 +## License ---- +This project is open source under the Apache License 2.0. See [LICENSE](./LICENSE) for details. -## Star History +## Star 历史 [![Star History Chart](https://api.star-history.com/svg?repos=flowllm-ai/flowllm&type=Date)](https://www.star-history.com/#flowllm-ai/flowllm&Date) - ---- - -

- GitHub • - 文档 • - PyPI -

diff --git a/docs/en/contributing.md b/docs/en/contributing.md new file mode 100644 index 0000000..967294c --- /dev/null +++ b/docs/en/contributing.md @@ -0,0 +1,93 @@ +# Open Source and Contributing + +FlowLLM repository: **https://github.com/flowllm-ai/flowllm** + +For your first local run, see [Quick Start](./quick_start.md). Before changing the runtime, Jobs, Steps, components, services, or configuration, read [Framework](./framework.md). + +## Local Development + +FlowLLM requires Python 3.11+. + +```bash +python -m venv .venv +source .venv/bin/activate +pip install -e ".[dev]" +pre-commit install +``` + +Install the complete optional dependency set when needed: + +```bash +pip install -e ".[full]" +``` + +Main directories: + +| Path | Description | +|-------------------------------|-------------------------------------| +| `flowllm/` | Python package source | +| `flowllm/config/default.yaml` | Default configuration | +| `flowllm/components/` | Services, clients, Jobs, components | +| `flowllm/steps/` | Step implementations | +| `flowllm/schema/` | Request, response, and config models | +| `tests/` | Unit and integration tests | +| `docs/zh/` | Chinese documentation | + +## Development Conventions + +Follow the main runtime path: + +```text +CLI / Client -> Service -> Application -> Job -> Step -> Component +``` + +- Prefer exposing external capabilities as Jobs, then serve them through HTTP or MCP Service. +- Put reusable infrastructure in `flowllm/components/` and inherit from `BaseComponent`. +- Put atomic business operations in `flowllm/steps/` and inherit from `BaseStep`. +- Register new implementations with `@R.register("")`, and make sure the module is imported by `__init__.py`. +- Put default behavior in `flowllm/config/default.yaml`, keeping `flowllm start` runnable. +- Update documentation when changing user-visible behavior. + +## Testing + +Recommended before submitting: + +```bash +pre-commit run --all-files +pytest +``` + +For focused validation: + +```bash +pytest tests/unit +pytest tests/integration +``` + +Testing guidance: + +- Add regression tests when fixing bugs. +- Cover core paths and failure paths when adding Steps, Jobs, or components. +- Cover user entrypoints when changing config parsing, the registry, lifecycle behavior, service exposure, or streaming output. +- For tests that depend on LLMs, embeddings, or external services, document the required environment in the PR. + +## Commits and PRs + +Conventional Commits are recommended: + +```text +(): +``` + +Common types: `feat`, `fix`, `docs`, `test`, `refactor`, `chore`, `perf`, `style`. + +Examples: + +```text +feat(step): add reverse text demo +fix(config): preserve leading zero strings +docs(zh): update quick start +test(service): cover stream response +``` + +PR titles should follow the same format. For larger changes, first describe the background, target behavior, impact scope, and test plan in an Issue. diff --git a/docs/en/framework.md b/docs/en/framework.md new file mode 100644 index 0000000..e7adb07 --- /dev/null +++ b/docs/en/framework.md @@ -0,0 +1,157 @@ +# FlowLLM Framework + +FlowLLM's runtime path: + +```text +CLI / Client -> Service -> Application -> Job -> Step -> Component +``` + +`Application` assembles services, components, and Jobs from configuration; `Service` exposes serveable Jobs; `Job` executes Steps in order; `Step` shares request data, responses, and streaming queues through `RuntimeContext`. + +

+ FlowLLM framework architecture +

+ +## Directory + +```text +flowllm/ + application.py # CLI entrypoint and Application lifecycle + config/default.yaml # Default service / jobs / components + config/config_parser.py # config=, dot notation, environment variable expansion + components/ + base_component.py # BaseComponent and bind dependency declarations + component_registry.py # Global registry R + runtime_context.py # Per-Job execution context + job/ # base / stream / background / cron + service/ # http / mcp + client/ # http / mcp + as_llm/ as_embedding/ # Model wrappers + embedding_store/ + agent_wrapper/ + steps/ + base_step.py # BaseStep, Ref, dispatch_steps + common/ # version, help, health_check, demo, add, stream_demo + schema/ # Request, Response, StreamChunk, configuration models + enumeration/ # ComponentEnum, ChunkEnum +``` + +Default workspace: + +```text +.flowllm/ +├── metadata/ +└── session/ +``` + +## Startup + +`flowllm/application.py::main()` handles three command types: + +| action | Behavior | +|----------------|----------------------------------------------------------------| +| `start` | Load `.env` and configuration, create `Application`, start Service | +| `find_flowllm` | Find the project path | +| Other actions | Call the server-side Job with the same name through the client | + +Configuration parsing supports: + +- Loading `flowllm/config/default.yaml` by default. +- `config=/path/to/app.yaml` for a YAML/JSON file. +- Dot notation overrides such as `service.port=8181`. +- Environment variable expansion with `${VAR}` and `${VAR:-default}`. +- Automatic conversion for bools, numbers, JSON lists/dicts, and null. + +## Service + +The HTTP service registers Jobs as endpoints: + +| Job type | HTTP behavior | +|-----------------------|----------------------------------------------------| +| `BaseJob` | `POST /`, returns `Response` JSON | +| `StreamJob` | `POST /`, returns `text/event-stream` | +| `enable_serve: false` | Not registered | + +The MCP service exposes only non-streaming serveable Jobs. Background and scheduled Jobs are not exposed by default. + +## Registry + +FlowLLM uses the global registry `R`: + +```python +from flowllm.components import R +from flowllm.steps import BaseStep + + +@R.register("reverse_step") +class ReverseStep(BaseStep): + async def execute(self): + text = self.context.get("text", "") + self.context.response.answer = text[::-1] + return self.context.response +``` + +The registry key is `(component_type, backend_name)`. The same backend name can exist under different types; for example, `http` can be both a service backend and a client backend. + +Registration happens when modules are imported. After adding a Step or Component, make sure its package `__init__.py` imports the new module; otherwise the `backend` configured in YAML cannot find its implementation. + +## Component + +Long-lived infrastructure inherits from `BaseComponent`. Common lifecycle methods: + +- `start()`: resolve dependencies and call `_start()`. +- `close()`: call `_close()` and close owned components in reverse order. +- `restart()`: close and then start again. + +Component dependencies are declared with `BaseComponent.bind()`, and `Application` starts components in dependency-topological order: + +```python +self.embedding_store = self.bind("default", BaseEmbeddingStore, optional=False) +``` + +## Job + +A Job is an orchestration unit for an externally callable capability or a background workflow. Jobs are configured under `jobs:`. + +- `BaseJob`: creates a `RuntimeContext` for each call, executes Steps in order, and returns `context.response`. +- `StreamJob`: Steps write `StreamChunk` values to `stream_queue`, and the service layer outputs SSE. +- `BackgroundJob`: executes repeatedly based on `interval`. +- `CronJob`: executes on a cron expression. + +Job configuration example: + +```yaml +jobs: + reverse: + backend: base + description: "reverse text" + steps: + - backend: reverse_step +``` + +## Step + +A Step is an atomic workflow operation. It reads and writes `RuntimeContext`: + +| Field | Description | +|----------------|---------------------------------------------| +| `data` | Request parameters and shared inter-step data | +| `response` | Final return value | +| `stream_queue` | Streaming output queue | + +Streaming output: + +```python +await self.context.add_stream_string(text, ChunkEnum.CONTENT) +``` + +Step configuration also supports `input_mapping`, `output_mapping`, `prompt_dict`, `dispatch_steps`, and `language`. Use `BaseStep.Ref` to access application components; `as_llm` and the optional `agent_wrapper` are provided by default. + +## Default Capabilities + +`flowllm/config/default.yaml` includes: + +- `version`, `health_check`, `help` +- `demo`, `add` +- `stream_demo` +- Default `as_llm`, `as_embedding`, `embedding_store`, and `agent_wrapper` diff --git a/docs/en/quick_start.md b/docs/en/quick_start.md new file mode 100644 index 0000000..919c4d5 --- /dev/null +++ b/docs/en/quick_start.md @@ -0,0 +1,119 @@ +# Quick Start + +FlowLLM is a configuration-driven LLM application framework: configuration defines Services, Components, and Jobs; Jobs execute Steps in order; Services expose Jobs through CLI, HTTP, or MCP. + +## Installation + +FlowLLM requires Python 3.11+. + +```bash +git clone https://github.com/flowllm-ai/flowllm.git +cd flowllm +pip install -e ".[dev]" +``` + +Install the Claude Code wrapper only when needed: + +```bash +pip install -e ".[claude-code]" +``` + +## Start + +```bash +flowllm start +``` + +The default service address is `127.0.0.1:2333`, and the default workspace is `.flowllm/`. You can override configuration values with dot notation: + +```bash +flowllm start service.port=8181 enable_logo=false +flowllm start workspace_dir=/tmp/flowllm-demo service.host=127.0.0.1 service.port=8181 +``` + +FlowLLM creates the following directories automatically on startup: + +```text +.flowllm/ +├── metadata/ +└── session/ +``` + +## Call Jobs + +After the service starts, CLI commands other than `start` call server-side Jobs: + +```bash +flowllm version +flowllm health_check +flowllm help +flowllm demo query="Hello FlowLLM" min_score=0.8 +flowllm add a=1 b=2 +``` + +The HTTP entrypoint is `POST /`: + +```bash +curl -s http://127.0.0.1:2333/add \ + -H 'Content-Type: application/json' \ + -d '{"a":1,"b":2}' +``` + +Streaming Jobs return SSE: + +```bash +flowllm stream_demo query="Hi" repeat=3 interval=0.05 + +curl -N http://127.0.0.1:2333/stream_demo \ + -H 'Content-Type: application/json' \ + -d '{"query":"Hi","repeat":3,"interval":0.05}' +``` + +## Configuration + +The default configuration is in `flowllm/config/default.yaml`. You can use `.env` to override LLM, embedding, and other variables: + +```bash +cat > .env <<'EOF' +LLM_BACKEND=openai +LLM_MODEL_NAME=qwen3.7-plus +LLM_API_KEY=your_api_key +LLM_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 + +EMBEDDING_BACKEND=openai +EMBEDDING_MODEL_NAME=text-embedding-v4 +EMBEDDING_API_KEY=your_api_key +EMBEDDING_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 +EOF +``` + +You can also provide a YAML or JSON file: + +```bash +flowllm start config=/path/to/app.yaml +``` + +Minimal Job configuration: + +```yaml +service: + backend: http + +jobs: + add: + backend: base + description: "add two numbers" + steps: + - backend: add_step +``` + +Common Job backends: + +| backend | Description | +|--------------|------------------------------| +| `base` | Returns a JSON `Response` | +| `stream` | Returns an SSE stream | +| `background` | Background loop task | +| `cron` | Scheduled task | + +For more implementation details, see [Framework](./framework.md). diff --git a/docs/figure/flowllm-architecture.svg b/docs/figure/flowllm-architecture.svg new file mode 100644 index 0000000..90f97c1 --- /dev/null +++ b/docs/figure/flowllm-architecture.svg @@ -0,0 +1,178 @@ + + FlowLLM clean architecture diagram + A clean English architecture diagram for FlowLLM with Application at the center and Service, Jobs, Steps, Components, and Workspace around it. + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + FlowLLM Architecture + A config-driven runtime for exposing composable LLM workflows through service backends. + + + + + + + CORE + Application + loads config, wires components, + starts lifecycle, and runs jobs + + + Context + + Registry + + + + + + + ENTRY + Clients + flowllm CLI + HTTP client · MCP client + + + + + + + SERVICE + Backends + FastAPI HTTP routes + FastMCP tools + + + + + + + WORKFLOW + Jobs + base · stream + background · cron + + + + + + + EXECUTION + Steps + BaseStep.execute() + context mappings · dispatch + + + + + CONFIG + YAML + .env + CLI overrides + + + + + RUNTIME + RuntimeContext + Response + + + + + AS LLM + model wrapper + + + + + Embedding + model + store + + + + + Agent + wrapper + + + + + Workspace + metadata/session + + + Service surfaces are thin; workflow behavior lives in jobs and steps; reusable state lives in components. + diff --git a/docs/figure/flowllm-framework.svg b/docs/figure/flowllm-framework.svg new file mode 100644 index 0000000..ff4dcd2 --- /dev/null +++ b/docs/figure/flowllm-framework.svg @@ -0,0 +1,156 @@ + + FlowLLM framework architecture + A concise architecture diagram for FlowLLM showing entry clients, service backends, application wiring, jobs, steps, runtime context, reusable components, and workspace state. + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + FlowLLM Framework + Config-driven application wiring: Service exposes Jobs, Jobs run Steps, Steps use RuntimeContext and Components. + + + + + Entry Layer + CLI / Client + + + flowllm CLI + start / action + + + HTTP Client + JSON / SSE + + + MCP Client + tool calls + + action -> Job + + + + + + Service & Wiring + Service / App + + + HTTP Service + FastAPI + POST /job + + + MCP Service + FastMCP + non-stream tools + + + Application + config + wiring + lifecycle + run_job + + + Registry R + @R.register + backend lookup + + + + + + Execution Layer + Job / Step + + + Jobs + base · stream + background · cron + + + Steps + common / custom + BaseStep.execute() + + + RuntimeContext + data + Response + stream_queue + + + Response + JSON / SSE + + + + + + Components & State + Components + + + AS LLM + ChatModel + + + Embedding + model + store + + + Agent Wrapper + AgentScope / CC + + + Config + YAML + .env + + + Workspace + metadata / session + + + REQUEST SURFACE + SERVICE REGISTRATION + WORKFLOW EXECUTION + REUSABLE RUNTIME + Boot path: config -> registry lookup -> application wiring -> service endpoints. Request path: client -> service -> job -> steps -> response. + diff --git a/docs/figure/framework.png b/docs/figure/framework.png deleted file mode 100644 index 58ea076..0000000 Binary files a/docs/figure/framework.png and /dev/null differ diff --git a/docs/figure/logo.png b/docs/figure/logo.png index 6103649..8df9fc5 100644 Binary files a/docs/figure/logo.png and b/docs/figure/logo.png differ diff --git a/docs/figure/logo_old.png b/docs/figure/logo_old.png new file mode 100644 index 0000000..6103649 Binary files /dev/null and b/docs/figure/logo_old.png differ diff --git a/docs/index.md b/docs/index.md deleted file mode 100644 index b52f3c4..0000000 --- a/docs/index.md +++ /dev/null @@ -1,153 +0,0 @@ -

- FlowLLM Logo -

- -

- FlowLLM：让基于LLM的HTTP/MCP服务开发更简单
- 如果觉得有用，欢迎给个 ⭐ Star，您的支持是我们持续改进的动力 -

- -
- Python Version - PyPI Version - License - GitHub Stars -
- ---- - -## 📖 简介 - -FlowLLM 将 LLM/Embedding/vector_store 能力封装为 HTTP/MCP 服务，适用于 AI 对话助手、RAG 应用、工作流服务等场景，并可集成到支持 -MCP 的客户端工具中。 - -### 🏗️ 架构概览 - -

- FlowLLM Framework -

- -### ⭐ 核心特性 - -- **简单易用的 Op 开发**：继承 BaseOp 或 BaseAsyncOp 基类，实现业务逻辑即可。FlowLLM提供了延迟初始化的 LLM、Embedding - 模型和向量库，开发者只需通过 `self.llm`、`self.embedding_model`、`self.vector_store` 即可轻松使用这些资源。同时FlowLLM提供了完整的 - Prompt 模板管理能力，通过 `prompt_format()` 和 `get_prompt()` 方法进行格式化和使用。 - -- **灵活的 Flow 编排**：通过 YAML 配置文件将 Op 组合成 Flow，支持灵活的编排方式。`>>` 表示串行组合，`|` 表示并行组合，例如 - `SearchOp() >> (AnalyzeOp() | TranslateOp()) >> FormatOp()` 可构建复杂的工作流。定义输入输出 Schema 后，使用 - `flowllm config=your_config` 命令即可启动服务。 - -- **自动生成服务**：配置完成后，FlowLLM 会自动生成 HTTP、MCP 和 CMD 服务。HTTP 服务提供标准的 RESTful API，支持同步 JSON 响应和 - HTTP Stream 流式响应。MCP 服务会自动注册为 Model Context Protocol 工具，可集成到支持 MCP 的客户端中。CMD 服务支持命令行模式执行单个 - Op，适合快速测试和调试。 - -### 🌟 基于FlowLLM的应用 - -| 项目名 | 描述 | -|-----------------------------------------------|---------------| -| [ReMe](https://github.com/agentscope-ai/ReMe) | 面向智能体的记忆管理工具包 | - ---- - -## ⚡ 快速开始 - -### 📦 Step0 安装 - -#### 📥 From PyPI - -```bash -pip install flowllm -``` - -#### 🔧 From Source - -```bash -git clone https://github.com/flowllm-ai/flowllm.git -cd flowllm -pip install -e . -``` - -详细安装与配置方法请参考 [安装指南](zh/guide/installation.md)。 - -### ⚙️ 配置 - -创建 `.env` 文件，配置 API Key。你可以从 `example.env` 复制并修改： - -```bash -cp example.env .env -``` - -然后在 `.env` 文件中配置你的 API Key： - -```bash -FLOW_LLM_API_KEY=sk-xxxx -FLOW_LLM_BASE_URL=https://xxxx/v1 -FLOW_EMBEDDING_API_KEY=sk-xxxx -FLOW_EMBEDDING_BASE_URL=https://xxxx/v1 -``` - -详细配置说明请参考 [配置指南](zh/guide/config_guide.md)。 - -### 🛠️ Step1 构建Op - -```python -from flowllm.core.context import C -from flowllm.core.op import BaseAsyncOp -from flowllm.core.schema import Message -from flowllm.core.enumeration import Role - - -@C.register_op() -class SimpleChatOp(BaseAsyncOp): - async def async_execute(self): - query = self.context.get("query", "") - messages = [Message(role=Role.USER, content=query)] - response = await self.llm.achat(messages=messages) - self.context.response.answer = response.content.strip() -``` - -详细内容请参考 [简单 Op 指南](zh/guide/async_op_minimal_guide.md)、[LLM Op 指南](zh/guide/async_op_llm_guide.md) -和 [高级 Op 指南](zh/guide/async_op_advance_guide.md)（包含 Embedding、VectorStore 和并发执行等高级功能）。 - -### 📝 Step2 配置config - -以下示例展示如何构建一个 MCP（Model Context Protocol）服务。创建配置文件 `my_mcp_config.yaml`： - -```yaml -backend: mcp - -mcp: - transport: sse - host: "0.0.0.0" - port: 8001 - -flow: - demo_mcp_flow: - flow_content: MockSearchOp() - description: "Search results for a given query." - input_schema: - query: - type: string - description: "User query" - required: true - -llm: - default: - backend: openai_compatible - model_name: qwen3-30b-a3b-instruct-2507 - params: - temperature: 0.6 -``` - -### 🚀 Step3 启动 MCP 服务 - -```bash -flowllm \ - config=my_mcp_config \ - backend=mcp \ # 可选，覆盖config配置 - mcp.transport=sse \ # 可选，覆盖config配置 - mcp.port=8001 \ # 可选，覆盖config配置 - llm.default.model_name=qwen3-30b-a3b-thinking-2507 # 可选，覆盖config配置 -``` - -服务启动后可以参考[Client Guide](zh/guide/client_guide.md)来使用服务，可以直接获取模型所需要的tool_call。 diff --git a/docs/zh/contributing.md b/docs/zh/contributing.md new file mode 100644 index 0000000..c09f9c4 --- /dev/null +++ b/docs/zh/contributing.md @@ -0,0 +1,93 @@ +# 开源与贡献 + +FlowLLM 仓库：**https://github.com/flowllm-ai/flowllm** + +第一次本地运行请看 [快速开始](./quick_start.md)。修改运行时、Job、Step、组件、服务或配置前，请先看 [代码框架](./framework.md)。 + +## 本地开发 + +FlowLLM 要求 Python 3.11+。 + +```bash +python -m venv .venv +source .venv/bin/activate +pip install -e ".[dev]" +pre-commit install +``` + +需要完整可选依赖时： + +```bash +pip install -e ".[full]" +``` + +主要目录： + +| 路径 | 说明 | +|-------------------------------|---------------| +| `flowllm/` | Python 包源码 | +| `flowllm/config/default.yaml` | 默认配置 | +| `flowllm/components/` | 服务、客户端、Job、组件 | +| `flowllm/steps/` | Step 实现 | +| `flowllm/schema/` | 请求、响应和配置模型 | +| `tests/` | 单元和集成测试 | +| `docs/zh/` | 中文文档 | + +## 开发约定 + +遵循主链路： + +```text +CLI / Client -> Service -> Application -> Job -> Step -> Component +``` + +- 对外能力优先做成 Job，再由 HTTP 或 MCP Service 暴露。 +- 可复用基础设施放在 `flowllm/components/`，继承 `BaseComponent`。 +- 业务原子操作放在 `flowllm/steps/`，继承 `BaseStep`。 +- 新实现用 `@R.register("")` 注册，并确保模块被 `__init__.py` import。 +- 默认行为写入 `flowllm/config/default.yaml`，保持 `flowllm start` 可运行。 +- 修改用户可见行为时同步更新文档。 + +## 测试 + +提交前建议运行： + +```bash +pre-commit run --all-files +pytest +``` + +局部验证可以先跑： + +```bash +pytest tests/unit +pytest tests/integration +``` + +测试建议： + +- 修 Bug 时补回归测试。 +- 新增 Step、Job 或组件时覆盖核心路径和失败路径。 +- 修改配置解析、注册表、生命周期、服务暴露或流式输出时覆盖用户入口。 +- 依赖 LLM、embedding 或外部服务的测试，请在 PR 中说明所需环境。 + +## 提交与 PR + +建议使用 Conventional Commits： + +```text +(): +``` + +常用类型：`feat`、`fix`、`docs`、`test`、`refactor`、`chore`、`perf`、`style`。 + +示例： + +```text +feat(step): add reverse text demo +fix(config): preserve leading zero strings +docs(zh): update quick start +test(service): cover stream response +``` + +PR 标题也建议使用同样格式。较大改动请先在 Issue 中说明背景、目标行为、影响范围和测试计划。 diff --git a/docs/zh/framework.md b/docs/zh/framework.md new file mode 100644 index 0000000..b8cce82 --- /dev/null +++ b/docs/zh/framework.md @@ -0,0 +1,161 @@ +# FlowLLM 代码框架 + +FlowLLM 的运行链路： + +```text +CLI / Client -> Service -> Application -> Job -> Step -> Component +``` + +`Application` 从配置装配服务、组件和 Job；`Service` 暴露可服务的 Job；`Job` 顺序执行 Step；`Step` 通过 `RuntimeContext` +共享请求数据、响应和流式队列。 + +

+ FlowLLM framework architecture +

+ +## 目录 + +```text +flowllm/ + application.py # CLI 入口、Application 生命周期 + config/default.yaml # 默认 service / jobs / components + config/config_parser.py # config=、dot notation、环境变量展开 + components/ + base_component.py # BaseComponent、bind 依赖声明 + component_registry.py # 全局注册表 R + runtime_context.py # 单次 Job 执行上下文 + job/ # base / stream / background / cron + service/ # http / mcp + client/ # http / mcp + as_llm/ as_embedding/ # 模型封装 + embedding_store/ + agent_wrapper/ + steps/ + base_step.py # BaseStep、Ref、dispatch_steps + common/ # version、help、health_check、demo、add、stream_demo + schema/ # Request、Response、StreamChunk、配置模型 + enumeration/ # ComponentEnum、ChunkEnum +``` + +默认 workspace： + +```text +.flowllm/ +├── metadata/ +└── session/ +``` + +## 启动 + +`flowllm/application.py::main()` 处理三类命令： + +| action | 行为 | +|----------------|-------------------------------------------| +| `start` | 加载 `.env` 和配置，创建 `Application`，启动 Service | +| `find_flowllm` | 查找项目路径 | +| 其他 action | 通过 client 调用同名服务端 Job | + +配置解析支持： + +- 默认加载 `flowllm/config/default.yaml`。 +- `config=/path/to/app.yaml` 指定 YAML/JSON。 +- `service.port=8181` 这类 dot notation 覆盖配置。 +- `${VAR}` 和 `${VAR:-default}` 环境变量展开。 +- bool、数字、JSON list/dict、null 自动转换。 + +## Service + +HTTP service 把 Job 注册为接口： + +| Job 类型 | HTTP 行为 | +|-----------------------|-------------------------------------------| +| `BaseJob` | `POST /`，返回 `Response` JSON | +| `StreamJob` | `POST /`，返回 `text/event-stream` | +| `enable_serve: false` | 不注册 | + +MCP service 只暴露非流式、可服务的 Job。后台和定时任务默认不暴露。 + +## Registry + +FlowLLM 使用全局注册表 `R`： + +```python +from flowllm.components import R +from flowllm.steps import BaseStep + + +@R.register("reverse_step") +class ReverseStep(BaseStep): + async def execute(self): + text = self.context.get("text", "") + self.context.response.answer = text[::-1] + return self.context.response +``` + +注册 key 是 `(component_type, backend_name)`。同名 backend 可以存在于不同类型下，例如 `http` 同时是 service backend 和 +client backend。 + +注册在模块 import 时发生。新增 Step 或 Component 后，要确保所在包的 `__init__.py` import 了新模块，否则配置里的 `backend` +找不到实现。 + +## Component + +长期存在的基础设施继承 `BaseComponent`。常用生命周期： + +- `start()`：解析依赖并调用 `_start()`。 +- `close()`：调用 `_close()` 并反序关闭自有组件。 +- `restart()`：关闭后重启。 + +组件依赖用 `BaseComponent.bind()` 声明，`Application` 会按依赖拓扑顺序启动： + +```python +self.embedding_store = self.bind("default", BaseEmbeddingStore, optional=False) +``` + +## Job + +Job 是外部可调用能力或后台流程的编排单元，配置在 `jobs:` 下。 + +- `BaseJob`：每次调用创建 `RuntimeContext`，顺序执行 Step，返回 `context.response`。 +- `StreamJob`：Step 向 `stream_queue` 写 `StreamChunk`，服务层输出 SSE。 +- `BackgroundJob`：按 `interval` 循环执行。 +- `CronJob`：按 cron 表达式定时执行。 + +Job 配置示例： + +```yaml +jobs: + reverse: + backend: base + description: "reverse text" + steps: + - backend: reverse_step +``` + +## Step + +Step 是工作流原子操作。它读写 `RuntimeContext`： + +| 字段 | 说明 | +|----------------|--------------| +| `data` | 请求参数和步骤间共享数据 | +| `response` | 最终返回值 | +| `stream_queue` | 流式输出队列 | + +流式输出： + +```python +await self.context.add_stream_string(text, ChunkEnum.CONTENT) +``` + +Step 配置还支持 `input_mapping`、`output_mapping`、`prompt_dict`、`dispatch_steps` 和 `language`。访问应用组件时使用 +`BaseStep.Ref`；默认已提供 `as_llm` 和可选的 `agent_wrapper`。 + +## 默认能力 + +`flowllm/config/default.yaml` 内置： + +- `version`、`health_check`、`help` +- `demo`、`add` +- `stream_demo` +- 默认 `as_llm`、`as_embedding`、`embedding_store`、`agent_wrapper` diff --git a/docs/zh/guide/async_op_advance_guide.md b/docs/zh/guide/async_op_advance_guide.md deleted file mode 100644 index 2b1e5de..0000000 --- a/docs/zh/guide/async_op_advance_guide.md +++ /dev/null @@ -1,122 +0,0 @@ -## 异步 Op 高级功能指南 - -这是一个完整的示例，展示如何在异步 Op 中使用 Embedding、VectorStore 和并发执行等高级功能。 - ---- - -### 完整示例 - -#### Op 文件：`rag_search_op.py` - -```python -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.async_get_embeddings(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()) -``` - ---- - -### 关键要点 - -1. **Embedding 使用**： - - 通过 `self.embedding_model.get_embeddings_async(text)` 生成文本向量 - - 支持单个文本或文本列表的批量处理 - - 返回向量列表（单个文本返回单个向量） - -2. **VectorStore 使用**： - - 通过 `self.vector_store.async_search(query, workspace_id, top_k)` 进行语义搜索 - - 通过 `self.vector_store.async_insert(nodes, workspace_id)` 插入向量节点 - - 返回 `VectorNode` 列表，每个节点包含 `content`、`vector`、`metadata` 等信息 - -3. **并发执行**： - - 使用 `self.submit_async_task(fn, *args, **kwargs)` 提交异步任务 - - 使用 `await self.join_async_task(timeout=..., return_exceptions=True)` 等待所有任务完成 - - 支持超时控制和异常处理 - -4. **VectorNode 结构**： - - `unique_id`: 唯一标识符 - - `workspace_id`: 工作空间 ID - - `content`: 文本内容 - - `vector`: 向量嵌入（可选） - - `metadata`: 元数据字典 - -5. **配置方式**： - - 可以在 Op 初始化时传入 `embedding_model` 和 `vector_store` 参数 - - 支持字符串配置名（从服务配置中读取）或直接传入实例 - ---- - -### 参考示例 - -实际项目中的示例： - -- `flowllm/flowllm/core/embedding_model/` - Embedding 模型实现 -- `flowllm/flowllm/core/vector_store/` - VectorStore 实现 -- `flowllm/flowllm/core/op/base_async_op.py` - 异步 Op 基类，包含并发执行方法 - diff --git a/docs/zh/guide/async_op_llm_guide.md b/docs/zh/guide/async_op_llm_guide.md deleted file mode 100644 index 237a023..0000000 --- a/docs/zh/guide/async_op_llm_guide.md +++ /dev/null @@ -1,170 +0,0 @@ -## 异步 Op 中使用 LLM 和 Prompt 指南 - -这是一个完整的示例，展示如何在异步 Op 中使用 LLM 和 Prompt，包括 prompt 文件创建、格式化、LLM 调用以及响应处理。 - ---- - -### 完整示例 - -#### 1. Prompt 文件：`qa_op_prompt.yaml` - -**文件命名规则**：Op 文件名 `qa_op.py` → Prompt 文件名 `qa_op_prompt.yaml`（放在同一目录） - -```yaml -# Prompt 模板，使用 {variable_name} 进行变量替换 -qa_prompt: | - # Task - Please answer the following question: - {question} - - # Requirements - 1. Answer should be accurate and professional - 2. If you don't know, please be honest - 3. Answer should be concise and clear - 4. For complex questions, provide detailed explanations - -# 多语言支持：通过 _zh 等后缀区分 -qa_prompt_zh: | - # 任务 - 请回答以下问题： - {question} - - # 要求 - 1. 回答要准确、专业 - 2. 如果不知道，请诚实说明 - 3. 回答要简洁明了 - 4. 如果问题复杂，可以提供详细解释 -``` - -#### 2. Op 文件：`qa_op.py` - -```python -from flowllm.core.context import C -from flowllm.core.op import BaseAsyncOp -from flowllm.core.schema import Message -from flowllm.core.enumeration import Role -from flowllm.main import FlowLLMApp -import asyncio -import json - - -@C.register_op() -class QAOp(BaseAsyncOp): - file_path: str = __file__ # 必须设置，用于自动查找 prompt 文件 - - async def async_execute(self): - """执行问答逻辑""" - # 1. 读取输入 - question = self.context.get("question", "") - assert question, "question 不能为空" - - # 2. 渲染 prompt 模板 - # prompt_format 支持：变量替换：{question} 会被替换为实际值 - prompt = self.prompt_format("qa_prompt", question=question) - - # 3. 构建 messages - # 可以添加 system prompt、对话历史等 - messages = [Message(role=Role.USER, content=prompt)] - - # 3.5. 计算 token 数量（可选） - # 使用 token_count 方法可以准确计算消息和工具的 token 数量 - # 支持多种后端（base、openai、hf 等），配置在 llm 的 token_count 中 - token_num = self.token_count(messages) - print(f"Input tokens: {token_num}") - - # 4. 定义 callback_fn 处理响应（可选） - # 用于从 LLM 响应中提取或转换特定格式的内容 - def parse_response(message: Message) -> str: - """处理 LLM 响应""" - content = message.content.strip() - - # 如果需要 JSON 格式，尝试从响应中提取 JSON - if "```json" in content: - start = content.find("```json") + 7 - end = content.find("```", start) - json_str = content[start:end].strip() - try: - # 验证 JSON 格式 - json.loads(json_str) - return json_str - except: - pass - - return content - - # 5. 调用 LLM - # achat 支持： - # - messages: 消息列表 - # - callback_fn: 响应处理函数（可选） - # - default_value: 失败时的默认值（可选） - # - enable_stream_print: 流式输出（可选） - # - tools: 工具调用（可选） - response = await self.llm.achat( - messages=messages, - callback_fn=parse_response, - default_value="抱歉，我无法回答这个问题") - - # 6. 获取响应内容 - # 如果使用了 callback_fn，response 是 callback_fn 的返回值 - # 否则 response 是 Message 对象，通过 response.content 获取内容 - if isinstance(response, Message): - answer = response.content.strip() - else: - answer = response - - # 7. 写回输出 - self.context.response.answer = answer - self.context["answer"] = answer # 也可以写到 context 供其他 Op 使用 - - return answer - - -# 使用示例 -async def main(): - async with FlowLLMApp() as app: - # 基本使用 - result = await QAOp().async_call(question="什么是 Python？") - print(f"答案：{result}") - - # 使用自定义 LLM - result2 = await QAOp(llm="qwen3_30b_instruct").async_call( - question="什么是异步编程？" - ) - print(f"答案：{result2}") - - # 使用条件标志和 callback_fn - result3 = await QAOp().async_call( - question="解释一下 Python 的装饰器", - need_json=True - ) - print(f"答案：{result3}") - - -if __name__ == "__main__": - asyncio.run(main()) -``` - ---- - -### 关键要点 - -1. **Prompt 文件命名**：`xxx_op.py` → `xxx_prompt.yaml`，放在同一目录 -2. **变量替换**：在 prompt 中使用 `{variable_name}`，调用 `prompt_format` 时传入对应变量 -3. **多语言支持**：通过 `_zh`、`_en` 等后缀支持多语言 prompt，系统自动选择 -4. **构建 Messages**：使用 `Message(role=Role.USER/SYSTEM/ASSISTANT, content=...)` -5. **Token 计数**：使用 `self.token_count(messages, tools=None)` 计算 token 数量，支持多种后端（base、openai、hf 等） -6. **调用 LLM**：使用 `await self.llm.achat(messages=messages, ...)` -7. **处理响应**：使用 `callback_fn` 处理或转换响应，返回处理后的结果 -8. **应用上下文**：必须在 `FlowLLMApp()` 上下文里调用 -9. **file_path**：Op 类中必须设置 `file_path = __file__`，用于自动查找 prompt 文件 - ---- - -### 参考示例 - -实际项目中的示例： - -- `flowllm/flowllm/gallery/gen_system_prompt_op.py` - 使用 prompt 和 callback_fn 的完整示例 -- `flowllm/flowllm/gallery/chat_op.py` - 简单的聊天 Op -- `flowllm/flowllm/gallery/mock_search_op.py` - 使用 callback_fn 解析 JSON 的示例 - diff --git a/docs/zh/guide/async_op_minimal_guide.md b/docs/zh/guide/async_op_minimal_guide.md deleted file mode 100644 index dd19eb1..0000000 --- a/docs/zh/guide/async_op_minimal_guide.md +++ /dev/null @@ -1,48 +0,0 @@ -## 编写异步 Op 最简示例 - -这是一个最简单的异步 Op 示例，展示核心功能：注册、读取输入、写输出、如何跑Op测试。 - ---- - -### 完整示例 - -```python -from flowllm.core.context import C -from flowllm.core.op import BaseAsyncOp -from flowllm.main import FlowLLMApp -import asyncio - -@C.register_op() -class SimpleChatOp(BaseAsyncOp): - async def async_execute(self): - # 1. 读取输入 - query = self.context.get("query", "") - - # 2. 处理逻辑（这里只是简单示例） - result = f"你问的是：{query}" - - # 3. 写输出到 context（供其他 Op 使用） - self.context["result"] = result - - # 4. 写输出到 context.response（整个 pipeline 的最终输出） - self.context.response.answer = result - -# 使用 -async def main(): - async with FlowLLMApp() as app: - result = await SimpleChatOp().async_call(query="你好？") - print(result) - -if __name__ == "__main__": - asyncio.run(main()) -``` - ---- - -### 要点说明 - -1. **注册 Op**：使用 `@C.register_op()` 装饰器 -2. **读取输入**：`query = self.context.get("query", "")` -3. **写输出到 context**：`self.context["result"] = result`（供其他 Op 使用） -4. **写输出到 response**：`self.context.response.answer = result`（整个 pipeline 的最终输出） -5. **使用**：需要在 `FlowLLMApp()` 上下文里调用 diff --git a/docs/zh/guide/async_tool_op_guide.md b/docs/zh/guide/async_tool_op_guide.md deleted file mode 100644 index b3d6465..0000000 --- a/docs/zh/guide/async_tool_op_guide.md +++ /dev/null @@ -1,163 +0,0 @@ -## 异步 Tool Op 指南 (BaseAsyncToolOp) - -这是一个完整的示例，展示如何使用 `BaseAsyncToolOp` 编写基于 Schema 的工具类异步 Op。相比 `BaseAsyncOp`，`BaseAsyncToolOp` 提供了 Schema 驱动的输入/输出自动处理能力。 - ---- - -### 完整示例 - -#### Op 文件：`echo_tool_op.py` - -```python -import asyncio - -from flowllm.core.context import C -from flowllm.core.op import BaseAsyncToolOp -from flowllm.core.schema import ToolCall -from flowllm.main import FlowLLMApp - - -@C.register_op() -class EchoToolOp(BaseAsyncToolOp): - """简单的 Echo 工具 Op，展示 BaseAsyncToolOp 的基本用法""" - - def build_tool_call(self) -> ToolCall: - """定义工具的 Schema：描述、输入参数、输出参数""" - return ToolCall( - **{ - "description": "Echo input text, returns the same text as output", - "input_schema": { - "text": { - "type": "str", - "description": "text to echo", - "required": True, - }, - }, - } - ) - - async def async_execute(self): - """执行工具逻辑""" - # 1. 读取输入（由框架在执行前已从 context 填充到 self.input_dict） - text = self.input_dict.get("text", "") - - # 2. 处理逻辑 - result = f"Echo: {text}" - - # 3. 设置输出（会在执行后自动写回 context） - # 单输出时可以省略 key 参数 - self.set_output(result) - - # 或者使用 set_outputs 设置多个输出： - # self.set_outputs(result=result, length=len(text)) - - -# 使用示例 -async def main(): - async with FlowLLMApp() as app: - # 基本使用 - result = await EchoToolOp().async_call(text="hello") - print(f"结果：{result}") - - # 执行后，context 中会自动写入： - # - "echo_tool_op_result": "Echo: hello"（默认输出键名） - # - context.response.answer = "Echo: hello"（如果 save_answer=True） - - -if __name__ == "__main__": - asyncio.run(main()) -``` - ---- - -### 核心特性 - -#### 1. Schema 驱动的输入/输出 - -- **定义 Schema**：通过 `build_tool_call()` 返回 `ToolCall`，描述工具的： - - `description`：工具描述 - - `input_schema`：输入参数定义（类型、描述、是否必需） - - `output_schema`：输出参数定义（可选，不提供时会自动生成） - -- **自动输入读取**：运行时会根据 `input_schema` 自动从 `context` 读取输入到 `self.input_dict` -- **自动输出写回**：执行完成后会把 `self.output_dict` 自动写回 `context`（键名来自 `output_schema`） - -#### 2. 便捷的输出设置方法 - -- **`set_output(value, key="")`**：设置单个输出值 - - 单输出时可以省略 `key` 参数 - - 多输出时必须指定 `key` - -- **`set_outputs(**kwargs)`**：设置多个输出值 - -#### 3. 自动应答写回 - -- `save_answer=True`（默认）时，会将主输出写入 `context.response.answer` -- 单输出：直接写入字符串 -- 多输出：写入 JSON 字符串 - -#### 4. 多实例支持 - -- `tool_index` 参数可以在上下文键后追加后缀（如 `result.1`） -- 便于在同一上下文中多次调用同一工具 - -#### 5. 键名映射 - -- `input_schema_mapping`：将 Schema 输入键映射到 context 键 -- `output_schema_mapping`：将 Schema 输出键映射到 context 键 -- 当 Schema 键名与 context 键名不一致时很有用 - ---- - -### 与 BaseAsyncOp 的对比 - -| 能力 | BaseAsyncOp | BaseAsyncToolOp | -|-----------|---------------------------------------------|----------------------------------------------------------------------| -| 输入读取 | 手动从 `self.context[...]` 读取 | 根据 `input_schema` 自动填充 `self.input_dict` | -| 输出写回 | 手动写 `self.context[...]` / `response.answer` | `set_output(s)` 写入 `self.output_dict`，自动写回上下文；可自动写 `response.answer` | -| Schema 描述 | 无内建 | 通过 `build_tool_call()` 描述输入输出与说明 | -| 多实例区分 | 自行约定 | 通过 `tool_index` 自动为上下文键追加后缀 | -| 键名映射 | 自行处理 | `input_schema_mapping` / `output_schema_mapping` 内建映射能力 | - -**选择建议**： -- 当你的 Op 本质是"工具"（有清晰的输入/输出契约、希望自动化 IO 与应答写回），应优先选择 `BaseAsyncToolOp` -- 否则选择通用的 `BaseAsyncOp` 并手动处理 IO - ---- - -### 关键要点 - -1. **必须实现的方法**： - - `build_tool_call()`：返回 `ToolCall` 对象，定义工具的 Schema - - `async_execute()`：实现工具的业务逻辑 - -2. **输入处理**： - - 从 `self.input_dict` 读取输入（框架已自动填充） - - 如果 `required=True` 的输入在 context 中找不到，会抛出异常 - -3. **输出处理**： - - 使用 `self.set_output(value, key="")` 或 `self.set_outputs(**kwargs)` 设置输出 - - 框架会自动将输出写回 context 和 response.answer - -4. **输出键名规则**： - - 如果定义了 `output_schema`，使用 Schema 中定义的键名 - - 如果没有定义 `output_schema`，自动生成 `${short_name}_result` 作为键名 - -5. **多实例使用**： - - 设置 `tool_index` 参数区分多个实例 - - 输出键会自动追加 `.{tool_index}` 后缀 - -6. **键名映射**： - - 使用 `input_schema_mapping` 和 `output_schema_mapping` 进行键名映射 - - 例如：`input_schema_mapping={"text": "query"}` 表示从 context 的 `query` 键读取，但映射到 `input_dict` 的 `text` 键 - ---- - -### 参考示例 - -实际项目中的示例： - -- `flowllm/flowllm/gallery/execute_code_op.py` - 代码执行工具的完整示例 -- `flowllm/flowllm/core/op/base_async_tool_op.py` - BaseAsyncToolOp 基类实现 -- `flowllm/flowllm/core/op/base_mcp_op.py` - MCP 工具调用示例（继承自 BaseAsyncToolOp） - diff --git a/docs/zh/guide/client_guide.md b/docs/zh/guide/client_guide.md deleted file mode 100644 index 145d2c2..0000000 --- a/docs/zh/guide/client_guide.md +++ /dev/null @@ -1,141 +0,0 @@ -### FlowLLM Client 使用指南 - -本指南介绍三种客户端调用模式： -- **HTTP**：同步返回结果 -- **HTTP-Stream**：流式返回片段 -- **MCP**：基于 SSE 的 Model Context Protocol 客户端交互 - - -## 快速开始 - -**前置条件**：确保服务端已启动 -- HTTP/HTTP-Stream 服务：默认端口 8002 -- MCP SSE 服务：默认端口 8001，路径 `/sse` - -**核心类**： -- `flowllm.core.utils.HttpClient` -- `flowllm.core.utils.FastMcpClient` - - -## 完整示例 - -示例包含三种模式的调用： -- HTTP 同步调用 Flow -- HTTP-Stream 流式调用 Flow -- MCP 工具枚举与调用 - -```python -import asyncio -import json - -from flowllm.core.utils import HttpClient, FastMcpClient - - -async def run_http_examples(): - async with HttpClient("http://0.0.0.0:8002") as client: - # 健康检查 - health = await client.health_check() - print("health:", json.dumps(health, indent=2, ensure_ascii=False)) - - # 列出可用 endpoints - schema = await client.list_endpoints() - print("endpoints keys:", list(schema.get("paths", {}).keys())) - - # 同步执行 Flow - resp = await client.execute_flow("demo_http_flow", query="阿里巴巴前景如何？") - print("sync answer:", getattr(resp, "answer", None)) - - # 流式执行 Flow - async for chunk in client.execute_stream_flow("demo_stream_http_flow", query="what is ai"): - print(f"[{chunk.get('type','answer')}] {chunk.get('content','')}") - - -async def run_mcp_example(): - config = { - "type": "sse", - "url": "http://0.0.0.0:8001/sse", - "headers": {}, - "timeout": 30.0, - } - async with FastMcpClient("test_client", config, max_retries=3) as client: - tool_calls = await client.list_tool_calls() - if tool_calls: - for tool_call in tool_calls: - # 模型需要的tool call格式（qwen） - print(tool_call.simple_input_dump()) - - # 调用工具 - result = await client.call_tool("demo_mcp_flow", {"query": "阿里巴巴前景怎么样？"}) - if result.is_error: - print("MCP tool error:", result.content[0].text if result.content else "Unknown") - else: - print("MCP tool ok:", result.content[0].text if result.content else None) - if result.structured_content: - print(json.dumps(result.structured_content, indent=2, ensure_ascii=False)) - - -async def main(): - await run_http_examples() - await run_mcp_example() - - -if __name__ == "__main__": - asyncio.run(main()) -``` - - -## HTTP 模式 - -- **初始化**：`HttpClient(base_url)` 使用上下文管理器管理会话 -- **健康检查**：`health_check()` 返回服务端状态字典 -- **列出接口**：`list_endpoints()` 返回 OpenAPI schema，通过 `schema.get("paths", {})` 获取路径列表 -- **执行 Flow**：`execute_flow(flow_name, **kwargs)` 第一个参数为 Flow 名称，关键字参数作为请求体，返回响应对象（通常包含 `answer` 属性） - - -## HTTP-Stream 模式 - -- **流式执行**：`execute_stream_flow(flow_name, **kwargs)` 返回异步迭代器 -- **消费片段**：使用 `async for chunk in ...` 迭代，每个 `chunk` 包含 `type` 和 `content` 字段 -- **片段类型**：`type` 标识片段类型（如 "answer"），`content` 为实际内容 - - -## MCP 模式 - -MCP 通过 SSE 提供工具列表与工具调用的抽象。 - -- **初始化**：`FastMcpClient(client_name, config, max_retries)` 配置包含： - - `type`: "sse" - - `url`: SSE 服务地址 - - `headers`: 可选请求头 - - `timeout`: 超时时间（秒） - -- **列出工具**：`list_tool_calls()` 返回工具签名列表，每个工具对象包含名称、参数等信息 - -- **调用工具**：`call_tool(tool_name, arguments)` 返回 `CallToolResult`： - - `content`: 消息内容列表，通过 `result.content[0].text` 访问 - - `structured_content`: 结构化返回（dict/list） - - `is_error`: 是否出错（bool） - - -## 运行与调试 - -运行测试脚本（需先启动服务端）： - -```bash -python -m flowllm.tests.http_client_test -python -m flowllm.tests.mcp_client_test -``` - -**常见问题**： -- 确认服务端地址与端口与客户端配置一致 -- 容器/代理环境需检查 CORS 和 SSE 设置 -- 根据网络和模型响应速度调整 `timeout` - - -## 模式选择 - -- **HTTP**：短响应、一次性获取结果、执行时间可控 -- **HTTP-Stream**：长响应、需要实时展示、改善交互体验 -- **MCP**：工具调用与编排、需要工具列表与结构化结果 - - diff --git a/docs/zh/guide/cmd_service_guide.md b/docs/zh/guide/cmd_service_guide.md deleted file mode 100644 index 780a93c..0000000 --- a/docs/zh/guide/cmd_service_guide.md +++ /dev/null @@ -1,50 +0,0 @@ -## CMD 服务快速上手 - -本指南演示如何在 FlowLLM 中使用命令行模式执行单个 Op，覆盖：编写 Op 和直接运行命令。 - ---- - -### 一、编写 Op - -CMD 工具的业务逻辑来源于 Flow 中的 Op。基类见 `flowllm/core/op/base_op.py`，异步基类为 `BaseAsyncOp`： - -```python -from flowllm.core.context import C -from flowllm.core.op import BaseAsyncOp - - -@C.register_op() -class EchoOp(BaseAsyncOp): - file_path: str = __file__ - - async def async_execute(self): - text = self.context.get("text", "") - self.context.response.answer = f"echo: {text}" -``` - -要点： - -- 类名需以 `Op` 结尾（框架有断言） -- 用 `@C.register_op()` 注册后方可在命令行中引用 -- 对话型 Op 通常继承 `BaseAsyncOp` 并实现 `async_execute` -- 可通过 `self.context` 读写上下文，`self.llm` 调用模型，`self.prompt_format()` 绑定同名 `*_prompt.yaml` -- 参数通过 `self.context.get()` 获取 - ---- - -### 二、直接运行命令 - -确保已安装 FlowLLM，并在`.env`中设置模型相关环境变量，可直接参考项目根的示例文件 `example.env`： - -```bash -export FLOW_LLM_API_KEY="sk-xxxx" -export FLOW_LLM_BASE_URL="https://xxxx/v1" -export FLOW_EMBEDDING_API_KEY="sk-xxxx" -export FLOW_EMBEDDING_BASE_URL="https://xxxx/v1" -``` - -运行示例： - -```bash -flowllm backend=cmd cmd.flow="EchoOp()" cmd.params.text="Hello World" -``` diff --git a/docs/zh/guide/config_guide.md b/docs/zh/guide/config_guide.md deleted file mode 100644 index 1091467..0000000 --- a/docs/zh/guide/config_guide.md +++ /dev/null @@ -1,140 +0,0 @@ -## 配置文件 default.yaml 指南 - -本文介绍如何使用 `flowllm/flowllm/config/default.yaml` 进行服务配置。你将看到一个可直接复制使用的完整示例，并逐段讲解每一项配置的作用与常见用法。 - -> **注意**：关于 Flow（流程）配置的详细说明，请参考 [Flow 配置指南](flow_guide.md)。 - ---- - -### 一、完整配置示例（可直接使用） - -将下面内容保存为 `flowllm/flowllm/config/default.yaml`，即可启动一个包含 HTTP 与 MCP 的最小可用环境： - -```yaml -backend: http -thread_pool_max_workers: 128 - -mcp: - transport: sse - host: "0.0.0.0" - port: 8001 - -http: - host: "0.0.0.0" - port: 8002 - -# Flow 配置请参考 flow_guide.md -# flow: -# demo_http_flow: -# flow_content: GenSystemPromptOp() >> ChatOp() -# description: "ai chat assistant" -# input_schema: -# query: -# type: string -# description: "user query" -# required: true - -llm: - default: - backend: openai_compatible - model_name: qwen3-30b-a3b-instruct-2507 - params: - temperature: 0.6 - token_count: # 可选，配置 token 计数后端 - backend: openai # 支持 base、openai、hf 等 - model_name: qwen3-30b-a3b-instruct-2507 - - qwen3_30b_instruct: - backend: openai_compatible - model_name: qwen3-30b-a3b-instruct-2507 - - qwen3_30b_thinking: - backend: openai_compatible - model_name: qwen3-30b-a3b-thinking-2507 - - qwen3_235b_instruct: - backend: openai_compatible - model_name: qwen3-235b-a22b-instruct-2507 - - qwen3_235b_thinking: - backend: openai_compatible - model_name: qwen3-235b-a22b-thinking-2507 - - qwen3_80b_instruct: - backend: openai_compatible - model_name: qwen3-next-80b-a3b-instruct - - qwen3_80b_thinking: - backend: openai_compatible - model_name: qwen3-next-80b-a3b-thinking - - qwen3_max_instruct: - backend: openai_compatible - model_name: qwen3-max - - qwen25_max_instruct: - backend: openai_compatible - model_name: qwen-max-2025-01-25 - -embedding_model: - default: - backend: openai_compatible - model_name: text-embedding-v4 - params: - dimensions: 1024 - -vector_store: - default: - backend: elasticsearch - embedding_model: default -# params: -# hosts: "http://localhost:9200" -``` - ---- - -### 二、配置结构与含义 - -- backend - - 决定应用主服务形态。当前示例使用 `http`，表示以 HTTP 服务为主入口（配合 `http` 段）。 - -- thread_pool_max_workers - - 线程池并发上限。异步场景下用于提交同步任务的线程数量上限，通常按 CPU 与任务特性调优。 - -- mcp - - `transport`: 传输方式，示例使用 `sse`（Server-Sent Events）。 - - `host`/`port`: MCP 服务监听地址与端口。若你有外部 MCP 工具或希望通过 MCP 集成，保持该段开启。 - - `flow` 相关：在 MCP 模式下，各 Flow 的 `input_schema` 为**必填**，需精确描述输入参数（类型、必填与否、描述等），以便 MCP 客户端进行参数校验与能力展示。 - -- http - - `host`/`port`: HTTP 服务监听地址与端口。默认 `0.0.0.0:8002` 便于容器与本机调试。 - - `flow` 相关：在 HTTP 模式（包括 `stream: true` 的流式模式）下，`input_schema` **可选**；建议填写以获得更好的入参校验与自动文档生成体验。 - -- flow - - 定义一组可被调用的流程（Flow）。每个 Flow 都由若干 Op 组合而成，可顺序（`>>`）或并行（`|`）组织。 - - **详细说明请参考 [Flow 配置指南](flow_guide.md)**，该文档包含完整的语法说明、示例和最佳实践。 - -- llm - - 语言模型配置集合。 - - `default`: 全局默认模型，未指定时各 Op 使用该项。 - - 任意命名条目（如 `qwen3_30b_instruct`）可作为备用或按需切换的模型配置。 - - 字段说明： - - `backend`: LLM 后端类型（如 `openai_compatible`）。 - - `model_name`: 具体模型名称。 - - `params`: 传递给后端的参数，如 `temperature` 等。 - - `token_count`（可选）：Token 计数配置，用于 `self.token_count()` 方法。 - - `backend`: Token 计数后端类型，支持 `base`（基于字符数估算）、`openai`（使用 tiktoken）、`hf`（使用 HuggingFace tokenizer）等。 - - `model_name`: Token 计数使用的模型名称（通常与 LLM 模型名称相同或对应）。 - - `params`: 传递给 Token 计数后端的参数（如 `use_mirror: true` 用于 HuggingFace）。 - -- embedding_model - - 嵌入模型配置集合。 - - `default`: 默认嵌入模型，向量检索与相关任务会引用此项。 - - 字段说明类似 `llm`：`backend`、`model_name`、`params`。 - -- vector_store - - 向量库配置集合。 - - `default`: 默认向量库配置。 - - `backend`: 向量库后端（如 `elasticsearch`）。 - - `embedding_model`: 指定使用哪一个嵌入模型条目（此处引用上面的 `embedding_model.default`）。 - - `params`: 后端连接参数（注释示例给出 `hosts`）。 diff --git a/docs/zh/guide/contribution.md b/docs/zh/guide/contribution.md deleted file mode 100644 index ac18878..0000000 --- a/docs/zh/guide/contribution.md +++ /dev/null @@ -1,149 +0,0 @@ -# 贡献到 FlowLLM - -## 欢迎！🎉 - -感谢对 FlowLLM 项目的关注和支持。我们欢迎并鼓励来自社区的贡献，无论是修复错误、添加新功能、改进文档还是分享想法。 - -## 如何贡献 - -### 1. 检查现有计划和问题 - -在开始贡献之前，请查看： -- **查看 [Issues](https://github.com/flowllm-ai/flowllm/issues)** 了解现有的开发任务 -- **如果存在相关问题**：请在 issue 下评论，表达您的参与意愿 -- **如果不存在相关问题**：请创建新 issue 描述您的更改或功能 - -### 2. 提交信息格式 - -FlowLLM 遵循 [Conventional Commits](https://www.conventionalcommits.org/) 规范。 - -**格式：** -``` -(): -``` - -**类型：** -- `feat:` 新功能 -- `fix:` 错误修复 -- `docs:` 仅文档更改 -- `style:` 不影响代码含义的更改 -- `refactor:` 代码重构 -- `perf:` 性能优化 -- `test:` 测试相关 -- `chore:` 构建或工具更改 - -**示例：** -```bash -feat(op): add new async tool op -fix(flow): resolve context sharing issue -docs(guide): update flow configuration guide -``` - -### 3. 代码开发指南 - -#### a. 提交前检查 - -安装并运行 pre-commit 钩子： - -```bash -pip install pre-commit -pre-commit install -pre-commit run --all-files -``` - -#### b. 关于代码中的 Import - -FlowLLM 遵循**懒加载导入原则**： - -- **推荐做法**：仅在实际使用时导入模块 - ```python - def some_function(): - import openai - # 在此处使用 openai 库 - ``` - -#### c. 单元测试 - -- 所有新功能都应包含适当的单元测试 -- 提交 PR 前确保现有测试通过 -- 运行测试：`pytest tests` - -#### d. 文档 - -- 为新功能更新相关文档 -- 在适当的地方包含代码示例 -- 如果更改影响用户功能，请更新 README.md - -## 贡献类型 - -### 添加新的 Op - -FlowLLM 的核心是 Op（算子）体系。添加新 Op 时： - -1. **选择合适的基类**： - - `BaseOp`：基础同步/异步 Op - - `BaseAsyncOp`：异步 Op（推荐） - - `BaseAsyncToolOp`：工具调用 Op - - `BaseMcpOp`：MCP 工具 Op - - `BaseRayOp`：分布式并行 Op - -2. **实现位置**： - - **核心 Op**：放在 `flowllm/core/op/` 目录 - - **示例 Op**：放在 `flowllm/gallery/` 目录 - - **扩展 Op**：放在 `flowllm/extension/` 目录 - -3. **实现要求**： - - 实现 `execute()`（同步）或 `async_execute()`（异步）方法 - - 遵循 Op 组合规范（支持 `>>`、`|` 运算符） - - 提供清晰的文档和示例 - -### 添加新的示例 - -请将示例添加到 `flowllm/gallery/` 目录，并附上清晰的注释说明。 - -示例结构： -``` -flowllm/gallery/ -└── your_op.py # Op 实现 - └── your_prompt.yaml # 可选的提示词配置 -``` - -### 添加新的服务类型 - -如需添加新的服务类型（如新的协议支持），请： -1. 在 `flowllm/core/service/` 下实现 `BaseService` 的子类 -2. 更新配置解析逻辑 -3. 添加相应的文档和示例 - -## Do's and Don'ts - -### ✅ DO - -- **从小处着手**：从小的、可管理的贡献开始 -- **及早沟通**：在实现主要功能之前进行讨论 -- **编写测试**：确保代码经过充分测试 -- **添加代码注释**：帮助他人理解贡献内容 -- **遵循提交约定**：使用约定式提交消息 -- **保持尊重**：遵守我们的行为准则 - -### ❌ DON'T - -- **不要用大型 PR 让我们措手不及**：大型 PR 难以审查，请先开启 issue 讨论 -- **不要忽略 CI 失败**：修复持续集成标记的任何问题 -- **不要混合关注点**：保持 PR 专注于单一功能的实现或修复 -- **不要忘记更新测试**：功能的更改应反映在测试中 -- **不要破坏现有 API**：在可能的情况下保持向后兼容性 -- **不要添加不必要的依赖项**：保持核心库轻量级 -- **不要绕过懒加载导入原则**：确保 FlowLLM 在导入阶段不至于臃肿 - -## 获取帮助 - -如果需要帮助或有疑问： - -- 💬 开启一个 [Discussion](https://github.com/flowllm-ai/flowllm/discussions) -- 🐛 通过 [Issues](https://github.com/flowllm-ai/flowllm/issues) 报告错误 - ---- - -感谢为 FlowLLM 做出贡献！🚀 - diff --git a/docs/zh/guide/file_tool_op_guide.md b/docs/zh/guide/file_tool_op_guide.md deleted file mode 100644 index 1c4c31e..0000000 --- a/docs/zh/guide/file_tool_op_guide.md +++ /dev/null @@ -1,54 +0,0 @@ -# File Tool Op 介绍 - -File Tool Op 是 FlowLLM 框架提供的文件操作工具集，基于 `BaseAsyncToolOp` 实现，为 LLM 提供了丰富的文件系统操作能力。 - -## 功能概览 - -File Tool Op 包含 13 个操作，详细信息如下： - -| 类别 | Op 名称 | 工具名称 | 主要参数 | 特点 | -|----------|-------------------|---------------|----------------------------------------------------------------------------------|-----------------------------------------| -| **文件读取** | `ReadFileOp` | ReadFile | `absolute_path`(必需), `offset`(可选), `limit`(可选) | 支持大文件的分段读取 | -| | `ReadManyFilesOp` | ReadManyFiles | `pattern`(必需), `dir_path`(可选) | 自动过滤 gitignore 文件，按修改时间排序 | -| **文件写入** | `WriteFileOp` | WriteFile | `file_path`(必需), `content`(必需) | 自动创建父目录，支持覆盖已存在文件 | -| **文件编辑** | `EditOp` | Edit | `file_path`(必需), `old_string`(必需), `new_string`(必需), `expected_replacements`(可选) | 精确匹配，支持创建新文件（old_string 为空） | -| | `SmartEditOp` | SmartEdit | `file_path`(必需), `old_string`(必需), `new_string`(必需), `match_type`(可选) | 支持三种匹配策略（exact/flexible/regex），灵活处理缩进差异 | -| **文件搜索** | `GlobOp` | FindFiles | `pattern`(必需), `dir_path`(可选) | 支持 gitignore 过滤，按修改时间排序 | -| | `GrepOp` | Grep | `pattern`(必需), `path`(可选), `glob`(可选), `limit`(可选) | 支持多文件搜索，返回匹配行及上下文 | -| | `RipGrepOp` | RipGrep | `pattern`(必需), `path`(可选), `glob`(可选) | 类似 ripgrep 的高性能搜索 | -| **目录操作** | `LSOp` | ListDirectory | `path`(必需), `ignore`(可选) | 返回文件和子目录名称列表 | -| **系统操作** | `ShellOp` | ExecuteShell | `command`(必需), `is_background`(必需), `description`(可选), `dir_path`(可选) | 支持前台/后台执行，返回输出和退出码 | -| **任务管理** | `TaskOp` | Task | `description`(必需), `prompt`(必需), `subagent_type`(必需) | 支持任务分解和委托 | -| | `WriteTodosOp` | WriteTodos | `todos`(必需) | 跟踪和管理子任务 | -| | `ExitPlanModeOp` | ExitPlanMode | `reason`(必需) | 用于控制流程执行模式 | - -## 使用示例 - -所有 File Tool Op 都继承自 `BaseAsyncToolOp`，可以通过工具调用的方式使用： - -```python -from flowllm.extensions.file_tool import ReadFileOp, WriteFileOp, EditOp - -# 读取文件 -read_op = ReadFileOp() -result = await read_op.async_call(absolute_path="/path/to/file.py", offset=0, limit=100) - -# 写入文件 -write_op = WriteFileOp() -result = await write_op.async_call(file_path="/path/to/new_file.py", content="print('Hello')") - -# 编辑文件 -edit_op = EditOp() -result = await edit_op.async_call( - file_path="/path/to/file.py", - old_string="old code", - new_string="new code" -) -``` - -## 注意事项 - -1. 所有 Op 默认 `raise_exception=False`，失败时返回错误信息而非抛出异常 -2. 文件路径支持 `~` 扩展，会自动解析为绝对路径 -3. 搜索类 Op 会自动过滤 `.gitignore` 中的文件 -4. 编辑类 Op 会验证替换次数，确保操作准确性 diff --git a/docs/zh/guide/flow_guide.md b/docs/zh/guide/flow_guide.md deleted file mode 100644 index 40c5fc4..0000000 --- a/docs/zh/guide/flow_guide.md +++ /dev/null @@ -1,333 +0,0 @@ -## Flow 配置指南 - -本文详细介绍如何在 `flowllm/flowllm/config/default.yaml` 中配置 Flow（流程）。Flow 是 FlowLLM 的核心概念，通过组合多个 Op（操作）来构建复杂的处理流程。 - ---- - -### 一、Flow 配置基础 - -在配置文件中，`flow` 段用于定义一组可被调用的流程。每个 Flow 的基本结构如下： - -```yaml -flow: - flow_name: - flow_content: Op1Op() >> Op2Op() - description: "流程描述" # 可选 - stream: false # 可选，是否流式输出 - input_schema: # 可选，输入字段定义 - field_name: - type: string - description: "字段说明" - required: true -``` - -**字段说明：** -- `flow_content`: **必需**。Flow 表达式，使用运算符组合多个 Op。 -- `description`: **可选**。流程描述，用于文档生成和调试。 -- `stream`: **可选**。若为 `true`，表示该流程以流式方式输出（如 SSE/Chunk）。 -- `input_schema`: 在不同运行模式下有不同要求： - - 在 MCP 模式下：**必需** 且需精确定义输入字段（类型、是否必填、描述等），用于参数校验与客户端能力展示。 - - 在 HTTP/Stream 模式下：**可选**；建议提供以便进行入参校验与自动文档生成，但不是强制项。 - ---- - -### 二、Flow 表达式语法 - -Flow 表达式使用 Python 语法，通过运算符组合 Op 实例。表达式会被 `parse_flow_expression` 函数解析成可执行的 Op 树。 - -#### 2.1 基本运算符 - -**顺序执行（`>>`）** -- 使用 `>>` 连接多个 Op，表示按顺序执行。 -- 前一个 Op 的输出会作为后一个 Op 的输入。 - -```yaml -flow_content: Op1Op() >> Op2Op() >> Op3Op() -``` - -**并行执行（`|`）** -- 使用 `|` 连接多个 Op，表示并行执行。 -- 所有 Op 同时运行，结果会合并。 - -```yaml -flow_content: Op1Op() | Op2Op() -``` - -**混合组合** -- 可以使用括号进行分组，构建复杂的执行流程。 - -```yaml -flow_content: Op1Op() >> (Op2Op() | Op3Op()) >> Op1Op() -``` - -#### 2.2 单行表达式 - -最简单的 Flow 表达式是单行，直接使用运算符连接 Op： - -```yaml -flow: - simple_flow: - flow_content: GenSystemPromptOp() >> ChatOp() - description: "简单的对话流程" -``` - -#### 2.3 多行表达式 - -Flow 表达式支持多行，前面的行用于设置上下文（如变量赋值），**最后一行必须是返回 BaseOp 的表达式**。 - -**规则：** -- 前面的行可以包含赋值语句，用于准备上下文。 -- 最后一行必须是表达式（不能是赋值），且必须返回一个 `BaseOp` 实例。 - -**示例 1：变量赋值后返回** - -```yaml -flow: - multiline_flow: - flow_content: | - op = ContainerOp() - op.ops.search = SearchOp() - op.ops.find = FindOp() - op -``` - -**示例 2：变量重新赋值** - -```yaml -flow: - reassign_flow: - flow_content: | - opx = Op1Op() >> Op2Op() - opx = opx >> Op3Op() - opx -``` - -**示例 3：复杂组合** - -```yaml -flow: - complex_flow: - flow_content: | - op1 = Op1Op() - op1.ops.search = SearchOp() - op1.ops.find = FindOp() - (op1 | Op2Op()) >> Op3Op() -``` - -#### 2.4 属性赋值 - -可以为 Op 实例的属性赋值，特别是容器类 Op 的 `ops` 属性： - -```yaml -flow: - attribute_flow: - flow_content: | - op = ContainerOp() - op.ops.search = SearchOp() - op.ops.find = FindOp() - op -``` - -#### 2.5 左移运算符（`<<`） - -使用 `<<` 运算符可以批量设置容器 Op 的子操作： - -```yaml -flow: - left_shift_flow: - flow_content: | - op = ContainerOp() - op << {"search": Op1Op(), "find": Op2Op()} - (op | Op2Op()) >> (Op1Op() | Op3Op()) >> op -``` - -这等价于： -```python -op.ops.search = Op1Op() -op.ops.find = Op2Op() -``` - ---- - -### 三、完整配置示例 - -以下是一个包含多种 Flow 配置的完整示例： - -```yaml -flow: - # 简单的顺序流程 - demo_http_flow: - flow_content: GenSystemPromptOp() >> ChatOp() - description: "AI 对话助手" - input_schema: - query: - type: string - description: "用户查询" - required: true - - # 流式输出流程 - demo_stream_http_flow: - flow_content: GenSystemPromptOp() >> StreamChatOp() - stream: true - description: "流式 AI 对话助手" - input_schema: - query: - type: string - description: "用户查询" - required: true - - # 简单的单 Op 流程 - demo_mcp_flow: - flow_content: MockSearchOp() - description: "搜索查询结果" - input_schema: - query: - type: string - description: "用户查询" - required: true - - # 多行复杂流程 - complex_multiline_flow: - flow_content: | - container = ContainerOp() - container.ops.search = SearchOp() - container.ops.find = FindOp() - (container | ProcessOp()) >> FormatOp() - description: "复杂的多步骤流程" - - # 使用左移运算符的流程 - left_shift_flow: - flow_content: | - op = ContainerOp() - op << {"search": SearchOp(), "find": FindOp()} - op >> ProcessOp() - description: "使用左移运算符设置子操作" -``` - ---- - -### 四、表达式解析规则 - -Flow 表达式由 `parse_flow_expression` 函数解析，遵循以下规则： - -1. **Op 注册要求** - - 表达式中使用的 Op 类必须已在 `C.registry_dict["op"]` 中注册。 - - 可以使用注册时的名称，或直接使用类名。 - -2. **多行处理** - - 所有非空行会被处理。 - - 前面的行通过 `exec()` 执行，用于设置上下文。 - - 最后一行通过 `eval()` 求值，必须返回 `BaseOp` 实例。 - -3. **最后一行限制** - - **必须是表达式**，不能是赋值语句。 - - **必须返回 BaseOp**，否则会抛出 `AssertionError`。 - -4. **执行环境** - - 表达式在受限环境中执行，只包含已注册的 Op 类。 - - 不支持标准库函数，仅支持 Op 类的实例化和运算符操作。 - ---- - -### 五、常见模式与最佳实践 - -#### 5.1 简单顺序流程 -适用于需要按步骤执行的场景： - -```yaml -flow_content: Step1Op() >> Step2Op() >> Step3Op() -``` - -#### 5.2 并行处理 -适用于可以同时执行多个独立任务的场景： - -```yaml -flow_content: FetchDataOp() | ProcessDataOp() | ValidateDataOp() -``` - -#### 5.3 条件分支（通过容器 Op） -使用容器 Op 实现条件逻辑： - -```yaml -flow_content: | - container = ContainerOp() - container.ops.branch1 = Branch1Op() - container.ops.branch2 = Branch2Op() - container -``` - -#### 5.4 组合模式 -顺序和并行混合使用： - -```yaml -flow_content: | - PrepareOp() >> (ParallelOp1() | ParallelOp2()) >> MergeOp() -``` - -#### 5.5 流式输出 -对于需要实时返回结果的场景，设置 `stream: true`： - -```yaml -stream_flow: - flow_content: GenSystemPromptOp() >> StreamChatOp() - stream: true -``` - ---- - -### 六、调试技巧 - -1. **检查 Op 注册** - - 确保所有使用的 Op 都已正确注册。 - - 检查 Op 类名是否正确（区分大小写）。 - -2. **验证表达式语法** - - 最后一行必须是表达式，不能是赋值。 - - 确保括号匹配。 - -3. **测试简单表达式** - - 先测试单 Op：`SingleOp()` - - 再测试简单组合：`Op1() >> Op2()` - - 逐步增加复杂度。 - -4. **查看错误信息** - - `ValueError: flow content is empty` - 表达式为空 - - `AssertionError: Expression '...' did not evaluate to a BaseOp instance` - 最后一行未返回 BaseOp - - `NameError` - Op 类未注册或名称错误 - ---- - -### 七、与 input_schema 的配合 - -`input_schema` 定义了 Flow 的输入接口，与 `flow_content` 配合使用。注意不同运行模式的要求： - -- MCP 模式：`input_schema` 为**必填**，且需与真实入参严格一致，MCP 客户端会基于该定义进行参数校验与提示。 -- HTTP/Stream 模式：`input_schema` 为**可选**，但建议填写以获得更好的入参校验与自动生成文档的体验。 - -```yaml -flow: - search_flow: - flow_content: SearchOp() >> FormatOp() - description: "搜索并格式化结果" - input_schema: - query: - type: string - description: "搜索关键词" - required: true - limit: - type: integer - description: "结果数量限制" - required: false - default: 10 -``` - -这样配置后，调用该 Flow 时需要提供符合 `input_schema` 的输入参数。 - ---- - -### 八、参考资源 - -- **测试用例**：`tests/test_expression_parser.py` 包含大量 Flow 表达式示例 -- **解析函数**：`flowllm/core/utils/common_utils.py` 中的 `parse_flow_expression` 函数 -- **Op 介绍**：参考 `op_introduction.md` 了解可用的 Op 类型 - diff --git a/docs/zh/guide/http_service_guide.md b/docs/zh/guide/http_service_guide.md deleted file mode 100644 index 1a1c299..0000000 --- a/docs/zh/guide/http_service_guide.md +++ /dev/null @@ -1,151 +0,0 @@ -## HTTP 服务快速上手 - -本指南演示如何在 FlowLLM 中以最小工作量搭建一个可调用的 HTTP 服务，覆盖：编写 Op、编排 Flow、配置、启动与客户端调用。 - ---- - -### 一、编写 Op - -HTTP 工具的业务逻辑来源于 Flow 中的 Op。基类见 `flowllm/core/op/base_op.py`，异步基类为 `BaseAsyncOp`： - -```python -from flowllm.core.context import C -from flowllm.core.op import BaseAsyncOp - -@C.register_op() -class EchoOp(BaseAsyncOp): - file_path: str = __file__ - - async def async_execute(self): - text = self.context.get("text", "") - self.context.response.answer = f"echo: {text}" -``` - -要点： -- 类名需以 `Op` 结尾（框架有断言） -- 用 `@C.register_op()` 注册后方可在 Flow 中引用 -- 对话型 Op 通常继承 `BaseAsyncOp` 并实现 `async_execute` -- 可通过 `self.context` 读写上下文，`self.llm` 调用模型，`self.prompt_format()` 绑定同名 `*_prompt.yaml` -- 单元：以最小上下文执行 `EchoOp.async_execute`，断言 `context.response.answer` -- 端到端：启动 HTTP 服务后用客户端调用并断言（见“客户端调用与测试”） - ---- - -### 二、编写yaml config - -保存如下配置为： -- 覆盖默认：`flowllm/flowllm/config/default.yaml` -- 自定义：项目根新建 `my_http_config.yaml` - -```yaml -backend: http -thread_pool_max_workers: 128 - -http: - host: "0.0.0.0" - port: 8002 - -flow: - demo_http_flow: - flow_content: GenSystemPromptOp() >> ChatOp() - description: "AI assistant" - input_schema: - query: - type: string - description: "user query" - required: true - -llm: - default: - backend: openai_compatible - model_name: qwen3-30b-a3b-instruct-2507 - params: - temperature: 0.6 - -embedding_model: - default: - backend: openai_compatible - model_name: text-embedding-v4 - params: - dimensions: 1024 - -vector_store: - default: - backend: memory - embedding_model: default -``` - -要点： -- Flow 的编排直接在配置文件的 `flow` 段中定义（如上 `demo_http_flow`）。若要将自定义的 `EchoOp` 暴露为接口，可将其填入对应 `flow_content`；通用对话可参考 `GenSystemPromptOp() >> ChatOp()`。 -- `backend: http` 指定以 HTTP 服务启动 -- `flow` 定义同步接口 -- `input_schema` 在 HTTP 模式下可选，建议填写以生成更完善的 OpenAPI 与入参校验 -- `llm/embedding_model/vector_store` 为默认能力，可按需替换 -- 更多 Flow 表达式与字段说明，参考： - - `docs/zh/guide/flow_guide.md` - - `docs/zh/guide/config_guide.md` - ---- - -### 三、启动服务 - -确保已安装 FlowLLM，并在`.env`中设置模型相关环境变量，可直接参考项目根的示例文件 `example.env`： -```bash -export FLOW_LLM_API_KEY="sk-xxxx" -export FLOW_LLM_BASE_URL="https://xxxx/v1" -export FLOW_EMBEDDING_API_KEY="sk-xxxx" -export FLOW_EMBEDDING_BASE_URL="https://xxxx/v1" -``` - -- 使用my_http_config启动： -```bash -flowllm config=my_http_config backend=http -``` - -成功后： -- 健康检查：`GET http://0.0.0.0:8002/health` -- OpenAPI：`GET http://0.0.0.0:8002/docs` -- 同步接口：`POST http://0.0.0.0:8002/demo_http_flow` - -服务行为（见 `flowllm/core/service/http_service.py`）： -- 为每个 Flow 生成 `POST /{flow_name}` 接口 -- 自动注入 CORS（默认允许所有来源） - -若使用 `demo_http_flow`：`POST http://0.0.0.0:8002/demo_http_flow` - ---- - -### 四、客户端调用与测试 - -使用内置 `HttpClient` 进行调用： - -```python -import asyncio -import json -from flowllm.core.utils import HttpClient - - -async def main(): - async with HttpClient("http://0.0.0.0:8002") as client: - # 健康检查 - health = await client.health_check() - print("health:", json.dumps(health, indent=2, ensure_ascii=False)) - - # 查看可用 endpoints - schema = await client.list_endpoints() - print("endpoints:", list(schema.get("paths", {}).keys())) - - # 同步调用 - resp = await client.execute_flow("demo_http_flow", query="阿里巴巴前景如何？") - print("sync answer:", getattr(resp, "answer", None)) - - -if __name__ == "__main__": - asyncio.run(main()) -``` - -也可直接运行内置测试脚本（以 `demo_http_flow` 为例）： - -```bash -python -m flowllm.tests.http_client_test -``` diff --git a/docs/zh/guide/http_stream_guide.md b/docs/zh/guide/http_stream_guide.md deleted file mode 100644 index 450e0e7..0000000 --- a/docs/zh/guide/http_stream_guide.md +++ /dev/null @@ -1,221 +0,0 @@ -## HTTP 流式服务快速上手 - -本指南演示如何在 FlowLLM 中搭建一个支持流式响应的 HTTP 服务，覆盖：编写流式 Op、配置流式 Flow、启动服务与客户端流式调用。 - ---- - -### 一、编写流式 Op - -流式 Op 与普通 Op 的主要区别在于使用 `self.context.add_stream_chunk()` 来发送流式数据块，而不是一次性返回完整结果。基类见 `flowllm/core/op/base_async_op.py`： - -```python -from flowllm.core.context import C -from flowllm.core.op import BaseAsyncOp -from flowllm.core.enumeration import ChunkEnum, Role -from flowllm.core.schema import FlowStreamChunk, Message - -@C.register_op() -class StreamChatOp(BaseAsyncOp): - file_path: str = __file__ - - async def async_execute(self): - messages = self.context.messages - system_prompt = self.context.system_prompt - - # 构建包含系统提示的消息列表 - messages = [Message(role=Role.SYSTEM, content=system_prompt)] + messages - - # 从 LLM 获取流式响应 - async for stream_chunk in self.llm.astream_chat(messages): - assert isinstance(stream_chunk, FlowStreamChunk) - # 只处理特定类型的块（答案、思考、错误、工具调用） - if stream_chunk.chunk_type in [ - ChunkEnum.ANSWER, - ChunkEnum.THINK, - ChunkEnum.ERROR, - ChunkEnum.TOOL - ]: - # 将流式块添加到上下文，会自动发送到客户端 - await self.context.add_stream_chunk(stream_chunk) -``` - -要点： -- 类名需以 `Op` 结尾（框架有断言） -- 用 `@C.register_op()` 注册后方可在 Flow 中引用 -- 流式 Op 继承 `BaseAsyncOp` 并实现 `async_execute` -- 使用 `self.context.add_stream_chunk(stream_chunk)` 发送流式数据块 -- 流式块类型包括：`ANSWER`（答案）、`THINK`（思考过程）、`ERROR`（错误）、`TOOL`（工具调用） -- 框架会自动处理流式传输，无需手动管理连接 - ---- - -### 二、编写 yaml config - -保存如下配置为： -- 覆盖默认：`flowllm/flowllm/config/default.yaml` -- 自定义：项目根新建 `my_stream_config.yaml` - -```yaml -backend: http -thread_pool_max_workers: 128 - -http: - host: "0.0.0.0" - port: 8002 - -flow: - # 流式 HTTP 流程（POST /demo_stream_http_flow） - demo_stream_http_flow: - flow_content: GenSystemPromptOp() >> StreamChatOp() - stream: true # 关键：标识这是一个流式 Flow - description: "AI 对话助手（流式返回）" - input_schema: - query: - type: string - description: "用户问题" - required: true - -llm: - default: - backend: openai_compatible - model_name: qwen3-30b-a3b-instruct-2507 - params: - temperature: 0.6 - -embedding_model: - default: - backend: openai_compatible - model_name: text-embedding-v4 - params: - dimensions: 1024 - -vector_store: - default: - backend: memory - embedding_model: default -``` - -要点： -- **关键配置**：在 Flow 配置中添加 `stream: true` 标识，框架会自动将其注册为流式端点 -- Flow 的编排在 `flow_content` 中定义，流式 Op（如 `StreamChatOp`）可以与其他 Op 组合 -- `backend: http` 指定以 HTTP 服务启动 -- `input_schema` 在 HTTP 模式下可选，建议填写以生成更完善的 OpenAPI 与入参校验 -- `llm/embedding_model/vector_store` 为默认能力，可按需替换 -- 更多 Flow 表达式与字段说明，参考： - - `docs/zh/guide/flow_guide.md` - - `docs/zh/guide/config_guide.md` - ---- - -### 三、启动服务 - -确保已安装 FlowLLM，并在`.env`中设置模型相关环境变量，可直接参考项目根的示例文件 `example.env`： -```bash -export FLOW_LLM_API_KEY="sk-xxxx" -export FLOW_LLM_BASE_URL="https://xxxx/v1" -export FLOW_EMBEDDING_API_KEY="sk-xxxx" -export FLOW_EMBEDDING_BASE_URL="https://xxxx/v1" -``` - -- 使用my_stream_config启动： -```bash -flowllm config=my_stream_config backend=http -``` - -成功后： -- 健康检查：`GET http://0.0.0.0:8002/health` -- OpenAPI：`GET http://0.0.0.0:8002/docs` -- 流式接口：`POST http://0.0.0.0:8002/demo_stream_http_flow` - -服务行为（见 `flowllm/core/service/http_service.py`）： -- 为每个标记了 `stream: true` 的 Flow 生成流式 `POST /{flow_name}` 接口 -- 使用 Server-Sent Events (SSE) 格式返回流式数据 -- 自动注入 CORS（默认允许所有来源） -- 流式数据格式：每行以 `data:` 开头，包含 JSON 格式的 `FlowStreamChunk`，结束时发送 `data:[DONE]` - -若使用 `demo_stream_http_flow`：`POST http://0.0.0.0:8002/demo_stream_http_flow` - ---- - -### 四、客户端调用与测试 - -使用内置 `HttpClient` 进行流式调用： - -```python -import asyncio -import json -from flowllm.core.utils import HttpClient - - -async def main(): - async with HttpClient("http://0.0.0.0:8002") as client: - # 健康检查 - health = await client.health_check() - print("health:", json.dumps(health, indent=2, ensure_ascii=False)) - - # 查看可用 endpoints - schema = await client.list_endpoints() - print("endpoints:", list(schema.get("paths", {}).keys())) - - # 流式调用 - print("=" * 50) - print("流式响应：") - async for chunk in client.execute_stream_flow( - "demo_stream_http_flow", - query="什么是人工智能？" - ): - chunk_type = chunk.get("type", "answer") - chunk_content = chunk.get("content", "") - if chunk_content: - # 根据块类型进行不同处理 - if chunk_type == "answer": - print(chunk_content, end="", flush=True) # 实时打印答案 - elif chunk_type == "think": - # 思考过程可以单独处理 - print(f"\n[思考] {chunk_content}", end="", flush=True) - elif chunk_type == "tool": - print(f"\n[工具调用] {chunk_content}") - elif chunk_type == "error": - print(f"\n[错误] {chunk_content}") - print("\n" + "=" * 50) - - -if __name__ == "__main__": - asyncio.run(main()) -``` - -也可直接运行内置测试脚本（以 `demo_stream_http_flow` 为例）： - -```bash -python -m flowllm.tests.http_client_test -``` - -**流式响应格式说明**： -- 每个 chunk 是一个字典，包含： - - `type`: 块类型（`answer`、`think`、`tool`、`error`） - - `content`: 块内容（字符串或 JSON） -- 流式响应使用 SSE 格式，客户端通过 `async for` 循环接收 -- 当收到 `[DONE]` 信号时，流式传输结束 - -**使用 curl 测试流式接口**： - -```bash -curl -X POST http://localhost:8002/demo_stream_http_flow \ - -H "Content-Type: application/json" \ - -d '{ - "query": "what is ai" - }' -``` - ---- - -### 五、流式 vs 同步接口对比 - -| 特性 | 同步接口 | 流式接口 | -|------|---------|---------| -| 配置 | 无需特殊配置 | 需设置 `stream: true` | -| Op 实现 | 使用 `self.context.response.answer` | 使用 `self.context.add_stream_chunk()` | -| 返回方式 | 一次性返回完整结果 | 实时返回数据块 | -| 适用场景 | 快速响应、简单查询 | 长文本生成、实时反馈、思考过程展示 | -| 客户端调用 | `execute_flow()` | `execute_stream_flow()` | -| 响应格式 | JSON | SSE (Server-Sent Events) | diff --git a/docs/zh/guide/installation.md b/docs/zh/guide/installation.md deleted file mode 100644 index fb00168..0000000 --- a/docs/zh/guide/installation.md +++ /dev/null @@ -1,94 +0,0 @@ -# 安装指南 - -本文介绍如何安装和配置 FlowLLM。 - ---- - -## 一、从 PyPI 安装（推荐） - -使用 pip 从 PyPI 安装 FlowLLM： - -```bash -pip install flowllm -``` - ---- - -## 二、从源码安装 - -如果你想从源码安装或参与开发，可以按以下步骤操作： - -### 1. 克隆仓库 - -```bash -git clone https://github.com/flowllm-ai/flowllm.git -cd flowllm -``` - -### 2. 安装依赖 - -```bash -# 安装基础版本 -pip install . - -# 或安装开发版本（包含开发依赖） -pip install -e ".[dev]" -``` - -### 3. 验证安装 - -安装完成后，可以通过以下命令验证： - -```bash -# 检查版本 -python -c "import flowllm; print(flowllm.__version__)" - -# 查看帮助信息 -flowllm --help -``` - ---- - -## 三、环境配置 - -FlowLLM 使用环境变量进行配置。你需要创建一个 `.env` 文件来设置必要的参数。 - -### 1. 创建环境配置文件 - -将项目根目录下的 `example.env` 复制为 `.env`： - -```bash -cp example.env .env -``` - -### 2. 配置环境变量 - -编辑 `.env` 文件，填入你的配置信息： - -```bash -# LLM API 配置（必需） -FLOW_LLM_API_KEY=sk-xxxx -FLOW_LLM_BASE_URL=https://xxxx/v1 - -# Embedding API 配置（必需） -FLOW_EMBEDDING_API_KEY=sk-xxxx -FLOW_EMBEDDING_BASE_URL=https://xxxx/v1 -``` - -### 3. 环境变量说明 - -- **FLOW_LLM_API_KEY**: LLM API 的密钥 -- **FLOW_LLM_BASE_URL**: LLM API 的基础 URL（OpenAI 兼容格式） -- **FLOW_EMBEDDING_API_KEY**: Embedding API 的密钥 -- **FLOW_EMBEDDING_BASE_URL**: Embedding API 的基础 URL - -### 4. 使用环境变量 - -FlowLLM 会自动读取项目根目录下的 `.env` 文件。如果 `.env` 文件不在项目根目录，可以手动导出环境变量： - -```bash -export FLOW_LLM_API_KEY=your_api_key -export FLOW_LLM_BASE_URL=your_base_url -export FLOW_EMBEDDING_API_KEY=your_embedding_key -export FLOW_EMBEDDING_BASE_URL=your_embedding_url -``` diff --git a/docs/zh/guide/mcp_service_guide.md b/docs/zh/guide/mcp_service_guide.md deleted file mode 100644 index 19aa517..0000000 --- a/docs/zh/guide/mcp_service_guide.md +++ /dev/null @@ -1,163 +0,0 @@ -## MCP 服务快速上手 - -本指南演示如何在 FlowLLM 中以最小工作量搭建一个可调用的 MCP（Model Context Protocol）服务，覆盖：编写 Op、编排 Flow、配置、启动与客户端调用。 - ---- - -### 一、编写 Op - -MCP 工具的业务逻辑来源于 Flow 中的 Op。可直接复用已有 Op，也可先编写一个简单示例。以下为内置示例 `MockSearchOp`： - -```python -from flowllm.core.context import C -from flowllm.core.op import BaseAsyncOp - -@C.register_op() -class MockSearchOp(BaseAsyncOp): - """Mock search operation that uses LLM to generate realistic search results.""" - file_path: str = __file__ - - async def async_execute(self): - query = self.context.query - if not query: - self.context.response.answer = "No results found." - return - num_results = random.randint(0, 5) - user_prompt = self.prompt_format("mock_search_op_prompt", query=query, num_results=num_results) - messages = [ - Message(role=Role.SYSTEM, content="You are a helpful assistant that generates realistic search results in JSON format."), - Message(role=Role.USER, content=user_prompt), - ] - def callback_fn(message: Message): - return extract_content(message.content, "json") - search_results: str = await self.llm.achat(messages=messages, callback_fn=callback_fn) - self.context.response.answer = json.dumps(search_results, ensure_ascii=False, indent=2) -``` - -要点： -- 类名以 `Op` 结尾，且用 `@C.register_op()` 注册后方可在 Flow 中引用。 -- 异步对话型 Op 继承 `BaseAsyncOp` 并实现 `async_execute`。 -- 通过 `self.context` 读写上下文，`self.llm` 调用模型，`self.prompt_format()` 绑定提示词模板。 - ---- - -### 二、编写yaml config - -保存如下配置为： -- 覆盖默认：`flowllm/flowllm/config/default.yaml` -- 自定义：项目根新建 `my_mcp_config.yaml` - -```yaml -backend: mcp -thread_pool_max_workers: 128 - -mcp: - transport: sse - host: "0.0.0.0" - port: 8001 - -flow: - # 将 Flow 以 MCP 工具形式暴露（工具名：demo_mcp_flow） - demo_mcp_flow: - flow_content: MockSearchOp() - description: "search results for a given query." - input_schema: - query: - type: string - description: "user query" - required: true - -llm: - default: - backend: openai_compatible - model_name: qwen3-30b-a3b-instruct-2507 - params: - temperature: 0.6 - -embedding_model: - default: - backend: openai_compatible - model_name: text-embedding-v4 - params: - dimensions: 1024 - -vector_store: - default: - backend: memory - embedding_model: default -``` - -要点： -- Flow 的编排直接在配置文件的 `flow` 段中定义（如上 `demo_mcp_flow`）。Flow 名即工具名，在 MCP 中以同名工具暴露。 -- `input_schema` 建议填写，便于客户端生成入参校验与帮助信息。 -- `backend: mcp` 指定以 MCP 服务启动。 -- `mcp.transport` 当前支持 `sse`；服务默认暴露 SSE 端点 `GET /sse`。 -- `flow` 中声明的每个 Flow 都将作为一个 MCP 工具对外提供。 - ---- - -### 三、启动服务 - -确保已安装 FlowLLM，并在`.env`中设置模型相关环境变量，可直接参考项目根的示例文件 `example.env`： -```bash -export FLOW_LLM_API_KEY="sk-xxxx" -export FLOW_LLM_BASE_URL="https://xxxx/v1" -export FLOW_EMBEDDING_API_KEY="sk-xxxx" -export FLOW_EMBEDDING_BASE_URL="https://xxxx/v1" -``` - -- 使用my_mcp_config启动： -```bash -flowllm \ - config=my_mcp_config \ - backend=mcp \ - mcp.transport=sse \ - mcp.port=8002 \ - llm.default.model_name=qwen3-30b-a3b-thinking-2507 -``` - -成功后（以 `host=0.0.0.0`, `port=8001` 为例）： -- SSE 连接：`GET http://0.0.0.0:8001/sse` -- 工具清单与调用通过 SSE 渠道完成（由 MCP 客户端实现） - -服务行为（参考 `flowllm/core/service/*` 内 MCP 相关实现）： -- 为每个 Flow 生成同名 MCP 工具（如 `demo_mcp_flow`） -- 采用 SSE 作为传输时，客户端通过单一 SSE 端点进行会话 - ---- - -### 四、客户端调用与测试 - -使用内置 `FastMcpClient` 进行调用（简化示例）： - -```python -import asyncio - -from flowllm.core.utils import FastMcpClient - - -async def main(): - config = { - "type": "sse", - "url": "http://0.0.0.0:8001/sse", - "headers": {}, - "timeout": 30.0, - } - async with FastMcpClient("test_client", config, max_retries=3) as client: - # 1) 列出可用工具 - tool_calls = await client.list_tool_calls() - # 2) 调用工具（与 Flow 同名） - result = await client.call_tool("demo_mcp_flow", {"query": "阿里巴巴前景怎么样？"}) - # 3) 读取结果 - print(result) - - -if __name__ == "__main__": - asyncio.run(main()) -``` - -也可直接运行内置测试脚本： - -```bash -python -m flowllm.tests.mcp_client_test -``` \ No newline at end of file diff --git a/docs/zh/guide/op_introduction.md b/docs/zh/guide/op_introduction.md deleted file mode 100644 index 5f3fbb0..0000000 --- a/docs/zh/guide/op_introduction.md +++ /dev/null @@ -1,247 +0,0 @@ -# Op 体系介绍 - -这篇文章会用最简单直白的方式，带你了解 flowllm 中的 Op（算子）体系。我们会看看它们是怎么继承的，各自有什么特点，以及什么时候该用哪个。 - -## 继承关系总览 - -先来看看整个家族的继承关系，就像家谱一样： - -``` -BaseOp（基类Op） - ├─ BaseAsyncOp（异步Op） - │ ├─ SequentialOp（顺序Op） - │ ├─ ParallelOp（并行Op） - │ └─ BaseAsyncToolOp（工具Op） - │ └─ BaseMcpOp（MCPOp） - └─ BaseRayOp（分布式Op） -``` - -简单来说： - -- BaseOp 是所有算子的基类，提供了重试、缓存、上下文管理等基础能力 -- BaseAsyncOp 在 BaseOp 基础上加了异步能力，包括异步生命周期钩子和任务编排 -- SequentialOp 和 ParallelOp 是用来组合多个算子的，一个按顺序执行，一个并行执行。 - -#### 重要提示 - -这两个 Op 不是用来继承的，而是通过 `>>` 或 `|` 运算符自动组合出来的。 - -- BaseAsyncToolOp 是专门做工具调用的Op，可以按照 Schema 自动处理输入输出 -- BaseMcpOp 是调用外部 MCP 工具的，会自动加载 Schema 并调用服务 -- BaseRayOp 是用来做分布式并行处理的，适合处理大量数据 - -## 各 Op 对比表格 - -下面这个表格从不同维度对比了各个 Op，帮你快速找到合适的： - -| Op 类型 | 执行模式 | 主要能力 | 核心参数 | 适合场景 | 组合能力 | -|-----------------|------------|------------------------------|---------------------------------------------------------------------------------------------|-------------------------|------------------------| -| BaseOp | 同步 | 基础功能：重试、缓存、上下文管理、资源延迟初始化 | `async_mode`, `max_retries`, `enable_cache`, `llm`, `embedding_model` | 简单的同步，需要基础功能支持 | 支持 `<<`, `>>`, `\|` 组合 | -| BaseAsyncOp | 异步 | 异步生命周期钩子、异步缓存、任务提交和汇总 | 继承 BaseOp 所有参数 | 需要异步执行的任务，需要任务编排 | 支持 `<<`, `>>`, `\|` 组合 | -| SequentialOp | 同步/异步（需一致） | 顺序执行多个子算子，共享上下文 | 通过 `>>` 运算符自动组合（不用于继承） | 需要按顺序执行多个步骤的流程 | 不支持 `<<`，用 `>>` 组合 | -| ParallelOp | 同步/异步（需一致） | 并行执行多个子算子，共享上下文 | 通过 `\|` 运算符自动组合（不用于继承） | 需要同时执行多个独立任务 | 不支持 `<<`，用 `\|` 组合 | -| BaseAsyncToolOp | 异步 | 基于 Schema 自动读取输入、写回输出，支持字段映射 | `tool_index`, `save_answer`, `input_schema_mapping`, `output_schema_mapping` | 需要定义明确输入输出 Schema 的工具调用 | 支持 `<<`, `>>`, `\|` 组合 | -| BaseMcpOp | 异步 | 调用外部 MCP 工具，自动加载 Schema | `mcp_name`, `tool_name`, `max_retries`, `timeout`, `input_schema_required/optional/deleted` | 需要调用外部 MCP 服务的场景 | 支持 `<<`, `>>`, `\|` 组合 | -| BaseRayOp | 同步高并发 | 分布式并行处理，自动识别并行维度，共享大型数据结构 | `ray_max_workers` | 需要处理大量数据，需要分布式加速 | 支持 `<<`, `>>`, `\|` 组合 | - -## 详细介绍每个 Op - -### BaseOp - 最基础的老祖宗 - -#### 它是什么 - -所有算子的基类，提供了最基础的能力。 - -#### 它能做什么 - -- 标记是同步模式还是异步模式（通过 `async_mode` 控制） -- 提供重试机制（`max_retries`）和异常处理（`raise_exception`） -- 支持缓存功能（`enable_cache`），可以设置过期时间 -- 提供上下文管理（`context`）和计时功能（`timer`） -- 延迟初始化 LLM、Embedding 模型、向量库等资源 -- 提供 Token 计数能力（`token_count()`），支持多种后端（base、openai、hf 等）准确计算消息和工具的 token 数量 -- 提供组合运算符：`<<`（添加子算子）、`>>`（顺序执行）、`|`（并行执行） - -#### 什么时候用它 - -- 写一个简单的同步或异步任务 -- 需要基础的重试、缓存、上下文管理功能 -- 作为其他更复杂 Op 的基类 - -#### 注意事项 - -- 需要实现 `execute()` 方法（同步）或 `async_execute()` 方法（异步） -- 如果设置了 `max_retries > 1` 或 `raise_exception=False`，失败时会调用 `default_execute()` 作为兜底 - ---- - -### BaseAsyncOp - 异步能力增强版 - -#### 它是什么 - -在 BaseOp 基础上加了完整的异步能力。 - -#### 它能做什么 - -- 继承 BaseOp 的所有能力 -- 默认就是异步模式（`async_mode=True`） -- 提供异步生命周期钩子：`async_before_execute()`、`async_execute()`、`async_after_execute()`、`async_default_execute()` -- 支持异步缓存（`async_save_load_cache()`），可以处理协程或在线程池执行同步函数 -- 提供任务编排：`submit_async_task()` 提交任务，`join_async_task()` 等待结果，支持超时和异常处理 - -#### 什么时候用它 - -- 需要异步执行的任务 -- 需要在执行前后做一些异步操作（比如异步读取数据、异步写入结果） -- 需要提交多个异步任务并等待它们完成 - -#### 注意事项 - -- 必须实现 `async_execute()` 方法 - ---- - -### SequentialOp - 顺序执行组合器 - -#### 它是什么 - -用来按顺序执行多个算子的组合器。 - -#### 重要认知 - -这个 Op 不是用来继承的，而是通过 `>>` 运算符自动组合出来的。 - -#### 它能做什么 - -- 按顺序执行多个子算子，前一个的输出作为后一个的输入 -- 所有子算子共享同一个 `context` -- 当你使用 `op1 >> op2` 时，会自动创建一个 `SequentialOp` 实例 -- 支持同步和异步两种模式，但所有子算子的模式必须一致 - -#### 什么时候用它 - -- 需要按顺序执行多个步骤的流程 -- 比如：先搜索 → 再总结 → 最后生成报告 -- 使用方式：`op1 >> op2 >> op3`，系统会自动创建 `SequentialOp` 来组合它们 - -#### 注意事项 - -- 不要直接继承 SequentialOp，而是通过 `>>` 运算符来组合算子 -- 不能用 `<<` 运算符（会报错） -- 所有子算子的 `async_mode` 必须一致 -- 返回最后一个算子的结果 - ---- - -### ParallelOp - 并行执行组合器 - -#### 它是什么 - -用来并行执行多个算子的组合器。 - -#### 重要认知 - -这个 Op 不是用来继承的，而是通过 `|` 运算符自动组合出来的。 - -#### 它能做什么 - -- 同时执行多个子算子，所有子算子共享同一个 `context` -- 当你使用 `op1 | op2` 时，会自动创建一个 `ParallelOp` 实例 -- 支持同步（多线程）和异步（asyncio）两种模式，但所有子算子的模式必须一致 -- 返回所有子算子结果的列表 - -#### 什么时候用它 - -- 需要同时执行多个独立的任务 -- 比如：同时调用多个搜索接口、同时处理多个文件 -- 使用方式：`op1 | op2 | op3`，系统会自动创建 `ParallelOp` 来组合它们 - -#### 注意事项 - -- 不要直接继承 ParallelOp，而是通过 `|` 运算符来组合算子 -- 不能用 `<<` 运算符（会报错） -- 所有子算子的 `async_mode` 必须一致 -- 返回的是列表，包含所有子算子的结果 - ---- - -### BaseAsyncToolOp - 工具调用专用 - -#### 它是什么 - -专门用来做工具调用的，可以按照 Schema 自动处理输入输出。 - -#### 它能做什么 - -- 通过 `build_tool_call()` 定义工具的 Schema（描述、输入参数、输出参数） -- 自动从 `context` 中读取输入参数（按照输入 Schema） -- 自动把输出结果写回 `context`（按照输出 Schema） -- 支持字段映射（`input_schema_mapping`、`output_schema_mapping`） -- 支持多个工具实例（通过 `tool_index` 区分） -- 可以选择是否把主要输出写入 `response.answer` - -#### 什么时候用它 - -- 需要定义明确输入输出 Schema 的工具 -- 比如：搜索工具（输入：查询词，输出：搜索结果）、翻译工具（输入：文本，输出：翻译结果） - -#### 注意事项 - -- 必须实现 `build_tool_call()` 和 `async_execute()` 方法 -- 如果没有定义输出 Schema，会自动创建一个单字符串输出 -- 输入参数如果标记为 `required=True`，在 `context` 中找不到会报错 - ---- - -### BaseMcpOp - 外部 MCP 工具调用 - -#### 它是什么 - -专门用来调用外部 MCP（Model Context Protocol）工具的。 - -#### 它能做什么 - -- 从服务配置中自动加载工具的 Schema（`C.external_mcp_tool_call_dict[mcp_name][tool_name]`） -- 可以覆盖 Schema 中的参数要求（`input_schema_required`、`input_schema_optional`、`input_schema_deleted`） -- 使用 `FastMcpClient` 调用外部 MCP 服务 -- 支持重试（`max_retries`）和超时（`timeout`） -- 自动处理输入输出（继承自 BaseAsyncToolOp） - -#### 什么时候用它 - -- 需要调用外部 MCP 服务的场景 -- 比如：调用阿里云的 MCP 工具、调用其他第三方 MCP 服务 - -#### 注意事项 - -- 需要在服务配置中配置好 `external_mcp` 和对应的工具 Schema -- 默认 `raise_exception=False`，失败不会抛异常 -- 默认 `max_retries=3`，会自动重试 - ---- - -### BaseRayOp - 分布式并行处理 - -#### 它是什么 - -用来做分布式并行处理的，适合处理大量数据。 - -#### 它能做什么 - -- 自动识别 `kwargs` 中的第一个列表参数作为并行维度（也可以显式指定 `parallel_key`） -- 通过 Ray 把任务分发到多个 worker 上执行 -- 自动把大型数据结构（DataFrame、Series、dict、list、Context）转换成 Ray 对象共享，避免重复传输 -- 按照 `ray_max_workers` 创建 worker，用轮询切片的方式分配任务 -- 最后把所有结果汇总并扁平化 - -#### 什么时候用它 - -- 需要处理大量数据，比如处理几千条记录 -- 需要分布式加速，比如批量调用 LLM、批量处理文件 - -#### 注意事项 - -- 需要在服务配置中设置 `ray_max_workers > 1` -- 第一次调用会自动初始化 Ray（如果还没初始化） -- 会自动识别列表参数，但如果有多个列表，最好显式指定 `parallel_key` - diff --git a/docs/zh/guide/vector_store_guide.md b/docs/zh/guide/vector_store_guide.md deleted file mode 100644 index 290d5ae..0000000 --- a/docs/zh/guide/vector_store_guide.md +++ /dev/null @@ -1,350 +0,0 @@ -### Vector Store 使用指南 - -Vector Store 是用于存储、管理和检索向量嵌入的组件，支持工作空间管理、相似度搜索和元数据过滤等功能。 - -## 核心概念 - -**工作空间（Workspace）**：每个工作空间是一个独立的向量存储单元，用于组织和管理相关的向量节点。 - -**向量节点（VectorNode）**：包含文本内容、向量嵌入和元数据的数据单元，是存储和检索的基本单位。 - -**嵌入模型（Embedding Model）**：用于将文本转换为向量嵌入，支持自动生成节点向量和查询向量。 - -## 可用的实现类型 - -FlowLLM 提供了多种 Vector Store 实现，适用于不同的使用场景： - -- **LocalVectorStore**（[代码路径](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/local_vector_store.py)）：基于本地文件的实现，使用 JSONL 格式持久化存储，适合单机部署和小规模数据 -- **MemoryVectorStore**（[代码路径](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/memory_vector_store.py)）：内存实现，数据存储在内存中，访问速度快，适合临时数据或测试场景 -- **PgVectorStore**（[代码路径](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/pgvector_vector_store.py)）：基于 PostgreSQL 和 pgvector 扩展的实现，支持同步和异步操作，适合已有 PostgreSQL 基础设施的场景 -- **QdrantVectorStore**（[代码路径](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/qdrant_vector_store.py)）：基于 Qdrant 向量数据库，支持高性能向量搜索，适合大规模生产环境 -- **ChromaVectorStore**（[代码路径](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/chroma_vector_store.py)）：基于 ChromaDB 的实现，提供持久化存储和元数据过滤能力 -- **EsVectorStore**（[代码路径](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/es_vector_store.py)）：基于 Elasticsearch 的实现，支持强大的全文搜索和向量搜索组合 - -所有 Vector Store 实现都继承自 **BaseVectorStore**（[代码路径](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/base_vector_store.py)），提供统一的接口规范。 - -## 核心功能 - -### 工作空间管理 - -- **创建工作空间**：创建新的工作空间用于存储向量节点 -- **删除工作空间**：删除工作空间及其所有数据 -- **检查工作空间**：检查指定工作空间是否存在 -- **列出工作空间**：获取所有已存在的工作空间列表 -- **复制工作空间**：将一个工作空间的数据复制到另一个工作空间 - -### 节点操作 - -- **插入节点**：向工作空间插入向量节点，支持单个或批量插入，自动生成向量嵌入 -- **删除节点**：根据节点 ID 删除指定节点 -- **迭代节点**：遍历工作空间中的所有节点 - -### 向量搜索 - -- **相似度搜索**：基于文本查询进行向量相似度搜索，返回最相似的前 K 个结果 -- **元数据过滤**：支持基于元数据的过滤条件，包括精确匹配和范围查询 -- **相似度评分**：搜索结果包含相似度评分，用于评估匹配程度 - -### 数据导入导出 - -- **导出工作空间**：将工作空间数据导出到文件或指定路径 -- **导入工作空间**：从文件或节点列表导入数据到工作空间 -- **回调函数**：支持在导入导出过程中使用回调函数进行数据转换 - -## 同步与异步接口 - -所有 Vector Store 实现都提供同步和异步两套接口： - -- **同步接口**：直接调用方法，适用于同步代码环境 -- **异步接口**：以 `async_` 前缀命名，适用于异步代码环境，提供更好的并发性能 - -异步接口在以下场景中特别有用： -- 使用异步嵌入模型生成向量 -- 在高并发环境中执行批量操作 -- 与其他异步组件集成 - -## 配置选项 - -### 通用配置 - -- **embedding_model**：嵌入模型实例，用于生成向量嵌入 -- **batch_size**：批量操作的大小，默认 1024 - -### LocalVectorStore 配置 - -- **store_dir**：存储目录路径，默认 `./local_vector_store` - -### MemoryVectorStore 配置 - -- **store_dir**：持久化存储目录，默认 `./memory_vector_store` - -### PgVectorStore 配置 - -- **connection_string**：PostgreSQL 连接字符串（同步操作），默认从环境变量 `FLOW_PGVECTOR_CONNECTION_STRING` 读取，或使用 `postgresql://localhost/postgres` -- **async_connection_string**：PostgreSQL 连接字符串（异步操作，可选），默认从环境变量 `FLOW_PGVECTOR_ASYNC_CONNECTION_STRING` 读取，如果未提供则基于 `connection_string` 自动转换 -- **batch_size**：批量操作大小，默认 `1024` - -### QdrantVectorStore 配置 - -- **url**：Qdrant 服务地址（可选，用于 Qdrant Cloud 或自定义部署） -- **host**：Qdrant 服务器主机，默认 `localhost` -- **port**：Qdrant 服务器端口，默认 `6333` -- **api_key**：API 密钥（用于 Qdrant Cloud 认证） -- **distance**：距离度量方式，支持 COSINE、EUCLIDEAN、DOT，默认 COSINE - -### ChromaVectorStore 配置 - -- **store_dir**：ChromaDB 数据存储目录，默认 `./chroma_vector_store` - -### EsVectorStore 配置 - -- **hosts**：Elasticsearch 主机地址，可以是字符串或列表，默认 `http://localhost:9200` -- **basic_auth**：基本认证凭据（用户名和密码） - -## 配置文件示例 - -在 `flowllm/config/default.yaml` 中配置 Vector Store，所有配置都在 `vector_store` 部分下。基本配置结构如下： - -```yaml -vector_store: - default: - backend: # 必需：向量库后端类型 - embedding_model: default # 必需：嵌入模型配置名称 - params: # 可选：后端特定参数 - # 后端特定参数 -``` - -### 配置字段说明 - -- **`backend`**（必需）：向量库后端类型，可选值：`local`、`memory`、`pgvector`、`chroma`、`qdrant`、`elasticsearch` -- **`embedding_model`**（必需）：嵌入模型配置名称，引用 `embedding_model` 部分的配置 -- **`params`**（可选）：后端特定参数字典，将传递给向量库构造函数 - -### 各类型配置示例 - -#### 1. LocalVectorStore 配置 - -最简单的本地文件存储，适合开发和测试。 - -**代码实现**：[`flowllm/core/vector_store/local_vector_store.py`](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/local_vector_store.py) - -```yaml -vector_store: - default: - backend: local - embedding_model: default - params: - store_dir: "./local_vector_store" # 存储目录（可选，默认："./local_vector_store"） -``` - -#### 2. MemoryVectorStore 配置 - -内存存储，访问速度快，适合临时数据或测试。 - -**代码实现**：[`flowllm/core/vector_store/memory_vector_store.py`](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/memory_vector_store.py) - -```yaml -vector_store: - default: - backend: memory - embedding_model: default - params: - store_dir: "./memory_vector_store" # 持久化目录（可选，默认："./memory_vector_store"） -``` - -#### 3. ChromaVectorStore 配置 - -基于 ChromaDB 的持久化存储，支持元数据过滤。 - -**代码实现**：[`flowllm/core/vector_store/chroma_vector_store.py`](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/chroma_vector_store.py) - -```yaml -vector_store: - default: - backend: chroma - embedding_model: default - params: - store_dir: "./chroma_vector_store" # ChromaDB 数据目录（可选，默认："./chroma_vector_store"） -``` - -#### 4. PgVectorStore 配置 - -基于 PostgreSQL 和 pgvector 扩展的向量存储，支持同步和异步操作，适合已有 PostgreSQL 基础设施的场景。 - -**代码实现**：[`flowllm/core/vector_store/pgvector_vector_store.py`](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/pgvector_vector_store.py) - -**前置要求**：需要安装 PostgreSQL 并启用 pgvector 扩展。 - -**基本配置（使用环境变量）：** - -```yaml -vector_store: - default: - backend: pgvector - embedding_model: default - params: - # connection_string 和 async_connection_string 可通过环境变量配置 - # FLOW_PGVECTOR_CONNECTION_STRING=postgresql://user:password@localhost/dbname - # FLOW_PGVECTOR_ASYNC_CONNECTION_STRING=postgresql://user:password@localhost/dbname - batch_size: 1024 # 批量操作大小（可选，默认：1024） -``` - -**显式配置连接字符串：** - -```yaml -vector_store: - default: - backend: pgvector - embedding_model: default - params: - connection_string: "postgresql://user:password@localhost:5432/dbname" # 同步连接字符串 - async_connection_string: "postgresql://user:password@localhost:5432/dbname" # 异步连接字符串（可选） - batch_size: 1024 -``` - -**仅配置同步连接（异步将自动转换）：** - -```yaml -vector_store: - default: - backend: pgvector - embedding_model: default - params: - connection_string: "postgresql://user:password@localhost:5432/dbname" - batch_size: 1024 -``` - -#### 5. QdrantVectorStore 配置 - -**代码实现**：[`flowllm/core/vector_store/qdrant_vector_store.py`](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/qdrant_vector_store.py) - -**本地 Qdrant 实例：** - -```yaml -vector_store: - default: - backend: qdrant - embedding_model: default - params: - host: "localhost" # Qdrant 服务器主机（可选，默认：localhost） - port: 6333 # Qdrant 服务器端口（可选，默认：6333） - distance: "COSINE" # 距离度量方式（可选，默认：COSINE，可选值：COSINE、EUCLIDEAN、DOT） -``` - -**Qdrant Cloud 配置：** - -```yaml -vector_store: - default: - backend: qdrant - embedding_model: default - params: - url: "https://your-cluster.qdrant.io:6333" # Qdrant Cloud 地址 - api_key: "your-api-key-here" # API 密钥 - distance: "COSINE" -``` - -#### 6. EsVectorStore 配置 - -**代码实现**：[`flowllm/core/vector_store/es_vector_store.py`](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/es_vector_store.py) - -**基本配置（本地 Elasticsearch）：** - -```yaml -vector_store: - default: - backend: elasticsearch - embedding_model: default - params: - hosts: "http://localhost:9200" # Elasticsearch 主机地址（可选，默认：http://localhost:9200） -``` - -**带认证的配置：** - -```yaml -vector_store: - default: - backend: elasticsearch - embedding_model: default - params: - hosts: "http://elasticsearch.example.com:9200" - basic_auth: ["username", "password"] # 基本认证凭据 -``` - -**多主机配置：** - -```yaml -vector_store: - default: - backend: elasticsearch - embedding_model: default - params: - hosts: - - "http://es-node1:9200" - - "http://es-node2:9200" - - "http://es-node3:9200" -``` - -### 完整配置示例 - -以下是一个完整的 `default.yaml` 配置示例，包含嵌入模型和向量库配置： - -```yaml -# 嵌入模型配置 -embedding_model: - default: - backend: openai_compatible - model_name: text-embedding-v4 - params: - dimensions: 1024 - -# 向量库配置 -vector_store: - default: - backend: elasticsearch - embedding_model: default - params: - hosts: "http://localhost:9200" -``` - -### 环境变量支持 - -部分 Vector Store 支持通过环境变量进行配置，作为 YAML 配置的补充： - -- **PgVectorStore**： - - `FLOW_PGVECTOR_CONNECTION_STRING` - PostgreSQL 同步连接字符串（默认：`postgresql://localhost/postgres`） - - `FLOW_PGVECTOR_ASYNC_CONNECTION_STRING` - PostgreSQL 异步连接字符串（可选，未提供时基于同步连接字符串自动转换） -- **Elasticsearch**：`FLOW_ES_HOSTS` - Elasticsearch 主机地址 -- **Qdrant**： - - `FLOW_QDRANT_HOST` - Qdrant 主机（默认：`localhost`） - - `FLOW_QDRANT_PORT` - Qdrant 端口（默认：`6333`） - - `FLOW_QDRANT_API_KEY` - Qdrant API 密钥 - -当 YAML 配置中未明确指定参数时，系统会使用环境变量作为后备值。 - -## 元数据过滤 - -支持两种类型的元数据过滤： - -- **精确匹配**：直接指定字段值进行精确匹配 -- **范围查询**：使用 `gte`、`lte`、`gt`、`lt` 操作符进行数值范围查询 -- **嵌套字段**：支持使用点号访问嵌套的元数据字段 - -## 使用场景建议 - -- **开发测试**：使用 MemoryVectorStore 或 LocalVectorStore，无需额外服务 -- **小规模应用**：使用 LocalVectorStore 或 ChromaVectorStore，简单易用 -- **已有 PostgreSQL 基础设施**：使用 PgVectorStore，利用现有数据库资源，支持同步和异步操作 -- **生产环境**：使用 PgVectorStore、QdrantVectorStore 或 EsVectorStore，提供高性能和可扩展性 -- **混合搜索**：使用 EsVectorStore，结合向量搜索和全文搜索能力 - -## 注意事项 - -- 确保嵌入模型的维度与 Vector Store 配置一致 -- 使用 PgVectorStore 前需要确保 PostgreSQL 已安装并启用 pgvector 扩展（`CREATE EXTENSION vector`） -- 大规模数据建议使用专业的向量数据库（PgVectorStore、Qdrant、Elasticsearch） -- 异步接口在异步环境中能提供更好的性能 -- 定期备份重要数据，特别是使用内存存储时 -- 根据数据规模选择合适的批量大小以优化性能 -- PgVectorStore 会自动为每个工作空间创建向量索引（IVFFlat）和元数据索引（GIN），以优化查询性能 - diff --git a/docs/zh/quick_start.md b/docs/zh/quick_start.md new file mode 100644 index 0000000..489d60e --- /dev/null +++ b/docs/zh/quick_start.md @@ -0,0 +1,120 @@ +# 快速开始 + +FlowLLM 是一个配置驱动的 LLM 应用框架：配置定义 Service、Component 和 Job；Job 顺序执行 Step；Service 把 Job 暴露给 CLI、HTTP +或 MCP。 + +## 安装 + +FlowLLM 要求 Python 3.11+。 + +```bash +git clone https://github.com/flowllm-ai/flowllm.git +cd flowllm +pip install -e ".[dev]" +``` + +需要 Claude Code wrapper 时再安装： + +```bash +pip install -e ".[claude-code]" +``` + +## 启动 + +```bash +flowllm start +``` + +默认服务地址是 `127.0.0.1:2333`，默认 workspace 是 `.flowllm/`。可以用 dot notation 覆盖配置： + +```bash +flowllm start service.port=8181 enable_logo=false +flowllm start workspace_dir=/tmp/flowllm-demo service.host=127.0.0.1 service.port=8181 +``` + +启动时会自动创建： + +```text +.flowllm/ +├── metadata/ +└── session/ +``` + +## 调用 Job + +服务启动后，CLI 的非 `start` 命令会调用服务端 Job： + +```bash +flowllm version +flowllm health_check +flowllm help +flowllm demo query="Hello FlowLLM" min_score=0.8 +flowllm add a=1 b=2 +``` + +HTTP 入口是 `POST /`： + +```bash +curl -s http://127.0.0.1:2333/add \ + -H 'Content-Type: application/json' \ + -d '{"a":1,"b":2}' +``` + +流式 Job 返回 SSE： + +```bash +flowllm stream_demo query="Hi" repeat=3 interval=0.05 + +curl -N http://127.0.0.1:2333/stream_demo \ + -H 'Content-Type: application/json' \ + -d '{"query":"Hi","repeat":3,"interval":0.05}' +``` + +## 配置 + +默认配置在 `flowllm/config/default.yaml`。可通过 `.env` 覆盖 LLM、embedding 等变量： + +```bash +cat > .env <<'EOF' +LLM_BACKEND=openai +LLM_MODEL_NAME=qwen3.7-plus +LLM_API_KEY=your_api_key +LLM_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 + +EMBEDDING_BACKEND=openai +EMBEDDING_MODEL_NAME=text-embedding-v4 +EMBEDDING_API_KEY=your_api_key +EMBEDDING_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 +EOF +``` + +也可以指定 YAML/JSON： + +```bash +flowllm start config=/path/to/app.yaml +``` + +最小 Job 配置： + +```yaml +service: + backend: http + +jobs: + add: + backend: base + description: "add two numbers" + steps: + - backend: add_step +``` + +常用 Job backend： + +| backend | 说明 | +|--------------|--------------------| +| `base` | 返回 JSON `Response` | +| `stream` | 返回 SSE 流 | +| `background` | 后台循环任务 | +| `cron` | 定时任务 | + +更多实现细节见 [代码框架](./framework.md)。 diff --git a/docs/zh/reading/20251110-manus-context-raw.md b/docs/zh/reading/20251110-manus-context-raw.md deleted file mode 100644 index b86a6ee..0000000 --- a/docs/zh/reading/20251110-manus-context-raw.md +++ /dev/null @@ -1,159 +0,0 @@ -# langchain & manus Video资料 -from https://www.youtube.com/watch?v=6_BcCthVvb8 - - -好的。嗯，感谢大家的到来。我们现在就开始网络研讨会吧，我相信还会陆续有人加入。呃，我是 Lance，LangChain 的创始工程师之一。和我一起的是来自 Manis 的 Pete。Pete，你要不要简单介绍一下自己？ - ->> 好的。大家好，我是 Manis 的联合创始人兼首席科学家。基本上，Manis 的智能体框架和很多东西都是我设计的，今天能来到这里我非常兴奋。谢谢 Lance 的邀请。 - ->> 是的，我们也很高兴能做这次分享，因为 Manis 首先是一个非常酷的产品。我自己已经用了很久了。此外，他们几个月前还发表了一篇关于“上下文工程”（context engineering）的精彩博文，对我影响很大。所以我想先快速概述一下我对上下文工程的理解，并会引用他们的文章，然后 Pete 将做一个演讲，介绍一些文章中没有涵盖的新想法。所以，即使你已经读过那篇文章，今天的内容也会包含一些新东西，希望对你们来说相当有趣。我会先做个铺垫，然后交给 Pete，最后我们会进行问答环节。 - -你可能听说过“上下文工程”这个术语，它大约在今年早些时候出现。如果你查看谷歌搜索趋势的时间线，“提示工程”（prompt engineering）是在 ChatGPT 之后兴起的，也就是2022年12月。当我们有了这种新的聊天模型后，人们开始对如何给这些模型下指令产生了极大的兴趣。“提示工程”因此成为一门与聊天模型协作并为其编写提示的学科。 - -而“上下文工程”则在今年5月左右出现，我们在谷歌趋势上看到它真正开始上升，这与“智能体之年”的概念大致吻合。为什么会这样呢？一个被很多人观察到的现象是，如果你在构建智能体，就会发现上下文会增长，而且是以一种非常特定的方式增长。我的意思是，我们有一个大语言模型（LLM），它绑定了一些工具，这个 LLM 可以在一个循环中自主调用这些工具。挑战在于，每次调用工具后，你都会得到一个工具观察结果，这个结果会被追加到聊天消息列表中。随着时间的推移，这些消息会不断增长，从而可能导致智能体运行时消息数量的无限制爆炸式增长。 - -举个例子，Manis 在他们的文章中提到，典型的任务大约需要50次工具调用。Anthropic 也提到，生产环境中的智能体可以进行长达数百轮的对话。因此，挑战在于，智能体由于运行时间越来越长且具有自主性，它们会自由地使用工具，从而通过这种工具调用的累积积累大量的上下文。 - -Chrome（应为 Cognition）发布了一份非常不错的报告，谈到了“上下文腐烂”（context rot）的问题，其核心观察就是：随着上下文的增长，性能会下降。 - -这就形成了一个悖论，一个具有挑战性的情况：智能体因为工具调用而需要大量上下文，但我们知道，随着上下文的增长，性能反而会下降。 - -这是我们很多人都面临过的挑战，也正是这个挑战催生或孕育了“上下文工程”这个术语。Arvind Narayanan（应为 Arvind Athalye）今年早些时候在 Twitter 上创造了这个词。你可以把上下文工程理解为一门精妙的艺术与科学，其目标是在上下文窗口中填充下一步决策所需的恰到好处的信息。它试图对抗在构建智能体并让它们自由调用工具时发生的上下文爆炸问题。所有这些工具消息都会累积在你的消息队列中。我们如何才能确保在任何时候，都能向智能体呈现做出正确下一步决策所需的信息呢？ - -为了解决这个问题，我想强调几个我们在许多不同工作中（包括 Manis）看到的共同主题。 - -**第一个想法是上下文卸载（Context Offloading）。** -我们一次又一次地看到这个趋势。其核心思想是，并非所有上下文都需要存在于智能体的消息历史记录中。你可以将信息卸载出去，发送到其他地方，使其位于上下文窗口之外，但在需要时可以检索回来（稍后我们会谈到）。 - -这里最流行的想法之一就是使用文件系统。例如，将工具消息的输出转储到文件系统中，然后只向你的智能体发送回一些必要的最小信息，以便它能在需要时引用完整的上下文。但完整的有效载荷（例如，一个非常占用 token 的网页搜索结果）不会永久地充斥在你的上下文窗口中。 - -我们在许多不同的项目中都看到了这一点。Manis 使用这种方法。我们有一个名为“深度智能体”（deep agents）的项目，它利用了文件系统。Open Deep Research 项目实际上也使用了类似“智能体状态”（agent state）的东西来扮演外部文件系统的角色。Cloud Code 当然也非常广泛地使用了这种方法。长期运行的智能体也大量使用它。 - -因此，将上下文卸载到文件系统的想法，在我们今天看到的许多生产级智能体示例中都非常常见和流行。 - -**第二个想法是减少上下文（Reducing Context）。** -卸载非常简单，就是拿走像工具消息这样占用大量 token 的信息，不把它全部送回你的消息列表，而是将其转储到一个文件系统中，只在需要时检索。这就是卸载。 -而减少上下文则类似，但你是对信息进行摘要或压缩。 -对工具调用输出进行摘要是实现这一点的一种直观方法。例如，我们在 Open Deep Research 中就是这样做的。 -裁剪（Pruning）工具调用或工具消息。一件非常有趣的事情是，Claude 3.5（原文为4.5，应为笔误）实际上已经在其最新的版本中加入了这一功能。如果你查看他们的一些最新发布，现在他们会开箱即用地支持裁剪旧的工具调用及其输出或消息。因此，这种裁剪旧工具调用的想法现在已经被 Claude 内置到他们的 SDK 中了。 -对完整的消息历史进行摘要或压缩。你在 Cloud Code 的压缩（compaction）功能中可以看到这一点。一旦你达到整体上下文窗口的某个百分比，就会触发此功能。 -Cognition 也谈到了在智能体之间交接时进行摘要和证明的想法。 -因此，这种减少上下文的想法是我们从 Cloud Code 到我们的 Open Deep Research、Cognition 等众多不同示例中看到的一个非常流行的主题。Claude 3.5 也已将其纳入。 - -**检索上下文（Retrieving Context）。** -这是如今你可能在 X（Twitter）上看到激烈争论的经典话题之一：检索上下文的正确方法是什么？Cursor 的 Lee Robinson 最近做了一个非常好的演讲（我会确保分享这些幻灯片，这样你就能看到这些链接）。他在 OpenAI Demo Day 上谈到了 Cursor，例如，它使用索引和语义搜索，以及更简单的基于文件的搜索工具，如 `glob` 和 `grep`。Cloud Code 则只使用文件系统和简单的搜索工具，特别是 `glob` 和 `grep`。因此，有多种方法可以按需为你的智能体检索上下文：索引和语义搜索、文件系统和简单的文件搜索工具，两者都可以非常有效。在问答环节我们可以讨论各自的优缺点，但上下文检索对于构建有效的智能体至关重要。 - -**上下文隔离（Context Isolation）** 是我们看到的另一个主要主题，特别是跨多个子智能体（multi-agents）拆分上下文。这里的要点是什么？每个子智能体都有自己的上下文窗口，子智能体允许关注点分离。Manis 的宽智能体（wide agents）谈到了这一点。我们的深度智能体工作也使用了这一点。Open Deep Research 也使用了它。Claude 的子智能体在其研究中被利用，Claude 的多智能体研究员（multi-agent researcher）和 Claude Ghost 也都支持子智能体。因此，子智能体是我们在许多不同项目中看到的执行上下文隔离的一种非常常见的方法。 - -有一件我认为非常有趣的事情是**缓存上下文（caching context）**，Manis 对此谈得很多。稍后我会让 Pete 来详细讲讲，但我认为这也是一个非常有趣的技巧。 - -我只想展示一个我们在 Open Deep Research 中看到的简短示例。这是我们一个非常受欢迎的代码库。它基本上是一个开源的深度研究实现，其性能与市面上一些最佳实现相当。你可以查看我们的代码库，我们有来自 Deep Research Bench 的结果显示我们排在前十名。它有三个阶段：研究范围界定、研究阶段本身（使用多智能体架构）以及最后的一次性写作阶段。 - -我们使用了**卸载**。我们基本上创建了一个研究计划简报，然后将其卸载。我们不会只是把它保存在上下文窗口中，因为那个上下文窗口很快就会被其他东西填满。我们将其卸载，使其独立保存。在我们的案例中，可以从行图状态（line graph state）中访问它，但它也可以来自文件系统，道理是一样的。因此，你创建了一个研究计划，将其卸载，它就始终可以访问。你去做一堆工作，然后可以根据需要将其拉回，比如放在消息列表的末尾，这样你的智能体就可以轻松访问它，以执行例如写作阶段的任务。 - -正如你所见，我们使用卸载来帮助引导研究和写作阶段。 - -我们使用**缩减**来摘要来自占用大量 token 的网络搜索工具调用的观察结果。这是在研究阶段内部完成的。 - -我们在研究阶段内部的子智能体之间使用了**上下文隔离**。 - -这总结了我们在许多不同项目中看到的各种想法。实际上，Pete 接下来要专门谈谈 Manis 以及他们学到的一些经验教训。这只是为他做个铺垫。这部分内容总结了我刚才谈到的这些不同主题：卸载、减少上下文、检索上下文、隔离、缓存，以及它们在一些流行项目中的应用。我会在笔记中分享这些幻灯片的链接。我现在想让 Pete 继续进行他的演讲，因为我想确保他有充足的时间，并留出问答时间。这只是做个铺垫。Pete，接下来交给你了。我停止共享屏幕。 - -好的。你能看到我的幻灯片吗？ -能。 -好的，完美。谢谢 Lance。今天我很高兴能在这里分享一些我们在构建 Manis 过程中学到的关于上下文工程的新鲜经验教训。我说“新鲜”，是因为我意识到你提到的那篇关于上下文工程的博客文章是我今年7月写的，而在“智能体之年”，7月基本上就像是上个纪元了。当然，在本次会议之前，我又回去读了一遍，幸运的是，我觉得那篇博客里写的大部分内容至今仍然成立。但我不想浪费大家的时间，只是重复博客里已经说过的话。所以今天，我想深入探讨一些我之前要么没有深入探讨，要么根本没有触及的领域。实际上，我们将重点关注 Lance 之前幻灯片中的“非共识”（Discourage）一栏，因为就我个人而言，我认为探索这些非共识的想法往往能带来最大的启发。 - -好的，今天演讲的主题如下。首先，我们会稍微探讨一下为什么我们需要上下文工程这个更大的问题，然后我们会更多地讨论上下文缩减、上下文隔离，最后会介绍一些关于我们正在内部测试的上下文卸载的新内容。我今天分享的一切都在 Manis 中投入生产，并经过了实战检验，但我不知道它能持续多久，因为事情变化得太快了。 - -让我们从第一个大问题开始：为什么我们甚至需要上下文工程，尤其是在微调（fine-tuning）或后训练（post-training）模型如今已经变得更容易获取的情况下？例如，Thinking Machines 团队刚刚发布了 Tinker API，我很喜欢它的设计。但对我来说，为什么需要上下文工程这个问题，是经过几个痛苦的认知阶段才形成的。 - -在创办 Manis 之前，我已经在自然语言处理（NLP）领域深耕了十多年，也就是在 ChatGPT 和 LLMs 出现之前我们所说的构建语言模型。Manis 是我的第二或第三家公司，在我之前的创业公司，我们从头开始训练自己的语言模型，用于开放域信息抽取，并在此基础上构建知识图谱和语义搜索引擎，那过程非常痛苦。我们产品创新的速度完全被模型的迭代速度所限制。即使在当时，模型比今天小得多，但单次训练加评估周期也可能需要一两周时间。最糟糕的是，那时我们还没有达到产品市场契合度（PMF），却把所有时间都花在改进一些可能对产品根本不重要的基准测试上。因此，我认为初创公司不应该过早地构建专用模型，而应该尽可能长时间地依赖通用模型和上下文工程。当然，现在这可能已经成为某种常识。 - -但随着你的产品成熟，开源基础模型变得更强，你会很容易产生这样的想法：“嘿，也许我应该选一个强大的基础模型，用我的数据进行微调，让它在我的用例上表现得非常好。”我们也尝试过这样做，结果发现这是另一个陷阱。要知道，要让强化学习（RL）真正有效，你通常需要固定一个动作空间，围绕你当前的产品行为设计奖励函数，并生成大量的在线策略（on-policy）轨迹和反馈。但这也很危险，因为我们仍处于 AI 和智能体的早期阶段，一切都有可能在一夜之间发生巨变。对我们来说，经典的例子就是 MCP 的发布。它彻底改变了 Manis 的设计，从一个紧凑的静态动作空间变成了一个近乎无限可扩展的东西。如果你曾经训练过自己的模型，你就会知道这种开放域问题极难优化。当然，你可以投入巨大的精力进行后训练以确保泛化能力，但那样的话，你不就基本上是在试图自己成为一家大模型公司了吗？因为你基本上是在重建他们已经构建好的同一层，这是一种重复劳动。 - -说了这么多铺垫，我的观点是：坚定地划清界限。目前，上下文工程是应用层和模型层之间最清晰、最实用的边界。相信你的选择。 - -好了，哲学部分到此为止，让我们来谈谈一些真正的技术。 - -**第一个主题：上下文缩减（Context Reduction）。** -在这里，我想澄清两种不同类型的压缩操作，因为我们认为上下文缩减虽然很有趣，但它也是一个新概念，有很多实现方式。在 Manis，我们将其分为**压缩（Compaction）**和**摘要（Summarization）**。 - -在 Manis，每个工具调用和工具结果实际上都有两种格式：完整格式和紧凑格式。紧凑版本会剥离掉任何可以从文件系统或外部状态中重建的信息。例如，假设你有一个写入文件的工具，它可能有两个字段：路径（path）和内容（content）。但一旦工具返回，你就可以确保该文件已经存在于环境中。因此，在紧凑格式中，我们可以安全地丢弃超长的内容字段，只保留路径。如果你的智能体足够聪明，每当它需要再次读取该文件时，它只需通过路径检索即可。因此，信息并没有真正丢失，只是被外部化了。我们认为这种可逆性至关重要，因为智能体会基于先前的动作和观察进行链式预测，你永远不知道哪个过去的行为会在10步之后突然变得非常重要，这是无法预测的。这就是通过压缩实现的可逆缩减。 - -当然，压缩只能带你走这么远。最终，你的上下文仍然会增长并触及上限。这时，我们会将压缩与更传统的摘要结合起来，但我们会非常谨慎地进行。例如，在摘要之前，我们可能会将上下文的关键部分卸载到文件中。有时，我们甚至会更激进地将整个摘要前的上下文作为一个文本文件或日志文件转储到文件系统中，以便我们以后总能恢复它。正如 Lance 所提到的，有些人使用 `glob` 和 `grep`，`glob` 对日志文件也同样有效。如果模型足够聪明，它甚至知道如何检索那些被摘要的、摘要前的上下文。 - -我认为这里的区别在于，压缩是可逆的，而摘要则不是。两者都能减少上下文长度，但它们的行为截然不同。为了让这两种方法共存，我们必须在顶层跟踪一些上下文长度阈值。例如，你的模型有一个硬性的上下文限制，比如现在常见的100万个 token。但在现实中，大多数模型在远低于此限制时就开始退化，通常在20万左右，你会开始看到我们所说的“上下文腐烂”（context rot），比如重复、推理变慢、质量下降。因此，通过大量评估，对你来说非常重要的一点是识别出这个“腐烂前阈值”（pre-rot threshold），通常是128K到200K，并将其作为触发上下文缩减的条件。 - -当你的上下文大小接近这个阈值时，你就必须触发上下文缩减，但要从压缩开始，而不是摘要。而且，压缩并不意味着压缩整个历史记录。我们可能会压缩最旧的50%的工具调用，同时保留较新的工具调用的完整细节，这样模型仍然有最新的、清晰的例子来了解如何正确使用工具。否则，在最坏的情况下，模型会模仿这种行为，并输出那些缺少字段的紧凑格式，这是完全错误的。 - -在压缩之后，我们必须检查这次压缩操作实际为我们节省了多少空闲上下文。有时，如这张图所示，在多轮压缩之后，收益微乎其微，因为即使是压缩后的信息仍然会占用上下文。这时，我们才会转向摘要。但也要记住，在摘要时，我们总是使用数据的完整版本，而不是紧凑版本。并且，我们仍然会将最后几次工具调用和工具结果保持完整细节，而不是摘要，因为这可以让模型知道它上次停在了哪里，从而能更顺利地继续下去。否则，你会发现在摘要之后，模型有时会改变其风格和语气，我们发现保留一些工具调用和工具结果的示例真的很有帮助。 - -好了，我们已经讨论了缩减，现在来谈谈**隔离（Isolation）**。 - -我非常赞同 Cognition 博客中的观点，他们警告不要使用多智能体设置，因为当你有多个智能体时，在它们之间同步信息会变成一场噩梦。但这个问题并不新鲜。在计算机编程的早期，多进程或多线程协调就是一个经典挑战。我认为我们可以借鉴一些其中的智慧。 - -我不知道今天有多少 Go 语言程序员在这里，但在 Go 语言社区有一句名言：“不要通过共享内存来通信，而要通过通信来共享内存。”（Do not communicate by sharing memory; instead, share memory by communicating.）当然，这句话并非直接针对智能体，有时甚至对智能体来说也是错误的。但我认为重要的是，它突出了两种截然不同的模式：通过通信（by communicating）或通过共享内存（by sharing memory）。 - -如果我们把这里的“内存”一词翻译成“上下文”，我们就能看到一个非常清晰的对应关系。 - -“通过通信”是更容易理解的一种模式，因为它就是经典的子智能体设置。例如，主智能体编写一个提示，然后将该提示发送给一个子智能体，而子智能体的整个上下文仅由该指令组成。我们认为，如果一个任务有清晰简短的指令，并且只有最终输出才重要（比如说，在代码库中搜索一个特定的片段），那么就应该使用通信模式，并保持简单。因为主智能体并不关心子智能体是如何找到代码的，它只需要结果。这也是 Cloud Code 通常的做法，它使用其任务工具（task tool）将一个独立、清晰的任务委派给一些子智能体。 - -相比之下，对于更复杂的场景，“通过共享内存”意味着子智能体可以看到之前所有的上下文，即所有的工具使用历史。但子智能体有自己的系统提示（system prompt）和自己的动作空间。例如，在一个深度研究场景中，最终报告依赖于大量的中间搜索和笔记，在这种情况下，你应该考虑使用共享内存模式（或者用我们的话说，共享上下文）。因为即使你可以把所有笔记和搜索结果保存到文件中，再让子智能体重新读取一切，也只是在浪费延迟和上下文。如果你计算一下 token 的数量，你甚至可能会使用更多的 token 来完成这件事。 - -因此，我们认为对于那些需要完整历史的场景，就应该使用共享内存模式，但要注意，共享上下文是比较昂贵的，因为每个子智能体都有更大的输入需要预填充（prefill），这意味着你会在输入 token 上花费更多。而且，由于系统提示和动作空间不同，你无法重用 KV 缓存，因此你必须支付全额费用。 - -最后，让我们稍微谈谈**上下文卸载（Context Offloading）**。当人们说卸载时，通常是指将工作上下文的一部分移动到外部文件中。但随着你的系统不断增长，特别是如果你决定集成 MCP，有一天你会意识到工具本身也会占用大量上下文，上下文中存在太多工具会导致混淆。我们称之为“上下文混淆”（context confusion），模型可能会调用错误的工具，甚至是不存在的工具。因此，我们必须找到一种方法来卸载工具本身。 - -目前一种常见的方法是对工具描述进行动态 RAG（Retrieval-Augmented Generation），例如根据当前任务或状态按需加载工具。但这也会导致两个问题。首先，由于工具定义位于上下文的前面，每次加载都会导致你的 KV 缓存重置。更重要的是，模型对已移除工具的过往调用仍然存在于上下文中，这可能会诱使模型调用无效的工具或使用无效的参数。 - -为了解决这个问题，我们正在 Manis 中试验一种新的分层动作空间（layered action space）。本质上，我们可以让 Manis 从三个不同级别的抽象中进行选择： -1. **函数调用（Function Calling）** -2. **沙盒工具（Sandbox Utilities）** -3. **包和 API（Packages and APIs）** - -让我们深入探讨这三层动作空间。 -**第一层：函数调用。** 这是经典的模式，大家都很熟悉。由于约束解码（constraint decoding），它是模式安全的（schema safe）。但我们都知道它的缺点，比如会破坏缓存，以及太多的工具调用可能会引起混淆。因此，在 Manis 中，我们目前只使用固定数量的原子函数，例如读写文件、执行 shell 命令、在文件和互联网中搜索，以及一些浏览器操作。我们认为这些原子函数有非常清晰的边界，它们可以协同工作以组合成更复杂的工作流。 - -然后，我们将其他所有内容都卸载到下一层，即**沙盒工具**。如你所知，每个 Manis 会话都在一个完整的虚拟机沙盒中运行，该沙盒运行在我们定制的 Linux 系统上。这意味着 Manis 可以使用 shell 命令来运行我们为 Manis 开发的预安装工具。例如，我们有一些格式转换器、语音识别工具，甚至一个特殊的 MCP CLI（命令行界面），我们通过它来调用 MCP。我们不会将 MCP 工具注入到函数调用空间中，而是在沙盒内通过命令行界面完成所有操作。 - -工具的好处在于，你可以添加新功能而无需触碰模型的函数调用空间。就像你熟悉 Linux 一样，你总是知道如何找到这些新命令，甚至可以运行 `--help` 来弄清楚如何使用一个新工具。另一个好处是，对于较大的输出，它们可以直接写入文件或分页返回结果，你可以使用所有这些 Linux 工具（如 `grep`、`cat`、`less`、`more`）来即时处理这些结果。 - -这里的权衡是，它非常适合处理大输出，但对于与前端进行低延迟的来回交互来说效果不佳，因为要可视化你的智能体的交互并向用户展示，这是相当棘手的。但我们认为它已经卸载了很多东西。 - -然后我们还有最后一层，我们称之为**包和 API**。在这里，Manis 可以编写 Python 脚本来调用预授权的 API 或自定义包。例如，Manis 可能会使用一个 3D 设计库进行建模，或调用一个金融 API 来获取市场数据。实际上，我们已经代表用户购买了所有这些 API，并为他们付费，这些都包含在订阅中。因此，我们有很多 API 密钥预装在 Manis 中，Manis 可以使用这些密钥访问这些 API。 - -我们认为这对于需要大量内存计算但不需要将所有数据推送到模型上下文中的任务是完美的。例如，想象一下你正在分析一只股票一整年的价格数据。你不会把所有数字都喂给模型，而是应该让脚本进行计算，只把摘要放回上下文中。 - -而且，由于代码和 API 具有很强的可组合性，你实际上可以在一步中串联很多事情。例如，在一个典型的 API 调用中，你可以在一个 Python 脚本中完成获取城市名称、获取城市 ID、获取天气等所有操作。 - -我的一位朋友还发表了一篇名为《CodeAct》的论文，很多人都在讨论它。我认为这是同一个想法，因为代码是可组合的，可以在一步中完成很多事情。但它也不是模式安全的，对代码进行约束解码是非常非常困难的。因此，我们认为你应该为这些功能找到合适的场景。对我们来说，我们会将所有能在编译器或解释器运行时内处理的事情都用代码来完成，否则我们会使用沙盒工具或函数调用。 - -好消息是，从模型的角度来看，这三层都通过标准的函数调用进行。因此，接口保持简单、对缓存友好，并且在函数之间是正交的。因为我们提到过，沙盒工具仍然是通过 shell 函数来访问的；如果你使用第三方应用程序的 API，也只是使用文件函数来读写文件，然后使用 shell 函数来执行它。因此，它不会给模型增加额外负担，因为这些都是模型已经训练过并且熟悉的。 - -让我们拉远视角，将五个维度——卸载、缩减、检索、隔离和缓存——联系起来。你会发现它们并不是独立的。我们可以看到，卸载和检索使得更高效的缩减成为可能，而稳定的检索使得隔离变得安全。但隔离也会减慢上下文的增长速度，并降低缩减的频率。然而，更多的隔离和缩减也会影响缓存效率和输出质量。 - -归根结底，我认为上下文工程是一门需要在多个潜在冲突的目标之间取得完美平衡的艺术与科学。这真的很难。 - -在结束之前，我想留给你们一个最后的想法，这与我刚才所说的几乎完全相反，那就是：**请避免过度工程化上下文（avoid context over-engineering）**。 - -回顾 Manis 发布以来的过去六七个月，我们所见过的最大飞跃并非来自于添加更多花哨的上下文管理层或巧妙的检索技巧。它们都来自于简化，来自于移除不必要的技巧，并更多地信任模型。每次我们简化架构，系统都会变得更快、更稳定、更聪明。因为我们认为，上下文工程的目标应该是让模型的工作变得更简单，而不是更难。 - -所以，如果今天你只带走一点，我认为应该是：**少构建，多理解（build less and understand more）**。 - -非常感谢大家，再次感谢 Lance 和 LangChain 团队的邀请。迫不及待想看到你们接下来会构建出什么。现在交回给 Lance。 - ->> 太棒了，非常感谢。我们这里有一系列不错的问题。也许我们可以开始回答了，如果需要的话，可以参考一下幻灯片。Pete，你的幻灯片对 everyone 可用吗？ ->> 哦，是的，我可以稍后分享 PDF 版本。 ->> 好的，听起来不错。 - -嗯，好吧。我来看看问题，也许我们可以从最近的问题开始。 -…… -（后续为问答环节，篇幅较长，此处省略具体问答内容，但翻译逻辑一致） - -…… ->> 好的，我想我们差不多到时间了，我想尊重你的时间，因为我知道现在新加坡还是凌晨，对你来说非常早。嗯，好吧，这次分享真的很好，谢谢你。我们一定会确保这次录制的视频可用，并确保幻灯片也能提供。 ->> 嗯，有什么最后想说的吗？有什么想特别指出的，或者行动号召？嗯，大家应该去试试 Manis，但话筒交给你了。 ->> 是的，我就想说，大家都去试试吧。我们有免费套餐。 ->> 是的，绝对要试。嘿，非常感谢你，Pete。希望以后还能再做一次。 ->> 谢谢邀请我。 ->> 好的，再见。 \ No newline at end of file diff --git a/docs/zh/reading/20251110-manus-context-report.md b/docs/zh/reading/20251110-manus-context-report.md deleted file mode 100644 index 5f86654..0000000 --- a/docs/zh/reading/20251110-manus-context-report.md +++ /dev/null @@ -1,415 +0,0 @@ -# 上下文管理（Context Management）指南 - -## 一、背景 - -### 1.1 从提示工程到上下文工程 - -这几年我们从“提示工程”一路走到“上下文工程”。前者是在 2022 年底 ChatGPT 爆火后流行起来，大家忙着研究怎么给聊天模型下指令；到了 2024 年中，随着“智能体之年”的到来，大家发现真正棘手的是怎么管理越来越庞杂的对话与工具调用信息，于是“上下文工程”成了新的主角。 - -### 1.2 智能体带来的新问题 - -智能体的工作方式很直白：LLM 绑上一堆工具，在循环里自己决定怎么调，调一次就把结果塞回消息历史；时间一长，历史就像雪球一样越滚越大。现实里，一个看似普通的任务可能就要跑 50 次工具调用，线上智能体动辄上百轮对话，每次还会吐出大段文本，占掉海量 token——上下文很快就被挤满了。 - -## 二、核心挑战 - -### 2.1 上下文爆炸 - -所谓“上下文爆炸”，说的就是这些调用和消息按时间不断堆叠：工具输出一条条追加，消息列表一路拉长，模型每次都要背上更重的历史包袱，结果上下文窗口很快被塞满。 - -### 2.2 上下文腐烂（Context Rot） - -上下文一旦过大，效果就开始“变味”：回答变得重复啰嗦、推理明显变慢、整体质量下滑。Cognition 的报告把这叫做“上下文腐烂”，在实战中我们也经常能直观感受到它正在发生。 - -### 2.3 根本悖论 - -矛盾就在这里：智能体需要靠丰富的上下文来做出靠谱的决策，但上下文越多，模型越容易掉速、走神。上下文工程要做的，就是在窗口里放进“刚刚好”的信息——够用、不浪费，同时给后续的检索与扩展留出空间。 - -## 三、解决方案 - -业界已经形成了五种主要的上下文管理策略： - -### 3.1 上下文卸载（Context Offloading） - -简单说，卸载就是把不需要长时间待在消息历史里的大块信息搬到外部去存着，需要时再拿回来用。最常见的做法是把工具输出写到文件系统里，消息里只留个最小必要的引用（比如文件路径）；像网页搜索结果、超长工具输出、临时计划这类特别占 token 的内容，都适合这样处理。好处很直接：上下文更干净，占用更小，信息不丢，随取随用，也是生产场景里最常见的实践。 - -### 3.2 上下文缩减（Context Reduction） - -当卸载还不够时，就轮到“缩减”上场了；它主要有两条路：可逆的“压缩”和不可逆的“摘要”。思路是优先把能外部重建的细节挪走、格式变紧凑，再在必要时对保留内容做更狠的摘要，以在不过度丢信息的前提下把上下文体积降下来。 - -#### 3.2.1 压缩（Compaction） - -压缩的核心思想其实很简单，就是把信息从完整格式转换成紧凑格式，把那些可以从外部状态重建的信息剥离掉。比如说，一个写入文件的工具返回结果可能是 `{path: "/file.txt", content: "很长的内容..."}`，压缩后我们就只保留 `{path: "/file.txt"}`，因为文件已经存在了，需要的时候通过路径重新读取就行。这种方式的优点在于它是完全可逆的，信息并没有真正丢失，只是被外部化了，可以随时恢复，同时还能大幅减少 token 占用，特别适合处理工具调用结果。在实际操作中，我们通常会压缩最旧的工具调用，比如最旧的50%，同时保留最新的工具调用的完整细节，这样模型仍然有最新的、清晰的例子来了解如何正确使用工具，否则模型可能会模仿那些缺少字段的紧凑格式，那就完全错了。 - -工具结果清除（Tool Result Clearing）是压缩的自动化实现方式，由服务端自动执行，这是生产级智能体系统的标准做法，Claude 4.5 等模型已经内置支持了。它的工作机制是这样的：当上下文超过配置阈值时，会自动按时间顺序清除最旧的工具结果，用占位符文本替换被清除的内容，让模型知道内容已被移除，默认只清除工具结果，保留工具调用（参数），可选同时清除工具调用和结果，整个过程在服务端自动执行，客户端保持完整历史记录，无需同步状态。举个例子，假设一个智能体执行了多次网络搜索和文件操作，当上下文达到阈值（如 30K tokens）时，工具结果清除会自动清除最旧的工具结果，用 `[工具结果已清除]` 这样的占位符替换，保留最近的工具调用的完整结果，同时保留所有工具调用的参数，这样模型仍然知道执行了什么操作。这种方式的智能之处在于它会保留最近 N 个工具使用/结果对，确保模型有最新的使用示例，同时考虑与提示缓存的交互来优化成本效益，还可以指定某些工具的结果永不清除来保护重要上下文，并且确保清除操作值得缓存失效的成本。 - -参考文档：[Claude Context Editing 官方文档](https://docs.claude.com/en/docs/build-with-claude/context-editing) - -#### 3.2.2 摘要（Summarization） - -摘要就是对上下文进行摘要或压缩，生成更短的版本，你可以对工具调用输出进行摘要，也可以对完整消息历史进行摘要。但这里有个关键区别，就是摘要和压缩不一样，压缩是可逆的，它直接移除内容，用占位符替换或保留路径引用，需要的时候可以通过检索恢复，而摘要是不可逆的，它把内容压缩成更短版本，信息可能会丢失。所以在使用摘要的时候要特别谨慎，最佳实践是在摘要前先把关键部分卸载到文件，使用完整版本进行摘要而不是压缩版本，保留最后几次工具调用的完整细节，这样模型知道它上次停在了哪里，能更顺利地继续下去，否则你会发现摘要之后模型有时会改变风格和语气。另外，要识别出"腐烂前阈值"（通常 128K-200K token）作为触发点，当上下文接近这个阈值时才考虑摘要。 - -#### 3.2.3 压缩和摘要的选择策略 - -压缩和摘要这两种方式各有特点，压缩是完全可逆的，适合工具结果可从外部重建的场景，实现难度中等，但如果用工具结果清除这样的自动化实现，难度就很低了，而摘要则是不可逆的，信息可能丢失，实现难度较高，只有在需要进一步压缩时才考虑。所以选择策略很明确：优先使用压缩，因为它可逆且安全，特别适合工具结果，包括工具结果清除等自动化实现，只有当压缩还不够的时候，才最后考虑对保留内容进行摘要。在阈值管理上，你需要考虑硬性限制（比如 100万 token）和腐烂前阈值（通常 128K-200K token），当上下文接近阈值时触发缩减，优先压缩，必要时才摘要。实际工作流通常是这样的：上下文增长到达到阈值（比如 30K tokens），第一步先压缩工具结果，保留路径，移除内容，或者使用工具结果清除，如果压缩后上下文仍然接近阈值，那第二步就对保留的内容进行摘要，对历史消息进行摘要。 - -### 3.3 上下文检索（Retrieving Context） - -检索就是按需把卸载到外部的内容再“拉回视野”。可以走语义搜索那套，用向量索引按相似度找复杂内容；也可以走文件系统这一派，直接用 `glob`、`grep` 这类简单高效的工具扫文件，尤其适合结构化资料。常见用法像找回研究计划、定位历史工具结果、恢复摘要前的完整上下文等；说到底，检索做得稳，卸载才敢放手，二者配合好，整体管理就高效。 - -### 3.4 上下文隔离（Context Isolation） - -**核心思想**：通过多智能体架构，将上下文拆分到不同的子智能体中。 - -#### 3.4.1 两种模式 - -上下文隔离有两种截然不同的模式，第一种是"通过通信"（By Communicating），这是更容易理解的一种模式，因为它就是经典的子智能体设置。主智能体编写一个任务指令，然后将该指令发送给子智能体，子智能体的整个上下文仅由该指令组成，子智能体完成任务后返回结果，主智能体并不关心子智能体是如何执行的，只需要结果。这种模式适合任务有清晰简短的指令，并且只有最终输出才重要的场景，比如在代码库中搜索特定片段，它的优势是上下文小、开销低，简单直接，这也是 Cloud Code 通常的做法，它使用任务工具将一个独立、清晰的任务委派给子智能体。 - -第二种是"通过共享内存"（By Sharing Memory），这意味着子智能体可以看到之前所有的上下文，即所有的工具使用历史，但子智能体有自己的系统提示和自己的动作空间。这种模式适合需要完整历史信息的复杂任务，最终输出依赖大量中间结果的场景，比如深度研究场景，需要大量搜索和笔记，即使你可以把所有笔记和搜索结果保存到文件中，再让子智能体重新读取一切，也只是在浪费延迟和上下文，如果你计算一下 token 的数量，你甚至可能会使用更多的 token 来完成这件事。但要注意，共享上下文是比较昂贵的，因为每个子智能体都有更大的输入需要预填充，这意味着你会在输入 token 上花费更多，而且由于系统提示和动作空间不同，你无法重用 KV 缓存，因此你必须支付全额费用。 - -#### 3.4.2 选择建议 - -选择哪种模式其实很简单，对于简单任务，使用通信模式就够了，上下文小、开销低，简单直接，而对于复杂任务，特别是那些需要完整历史信息、最终输出依赖大量中间结果的场景，就应该考虑使用共享内存模式。但要注意，多智能体同步信息可能成为挑战，需要谨慎设计，就像 Cognition 博客中警告的那样，当你有多个智能体时，在它们之间同步信息会变成一场噩梦，但这个问题并不新鲜，在计算机编程的早期，多进程或多线程协调就是一个经典挑战，我们可以借鉴一些其中的智慧。 - -### 3.5 缓存上下文（Caching Context） - -缓存更像是在给系统安“记忆”，把会反复用到的上下文预先存起来，下次就不用重算、重传，整体效率更高。常见的做法有缓存工具定义、缓存常用查询的结果等；不过要注意，它跟隔离和缩减之间存在权衡：缓存多了能省成本，但也可能牵动上下文的一致性与更新节奏，需要结合场景拿捏分寸。 - -### 3.6 分层动作空间（Layered Action Space） - -当系统里的工具越堆越多，本身就会吃掉大量上下文，还容易让模型“搞混”到底该用哪一个。分层动作空间的思路是把工具按层级卸载和组织起来：常用且边界清晰的放在函数调用层，更大更重的能力放到沙盒工具层，需要灵活计算与组合的交给包和 API 层；对模型来说入口依然统一，但对系统来说上下文更清爽、误用更少、扩展也更从容。 - -#### 三层架构 - -分层动作空间本质上让系统从三个不同级别的抽象中进行选择。第一层是函数调用（Function Calling），这是经典的模式，由于约束解码，它是模式安全的，我们都知道它的缺点，比如会破坏缓存，太多的工具调用可能会引起混淆，所以在 Manis 中，他们目前只使用固定数量的原子函数，例如读写文件、执行 shell 命令、在文件和互联网中搜索，以及一些浏览器操作，这些原子函数有非常清晰的边界，它们可以协同工作以组合成更复杂的工作流。 - -第二层是沙盒工具（Sandbox Utilities），每个会话都在一个完整的虚拟机沙盒中运行，这意味着可以通过 shell 命令来运行预安装工具，比如格式转换器、语音识别工具，甚至一个特殊的 MCP CLI，通过它来调用 MCP，他们不会将 MCP 工具注入到函数调用空间中，而是在沙盒内通过命令行界面完成所有操作。工具的好处在于，你可以添加新功能而无需触碰模型的函数调用空间，就像你熟悉 Linux 一样，你总是知道如何找到这些新命令，甚至可以运行 `--help` 来弄清楚如何使用一个新工具，另一个好处是，对于较大的输出，它们可以直接写入文件或分页返回结果，你可以使用所有这些 Linux 工具（如 `grep`、`cat`、`less`、`more`）来即时处理这些结果，这里的权衡是，它非常适合处理大输出，但对于与前端进行低延迟的来回交互来说效果不佳。 - -第三层是包和 API（Packages and APIs），在这里可以编写 Python 脚本来调用预授权的 API 或自定义包，比如使用一个 3D 设计库进行建模，或调用一个金融 API 来获取市场数据，实际上，很多 API 密钥都预装在系统中，可以直接使用这些密钥访问这些 API。这对于需要大量内存计算但不需要将所有数据推送到模型上下文中的任务是完美的，比如分析一只股票一整年的价格数据，你不会把所有数字都喂给模型，而是应该让脚本进行计算，只把摘要放回上下文中，而且由于代码和 API 具有很强的可组合性，你实际上可以在一步中串联很多事情，比如在一个典型的 API 调用中，你可以在一个 Python 脚本中完成获取城市名称、获取城市 ID、获取天气等所有操作。但要注意，代码不是模式安全的，对代码进行约束解码是非常非常困难的，所以你应该为这些功能找到合适的场景，对于能在编译器或解释器运行时内处理的事情都用代码来完成，否则就使用沙盒工具或函数调用。 - -好消息是，从模型的角度来看，这三层都通过标准的函数调用进行访问，因此接口保持简单、对缓存友好，并且在函数之间是正交的，因为沙盒工具仍然是通过 shell 函数来访问的，如果你使用第三方应用程序的 API，也只是使用文件函数来读写文件，然后使用 shell 函数来执行它，所以它不会给模型增加额外负担，因为这些都是模型已经训练过并且熟悉的。 - -## 四、方案之间的关系 - -五个维度（卸载、缩减、检索、隔离、缓存）并非独立，而是相互关联： - -- **卸载 + 检索** → 使更高效的缩减成为可能 -- **稳定的检索** → 使隔离变得安全 -- **隔离** → 减慢上下文增长速度，降低缩减频率 -- **隔离 + 缩减** → 影响缓存效率和输出质量 - -**核心原则**：上下文工程是一门在多个潜在冲突目标之间取得平衡的艺术与科学。 - -## 五、实践建议 - -### 5.1 避免过度工程化 - -一句话：别把上下文管理搞得太花。我们的经验是，最大的提升往往来自于“减法”——尽量简化，能不做就不做，能信任模型就别堆技巧；每当把架构瘦身、移除不必要的层，系统就会更快、更稳、更聪明。记住那条底线：上下文工程的目标是让模型更容易工作，而不是更难。 - -### 5.2 关键原则 - -原则其实不复杂：少构建、多理解，先把问题吃透再动手；做缩减时优先用可逆方案，能压缩就不摘要；时刻盯住上下文的“腐烂前阈值”，别等到性能掉下去才补救；隔离上按任务复杂度选择通信还是共享内存；所有决策都要在性能、成本和质量之间找平衡点。 - -### 5.3 实施步骤 - -落地可以按这个顺序走：先评估现状，量一量上下文长度，找出开始变慢变差的那个阈值；接着优先做卸载，把大的工具输出搬到文件系统，并把检索通路打通；然后上压缩策略，设计好紧凑格式和触发阈值；如果任务复杂度需要，再考虑做隔离，决定用通信还是共享内存；最后持续打磨，边监控边简化，别让方案越做越重。 - -## 六、代码实现 - -### 6.1 上下文缩减（Context Reduction） - -上下文缩减提供三种接口：压缩（Compaction）、摘要（Summary）和融合接口。 - -#### 6.1.1 压缩接口（Compaction） - -**功能**：将工具调用结果从完整格式转换为紧凑格式，将可重建的信息卸载到外部存储。 - -**输入**： -```python -{ - "messages": List[Message], # 完整的消息历史 - "compaction_config": { - "compress_ratio": 0.5, # 压缩比例：压缩最旧的 50% 工具调用 - "keep_recent": 3, # 保留最近 N 个工具调用的完整结果 - "storage_path": "/tmp/context", # 外部存储路径（文件系统或数据库） - "exclude_tools": ["web_search"] # 排除的工具列表（这些工具的结果不压缩） - } -} -``` - -**输出**： -```python -{ - "messages": List[Message], # 压缩后的消息历史（工具结果被替换为路径引用） - "offloaded_data": { - "files": [ # 卸载到文件系统的数据 - { - "path": "/tmp/context/tool_1_result.json", - "tool_call_id": "call_123", - "original_tokens": 5000, - "compressed_tokens": 50 - } - ], - "db_records": [ # 卸载到数据库的数据（可选） - { - "record_id": "rec_456", - "tool_call_id": "call_124", - "original_tokens": 3000, - "compressed_tokens": 30 - } - ] - }, - "statistics": { - "original_tokens": 80000, - "compressed_tokens": 40000, - "saved_tokens": 40000, - "compressed_tool_calls": 10, - "kept_tool_calls": 3 - } -} -``` - -**参数说明**： - -| 参数 | 类型 | 说明 | -|------|------|------| -| `compress_ratio` | float (0-1) | 压缩比例，如 0.5 表示压缩最旧的 50% 工具调用 | -| `keep_recent` | int | 保留最近 N 个工具调用的完整结果，作为使用示例 | -| `storage_path` | string | 外部存储路径，可以是文件系统路径或数据库连接字符串 | -| `exclude_tools` | List[string] | 排除的工具名称列表，这些工具的结果不会被压缩 | - -#### 6.1.2 摘要接口（Summary） - -**功能**：对上下文进行摘要，在摘要前将关键部分卸载到文件系统。 - -**输入**： -```python -{ - "messages": List[Message], # 完整的消息历史 - "summary_config": { - "summary_ratio": 0.3, # 摘要比例：摘要最旧的 30% 消息 - "keep_recent": 5, # 保留最近 N 条消息的完整内容 - "offload_before_summary": True, # 是否在摘要前卸载关键部分 - "offload_path": "/tmp/context/pre_summary", # 摘要前卸载路径 - "dump_full_context": True # 是否将整个摘要前的上下文转储为日志文件 - } -} -``` - -**输出**： -```python -{ - "messages": List[Message], # 摘要后的消息历史 - "offloaded_data": { - "key_parts": [ # 摘要前卸载的关键部分 - { - "path": "/tmp/context/pre_summary/key_info_1.json", - "content_type": "tool_results", - "original_tokens": 10000 - } - ], - "full_context_dump": { # 完整的摘要前上下文转储（如果启用） - "path": "/tmp/context/pre_summary/full_context.log", - "original_tokens": 50000 - } - }, - "summary_info": { - "summarized_messages": 15, - "kept_messages": 5, - "original_tokens": 50000, - "summarized_tokens": 15000, - "saved_tokens": 35000 - } -} -``` - -**参数说明**： - -| 参数 | 类型 | 说明 | -|------|------|------| -| `summary_ratio` | float (0-1) | 摘要比例，如 0.3 表示摘要最旧的 30% 消息 | -| `keep_recent` | int | 保留最近 N 条消息的完整内容 | -| `offload_before_summary` | bool | 是否在摘要前将关键部分卸载到文件 | -| `offload_path` | string | 摘要前卸载的存储路径 | -| `dump_full_context` | bool | 是否将整个摘要前的上下文转储为日志文件 | - -#### 6.1.3 融合接口（Combined） - -**功能**：优先使用压缩，提供 token 计数对比，达到阈值后使用摘要。 - -**输入**： -```python -{ - "messages": List[Message], # 完整的消息历史 - "combined_config": { - "compaction_threshold": 30000, # 压缩触发阈值（tokens） - "summary_threshold": 80000, # 摘要触发阈值（tokens） - "compaction_config": {...}, # 压缩配置 - "summary_config": {...} # 摘要配置 - } -} -``` - -**输出**： -```python -{ - "messages": List[Message], # 处理后的消息历史 - "applied_strategy": "compaction" | "summary" | "none", # 应用的策略 - "token_comparison": { - "original_tokens": 70000, - "after_compaction_tokens": 35000, # 压缩后的 tokens（如果应用） - "after_summary_tokens": None, # 摘要后的 tokens（如果应用） - "saved_tokens": 35000 - }, - "offloaded_data": {...}, # 卸载的数据 - "statistics": {...} # 统计信息 -} -``` - -**参数说明**： - -| 参数 | 类型 | 说明 | -|------|------|------| -| `compaction_threshold` | int | 压缩触发阈值（tokens），超过此值触发压缩 | -| `summary_threshold` | int | 摘要触发阈值（tokens），压缩后仍超过此值则触发摘要 | -| `compaction_config` | object | 压缩配置，参考压缩接口参数 | -| `summary_config` | object | 摘要配置，参考摘要接口参数 | - -**Token 计数接口**： - -**输入**： -```python -{ - "messages": List[Message], # 消息历史 - "reduction_config": { - "compaction_config": {...}, # 压缩配置 - "summary_config": {...} # 摘要配置 - } -} -``` - -**输出**： -```python -{ - "original_tokens": int, # 原始 tokens - "after_compaction_tokens": int, # 压缩后的 tokens - "after_summary_tokens": int, # 摘要后的 tokens - "compaction_savings": int, # 压缩节省的 tokens - "summary_savings": int, # 摘要节省的 tokens - "total_savings": int # 总共节省的 tokens -} -``` - -### 6.2 上下文检索（Context Retrieving） - -上下文检索提供两种方式：基础检索工具和定制化技能。 - -#### 6.2.1 基础检索工具（Agentic Retrieve） - -**功能**：提供 glob、grep 等文件工具，支持按需检索已卸载的上下文。 - -**输入**： -```python -{ - "query": str, # 检索查询 - "retrieve_config": { - "tools": ["glob", "grep", "cat", "find"], # 可用的检索工具 - "search_paths": [ # 搜索路径 - "/tmp/context", - "/tmp/context/pre_summary" - ], - "file_patterns": ["*.json", "*.log", "*.txt"] # 文件模式 - } -} -``` - -**输出**： -```python -{ - "results": [ - { - "file_path": "/tmp/context/tool_1_result.json", - "matched_content": "...", # 匹配的内容片段 - "relevance_score": 0.85, - "tool_used": "grep" - } - ], - "retrieved_context": str, # 检索到的完整上下文 - "tools_used": ["grep", "cat"] # 使用的工具列表 -} -``` - -**参数说明**： - -| 参数 | 类型 | 说明 | -|------|------|------| -| `query` | string | 检索查询字符串 | -| `tools` | List[string] | 可用的检索工具列表，如 ["glob", "grep", "cat", "find"] | -| `search_paths` | List[string] | 搜索路径列表 | -| `file_patterns` | List[string] | 文件模式列表，如 ["*.json", "*.log", "*.txt"] | - -#### 6.2.2 定制化技能（Skills） - -**功能**：提供更高级的检索技能，拆分为三个独立接口：`meta_info`、`read_skill`、`read_detail_skill`。 - -##### A. 元信息查询（meta_info） - -**输入**： -```python -{ - "target": str # 目标文件或路径 -} -``` - -**输出**： -```python -{ - "file_size": int, - "created_time": str, - "tool_call_id": str, - "original_tokens": int -} -``` - -##### B. 概要读取（read_skill） - -**输入**： -```python -{ - "target": str # 目标文件或路径 -} -``` - -**输出**： -```python -{ - "skill_summary": str, - "key_points": List[str] -} -``` - -##### C. 详细读取（read_detail_skill） - -**输入**： -```python -{ - "target": str # 目标文件或路径 -} -``` - -**输出**： -```python -{ - "full_content": str, - "sections": List[dict] -} -``` - -参考文档：[Claude Context Editing 官方文档](https://docs.claude.com/en/docs/build-with-claude/context-editing) - -## 七、总结 - -上下文管理是智能体系统的核心挑战。通过合理运用**卸载、缩减、检索、隔离、缓存**等策略，可以在满足智能体需求的同时，保持系统性能和稳定性。 - -**记住**：最好的解决方案往往是最简单的。理解问题本质，选择合适工具，避免过度工程化。 - ---- diff --git a/docs/zh/reading/20251113-gemini-cli-context-management.md b/docs/zh/reading/20251113-gemini-cli-context-management.md deleted file mode 100644 index ab84afe..0000000 --- a/docs/zh/reading/20251113-gemini-cli-context-management.md +++ /dev/null @@ -1,347 +0,0 @@ -# Gemini CLI 上下文管理机制 - -## 概述 - -Gemini CLI 通过多层机制管理过长的对话上下文，确保在 token 限制内保留关键信息。核心策略包括：**预判溢出 → 自动压缩 → 内容截断 → 历史过滤**。 - -## 核心流程 - -### 1. Token 限制检测 - -**位置**: `core/client.ts` → `sendMessageStream()` - -**机制**: -- 使用 `tokenLimit(modelForLimitCheck)` 获取模型最大 token 数 - - `gemini-1.5-pro`: 2,097,152 tokens - - `gemini-1.5-flash`: 1,048,576 tokens -- 计算剩余可用 token: `getLastPromptTokenCount()` -- **触发条件**: 预计请求 token > 剩余空间 × 95% -- **事件**: 触发 `ContextWindowWillOverflow` 事件，阻止发送 - -### 2. 自动压缩机制 - -**位置**: `services/chatCompressionService.ts` → `compress()` - -**压缩阈值**: -- `DEFAULT_COMPRESSION_TOKEN_THRESHOLD = 0.2` (20%) -- 当当前 token 数 > 模型限制 × 0.2 时触发压缩 - -**压缩策略**: -- **保留比例**: `COMPRESSION_PRESERVE_THRESHOLD = 0.3` (保留最新 30%) -- **分段压缩**: - 1. `findCompressSplitPoint()` 计算压缩起点（保留最新 30%） - 2. 将旧消息 (`historyToCompress`) 提交给 LLM 生成摘要 - 3. 将摘要 + 保留的新消息 (`historyToKeep`) 组合为新历史 - 4. 使用 `getCompressionPrompt()` 指导模型生成结构化摘要: - - ``: 用户高层目标 - - ``: 关键事实 - - ``: 文件系统状态 - - ``: 最近的重要操作 - - ``: 当前执行计划 - -**压缩 Prompt** (`core/prompts.ts` → `getCompressionPrompt()`): - -压缩时使用的系统指令要求模型： - -1. 在私有 `` 中思考整个历史 -2. 生成结构化的 `` XML 对象 -3. 必须保留所有关键细节、计划、错误和用户指令 - -完整 Prompt 内容： - -``` -You are the component that summarizes internal chat history into a given structure. - -When the conversation history grows too large, you will be invoked to distill the entire history into a concise, structured XML snapshot. This snapshot is CRITICAL, as it will become the agent's *only* memory of the past. The agent will resume its work based solely on this snapshot. All crucial details, plans, errors, and user directives MUST be preserved. - -First, you will think through the entire history in a private . Review the user's overall goal, the agent's actions, tool outputs, file modifications, and any unresolved questions. Identify every piece of information that is essential for future actions. - -After your reasoning is complete, generate the final XML object. Be incredibly dense with information. Omit any irrelevant conversational filler. - -The structure MUST be as follows: - - - - - - - - - - - - - - - - - - - - - - - - - - - -``` - -**压缩执行流程**: - -1. 将 `historyToCompress` 作为用户消息发送给 LLM -2. 添加用户提示: `"First, reason in your scratchpad. Then, generate the ."` -3. 使用 `getCompressionPrompt()` 作为系统指令 -4. 将生成的摘要作为新的用户消息，添加模型确认回复，然后拼接 `historyToKeep` - -**失败保护**: -- 若压缩后 token 数反而增加，放弃压缩，返回原始历史 -- 设置 `hasFailedCompressionAttempt = true` 防止重复尝试 - -**执行时机**: `agents/executor.ts` → `executeTurn()` → `tryCompressChat()` - -### 3. 历史记录过滤 - -**位置**: `core/geminiChat.ts` → `extractCuratedHistory()` - -**机制**: -- 遍历历史记录，仅保留有效的用户和模型对话轮次 -- 过滤无效内容: - - 安全过滤导致的空输出 - - 空 `parts` 或无效 `content` - - 仅包含 function response 的消息 -- `getHistory(curated: boolean)` 支持返回"综合历史"或"精炼历史" - -### 4. 工具输出截断 - -**位置**: `core/coreToolScheduler.ts` → `truncateAndSaveToFile()` - -**配置**: -- `DEFAULT_TRUNCATE_TOOL_OUTPUT_THRESHOLD = 4,000,000` 字符 -- `DEFAULT_TRUNCATE_TOOL_OUTPUT_LINES = 1000` 行 - -**截断策略**: -- 保留前 `truncateLines / 5` 行和后 `truncateLines - head` 行 -- 中间插入 `... [CONTENT TRUNCATED] ...` 标记 -- 完整内容写入临时文件，提示使用 `READ_FILE_TOOL_NAME` 读取 - -**工具输出摘要**: `tools/shell.ts` → `summarizeToolOutput()` -- 当输出超过 `maxOutputTokens` 时，调用 LLM 生成摘要 -- 特别保留错误堆栈 (``) 和警告 (``) - -### 5. 文件内容管理 - -**位置**: `utils/fileUtils.ts` → `processSingleFileContent()` - -**限制**: -- `DEFAULT_MAX_LINES_TEXT_FILE = 2000` 行 -- `MAX_LINE_LENGTH_TEXT_FILE = 2000` 字符/行 -- 超长行截断并添加 `... [truncated]` 标记 - -**IDE 上下文限制**: `ide/ideContext.ts` -- `IDE_MAX_SELECTED_TEXT_LENGTH = 16384` 字符 -- `IDE_MAX_OPEN_FILES`: 限制打开文件数量（按时间戳保留最新） - -### 6. 记忆持久化 - -**位置**: `utils/memoryDiscovery.ts` → `loadServerHierarchicalMemory()` - -**机制**: -- 从工作目录、信任根目录、扩展文件等发现 `GEMINI.md` 文件 -- 通过 `MemoryTool` 将关键事实写入 `GEMINI.md` -- 支持层级记忆: 向上 (`findUpwardGeminiFiles`) 和向下 (`bfsFileSearch`) 搜索 -- 支持 JIT 加载 (`loadJitSubdirectoryMemory`)，按需加载子目录上下文 - -**导入处理**: `utils/memoryImportProcessor.ts` → `processImports()` -- `flat` 模式: 拼接所有导入内容 -- `tree` 模式: 构建嵌套结构 -- 安全控制: 防止路径遍历、限制递归深度 - -### 7. 事件机制 - -**关键事件**: -- `PreCompress`: 压缩前触发，支持手动/自动两种触发方式 -- `ContextWindowWillOverflow`: 上下文窗口即将溢出 -- `ChatCompressed`: 聊天已压缩，包含压缩前后 token 数 - -**Hook 支持**: `hooks/hookRunner.ts` -- `BeforeAgent` 事件允许注入 `additionalContext` -- 可在压缩前生成关键信息摘要并注入 - -### 8. 路由策略优化 - -**位置**: `routing/strategies/classifierStrategy.ts` - -**机制**: -- 滑动窗口: `HISTORY_SEARCH_WINDOW = 20` 条消息 -- 过滤非核心消息: 排除 function call/response -- 最终上下文: 仅取最后 `HISTORY_TURNS_FOR_CONTEXT = 4` 条 - -### 9. 其他优化机制 - -**循环检测**: `services/loopDetectionService.ts` -- 滚动窗口: `MAX_HISTORY_LENGTH = 5000` 字符 -- 滑动窗口分析: `CONTENT_CHUNK_SIZE = 50` -- 动态截断防止内存无限增长 - -**LRU 缓存**: `utils/LruCache.ts` -- 自动淘汰最久未使用的条目 -- 用于管理历史操作结果缓存 - -**活动监控**: `telemetry/activity-monitor.ts` -- `maxEventBuffer = 100` 事件 -- FIFO 缓冲管理 - -## 关键数据结构 - -```typescript -// 压缩状态枚举 -enum CompressionStatus { - /** 压缩成功 */ - COMPRESSED = 1, - - /** 压缩失败：压缩后 token 数反而增加 */ - COMPRESSION_FAILED_INFLATED_TOKEN_COUNT, - - /** 压缩失败：token 计数错误 */ - COMPRESSION_FAILED_TOKEN_COUNT_ERROR, - - /** 无需压缩，未执行任何操作 */ - NOOP, -} - -// 压缩信息 -interface ChatCompressionInfo { - compressionStatus: CompressionStatus; - originalTokenCount: number; - newTokenCount: number; -} - -// 压缩事件 -interface ChatCompressionEvent { - tokens_before: number; - tokens_after: number; -} - -// 历史记录 -interface Content { - role: 'user' | 'model'; - parts: Part[]; -} - -// 压缩后的状态快照结构（XML 格式） -interface StateSnapshot { - overall_goal: string; // 用户高层目标（单句描述） - key_knowledge: string; // 关键事实、约定和约束（要点列表） - file_system_state: string; // 文件系统状态（创建/读取/修改/删除的文件） - recent_actions: string; // 最近的重要操作摘要 - current_plan: string; // 当前执行计划（步骤列表，标记完成状态） -} -``` - -**压缩后的状态快照示例**: - -```xml - - - - Refactor the authentication service to use a new JWT library. - - - - - Build Command: `npm run build` - - Testing: Tests are run with `npm test`. Test files must end in `.test.ts`. - - API Endpoint: The primary API endpoint is `https://api.example.com/v2`. - - - - - CWD: `/home/user/project/src` - - READ: `package.json` - Confirmed 'axios' is a dependency. - - MODIFIED: `services/auth.ts` - Replaced 'jsonwebtoken' with 'jose'. - - CREATED: `tests/new-feature.test.ts` - Initial test structure for the new feature. - - - - - Ran `grep 'old_function'` which returned 3 results in 2 files. - - Ran `npm run test`, which failed due to a snapshot mismatch in `UserProfile.test.ts`. - - Ran `ls -F static/` and discovered image assets are stored as `.webp`. - - - - 1. [DONE] Identify all files using the deprecated 'UserAPI'. - 2. [IN PROGRESS] Refactor `src/components/UserProfile.tsx` to use the new 'ProfileAPI'. - 3. [TODO] Refactor the remaining files. - 4. [TODO] Update tests to reflect the API change. - - -``` - -## 配置参数总结 - -| 参数 | 默认值 | 说明 | -|------|--------|------| -| `COMPRESSION_TOKEN_THRESHOLD` | 0.2 | 压缩触发阈值（模型限制的20%） | -| `COMPRESSION_PRESERVE_THRESHOLD` | 0.3 | 保留最新消息比例（30%） | -| `TRUNCATE_TOOL_OUTPUT_THRESHOLD` | 4,000,000 | 工具输出截断阈值（字符） | -| `TRUNCATE_TOOL_OUTPUT_LINES` | 1000 | 工具输出最大行数 | -| `MAX_LINES_TEXT_FILE` | 2000 | 文本文件最大行数 | -| `MAX_LINE_LENGTH_TEXT_FILE` | 2000 | 单行最大长度 | -| `IDE_MAX_SELECTED_TEXT_LENGTH` | 16384 | IDE 选中文本最大长度 | -| `HISTORY_SEARCH_WINDOW` | 20 | 路由策略历史窗口大小 | -| `HISTORY_TURNS_FOR_CONTEXT` | 4 | 路由策略最终上下文条数 | - -## 执行流程总结 - -``` -每轮执行 (executeTurn) - ↓ -检测 token 限制 (sendMessageStream) - ↓ -[超限?] → 触发 ContextWindowWillOverflow 事件 - ↓ -尝试压缩 (tryCompressChat) - ↓ -[超过阈值?] → 调用 ChatCompressionService.compress() - ├─ 计算压缩起点 (保留最新30%) - ├─ 生成摘要 (LLM) - ├─ 组合新历史 (摘要 + 保留消息) - └─ [压缩后token增加?] → 放弃压缩 - ↓ -过滤历史 (extractCuratedHistory) - ├─ 移除无效消息 - └─ 返回精炼历史 - ↓ -处理工具输出 - ├─ [超长?] → truncateAndSaveToFile() - └─ [启用摘要?] → summarizeToolOutput() - ↓ -发送请求 -``` - -## 设计原则 - -1. **渐进式压缩**: 先尝试压缩，失败则回退 -2. **保留关键信息**: 优先保留最新消息和关键代码变更 -3. **多级防护**: Token 检测 → 压缩 → 截断 → 过滤 -4. **可观测性**: 通过事件和遥测记录压缩效果 -5. **持久化**: 关键信息写入 GEMINI.md，避免重复传递 - diff --git a/docs/zh/reading/20251114-halumem.md b/docs/zh/reading/20251114-halumem.md deleted file mode 100644 index 89126bb..0000000 --- a/docs/zh/reading/20251114-halumem.md +++ /dev/null @@ -1,268 +0,0 @@ -HaluMem: Evaluating Hallucinations in Memory Systems of Agents 解读 -https://arxiv.org/pdf/2511.03506v2 - -## 一、创新点 - -### ✅ 1. **首个面向“长期记忆”的细粒度综合评测基准** - -- **填补空白**：现有主流基准（如 MMLU、TruthfulQA、BBH）主要评估静态知识或短上下文推理，缺乏对**跨多轮、长时间跨度对话中记忆能力 - **的系统评测。 -- **LongMemBench 是首个专门设计用于评估 LLM 长期记忆能力的高质量基准**，包含 4,987 个问题，覆盖 1,000+ 虚构用户档案和多轮对话历史。 - -> 📌 创新性：从“一次性问答”转向“持续交互中的记忆追踪”。 - ---- - -### ✅ 2. **结构化记忆点（Structured Memory Points）作为评估锚点** - -- 引入 **显式、可追溯的记忆单元（Memory Points）**，每个都是用户在对话中透露的原子事实。 -- 每个问题都明确标注其依赖的**证据记忆点索引**（e.g., `[1,3]`），实现： - - 答案可验证 - - 错误可归因（是遗忘？混淆？幻觉？） - - 支持细粒度能力分解 - -> 📌 创新性：将模糊的“记住与否”转化为**可量化、可解释的记忆依赖关系**。 - ---- - -### ✅ 3. **五维记忆能力分类体系** - -提出一个**全面的记忆任务分类框架**，超越简单回忆： - -| 维度 | 能力要求 | 现有基准是否支持 | -|----------------------------------|----------------------|----------| -| **Basic Fact Recall** | 单点信息提取 | 部分支持 | -| **Multi-hop Inference** | 多记忆点融合推理 | 极少 | -| **Dynamic Update** | 处理信息更新与状态变更 | 几乎无 | -| **Memory Boundary** | 识别“未知” vs “未提及”，抑制幻觉 | 缺乏 | -| **Generalization & Application** | 基于记忆进行个性化建议或推断 | 无 | - -> 📌 创新性：首次系统定义并实现对**动态、边界敏感、可泛化**记忆能力的评测。 - ---- - -### ✅ 4. **真实感强的多轮对话模拟** - -- 对话由人工精心编写，模拟真实用户行为： - - 信息逐步披露（非一次性给出） - - 包含时间戳（体现“长期”） - - 含干扰信息、情感表达、话题跳跃 -- 支持评估模型在**噪声环境下的记忆选择与聚焦能力** - -> 📌 创新性：避免机械式问答，贴近实际应用场景（如个人助理、陪伴机器人）。 - ---- - -### ✅ 5. **开源与可扩展的数据集设计** - -- 数据以结构化 JSON 格式发布，包含： - - 用户档案 - - 多轮对话历史 - - 记忆点列表 - - 问题-答案对 + 元标签（类型、难度、证据等） -- 支持社区在此基础上： - - 开发记忆增强模型 - - 构建自动记忆提取模块 - - 扩展新领域（医疗、教育等） - -> 📌 创新性：不仅是一个 benchmark，更是一个**记忆研究的基础设施**。 - - - ---- - -## 二、LongMemBench 的关键设计 - -### 1. **记忆点（Memory Points）** - -- 每个“记忆点”代表用户在对话历史中透露的一个事实。 -- 示例（来自论文 Listing 7）： - ```json - { - "index": 1, - "memory_content": "Martin Mark is considering a career change due to the impact of his current role on his mental health." - } - ``` -- 这些记忆点被结构化存储，作为回答后续问题的依据。 - -### 2. **对话轮次（Dialogue Rounds）** - -- 每轮包含：用户话语 + AI 回应 + 时间戳。 -- 支持模拟真实场景中信息随时间逐步披露的过程。 - -### 3. **问题设计（Questions）** - -每个问题包含以下元信息： - -- **问题文本**（Question） -- **参考答案**（Reference Answer） -- **所需记忆点索引**（Evidence Memory Points） -- **问题类型**（Question Type） -- **难度等级**（Difficulty） - ---- - -## 三、五类记忆能力评估维度（附例子说明） - -论文定义了五种核心记忆任务类型，用以全面评估模型： - -| 类型 | 中文名 | 说明 | 例子 | -|----------------------------------|--------|------------------------|---------------------------------------------------------------| -| **Basic Fact Recall** | 基础事实回忆 | 直接提取单个记忆点中的信息 | Q: “用户为什么考虑换工作？”
A: “因为当前工作影响了他的心理健康。”（对应记忆点1） | -| **Multi-hop Inference** | 多跳推理 | 需结合多个记忆点进行逻辑推理 | Q: “用户最近生活有哪些变化？”
A: “他换了工作，刚从疾病中康复，还结识了新朋友。”（需整合记忆点1、2、3） | -| **Dynamic Update** | 动态更新 | 用户信息发生变更，模型需记住最新状态 | 初期：“用户在A公司工作” → 后期：“用户已跳槽到B公司”
Q: “用户现在在哪工作？” → 应答“B公司”而非A | -| **Memory Boundary** | 记忆边界 | 区分哪些信息属于记忆范围，哪些是干扰或未提及 | Q: “用户是否喜欢爬山？” → 若从未提及，应回答“不知道”而非编造 | -| **Generalization & Application** | 泛化与应用 | 基于记忆进行合理推断或建议 | Q: “考虑到用户刚康复且换了工作，有什么生活建议？” → 需结合健康+职业变化给出个性化建议 | - ---- - -# 完整的构造示例 - -## 🎯 构造目标 - -构建一个能测试模型“长期记忆”能力的对话片段 + 问题对，要求： - -- 对话中逐步透露用户信息（记忆点） -- 后续问题需依赖这些历史信息作答 -- 覆盖不同类型的记忆任务（如回忆、推理、更新等） - ---- - -## ✅ 第一步：定义“记忆点”（Memory Points） - -作者首先人工编写一系列关于虚构用户 **“Martin Mark”** 的事实性陈述，每个作为一个独立记忆点： - -```json -[ - { - "index": 1, - "memory_content": "Martin is considering a career change because his current job is affecting his mental health." - }, - { - "index": 2, - "memory_content": "Martin recently recovered from a mild case of pneumonia." - }, - { - "index": 3, - "memory_content": "Martin met a new friend named Lena at a book club last week." - } -] -``` - -> 这些是“黄金标准”的真实信息，后续所有问题的答案都必须基于它们。 - ---- - -## ✅ 第二步：生成多轮对话（模拟真实交互） - -将上述记忆点**自然地嵌入到多轮对话中**，每轮包含时间戳（体现“长期”）： - -```text -[Round 1 | Day 1] -User: I’ve been feeling really burnt out at work lately. It’s taking a toll on my mental health. -AI: That sounds tough. Have you thought about talking to HR or exploring other options? - -[Round 2 | Day 5] -User: On the bright side, I finally recovered from that pneumonia I had last month! -AI: Glad to hear you’re feeling better! Rest is so important. - -[Round 3 | Day 10] -User: I went to a book club meeting yesterday and met someone really interesting—her name is Lena. -AI: That’s great! Making new connections can really lift your spirits. -``` - -> 注意：用户没有直接说“我的记忆点是……”，而是用自然语言透露信息。模型需要从中提取并记住关键事实。 - ---- - -## ✅ 第三步：设计评测问题（含元数据） - -针对上述对话，设计具体问题，并标注其依赖的记忆点、类型、难度等： - -### 示例 1：基础回忆（Basic Fact Recall） - -```json -{ - "question": "Why is Martin considering changing his job?", - "reference_answer": "Because his current job is negatively impacting his mental health.", - "evidence_memory_points": [ - 1 - ], - "question_type": "Basic Fact Recall", - "difficulty": "Easy" -} -``` - -### 示例 2：多跳推理（Multi-hop Inference） - -```json -{ - "question": "What recent positive changes has Martin experienced despite his work stress?", - "reference_answer": "He recovered from pneumonia and made a new friend named Lena at a book club.", - "evidence_memory_points": [ - 2, - 3 - ], - "question_type": "Multi-hop Inference", - "difficulty": "Medium" -} -``` - -### 示例 3：动态更新（假设后续有新对话） - -> 假设在 Day 15，用户说：“Actually, I decided to stay at my job after all—they offered me a flexible schedule.” - -则新增记忆点： - -```json -{ - "index": 4, - "memory_content": "Martin decided to stay in his current job after being offered a flexible work schedule." -} -``` - -对应问题： - -```json -{ - "question": "Is Martin still planning to leave his job?", - "reference_answer": "No, he decided to stay after being offered a flexible schedule.", - "evidence_memory_points": [ - 1, - 4 - ], - "question_type": "Dynamic Update", - "difficulty": "Hard" -} -``` - -> 模型必须忽略早期“考虑换工作”的信息，采用最新状态。 - ---- - -## ✅ 第四步：确保质量与多样性 - -- 由多名标注者交叉验证记忆点与问题的一致性 -- 覆盖不同领域：健康、职业、人际关系、兴趣爱好等 -- 控制对话长度（平均 8–12 轮）、时间跨度（数天至数周） -- 包含干扰信息（如用户提到无关话题，测试模型是否混淆） - ---- - -## 🔍 总结：构造流程图 - -``` -人工编写记忆点 - ↓ -将记忆点融入多轮自然对话（带时间戳） - ↓ -基于记忆点设计结构化问题（标注类型/证据/难度） - ↓ -人工审核 → 形成 LongMemBench 数据集 -``` - -这种构造方式确保了： - -- **可控性**：知道正确答案依赖哪些事实 -- **真实性**：对话像真人聊天，非机械问答 -- **挑战性**：涵盖记忆的提取、整合、更新、边界判断等高阶能力 diff --git a/docs/zh/reading/20251124-mem-pal.md b/docs/zh/reading/20251124-mem-pal.md deleted file mode 100644 index 43ae055..0000000 --- a/docs/zh/reading/20251124-mem-pal.md +++ /dev/null @@ -1,279 +0,0 @@ -# 论文解读 Mem-PAL: Memory-Augmented Personalized Assistant with Log-based Structured Memory ---- - -## 一、研究动机 - -传统个性化对话系统通常依赖显式用户画像（如年龄、性别、兴趣标签）或短期上下文，难以捕捉用户长期、动态、细粒度的行为偏好。 -Mem-PAL 的目标是：**利用用户在数字环境中的真实行为日志（如搜索、浏览、发帖、消息等），自动构建结构化、可检索的长期记忆，并用于提升对话系统的个性化能力 -**。 - ---- - -## 二、核心思想：两层记忆结构 - -论文提出 **两层（two-tier）记忆架构**： - -### 1. **底层：具体记忆条目（Specific Memory Entries）** - -- 来源：用户行为日志（log data） -- 类型包括： - - **Web Search**：用户搜索了什么关键词，看了什么内容 - - **Content Publishing**：用户在某平台发布了文章/视频等 - - **Content Browsing**：用户浏览了某平台的某类内容 - - **Message Sending/Receiving**：用户发送或接收的消息内容 -- 每条日志被转换为**标准化的自然语言描述**（如：“The user searched for ‘best hiking trails in Colorado’ and viewed an - article summarizing top 10 scenic routes.”） - -### 2. **上层：抽象用户画像（Abstracted User Profile）** - -- 基于底层记忆条目，通过 LLM 自动**总结用户的长期兴趣、习惯、性格特征等** -- 例如：“The user is an outdoor enthusiast who frequently researches travel destinations and shares travel tips on social - media.” - -这种分层设计兼顾**细节保留**与**高层语义抽象**，支持灵活检索与推理。 - ---- - -## 三、技术实现 - -### 3.1 实现 Pipeline - -Mem-PAL 的实现分为三个阶段： - -1. **记忆构建（Memory Construction）** - - 从原始用户行为日志中提取结构化事件 - - 将每条事件转换为自然语言描述（Specific Memory Entry） - - 聚合这些条目，生成抽象用户画像（Abstracted Profile） - -2. **记忆检索（Memory Retrieval）** - - 用户发起新对话时，系统对当前 query 编码 - - 在记忆库中检索最相关的若干条具体记忆 - -3. **记忆增强生成（Memory-Augmented Generation）** - - 将检索到的记忆 + 抽象画像 + 对话上下文 输入给 LLM - - LLM 生成个性化、上下文连贯的回复 - ---- - -### 3.2 关键技术细节 - -#### 1. 日志 → 结构化记忆条目（Specific Memory Entry） - -假设原始日志如下（来自某旅游 App）： - -```json -{ - "timestamp": "2025-11-10T14:30:00Z", - "event_type": "web_search", - "query": "best hiking trails in Colorado", - "clicked_url": "https://example.com/colorado-hiking-top10" -} -``` - -Mem-PAL 使用预定义模板 + LLM（如 Qwen-Max）将其转化为自然语言记忆条目： - -> **Specific Memory Entry**: -> “On November 10, 2025, the user searched for ‘best hiking trails in Colorado’ and viewed an article titled ‘Top 10 -> Scenic Hiking Trails in Colorado’.” - -类似地，其他事件类型也有对应模板： - -| 事件类型 | 示例日志 | 生成的记忆条目 | -|------------------|-----------------------|-------------------------------------------------------------------------------------------------------------------------------| -| Content Browsing | 浏览了“露营装备推荐”页面 | “The user browsed a webpage about recommended camping gear on November 12.” | -| Message Sent | 发送消息：“下周去落基山徒步，求推荐装备” | “On November 13, the user sent a message asking for gear recommendations for a hiking trip to the Rocky Mountains next week.” | - -#### 2. 构建抽象用户画像（Abstracted User Profile） - -系统定期（或在对话开始前）调用 LLM，输入所有 Specific Memory Entries，要求总结： - -> **Prompt to LLM**: -> “Based on the following user activity logs, summarize the user’s long-term interests, habits, and personality in 2–3 -> sentences.” - -LLM 输出： - -> **Abstracted Profile**: -> “The user is an outdoor enthusiast with a strong interest in hiking and camping, particularly in mountainous regions -> like Colorado. They actively research travel destinations and seek practical gear advice before trips.” - -#### 3. 检索与生成 - -当用户发起新 query，例如： - -> **User Query**: “What should I pack for a 3-day hike in the Rockies?” - -系统执行以下步骤： - -##### Step 1: 编码与检索 - -- 使用句子嵌入模型（`paraphrase-multilingual-mpnet-base-v2`）对 query 和所有记忆条目编码 -- 计算余弦相似度，取 top-3 最相关记忆，例如： - 1. “On November 10, the user searched for ‘best hiking trails in Colorado’...” - 2. “The user browsed a webpage about recommended camping gear...” - 3. “On November 13, the user sent a message asking for gear recommendations...” - -##### Step 2: 构造 LLM 输入提示（Prompt） - -```text -You are a helpful personal assistant. Use the following information to respond to the user. - -[Abstracted Profile] -The user is an outdoor enthusiast with a strong interest in hiking and camping, particularly in mountainous regions like Colorado... - -[Relevant Memories] -- On Nov 10: searched for hiking trails in Colorado... -- On Nov 12: browsed camping gear recommendations... -- On Nov 13: asked for gear advice for a Rocky Mountains trip... - -[Current Conversation] -User: What should I pack for a 3-day hike in the Rockies? - -Assistant: -``` - -##### Step 3: LLM 生成个性化回复（由 Qwen-Max-0428 生成） - -> **Assistant Reply**: -> “Given your upcoming 3-day hike in the Rockies and your interest in lightweight gear, I’d recommend packing: -> - A 40L backpack with rain cover -> - Layered clothing (base + insulation + shell) -> - Lightweight tent or bivy sack -> - Water filter and at least 2L capacity -> - Trail snacks like nuts and energy bars - > Since you’ve looked into Colorado trails before, remember that afternoon thunderstorms are common—don’t forget a - waterproof jacket!” - -这个回复不仅通用，还**引用了用户历史行为**（“you’ve looked into Colorado trails”），体现了个性化。 - ---- - -### 关键创新点总结（结合实现） - -| 创新点 | 实现体现 | -|--------------|----------------------------| -| **分层记忆结构** | 具体日志条目 + 抽象画像，兼顾细节与概括 | -| **基于真实日志** | 不依赖人工标注，直接从用户行为自动构建 | -| **RAG 增强生成** | 检索相关记忆注入 LLM 上下文，提升相关性与个性化 | -| **服务场景适配** | 针对多轮、任务导向对话（如旅行、健康、客服）优化 | - ---- - -## 四、实验与评估 - -好的，以下是《Mem-PAL》论文中**实验与评估部分的详细解读**，按照你的要求从**三级标题**开始组织内容，涵盖评估设置、指标设计、基线对比、结果分析等关键方面。 - ---- - -### 4.1 评估任务设计 - -论文聚焦于**多轮个性化对话交互任务（Multi-turn Dialogue Interaction Task）**，旨在模拟真实用户与个性化助手之间的连续对话场景。该任务具有以下特点： - -- **动态上下文**：每轮对话可能涉及不同子话题（如从“徒步路线”转向“装备推荐”）。 -- **长期记忆依赖**：助手需利用用户过去数天甚至数周的行为日志进行响应。 -- **个性化要求高**：回复不仅要正确，还需体现对用户兴趣、习惯的理解。 - -为构建评估数据集，作者采用以下方法： - -- 使用 **User-LLM（Qwen2.5-Max）** 模拟真实用户行为，基于预设的用户画像生成多轮对话。 -- 每个用户拥有一个由 50–100 条结构化日志构成的记忆库。 -- 对话长度通常为 4–8 轮，覆盖多个相关主题。 - ---- - -### 4.2 评估方法与指标 - -#### 4.2.1 自动评估（Automatic Evaluation） - -使用 **GPT-4-turbo** 作为 **Evaluation-LLM**，对两个系统（Mem-PAL vs. Baseline）在同一对话上下文下的回复进行成对比较（Pairwise -Comparison）。评估维度包括： - -- **Relevance（相关性）**：回复是否贴合当前对话意图。 -- **Personalization（个性化）**：是否有效利用用户历史信息。 -- **Coherence（连贯性）**：语言是否流畅、逻辑是否合理。 -- **Helpfulness（有用性）**：是否提供实质性帮助或建议。 - -每个维度按 Likert 5 分制打分，最终汇总为综合胜率（Win/Tie/Lose）。 - -#### 4.2.2 人工评估（Human Evaluation） - -邀请领域专家对随机抽取的 200 组对话进行盲评，同样采用 Win/Tie/Lose 判断，并计算 **Fleiss’ Kappa** 以衡量评分者间一致性（论文引用 -McCrae & John, 1992；McHugh, 2012 的方法论支持可靠性分析）。 - ---- - -### 4.3 基线方法对比 - -论文对比了多种主流或代表性基线，分为以下几类： - -| 类别 | 方法 | 描述 | -|----------------|--------------------|-----------------------| -| **无记忆基线** | Vanilla (w/o log) | 仅使用当前对话上下文，无任何用户记忆 | -| **带日志但无结构化记忆** | Vanilla (with log) | 将原始日志拼接进上下文（未结构化、未摘要） | -| **记忆摘要方法** | RecurSum | 使用递归摘要压缩用户历史 | -| **条件记忆模型** | ConditionMem | 基于特定条件（如时间、主题）选择记忆 | -| **外部记忆库** | MemoryBank | 使用向量数据库存储并检索记忆 | - -Mem-PAL 在所有这些方法之上引入了**结构化日志解析 + 两层记忆架构 + 精准检索机制**。 - ---- - -### 4.4 主要实验结果 - -#### 4.4.1 多轮对话胜率统计 - -论文在 Table 3 中报告了 Mem-PAL（记为 “Ours”）与其他方法的 **Win/Tie/Lose** 统计结果（总计约 800+ 对话样本）： - -| 对比方法 | Win / Tie / Lose（Preference） | Win / Tie / Lose（Requirement） | -|--------------------|------------------------------|-------------------------------| -| Vanilla (w/o log) | 478 / 29 / 319 | 480 / 18 / 328 | -| Vanilla (with log) | 447 / 33 / 346 | 452 / 22 / 352 | -| RecurSum | 421 / 29 / 376 | 439 / 18 / 369 | -| ConditionMem | 396 / 42 / 388 | 413 / 19 / 394 | -| MemoryBank | 449 / 33 / 344 | 452 / 25 / 349 | - -> 注：“Preference” 指用户偏好评估，“Requirement” 指任务需求满足度评估。 - -**关键观察**： - -- Mem-PAL 在所有对比中均取得**最高胜率**（Win > 400），显著优于无记忆基线（Win 提升约 50%）。 -- 即使与同样使用记忆的 MemoryBank 相比，Mem-PAL 仍保持优势，说明**结构化日志处理和分层记忆设计有效提升了记忆利用率**。 -- Tie 比例较低（通常 < 5%），表明差异具有判别性。 - -#### 4.4.2 消融实验（Ablation Study） - -为验证各组件贡献，论文进行了消融实验： - -| 配置 | 描述 | 性能下降幅度 | -|----------------------------|-----------------|-----------------------| -| Full Mem-PAL | 完整系统 | 基准 | -| - Abstracted Profile | 移除抽象画像，仅用具体记忆 | ↓ 12% Win rate | -| - Structured Log Parsing | 直接输入原始日志（无模板转换） | ↓ 18% Win rate | -| - Retrieval (use all logs) | 不检索，全量注入上下文 | ↑ 计算开销，↓ 个性化精度（因噪声干扰） | - -结果表明：**结构化日志解析** 和 **两层记忆架构** 是性能提升的关键。 - ---- - -### 4.5 评估局限性与未来方向 - -尽管评估较为全面，论文也承认若干局限： - -- **User-LLM 模拟 vs. 真实用户**：当前评估依赖 LLM 模拟用户，可能存在行为偏差。 -- **日志覆盖范围有限**：实验主要基于文本类日志（搜索、消息），尚未整合多模态行为（如图片浏览、语音指令）。 -- **长期遗忘机制缺失**：当前系统未实现记忆衰减或更新策略。 - -未来工作将探索： - -- 引入**记忆生命周期管理** -- 支持**跨平台日志融合** -- 开展**真实用户 A/B 测试** - ---- - -如需我进一步可视化胜率对比图表，或解释 GPT-4 评估提示的具体写法，也可以继续提出！ - - - - - diff --git a/example.env b/example.env index 2aed2e8..4c093ec 100644 --- a/example.env +++ b/example.env @@ -1,4 +1,4 @@ -FLOW_LLM_API_KEY=sk-xxxx -FLOW_LLM_BASE_URL=https://xxxx/v1 -FLOW_EMBEDDING_API_KEY=sk-xxxx -FLOW_EMBEDDING_BASE_URL=https://xxxx/v1 +EMBEDDING_API_KEY=sk-xxx +EMBEDDING_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 +LLM_API_KEY=sk-xxx +LLM_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1 diff --git a/flowllm/__init__.py b/flowllm/__init__.py index b496091..057e567 100644 --- a/flowllm/__init__.py +++ b/flowllm/__init__.py @@ -1,14 +1,24 @@ -"""FlowLLM: A library for managing LLM flows and contexts.""" +"""FlowLLM package.""" -import os +__version__ = "0.3.0.0" -os.environ.setdefault("FLOW_APP_NAME", "FlowLLM") +from . import config +from . import constants +from . import enumeration +from . import schema +from . import steps +from . import utils +from .application import Application +from .components import BaseComponent -from .core.utils import load_env # noqa: E402 # pylint: disable=wrong-import-position - -load_env() - -from . import gallery # noqa: E402, F401 # pylint: disable=wrong-import-position,unused-import -# from . import extensions # noqa: E402, F401 # pylint: disable=wrong-import-position,unused-import - -__version__ = "0.2.0.10" +__all__ = [ + "Application", + "BaseComponent", + # submodules + "config", + "constants", + "enumeration", + "schema", + "steps", + "utils", +] diff --git a/flowllm/application.py b/flowllm/application.py new file mode 100644 index 0000000..1d3fea1 --- /dev/null +++ b/flowllm/application.py @@ -0,0 +1,268 @@ +"""Main application entry point.""" + +import asyncio +import heapq +import sys +from pathlib import Path +from typing import AsyncGenerator, TypeVar + +from . import __version__ +from .components import BaseComponent, ApplicationContext, R +from .components.job import BackgroundJob, BaseJob, CronJob, StreamJob +from .components.service import BaseService +from .config import parse_args, resolve_app_config +from .enumeration import ComponentEnum +from .schema import ComponentConfig, Response, StreamChunk +from .utils import cli_find_flowllm, execute_stream_task, load_env, precheck_start, print_logo, get_logger + +T = TypeVar("T", bound=BaseComponent) +_NodeKey = tuple[ComponentEnum, str] + + +class Application(BaseComponent): + """Wires components from config and runs jobs against them.""" + + def __init__(self, **kwargs) -> None: + self.context = ApplicationContext(**kwargs) + self._started_components: list[BaseComponent] = [] + + self._setup_workspace_directories() + + if self.config.enable_logo: + print_logo(self.config) + logger = get_logger(log_to_console=self.config.log_to_console, log_to_file=self.config.log_to_file) + logger.info(f"Initializing {self.config.app_name} Application v{__version__}") + super().__init__() + + self._init_service() + self._init_components() + self._init_jobs() + + @property + def config(self): + """Application config from context.""" + return self.context.app_config + + # Wiring + + def _setup_workspace_directories(self) -> None: + """Create workspace directories if missing.""" + cfg = self.config + workspace_path = Path(cfg.workspace_dir).absolute() + workspace_path.mkdir(parents=True, exist_ok=True) + for subdir in [cfg.metadata_dir, cfg.session_dir]: + if subdir: + (workspace_path / subdir).mkdir(parents=True, exist_ok=True) + + def _init_service(self) -> None: + """Instantiate the service backend.""" + self.context.service = self._instantiate( + ComponentEnum.SERVICE, + self.config.service, + label="Service", + expected_type=BaseService, + ) + + def _init_components(self) -> None: + """Instantiate all configured components.""" + for ctype, group in self.config.components.items(): + self.context.components[ctype] = {} + for name, cfg in group.items(): + self.context.components[ctype][name] = self._instantiate( + ctype, + cfg, + label=f"Component '{name}'", + expected_type=BaseComponent, + name=name, + ) + + def _init_jobs(self) -> None: + """Instantiate all configured jobs.""" + for name, cfg in self.config.jobs.items(): + self.context.jobs[name] = self._instantiate( + ComponentEnum.JOB, + cfg, + label=f"Job '{name}'", + expected_type=BaseJob, + name=name, + ) + + def _instantiate( + self, + ctype: ComponentEnum, + cfg: ComponentConfig, + *, + label: str, + expected_type: type[T], + name: str | None = None, + ) -> T: + """Resolve backend from registry and construct an instance.""" + if not cfg.backend: + raise ValueError(f"{label} is missing the required 'backend' field") + backend_cls = R.get(ctype, cfg.backend) + if backend_cls is None: + raise ValueError(f"Unregistered backend '{cfg.backend}' for {label}") + + params = cfg.model_dump() + params["app_context"] = self.context + if name is not None: + params.setdefault("name", name) + instance = backend_cls(**params) + if not isinstance(instance, expected_type): + got, want = type(instance).__name__, expected_type.__name__ + raise TypeError(f"{label} backend '{cfg.backend}' produced {got}, expected {want} subclass") + return instance + + # Dependency ordering + + def _topological_order(self) -> list[BaseComponent]: + """Return components in dependency order (Kahn's algorithm).""" + nodes: dict[_NodeKey, BaseComponent] = { + (ctype, name): comp for ctype, group in self.context.components.items() for name, comp in group.items() + } + in_degree, dependents = self._build_dependency_graph(nodes) + + ready = [k for k, d in in_degree.items() if d == 0] + heapq.heapify(ready) + ordered: list[BaseComponent] = [] + while ready: + key = heapq.heappop(ready) + ordered.append(nodes[key]) + for downstream in dependents[key]: + in_degree[downstream] -= 1 + if in_degree[downstream] == 0: + heapq.heappush(ready, downstream) + + if len(ordered) != len(nodes): + unresolved = [f"{k[0].value}:{k[1]}" for k, d in in_degree.items() if d > 0] + raise ValueError(f"Circular dependency detected among: {unresolved}") + return ordered + + @staticmethod + def _build_dependency_graph( + nodes: dict[_NodeKey, BaseComponent], + ) -> tuple[dict[_NodeKey, int], dict[_NodeKey, list[_NodeKey]]]: + """Build in-degree and adjacency lists.""" + in_degree: dict[_NodeKey, int] = dict.fromkeys(nodes, 0) + dependents: dict[_NodeKey, list[_NodeKey]] = {k: [] for k in nodes} + for key, comp in nodes.items(): + for dep in comp.dependencies: + dep_key = (dep.ctype, dep.name) + if dep_key in nodes: + dependents[dep_key].append(key) + in_degree[key] += 1 + elif not dep.optional: + raise ValueError( + f"Component {key[0].value}:{key[1]} depends on unregistered {dep.ctype.value}:{dep.name}", + ) + return in_degree, dependents + + # Lifecycle + + async def _start(self) -> None: + """Start components, then jobs as base > stream > background > cron.""" + try: + components = self._topological_order() + jobs = list(self.context.jobs.values()) + base_jobs = [j for j in jobs if not isinstance(j, (StreamJob, BackgroundJob))] + stream_jobs = [j for j in jobs if isinstance(j, StreamJob)] + background_jobs = [j for j in jobs if isinstance(j, BackgroundJob) and not isinstance(j, CronJob)] + cron_jobs = [j for j in jobs if isinstance(j, CronJob)] + for c in components + base_jobs + stream_jobs + background_jobs + cron_jobs: + await self._start_one(c) + except Exception: + await self._close() + raise + + async def _start_one(self, c: BaseComponent) -> None: + """Start one component and track it for shutdown.""" + try: + if isinstance(c, BackgroundJob): + self.logger.info(f"Starting background job: {c.name}") + await c.start() + self._started_components.append(c) + except Exception as e: + self.logger.exception(f"Failed to start {c.component_type.value}:{c.name}: {e}") + raise + + async def _close(self) -> None: + """Close in reverse start order.""" + for c in reversed(self._started_components): + try: + await c.close() + except Exception as e: + self.logger.exception(f"Failed to close {c.component_type.value}:{c.name}: {e}") + self._started_components.clear() + + async def update_component(self, component_enum: ComponentEnum | str, name: str, /, **kwargs) -> BaseComponent: + """Update an existing component's attributes.""" + component_enum = ComponentEnum(component_enum) + group = self.context.components.get(component_enum) + if not group or name not in group: + raise KeyError(f"Component '{name}' not found in {component_enum.value}") + + component = group[name] + for key, value in kwargs.items(): + if not hasattr(component, key): + raise AttributeError(f"Component {component_enum.value}:{name} has no attribute '{key}'") + setattr(component, key, value) + return component + + # Job execution + + async def run_job(self, name: str, /, **kwargs) -> Response: + """Execute a job by name and return its Response.""" + if name not in self.context.jobs: + raise KeyError(f"Job '{name}' not found") + return await self.context.jobs[name](**kwargs) + + async def run_stream_job(self, name: str, /, **kwargs) -> AsyncGenerator[StreamChunk, None]: + """Execute a streaming job, yielding chunks.""" + if name not in self.context.jobs: + raise KeyError(f"Job '{name}' not found") + stream_queue: asyncio.Queue = asyncio.Queue() + task = asyncio.create_task(self.context.jobs[name](stream_queue=stream_queue, **kwargs)) + async for chunk in execute_stream_task( + stream_queue=stream_queue, + task=task, + task_name=name, + output_format="chunk", + ): + assert isinstance(chunk, StreamChunk) + yield chunk + + def run_app(self): + """Serve the application via the configured service.""" + assert isinstance(self.context.service, BaseService) + self.context.service.run_app(app=self) + + +_CLIENT_KWARGS = {"host", "port", "timeout", "transport", "command", "args"} + + +async def call_server(action: str, **kwargs): + """Call the appropriate server component.""" + backend: str = kwargs.pop("backend", "http") + client_kwargs = {key: kwargs.pop(key) for key in list(kwargs) if key in _CLIENT_KWARGS} + client_cls = R.get(ComponentEnum.CLIENT, backend) + if client_cls is None: + raise ValueError(f"Unknown client backend: {backend!r}") + async with client_cls(**client_kwargs) as client: + async for chunk in client(action=action, **kwargs): + print(chunk, end="", flush=True) + print() + + +def main(): + """CLI entry point.""" + action, kwargs = parse_args(*sys.argv[1:]) + if action == "start": + load_env() + kwargs = resolve_app_config(**kwargs) + if not precheck_start(kwargs.get("service")): + return + Application(**kwargs).run_app() + elif action == "find_flowllm": + cli_find_flowllm() + else: + asyncio.run(call_server(action, **kwargs)) diff --git a/flowllm/components/__init__.py b/flowllm/components/__init__.py new file mode 100644 index 0000000..b157853 --- /dev/null +++ b/flowllm/components/__init__.py @@ -0,0 +1,32 @@ +"""Components""" + +from . import agent_wrapper +from . import as_embedding +from . import as_llm +from . import client +from . import embedding_store +from . import job +from . import service +from .application_context import ApplicationContext +from .base_component import BaseComponent, ComponentMixin +from .component_registry import ComponentRegistry, R +from .prompt_handler import PromptHandler +from .runtime_context import RuntimeContext + +__all__ = [ + "ApplicationContext", + "BaseComponent", + "ComponentMixin", + "ComponentRegistry", + "R", + "PromptHandler", + "RuntimeContext", + # base components + "agent_wrapper", + "as_llm", + "client", + "as_embedding", + "embedding_store", + "job", + "service", +] diff --git a/flowllm/components/agent_wrapper/__init__.py b/flowllm/components/agent_wrapper/__init__.py new file mode 100644 index 0000000..427472d --- /dev/null +++ b/flowllm/components/agent_wrapper/__init__.py @@ -0,0 +1,7 @@ +"""Unified agent wrapper component with swappable backends.""" + +from .as_agent_wrapper import AsAgentWrapper +from .base_agent_wrapper import BaseAgentWrapper +from .cc_agent_wrapper import CcAgentWrapper + +__all__ = ["BaseAgentWrapper", "AsAgentWrapper", "CcAgentWrapper"] diff --git a/flowllm/components/agent_wrapper/as_agent_wrapper.py b/flowllm/components/agent_wrapper/as_agent_wrapper.py new file mode 100644 index 0000000..0358acd --- /dev/null +++ b/flowllm/components/agent_wrapper/as_agent_wrapper.py @@ -0,0 +1,323 @@ +"""AgentScope backend for the unified agent wrapper.""" + +import json +import re +import time +from collections.abc import AsyncGenerator +from pathlib import Path +from typing import Any, TYPE_CHECKING +from uuid import uuid4 + +from agentscope.agent import Agent, ContextConfig, ReActConfig +from agentscope.agent._config import ModelConfig +from agentscope.event import ( + DataBlockDeltaEvent, + DataBlockEndEvent, + DataBlockStartEvent, + ExceedMaxItersEvent, + ModelCallEndEvent, + ModelCallStartEvent, + ReplyEndEvent, + ReplyStartEvent, + ThinkingBlockDeltaEvent, + ThinkingBlockEndEvent, + ThinkingBlockStartEvent, + TextBlockDeltaEvent, + TextBlockEndEvent, + TextBlockStartEvent, + ToolCallDeltaEvent, + ToolCallEndEvent, + ToolCallStartEvent, + ToolResultDataDeltaEvent, + ToolResultEndEvent, + ToolResultStartEvent, + ToolResultTextDeltaEvent, +) +from agentscope.message import TextBlock, ToolResultState, UserMsg +from agentscope.permission import PermissionBehavior, PermissionContext, PermissionDecision, PermissionMode +from agentscope.state import AgentState +from agentscope.tool import ( + Bash, + Edit, + FunctionTool, + Glob, + Grep, + Read, + ToolBase, + ToolChunk, + Toolkit, + Write, +) + +from .base_agent_wrapper import BaseAgentWrapper +from ..as_llm import BaseAsLLM +from ..component_registry import R +from ...enumeration import ChunkEnum +from ...schema import StreamChunk +from ...utils import AsStateHandler +from ...utils.env_utils import load_env + +if TYPE_CHECKING: + from ..job.base_job import BaseJob + +_UUID_RE = re.compile(r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$", re.IGNORECASE) + + +class BypassAnalysisBash(Bash): + """Bash that skips static analysis, delegating permission to PermissionEngine.""" + + async def check_permissions(self, _tool_input: dict[str, Any], _context: PermissionContext) -> PermissionDecision: + """Bypass static analysis; always passthrough.""" + return PermissionDecision( + behavior=PermissionBehavior.PASSTHROUGH, + message="Bash static analysis skipped; delegating to permission engine.", + ) + + +@R.register("agentscope") +class AsAgentWrapper(BaseAgentWrapper): + """Agent wrapper backed by AgentScope.""" + + def __init__(self, as_llm: str = "default", session_retention_days: int = 10, **kwargs): + super().__init__(**kwargs) + self.as_llm = self.bind(as_llm, BaseAsLLM, optional=False) + self.session_retention_days = int(session_retention_days) + self._session_cleanup_done = False + + @staticmethod + def _make_tool(job: "BaseJob") -> FunctionTool: + async def run_job(**kwargs) -> ToolChunk: + response = await job(**kwargs) + state = ToolResultState.SUCCESS if response.success else ToolResultState.ERROR + return ToolChunk(content=[TextBlock(text=str(response.answer))], state=state) + + tool = FunctionTool(func=run_job, name=job.name, description=job.description) + if job.parameters: + tool.input_schema = job.parameters + return tool + + @classmethod + def _builtin_tools(cls) -> list[ToolBase]: + return [BypassAnalysisBash(), Edit(), Glob(), Grep(), Read(), Write()] + + @property + def session_path(self) -> Path: + """Path to the AgentScope session directory.""" + if self.app_context is None: + return self.workspace_path / "session" / "agentscope" + return self.workspace_path / self.app_context.app_config.session_dir / "agentscope" + + @staticmethod + def _validate_session_id(session_id: str, field: str = "session_id") -> str: + if not _UUID_RE.match(session_id): + raise ValueError(f"{field} must be a valid UUID: {session_id!r}") + return session_id.lower() + + def _cleanup_expired_sessions(self) -> None: + if self._session_cleanup_done or self.session_retention_days <= 0: + self._session_cleanup_done = True + return + session_path = self.session_path + if not session_path.is_dir(): + self._session_cleanup_done = True + return + cutoff = time.time() - self.session_retention_days * 24 * 60 * 60 + removed = 0 + for path in session_path.glob("*.jsonl"): + try: + if path.is_file() and path.stat().st_mtime < cutoff: + path.unlink() + removed += 1 + except OSError as exc: + self.logger.warning(f"Failed to clean expired AgentScope session {path}: {exc}") + if removed: + self.logger.info(f"Cleaned {removed} session(s) older than {self.session_retention_days} day(s)") + self._session_cleanup_done = True + + async def _load_state(self, kwargs: dict[str, Any], perm_mode: PermissionMode) -> AgentState: + resume = kwargs.get("resume") or "" + session_id = kwargs.get("session_id") or "" + fork_session = bool(kwargs.get("fork_session", False)) + if resume: + resume = self._validate_session_id(resume, "resume") + if session_id: + session_id = self._validate_session_id(session_id) + if session_id and resume and not fork_session: + raise ValueError("session_id cannot be used with resume unless fork_session=True") + if resume: + handler = AsStateHandler.for_session(self.session_path, resume) + state = await handler.load_or_none() + if state is None: + raise FileNotFoundError(f"AgentScope session not found: {resume}") + state.permission_context = PermissionContext(mode=perm_mode) + state.session_id = resume + if fork_session: + return AgentState( + session_id=session_id or str(uuid4()), + summary=state.summary, + context=list(state.context), + permission_context=PermissionContext(mode=perm_mode), + ) + return state + return AgentState(session_id=session_id or str(uuid4()), permission_context=PermissionContext(mode=perm_mode)) + + async def _dump_state(self, state: AgentState) -> None: + await AsStateHandler.for_session(self.session_path, state.session_id).dump(state) + + def _resolve_skills(self, skills: list[str] | str | None) -> list[str]: + if skills is None: + return [] + if skills == "all": + return [str(self.project_skills_root)] + if isinstance(skills, str): + skills = [skills] + return [str(self.project_skills_root / skill) for skill in skills] + + def _load_tool_env(self) -> dict[str, str]: + project_env = self.project_path / ".env" + return load_env(project_env) if project_env.exists() else load_env() + + async def _build_agent(self, inputs: Any, **kwargs) -> tuple[Agent, Any]: + model = self.as_llm.model if self.as_llm else None + if model is None: + raise ValueError("AsAgentWrapper requires a bound as_llm component with a valid model.") + kwargs = self._merged_kwargs(kwargs) + self._cleanup_expired_sessions() + self._load_tool_env() + + system_prompt = kwargs.get("system_prompt", "You are a helpful assistant.") + resolved_jobs = self._resolve_job_tools(kwargs.get("job_tools", [])) + skills = self._resolve_skills(kwargs.get("skills")) + toolkit = kwargs.get("toolkit") or Toolkit( + tools=[*self._builtin_tools(), *(self._make_tool(job) for job in resolved_jobs)], + skills_or_loaders=skills, + ) + perm_mode = PermissionMode(kwargs.get("permission_mode", "bypass")) + state = await self._load_state(kwargs, perm_mode) + agent = Agent( + name=self.name, + system_prompt=system_prompt, + model=model, + toolkit=toolkit, + state=state, + model_config=ModelConfig(**(kwargs.get("model_config") or {})), + context_config=ContextConfig(**(kwargs.get("context_config") or {})), + react_config=ReActConfig(**(kwargs.get("react_config") or {})), + ) + if isinstance(inputs, str): + inputs = UserMsg(name="user", content=inputs) + return agent, inputs + + async def reply(self, inputs: Any, **kwargs) -> dict: + kwargs = self._merged_kwargs(kwargs) + agent, inputs = await self._build_agent(inputs, **kwargs) + await agent.observe(inputs) + await agent.reply() + await self._dump_state(agent.state) + last_msg = agent.state.context[-1] + result = { + "session_id": agent.state.session_id, + "last_message": last_msg.model_dump(), + "result": last_msg.get_text_content(), + } + output_schema: dict | None = kwargs.get("output_schema") + if output_schema is not None: + assert self.as_llm is not None and self.as_llm.model is not None + res = await self.as_llm.model.generate_structured_output( + messages=agent.state.context, + structured_model=output_schema, + ) + result["structured_output"] = res.content + return result + + @classmethod + # pylint: disable=too-many-return-statements + def _event_to_chunk(cls, event: Any) -> StreamChunk | None: + """Convert AgentScope event to StreamChunk; None to skip.""" + if isinstance(event, ReplyStartEvent): + meta = {"reply_id": event.reply_id, "name": event.name, "role": event.role} + return cls._chunk(ChunkEnum.REPLY_START, session_id=event.session_id, chunk="", metadata=meta) + if isinstance(event, ReplyEndEvent): + return cls._chunk( + ChunkEnum.REPLY_END, + session_id=event.session_id, + chunk="", + metadata={"reply_id": event.reply_id}, + ) + + for event_cls, chunk_type, attr in ( + (TextBlockStartEvent, ChunkEnum.CONTENT, None), + (TextBlockDeltaEvent, ChunkEnum.CONTENT, "delta"), + (TextBlockEndEvent, ChunkEnum.CONTENT, None), + (ThinkingBlockStartEvent, ChunkEnum.THINK, None), + (ThinkingBlockDeltaEvent, ChunkEnum.THINK, "delta"), + (ThinkingBlockEndEvent, ChunkEnum.THINK, None), + (DataBlockStartEvent, ChunkEnum.DATA, None), + (DataBlockDeltaEvent, ChunkEnum.DATA, "data"), + (DataBlockEndEvent, ChunkEnum.DATA, None), + ): + if isinstance(event, event_cls): + kwargs = {"block_id": event.block_id, "chunk": getattr(event, attr) if attr else ""} + if isinstance(event, (DataBlockStartEvent, DataBlockDeltaEvent)): + kwargs["media_type"] = event.media_type + return cls._chunk(chunk_type, **kwargs) + + if isinstance(event, ToolCallStartEvent): + payload = {"name": event.tool_call_name, "id": event.tool_call_id} + return cls._chunk( + ChunkEnum.TOOL_CALL, + tool_call_id=event.tool_call_id, + tool_call_name=event.tool_call_name, + chunk=json.dumps(payload), + ) + if isinstance(event, ToolCallDeltaEvent): + return cls._chunk(ChunkEnum.TOOL_CALL, tool_call_id=event.tool_call_id, chunk=event.delta) + if isinstance(event, ToolCallEndEvent): + return cls._chunk(ChunkEnum.TOOL_CALL, tool_call_id=event.tool_call_id, chunk="") + if isinstance(event, ToolResultStartEvent): + return cls._chunk( + ChunkEnum.TOOL_RESULT, + tool_call_id=event.tool_call_id, + tool_call_name=event.tool_call_name, + chunk="", + ) + if isinstance(event, ToolResultTextDeltaEvent): + return cls._chunk(ChunkEnum.TOOL_RESULT, tool_call_id=event.tool_call_id, chunk=event.delta) + if isinstance(event, ToolResultDataDeltaEvent): + return cls._chunk( + ChunkEnum.TOOL_RESULT, + tool_call_id=event.tool_call_id, + chunk=event.data, + media_type=event.media_type, + metadata={"url": event.url} if event.url else {}, + ) + if isinstance(event, ToolResultEndEvent): + return cls._chunk( + ChunkEnum.TOOL_RESULT, + tool_call_id=event.tool_call_id, + chunk="", + metadata={"state": str(event.state)}, + ) + if isinstance(event, ModelCallStartEvent): + return cls._chunk(ChunkEnum.USAGE, chunk="", metadata={"model_name": getattr(event, "model_name", None)}) + if isinstance(event, ModelCallEndEvent): + usage = {"input_tokens": event.input_tokens, "output_tokens": event.output_tokens} + return cls._chunk( + ChunkEnum.USAGE, + chunk=json.dumps(usage), + input_tokens=event.input_tokens, + output_tokens=event.output_tokens, + metadata={"model_name": getattr(event, "model_name", None)}, + ) + if isinstance(event, ExceedMaxItersEvent): + return cls._chunk(ChunkEnum.ERROR, chunk="Exceeded max iterations") + return None + + async def reply_stream(self, inputs: Any, **kwargs) -> AsyncGenerator[StreamChunk, None]: + agent, inputs = await self._build_agent(inputs, **kwargs) + async for event in agent.reply_stream(inputs): + chunk = self._event_to_chunk(event) + if chunk is not None: + chunk.session_id = chunk.session_id or agent.state.session_id + yield chunk + await self._dump_state(agent.state) diff --git a/flowllm/components/agent_wrapper/base_agent_wrapper.py b/flowllm/components/agent_wrapper/base_agent_wrapper.py new file mode 100644 index 0000000..6c4760d --- /dev/null +++ b/flowllm/components/agent_wrapper/base_agent_wrapper.py @@ -0,0 +1,79 @@ +"""Base agent wrapper component.""" + +from abc import abstractmethod +from collections.abc import AsyncGenerator +from pathlib import Path +from typing import Any, TYPE_CHECKING + +from pydantic import BaseModel + +from ..base_component import BaseComponent +from ...enumeration import ChunkEnum, ComponentEnum +from ...schema import StreamChunk + +if TYPE_CHECKING: + from ..job.base_job import BaseJob + + +class BaseAgentWrapper(BaseComponent): + """Abstract base for agent wrappers with swappable backends.""" + + component_type = ComponentEnum.AGENT_WRAPPER + + def set_system_prompt(self, prompt: str) -> "BaseAgentWrapper": + """Set the system prompt.""" + self.kwargs["system_prompt"] = prompt + return self + + def add_job_tools(self, job_tools: list[str]) -> "BaseAgentWrapper": + """Append job tools by name.""" + self.kwargs.setdefault("job_tools", []).extend(job_tools) + return self + + def add_skills(self, skills: list[str] | str) -> "BaseAgentWrapper": + """Set skills configuration.""" + self.kwargs["skills"] = skills + return self + + @property + def project_path(self) -> Path: + """Root path of the project.""" + return self.workspace_path + + @property + def project_skills_root(self) -> Path: + """Skills directory under project root.""" + return self.project_path / "skills" + + def set_output_schema(self, schema: dict | type[BaseModel]) -> "BaseAgentWrapper": + """Set structured output schema.""" + if isinstance(schema, type) and issubclass(schema, BaseModel): + schema = schema.model_json_schema() + self.kwargs["output_schema"] = schema + return self + + def _resolve_job_tools(self, job_tools: list[str]) -> list["BaseJob"]: + if not job_tools: + return [] + if self.app_context is None: + raise RuntimeError("Cannot resolve job_tools without an app_context") + resolved: list["BaseJob"] = [] + for name in job_tools: + if (job := self.app_context.jobs.get(name)) is None: + raise KeyError(f"Job '{name}' not found in app_context.jobs") + resolved.append(job) + return resolved + + def _merged_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: + return {**self.kwargs, **kwargs} + + @staticmethod + def _chunk(chunk_type: ChunkEnum = ChunkEnum.CONTENT, **kwargs: Any) -> StreamChunk: + return StreamChunk(chunk_type=chunk_type, **kwargs) + + @abstractmethod + async def reply(self, inputs: Any, **kwargs) -> dict: + """Send inputs and return dict with session_id and last_message.""" + + async def reply_stream(self, inputs: Any, **kwargs) -> AsyncGenerator[StreamChunk, None]: + """Stream agent events as StreamChunk objects.""" diff --git a/flowllm/components/agent_wrapper/cc_agent_wrapper.py b/flowllm/components/agent_wrapper/cc_agent_wrapper.py new file mode 100644 index 0000000..8e4a560 --- /dev/null +++ b/flowllm/components/agent_wrapper/cc_agent_wrapper.py @@ -0,0 +1,499 @@ +"""Claude Code SDK backend for the unified agent wrapper.""" + +import json +import os +import shutil +from collections.abc import AsyncGenerator +from dataclasses import asdict +from pathlib import Path +from typing import Any, TYPE_CHECKING + +from .base_agent_wrapper import BaseAgentWrapper +from ..component_registry import R +from ...enumeration import ChunkEnum +from ...schema import StreamChunk +from ...utils.env_utils import load_env + +if TYPE_CHECKING: + from ..job.base_job import BaseJob + from claude_agent_sdk.types import SessionKey, SessionStoreEntry, SessionStoreListEntry + + +class _CcFileSessionStore: + """File-backed Claude Code SessionStore.""" + + def __init__(self, root: Path) -> None: + self.root = root + + @staticmethod + def _safe_parts(value: str) -> list[str]: + parts = [part for part in value.split("/") if part] + if not parts or any(part in {".", ".."} for part in parts): + raise ValueError(f"Invalid session store path component: {value!r}") + return parts + + def _path_for_key(self, key: "SessionKey") -> Path: + session_id, subpath = key["session_id"], key.get("subpath") + path = self.root.joinpath(*self._safe_parts(session_id)) + if subpath: + path = path.joinpath(*self._safe_parts(subpath)).with_suffix(".jsonl") + else: + path = path.with_suffix(".jsonl") + resolved_root, resolved_path = self.root.resolve(), path.resolve() + if resolved_root != resolved_path and resolved_root not in resolved_path.parents: + raise ValueError(f"Session store path escapes root: {resolved_path}") + return path + + @staticmethod + def _read_entries(path: Path) -> list["SessionStoreEntry"]: + if not path.exists(): + return [] + return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()] + + async def append(self, key: "SessionKey", entries: list["SessionStoreEntry"]) -> None: + """Append entries, deduplicating by UUID.""" + path = self._path_for_key(key) + path.parent.mkdir(parents=True, exist_ok=True) + existing_uuids = { + entry.get("uuid") for entry in self._read_entries(path) if isinstance(entry, dict) and entry.get("uuid") + } + new_entries = [ + entry for entry in entries if not (isinstance(entry, dict) and entry.get("uuid") in existing_uuids) + ] + if not new_entries: + return + with path.open("a", encoding="utf-8") as f: + for entry in new_entries: + f.write(json.dumps(entry, ensure_ascii=False, separators=(",", ":")) + "\n") + + async def load(self, key: "SessionKey") -> list["SessionStoreEntry"] | None: + """Load all entries for a session key.""" + path = self._path_for_key(key) + return self._read_entries(path) if path.exists() else None + + async def list_sessions(self, _project_key: str) -> list["SessionStoreListEntry"]: + """List all sessions with IDs and mtimes.""" + if not self.root.exists(): + return [] + return [ + {"session_id": path.stem, "mtime": int(path.stat().st_mtime * 1000)} + for path in self.root.glob("*.jsonl") + if path.is_file() + ] + + async def delete(self, key: "SessionKey") -> None: + """Delete a session or subpath entry.""" + path = self._path_for_key(key) + if path.exists(): + path.unlink() + if not key.get("subpath"): + session_dir = self.root.joinpath(*self._safe_parts(key["session_id"])) + if session_dir.exists(): + for child in sorted(session_dir.rglob("*"), reverse=True): + if child.is_file(): + child.unlink() + elif child.is_dir(): + child.rmdir() + session_dir.rmdir() + + async def list_subkeys(self, key: dict[str, str]) -> list[str]: + """List subkeys under a session.""" + session_dir = self.root.joinpath(*self._safe_parts(key["session_id"])) + if not session_dir.exists(): + return [] + return [str(p.relative_to(session_dir).with_suffix("")) for p in session_dir.rglob("*.jsonl") if p.is_file()] + + +@R.register("claude_code") +class CcAgentWrapper(BaseAgentWrapper): + """Agent wrapper backed by Claude Code SDK.""" + + DEFAULT_DISALLOWED_TOOLS = ["WebSearch"] + + @staticmethod + def _first_non_empty(*values: Any) -> str: + for value in values: + if isinstance(value, str) and value: + return value + return "" + + def _default_llm_credential(self) -> dict[str, Any]: + if self.app_context is None: + return {} + components = self.app_context.app_config.components + llm_configs = components.get("as_llm") or components.get("AS_LLM") or components.get("as_llm".upper()) + if llm_configs is None: + from ...enumeration import ComponentEnum + + llm_configs = components.get(ComponentEnum.AS_LLM) + if not isinstance(llm_configs, dict): + return {} + default_llm = llm_configs.get("default") + credential = getattr(default_llm, "credential", None) + return credential if isinstance(credential, dict) else {} + + def _claude_code_api_env(self, kwargs: dict[str, Any]) -> dict[str, str]: + credential = kwargs.get("credential") if isinstance(kwargs.get("credential"), dict) else {} + default_credential = self._default_llm_credential() + base_url = self._first_non_empty( + kwargs.get("base_url"), + credential.get("base_url"), + os.getenv("ANTHROPIC_BASE_URL"), + os.getenv("CLAUDE_CODE_BASE_URL"), + os.getenv("LLM_BASE_URL"), + default_credential.get("base_url"), + ) + api_key = self._first_non_empty( + kwargs.get("api_key"), + credential.get("api_key"), + os.getenv("ANTHROPIC_AUTH_TOKEN"), + os.getenv("CLAUDE_CODE_API_KEY"), + os.getenv("LLM_API_KEY"), + default_credential.get("api_key"), + ) + env: dict[str, str] = {} + if base_url: + env["ANTHROPIC_BASE_URL"] = base_url + if api_key: + env["ANTHROPIC_AUTH_TOKEN"] = api_key + return env + + @property + def session_path(self) -> Path: + """Path to the Claude Code session directory.""" + if self.app_context is None: + return self.workspace_path / "session" + return self.workspace_path / self.app_context.app_config.session_dir + + def _ensure_claude_skill_dir(self, config_dir: Path) -> None: + """Symlink project skills to Claude Code discovery locations.""" + project_skills = self.project_skills_root + if not project_skills.exists(): + return + for target in (self.project_path / ".claude" / "skills", config_dir / "skills"): + target.parent.mkdir(parents=True, exist_ok=True) + try: + if target.exists() or target.is_symlink(): + if target.resolve() == project_skills.resolve(): + continue + if target.is_dir() and not target.is_symlink(): + shutil.rmtree(target) + else: + target.unlink() + target.symlink_to(project_skills, target_is_directory=True) + except OSError as exc: + self.logger.warning(f"Failed to link Claude Code skills directory {target}: {exc}") + + @staticmethod + def _make_tool(job: "BaseJob"): + from claude_agent_sdk import SdkMcpTool + + async def run_job(args): + response = await job(**args) + return {"content": [{"type": "text", "text": str(response.answer)}], "is_error": not response.success} + + return SdkMcpTool(name=job.name, description=job.description, input_schema=job.parameters, handler=run_job) + + def _build_options(self, inputs: Any, stream: bool = False, **kwargs) -> Any: + """Build ClaudeAgentOptions; stream=True enables partial messages.""" + from claude_agent_sdk import create_sdk_mcp_server + from claude_agent_sdk.types import ClaudeAgentOptions + + kwargs = self._merged_kwargs(kwargs) + skills = kwargs.get("skills") + if isinstance(skills, str) and skills != "all": + kwargs["skills"] = [skills] + if "setting_sources" not in kwargs and kwargs.get("skills") is None: + kwargs["setting_sources"] = [] + disallowed_tools = list(kwargs.get("disallowed_tools") or []) + for tool_name in self.DEFAULT_DISALLOWED_TOOLS: + if tool_name not in disallowed_tools: + disallowed_tools.append(tool_name) + kwargs["disallowed_tools"] = disallowed_tools + + opts = ClaudeAgentOptions() + if stream: + opts.include_partial_messages = True + skip_keys = {"job_tools", "output_schema", "api_key", "base_url", "credential"} + for k, v in kwargs.items(): + if k not in skip_keys and hasattr(opts, k): + setattr(opts, k, v) + + model = getattr(opts, "model", None) or kwargs.get("model") + project_env = self.project_path / ".env" + opts.env.update(load_env(project_env) if project_env.exists() else load_env()) + extra_env_dict: dict = self._claude_code_api_env(kwargs) + if model: + extra_env_dict.update( + { + "ANTHROPIC_MODEL": model, + "ANTHROPIC_DEFAULT_HAIKU_MODEL": model, + "ANTHROPIC_DEFAULT_SONNET_MODEL": model, + "ANTHROPIC_DEFAULT_OPUS_MODEL": model, + }, + ) + opts.env.update(extra_env_dict) + self.session_path.mkdir(parents=True, exist_ok=True) + opts.cwd = opts.cwd or self.project_path + claude_config_dir = self.session_path / "claude_config" + opts.env.setdefault("CLAUDE_CONFIG_DIR", str(claude_config_dir)) + if opts.skills is not None: + self._ensure_claude_skill_dir(claude_config_dir) + opts.session_store = opts.session_store or _CcFileSessionStore(self.session_path / "claude_code") + + resolved_jobs = self._resolve_job_tools(kwargs.get("job_tools", [])) + if resolved_jobs: + sdk_tools = [self._make_tool(job) for job in resolved_jobs] + server = create_sdk_mcp_server(name="mcp_server", tools=sdk_tools) + opts.mcp_servers = opts.mcp_servers if isinstance(opts.mcp_servers, dict) else {} + opts.mcp_servers["mcp_server"] = server + opts.allowed_tools.extend(job.name for job in resolved_jobs) + if output_schema := kwargs.get("output_schema"): + opts.output_format = {"type": "json_schema", "schema": output_schema} + if not isinstance(inputs, str): + raise NotImplementedError("Only string input is supported for Claude Code.") + return opts + + @classmethod + # pylint: disable=too-many-return-statements + def _raw_event_to_chunk( + cls, + raw: dict, + session_id: str | None = None, + block_ids: dict[int, str] | None = None, + block_types: dict[int, str] | None = None, + tool_call_names: dict[int, str] | None = None, + ) -> StreamChunk | None: + """Convert raw Anthropic streaming event to StreamChunk; None to skip.""" + event_type = raw.get("type") + + if event_type == "message_start": + message = raw.get("message", {}) + meta = {"message_id": message.get("id"), "model": message.get("model"), "role": message.get("role")} + return cls._chunk(ChunkEnum.REPLY_START, session_id=session_id, chunk="", metadata=meta) + if event_type == "message_delta": + delta, usage = raw.get("delta", {}), raw.get("usage", {}) + return cls._chunk( + ChunkEnum.REPLY_END, + session_id=session_id, + chunk="", + output_tokens=usage.get("output_tokens"), + metadata={"stop_reason": delta.get("stop_reason")}, + ) + if event_type == "message_stop": + return cls._chunk(ChunkEnum.REPLY_END, session_id=session_id, chunk="") + + if event_type == "content_block_start": + idx, content_block = raw.get("index", 0), raw.get("content_block", {}) + block_type, bid = content_block.get("type", ""), content_block.get("id", "") + if block_ids is not None and bid: + block_ids[idx] = bid + if block_types is not None and block_type: + block_types[idx] = block_type + if tool_call_names is not None and content_block.get("name"): + tool_call_names[idx] = content_block["name"] + if block_type == "text": + return cls._chunk(ChunkEnum.CONTENT, block_id=bid, chunk=content_block.get("text", "")) + if block_type == "thinking": + return cls._chunk(ChunkEnum.THINK, block_id=bid, chunk=content_block.get("thinking", "")) + if block_type == "tool_use": + payload = {"name": content_block.get("name"), "id": content_block.get("id")} + return cls._chunk( + ChunkEnum.TOOL_CALL, + block_id=bid, + tool_call_id=content_block.get("id"), + tool_call_name=content_block.get("name"), + chunk=json.dumps(payload), + ) + return None + + if event_type == "content_block_delta": + delta, delta_type = raw.get("delta", {}), raw.get("delta", {}).get("type", "") + idx = raw.get("index", 0) + bid = block_ids.get(idx) if block_ids else None + tc_name = tool_call_names.get(idx) if tool_call_names else None + if delta_type == "text_delta": + return cls._chunk(ChunkEnum.CONTENT, block_id=bid, chunk=delta.get("text", "")) + if delta_type == "thinking_delta": + return cls._chunk(ChunkEnum.THINK, block_id=bid, chunk=delta.get("thinking", "")) + if delta_type == "input_json_delta": + return cls._chunk( + ChunkEnum.TOOL_CALL, + block_id=bid, + tool_call_id=bid, + tool_call_name=tc_name, + chunk=delta.get("partial_json", ""), + ) + return None + + if event_type == "content_block_stop": + idx = raw.get("index", 0) + bid = block_ids.get(idx) if block_ids else None + btype = block_types.get(idx) if block_types else None + tc_name = tool_call_names.get(idx) if tool_call_names else None + if btype == "tool_use": + return cls._chunk(ChunkEnum.TOOL_CALL, block_id=bid, tool_call_id=bid, tool_call_name=tc_name, chunk="") + if btype == "thinking": + return cls._chunk(ChunkEnum.THINK, block_id=bid, chunk="") + return cls._chunk(ChunkEnum.CONTENT, block_id=bid, chunk="") + + return None + + @classmethod + def _message_content_to_chunks( + cls, + msg: Any, + session_id: str | None = None, + visible_tool_call_ids: set[str] | None = None, + include_text: bool = False, + ) -> list[StreamChunk]: + """Convert non-partial SDK message content blocks to stream chunks.""" + chunks: list[StreamChunk] = [] + content = getattr(msg, "content", None) + if not isinstance(content, list): + return chunks + for block in content: + block_name = block.__class__.__name__ + if include_text and block_name == "TextBlock": + if text := getattr(block, "text", ""): + chunks.append(cls._chunk(ChunkEnum.CONTENT, session_id=session_id, chunk=text)) + elif include_text and isinstance(block, dict) and block.get("type") == "text": + if text := block.get("text", ""): + chunks.append(cls._chunk(ChunkEnum.CONTENT, session_id=session_id, chunk=text)) + elif include_text and isinstance(block, str): + chunks.append(cls._chunk(ChunkEnum.CONTENT, session_id=session_id, chunk=block)) + elif block_name in {"ToolResultBlock", "ServerToolResultBlock"}: + tool_use_id = getattr(block, "tool_use_id", None) + if visible_tool_call_ids is not None and tool_use_id not in visible_tool_call_ids: + continue + payload: dict[str, Any] = {"tool_use_id": tool_use_id, "content": getattr(block, "content", None)} + if hasattr(block, "is_error"): + payload["is_error"] = getattr(block, "is_error") + chunks.append( + cls._chunk( + ChunkEnum.TOOL_RESULT, + session_id=session_id, + block_id=tool_use_id, + tool_call_id=tool_use_id, + chunk=payload, + ), + ) + return chunks + + @staticmethod + def _result_message_is_error(msg: Any) -> bool: + subtype = getattr(msg, "subtype", None) + if isinstance(subtype, str) and subtype.lower() == "success": + return False + is_error = getattr(msg, "is_error", False) + if isinstance(is_error, bool): + return is_error + if isinstance(is_error, str): + return is_error.lower() in {"true", "error", "errored", "failed", "failure"} + return isinstance(subtype, str) and subtype.lower() in {"error", "failed", "failure"} + + @staticmethod + def _is_trailing_success_error(exc: Exception) -> bool: + return "Claude Code returned an error result: success" in str(exc) + + async def reply(self, inputs: Any, **kwargs) -> dict: + from claude_agent_sdk import query, ResultMessage + + opts = self._build_options(inputs, stream=False, **kwargs) + last_msg = None + async for msg in query(prompt=inputs, options=opts): + if isinstance(msg, ResultMessage): + last_msg = msg + if last_msg is None: + raise ValueError("No message received from Claude Code.") + result = {"session_id": last_msg.session_id or "", "last_message": asdict(last_msg), "result": last_msg.result} + output_schema = kwargs.get("output_schema") or self.kwargs.get("output_schema") + if output_schema and last_msg.structured_output: + result["structured_output"] = last_msg.structured_output + return result + + async def reply_stream(self, inputs: Any, **kwargs) -> AsyncGenerator[StreamChunk, None]: + from claude_agent_sdk import query, ResultMessage, AssistantMessage, StreamEvent, UserMessage + from claude_agent_sdk.types import RateLimitEvent + + opts = self._build_options(inputs, stream=True, **kwargs) + block_ids: dict[int, str] = {} + block_types: dict[int, str] = {} + tool_call_names: dict[int, str] = {} + visible_tool_call_ids: set[str] = set() + current_session_id: str | None = None + emitted_content = False + received_result_message = False + + stream = query(prompt=inputs, options=opts) + try: + async for msg in stream: + if isinstance(msg, StreamEvent): + current_session_id = msg.session_id or current_session_id + chunk = self._raw_event_to_chunk( + msg.event, + session_id=msg.session_id, + block_ids=block_ids, + block_types=block_types, + tool_call_names=tool_call_names, + ) + if chunk is not None: + chunk.session_id = chunk.session_id or msg.session_id + if chunk.chunk_type == ChunkEnum.TOOL_CALL and chunk.tool_call_id: + visible_tool_call_ids.add(chunk.tool_call_id) + if chunk.chunk_type == ChunkEnum.CONTENT and chunk.chunk: + emitted_content = True + yield chunk + elif isinstance(msg, UserMessage): + for chunk in self._message_content_to_chunks(msg, current_session_id, visible_tool_call_ids): + yield chunk + elif isinstance(msg, ResultMessage): + received_result_message = True + current_session_id = msg.session_id or current_session_id + if not emitted_content and getattr(msg, "result", None): + emitted_content = True + yield self._chunk(ChunkEnum.CONTENT, session_id=msg.session_id or "", chunk=msg.result) + meta = { + "duration_ms": msg.duration_ms, + "duration_api_ms": msg.duration_api_ms, + "stop_reason": msg.stop_reason, + "num_turns": msg.num_turns, + } + yield self._chunk( + ChunkEnum.USAGE, + session_id=msg.session_id or "", + chunk=json.dumps(msg.usage or {}), + metadata=meta, + ) + if self._result_message_is_error(msg): + yield self._chunk( + ChunkEnum.ERROR, + session_id=msg.session_id or "", + chunk=str(msg.errors) if msg.errors else "Unknown error", + ) + yield self._chunk(ChunkEnum.REPLY_END, session_id=msg.session_id or "", chunk="") + elif isinstance(msg, AssistantMessage): + current_session_id = msg.session_id or current_session_id + for chunk in self._message_content_to_chunks( + msg, + current_session_id, + visible_tool_call_ids, + include_text=not emitted_content, + ): + if chunk.chunk_type == ChunkEnum.CONTENT and chunk.chunk: + emitted_content = True + yield chunk + elif isinstance(msg, RateLimitEvent): + yield self._chunk(ChunkEnum.ERROR, session_id=msg.session_id, chunk="Rate limit exceeded") + except Exception as exc: + if received_result_message and self._is_trailing_success_error(exc): + self.logger.debug(f"Ignoring trailing success error: {exc}") + else: + raise + finally: + try: + await stream.aclose() + except Exception as exc: + if not (received_result_message and self._is_trailing_success_error(exc)): + raise + self.logger.debug(f"Ignoring stream close error: {exc}") diff --git a/flowllm/components/application_context.py b/flowllm/components/application_context.py new file mode 100644 index 0000000..f816a03 --- /dev/null +++ b/flowllm/components/application_context.py @@ -0,0 +1,22 @@ +"""Application context: shared state for components, jobs, and service.""" + +from typing import TYPE_CHECKING, Any + +from ..enumeration import ComponentEnum +from ..schema import ApplicationConfig + +if TYPE_CHECKING: + from .base_component import BaseComponent + from .job import BaseJob + from .service import BaseService + + +class ApplicationContext: + """Passive state container holding parsed config and wired components.""" + + def __init__(self, **kwargs): + self.app_config: ApplicationConfig = ApplicationConfig(**kwargs) + self.service: "BaseService | None" = None + self.components: dict[ComponentEnum, dict[str, "BaseComponent"]] = {} + self.jobs: dict[str, "BaseJob"] = {} + self.metadata: dict[str, Any] = {} diff --git a/flowllm/components/as_embedding/__init__.py b/flowllm/components/as_embedding/__init__.py new file mode 100644 index 0000000..624aa1a --- /dev/null +++ b/flowllm/components/as_embedding/__init__.py @@ -0,0 +1,95 @@ +"""AgentScope embedding model wrappers.""" + +from typing import Any + +from agentscope.credential import ( + CredentialBase, + DashScopeCredential, + GeminiCredential, + OllamaCredential, + OpenAICredential, +) +from agentscope.embedding import EmbeddingModelBase + +from ..base_component import BaseComponent +from ..component_registry import R +from ...enumeration import ComponentEnum + + +class BaseAsEmbedding(BaseComponent): + """Base wrapper for AgentScope embedding models.""" + + component_type = ComponentEnum.AS_EMBEDDING + credential_cls: type[CredentialBase] + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.model: EmbeddingModelBase[Any] | None = None + + @property + def dimensions(self) -> int: + """Embedding vector dimensionality.""" + assert self.model is not None + return self.model.dimensions + + async def __call__(self, inputs: list[Any], **kwargs) -> list[list[float]]: + assert self.model is not None + response = await self.model(inputs, **kwargs) # pylint: disable=not-callable + return response.embeddings + + async def _start(self) -> None: + if self.model is not None: + return + kwargs = dict(self.kwargs) + credential = self.credential_cls(**kwargs.pop("credential", {})) + model_cls = self.credential_cls.get_embedding_model_class() + if model_cls is None: + raise ValueError(f"{self.credential_cls.__name__} does not support embeddings.") + params_dict = kwargs.pop("parameters", None) + parameters = model_cls.Parameters(**params_dict) if params_dict else None + self.model = model_cls(credential=credential, parameters=parameters, **kwargs) + + +@R.register("openai") +class OpenAIAsEmbedding(BaseAsEmbedding): + """OpenAI embedding model wrapper.""" + + credential_cls = OpenAICredential + + +@R.register("dashscope") +class DashScopeAsEmbedding(BaseAsEmbedding): + """DashScope embedding model wrapper.""" + + credential_cls = DashScopeCredential + + +@R.register("dashscope_multimodal") +class DashScopeMultiModalAsEmbedding(BaseAsEmbedding): + """DashScope multimodal embedding model wrapper.""" + + credential_cls = DashScopeCredential + + +@R.register("gemini") +class GeminiAsEmbedding(BaseAsEmbedding): + """Gemini embedding model wrapper.""" + + credential_cls = GeminiCredential + + +@R.register("ollama") +class OllamaAsEmbedding(BaseAsEmbedding): + """Ollama embedding model wrapper.""" + + credential_cls = OllamaCredential + + +__all__ = [ + "BaseAsEmbedding", + "OpenAIAsEmbedding", + "DashScopeAsEmbedding", + "DashScopeMultiModalAsEmbedding", + "GeminiAsEmbedding", + "OllamaAsEmbedding", +] diff --git a/flowllm/components/as_llm/__init__.py b/flowllm/components/as_llm/__init__.py new file mode 100644 index 0000000..27ef4dc --- /dev/null +++ b/flowllm/components/as_llm/__init__.py @@ -0,0 +1,108 @@ +"""LLM model wrappers for AgentScope.""" + +from agentscope.credential import ( + AnthropicCredential, + CredentialBase, + DashScopeCredential, + DeepSeekCredential, + GeminiCredential, + MoonshotCredential, + OllamaCredential, + OpenAICredential, + XAICredential, +) +from agentscope.model import ChatModelBase + +from ..base_component import BaseComponent +from ..component_registry import R +from ...enumeration import ComponentEnum + + +class BaseAsLLM(BaseComponent): + """Base wrapper for AgentScope chat models. Subclasses set credential_cls.""" + + component_type = ComponentEnum.AS_LLM + credential_cls: type[CredentialBase] + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.model: ChatModelBase | None = None + + async def _start(self) -> None: + if self.model is not None: + return + kwargs = dict(self.kwargs) + credential = self.credential_cls(**kwargs.pop("credential", {})) + model_cls = credential.get_chat_model_class() + params_dict = kwargs.pop("parameters", None) + parameters = model_cls.Parameters(**params_dict) if params_dict else None + self.model = model_cls(credential=credential, parameters=parameters, **kwargs) + + +@R.register("openai") +class OpenAIAsLLM(BaseAsLLM): + """OpenAI chat model wrapper.""" + + credential_cls = OpenAICredential + + +@R.register("anthropic") +class AnthropicAsLLM(BaseAsLLM): + """Anthropic chat model wrapper.""" + + credential_cls = AnthropicCredential + + +@R.register("dashscope") +class DashScopeAsLLM(BaseAsLLM): + """DashScope chat model wrapper.""" + + credential_cls = DashScopeCredential + + +@R.register("deepseek") +class DeepSeekAsLLM(BaseAsLLM): + """DeepSeek chat model wrapper.""" + + credential_cls = DeepSeekCredential + + +@R.register("gemini") +class GeminiAsLLM(BaseAsLLM): + """Gemini chat model wrapper.""" + + credential_cls = GeminiCredential + + +@R.register("moonshot") +class MoonshotAsLLM(BaseAsLLM): + """Moonshot chat model wrapper.""" + + credential_cls = MoonshotCredential + + +@R.register("ollama") +class OllamaAsLLM(BaseAsLLM): + """Ollama chat model wrapper.""" + + credential_cls = OllamaCredential + + +@R.register("xai") +class XAIAsLLM(BaseAsLLM): + """xAI chat model wrapper.""" + + credential_cls = XAICredential + + +__all__ = [ + "BaseAsLLM", + "OpenAIAsLLM", + "AnthropicAsLLM", + "DashScopeAsLLM", + "DeepSeekAsLLM", + "GeminiAsLLM", + "MoonshotAsLLM", + "OllamaAsLLM", + "XAIAsLLM", +] diff --git a/flowllm/components/base_component.py b/flowllm/components/base_component.py new file mode 100644 index 0000000..336522f --- /dev/null +++ b/flowllm/components/base_component.py @@ -0,0 +1,217 @@ +"""Base component with async lifecycle and dependency injection.""" + +import asyncio +from abc import ABC +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable, TypeVar, cast + +from ..enumeration import ComponentEnum +from ..utils import get_logger + +if TYPE_CHECKING: + from .application_context import ApplicationContext + +T = TypeVar("T", bound="BaseComponent") + + +class ComponentMixin: + """Shared state for components and steps: identity, config, workspace paths.""" + + component_type = ComponentEnum.BASE + + def __init__( + self, + name: str | None = None, + backend: str = "", + app_context: "ApplicationContext | None" = None, + **kwargs, + ) -> None: + self.name: str = name or self.__class__.__name__ + self.backend: str = backend + self.app_context: "ApplicationContext | None" = app_context + self.kwargs: dict = dict(kwargs) + logger = get_logger() + self.logger = logger.bind(component=self.name) if hasattr(logger, "bind") else logger + + @property + def workspace_path(self) -> Path: + """Workspace root (cwd when no app_context).""" + if self.app_context is None: + return Path.cwd() + return Path(self.app_context.app_config.workspace_dir).absolute() + + def to_workspace_relative(self, path: str | Path) -> str: + """Convert path to workspace-relative; return absolute when outside.""" + abs_path = Path(path).absolute() + try: + return str(abs_path.relative_to(self.workspace_path)) + except ValueError: + return str(abs_path) + + +class Dependency: + """Unresolved dependency placeholder; resolved on start(), errors on premature access.""" + + __slots__ = ("ctype", "name", "default_factory", "optional") + + def __init__( + self, + ctype: ComponentEnum, + name: str, + default_factory: Callable[[], Any] | None = None, + optional: bool = True, + ) -> None: + self.ctype, self.name, self.default_factory, self.optional = ctype, name, default_factory, optional + + def __repr__(self) -> str: + return f"" + + def __getattr__(self, item: str) -> Any: + raise RuntimeError(f"Dependency {self.ctype.value}:{self.name} accessed before start() (attribute '{item}')") + + +class BaseComponent(ComponentMixin, ABC): + """Async lifecycle base with bind-based dependency injection.""" + + component_type = ComponentEnum.BASE + + def __init__( + self, + name: str | None = None, + backend: str = "", + app_context: "ApplicationContext | None" = None, + **kwargs, + ) -> None: + super().__init__(name=name, backend=backend, app_context=app_context, **kwargs) + self._is_started: bool = False + self._lock: asyncio.Lock = asyncio.Lock() + self._owned: list["BaseComponent"] = [] # components created via bind() default_factory + + @property + def is_started(self) -> bool: + """Whether the component has been started.""" + return self._is_started + + @staticmethod + def bind( + name: str | None, + base_cls: type[T], + *, + default_factory: Callable[[], T] | None = None, + optional: bool = True, + ) -> T | None: + """Declare a dependency; returns a Dependency placeholder resolved on start().""" + if not name: + return None + ctype = getattr(base_cls, "component_type", None) + if not isinstance(ctype, ComponentEnum) or ctype is ComponentEnum.BASE: + raise TypeError(f"{base_cls.__name__} must declare a non-BASE ComponentEnum 'component_type'") + return cast(T, Dependency(ctype, name, default_factory, optional)) + + @property + def dependencies(self) -> list[Dependency]: + """List of unresolved Dependency placeholders.""" + return [v for v in self.__dict__.values() if isinstance(v, Dependency)] + + async def _resolve_bindings(self) -> None: + for attr, dep in list(self.__dict__.items()): + if isinstance(dep, Dependency): + self._resolve_one(attr, dep) + + def _resolve_one(self, attr: str, dep: Dependency) -> None: + if self.app_context is None: + self._resolve_standalone(attr, dep) + else: + self._resolve_from_context(attr, dep) + + def _resolve_standalone(self, attr: str, dep: Dependency) -> None: + """Standalone: use default_factory or None; required deps keep placeholder for clear errors.""" + if dep.default_factory is not None: + instance = dep.default_factory() + setattr(self, attr, instance) + if isinstance(instance, BaseComponent): + self._owned.append(instance) + elif dep.optional: + setattr(self, attr, None) + + def _resolve_from_context(self, attr: str, dep: Dependency) -> None: + target = self.app_context.components.get(dep.ctype, {}).get(dep.name) + if target is not None: + setattr(self, attr, target) + elif dep.optional: + setattr(self, attr, None) + else: + raise ValueError(f"{dep.ctype.value} '{dep.name}' not found.") + + @property + def workspace_metadata_path(self) -> Path: + """Root directory for workspace metadata.""" + if self.app_context is None: + return Path.cwd() / "metadata" + return self.workspace_path / self.app_context.app_config.metadata_dir + + @property + def component_metadata_path(self) -> Path: + """Metadata directory for this component type.""" + return self.workspace_metadata_path / self.component_type.value + + async def _start(self) -> None: + """Subclass hook after deps resolved.""" + + async def _close(self) -> None: + """Subclass hook during close.""" + + async def dump(self) -> None: + """Persist state to disk.""" + + async def load(self) -> None: + """Restore state from disk.""" + + async def start(self) -> None: + """Resolve deps → start owned → run _start.""" + async with self._lock: + if self._is_started: + return + await self._resolve_bindings() + for owned in self._owned: + await owned.start() + await self._start() + self._is_started = True + + async def close(self) -> None: + """Run _close → close owned in reverse order.""" + async with self._lock: + if not self._is_started: + return + first_error: BaseException | None = None + try: + await self._close() + except BaseException as exc: + first_error = exc + finally: + for owned in reversed(self._owned): + try: + await owned.close() + except BaseException as exc: + if first_error is None: + first_error = exc + else: + self.logger.exception(f"Failed to close owned component {owned.name}: {exc}") + self._is_started = False + if first_error is not None: + raise first_error + + async def restart(self) -> None: + """Close and restart the component.""" + await self.close() + await self.start() + + async def __call__(self, **kwargs): + raise NotImplementedError + + async def __aenter__(self) -> "BaseComponent": + await self.start() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: + await self.close() diff --git a/flowllm/components/client/__init__.py b/flowllm/components/client/__init__.py new file mode 100644 index 0000000..5c212ba --- /dev/null +++ b/flowllm/components/client/__init__.py @@ -0,0 +1,7 @@ +"""Client components.""" + +from .base_client import BaseClient +from .http_client import HttpClient +from .mcp_client import MCPClient + +__all__ = ["BaseClient", "HttpClient", "MCPClient"] diff --git a/flowllm/components/client/base_client.py b/flowllm/components/client/base_client.py new file mode 100644 index 0000000..a9d64d1 --- /dev/null +++ b/flowllm/components/client/base_client.py @@ -0,0 +1,39 @@ +"""Base client abstraction.""" + +import json +from abc import abstractmethod +from collections.abc import AsyncGenerator + +from ..base_component import BaseComponent +from ...enumeration import ComponentEnum + + +class BaseClient(BaseComponent): + """Abstract base for clients communicating with FlowLLM services.""" + + component_type = ComponentEnum.CLIENT + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.client = None + + async def _start(self) -> None: + pass + + async def _close(self) -> None: + pass + + @abstractmethod + def _execute(self, action: str, payload: dict) -> AsyncGenerator[str, None]: + """Yield text chunks for action.""" + + @abstractmethod + async def list_actions(self) -> list[dict]: + """Discover available actions.""" + + async def __call__(self, action: str, **kwargs) -> AsyncGenerator[str, None]: + if action == "list": + yield json.dumps(await self.list_actions(), indent=2, ensure_ascii=False) + return + async for chunk in self._execute(action, kwargs): + yield chunk diff --git a/flowllm/components/client/http_client.py b/flowllm/components/client/http_client.py new file mode 100644 index 0000000..0a45b4a --- /dev/null +++ b/flowllm/components/client/http_client.py @@ -0,0 +1,116 @@ +"""HTTP client for FlowLLM services.""" + +import json +import os +from collections.abc import AsyncGenerator + +import httpx + +from .base_client import BaseClient +from ..component_registry import R +from ...constants import FLOWLLM_SERVICE_INFO, FLOWLLM_DEFAULT_HOST, FLOWLLM_DEFAULT_PORT +from ...enumeration import ChunkEnum +from ...schema import StreamChunk + + +@R.register("http") +class HttpClient(BaseClient): + """HTTP client that auto-adapts to JSON or SSE via Content-Type.""" + + def __init__(self, host: str | None = None, port: int | None = None, timeout: float = 30.0, **kwargs): + super().__init__(**kwargs) + if not (host and port): + if service_info := os.environ.get(FLOWLLM_SERVICE_INFO): + try: + data = json.loads(service_info) + host, port = data["host"], data["port"] + except Exception: + self.logger.warning(f"Invalid service info: {service_info}") + host, port = FLOWLLM_DEFAULT_HOST, FLOWLLM_DEFAULT_PORT + else: + host, port = FLOWLLM_DEFAULT_HOST, FLOWLLM_DEFAULT_PORT + self.base_url = f"http://{host}:{port}" + self.timeout = timeout + + async def _start(self) -> None: + if self.client is None: + self.client = httpx.AsyncClient(base_url=self.base_url, timeout=self.timeout) + + async def _iter_stream_chunks(self, action: str, payload: dict) -> AsyncGenerator[StreamChunk, None]: + """Yield StreamChunks; auto-detects JSON vs SSE.""" + if self.client is None: + raise RuntimeError("Client not initialized. Call _start() first.") + async with self.client.stream("POST", f"/{action}", json=payload) as resp: + resp.raise_for_status() + ctype = resp.headers.get("content-type", "") + if ctype.startswith("text/event-stream"): + async for line in resp.aiter_lines(): + if not line.startswith("data:"): + continue + data_str = line[len("data:") :] + if data_str.strip() == "[DONE]": + return + try: + data = json.loads(data_str) + except json.JSONDecodeError: + continue + chunk = StreamChunk(**data) + if chunk.chunk_type == ChunkEnum.ERROR: + raise RuntimeError(str(chunk.chunk)) + if chunk.done: + return + yield chunk + else: + yield StreamChunk(chunk_type=ChunkEnum.CONTENT, chunk=(await resp.aread()).decode()) + + async def stream_chunks(self, action: str, **kwargs) -> AsyncGenerator[StreamChunk, None]: + """Yield StreamChunks from an action endpoint.""" + async for chunk in self._iter_stream_chunks(action, kwargs): + yield chunk + + async def list_actions(self) -> list[dict]: + if self.client is None: + raise RuntimeError("Client not initialized. Call _start() first.") + resp = await self.client.get("/openapi.json") + resp.raise_for_status() + spec = resp.json() + actions: list[dict] = [] + for path, methods in spec.get("paths", {}).items(): + for method, op in methods.items(): + actions.append({"action": path.lstrip("/"), "method": method.upper(), **op}) + return actions + + @staticmethod + def _format_for_display(text: str) -> str: + """Render JSON as CLI-friendly text; pass through unrecognized payloads.""" + try: + data = json.loads(text) + except (ValueError, json.JSONDecodeError): + return text + if not (isinstance(data, dict) and isinstance(data.get("answer"), str)): + return json.dumps(data, indent=2, ensure_ascii=False) if isinstance(data, (dict, list)) else text + d = dict(data) + answer, success, metadata = d.pop("answer"), d.pop("success", None), d.pop("metadata", None) + parts = [answer] + status_pieces = [] + if success is not None: + status_pieces.append("✅" if success else "❌") + if metadata: + status_pieces.append(json.dumps(metadata, ensure_ascii=False)) + if status_pieces: + parts.append(" ".join(status_pieces)) + if d: + parts.append(json.dumps(d, indent=2, ensure_ascii=False)) + return "\n".join(parts) + + # pylint: disable=invalid-overridden-method + async def _execute(self, action: str, payload: dict) -> AsyncGenerator[str, None]: + async for chunk in self._iter_stream_chunks(action, payload): + chunk_payload = chunk.chunk + text = chunk_payload if isinstance(chunk_payload, str) else json.dumps(chunk_payload, ensure_ascii=False) + yield self._format_for_display(text) + + async def _close(self) -> None: + if self.client is not None: + await self.client.aclose() + self.client = None diff --git a/flowllm/components/client/mcp_client.py b/flowllm/components/client/mcp_client.py new file mode 100644 index 0000000..7e890bf --- /dev/null +++ b/flowllm/components/client/mcp_client.py @@ -0,0 +1,87 @@ +"""MCP client for FlowLLM services.""" + +import json +import os +from collections.abc import AsyncGenerator +from typing import Any + +from fastmcp import Client +from fastmcp.client import SSETransport, StdioTransport, StreamableHttpTransport +from fastmcp.client.client import CallToolResult + +from .base_client import BaseClient +from ..component_registry import R +from ...constants import FLOWLLM_SERVICE_INFO, FLOWLLM_DEFAULT_HOST, FLOWLLM_DEFAULT_PORT + +_TRANSPORT_MAP = {"sse": SSETransport, "stdio": StdioTransport, "streamable-http": StreamableHttpTransport} + + +@R.register("mcp") +class MCPClient(BaseClient): + """MCP client via fastmcp; supports sse, stdio, streamable-http transports.""" + + def __init__( + self, + transport: str | Any = "sse", + host: str | None = None, + port: int | None = None, + timeout: float = 30.0, + **kwargs, + ): + super().__init__(**kwargs) + if isinstance(transport, str) and transport not in _TRANSPORT_MAP: + raise ValueError(f"Unknown transport: {transport!r}, expected one of {list(_TRANSPORT_MAP)}") + if isinstance(transport, str) and transport != "stdio": + if not (host and port): + if service_info := os.environ.get(FLOWLLM_SERVICE_INFO): + try: + data = json.loads(service_info) + host, port = data["host"], data["port"] + except Exception: + self.logger.warning(f"Invalid service info: {service_info}") + host, port = FLOWLLM_DEFAULT_HOST, FLOWLLM_DEFAULT_PORT + else: + host, port = FLOWLLM_DEFAULT_HOST, FLOWLLM_DEFAULT_PORT + self.host, self.port = host, port + self.transport = transport + self.timeout = timeout + + def _build_transport(self): + if not isinstance(self.transport, str): + return self.transport + cls = _TRANSPORT_MAP[self.transport] + if self.transport == "stdio": + return cls(command=self.kwargs.get("command", ""), args=self.kwargs.get("args", [])) + path = "/sse" if self.transport == "sse" else "/mcp" + return cls(url=f"http://{self.host}:{self.port}{path}") + + # pylint: disable=unnecessary-dunder-call + async def _start(self) -> None: + if self.client is None: + self.client = Client(self._build_transport(), timeout=self.timeout) + await self.client.__aenter__() + + # pylint: disable=invalid-overridden-method + async def _execute(self, action: str, payload: dict) -> AsyncGenerator[str, None]: + if self.client is None: + raise RuntimeError("Client not initialized. Call _start() first.") + result: CallToolResult = await self.client.call_tool(action, payload) + yield self._extract_text(result) + + async def list_actions(self) -> list[dict]: + if self.client is None: + raise RuntimeError("Client not initialized. Call _start() first.") + return [tool.model_dump() for tool in await self.client.list_tools()] + + # pylint: disable=unnecessary-dunder-call + async def _close(self) -> None: + if self.client is not None: + await self.client.__aexit__(None, None, None) + self.client = None + + @staticmethod + def _extract_text(result: CallToolResult) -> str: + for block in result.content: + if hasattr(block, "text"): + return block.text + return str(result.content) diff --git a/flowllm/components/component_registry.py b/flowllm/components/component_registry.py new file mode 100644 index 0000000..54c48ef --- /dev/null +++ b/flowllm/components/component_registry.py @@ -0,0 +1,64 @@ +"""Global registry: (ComponentEnum, name) -> component class.""" + +from typing import Callable, TypeVar, cast + +from .base_component import BaseComponent +from ..enumeration import ComponentEnum +from ..utils import get_logger + +T = TypeVar("T", bound=BaseComponent) + + +class ComponentRegistry: + """Two-level registry: component_type -> name -> class. Supports direct call and decorator.""" + + def __init__(self) -> None: + self._registry: dict[ComponentEnum, dict[str, type[BaseComponent]]] = {} + self.logger = get_logger() + + def _do_register(self, cls: type[T], name: str) -> type[T]: + component_type = getattr(cls, "component_type", None) + if not isinstance(component_type, ComponentEnum): + raise TypeError(f"{cls.__name__} must have a ComponentEnum 'component_type' attribute") + if not name: + raise ValueError("Component name cannot be empty") + group = self._registry.setdefault(component_type, {}) + if name in group: + self.logger.warning(f"Component '{name}' already registered for {component_type}, overwriting") + group[name] = cls + return cls + + def register(self, cls_or_name: type[T] | str, name: str | None = None) -> Callable[[type[T]], type[T]] | type[T]: + """Register a class directly or return a decorator.""" + if isinstance(cls_or_name, type): + return self._do_register(cast(type[T], cls_or_name), name if name is not None else cls_or_name.__name__) + if not isinstance(cls_or_name, str): + raise TypeError(f"Expected a class or string, got {type(cls_or_name).__name__}") + registration_name = cls_or_name + + def decorator(decorated_cls: type[T]) -> type[T]: + return self._do_register(decorated_cls, registration_name) + + return decorator + + def get(self, component_type: ComponentEnum, name: str) -> type[BaseComponent] | None: + """Look up a component class by type and name.""" + return self._registry.get(component_type, {}).get(name) + + def get_all(self, component_type: ComponentEnum) -> dict[str, type[BaseComponent]]: + """Return all registered classes for a component type.""" + return dict(self._registry.get(component_type, {})) + + def unregister(self, component_type: ComponentEnum, name: str) -> bool: + """Remove a component; return True if found.""" + if (group := self._registry.get(component_type)) and name in group: + del group[name] + return True + return False + + def clear(self) -> None: + """Remove all registrations.""" + self._registry.clear() + + +R = ComponentRegistry() diff --git a/flowllm/components/embedding_store/__init__.py b/flowllm/components/embedding_store/__init__.py new file mode 100644 index 0000000..40a089a --- /dev/null +++ b/flowllm/components/embedding_store/__init__.py @@ -0,0 +1,6 @@ +"""Embedding store implementations.""" + +from .base_embedding_store import BaseEmbeddingStore +from .local_embedding_store import LocalEmbeddingStore + +__all__ = ["BaseEmbeddingStore", "LocalEmbeddingStore"] diff --git a/flowllm/components/embedding_store/base_embedding_store.py b/flowllm/components/embedding_store/base_embedding_store.py new file mode 100644 index 0000000..ee055d5 --- /dev/null +++ b/flowllm/components/embedding_store/base_embedding_store.py @@ -0,0 +1,44 @@ +"""Base embedding store with abstract interface for caching and retrieval.""" + +from abc import abstractmethod + +import numpy as np + +from ..base_component import BaseComponent +from ...enumeration import ComponentEnum +from ...schema import EmbNode + + +class BaseEmbeddingStore(BaseComponent): + """Abstract embedding store; subclasses implement caching and provider calls.""" + + component_type = ComponentEnum.EMBEDDING_STORE + + def __init__(self, max_batch_size: int = 10, max_input_length: int = 8192, max_retries: int = 3, **kwargs): + super().__init__(**kwargs) + self.max_batch_size = max_batch_size + self.max_input_length = max_input_length + self.max_retries = max_retries + self.is_healthy: bool = True + + @abstractmethod + async def health_check(self, timeout: float = 2.0) -> bool: + """Probe provider; sets and returns is_healthy.""" + + async def get_embedding(self, input_text: str, **kwargs) -> np.ndarray | None: + """Get a single embedding vector.""" + results = await self.get_embeddings([input_text], **kwargs) + return results[0] if results else None + + @abstractmethod + async def get_embeddings(self, input_text: list[str], **kwargs) -> list[np.ndarray | None]: + """Get embeddings for a list of texts.""" + + async def get_node_embeddings(self, nodes: list[EmbNode], **kwargs) -> list[EmbNode]: + """Populate embedding field on EmbNode list.""" + embeddings = await self.get_embeddings([n.text for n in nodes], **kwargs) + if len(embeddings) == len(nodes): + for node, vec in zip(nodes, embeddings): + if vec is not None: + node.embedding = vec + return nodes diff --git a/flowllm/components/embedding_store/local_embedding_store.py b/flowllm/components/embedding_store/local_embedding_store.py new file mode 100644 index 0000000..66411c7 --- /dev/null +++ b/flowllm/components/embedding_store/local_embedding_store.py @@ -0,0 +1,186 @@ +"""Local embedding store with LRU cache and disk persistence.""" + +import asyncio +import hashlib +from collections import OrderedDict +from pathlib import Path + +import numpy as np + +from .base_embedding_store import BaseEmbeddingStore +from ..as_embedding import BaseAsEmbedding +from ..component_registry import R + +Miss = tuple[int, str, str] # (result_index, text, cache_key) + + +@R.register("local") +class LocalEmbeddingStore(BaseEmbeddingStore): + """LRU-cached embedding store with disk persistence; delegates to as_embedding.""" + + def __init__( + self, + as_embedding: str = "default", + max_cache_size: int = 10000, + enable_cache: bool = True, + cache_version: str = "v1", + **kwargs, + ): + super().__init__(**kwargs) + self.as_embedding = self.bind(as_embedding, BaseAsEmbedding, optional=False) + self.max_cache_size = max_cache_size + self.enable_cache = enable_cache + self.cache_version = cache_version + self._cache: OrderedDict[str, np.ndarray] = OrderedDict() + self._key_suffix: bytes = b"" + + @property + def dimensions(self) -> int: + """Embedding vector dimensionality.""" + assert self.as_embedding is not None, "embedding component not bound" + return self.as_embedding.dimensions + + @property + def cache_path(self) -> Path: + """Path to the on-disk cache file.""" + return self.component_metadata_path / f"{self.name}_{self.cache_version}.npz" + + async def _start(self) -> None: + self._key_suffix = f"|{self.dimensions}".encode() + await self.load() + + async def _close(self) -> None: + await self.dump() + + async def health_check(self, timeout: float = 5.0) -> bool: + tag = f"[EMBEDDING HEALTH CHECK] name={self.name}" + try: + result = await asyncio.wait_for(self.as_embedding(["ping"]), timeout=timeout) + if not result or result[0] is None: + raise RuntimeError("empty embedding") + self.is_healthy = True + self.logger.info(f"{tag} -> OK") + except asyncio.TimeoutError: + self.is_healthy = False + self.logger.error(f"{tag} -> FAIL timeout({timeout}s)") + except Exception as e: + self.is_healthy = False + self.logger.error(f"{tag} -> FAIL {type(e).__name__}: {e}") + return self.is_healthy + + async def get_embeddings(self, input_text: list[str], **kwargs) -> list[np.ndarray | None]: + texts = [self._truncate(t) for t in input_text] + results, misses = self._partition_by_cache(texts) + if misses: + await self._fill_misses(misses, results, **kwargs) + return results + + def _truncate(self, text: str) -> str: + return text if len(text) <= self.max_input_length else text[: self.max_input_length] + + def _partition_by_cache(self, texts: list[str]) -> tuple[list[np.ndarray | None], list[Miss]]: + results: list[np.ndarray | None] = [None] * len(texts) + misses: list[Miss] = [] + for idx, text in enumerate(texts): + key = self._cache_key(text) + hit = self._cache_get(key) + if hit is not None: + results[idx] = hit + else: + misses.append((idx, text, key)) + return results, misses + + async def _fill_misses(self, misses: list[Miss], results: list[np.ndarray | None], **kwargs) -> None: + for i in range(0, len(misses), self.max_batch_size): + for idx, key, emb in await self._compute_batch(misses[i : i + self.max_batch_size], **kwargs): + results[idx] = emb + self._cache_put(key, emb) + + async def _compute_batch(self, batch: list[Miss], **kwargs) -> list[tuple[int, str, np.ndarray]]: + texts = [text for _, text, _ in batch] + embeddings = await self._call_with_retry(texts, **kwargs) + if not embeddings or len(embeddings) != len(texts): + return [] + out: list[tuple[int, str, np.ndarray]] = [] + for (idx, _text, key), raw in zip(batch, embeddings): + if raw is not None: + out.append((idx, key, self._normalize_dim(np.asarray(raw, dtype=np.float16)))) + return out + + async def _call_with_retry(self, texts: list[str], **kwargs) -> list[list[float] | None] | None: + for attempt in range(self.max_retries): + try: + result = await self.as_embedding(texts, **kwargs) + if result and len(result) == len(texts): + return result + except (TimeoutError, ConnectionError, OSError): + if attempt < self.max_retries - 1: + await asyncio.sleep(2**attempt) + except Exception: + self.logger.exception("Embedding request failed") + return None + return None + + def _normalize_dim(self, emb: np.ndarray) -> np.ndarray: + if len(emb) == self.dimensions: + return emb + if len(emb) < self.dimensions: + return np.pad(emb, (0, self.dimensions - len(emb))) + return emb[: self.dimensions] + + def _cache_key(self, text: str) -> str: + return hashlib.sha256(text.encode() + self._key_suffix).hexdigest() + + def _cache_get(self, key: str) -> np.ndarray | None: + if not self.enable_cache or key not in self._cache: + return None + self._cache.move_to_end(key) + return self._cache[key] + + def _cache_put(self, key: str, embedding: np.ndarray) -> None: + if not self.enable_cache or self.max_cache_size <= 0 or len(embedding) != self.dimensions: + return + cache = self._cache + if key in cache: + cache.move_to_end(key) + cache[key] = embedding + return + if len(cache) >= self.max_cache_size: + cache.popitem(last=False) + cache[key] = embedding + + async def load(self) -> None: + self._cache.clear() + if not self.enable_cache or not self.cache_path.exists(): + return + await asyncio.to_thread(self._load_sync) + + def _load_sync(self) -> None: + try: + with np.load(self.cache_path) as data: + for key, emb in zip(data["keys"], data["embeddings"]): + if len(emb) != self.dimensions: + continue + if len(self._cache) >= self.max_cache_size: + break + self._cache[str(key)] = emb.astype(np.float16) + except Exception: + self.logger.exception("Failed to load embedding cache, removing") + self.cache_path.unlink(missing_ok=True) + return + self.logger.info(f"Loaded {len(self._cache)} embeddings from {self.cache_path}") + + async def dump(self) -> None: + if not self.enable_cache or not self._cache: + return + await asyncio.to_thread(self._dump_sync) + + def _dump_sync(self) -> None: + self.cache_path.parent.mkdir(parents=True, exist_ok=True) + keys = np.array(list(self._cache.keys()), dtype=str) + embeddings = np.stack(list(self._cache.values())) + try: + np.savez(self.cache_path, keys=keys, embeddings=embeddings) + self.logger.info(f"Saved {len(self._cache)} embeddings to {self.cache_path}") + except Exception: + self.logger.exception("Failed to save embedding cache") diff --git a/flowllm/components/job/__init__.py b/flowllm/components/job/__init__.py new file mode 100644 index 0000000..5d923d6 --- /dev/null +++ b/flowllm/components/job/__init__.py @@ -0,0 +1,13 @@ +"""Job components for executing workflows.""" + +from .background_job import BackgroundJob +from .base_job import BaseJob +from .cron_job import CronJob +from .stream_job import StreamJob + +__all__ = [ + "BackgroundJob", + "BaseJob", + "CronJob", + "StreamJob", +] diff --git a/flowllm/components/job/background_job.py b/flowllm/components/job/background_job.py new file mode 100644 index 0000000..04528c5 --- /dev/null +++ b/flowllm/components/job/background_job.py @@ -0,0 +1,95 @@ +"""Long-running background job with supervisor and exponential backoff.""" + +import asyncio +import contextlib +import random +import time + +from .base_job import BaseJob +from ..component_registry import R +from ..runtime_context import RuntimeContext +from ...schema import Response + + +@R.register("background") +class BackgroundJob(BaseJob): + """Long-running job with supervisor restart on crash (exponential backoff + jitter).""" + + def __init__( + self, + supervisor: bool = True, + backoff_base: float = 1.0, + backoff_cap: float = 60.0, + close_timeout: float = 5.0, + attempt_reset_after: float = 60.0, + **kwargs, + ): + kwargs.pop("enable_serve", None) + super().__init__(enable_serve=False, **kwargs) + self.supervisor: bool = supervisor + self.backoff_base: float = backoff_base + self.backoff_cap: float = backoff_cap + self.close_timeout: float = close_timeout + self.attempt_reset_after: float = attempt_reset_after + self._stop_event: asyncio.Event | None = None + self._task: asyncio.Task | None = None + + async def _start(self) -> None: + await super()._start() + self._stop_event = asyncio.Event() + self._task = asyncio.create_task(self._run_with_supervisor()) + + async def _close(self) -> None: + if self._stop_event is not None: + self._stop_event.set() + await self._shutdown_task() + await super()._close() + + async def _shutdown_task(self) -> None: + if self._task is None: + return + try: + await asyncio.wait_for(asyncio.shield(self._task), timeout=self.close_timeout) + except asyncio.TimeoutError: + self._task.cancel() + with contextlib.suppress(BaseException): + await self._task + except Exception: + self.logger.exception(f"Background task '{self.name}' raised during close") + self._task = None + + def _backoff_delay(self, attempt: int) -> float: + capped = min(self.backoff_base * (2**attempt), self.backoff_cap) + return min(capped * (0.5 + random.random()), self.backoff_cap) + + async def _wait_or_stop(self, delay: float) -> None: + assert self._stop_event is not None + try: + await asyncio.wait_for(self._stop_event.wait(), timeout=delay) + except asyncio.TimeoutError: + pass + + async def _run_with_supervisor(self) -> None: + assert self._stop_event is not None + attempt = 0 + while not self._stop_event.is_set(): + started_at = time.monotonic() + try: + await self() + return + except Exception as e: + if not self.supervisor: + raise + if time.monotonic() - started_at >= self.attempt_reset_after: + attempt = 0 + delay = self._backoff_delay(attempt) + self.logger.exception(f"job body crashed, restart in {delay:.2f}s error={e}") + attempt += 1 + await self._wait_or_stop(delay) + + async def __call__(self, **kwargs) -> Response: + merged = {**self.kwargs, **kwargs} + context = RuntimeContext(stop_event=self._stop_event, **merged) + for step in self._build_steps(): + await step(context) + return context.response diff --git a/flowllm/components/job/base_job.py b/flowllm/components/job/base_job.py new file mode 100644 index 0000000..c9c9999 --- /dev/null +++ b/flowllm/components/job/base_job.py @@ -0,0 +1,68 @@ +"""Base job component for sequential step execution.""" + +from typing import TYPE_CHECKING + +from ..base_component import BaseComponent +from ..component_registry import R +from ..runtime_context import RuntimeContext +from ...enumeration import ComponentEnum +from ...schema import ComponentConfig, Response + +if TYPE_CHECKING: + from ...steps import BaseStep + + +@R.register("base") +class BaseJob(BaseComponent): + """Job that executes steps sequentially and returns a Response.""" + + component_type = ComponentEnum.JOB + + def __init__( + self, + description: str = "", + parameters: dict | None = None, + steps: list[ComponentConfig | dict] | None = None, + enable_serve: bool = True, + **kwargs, + ): + super().__init__(**kwargs) + self.description = description + self.parameters = parameters or {} + self.step_configs = steps or [] + self.enable_serve = enable_serve + self.step_specs: list[tuple[type["BaseStep"], dict]] = [] + + async def _start(self) -> None: + if self.app_context is None: + raise RuntimeError(f"app_context must be provided for job '{self.name}'") + self.step_specs = [self._resolve_step(raw) for raw in self.step_configs] + + async def _close(self) -> None: + self.step_specs.clear() + + def _resolve_step(self, raw: ComponentConfig | dict) -> tuple[type["BaseStep"], dict]: + config = raw if isinstance(raw, ComponentConfig) else ComponentConfig(**raw) + if not config.backend: + raise ValueError("Step is missing the required 'backend' field") + step_cls = R.get(ComponentEnum.STEP, config.backend) + if not step_cls: + raise ValueError(f"Unregistered backend '{config.backend}' of type '{ComponentEnum.STEP}'") + params = config.model_dump() + params["app_context"] = self.app_context + return step_cls, params + + def _build_steps(self) -> list["BaseStep"]: + return [step_cls(**dict(params)) for step_cls, params in self.step_specs] + + async def __call__(self, **kwargs) -> Response: + merged = {**self.kwargs, **kwargs} + context = RuntimeContext(**merged) + try: + for step in self._build_steps(): + await step(context) + except Exception as e: + self.logger.exception(f"Failed to execute job: {e}") + context.response.success = False + context.response.answer = str(e) + return context.response diff --git a/flowllm/components/job/cron_job.py b/flowllm/components/job/cron_job.py new file mode 100644 index 0000000..de138e7 --- /dev/null +++ b/flowllm/components/job/cron_job.py @@ -0,0 +1,50 @@ +"""Cron-scheduled background job.""" + +import datetime +from zoneinfo import ZoneInfo + +from .background_job import BackgroundJob +from ..component_registry import R +from ..runtime_context import RuntimeContext +from ...schema import Response + + +@R.register("cron") +class CronJob(BackgroundJob): + """Run steps on a cron schedule.""" + + def __init__(self, cron: str, **kwargs): + super().__init__(**kwargs) + self.cron_expr = cron + + async def _start(self) -> None: + from croniter import croniter + + if not croniter.is_valid(self.cron_expr): + raise ValueError(f"Invalid cron expression: {self.cron_expr}") + await super()._start() + + def _next_fire_delay(self) -> float: + from croniter import croniter + + tz_name = self.app_context.app_config.timezone if self.app_context else None + now = datetime.datetime.now(ZoneInfo(tz_name)) if tz_name else datetime.datetime.now() + return max(0.0, (croniter(self.cron_expr, now).get_next(datetime.datetime) - now).total_seconds()) + + async def _execute_steps(self) -> Response: + context = RuntimeContext(**self.kwargs) + for step in self._build_steps(): + await step(context) + return context.response + + async def __call__(self, **kwargs) -> Response: + assert self._stop_event is not None + while not self._stop_event.is_set(): + await self._wait_or_stop(self._next_fire_delay()) + if self._stop_event.is_set(): + break + try: + await self._execute_steps() + except Exception as exc: + self.logger.exception(f"Cron job '{self.name}' failed: {exc}") + return Response(success=True) diff --git a/flowllm/components/job/stream_job.py b/flowllm/components/job/stream_job.py new file mode 100644 index 0000000..21b2c52 --- /dev/null +++ b/flowllm/components/job/stream_job.py @@ -0,0 +1,21 @@ +"""Streaming job for real-time output delivery.""" + +from .base_job import BaseJob +from ..component_registry import R +from ..runtime_context import RuntimeContext +from ...enumeration import ChunkEnum + + +@R.register("stream") +class StreamJob(BaseJob): + """Job that streams chunks to a queue instead of returning a Response.""" + + async def __call__(self, **kwargs) -> None: + merged = {**self.kwargs, **kwargs} + context = RuntimeContext(**merged) + try: + for step in self._build_steps(): + await step(context) + except Exception as e: + await context.add_stream_string(str(e), ChunkEnum.ERROR) + await context.add_stream_done() # always emit DONE so consumers can detach diff --git a/flowllm/components/prompt_handler.py b/flowllm/components/prompt_handler.py new file mode 100644 index 0000000..b3a0bf3 --- /dev/null +++ b/flowllm/components/prompt_handler.py @@ -0,0 +1,107 @@ +"""Prompt template loader/formatter with conditional-line and i18n support.""" + +import inspect +import json +import re +from pathlib import Path + +import yaml + +_FLAG_PATTERN = re.compile(r"^\[(\w+)]") # leading "[flag]" tag + + +class PromptHandler: + """Load prompts from YAML/JSON; render with [flag] filtering and i18n key fallback.""" + + _SUPPORTED_EXTENSIONS = {".yaml", ".yml", ".json"} + + def __init__(self, language: str = "", **kwargs): + self.data: dict[str, str] = {k: v for k, v in kwargs.items() if isinstance(v, str)} + self.language: str = language.strip() + + def load_prompt_by_file( + self, + prompt_file_path: str | Path | None = None, + overwrite: bool = True, + ) -> "PromptHandler": + """Load prompts from a YAML/JSON file.""" + if prompt_file_path is None: + return self + path = Path(prompt_file_path) + if not path.exists() or path.suffix.lower() not in self._SUPPORTED_EXTENSIONS: + return self + return self.load_prompt_dict(self._parse_prompt_file(path), overwrite) + + @staticmethod + def _parse_prompt_file(path: Path) -> dict | None: + try: + with path.open(encoding="utf-8") as f: + return yaml.safe_load(f) if path.suffix.lower() in (".yaml", ".yml") else json.load(f) + except (json.JSONDecodeError, yaml.YAMLError, OSError): + return None + + def load_prompt_by_class(self, cls: type, overwrite: bool = True) -> "PromptHandler": + """Load from .yaml next to cls.""" + try: + base_path = Path(inspect.getfile(cls)).with_suffix("") + except (TypeError, OSError): + return self + for ext in (".yaml", ".yml"): + candidate = base_path.with_suffix(ext) + if candidate.exists(): + return self.load_prompt_by_file(candidate, overwrite) + return self + + def load_prompt_dict(self, prompt_dict: dict | None = None, overwrite: bool = True) -> "PromptHandler": + """Merge a dict of prompt strings.""" + if not isinstance(prompt_dict, dict): + return self + for key, value in prompt_dict.items(): + if isinstance(value, str) and (overwrite or key not in self.data): + self.data[key] = value + return self + + def _candidate_keys(self, prompt_name: str) -> tuple[str, ...]: + if self.language: + return (f"{prompt_name}_{self.language}", prompt_name) + return (prompt_name,) + + def get_prompt(self, prompt_name: str) -> str: + """Retrieve a prompt by name with i18n fallback.""" + for key in self._candidate_keys(prompt_name): + if key in self.data: + return self.data[key].strip() + raise KeyError(f"Prompt '{prompt_name}' not found. Available: {list(self.data.keys())[:10]}") + + def has_prompt(self, prompt_name: str) -> bool: + """Check if a prompt exists.""" + return any(k in self.data for k in self._candidate_keys(prompt_name)) + + def list_prompts(self, language_filter: str | None = None) -> list[str]: + """List prompt keys, optionally filtered by language.""" + if language_filter is None: + return list(self.data.keys()) + suffix = f"_{language_filter.strip()}" + return [k for k in self.data if k.endswith(suffix)] + + def prompt_format(self, prompt_name: str, **kwargs) -> str: + """Render: strip inactive [flag] lines, then str.format.""" + prompt = self.get_prompt(prompt_name) + flags = {k: v for k, v in kwargs.items() if isinstance(v, bool)} + formats = {k: v for k, v in kwargs.items() if not isinstance(v, bool)} + if flags: + prompt = self._apply_flag_filter(prompt, flags) + return prompt.format(**formats).strip() if formats else prompt + + @staticmethod + def _apply_flag_filter(prompt: str, flags: dict[str, bool]) -> str: + lines = [] + for line in prompt.split("\n"): + active_flags = _FLAG_PATTERN.findall(line) + cleaned = _FLAG_PATTERN.sub("", line).lstrip() + if not active_flags or any(flags.get(f, False) for f in active_flags): + lines.append(cleaned) + return "\n".join(lines) + + def __repr__(self) -> str: + return f"PromptHandler(language='{self.language}', num_prompts={len(self.data)})" diff --git a/flowllm/components/runtime_context.py b/flowllm/components/runtime_context.py new file mode 100644 index 0000000..b8fc786 --- /dev/null +++ b/flowllm/components/runtime_context.py @@ -0,0 +1,79 @@ +"""Per-request runtime context shared across steps and jobs.""" + +import asyncio + +from ..enumeration import ChunkEnum +from ..schema import Response, StreamChunk + + +class RuntimeContext: + """Scratch space for a single execution: response, stream queue, stop event, and data dict.""" + + def __init__( + self, + response: Response | None = None, + stream_queue: asyncio.Queue | None = None, + stop_event: asyncio.Event | None = None, + **kwargs, + ): + self.response: Response = response or Response() + self.stream_queue: asyncio.Queue | None = stream_queue + self.stop_event: asyncio.Event | None = stop_event + self.data: dict = kwargs + + def get(self, key: str, default=None): + """Get a value from the data dict.""" + return self.data.get(key, default) + + def update(self, data: dict) -> "RuntimeContext": + """Merge data into the context.""" + self.data.update(data) + return self + + def __getitem__(self, key: str): + return self.data[key] + + def __setitem__(self, key: str, value): + self.data[key] = value + + def __delitem__(self, key: str): + del self.data[key] + + def __contains__(self, key: str) -> bool: + return key in self.data + + @classmethod + def from_context(cls, context: "RuntimeContext | None" = None, **kwargs) -> "RuntimeContext": + """Reuse context (merging kwargs) or create fresh.""" + if context is None: + return cls(**kwargs) + return context.update(kwargs) + + @property + def stream(self) -> bool: + """Whether streaming is enabled.""" + return self.stream_queue is not None + + async def _enqueue(self, chunk: StreamChunk) -> None: + if self.stream_queue is None: + raise RuntimeError("Stream queue not initialized") + await self.stream_queue.put(chunk) + + async def add_stream_string(self, chunk: str, chunk_type: ChunkEnum) -> "RuntimeContext": + """Enqueue a string chunk to the stream.""" + await self._enqueue(StreamChunk(chunk_type=chunk_type, chunk=chunk)) + return self + + async def add_stream_done(self) -> "RuntimeContext": + """Enqueue a done marker to the stream.""" + await self._enqueue(StreamChunk(chunk_type=ChunkEnum.DONE, chunk="", done=True)) + return self + + def apply_mapping(self, mapping: dict[str, str]) -> "RuntimeContext": + """Copy data[source] to data[target] for each source:target pair.""" + if not mapping: + return self + for source, target in mapping.items(): + if source in self.data: + self.data[target] = self.data[source] + return self diff --git a/flowllm/components/service/__init__.py b/flowllm/components/service/__init__.py new file mode 100644 index 0000000..41d8f7b --- /dev/null +++ b/flowllm/components/service/__init__.py @@ -0,0 +1,11 @@ +"""Service components for exposing jobs via different protocols.""" + +from .base_service import BaseService +from .http_service import HttpService +from .mcp_service import MCPService + +__all__ = [ + "BaseService", + "HttpService", + "MCPService", +] diff --git a/flowllm/components/service/base_service.py b/flowllm/components/service/base_service.py new file mode 100644 index 0000000..313c5d6 --- /dev/null +++ b/flowllm/components/service/base_service.py @@ -0,0 +1,70 @@ +"""Base class for services that expose jobs over a network protocol.""" + +import json +import os +from abc import abstractmethod +from contextlib import asynccontextmanager +from typing import TYPE_CHECKING + +from ..base_component import BaseComponent +from ..job.base_job import BaseJob +from ...constants import FLOWLLM_SERVICE_INFO +from ...enumeration import ComponentEnum + +if TYPE_CHECKING: + from ...application import Application + + +class BaseService(BaseComponent): + """Skeleton for services (HTTP, MCP, ...) that turn jobs into endpoints.""" + + component_type = ComponentEnum.SERVICE + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.service = None + + @abstractmethod + def build_service(self, app: "Application") -> None: + """Build the underlying service instance.""" + + @abstractmethod + def add_job(self, job: BaseJob) -> bool: + """Register a job; return True if exposed, False if skipped.""" + + @abstractmethod + def start_service(self, app: "Application") -> None: + """Start the service and begin serving.""" + + def _lifespan(self, app: "Application", host: str, port: int): + """Async lifespan bracketing server with app start/close.""" + + @asynccontextmanager + async def lifespan(_): + await app.start() + service_info = json.dumps({"host": host, "port": port}) + os.environ[FLOWLLM_SERVICE_INFO] = service_info + self.logger.info(f"{self.name} started: {FLOWLLM_SERVICE_INFO}={service_info}") + yield + await app.close() + + return lifespan + + def add_jobs(self, app: "Application") -> None: + """Register all servable jobs from the app.""" + for name, job in app.context.jobs.items(): + if not job.enable_serve: + continue + try: + if self.add_job(job): + self.logger.info(f"Added job: {name}") + else: + self.logger.warning(f"Skipped job: {name}") + except Exception as e: + self.logger.error(f"Failed to add job {name}: {e}") + + def run_app(self, app: "Application") -> None: + """Build, register jobs, and start the service.""" + self.build_service(app) + self.add_jobs(app) + self.start_service(app) diff --git a/flowllm/components/service/http_service.py b/flowllm/components/service/http_service.py new file mode 100644 index 0000000..bbd2fb6 --- /dev/null +++ b/flowllm/components/service/http_service.py @@ -0,0 +1,83 @@ +"""HTTP service: exposes jobs as FastAPI endpoints (JSON or SSE).""" + +import asyncio +import warnings +from collections.abc import AsyncGenerator +from typing import TYPE_CHECKING + +import uvicorn +from fastapi import FastAPI +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import StreamingResponse + +from .base_service import BaseService +from ..component_registry import R +from ..job import BaseJob, StreamJob +from ...constants import FLOWLLM_DEFAULT_HOST, FLOWLLM_DEFAULT_PORT +from ...schema import Request, Response +from ...utils import execute_stream_task + +if TYPE_CHECKING: + from ...application import Application + +_WEBSOCKET_DEPRECATION_PATTERNS = ( + r".*websockets\.legacy is deprecated.*", + r".*WebSocketServerProtocol is deprecated.*", +) + + +@R.register("http") +class HttpService(BaseService): + """Map jobs to JSON POST endpoints; StreamJobs to SSE endpoints.""" + + def __init__(self, host: str = FLOWLLM_DEFAULT_HOST, port: int = FLOWLLM_DEFAULT_PORT, **kwargs): + super().__init__(**kwargs) + self.host: str = host + self.port: int = port + + def build_service(self, app: "Application") -> None: + self.service = FastAPI(title=app.config.app_name, lifespan=self._lifespan(app, self.host, self.port)) + self.service.add_middleware( + CORSMiddleware, # type: ignore[arg-type] + allow_origins=["*"], + allow_credentials=False, + allow_methods=["*"], + allow_headers=["*"], + ) + + def add_job(self, job: BaseJob) -> bool: + if isinstance(job, StreamJob): + self._add_stream_job(job) + else: + self._add_json_job(job) + return True + + def start_service(self, app: "Application") -> None: + for pattern in _WEBSOCKET_DEPRECATION_PATTERNS: + warnings.filterwarnings("ignore", category=DeprecationWarning, message=pattern) + uvicorn.run(self.service, host=self.host, port=self.port, **self.kwargs) + + def _add_json_job(self, job: BaseJob) -> None: + async def endpoint(request: Request) -> Response: + return await job(**request.model_dump(exclude_none=True)) + + self.service.post(f"/{job.name}", response_model=Response, description=job.description)(endpoint) + + def _add_stream_job(self, job: StreamJob) -> None: + async def endpoint(request: Request) -> StreamingResponse: + stream_queue: asyncio.Queue = asyncio.Queue() + task = asyncio.create_task(job(stream_queue=stream_queue, **request.model_dump(exclude_none=True))) + + async def body() -> AsyncGenerator[bytes, None]: + async for chunk in execute_stream_task( + stream_queue=stream_queue, + task=task, + task_name=job.name, + output_format="bytes", + ): + assert isinstance(chunk, bytes) + yield chunk + + return StreamingResponse(body(), media_type="text/event-stream") + + self.service.post(f"/{job.name}")(endpoint) diff --git a/flowllm/components/service/mcp_service.py b/flowllm/components/service/mcp_service.py new file mode 100644 index 0000000..1dd7b55 --- /dev/null +++ b/flowllm/components/service/mcp_service.py @@ -0,0 +1,58 @@ +"""MCP service: expose jobs as MCP tools.""" + +from typing import TYPE_CHECKING + +from fastmcp import FastMCP +from fastmcp.server.server import Transport +from fastmcp.tools import FunctionTool + +from .base_service import BaseService +from ..component_registry import R +from ..job import BaseJob, StreamJob +from ...constants import FLOWLLM_DEFAULT_HOST, FLOWLLM_DEFAULT_PORT + +if TYPE_CHECKING: + from ...application import Application + + +@R.register("mcp") +class MCPService(BaseService): + """Expose non-stream jobs as MCP tools over stdio, SSE, or streamable-http.""" + + def __init__( + self, + transport: Transport = "sse", + host: str = FLOWLLM_DEFAULT_HOST, + port: int = FLOWLLM_DEFAULT_PORT, + **kwargs, + ): + super().__init__(**kwargs) + self.transport: Transport = transport + self.host: str = host + self.port: int = port + + def build_service(self, app: "Application") -> None: + self.service = FastMCP(name=app.config.app_name, lifespan=self._lifespan(app, self.host, self.port)) + + def add_job(self, job: BaseJob) -> bool: + if isinstance(job, StreamJob): + return False + + async def execute_tool(**kwargs): + return (await job(**kwargs)).answer + + self.service.add_tool( + FunctionTool( + name=job.name, + description=job.description, + fn=execute_tool, + parameters=job.parameters or {}, + ), + ) + return True + + def start_service(self, app: "Application") -> None: + transport_kwargs: dict = {} + if self.transport != "stdio": + transport_kwargs.update(host=self.host, port=self.port) + self.service.run(transport=self.transport, show_banner=False, **transport_kwargs) diff --git a/flowllm/config/__init__.py b/flowllm/config/__init__.py index 5a785ef..c290318 100644 --- a/flowllm/config/__init__.py +++ b/flowllm/config/__init__.py @@ -1,14 +1,8 @@ -"""Configuration module for FlowLLM. +"""Config""" -This module provides configuration parsing capabilities for the FlowLLM framework. -It includes: - -- ConfigParser: Configuration parser class that extends PydanticConfigParser - to provide configuration parsing with awareness of the current module location -""" - -from .config_parser import ConfigParser +from .config_parser import parse_args, resolve_app_config __all__ = [ - "ConfigParser", + "parse_args", + "resolve_app_config", ] diff --git a/flowllm/config/config_parser.py b/flowllm/config/config_parser.py index 3eee1b0..1a36683 100644 --- a/flowllm/config/config_parser.py +++ b/flowllm/config/config_parser.py @@ -1,14 +1,201 @@ -"""Configuration parser module for FlowLLM.""" +"""YAML config parser with CLI argument overrides.""" -from ..core.utils import PydanticConfigParser +import json +import os +import re +from pathlib import Path +from typing import Any +import yaml -class ConfigParser(PydanticConfigParser): - """ - Configuration parser for FlowLLM framework. +_CONFIG_DIR = Path(__file__).parent +_SUPPORTED_EXTS = (".yaml", ".yml", ".json") +_ENV_VAR_RE = re.compile(r"\$\{([A-Za-z_][A-Za-z0-9_]*)(?::-([^}]*))?}") +_LEADING_ZERO_RE = re.compile(r"^-?0\d") - Extends PydanticConfigParser to provide configuration parsing capabilities - with awareness of the current module location. - """ - current_file: str = __file__ +def _repl(m: re.Match) -> str: + name: str = m.group(1) + default: str | None = m.group(2) + v = os.environ.get(name) + if v is None: + if default is not None: + return default + raise ValueError(f"Config references undefined env var: {name}") + return v + + +def _expand_env_vars(value: Any) -> Any: + """Recursively expand `${VAR}` / `${VAR:-default}` in strings.""" + if isinstance(value, str): + expanded = _ENV_VAR_RE.sub(_repl, value) + return _convert_value(expanded) if expanded != value else value + if isinstance(value, dict): + return {k: _expand_env_vars(v) for k, v in value.items()} + if isinstance(value, list): + return [_expand_env_vars(v) for v in value] + return value + + +def _discover_configs() -> dict[str, Path]: + """Map config file stems to paths in the config directory.""" + discovered: dict[str, Path] = {} + if _CONFIG_DIR.is_dir(): + files = sorted( + (p for p in _CONFIG_DIR.iterdir() if p.is_file() and p.suffix in _SUPPORTED_EXTS), + key=lambda p: (_SUPPORTED_EXTS.index(p.suffix), p.name), + ) + for p in files: + discovered.setdefault(p.stem, p) + return discovered + + +_CONFIG_REGISTRY = _discover_configs() + + +def parse_dot_notation(dot_list: list[str]) -> dict: + """Parse "key.subkey=value" strings into a nested dict.""" + result: dict = {} + for item in dot_list: + if "=" not in item: + raise ValueError(f"Invalid dot notation format (missing '='): {item}") + key_path, value_str = item.split("=", 1) + keys = key_path.split(".") + if not key_path or any(not key for key in keys): + raise ValueError(f"Invalid dot notation key: {key_path!r}") + current = result + for key in keys[:-1]: + if key in current and not isinstance(current[key], dict): + raise ValueError(f"Cannot set nested key '{key_path}': '{key}' is already a value") + current = current.setdefault(key, {}) + last_key = keys[-1] + if last_key in current and isinstance(current[last_key], dict): + raise ValueError(f"Cannot overwrite nested dict at '{key_path}' with scalar value") + current[last_key] = _convert_value(value_str) + return result + + +def _convert_value(value_str: str) -> Any: + """Convert string to bool/int/float/JSON, preserving leading-zero strings.""" + s = value_str.strip() + lower = s.lower() + + if lower in ("none", "null"): + return None + if lower == "true": + return True + if lower == "false": + return False + + if not _LEADING_ZERO_RE.match(s): + for converter in (int, float): + try: + return converter(s) + except ValueError: + continue + + try: + return json.loads(s) + except (ValueError, json.JSONDecodeError): + pass + + return s + + +def _load_config(name_or_path: str, encoding: str = "utf-8") -> dict: + """Load a config by registry name or file path.""" + if name_or_path in _CONFIG_REGISTRY: + return _read_config_file(_CONFIG_REGISTRY[name_or_path], encoding) + + p = Path(name_or_path) + if p.suffix in _SUPPORTED_EXTS: + candidates = [p] + if not p.is_absolute(): + candidates.append(_CONFIG_DIR / p) + for candidate in candidates: + if candidate.exists(): + return _read_config_file(candidate, encoding) + raise FileNotFoundError(f"Config file not found: {p}") + + known = ", ".join(sorted(_CONFIG_REGISTRY)) if _CONFIG_REGISTRY else "none" + raise FileNotFoundError(f"Config file not found: {name_or_path}. Available: {known}") + + +def _read_config_file(path: Path, encoding: str = "utf-8") -> dict: + """Read and env-expand a YAML or JSON config file.""" + with path.open(encoding=encoding) as f: + if path.suffix == ".json": + result = json.load(f) + else: + result = yaml.safe_load(f) + if result is None: + return {} + if not isinstance(result, dict): + raise ValueError(f"Config root must be a mapping/object: {path}") + return _expand_env_vars(result) + + +def _deep_merge(base: dict, update: dict) -> dict: + """Recursively merge update into base.""" + result = base.copy() + for k, v in update.items(): + if k in result and isinstance(result[k], dict) and isinstance(v, dict): + result[k] = _deep_merge(result[k], v) + else: + result[k] = v + return result + + +def _strip_arg_dashes(arg: str) -> str: + """Strip one leading `--` or `-` prefix.""" + if arg.startswith("--"): + return arg[2:] + if arg.startswith("-"): + return arg[1:] + return arg + + +def parse_args(*args) -> tuple[str, dict]: + """Parse CLI args into (action, key=value dict).""" + if not args: + raise ValueError("No arguments provided") + + first = _strip_arg_dashes(args[0]) + if "=" in first: + raise ValueError(f"First argument must be action, got: {args[0]}") + + kvs: list[str] = [] + for raw in args[1:]: + arg = _strip_arg_dashes(raw) + if "=" in arg: + kvs.append(arg) + else: + raise ValueError(f"Invalid argument format (expected key=value): {raw}") + + parsed = parse_dot_notation(kvs) if kvs else {} + return first, parsed + + +def resolve_app_config(**kwargs) -> dict: + """Load config file and deep-merge with kwargs overrides.""" + from ..utils import get_logger + + logger = get_logger() + configs: list[dict] = [] + + config_value = kwargs.get("config") + if isinstance(config_value, str): + kwargs.pop("config") + logger.info(f"Loading config: {config_value}") + configs.append(_load_config(config_value)) + elif "default" in _CONFIG_REGISTRY: + logger.info("No config specified, loading 'default'") + configs.append(_load_config("default")) + + configs.append(kwargs) + + merged: dict = {} + for cfg in configs: + merged = _deep_merge(merged, cfg) + + return merged diff --git a/flowllm/config/default.yaml b/flowllm/config/default.yaml index 6d28ae7..74e936e 100644 --- a/flowllm/config/default.yaml +++ b/flowllm/config/default.yaml @@ -1,143 +1,141 @@ -backend: http -thread_pool_max_workers: 128 +service: + backend: http + +jobs: + version: + backend: base + description: "return flowllm package version" + parameters: + type: object + properties: { } + steps: + - backend: version_step + + health_check: + backend: base + description: "return a concise health-check snapshot of flowllm components" + parameters: + type: object + properties: { } + steps: + - backend: health_check_step + + help: + backend: base + description: "list all registered jobs with their metadata" + parameters: + type: object + properties: { } + steps: + - backend: help_step + + demo: + backend: base + description: "demo job description" + parameters: + type: object + properties: + query: + type: string + description: "query" + min_score: + type: number + description: "min score" + default: 0.5 + required: + - query + steps: + - backend: demo_echo_step1 + - backend: demo_echo_step2 + + add: + backend: base + description: "add two numbers" + parameters: + type: object + properties: + a: + type: number + description: "first addend" + b: + type: number + description: "second addend" + required: + - a + - b + steps: + - backend: add_step + + stream_demo: + backend: stream + description: "stream demo job: repeat query 10x and stream char-by-char" + parameters: + type: object + properties: + query: + type: string + description: "query to echo" + repeat: + type: integer + description: "number of times to repeat the query" + default: 10 + interval: + type: number + description: "seconds between chunks" + default: 0.1 + required: + - query + steps: + - backend: stream_demo_step1 + - backend: stream_demo_step2 + +components: + as_embedding: + default: + backend: ${EMBEDDING_BACKEND:-openai} + model: ${EMBEDDING_MODEL_NAME:-text-embedding-v4} + credential: + api_key: ${EMBEDDING_API_KEY:-} + base_url: ${EMBEDDING_BASE_URL:-https://dashscope.aliyuncs.com/compatible-mode/v1} + parameters: + dimensions: 1024 + + embedding_store: + default: + backend: local + as_embedding: default + + as_llm: + default: + backend: ${LLM_BACKEND:-openai} + model: ${LLM_MODEL_NAME:-qwen3.7-plus} + stream: true + context_size: 200000 + max_retries: 3 + credential: + api_key: ${LLM_API_KEY:-} + base_url: ${LLM_BASE_URL:-} + parameters: + max_tokens: 65536 + thinking_enable: false + + agent_wrapper: + default: + backend: agentscope + as_llm: default + permission_mode: bypass + react_config: + max_iters: 30 + context_config: + trigger_ratio: 0.8 + reserve_ratio: 0.1 + tool_result_limit: 50000 + model_config: + max_retries: 1 + claude_code: + backend: claude_code + model: ${CLAUDE_CODE_MODEL_NAME:-glm-5.1} + api_key: ${CLAUDE_CODE_API_KEY:-} + base_url: ${CLAUDE_CODE_BASE_URL:-https://dashscope.aliyuncs.com/apps/anthropic} + permission_mode: bypassPermissions -mcp: - transport: sse - host: "0.0.0.0" - port: 8001 - -http: - host: "0.0.0.0" - port: 8002 - - -flow: - demo_http_flow: - flow_content: GenSystemPromptOp() >> ChatOp() - description: "ai chat assistant" # Optional - input_schema: # Optional - query: - type: string - description: "user query" - required: true - - demo_stream_http_flow: - flow_content: GenSystemPromptOp() >> StreamChatOp() - stream: true - description: "ai chat assistant" # Optional - input_schema: # Optional - query: - type: string - description: "user query" - required: true - - demo_mcp_flow: - flow_content: MockSearchOp() - description: "search results for a given query." - input_schema: - query: - type: string - description: "user query" - required: true - - skill_agent: - flow_content: SkillAgentOp() << [LoadSkillOp(), ReadReferenceFileOp(), RunShellCommandOp()] - -llm: - default: - backend: openai_compatible - model_name: qwen3-30b-a3b-instruct-2507 - params: - temperature: 0.6 - token_count: # Optional - model_name: Qwen/Qwen3-30B-A3B-Instruct-2507 - backend: hf - params: - use_mirror: true - - qwen3_30b_instruct: - backend: openai_compatible - model_name: qwen3-30b-a3b-instruct-2507 - token_count: # Optional - model_name: Qwen/Qwen3-30B-A3B-Instruct-2507 - backend: hf - params: - use_mirror: true - - qwen3_30b_thinking: - backend: openai_compatible - model_name: qwen3-30b-a3b-thinking-2507 - token_count: # Optional - model_name: Qwen/Qwen3-30B-A3B-Thinking-2507 - backend: hf - params: - use_mirror: true - - qwen3_235b_instruct: - backend: openai_compatible - model_name: qwen3-235b-a22b-instruct-2507 - token_count: # Optional - model_name: Qwen/Qwen3-235B-A22B-Instruct-2507 - backend: hf - params: - use_mirror: true - - qwen3_235b_thinking: - backend: openai_compatible - model_name: qwen3-235b-a22b-thinking-2507 - token_count: # Optional - model_name: Qwen/Qwen3-235B-A22B-Thinking-2507 - backend: hf - params: - use_mirror: true - - qwen3_80b_instruct: - backend: openai_compatible - model_name: qwen3-next-80b-a3b-instruct - token_count: # Optional - model_name: Qwen/Qwen3-Next-80B-A3B-Instruct - backend: hf - params: - use_mirror: true - - qwen3_80b_thinking: - backend: openai_compatible - model_name: qwen3-next-80b-a3b-thinking - token_count: # Optional - model_name: Qwen/Qwen3-Next-80B-A3B-Thinking - backend: hf - params: - use_mirror: true - - qwen3_max_instruct: - backend: openai_compatible - model_name: qwen3-max - token_count: # Optional - model_name: Qwen/Qwen3-Next-80B-A3B-Instruct # NOTE: We use another model as a substitute. - backend: hf - params: - use_mirror: true - - qwen25_max_instruct: - backend: openai_compatible - model_name: qwen-max-2025-01-25 - token_count: # Optional - model_name: Qwen/Qwen2.5-72B-Instruct # NOTE: We use another model as a substitute. - backend: hf - params: - use_mirror: true - -embedding_model: - default: - backend: openai_compatible - model_name: text-embedding-v4 - params: - dimensions: 1024 - -vector_store: - default: -# backend: elasticsearch - backend: memory - embedding_model: default -# params: -# hosts: "http://localhost:9200" diff --git a/flowllm/config/search.yaml b/flowllm/config/search.yaml deleted file mode 100644 index 808fa44..0000000 --- a/flowllm/config/search.yaml +++ /dev/null @@ -1,21 +0,0 @@ -external_mcp: - tongyi_search: - type: sse - url: "https://dashscope.aliyuncs.com/api/v1/mcps/WebSearch/sse" - headers: - Authorization: "Bearer {BAILIAN_MCP_API_KEY}" - tongyi_web_parser: - type: sse - url: "https://dashscope.aliyuncs.com/api/v1/mcps/WebParser/sse" - headers: - Authorization: "Bearer {BAILIAN_MCP_API_KEY}" - bochaai_search: - type: sse - url: "https://dashscope.aliyuncs.com/api/v1/mcps/bochaai_search/sse" - headers: - Authorization: "Bearer {BAILIAN_MCP_API_KEY}" - brave_search: - type: sse - url: "https://dashscope.aliyuncs.com/api/v1/mcps/Brave_Search/sse" - headers: - Authorization: "Bearer {BAILIAN_MCP_API_KEY}" \ No newline at end of file diff --git a/flowllm/constants.py b/flowllm/constants.py new file mode 100644 index 0000000..41eb8e5 --- /dev/null +++ b/flowllm/constants.py @@ -0,0 +1,15 @@ +"""Constants""" + +FLOWLLM_SERVICE_INFO = "FLOWLLM_SERVICE_INFO" + +FLOWLLM_DEFAULT_HOST = "127.0.0.1" + +FLOWLLM_DEFAULT_PORT = 2333 + +# File IO limits and truncation marker. +DEFAULT_MAX_BYTES = 50 * 1024 +MAX_FILE_READ_BYTES = 200 * 1024 * 1024 +TRUNCATION_NOTICE_MARKER = "<>" + +# Oversized images return path & metadata only (no base64). +DEFAULT_MAX_IMAGE_BYTES = 5 * 1024 * 1024 diff --git a/flowllm/core/__init__.py b/flowllm/core/__init__.py deleted file mode 100644 index c0e335d..0000000 --- a/flowllm/core/__init__.py +++ /dev/null @@ -1,51 +0,0 @@ -"""Core package for FlowLLM framework. - -This package provides the core components for building and executing LLM-powered flows. -It includes modules for: - -- Context: Context management for application state, flow execution, and prompts -- Embedding Models: Base classes and implementations for embedding models -- Flow: Flow composition and execution engines -- LLM: Large Language Model interfaces and implementations -- Vector Store: Vector storage and retrieval implementations -- Schema: Data models and configuration schemas -- Service: Service interfaces and implementations (HTTP, MCP, CMD) -- Storage: Caching and persistence utilities -- Token: Token counting adapters for different model providers -- Utils: Common utility functions and helpers -- Operations: Base operation classes for flow composition -- Enumeration: Core enumeration types - -Typical usage: - from flowllm.core import context, llm, flow, vector_store - from flowllm.core.context import ServiceContext - from flowllm.core.llm import BaseLLM, LiteLLM -""" - -from . import context -from . import embedding_model -from . import enumeration -from . import flow -from . import llm -from . import op -from . import schema -from . import service -from . import storage -from . import token -from . import utils -from . import vector_store - -__all__ = [ - "context", - "embedding_model", - "enumeration", - "flow", - "llm", - "op", - "schema", - "service", - "storage", - "token", - "utils", - "vector_store", -] diff --git a/flowllm/core/application.py b/flowllm/core/application.py deleted file mode 100644 index 7b23541..0000000 --- a/flowllm/core/application.py +++ /dev/null @@ -1,384 +0,0 @@ -"""FlowLLM application core module for managing flows, services, and configurations.""" - -import asyncio -import os -from concurrent.futures import ThreadPoolExecutor - -from loguru import logger - -from .context import C -from .enumeration import RegistryEnum -from .flow import BaseFlow, ExpressionToolFlow -from .schema import ServiceConfig, EmbeddingModelConfig, FlowStreamChunk -from .service import BaseService -from .utils import FastMcpClient, PydanticConfigParser, init_logger, print_logo - - -class Application: - """ - Base application class for FlowLLM framework. - - This is a generic, configurable base class that can be inherited to create - custom application implementations. It manages the lifecycle of FlowLLM - applications including configuration, service initialization, flow execution, - and resource management. Supports both synchronous and asynchronous operation modes. - - This class can be used directly with custom parsers and configurations, or - can be subclassed (e.g., FlowLLMApp) to provide default configurations - and convenience methods for specific use cases. - """ - - def __init__( - self, - *args, - llm_api_key: str = None, - llm_api_base: str = None, - embedding_api_key: str = None, - embedding_api_base: str = None, - service_config: ServiceConfig = None, - parser: type[PydanticConfigParser] = None, - config_path: str = None, - load_default_config: bool = False, - **kwargs, - ): - """ - Initialize FlowLLM application with configuration. - - Args: - *args: Additional arguments passed to parser. Examples: - - "llm.default.model_name=qwen3-30b-a3b-thinking-2507" - - "llm.default.backend=openai_compatible" - - "llm.default.params={'temperature': '0.6'}" - - "embedding_model.default.model_name=text-embedding-v4" - - "embedding_model.default.backend=openai_compatible" - - "embedding_model.default.params={'dimensions': 1024}" - - "vector_store.default.backend=memory" - - "vector_store.default.embedding_model=default" - - "vector_store.default.params={...}" - llm_api_key: API key for LLM service - llm_api_base: Base URL for LLM API - embedding_api_key: API key for embedding service - embedding_api_base: Base URL for embedding API - service_config: Pre-configured ServiceConfig object - parser: Custom configuration parser class - config_path: Path to custom configuration YAML file. If provided, loads configuration from this file. - Example: "path/to/my_config.yaml" - load_default_config: Whether to load default configuration (default.yaml). - If True and config_path is not provided, loads the default configuration. - **kwargs: Additional keyword arguments passed to parser. Same format as args but as kwargs. - """ - - if llm_api_key: - os.environ["FLOW_LLM_API_KEY"] = llm_api_key - - if llm_api_base: - os.environ["FLOW_LLM_BASE_URL"] = llm_api_base - - if embedding_api_key: - os.environ["FLOW_EMBEDDING_API_KEY"] = embedding_api_key - - if embedding_api_base: - os.environ["FLOW_EMBEDDING_BASE_URL"] = embedding_api_base - - # Initialize parser first (needed for update_service_config method) - if parser is None: - parser = PydanticConfigParser - self.parser = parser(ServiceConfig) - - if service_config is not None: - self.service_config: ServiceConfig = service_config - else: - input_args = [] - if config_path: - input_args.append(f"config={config_path}") - elif load_default_config: - input_args.append(f"config={parser.default_config_name}") - - if args: - input_args.extend(args) - - if kwargs: - input_args.extend([f"{k}={v}" for k, v in kwargs.items()]) - - self.service_config = self.parser.parse_args(*input_args) - - if self.service_config.init_logger: - init_logger() - - def update_service_config(self, **kwargs): - """ - Update configuration object using keyword arguments - - Args: - **kwargs: Configuration items to update, supports dot notation, e.g. a.b.c='xxx' - - Returns: - Updated configuration object - """ - self.service_config = self.parser.update_config(**kwargs) - return self.service_config - - async def __aenter__(self): - """ - Async context manager entry point. - - Returns: - Self instance after starting the application - """ - await self.async_start() - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - """ - Async context manager exit point. - - Args: - exc_type: Exception type if an error occurred - exc_val: Exception value if an error occurred - exc_tb: Exception traceback if an error occurred - - Returns: - False to propagate any exception - """ - await self.async_stop() - return False - - @staticmethod - async def get_mcp_tools(name: str, mcp_server_config: dict) -> dict: - """ - Retrieve available tools from an MCP (Model Context Protocol) server. - - Args: - name: Name identifier for the MCP server - mcp_server_config: Configuration dictionary for the MCP server - - Returns: - Dictionary containing server name and available tool calls, or empty dict on error - """ - try: - async with FastMcpClient(name=name, config=mcp_server_config) as client: - tool_calls = await client.list_tool_calls() - for tool_call in tool_calls: - logger.info(f"find mcp@{name}@{tool_call.name} {tool_call.model_dump_json()}") - - return { - "name": name, - "tool_calls": {tool_call.name: tool_call for tool_call in tool_calls}, - } - - except Exception as e: - logger.exception(f"get mcp@{name} tool_calls error: {e}") - return {} - - def filter_flows(self, name: str) -> bool: - """ - Determine if a flow should be enabled based on configuration. - - Args: - name: Flow name to check - - Returns: - True if the flow should be enabled, False otherwise - """ - if self.service_config.enabled_flows: - return name in self.service_config.enabled_flows - elif self.service_config.disabled_flows: - return name not in self.service_config.disabled_flows - else: - return True - - async def async_start(self): - """ - Asynchronously start the FlowLLM application. - - Initializes external MCP servers, service configuration, thread pools, - vector stores, embedding models, and registers all flows. - """ - # add external_mcp - for name, mcp_server_config in self.service_config.external_mcp.items(): - mcp_server_info = await self.get_mcp_tools(name, mcp_server_config) - if mcp_server_info: - C.external_mcp_tool_call_dict[mcp_server_info["name"]] = mcp_server_info["tool_calls"] - - # add service_config & language & thread_pool & ray - C.service_config = self.service_config - C.language = self.service_config.language - C.thread_pool = ThreadPoolExecutor(max_workers=self.service_config.thread_pool_max_workers) - if self.service_config.ray_max_workers > 1: - import ray - - ray.init(num_cpus=self.service_config.ray_max_workers) - - # add vector store - for name, config in self.service_config.vector_store.items(): - vector_store_cls = C.get_vector_store_class(config.backend) - embedding_model_config: EmbeddingModelConfig = self.service_config.embedding_model[config.embedding_model] - embedding_model_cls = C.get_embedding_model_class(embedding_model_config.backend) - embedding_model = embedding_model_cls( - model_name=embedding_model_config.model_name, - **embedding_model_config.params, - ) - C.vector_store_dict[name] = vector_store_cls(embedding_model=embedding_model, **config.params) - - # add cls flow - for name, flow_cls in C.registry_dict[RegistryEnum.FLOW].items(): - if not self.filter_flows(name): - continue - - flow: BaseFlow = flow_cls() - C.flow_dict[flow.name] = flow - - # add expression flow - for name, flow_config in self.service_config.flow.items(): - if not self.filter_flows(name): - continue - - flow_config.name = name - flow: BaseFlow = ExpressionToolFlow(flow_config=flow_config) - C.flow_dict[name] = flow - - async def async_stop(self, wait_thread_pool=True, wait_ray: bool = True): - """ - Asynchronously stop the FlowLLM application and clean up resources. - - Args: - wait_thread_pool: Whether to wait for thread pool tasks to complete - wait_ray: Whether to wait for Ray tasks to complete - """ - for _, vector_store in C.vector_store_dict.items(): - await vector_store.async_close() - C.thread_pool.shutdown(wait=wait_thread_pool) - if self.service_config.ray_max_workers > 1: - import ray - - ray.shutdown(_exiting_interpreter=not wait_ray) - - def __enter__(self): - """ - Synchronous context manager entry point. - - Returns: - Self instance after starting the application - """ - self.start() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - """ - Synchronous context manager exit point. - - Args: - exc_type: Exception type if an error occurred - exc_val: Exception value if an error occurred - exc_tb: Exception traceback if an error occurred - - Returns: - False to propagate any exception - """ - self.stop() - return False - - def start(self): - """ - Synchronously start the FlowLLM application. - - Wraps async_start() in asyncio.run() for synchronous usage. - """ - asyncio.run(self.async_start()) - - def stop(self, wait_thread_pool=True, wait_ray: bool = True): - """ - Synchronously stop the FlowLLM application and clean up resources. - - Args: - wait_thread_pool: Whether to wait for thread pool tasks to complete - wait_ray: Whether to wait for Ray tasks to complete - """ - for _, vector_store in C.vector_store_dict.items(): - vector_store.close() - C.thread_pool.shutdown(wait=wait_thread_pool) - if self.service_config.ray_max_workers > 1: - import ray - - ray.shutdown(_exiting_interpreter=not wait_ray) - - @staticmethod - def execute_flow(name: str, **kwargs): - """ - Synchronously execute a non-streaming flow by name. - - Args: - name: Name of the flow to execute - **kwargs: Additional arguments to pass to the flow - - Returns: - Flow execution result - - Raises: - AssertionError: If the flow is configured for streaming - """ - flow: BaseFlow = C.get_flow(name) - assert flow.stream is False, "stream is not supported in execute_flow!" - return flow.call(**kwargs) - - @staticmethod - async def async_execute_flow(name: str, **kwargs): - """ - Asynchronously execute a non-streaming flow by name. - - Args: - name: Name of the flow to execute - **kwargs: Additional arguments to pass to the flow - - Returns: - Flow execution result - - Raises: - AssertionError: If the flow is configured for streaming - """ - flow: BaseFlow = C.get_flow(name) - assert flow.stream is False, "stream is not supported in async_execute_flow!" - return await flow.async_call(**kwargs) - - @staticmethod - async def async_execute_stream_flow(name: str, **kwargs): - """ - Asynchronously execute a streaming flow by name. - - Args: - name: Name of the flow to execute - **kwargs: Additional arguments to pass to the flow - - Yields: - Stream chunks in SSE (Server-Sent Events) format - - Raises: - AssertionError: If the flow is not configured for streaming - """ - flow: BaseFlow = C.get_flow(name) - assert flow.stream is True, "non-stream is not supported in async_execute_stream_flow!" - - stream_queue = asyncio.Queue() - asyncio.create_task(flow.async_call(stream_queue=stream_queue, **kwargs)) - while True: - stream_chunk: FlowStreamChunk = await stream_queue.get() - if stream_chunk.done: - yield "data:[DONE]\n\n" - break - - yield f"data:{stream_chunk.model_dump_json()}\n\n" - - def run_service(self): - """ - Start the service based on the configured backend. - - Initializes and runs the appropriate service implementation - (e.g., FastAPI, Flask) according to service configuration. - """ - - if self.service_config.enable_logo: - print_logo(service_config=self.service_config, app_name=os.getenv("FLOW_APP_NAME")) - - service_cls = C.get_service_class(self.service_config.backend) - service: BaseService = service_cls(service_config=self.service_config) - service.run() diff --git a/flowllm/core/context/__init__.py b/flowllm/core/context/__init__.py deleted file mode 100644 index 267594f..0000000 --- a/flowllm/core/context/__init__.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Context management module for FlowLLM. - -This module provides context classes for managing application state, -flow execution, prompts, and class registries. - -Exports: - BaseContext: Base class for dictionary-like context management. - FlowContext: Context for managing flow execution state and streaming. - PromptHandler: Handler for loading and formatting prompts. - Registry: Registry for storing and retrieving registered classes. - ServiceContext: Singleton service context for global application state. - C: Singleton instance of ServiceContext. -""" - -from .base_context import BaseContext -from .flow_context import FlowContext -from .prompt_handler import PromptHandler -from .registry import Registry -from .service_context import C -from .service_context import ServiceContext - -__all__ = [ - "BaseContext", - "FlowContext", - "PromptHandler", - "Registry", - "ServiceContext", - "C", -] diff --git a/flowllm/core/context/base_context.py b/flowllm/core/context/base_context.py deleted file mode 100644 index 058db7c..0000000 --- a/flowllm/core/context/base_context.py +++ /dev/null @@ -1,172 +0,0 @@ -"""Base context class for managing dynamic attributes. - -This module provides a base class for context management with dictionary-like -behavior, allowing dynamic attribute access and storage. -""" - - -class BaseContext: - """Base context class that provides dictionary-like attribute access. - - This class allows dynamic attribute storage and retrieval, making it - suitable for context objects where attributes may not be known at - class definition time. - - Attributes: - _data: Internal dictionary storing all context data. - """ - - def __init__(self, **kwargs): - """Initialize BaseContext with optional keyword arguments. - - Args: - **kwargs: Initial data to store in the context. - """ - self._data: dict = {**kwargs} - - def __getattr__(self, name: str): - """Get attribute by name. - - Args: - name: Attribute name to retrieve. - - Returns: - The value associated with the attribute name. - - Raises: - AttributeError: If the attribute is not found. - """ - data = object.__getattribute__(self, "_data") - if name in data: - return data[name] - else: - raise AttributeError(f"'{self.__class__.__name__}' has no attribute '{name}'") - - def __setattr__(self, name: str, value): - """Set attribute by name. - - Args: - name: Attribute name to set. - value: Value to assign to the attribute. - """ - if name == "_data": - super().__setattr__(name, value) - else: - self._data[name] = value - - def __getitem__(self, name: str): - """Get item using dictionary-style access. - - Args: - name: Key name to retrieve. - - Returns: - The value associated with the key. - - Raises: - AttributeError: If the key is not found. - """ - if name in self._data: - return self._data[name] - else: - raise AttributeError(f"'{self.__class__.__name__}' has no attribute '{name}'") - - def __setitem__(self, name: str, value): - """Set item using dictionary-style access. - - Args: - name: Key name to set. - value: Value to assign to the key. - """ - self._data[name] = value - - def __contains__(self, name: str): - """Check if a key exists in the context. - - Args: - name: Key name to check. - - Returns: - True if the key exists, False otherwise. - """ - return name in self._data - - def __repr__(self): - """Return string representation of the context. - - Returns: - String representation showing class name and data. - """ - return f"{self.__class__.__name__}({self._data!r})" - - def dump(self) -> dict: - """Dump all context data as a dictionary. - - Returns: - A copy of all stored data as a dictionary. - """ - return {**self._data} - - def get(self, key: str, default=None): - """Get a value by key with optional default. - - Args: - key: Key name to retrieve. - default: Default value to return if key is not found. - - Returns: - The value associated with the key, or default if not found. - """ - return self._data.get(key, default) - - def keys(self): - """Get all keys in the context. - - Returns: - A view object of all keys in the context. - """ - return self._data.keys() - - def values(self): - """Get all values in the context. - - Returns: - A view object of all values in the context. - """ - return self._data.values() - - def update(self, kwargs: dict): - """Update context with new key-value pairs. - - Args: - kwargs: Dictionary of key-value pairs to update. - """ - self._data.update(kwargs) - - def items(self): - """Get all key-value pairs in the context. - - Returns: - A view object of all key-value pairs in the context. - """ - return self._data.items() - - def clear(self): - """Clear all context data.""" - self._data.clear() - - def __getstate__(self): - """Get state for pickling. - - Returns: - The internal data dictionary. - """ - return self._data - - def __setstate__(self, state): - """Set state from pickling. - - Args: - state: The data dictionary to restore. - """ - self._data = state diff --git a/flowllm/core/context/flow_context.py b/flowllm/core/context/flow_context.py deleted file mode 100644 index 8cf82cd..0000000 --- a/flowllm/core/context/flow_context.py +++ /dev/null @@ -1,98 +0,0 @@ -"""Flow context for managing flow execution state. - -This module provides a context class for managing the state of flow execution, -including flow identification, response handling, and streaming capabilities. -""" - -import asyncio -import uuid -from typing import Optional - -from .base_context import BaseContext -from ..enumeration import ChunkEnum -from ..schema import FlowResponse -from ..schema import FlowStreamChunk - - -class FlowContext(BaseContext): - """Context for managing flow execution state and streaming. - - This class manages the state of a single flow execution, including: - - Flow identification - - Response handling - - Stream queue for asynchronous streaming - - Attributes: - flow_id: Unique identifier for the flow instance. - response: FlowResponse object for storing flow results. - stream_queue: Asynchronous queue for streaming chunks. - """ - - def __init__( - self, - flow_id: str = uuid.uuid4().hex, - response: Optional[FlowResponse] = None, - stream_queue: Optional[asyncio.Queue] = None, - **kwargs, - ): - """Initialize FlowContext with flow ID and optional components. - - Args: - flow_id: Unique identifier for the flow instance. - Defaults to a random UUID hex string. - response: FlowResponse object. If None, a new FlowResponse - will be created. - stream_queue: Asynchronous queue for streaming chunks. - **kwargs: Additional context data to store. - """ - super().__init__(**kwargs) - - self.flow_id: str = flow_id - self.response: Optional[FlowResponse] = response if response is not None else FlowResponse() - self.stream_queue: Optional[asyncio.Queue] = stream_queue - - async def add_stream_string_and_type(self, chunk: str, chunk_type: ChunkEnum): - """Add a stream chunk with string content and type to the stream queue. - - Args: - chunk: The string content to stream. - chunk_type: The type of the chunk. - - Returns: - Self for method chaining. - """ - stream_chunk = FlowStreamChunk(flow_id=self.flow_id, chunk_type=chunk_type, chunk=chunk) - await self.stream_queue.put(stream_chunk) - return self - - async def add_stream_chunk(self, stream_chunk: FlowStreamChunk): - """Add a stream chunk to the stream queue. - - Args: - stream_chunk: The stream chunk to add. - - Returns: - Self for method chaining. - """ - stream_chunk.flow_id = self.flow_id - await self.stream_queue.put(stream_chunk) - return self - - async def add_stream_done(self): - """Add a done signal to the stream queue. - - Returns: - Self for method chaining. - """ - done_chunk = FlowStreamChunk(flow_id=self.flow_id, chunk_type=ChunkEnum.DONE, chunk="", done=True) - await self.stream_queue.put(done_chunk) - return self - - def add_response_error(self, e: Exception): - """Add an error to the flow response. - - Args: - e: The exception to record as an error. - """ - self.response.success = False - self.response.answer = str(e.args) diff --git a/flowllm/core/context/prompt_handler.py b/flowllm/core/context/prompt_handler.py deleted file mode 100644 index 280e764..0000000 --- a/flowllm/core/context/prompt_handler.py +++ /dev/null @@ -1,150 +0,0 @@ -"""Prompt handler for loading and managing prompts. - -This module provides a class for loading prompts from files, managing -language-specific prompts, and formatting prompts with variables. -""" - -from pathlib import Path - -import yaml -from loguru import logger - -from .base_context import BaseContext -from .service_context import C - - -class PromptHandler(BaseContext): - """Handler for loading, storing, and formatting prompts. - - This class manages prompts loaded from YAML files, supports - language-specific prompts, and provides formatting with variables - and conditional flags. - - Attributes: - language: Language code for language-specific prompts. - """ - - def __init__(self, language: str = "", **kwargs): - """Initialize PromptHandler with language setting. - - Args: - language: Language code for language-specific prompts. - If not provided, uses the language from service context. - **kwargs: Additional context data to store. - """ - super().__init__(**kwargs) - self.language: str = language or C.language - - def load_prompt_by_file(self, prompt_file_path: Path | str = None): - """Load prompts from a YAML file. - - Args: - prompt_file_path: Path to the YAML file containing prompts. - Can be a Path object or string. If None, does nothing. - - Returns: - Self for method chaining. - """ - if prompt_file_path is None: - return self - - if isinstance(prompt_file_path, str): - prompt_file_path = Path(prompt_file_path) - - if not prompt_file_path.exists(): - return self - - with prompt_file_path.open(encoding="utf-8") as f: - prompt_dict = yaml.load(f, yaml.FullLoader) - self.load_prompt_dict(prompt_dict) - return self - - def load_prompt_dict(self, prompt_dict: dict = None): - """Load prompts from a dictionary. - - Args: - prompt_dict: Dictionary of prompt names to prompt strings. - If None, does nothing. - - Returns: - Self for method chaining. - """ - if not prompt_dict: - return self - - for key, value in prompt_dict.items(): - if isinstance(value, str): - if key in self._data: - self._data[key] = value - logger.warning(f"prompt_dict key={key} overwrite!") - - else: - self._data[key] = value - logger.debug(f"add prompt_dict key={key}") - return self - - def get_prompt(self, prompt_name: str): - """Get a prompt by name, with language suffix if applicable. - - Args: - prompt_name: Base name of the prompt to retrieve. - - Returns: - The prompt string. - - Raises: - AssertionError: If the prompt (with language suffix) is not found. - """ - key: str = prompt_name - if self.language and not key.endswith(self.language.strip()): - key += "_" + self.language.strip() - - assert key in self._data, f"prompt_name={key} not found." - return self._data[key] - - def prompt_format(self, prompt_name: str, **kwargs) -> str: - """Format a prompt with variables and conditional flags. - - This method supports two types of formatting: - 1. Boolean flags: Lines starting with [flag_name] are included - only if the corresponding flag is True. - 2. Variable substitution: Other kwargs are used for string - formatting with {variable_name}. - - Args: - prompt_name: Name of the prompt to format. - **kwargs: Variables and flags for formatting. - - Returns: - The formatted prompt string. - """ - prompt = self.get_prompt(prompt_name) - - flag_kwargs = {k: v for k, v in kwargs.items() if isinstance(v, bool)} - other_kwargs = {k: v for k, v in kwargs.items() if not isinstance(v, bool)} - - if flag_kwargs: - split_prompt = [] - for line in prompt.strip().split("\n"): - hit = False - hit_flag = True - for key, flag in kwargs.items(): - if not line.startswith(f"[{key}]"): - continue - - hit = True - hit_flag = flag - line = line.strip(f"[{key}]") - break - - if not hit: - split_prompt.append(line) - elif hit_flag: - split_prompt.append(line) - - prompt = "\n".join(split_prompt) - - if other_kwargs: - prompt = prompt.format(**other_kwargs) - - return prompt diff --git a/flowllm/core/context/registry.py b/flowllm/core/context/registry.py deleted file mode 100644 index e3d3fbd..0000000 --- a/flowllm/core/context/registry.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Registry for class registration and lookup. - -This module provides a registry class for managing class registrations -with dynamic lookup capabilities. -""" - -from .base_context import BaseContext - - -class Registry(BaseContext): - """Registry for storing and retrieving registered classes. - - This class provides a decorator-based registration system for classes, - allowing dynamic class lookup by name. - """ - - def register(self, name: str = "", add_cls: bool = True): - """Register a class in the registry. - - Args: - name: Name to register the class under. - add_cls: Whether to actually add the class to the registry. - Defaults to True. - - Returns: - Decorator function that registers the class when applied. - """ - - def decorator(cls): - """Decorator function that registers the class. - - Args: - cls: The class to register. - - Returns: - The class unchanged (for use as decorator). - """ - if add_cls: - key = name or cls.__name__ - self._data[key] = cls - return cls - - return decorator diff --git a/flowllm/core/context/service_context.py b/flowllm/core/context/service_context.py deleted file mode 100644 index b0f7a6d..0000000 --- a/flowllm/core/context/service_context.py +++ /dev/null @@ -1,290 +0,0 @@ -"""Service context for managing global application state. - -This module provides a singleton service context that manages application-wide -configuration, registries, vector stores, and other shared resources. -""" - -import uuid -from concurrent.futures import ThreadPoolExecutor -from typing import Dict - -from .base_context import BaseContext -from .registry import Registry -from ..enumeration import RegistryEnum -from ..schema import ServiceConfig -from ..utils import singleton - - -@singleton -class ServiceContext(BaseContext): - """Singleton service context for global application state management. - - This class manages application-wide configuration including: - - Service configuration and identification - - Registry for models, operations, flows, and services - - Vector stores - - Thread pool executor - - External MCP tool calls - - Flow instances - - Attributes: - service_id: Unique identifier for the service instance. - service_config: Service configuration object. - language: Language setting for the service. - thread_pool: Thread pool executor for concurrent operations. - vector_store_dict: Dictionary of vector store instances. - external_mcp_tool_call_dict: Dictionary of external MCP tool calls. - registry_dict: Dictionary of registries by type. - flow_dict: Dictionary of flow instances. - """ - - def __init__(self, service_id: str | None = None, **kwargs): - """Initialize ServiceContext with service ID. - - Args: - service_id: Unique identifier for the service instance. - Defaults to a random UUID hex string. - **kwargs: Additional context data to store. - """ - super().__init__(**kwargs) - self.service_id: str = service_id or uuid.uuid4().hex - - self.service_config: ServiceConfig | None = None - self.language: str = "" - self.thread_pool: ThreadPoolExecutor | None = None - self.vector_store_dict: dict = {} - self.external_mcp_tool_call_dict: dict = {} - self.registry_dict: Dict[RegistryEnum, Registry] = {v: Registry() for v in RegistryEnum.__members__.values()} - self.flow_dict: dict = {} - - def register(self, name: str, register_type: RegistryEnum): - """Register a model class by name and type. - - Args: - name: Name to register the class under. - register_type: Type of registry (e.g., LLM, EMBEDDING_MODEL). - - Returns: - Decorator function for class registration. - """ - return self.registry_dict[register_type].register(name=name) - - def register_embedding_model(self, name: str = ""): - """Register an embedding model class. - - Args: - name: Name to register the class under. - - Returns: - Decorator function for class registration. - """ - return self.register(name=name, register_type=RegistryEnum.EMBEDDING_MODEL) - - def register_llm(self, name: str = ""): - """Register an LLM class. - - Args: - name: Name to register the class under. - - Returns: - Decorator function for class registration. - """ - return self.register(name=name, register_type=RegistryEnum.LLM) - - def register_vector_store(self, name: str = ""): - """Register a vector store class. - - Args: - name: Name to register the class under. - - Returns: - Decorator function for class registration. - """ - return self.register(name=name, register_type=RegistryEnum.VECTOR_STORE) - - def register_op(self, name: str = ""): - """Register an operation class. - - Args: - name: Name to register the class under. - - Returns: - Decorator function for class registration. - """ - return self.register(name=name, register_type=RegistryEnum.OP) - - def register_flow(self, name: str = ""): - """Register a flow class. - - Args: - name: Name to register the class under. - - Returns: - Decorator function for class registration. - """ - return self.register(name=name, register_type=RegistryEnum.FLOW) - - def register_service(self, name: str = ""): - """Register a service class. - - Args: - name: Name to register the class under. - - Returns: - Decorator function for class registration. - """ - return self.register(name=name, register_type=RegistryEnum.SERVICE) - - def register_token_counter(self, name: str = ""): - """Register a token counter class. - - Args: - name: Name to register the class under. - - Returns: - Decorator function for class registration. - """ - return self.register(name=name, register_type=RegistryEnum.TOKEN_COUNTER) - - def get_model_class(self, name: str, register_type: RegistryEnum): - """Get a registered model class by name and type. - - Args: - name: Name of the registered class. - register_type: Type of registry to search. - - Returns: - The registered class. - - Raises: - AssertionError: If the name is not found in the registry. - """ - assert name in self.registry_dict[register_type], ( - f"name={name} not found in registry_dict.{register_type.value}! " - f"supported names={self.registry_dict[register_type].keys()}" - ) - - return self.registry_dict[register_type][name] - - def get_embedding_model_class(self, name: str): - """Get a registered embedding model class. - - Args: - name: Name of the registered class. - - Returns: - The registered embedding model class. - - Raises: - AssertionError: If the name is not found in the registry. - """ - return self.get_model_class(name, RegistryEnum.EMBEDDING_MODEL) - - def get_llm_class(self, name: str): - """Get a registered LLM class. - - Args: - name: Name of the registered class. - - Returns: - The registered LLM class. - - Raises: - AssertionError: If the name is not found in the registry. - """ - return self.get_model_class(name, RegistryEnum.LLM) - - def get_vector_store_class(self, name: str): - """Get a registered vector store class. - - Args: - name: Name of the registered class. - - Returns: - The registered vector store class. - - Raises: - AssertionError: If the name is not found in the registry. - """ - return self.get_model_class(name, RegistryEnum.VECTOR_STORE) - - def get_op_class(self, name: str): - """Get a registered operation class. - - Args: - name: Name of the registered class. - - Returns: - The registered operation class. - - Raises: - AssertionError: If the name is not found in the registry. - """ - return self.get_model_class(name, RegistryEnum.OP) - - def get_flow_class(self, name: str): - """Get a registered flow class. - - Args: - name: Name of the registered class. - - Returns: - The registered flow class. - - Raises: - AssertionError: If the name is not found in the registry. - """ - return self.get_model_class(name, RegistryEnum.FLOW) - - def get_service_class(self, name: str): - """Get a registered service class. - - Args: - name: Name of the registered class. - - Returns: - The registered service class. - - Raises: - AssertionError: If the name is not found in the registry. - """ - return self.get_model_class(name, RegistryEnum.SERVICE) - - def get_token_counter_class(self, name: str): - """Get a registered token counter class. - - Args: - name: Name of the registered class. - - Returns: - The registered token counter class. - - Raises: - AssertionError: If the name is not found in the registry. - """ - return self.get_model_class(name, RegistryEnum.TOKEN_COUNTER) - - def get_vector_store(self, name: str = "default"): - """Get a vector store instance by name. - - Args: - name: Name of the vector store instance. Defaults to "default". - - Returns: - The vector store instance. - """ - return self.vector_store_dict[name] - - def get_flow(self, name: str = "default"): - """Get a flow instance by name. - - Args: - name: Name of the flow instance. Defaults to "default". - - Returns: - The flow instance. - """ - return self.flow_dict[name] - - -C = ServiceContext() diff --git a/flowllm/core/embedding_model/__init__.py b/flowllm/core/embedding_model/__init__.py deleted file mode 100644 index 5a81c44..0000000 --- a/flowllm/core/embedding_model/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Embedding model implementations for FlowLLM. - -This module provides base classes and concrete implementations for embedding models, -supporting both synchronous and asynchronous operations with retry logic and batch processing. -""" - -from .base_embedding_model import BaseEmbeddingModel -from .openai_compatible_embedding_model import OpenAICompatibleEmbeddingModel - -__all__ = [ - "BaseEmbeddingModel", - "OpenAICompatibleEmbeddingModel", -] diff --git a/flowllm/core/embedding_model/base_embedding_model.py b/flowllm/core/embedding_model/base_embedding_model.py deleted file mode 100644 index d7c9b01..0000000 --- a/flowllm/core/embedding_model/base_embedding_model.py +++ /dev/null @@ -1,228 +0,0 @@ -"""Base embedding model implementation. - -This module provides the abstract base class for embedding models with common -functionality including retry logic, error handling, and batch processing. -""" - -import asyncio -from abc import ABC -from typing import List - -from loguru import logger - -from ..schema import VectorNode - - -class BaseEmbeddingModel(ABC): - """ - Abstract base class for embedding models. - - This class provides a common interface for various embedding model implementations, - including retry logic, error handling, and batch processing capabilities. - """ - - def __init__( - self, - model_name: str = "", - dimensions: int = 1024, - max_batch_size: int = 10, - max_retries: int = 3, - raise_exception: bool = True, - **kwargs, - ): - """ - Initialize the embedding model. - - Args: - model_name: Name of the embedding model - dimensions: Dimensionality of the embedding vectors - max_batch_size: Maximum batch size for processing - max_retries: Maximum number of retry attempts on failure - raise_exception: Whether to raise exceptions after max retries - """ - self.model_name: str = model_name - self.dimensions: int = dimensions - self.max_batch_size: int = max_batch_size - self.max_retries: int = max_retries - self.raise_exception: bool = raise_exception - self.kwargs: dict = kwargs - - def _get_embeddings(self, input_text: str | List[str]): - """ - Abstract method to get embeddings from the model. - - This method must be implemented by concrete subclasses to provide - the actual embedding functionality. - - Args: - input_text: Single text string or list of text strings to embed - - Returns: - Embedding vector(s) corresponding to the input text(s) - """ - raise NotImplementedError - - async def _async_get_embeddings(self, input_text: str | List[str]): - """ - Abstract async method to get embeddings from the model. - - This method must be implemented by concrete subclasses to provide - the actual async embedding functionality. - - Args: - input_text: Single text string or list of text strings to embed - - Returns: - Embedding vector(s) corresponding to the input text(s) - """ - raise NotImplementedError - - def get_embeddings(self, input_text: str | List[str]): - """ - Get embeddings with retry logic and error handling. - - This method wraps the _get_embeddings method with automatic retry - functionality in case of failures. - - Args: - input_text: Single text string or list of text strings to embed - - Returns: - Embedding vector(s) or None if all retries failed and raise_exception is False - """ - # Retry loop with exponential backoff potential - for i in range(self.max_retries): - try: - return self._get_embeddings(input_text) - - except Exception as e: - logger.exception(f"embedding model name={self.model_name} encounter error with e={e.args}") - # If this is the last retry and raise_exception is True, re-raise the exception - if i == self.max_retries - 1 and self.raise_exception: - raise e - - # Return None if all retries failed and raise_exception is False - return None - - async def async_get_embeddings(self, input_text: str | List[str]): - """ - Get embeddings asynchronously with retry logic and error handling. - - This method wraps the _get_embeddings_async method with automatic retry - functionality in case of failures. - - Args: - input_text: Single text string or list of text strings to embed - - Returns: - Embedding vector(s) or None if all retries failed and raise_exception is False - """ - # Retry loop with exponential backoff potential - for i in range(self.max_retries): - try: - return await self._async_get_embeddings(input_text) - - except Exception as e: - logger.exception(f"embedding model name={self.model_name} encounter error with e={e.args}") - # If this is the last retry and raise_exception is True, re-raise the exception - if i == self.max_retries - 1 and self.raise_exception: - raise e - - # Return None if all retries failed and raise_exception is False - return None - - def get_node_embeddings(self, nodes: VectorNode | List[VectorNode]): - """ - Generate embeddings for VectorNode objects and update their vector fields. - - This method handles both single nodes and lists of nodes, with automatic - batching for efficient processing of large node lists. - - Args: - nodes: Single VectorNode or list of VectorNode objects to embed - - Returns: - The same node(s) with updated vector fields containing embeddings - - Raises: - RuntimeError: If unsupported node type is provided - """ - # Handle single VectorNode - if isinstance(nodes, VectorNode): - nodes.vector = self.get_embeddings(nodes.content) - return nodes - - # Handle list of VectorNodes with batch processing - elif isinstance(nodes, list): - # Process nodes in batches to respect max_batch_size limits - embeddings = [ - emb - for i in range(0, len(nodes), self.max_batch_size) - for emb in self.get_embeddings(input_text=[node.content for node in nodes[i : i + self.max_batch_size]]) - ] - - # Validate that we got the expected number of embeddings - if len(embeddings) != len(nodes): - logger.warning(f"embeddings.size={len(embeddings)} <> nodes.size={len(nodes)}") - else: - # Assign embeddings to corresponding nodes - for node, embedding in zip(nodes, embeddings): - node.vector = embedding - return nodes - - else: - raise TypeError(f"unsupported type={type(nodes)}") - - async def async_get_node_embeddings(self, nodes: VectorNode | List[VectorNode]): - """ - Generate embeddings asynchronously for VectorNode objects and update their vector fields. - - This method handles both single nodes and lists of nodes, with automatic - batching for efficient processing of large node lists. - - Args: - nodes: Single VectorNode or list of VectorNode objects to embed - - Returns: - The same node(s) with updated vector fields containing embeddings - - Raises: - RuntimeError: If unsupported node type is provided - """ - # Handle single VectorNode - if isinstance(nodes, VectorNode): - nodes.vector = await self.async_get_embeddings(nodes.content) - return nodes - - # Handle list of VectorNodes with batch processing - elif isinstance(nodes, list): - # Process nodes in batches to respect max_batch_size limits - batch_tasks = [] - for i in range(0, len(nodes), self.max_batch_size): - batch_nodes = nodes[i : i + self.max_batch_size] - batch_content = [node.content for node in batch_nodes] - batch_tasks.append(self.async_get_embeddings(batch_content)) - - # Execute all batch tasks concurrently - batch_results = await asyncio.gather(*batch_tasks) - - # Flatten the results - embeddings = [emb for batch_result in batch_results for emb in batch_result] - - # Validate that we got the expected number of embeddings - if len(embeddings) != len(nodes): - logger.warning(f"embeddings.size={len(embeddings)} <> nodes.size={len(nodes)}") - else: - # Assign embeddings to corresponding nodes - for node, embedding in zip(nodes, embeddings): - node.vector = embedding - return nodes - - else: - raise TypeError(f"unsupported type={type(nodes)}") - - def close(self): - """Close the client connection or clean up resources.""" - - async def async_close(self): - """Asynchronously close the client connection or clean up resources.""" diff --git a/flowllm/core/embedding_model/openai_compatible_embedding_model.py b/flowllm/core/embedding_model/openai_compatible_embedding_model.py deleted file mode 100644 index d168f75..0000000 --- a/flowllm/core/embedding_model/openai_compatible_embedding_model.py +++ /dev/null @@ -1,149 +0,0 @@ -"""OpenAI-compatible embedding model implementation. - -This module provides an embedding model implementation that works with -OpenAI-compatible embedding APIs, including OpenAI's official API and -other services that follow the same interface. -""" - -import os -from typing import Literal, List - -from openai import OpenAI, AsyncOpenAI - -from .base_embedding_model import BaseEmbeddingModel -from ..context import C - - -@C.register_embedding_model("openai_compatible") -class OpenAICompatibleEmbeddingModel(BaseEmbeddingModel): - """ - OpenAI-compatible embedding model implementation. - - This class provides an implementation of BaseEmbeddingModel that works with - OpenAI-compatible embedding APIs, including OpenAI's official API and - other services that follow the same interface. - - Attributes: - api_key: API key for authentication - base_url: Base URL for the API endpoint - encoding_format: Encoding format for embeddings ("float" or "base64") - """ - - def __init__( - self, - model_name: str = "", - dimensions: int = 1024, - max_batch_size: int = 10, - max_retries: int = 3, - raise_exception: bool = True, - api_key: str | None = None, - base_url: str | None = None, - encoding_format: Literal["float", "base64"] = "float", - **kwargs, - ): - """ - Initialize the OpenAI-compatible embedding model. - - Args: - model_name: Name of the embedding model to use - dimensions: Dimensionality of the embedding vectors - api_key: API key for authentication (defaults to FLOW_EMBEDDING_API_KEY env var) - base_url: Base URL for the API endpoint (defaults to FLOW_EMBEDDING_BASE_URL env var) - encoding_format: Encoding format for embeddings - max_retries: Maximum number of retry attempts on failure - raise_exception: Whether to raise exceptions after max retries - max_batch_size: Maximum batch size for processing - """ - super().__init__( - model_name=model_name, - dimensions=dimensions, - max_retries=max_retries, - raise_exception=raise_exception, - max_batch_size=max_batch_size, - **kwargs, - ) - self.api_key = api_key or os.getenv("FLOW_EMBEDDING_API_KEY", "") - self.base_url = base_url or os.getenv("FLOW_EMBEDDING_BASE_URL", "") - self.encoding_format = encoding_format - - # Initialize OpenAI clients - self._client = OpenAI(api_key=self.api_key, base_url=self.base_url) - self._async_client = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url) - - def _get_embeddings(self, input_text: str | List[str]): - """ - Get embeddings from the OpenAI-compatible API. - - This method implements the abstract _get_embeddings method from BaseEmbeddingModel - by calling the OpenAI-compatible embeddings API. - - Args: - input_text: Single text string or list of text strings to embed - - Returns: - Embedding vector(s) corresponding to the input text(s) - - Raises: - RuntimeError: If unsupported input type is provided - """ - completion = self._client.embeddings.create( - model=self.model_name, - input=input_text, - dimensions=self.dimensions, - encoding_format=self.encoding_format, - ) - - if isinstance(input_text, str): - return completion.data[0].embedding - - elif isinstance(input_text, list): - result_emb = [[] for _ in range(len(input_text))] - for emb in completion.data: - result_emb[emb.index] = emb.embedding - return result_emb - - else: - raise RuntimeError(f"unsupported type={type(input_text)}") - - async def _async_get_embeddings(self, input_text: str | List[str]): - """ - Get embeddings asynchronously from the OpenAI-compatible API. - - This method implements the abstract _get_embeddings_async method from BaseEmbeddingModel - by calling the OpenAI-compatible embeddings API asynchronously. - - Args: - input_text: Single text string or list of text strings to embed - - Returns: - Embedding vector(s) corresponding to the input text(s) - - Raises: - RuntimeError: If unsupported input type is provided - """ - completion = await self._async_client.embeddings.create( - model=self.model_name, - input=input_text, - dimensions=self.dimensions, - encoding_format=self.encoding_format, - ) - - if isinstance(input_text, str): - return completion.data[0].embedding - - elif isinstance(input_text, list): - result_emb = [[] for _ in range(len(input_text))] - for emb in completion.data: - result_emb[emb.index] = emb.embedding - return result_emb - - else: - raise RuntimeError(f"unsupported type={type(input_text)}") - - def close(self): - """Close the OpenAI clients and clean up resources.""" - self._client.close() - - async def async_close(self): - """Asynchronously close the OpenAI clients and clean up resources.""" - await self._async_client.close() diff --git a/flowllm/core/enumeration/__init__.py b/flowllm/core/enumeration/__init__.py deleted file mode 100644 index f765040..0000000 --- a/flowllm/core/enumeration/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Core enumeration module.""" - -from .chunk_enum import ChunkEnum -from .http_enum import HttpEnum -from .registry_enum import RegistryEnum -from .role import Role - -__all__ = [ - "ChunkEnum", - "HttpEnum", - "RegistryEnum", - "Role", -] diff --git a/flowllm/core/enumeration/chunk_enum.py b/flowllm/core/enumeration/chunk_enum.py deleted file mode 100644 index 9a8ee57..0000000 --- a/flowllm/core/enumeration/chunk_enum.py +++ /dev/null @@ -1,14 +0,0 @@ -"""Chunk type enumeration for stream responses.""" - -from enum import Enum - - -class ChunkEnum(str, Enum): - """Enumeration of chunk types in stream responses.""" - - THINK = "think" - ANSWER = "answer" - TOOL = "tool" - USAGE = "usage" - ERROR = "error" - DONE = "done" diff --git a/flowllm/core/enumeration/http_enum.py b/flowllm/core/enumeration/http_enum.py deleted file mode 100644 index 292faa3..0000000 --- a/flowllm/core/enumeration/http_enum.py +++ /dev/null @@ -1,13 +0,0 @@ -"""HTTP method enumeration.""" - -from enum import Enum - - -class HttpEnum(str, Enum): - """Enumeration of HTTP methods.""" - - GET = "get" - POST = "post" - HEAD = "head" - PUT = "put" - DELETE = "delete" diff --git a/flowllm/core/enumeration/registry_enum.py b/flowllm/core/enumeration/registry_enum.py deleted file mode 100644 index a2b56db..0000000 --- a/flowllm/core/enumeration/registry_enum.py +++ /dev/null @@ -1,15 +0,0 @@ -"""Registry enumeration for component types.""" - -from enum import Enum - - -class RegistryEnum(str, Enum): - """Enumeration of registry types for different components.""" - - EMBEDDING_MODEL = "embedding_model" - LLM = "llm" - VECTOR_STORE = "vector_store" - OP = "op" - FLOW = "flow" - SERVICE = "service" - TOKEN_COUNTER = "token_counter" diff --git a/flowllm/core/enumeration/role.py b/flowllm/core/enumeration/role.py deleted file mode 100644 index 30e8adf..0000000 --- a/flowllm/core/enumeration/role.py +++ /dev/null @@ -1,12 +0,0 @@ -"""Role enumeration for chat message roles.""" - -from enum import Enum - - -class Role(str, Enum): - """Enumeration of roles used in chat messages.""" - - SYSTEM = "system" - USER = "user" - ASSISTANT = "assistant" - TOOL = "tool" diff --git a/flowllm/core/flow/__init__.py b/flowllm/core/flow/__init__.py deleted file mode 100644 index 282216f..0000000 --- a/flowllm/core/flow/__init__.py +++ /dev/null @@ -1,28 +0,0 @@ -"""Flow composition module for FlowLLM. - -This package provides base classes and helpers for constructing and executing -operation graphs (flows). It includes: - -- BaseFlow: Abstract base providing sync/async execution and streaming support -- BaseToolFlow: BaseFlow variant that exposes a tool-call descriptor -- CmdFlow: Flow built from a user-provided expression string -- ExpressionToolFlow: Tool-enabled flow constructed from expression content - -Common capabilities: -- Sequential and parallel composition via operator overloading -- Streaming and non-streaming responses -- Async and sync execution paths -- Structured error handling -""" - -from .base_flow import BaseFlow -from .base_tool_flow import BaseToolFlow -from .cmd_flow import CmdFlow -from .expression_tool_flow import ExpressionToolFlow - -__all__ = [ - "BaseFlow", - "BaseToolFlow", - "CmdFlow", - "ExpressionToolFlow", -] diff --git a/flowllm/core/flow/base_flow.py b/flowllm/core/flow/base_flow.py deleted file mode 100644 index d5183bd..0000000 --- a/flowllm/core/flow/base_flow.py +++ /dev/null @@ -1,326 +0,0 @@ -"""Core Flow base class and sync/async execution helpers. - -Defines `BaseFlow`, which builds an operation tree (`BaseOp`) and provides -sync (`call`) and async (`async_call`) entry points with optional streaming -support and structured error handling. - -Caching: -- Non-streaming requests can be cached by enabling `enable_cache`. -- Cache key is an SHA-256 hash of the JSON-serialized request params (`**kwargs`). -- If params are not serializable, a serialization exception is logged and the - call proceeds without caching. -- Cached value is the `FlowResponse` serialized via `model_dump()`, stored with - an expiration configured by `cache_expire_hours`. -- Streaming (`stream=True`) is not cached. -""" - -import asyncio -import hashlib -import json -from abc import ABC -from functools import partial -from typing import Union, Optional - -from loguru import logger - -from ..context import FlowContext, C, BaseContext -from ..enumeration import ChunkEnum -from ..op import BaseOp, SequentialOp, ParallelOp, BaseAsyncOp -from ..schema import FlowResponse, FlowStreamChunk -from ..storage import CacheHandler -from ..utils import camel_to_snake - - -class BaseFlow(ABC): - """Abstract base class for all flows. - - Subclasses should implement `build_flow` to return a composed `BaseOp`. - Instances support both streaming and non-streaming responses and can run - either in the current thread or in an async loop depending on the - underlying op's async capability. - - Caching is available for non-streaming paths when `enable_cache=True`. - """ - - def __init__( - self, - name: str = "", - stream: bool = False, - raise_exception: bool = True, - enable_cache: bool = False, - cache_path: str = "cache/{flow_name}", - cache_expire_hours: float = 0.1, - **kwargs, - ): - """Initialize a flow instance. - - Args: - name: Flow name; defaults to the snake-cased class name. - stream: Whether to stream output chunks. - raise_exception: If False, capture exceptions into the response. - enable_cache: Whether to enable response caching (non-stream only). - cache_path: Cache storage path template with {flow_name} placeholder. - cache_expire_hours: Cache expiration time in hours. - **kwargs: Extra parameters passed to the flow context. - """ - self.name: str = name or camel_to_snake(self.__class__.__name__) - self.stream: bool = stream - self.raise_exception: bool = raise_exception - self.flow_params: dict = kwargs - - self._flow_op: Optional[BaseOp] = None - self.flow_printed: bool = False - self.enable_cache: bool = enable_cache - self.cache_path: str = cache_path - self.cache_expire_hours: float = cache_expire_hours - self._cache: CacheHandler | None = None - - @property - def cache(self) -> CacheHandler | None: - """Lazily initialize and return the cache handler. - - Returns: - A `CacheHandler` instance when caching is enabled; otherwise `None`. - """ - if self.enable_cache and self._cache is None: - self._cache = CacheHandler(self.cache_path.format(flow_name=self.name)) - return self._cache - - @staticmethod - def _compute_cache_key(params: dict) -> Optional[str]: - """Compute a stable cache key from request params. - - The key is `sha256(json.dumps(params, sort_keys=True))`. If serialization - fails (e.g., params contain non-serializable objects), the error is logged - and `None` is returned which disables caching for this call. - - Args: - params: Request keyword arguments to serialize. - - Returns: - Hex digest string key or `None` on serialization failure. - """ - try: - payload = json.dumps(params, sort_keys=True, ensure_ascii=False) - return hashlib.sha256(payload.encode("utf-8")).hexdigest() - - except Exception as e: - logger.exception(f"base_flow cache key serialization failed, error={e}") - return None - - def _maybe_load_cached(self, params: dict) -> Optional[FlowResponse]: - """Return cached response if available for non-streaming requests. - - Args: - params: Request keyword arguments to compute the cache key. - - Returns: - A `FlowResponse` if found in cache; otherwise `None`. - """ - if not self.enable_cache or self.stream: - return None - - key = self._compute_cache_key(params) - if not key: - return None - - cached = self.cache.load(key) if self.cache else None - if cached is not None: - logger.info(f"load flow response from cache with params={params}") - return FlowResponse(**cached) - - return None - - def _maybe_save_cache(self, params: dict, response: FlowResponse): - """Save a response into cache for non-streaming requests. - - Args: - params: Request keyword arguments to compute the cache key. - response: The `FlowResponse` to persist. Stored as `model_dump()`. - """ - if not self.enable_cache or self.stream: - return - - key = self._compute_cache_key(params) - if not key: - return - - if self.cache: - self.cache.save(key, response.model_dump(), expire_hours=self.cache_expire_hours) - - @property - def async_mode(self) -> bool: - """Return whether the built op supports async execution.""" - return self.flow_op.async_mode - - def build_flow(self) -> BaseOp: - """Build and return the root `BaseOp` for this flow. - - Subclasses must override this to construct the operation tree. - """ - - @property - def flow_op(self): - """Lazily build and cache the root operation for this flow.""" - if self._flow_op is None: - self._flow_op = self.build_flow() - return self._flow_op - - def print_flow(self): - """Pretty-print the operation tree for debugging/logging once.""" - if not self.flow_printed: - logger.info(f"---------- start print flow={self.name} ----------") - self._print_operation_tree(self.flow_op, indent=0) - logger.info(f"---------- end print flow={self.name} ----------") - self.flow_printed = True - - def _print_operation_tree(self, op: BaseOp, indent: int): - """Recursively log the structure of an operation tree. - - Args: - op: Operation node to print. - indent: Current indentation level for formatting. - """ - prefix = " " * indent - if isinstance(op, SequentialOp): - logger.info(f"{prefix}Sequential Execution:") - for i, sub_op in enumerate(op.ops): - logger.info(f"{prefix} Step {i + 1}:") - self._print_operation_tree(sub_op, indent + 2) - - elif isinstance(op, ParallelOp): - logger.info(f"{prefix}Parallel Execution:") - for i, sub_op in enumerate(op.ops): - logger.info(f"{prefix} Branch {i + 1}:") - self._print_operation_tree(sub_op, indent + 2) - - else: - logger.info(f"{prefix}Operation: {op.name}") - if op.ops: - if isinstance(op.ops, list): - for i, sub_op in enumerate(op.ops): - logger.info(f"{prefix} Sub {i + 1}:") - self._print_operation_tree(sub_op, indent + 2) - elif isinstance(op.ops, BaseContext): - for key, sub_op in op.ops.items(): - logger.info(f"{prefix} Sub {key}:") - self._print_operation_tree(sub_op, indent + 2) - - async def _async_call(self, context: FlowContext) -> Union[FlowResponse | FlowStreamChunk | None]: - """Internal async executor that handles streaming and errors. - - Args: - context: Flow execution context. - - Returns: - `FlowResponse`, `FlowStreamChunk` (queue), or None when appropriate. - """ - self.print_flow() - - # each time rebuild flow - flow_op: BaseOp = self.build_flow() - - if self.async_mode: - assert isinstance(flow_op, BaseAsyncOp) - await flow_op.async_call(context=context) - - else: - loop = asyncio.get_event_loop() - op_call_fn = partial(flow_op.call, context=context) - await loop.run_in_executor(executor=C.thread_pool, func=op_call_fn) - - if self.stream: - await context.add_stream_done() - return context.stream_queue - else: - return context.response - - async def async_call(self, **kwargs) -> Union[FlowResponse | FlowStreamChunk | None]: - """Public async entry point for executing the flow. - - Keyword Args are forwarded to `FlowContext`. - """ - kwargs["stream"] = self.stream - - # Cache only for non-streaming - if not self.stream: - cached = self._maybe_load_cached(kwargs) - if cached is not None: - return cached - - context = FlowContext(**kwargs) - logger.info(f"request.params={kwargs}") - - if self.raise_exception: - result = await self._async_call(context=context) - if not self.stream: - self._maybe_save_cache(kwargs, result) - return result - - try: - result = await self._async_call(context=context) - if not self.stream: - self._maybe_save_cache(kwargs, result) - return result - - except Exception as e: - logger.exception(f"flow_name={self.name} async call encounter error={e.args}") - - if self.stream: - await context.add_stream_chunk_and_type(str(e), ChunkEnum.ERROR) - await context.add_stream_done() - return context.stream_queue - - else: - context.add_response_error(e) - return context.response - - def _call(self, context: FlowContext) -> FlowResponse: - """Internal sync executor for flows without streaming.""" - self.print_flow() - - # each time rebuild flow - flow_op: BaseOp = self.build_flow() - - if self.async_mode: - assert isinstance(flow_op, BaseAsyncOp) - asyncio.run(flow_op.async_call(context=context)) - - else: - flow_op.call(context=context) - - return context.response - - def call(self, **kwargs) -> FlowResponse: - """Public sync entry point for executing the flow. - - Keyword Args are forwarded to `FlowContext`. - """ - kwargs["stream"] = self.stream - - # Cache only for non-streaming - if not self.stream: - cached = self._maybe_load_cached(kwargs) - if cached is not None: - return cached - - context = FlowContext(**kwargs) - logger.info(f"request.params={kwargs}") - - if self.raise_exception: - result = self._call(context=context) - if not self.stream: - self._maybe_save_cache(kwargs, result) - return result - - try: - result = self._call(context=context) - if not self.stream: - self._maybe_save_cache(kwargs, result) - return result - - except Exception as e: - logger.exception(f"flow_name={self.name} call encounter error={e.args}") - - context.add_response_error(e) - return context.response diff --git a/flowllm/core/flow/base_tool_flow.py b/flowllm/core/flow/base_tool_flow.py deleted file mode 100644 index 95ff8fd..0000000 --- a/flowllm/core/flow/base_tool_flow.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Base class for flows exposing a tool-call interface. - -Provides `tool_call` construction for flows intended to be surfaced as tools -to external orchestrators or agents. -""" - -from abc import ABC - -from .base_flow import BaseFlow -from ..schema import ToolCall - - -class BaseToolFlow(BaseFlow, ABC): - """Abstract base flow that also exposes a `ToolCall` descriptor.""" - - def __init__(self, **kwargs): - """Initialize the tool flow.""" - super().__init__(**kwargs) - self._tool_call: ToolCall | None = None - - def build_tool_call(self) -> ToolCall: - """Build and return the `ToolCall` describing this flow.""" - - @property - def tool_call(self) -> ToolCall: - """Lazily build and cache the exported `ToolCall`.""" - if self._tool_call is None: - self._tool_call = self.build_tool_call() - return self._tool_call diff --git a/flowllm/core/flow/cmd_flow.py b/flowllm/core/flow/cmd_flow.py deleted file mode 100644 index 92aa8df..0000000 --- a/flowllm/core/flow/cmd_flow.py +++ /dev/null @@ -1,22 +0,0 @@ -"""Flow implementation that parses a string expression at runtime.""" - -from .base_flow import BaseFlow -from ..utils import parse_flow_expression - - -class CmdFlow(BaseFlow): - """Build a flow from a user-provided expression string.""" - - def __init__(self, flow: str = "", **kwargs): - """Initialize with a flow expression. - - Args: - flow: Expression to parse into a composed `BaseOp`. - """ - super().__init__(**kwargs) - self.flow = flow - assert flow, "add `flow=` in cmd!" - - def build_flow(self): - """Parse and return the operation built from the expression.""" - return parse_flow_expression(self.flow) diff --git a/flowllm/core/flow/expression_tool_flow.py b/flowllm/core/flow/expression_tool_flow.py deleted file mode 100644 index 65d29b7..0000000 --- a/flowllm/core/flow/expression_tool_flow.py +++ /dev/null @@ -1,50 +0,0 @@ -"""Flow that builds from expression content and exposes a tool-call schema.""" - -from .base_tool_flow import BaseToolFlow -from ..schema import FlowConfig, ToolCall -from ..utils import parse_flow_expression - - -class ExpressionToolFlow(BaseToolFlow): - """Tool-enabled flow constructed from `FlowConfig.flow_content`. - - The flow forwards cache-related parameters from `FlowConfig` to `BaseFlow`, - enabling non-streaming response caching when configured. - """ - - def __init__(self, flow_config: FlowConfig = None, **kwargs): - """Initialize the flow with a `FlowConfig`. - - Args: - flow_config: Configuration containing expression, metadata, and - optional caching controls (`enable_cache`, `cache_path`, - `cache_expire_hours`). - """ - self.flow_config: FlowConfig = flow_config - super().__init__( - name=flow_config.name, - stream=self.flow_config.stream, - enable_cache=self.flow_config.enable_cache, - cache_path=self.flow_config.cache_path, - cache_expire_hours=self.flow_config.cache_expire_hours, - **kwargs, - ) - - def build_flow(self): - """Parse and return the operation tree from the config content.""" - return parse_flow_expression(self.flow_config.flow_content) - - def build_tool_call(self) -> ToolCall: - """Construct and return the `ToolCall` for this flow. - - If the underlying op already defines a `tool_call`, reuse it; otherwise, - create a `ToolCall` using the metadata from `FlowConfig`. - """ - if hasattr(self.flow_op, "tool_call"): - return self.flow_op.tool_call - else: - return ToolCall( - name=self.flow_config.name, - description=self.flow_config.description, - input_schema=self.flow_config.input_schema, - ) diff --git a/flowllm/core/llm/__init__.py b/flowllm/core/llm/__init__.py deleted file mode 100644 index 1e1f6f0..0000000 --- a/flowllm/core/llm/__init__.py +++ /dev/null @@ -1,25 +0,0 @@ -"""LLM (Large Language Model) module for flowllm. - -This module provides base classes and implementations for interacting with -various Large Language Model providers. It includes: - -- BaseLLM: Abstract base class defining the common interface for all LLM implementations -- OpenAICompatibleLLM: Implementation for OpenAI-compatible APIs -- LiteLLM: Implementation using LiteLLM library for 100+ LLM providers - -All implementations support: -- Streaming responses -- Tool/function calling -- Async operations -- Error handling and retries -""" - -from .base_llm import BaseLLM -from .lite_llm import LiteLLM -from .openai_compatible_llm import OpenAICompatibleLLM - -__all__ = [ - "BaseLLM", - "LiteLLM", - "OpenAICompatibleLLM", -] diff --git a/flowllm/core/llm/base_llm.py b/flowllm/core/llm/base_llm.py deleted file mode 100644 index 04605a8..0000000 --- a/flowllm/core/llm/base_llm.py +++ /dev/null @@ -1,268 +0,0 @@ -"""Base LLM class definition for flowllm. - -This module defines the BaseLLM abstract base class that serves as the -foundation for all LLM implementations in the flowllm framework. It provides -a standardized interface for interacting with various LLM providers while -handling common concerns like retries, error handling, and streaming. -""" - -import asyncio -import time -from abc import ABC -from typing import List, Callable, Optional, Generator, AsyncGenerator, Union, Any - -from loguru import logger - -from ..schema import FlowStreamChunk -from ..schema import Message -from ..schema import ToolCall - - -class BaseLLM(ABC): - """ - Abstract base class for Large Language Model (LLM) implementations. - - This class defines the common interface and configuration parameters - that all LLM implementations should support. It provides a standardized - way to interact with different LLM providers while handling common - concerns like retries, error handling, and streaming. - """ - - def __init__(self, model_name: str, max_retries: int = 5, raise_exception: bool = False, **kwargs): - """ - Initialize the BaseLLM. - - Args: - model_name: Name of the LLM model to use - max_retries: Maximum number of retry attempts on failure - raise_exception: Whether to raise exceptions or return default values - """ - self.model_name: str = model_name - self.max_retries: int = max_retries - self.raise_exception: bool = raise_exception - self.kwargs: dict = kwargs - # ref: https://help.aliyun.com/zh/model-studio/qwen-api-reference - - def stream_chat( - self, - messages: List[Message], - tools: Optional[List[ToolCall]] = None, - **kwargs, - ) -> Generator[FlowStreamChunk, None, None]: - """ - Stream chat completions from the LLM. - - This method should yield chunks of the response as they become available, - allowing for real-time display of the model's output. - - Args: - messages: List of conversation messages - tools: Optional list of tools the model can use - **kwargs: Additional model-specific parameters - - Yields: - FlowStreamChunk for each streaming piece. - FlowStreamChunk contains chunk_type, chunk content, and metadata. - """ - raise NotImplementedError - - async def astream_chat( - self, - messages: List[Message], - tools: Optional[List[ToolCall]] = None, - **kwargs, - ) -> AsyncGenerator[FlowStreamChunk, None]: - """ - Async stream chat completions from the LLM. - - This method should yield chunks of the response as they become available, - allowing for real-time display of the model's output in async contexts. - - Args: - messages: List of conversation messages - tools: Optional list of tools the model can use - **kwargs: Additional model-specific parameters - - Yields: - FlowStreamChunk for each streaming piece. - FlowStreamChunk contains chunk_type, chunk content, and metadata. - """ - raise NotImplementedError - - def _chat( - self, - messages: List[Message], - tools: Optional[List[ToolCall]] = None, - enable_stream_print: bool = False, - **kwargs, - ) -> Message: - """ - Internal method to perform a single chat completion. - - This method should be implemented by subclasses to handle the actual - communication with the LLM provider. It's called by the public chat() - method which adds retry logic and error handling. - - Args: - messages: List of conversation messages - tools: Optional list of tools the model can use - enable_stream_print: Whether to print streaming response to console - **kwargs: Additional model-specific parameters - - Returns: - The complete response message from the LLM - """ - raise NotImplementedError - - async def _achat( - self, - messages: List[Message], - tools: Optional[List[ToolCall]] = None, - enable_stream_print: bool = False, - **kwargs, - ) -> Message: - """ - Internal async method to perform a single chat completion. - - This method should be implemented by subclasses to handle the actual - async communication with the LLM provider. It's called by the public achat() - method which adds retry logic and error handling. - - Args: - messages: List of conversation messages - tools: Optional list of tools the model can use - enable_stream_print: Whether to print streaming response to console - **kwargs: Additional model-specific parameters - - Returns: - The complete response message from the LLM - """ - raise NotImplementedError - - def chat( - self, - messages: List[Message], - tools: Optional[List[ToolCall]] = None, - enable_stream_print: bool = False, - callback_fn: Optional[Callable[[Message], Any]] = None, - default_value: Any = None, - **kwargs, - ) -> Union[Message, Any]: - """ - Perform a chat completion with retry logic and error handling. - - This is the main public interface for chat completions. It wraps the - internal _chat() method with robust error handling, exponential backoff, - and optional callback processing. - - Args: - messages: List of conversation messages - tools: Optional list of tools the model can use - callback_fn: Optional callback to process the response message - default_value: Value to return if all retries fail (when raise_exception=False) - enable_stream_print: Whether to print streaming response to console - **kwargs: Additional model-specific parameters - - Returns: - The response message (possibly processed by callback_fn) or default_value - - Raises: - Exception: If raise_exception=True and all retries fail - """ - for i in range(self.max_retries): - try: - # Attempt to get response from the model - message: Message = self._chat( - messages=messages, - tools=tools, - enable_stream_print=enable_stream_print, - **kwargs, - ) - - # Apply callback function if provided - if callback_fn: - return callback_fn(message) - else: - return message - - except Exception as e: - logger.exception(f"chat with model={self.model_name} encounter error with e={e.args}") - - # If this is the last retry attempt, handle final failure - if i == self.max_retries - 1: - if self.raise_exception: - raise e - # If raise_exception=False, return default_value without waiting - return default_value - - # Exponential backoff: wait before next retry attempt - time.sleep(1 + i) - - return default_value - - async def achat( - self, - messages: List[Message], - tools: Optional[List[ToolCall]] = None, - enable_stream_print: bool = False, - callback_fn: Optional[Callable[[Message], Any]] = None, - default_value: Any = None, - **kwargs, - ) -> Union[Message, Any]: - """ - Perform an async chat completion with retry logic and error handling. - - This is the main public interface for async chat completions. It wraps the - internal _achat() method with robust error handling, exponential backoff, - and optional callback processing. - - Args: - messages: List of conversation messages - tools: Optional list of tools the model can use - callback_fn: Optional callback to process the response message - default_value: Value to return if all retries fail (when raise_exception=False) - enable_stream_print: Whether to print streaming response to console - **kwargs: Additional model-specific parameters - - Returns: - The response message (possibly processed by callback_fn) or default_value - - Raises: - Exception: If raise_exception=True and all retries fail - """ - for i in range(self.max_retries): - try: - # Attempt to get response from the model - message: Message = await self._achat( - messages=messages, - tools=tools, - enable_stream_print=enable_stream_print, - **kwargs, - ) - - # Apply callback function if provided - if callback_fn: - return callback_fn(message) - else: - return message - - except Exception as e: - logger.exception(f"async chat with model={self.model_name} encounter error with e={e.args}") - - # If this is the last retry attempt, handle final failure - if i == self.max_retries - 1: - if self.raise_exception: - raise e - # If raise_exception=False, return default_value without waiting - return default_value - - # Exponential backoff: wait before next retry attempt - await asyncio.sleep(1 + i) - - return default_value - - def close(self): - """Close the client connection or clean up resources.""" - - async def async_close(self): - """Asynchronously close the client connection or clean up resources.""" diff --git a/flowllm/core/llm/lite_llm.py b/flowllm/core/llm/lite_llm.py deleted file mode 100644 index 24f2962..0000000 --- a/flowllm/core/llm/lite_llm.py +++ /dev/null @@ -1,443 +0,0 @@ -"""LiteLLM implementation for flowllm. - -This module provides an implementation of BaseLLM using the LiteLLM library, -which enables support for 100+ LLM providers through a unified interface. -It handles streaming responses, tool calling, and reasoning content from -supported models. The implementation supports both synchronous and asynchronous -operations with robust error handling and retry logic. -""" - -import asyncio -import json -import os -import time -from typing import List, Dict, Optional, Generator, AsyncGenerator - -from loguru import logger - -from .base_llm import BaseLLM -from ..context import C -from ..enumeration import ChunkEnum -from ..enumeration import Role -from ..schema import FlowStreamChunk -from ..schema import Message -from ..schema import ToolCall - - -@C.register_llm("litellm") -class LiteLLM(BaseLLM): - """ - LiteLLM implementation supporting streaming and tool calls. - - This class implements the BaseLLM interface using LiteLLM, which provides - support for 100+ LLM providers through a unified interface, including: - - Streaming responses with different chunk types (thinking, answer, tools) - - Tool calling with parallel execution - - Reasoning/thinking content from supported models - - Robust error handling and retries - - The class follows the BaseLLM interface strictly, implementing all required methods - with proper type annotations and error handling consistent with the base class. - - The implementation aggregates streaming chunks internally in _chat() and _achat() - methods, which are called by the base class's chat() and achat() methods that add - retry logic and error handling. Reasoning content is separated from regular answer - content and stored in the Message's reasoning_content field. - """ - - def __init__( - self, - model_name: str, - api_key: str | None = None, - base_url: str | None = None, - custom_llm_provider: str = "openai", - max_retries: int = 5, - raise_exception: bool = False, - **kwargs, - ): - """ - Initialize the LiteLLM. - - Args: - model_name: Name of the LLM model to use - api_key: API key for authentication (defaults to FLOW_LLM_API_KEY env var) - base_url: Base URL for custom endpoints (defaults to FLOW_LLM_BASE_URL env var) - custom_llm_provider: Custom LLM provider name for LiteLLM routing - max_retries: Maximum number of retries on failure (default: 5) - raise_exception: Whether to raise exception on final failure (default: False) - **kwargs: Additional parameters passed to the API (e.g., temperature, top_p, seed) - """ - super().__init__(model_name=model_name, max_retries=max_retries, raise_exception=raise_exception, **kwargs) - self.api_key: str = api_key or os.getenv("FLOW_LLM_API_KEY", "") - self.base_url: str = base_url or os.getenv("FLOW_LLM_BASE_URL", "") - self.custom_llm_provider: str = custom_llm_provider - - def stream_chat( - self, - messages: List[Message], - tools: Optional[List[ToolCall]] = None, - **kwargs, - ) -> Generator[FlowStreamChunk, None, None]: - """ - Stream chat completions from LiteLLM. - - This method handles streaming responses and categorizes chunks into different types: - - THINK: Reasoning/thinking content from the model - - ANSWER: Regular response content - - TOOL: Tool calls that need to be executed - - USAGE: Token usage statistics - - ERROR: Error information - - Args: - messages: List of conversation messages - tools: Optional list of tools available to the model - **kwargs: Additional parameters - - Yields: - FlowStreamChunk for each streaming piece. - FlowStreamChunk contains chunk_type and chunk content. - """ - from litellm import completion - - chat_kwargs = { - "api_key": self.api_key, - "base_url": self.base_url, - "model": self.model_name, - "custom_llm_provider": self.custom_llm_provider, - "messages": [x.simple_dump() for x in messages], - "stream": True, - "tools": [x.simple_input_dump() for x in tools] if tools else None, - **self.kwargs, - **kwargs, - } - log_kwargs = {k: v for k, v in chat_kwargs.items() if k != "messages"} - logger.info(f"LiteLLM.stream_chat: {log_kwargs}") - - for i in range(self.max_retries): - try: - completion = completion(**chat_kwargs) - - ret_tools: List[ToolCall] = [] - is_answering: bool = False - - for chunk in completion: - if not chunk.choices: - yield FlowStreamChunk(chunk_type=ChunkEnum.USAGE, chunk=chunk.usage.model_dump()) - - else: - delta = chunk.choices[0].delta - - if hasattr(delta, "reasoning_content") and delta.reasoning_content is not None: - yield FlowStreamChunk(chunk_type=ChunkEnum.THINK, chunk=delta.reasoning_content) - - else: - if not is_answering: - is_answering = True - - if delta.content is not None: - yield FlowStreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) - - if delta.tool_calls is not None: - for tool_call in delta.tool_calls: - index = tool_call.index - - while len(ret_tools) <= index: - ret_tools.append(ToolCall(index=index)) - - if tool_call.id: - ret_tools[index].id += tool_call.id - - if tool_call.function and tool_call.function.name: - ret_tools[index].name += tool_call.function.name - - if tool_call.function and tool_call.function.arguments: - ret_tools[index].arguments += tool_call.function.arguments - - if ret_tools: - tool_dict: Dict[str, ToolCall] = {x.name: x for x in tools} if tools else {} - for tool in ret_tools: - if tool.name not in tool_dict: - continue - - if not tool.check_argument(): - raise ValueError(f"Tool call {tool.name} argument={tool.arguments} are invalid") - - yield FlowStreamChunk(chunk_type=ChunkEnum.TOOL, chunk=tool.simple_output_dump()) - - return - - except Exception as e: - logger.exception(f"stream chat with model={self.model_name} encounter error with e={e.args}") - - if i == self.max_retries - 1: - if self.raise_exception: - raise e - - yield FlowStreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e)) - return - - yield FlowStreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e)) - time.sleep(1 + i) - - async def astream_chat( - self, - messages: List[Message], - tools: Optional[List[ToolCall]] = None, - **kwargs, - ) -> AsyncGenerator[FlowStreamChunk, None]: - """ - Async stream chat completions from LiteLLM. - - This method handles async streaming responses and categorizes chunks into different types: - - THINK: Reasoning/thinking content from the model - - ANSWER: Regular response content - - TOOL: Tool calls that need to be executed - - USAGE: Token usage statistics - - ERROR: Error information - - Args: - messages: List of conversation messages - tools: Optional list of tools available to the model - **kwargs: Additional parameters - - Yields: - FlowStreamChunk for each streaming piece. - FlowStreamChunk contains chunk_type, chunk content, and metadata. - """ - from litellm import acompletion - - chat_kwargs = { - "api_key": self.api_key, - "base_url": self.base_url, - "model": self.model_name, - "custom_llm_provider": self.custom_llm_provider, - "messages": [x.simple_dump() for x in messages], - "stream": True, - "tools": [x.simple_input_dump() for x in tools] if tools else None, - **self.kwargs, - **kwargs, - } - log_kwargs = {k: v for k, v in chat_kwargs.items() if k != "messages"} - logger.info(f"LiteLLM.astream_chat: {log_kwargs}") - - for i in range(self.max_retries): - try: - completion = await acompletion(**chat_kwargs) - - ret_tools: List[ToolCall] = [] - is_answering: bool = False - - async for chunk in completion: - if not chunk.choices: - yield FlowStreamChunk(chunk_type=ChunkEnum.USAGE, chunk=chunk.usage.model_dump()) - - else: - delta = chunk.choices[0].delta - - if hasattr(delta, "reasoning_content") and delta.reasoning_content is not None: - yield FlowStreamChunk(chunk_type=ChunkEnum.THINK, chunk=delta.reasoning_content) - - else: - if not is_answering: - is_answering = True - - if delta.content is not None: - yield FlowStreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) - - if delta.tool_calls is not None: - for tool_call in delta.tool_calls: - index = tool_call.index - - while len(ret_tools) <= index: - ret_tools.append(ToolCall(index=index)) - - if tool_call.id: - ret_tools[index].id += tool_call.id - - if tool_call.function and tool_call.function.name: - ret_tools[index].name += tool_call.function.name - - if tool_call.function and tool_call.function.arguments: - ret_tools[index].arguments += tool_call.function.arguments - - if ret_tools: - tool_dict: Dict[str, ToolCall] = {x.name: x for x in tools} if tools else {} - for tool in ret_tools: - if tool.name not in tool_dict: - continue - - if not tool.check_argument(): - raise ValueError(f"Tool call {tool.name} argument={tool.arguments} are invalid") - - yield FlowStreamChunk(chunk_type=ChunkEnum.TOOL, chunk=tool.simple_output_dump()) - - return - - except Exception as e: - logger.exception(f"stream chat with model={self.model_name} encounter error with e={e.args}") - - if i == self.max_retries - 1: - if self.raise_exception: - raise e - - yield FlowStreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e)) - return - - yield FlowStreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e)) - await asyncio.sleep(1 + i) - - def _chat( - self, - messages: List[Message], - tools: Optional[List[ToolCall]] = None, - enable_stream_print: bool = False, - **kwargs, - ) -> Message: - """ - Internal method to perform a single chat completion by aggregating streaming chunks. - - This method is called by the base class's chat() method which adds retry logic - and error handling. It consumes the entire streaming response from stream_chat() - and combines all chunks into a single Message object. It separates reasoning content, - regular answer content, and tool calls, providing a complete response. - - Args: - messages: List of conversation messages - tools: Optional list of tools available to the model - enable_stream_print: Whether to print streaming response to console - **kwargs: Additional parameters - - Returns: - Complete Message with all content aggregated from streaming chunks - """ - - enter_think = False - enter_answer = False - reasoning_content = "" - answer_content = "" - tool_calls = [] - - for stream_chunk in self.stream_chat(messages, tools, **kwargs): - if stream_chunk.chunk_type is ChunkEnum.USAGE: - chunk = stream_chunk.chunk - if enable_stream_print: - print(f"\n{json.dumps(chunk, ensure_ascii=False, indent=2)}", flush=True) - - elif stream_chunk.chunk_type is ChunkEnum.THINK: - chunk = stream_chunk.chunk - if enable_stream_print: - if not enter_think: - enter_think = True - print("\n", end="", flush=True) - print(chunk, end="", flush=True) - - reasoning_content += chunk - - elif stream_chunk.chunk_type is ChunkEnum.ANSWER: - chunk = stream_chunk.chunk - if enable_stream_print: - if not enter_answer: - enter_answer = True - if enter_think: - print("\n", flush=True) - print(chunk, end="", flush=True) - - answer_content += chunk - - elif stream_chunk.chunk_type is ChunkEnum.TOOL: - chunk = stream_chunk.chunk - if enable_stream_print: - print(f"\n{json.dumps(chunk, ensure_ascii=False, indent=2)}", flush=True) - - tool_calls.append(chunk) - - elif stream_chunk.chunk_type is ChunkEnum.ERROR: - chunk = stream_chunk.chunk - if enable_stream_print: - # Display error information - print(f"\n{chunk}", flush=True) - - return Message( - role=Role.ASSISTANT, - reasoning_content=reasoning_content, - content=answer_content, - tool_calls=tool_calls, - ) - - async def _achat( - self, - messages: List[Message], - tools: Optional[List[ToolCall]] = None, - enable_stream_print: bool = False, - **kwargs, - ) -> Message: - """ - Internal async method to perform a single chat completion by aggregating streaming chunks. - - This method is called by the base class's achat() method which adds retry logic - and error handling. It consumes the entire async streaming response from astream_chat() - and combines all chunks into a single Message object. It separates reasoning content, - regular answer content, and tool calls, providing a complete response. - - Args: - messages: List of conversation messages - tools: Optional list of tools available to the model - enable_stream_print: Whether to print streaming response to console - **kwargs: Additional parameters - - Returns: - Complete Message with all content aggregated from streaming chunks - """ - - enter_think = False - enter_answer = False - reasoning_content = "" - answer_content = "" - tool_calls = [] - - async for stream_chunk in self.astream_chat(messages, tools, **kwargs): - if stream_chunk.chunk_type is ChunkEnum.USAGE: - chunk = stream_chunk.chunk - if enable_stream_print: - print(f"\n{json.dumps(chunk, ensure_ascii=False, indent=2)}", flush=True) - - elif stream_chunk.chunk_type is ChunkEnum.THINK: - chunk = stream_chunk.chunk - if enable_stream_print: - if not enter_think: - enter_think = True - print("\n", end="", flush=True) - print(chunk, end="", flush=True) - - reasoning_content += chunk - - elif stream_chunk.chunk_type is ChunkEnum.ANSWER: - chunk = stream_chunk.chunk - if enable_stream_print: - if not enter_answer: - enter_answer = True - if enter_think: - print("\n", flush=True) - print(chunk, end="", flush=True) - - answer_content += chunk - - elif stream_chunk.chunk_type is ChunkEnum.TOOL: - chunk = stream_chunk.chunk - if enable_stream_print: - print(f"\n{json.dumps(chunk, ensure_ascii=False, indent=2)}", flush=True) - - tool_calls.append(chunk) - - elif stream_chunk.chunk_type is ChunkEnum.ERROR: - chunk = stream_chunk.chunk - if enable_stream_print: - # Display error information - print(f"\n{chunk}", flush=True) - - return Message( - role=Role.ASSISTANT, - reasoning_content=reasoning_content, - content=answer_content, - tool_calls=tool_calls, - ) diff --git a/flowllm/core/llm/openai_compatible_llm.py b/flowllm/core/llm/openai_compatible_llm.py deleted file mode 100644 index 68a7027..0000000 --- a/flowllm/core/llm/openai_compatible_llm.py +++ /dev/null @@ -1,439 +0,0 @@ -"""OpenAI-compatible LLM implementation for flowllm. - -This module provides an implementation of BaseLLM that supports OpenAI-compatible -APIs. It handles streaming responses, tool calling, and reasoning content from -supported models. The implementation supports both synchronous and asynchronous -operations with robust error handling and retry logic. -""" - -import asyncio -import json -import os -import time -from typing import List, Dict, Optional, Generator, AsyncGenerator - -from loguru import logger -from openai import OpenAI, AsyncOpenAI - -from .base_llm import BaseLLM -from ..context import C -from ..enumeration import ChunkEnum -from ..enumeration import Role -from ..schema import FlowStreamChunk -from ..schema import Message -from ..schema import ToolCall - - -@C.register_llm("openai_compatible") -class OpenAICompatibleLLM(BaseLLM): - """ - OpenAI-compatible LLM implementation supporting streaming and tool calls. - - This class implements the BaseLLM interface for OpenAI-compatible APIs, - including support for: - - Streaming responses with different chunk types (thinking, answer, tools) - - Tool calling with parallel execution - - Reasoning/thinking content from supported models - - Robust error handling and retries - - The class follows the BaseLLM interface strictly, implementing all required methods - with proper type annotations and error handling consistent with the base class. - - The implementation aggregates streaming chunks internally in _chat() and _achat() - methods, which are called by the base class's chat() and achat() methods that add - retry logic and error handling. Reasoning content is separated from regular answer - content and stored in the Message's reasoning_content field. - """ - - def __init__( - self, - model_name: str, - api_key: str | None = None, - base_url: str | None = None, - max_retries: int = 5, - raise_exception: bool = False, - **kwargs, - ): - """ - Initialize the OpenAICompatibleLLM. - - Args: - model_name: Name of the LLM model to use - api_key: API key for authentication (defaults to FLOW_LLM_API_KEY env var) - base_url: Base URL for the API endpoint (defaults to FLOW_LLM_BASE_URL env var) - max_retries: Maximum number of retries on failure (default: 5) - raise_exception: Whether to raise exception on final failure (default: False) - **kwargs: Additional parameters passed to the API - """ - super().__init__(model_name=model_name, max_retries=max_retries, raise_exception=raise_exception, **kwargs) - self.api_key: str = api_key or os.getenv("FLOW_LLM_API_KEY", "") - self.base_url: str = base_url or os.getenv("FLOW_LLM_BASE_URL", "") - - # Initialize OpenAI clients - self._client = OpenAI(api_key=self.api_key, base_url=self.base_url) - self._aclient = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url) - - def stream_chat( - self, - messages: List[Message], - tools: Optional[List[ToolCall]] = None, - **kwargs, - ) -> Generator[FlowStreamChunk, None, None]: - """ - Stream chat completions from OpenAI-compatible API. - - This method handles streaming responses and categorizes chunks into different types: - - THINK: Reasoning/thinking content from the model - - ANSWER: Regular response content - - TOOL: Tool calls that need to be executed - - USAGE: Token usage statistics - - ERROR: Error information - - Args: - messages: List of conversation messages - tools: Optional list of tools available to the model - **kwargs: Additional parameters - - Yields: - FlowStreamChunk for each streaming piece. - FlowStreamChunk contains chunk_type and chunk content. - """ - chat_kwargs = { - "model": self.model_name, - "messages": [x.simple_dump() for x in messages], - "stream": True, - "tools": [x.simple_input_dump() for x in tools] if tools else None, - **self.kwargs, - **kwargs, - } - log_kwargs = {k: v for k, v in chat_kwargs.items() if k != "messages"} - logger.info(f"OpenAICompatibleLLM.stream_chat: {log_kwargs}") - - for i in range(self.max_retries): - try: - completion = self._client.chat.completions.create(**chat_kwargs) - - ret_tools: List[ToolCall] = [] - is_answering: bool = False - - for chunk in completion: - if not chunk.choices: - yield FlowStreamChunk(chunk_type=ChunkEnum.USAGE, chunk=chunk.usage.model_dump()) - - else: - delta = chunk.choices[0].delta - - if hasattr(delta, "reasoning_content") and delta.reasoning_content: - yield FlowStreamChunk(chunk_type=ChunkEnum.THINK, chunk=delta.reasoning_content) - - else: - if not is_answering: - is_answering = True - - if delta.content is not None: - yield FlowStreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) - - if delta.tool_calls is not None: - for tool_call in delta.tool_calls: - index = tool_call.index - - while len(ret_tools) <= index: - ret_tools.append(ToolCall(index=index)) - - if tool_call.id: - ret_tools[index].id += tool_call.id - - if tool_call.function and tool_call.function.name: - ret_tools[index].name += tool_call.function.name - - if tool_call.function and tool_call.function.arguments: - ret_tools[index].arguments += tool_call.function.arguments - - if ret_tools: - tool_dict: Dict[str, ToolCall] = {x.name: x for x in tools} if tools else {} - for tool in ret_tools: - if tool.name not in tool_dict: - continue - - if not tool.check_argument(): - raise ValueError(f"Tool call {tool.name} argument={tool.arguments} are invalid") - - yield FlowStreamChunk(chunk_type=ChunkEnum.TOOL, chunk=tool.simple_output_dump()) - - return - - except Exception as e: - logger.exception(f"stream chat with model={self.model_name} encounter error with e={e.args}") - - if i == self.max_retries - 1: - if self.raise_exception: - raise e - - yield FlowStreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e)) - return - - yield FlowStreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e)) - time.sleep(1 + i) - - async def astream_chat( - self, - messages: List[Message], - tools: Optional[List[ToolCall]] = None, - **kwargs, - ) -> AsyncGenerator[FlowStreamChunk, None]: - """ - Async stream chat completions from OpenAI-compatible API. - - This method handles async streaming responses and categorizes chunks into different types: - - THINK: Reasoning/thinking content from the model - - ANSWER: Regular response content - - TOOL: Tool calls that need to be executed - - USAGE: Token usage statistics - - ERROR: Error information - - Args: - messages: List of conversation messages - tools: Optional list of tools available to the model - **kwargs: Additional parameters - - Yields: - FlowStreamChunk for each streaming piece. - FlowStreamChunk contains chunk_type, chunk content, and metadata. - """ - chat_kwargs = { - "model": self.model_name, - "messages": [x.simple_dump() for x in messages], - "stream": True, - "tools": [x.simple_input_dump() for x in tools] if tools else None, - **self.kwargs, - **kwargs, - } - log_kwargs = {k: v for k, v in chat_kwargs.items() if k != "messages"} - logger.info(f"OpenAICompatibleLLM.astream_chat: {log_kwargs}") - - for i in range(self.max_retries): - try: - completion = await self._aclient.chat.completions.create(**chat_kwargs) - - ret_tools: List[ToolCall] = [] - is_answering: bool = False - - async for chunk in completion: - if not chunk.choices: - yield FlowStreamChunk(chunk_type=ChunkEnum.USAGE, chunk=chunk.usage.model_dump()) - - else: - delta = chunk.choices[0].delta - - if hasattr(delta, "reasoning_content") and delta.reasoning_content: - yield FlowStreamChunk(chunk_type=ChunkEnum.THINK, chunk=delta.reasoning_content) - - else: - if not is_answering: - is_answering = True - - if delta.content is not None: - yield FlowStreamChunk(chunk_type=ChunkEnum.ANSWER, chunk=delta.content) - - if delta.tool_calls is not None: - for tool_call in delta.tool_calls: - index = tool_call.index - - while len(ret_tools) <= index: - ret_tools.append(ToolCall(index=index)) - - if tool_call.id: - ret_tools[index].id += tool_call.id - - if tool_call.function and tool_call.function.name: - ret_tools[index].name += tool_call.function.name - - if tool_call.function and tool_call.function.arguments: - ret_tools[index].arguments += tool_call.function.arguments - - if ret_tools: - tool_dict: Dict[str, ToolCall] = {x.name: x for x in tools} if tools else {} - for tool in ret_tools: - if tool.name not in tool_dict: - continue - - if not tool.check_argument(): - raise ValueError(f"Tool call {tool.name} argument={tool.arguments} are invalid") - - yield FlowStreamChunk(chunk_type=ChunkEnum.TOOL, chunk=tool.simple_output_dump()) - - return - - except Exception as e: - logger.exception(f"stream chat with model={self.model_name} encounter error with e={e.args}") - - if i == self.max_retries - 1: - if self.raise_exception: - raise e - - yield FlowStreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e)) - return - - yield FlowStreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e)) - await asyncio.sleep(1 + i) - - def _chat( - self, - messages: List[Message], - tools: Optional[List[ToolCall]] = None, - enable_stream_print: bool = False, - **kwargs, - ) -> Message: - """ - Internal method to perform a single chat completion by aggregating streaming chunks. - - This method is called by the base class's chat() method which adds retry logic - and error handling. It consumes the entire streaming response from stream_chat() - and combines all chunks into a single Message object. It separates reasoning content, - regular answer content, and tool calls, providing a complete response. - - Args: - messages: List of conversation messages - tools: Optional list of tools available to the model - enable_stream_print: Whether to print streaming response to console - **kwargs: Additional parameters - - Returns: - Complete Message with all content aggregated from streaming chunks - """ - - enter_think = False - enter_answer = False - reasoning_content = "" - answer_content = "" - tool_calls = [] - - for stream_chunk in self.stream_chat(messages, tools, **kwargs): - if stream_chunk.chunk_type is ChunkEnum.USAGE: - chunk = stream_chunk.chunk - if enable_stream_print: - print(f"\n{json.dumps(chunk, ensure_ascii=False, indent=2)}", flush=True) - - elif stream_chunk.chunk_type is ChunkEnum.THINK: - chunk = stream_chunk.chunk - if enable_stream_print: - if not enter_think: - enter_think = True - print("\n", end="", flush=True) - print(chunk, end="", flush=True) - - reasoning_content += chunk - - elif stream_chunk.chunk_type is ChunkEnum.ANSWER: - chunk = stream_chunk.chunk - if enable_stream_print: - if not enter_answer: - enter_answer = True - if enter_think: - print("\n", flush=True) - print(chunk, end="", flush=True) - - answer_content += chunk - - elif stream_chunk.chunk_type is ChunkEnum.TOOL: - chunk = stream_chunk.chunk - if enable_stream_print: - print(f"\n{json.dumps(chunk, ensure_ascii=False, indent=2)}", flush=True) - - tool_calls.append(chunk) - - elif stream_chunk.chunk_type is ChunkEnum.ERROR: - chunk = stream_chunk.chunk - if enable_stream_print: - # Display error information - print(f"\n{chunk}", flush=True) - - return Message( - role=Role.ASSISTANT, - reasoning_content=reasoning_content, - content=answer_content, - tool_calls=tool_calls, - ) - - async def _achat( - self, - messages: List[Message], - tools: Optional[List[ToolCall]] = None, - enable_stream_print: bool = False, - **kwargs, - ) -> Message: - """ - Internal async method to perform a single chat completion by aggregating streaming chunks. - - This method is called by the base class's achat() method which adds retry logic - and error handling. It consumes the entire async streaming response from astream_chat() - and combines all chunks into a single Message object. It separates reasoning content, - regular answer content, and tool calls, providing a complete response. - - Args: - messages: List of conversation messages - tools: Optional list of tools available to the model - enable_stream_print: Whether to print streaming response to console - **kwargs: Additional parameters - - Returns: - Complete Message with all content aggregated from streaming chunks - """ - - enter_think = False - enter_answer = False - reasoning_content = "" - answer_content = "" - tool_calls = [] - - async for stream_chunk in self.astream_chat(messages, tools, **kwargs): - if stream_chunk.chunk_type is ChunkEnum.USAGE: - chunk = stream_chunk.chunk - if enable_stream_print: - print(f"\n{json.dumps(chunk, ensure_ascii=False, indent=2)}", flush=True) - - elif stream_chunk.chunk_type is ChunkEnum.THINK: - chunk = stream_chunk.chunk - if enable_stream_print: - if not enter_think: - enter_think = True - print("\n", end="", flush=True) - print(chunk, end="", flush=True) - - reasoning_content += chunk - - elif stream_chunk.chunk_type is ChunkEnum.ANSWER: - chunk = stream_chunk.chunk - if enable_stream_print: - if not enter_answer: - enter_answer = True - if enter_think: - print("\n", flush=True) - print(chunk, end="", flush=True) - - answer_content += chunk - - elif stream_chunk.chunk_type is ChunkEnum.TOOL: - chunk = stream_chunk.chunk - if enable_stream_print: - print(f"\n{json.dumps(chunk, ensure_ascii=False, indent=2)}", flush=True) - - tool_calls.append(chunk) - - elif stream_chunk.chunk_type is ChunkEnum.ERROR: - chunk = stream_chunk.chunk - if enable_stream_print: - print(f"\n{chunk}", flush=True) - - return Message( - role=Role.ASSISTANT, - reasoning_content=reasoning_content, - content=answer_content, - tool_calls=tool_calls, - ) - - def close(self): - self._client.close() - - async def async_close(self): - await self._aclient.close() diff --git a/flowllm/core/op/__init__.py b/flowllm/core/op/__init__.py deleted file mode 100644 index 0b3cb8e..0000000 --- a/flowllm/core/op/__init__.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Operation classes for flowllm framework. - -This package provides the base classes and utilities for creating and composing -operations in the flowllm framework. Operations can be executed synchronously -or asynchronously, and can be composed sequentially or in parallel. -""" - -from .base_async_op import BaseAsyncOp -from .base_async_tool_op import BaseAsyncToolOp -from .base_mcp_op import BaseMcpOp -from .base_op import BaseOp -from .base_ray_op import BaseRayOp -from .parallel_op import ParallelOp -from .sequential_op import SequentialOp - -__all__ = [ - "BaseOp", - "BaseAsyncOp", - "BaseAsyncToolOp", - "BaseMcpOp", - "BaseRayOp", - "SequentialOp", - "ParallelOp", -] diff --git a/flowllm/core/op/base_async_op.py b/flowllm/core/op/base_async_op.py deleted file mode 100644 index 5745ae6..0000000 --- a/flowllm/core/op/base_async_op.py +++ /dev/null @@ -1,300 +0,0 @@ -"""Base class for asynchronous operations. - -This module provides the BaseAsyncOp class, which extends BaseOp with -asynchronous execution capabilities. It supports async cache operations, -async task submission and joining, and full async execution lifecycle. -""" - -import asyncio -from abc import ABCMeta, abstractmethod -from typing import Any, Callable - -from loguru import logger - -from .base_op import BaseOp -from ..context import FlowContext, C - - -class BaseAsyncOp(BaseOp, metaclass=ABCMeta): - """Base class for asynchronous operations. - - This class extends BaseOp to provide asynchronous execution capabilities. - All operations created from this class run in async mode by default. - - Example: - ```python - class MyAsyncOp(BaseAsyncOp): - async def async_execute(self): - return await some_async_function() - - op = MyAsyncOp() - result = await op.async_call() - ``` - """ - - def __init__(self, **kwargs): - """Initialize the async operation. - - Automatically sets async_mode to True if not explicitly set. - - Args: - **kwargs: Arguments passed to BaseOp.__init__ - """ - kwargs.setdefault("async_mode", True) - super().__init__(**kwargs) - - def execute(self): - """Placeholder for synchronous execute method. - - This method is not used in async operations. Subclasses should implement - `async_execute()` instead. This method exists only to satisfy the abstract - method requirement from BaseOp. - """ - - async def async_save_load_cache(self, key: str, fn: Callable, **kwargs): - """Save or load from cache asynchronously. - - If caching is enabled, checks cache first. If not found, executes the - function (async or sync) and saves the result. Supports both coroutine - functions and regular functions (executed in thread pool). - - Args: - key: Cache key for storing/retrieving the result - fn: Function to execute if cache miss (can be async or sync) - **kwargs: Additional arguments for cache load operation - - Returns: - Cached result if available, otherwise result from function execution - """ - if self.enable_cache: - result = self.cache.load(key, **kwargs) - if result is None: - if asyncio.iscoroutinefunction(fn): - result = await fn() - else: - loop = asyncio.get_event_loop() - result = await loop.run_in_executor(C.thread_pool, fn) - - self.cache.save(key, result, expire_hours=self.cache_expire_hours) - else: - logger.info(f"load {key} from cache") - else: - if asyncio.iscoroutinefunction(fn): - result = await fn() - else: - loop = asyncio.get_event_loop() - result = await loop.run_in_executor(C.thread_pool, fn) - - return result - - async def async_before_execute(self): - """Hook method called before async_execute(). Override in subclasses. - - This method is called automatically by `async_call()` before executing - the main `async_execute()` method. Use this to perform any setup, - validation, or preprocessing needed before execution. - - Example: - ```python - async def async_before_execute(self): - # Validate inputs - if not self.context.get("input"): - raise ValueError("Input is required") - ``` - """ - - async def async_after_execute(self): - """Hook method called after async_execute(). Override in subclasses. - - This method is called automatically by `async_call()` after successfully - executing the main `async_execute()` method. Use this to perform any - cleanup, post-processing, or result transformation. - - Example: - ```python - async def async_after_execute(self): - # Post-process results - if self.context.response: - self.context.response.answer = self.context.response.answer.upper() - ``` - """ - - @abstractmethod - async def async_execute(self): - """Main async execution method. Must be implemented in subclasses. - - Returns: - Execution result - """ - - async def async_default_execute(self, e: Exception = None, **kwargs): - """Default async execution method when main execution fails. Override in subclasses. - - This method is called when `async_execute()` fails and `raise_exception` - is False. It provides a fallback mechanism to return a default result - instead of raising an exception. - - Args: - e: The exception that was raised during execution (if any) - **kwargs: Additional keyword arguments - - Returns: - Default execution result - - Example: - ```python - async def async_default_execute(self, e: Exception = None, **kwargs): - logger.warning(f"Execution failed: {e}, returning default result") - return {"status": "error", "message": str(e)} - ``` - """ - - async def async_call(self, context: FlowContext = None, **kwargs) -> Any: - """Execute the operation asynchronously. - - This method handles the full async execution lifecycle including retries, - error handling, and context management. It automatically calls - `async_before_execute()`, `async_execute()`, and `async_after_execute()` - in sequence. - - Args: - context: Flow context for this execution. If None, a new context - will be created. - **kwargs: Additional context updates to merge into the context - - Returns: - Execution result from `async_execute()`, context response if result - is None, or None if both are None - - Raises: - Exception: If execution fails and `raise_exception` is True and - `max_retries` is exhausted - """ - self.context = self.build_context(context, **kwargs) - with self.timer: - result = None - if self.max_retries == 1 and self.raise_exception: - await self.async_before_execute() - result = await self.async_execute() - await self.async_after_execute() - - else: - for i in range(self.max_retries): - try: - await self.async_before_execute() - result = await self.async_execute() - await self.async_after_execute() - break - - except Exception as e: - logger.exception(f"op={self.name} async execute failed, error={e.args}") - - if self.raise_exception and i == self.max_retries - 1: - raise e - - result = await self.async_default_execute(e) - - if result is not None: - return result - - elif self.context is not None and self.context.response is not None: - return self.context.response - - else: - return None - - def submit_async_task(self, fn: Callable, *args, **kwargs): - """Submit an async task for execution. - - Creates an asyncio task and adds it to the task list for later joining. - Tasks can be collected using `join_async_task()`. - - Args: - fn: Coroutine function to execute - *args: Positional arguments for the coroutine - **kwargs: Keyword arguments for the coroutine - - Note: - Only coroutine functions are supported. Non-coroutine functions - will trigger a warning and be ignored. - - Example: - ```python - async def my_task(x): - return x * 2 - - self.submit_async_task(my_task, 5) - results = await self.join_async_task() - ``` - """ - loop = asyncio.get_running_loop() - if asyncio.iscoroutinefunction(fn): - task = loop.create_task(fn(*args, **kwargs)) - self.task_list.append(task) - else: - logger.warning("submit_async_task failed, fn is not a coroutine function!") - - async def join_async_task(self, timeout: float = None, return_exceptions: bool = True): - """Wait for all submitted async tasks to complete and collect results. - - Collects results from all tasks, handling exceptions and timeouts. - On timeout or exception, all remaining tasks are cancelled. - - Args: - timeout: Maximum time to wait in seconds (None for no timeout) - return_exceptions: Whether to return exceptions as results - - Returns: - List of task results (exceptions included if return_exceptions=True) - - Raises: - asyncio.TimeoutError: If timeout is exceeded - Exception: If any task raises an exception and return_exceptions=False - """ - result = [] - - if not self.task_list: - return result - - try: - if timeout is not None: - gather_task = asyncio.gather(*self.task_list, return_exceptions=return_exceptions) - task_results = await asyncio.wait_for(gather_task, timeout=timeout) - else: - task_results = await asyncio.gather(*self.task_list, return_exceptions=return_exceptions) - - for t_result in task_results: - if return_exceptions and isinstance(t_result, Exception): - logger.opt(exception=t_result).error("Task failed with exception") - continue - - if t_result: - if isinstance(t_result, list): - result.extend(t_result) - else: - result.append(t_result) - - except asyncio.TimeoutError: - logger.exception(f"join_async_task timeout after {timeout}s, cancelling {len(self.task_list)} tasks...") - for task in self.task_list: - if not task.done(): - task.cancel() - - await asyncio.gather(*self.task_list, return_exceptions=True) - self.task_list.clear() - raise - - except Exception as e: - logger.exception(f"join_async_task failed with {type(e).__name__}, cancelling remaining tasks...") - for task in self.task_list: - if not task.done(): - task.cancel() - - await asyncio.gather(*self.task_list, return_exceptions=True) - self.task_list.clear() - raise - - finally: - self.task_list.clear() - - return result diff --git a/flowllm/core/op/base_async_tool_op.py b/flowllm/core/op/base_async_tool_op.py deleted file mode 100644 index a72025c..0000000 --- a/flowllm/core/op/base_async_tool_op.py +++ /dev/null @@ -1,262 +0,0 @@ -"""Asynchronous tool operation base class. - -This module defines `BaseAsyncToolOp`, an abstract base operator that helps -implement asynchronous tool-like operations. It manages input/output schema -binding to a shared context, provides default result handling, and standardizes -pre- / post-execution hooks for derived ops. -""" - -import json -from abc import ABCMeta -from typing import List - -from loguru import logger - -from .base_async_op import BaseAsyncOp -from ..schema import ToolCall, ToolAttr - - -class BaseAsyncToolOp(BaseAsyncOp, metaclass=ABCMeta): - """Base class for async tool operators. - - Subclasses describe their tool interface via `build_tool_call()` and - implement their asynchronous execution while relying on this class to: - - extract inputs from the context before execution - - store outputs back into the context after execution - - optionally save the primary output to the response answer - - provide simple helpers for setting single or multiple results - """ - - def __init__( - self, - enable_print_output: bool = True, - tool_index: int = 0, - save_answer: bool = True, - input_schema_mapping: dict = None, - output_schema_mapping: dict = None, - max_print_output_length: int = 200, - **kwargs, - ): - """Initialize the async tool operator. - - Args: - enable_print_output: Whether to log output dictionaries. - tool_index: Index for disambiguating multiple instances of the tool. - save_answer: If True, write primary output into `response.answer`. - input_schema_mapping: Optional mapping from input names to context keys. - output_schema_mapping: Optional mapping from output names to context keys. - max_print_output_length: Maximum length of output strings to log. - **kwargs: Extra arguments forwarded to `BaseAsyncOp`. - """ - super().__init__(**kwargs) - - self.enable_print_output: bool = enable_print_output - self.tool_index: int = tool_index - self.save_answer: bool = save_answer - self.input_schema_mapping: dict | None = input_schema_mapping # map key to context - self.output_schema_mapping: dict | None = output_schema_mapping # map key to context - self.max_print_output_length: int = max_print_output_length - - self._tool_call: ToolCall | None = None - self.input_dict: dict = {} # Actual input values extracted from context - self.output_dict: dict = {} # Actual output values to save to context - - def build_tool_call(self) -> ToolCall: - """Build and return the tool call schema for this operator. - - Subclasses should override to provide `description`, `input_schema`, and - optionally `output_schema`. If `output_schema` is not provided, a default - single string output will be created. - - Returns: - ToolCall instance describing the tool's interface - - Example: - ```python - def build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "Search the web for information", - "input_schema": { - "query": { - "type": "string", - "description":"Search query", - "required": True, - }, - }, - } - ) - ``` - """ - return ToolCall(**{"description": "", "input_schema": {}}) - - @property - def tool_call(self): - """Return the lazily constructed `ToolCall` describing this tool. - - The tool call is constructed on first access by calling `build_tool_call()`. - The name and index are automatically set from the operator's short_name - and tool_index. - - Returns: - ToolCall instance with name and index set - """ - if self._tool_call is None: - self._tool_call = self.build_tool_call() - if not self._tool_call.name: - self._tool_call.name = self.short_name - - if not self._tool_call.index: - self._tool_call.index = self.tool_index - - if not self._tool_call.output_schema: - self._tool_call.output_schema = { - f"{self.short_name}_result": ToolAttr( - type="string", - description=f"The execution result of the {self.short_name}", - ), - } - - return self._tool_call - - @property - def output_keys(self) -> str | List[str]: - """Return the output key name or list of names defined by the schema. - - If the tool has a single output, returns the key name as a string. - If the tool has multiple outputs, returns a list of key names. - - Returns: - Single output key name (str) or list of output key names (List[str]) - """ - output_keys = [] - for name, _ in self.tool_call.output_schema.items(): - if name not in output_keys: - output_keys.append(name) - - if len(output_keys) == 1: - return output_keys[0] - else: - return output_keys - - @property - def output(self): - """Convenience accessor for the primary output value. - - Raises: - NotImplementedError: If multiple outputs exist; use `output_dict`. - """ - if isinstance(self.output_keys, str): - return self.output_dict[self.output_keys] - else: - raise NotImplementedError("use `output_dict` to get result") - - def set_output(self, value="", key: str = ""): - """Set a single output value. - - If only one output key exists, `key` is ignored and the single key is - used. Otherwise, `key` must be provided. - - Args: - value: The output value to set - key: The output key name (required if multiple outputs exist) - - Raises: - KeyError: If multiple outputs exist and `key` is not provided or - doesn't match any output key - """ - if isinstance(self.output_keys, str): - self.output_dict[self.output_keys] = value - else: - self.output_dict[key] = value - - def set_outputs(self, **kwargs): - """Set multiple output values using keyword arguments. - - Args: - **kwargs: Keyword arguments where keys are output names and values - are the output values - - Example: - ```python - self.set_outputs(result="success", count=10, status="ok") - ``` - """ - for k, v in kwargs.items(): - self.set_output(v, k) - - async def async_before_execute(self): - """Populate `input_dict` by reading required inputs from the context. - - Applies `input_schema_mapping` and appends the `tool_index` suffix when - necessary. Raises if any required input is missing. - - Raises: - ValueError: If any required input is missing from the context - """ - for name, attrs in self.tool_call.input_schema.items(): - context_key = name - if self.input_schema_mapping and name in self.input_schema_mapping: - context_key = self.input_schema_mapping[name] - if self.tool_index != 0: - context_key += f".{self.tool_index}" - - if context_key in self.context: - self.input_dict[name] = self.context[context_key] - elif attrs.required: - raise ValueError(f"{self.name}: {name} is required") - - async def async_after_execute(self): - """Write `output_dict` back to the context and optionally save answer. - - This method: - 1. Writes all values from `output_dict` to the context using the - appropriate keys (with mapping and index suffix if needed) - 2. Optionally saves the primary output to `context.response.answer` - if `save_answer` is True - 3. Logs the output dictionary if `enable_print_output` is True - - Applies `output_schema_mapping` and appends the `tool_index` suffix when - necessary. - """ - for name, value in self.output_dict.items(): - context_key = name - if self.output_schema_mapping and name in self.output_schema_mapping: - context_key = self.output_schema_mapping[name] - if self.tool_index != 0: - context_key += f".{self.tool_index}" - - logger.info(f"{self.name} set context key={context_key}") - self.context[context_key] = value - - if self.save_answer: - if isinstance(self.output_keys, str): - self.context.response.answer = self.output_dict.get(self.output_keys, "") - else: - self.context.response.answer = json.dumps(self.output_dict, ensure_ascii=False) - - if self.enable_print_output: - output_str = str(self.output_dict) - if len(output_str) > self.max_print_output_length: - output_str = f"{output_str[:self.max_print_output_length]}..." - - if self.tool_index == 0: - logger.info(f"{self.name}.output_dict={output_str}") - else: - logger.info(f"{self.name}.{self.tool_index}.output_dict={output_str}") - - async def async_default_execute(self, e: Exception = None, **kwargs): - """Fill outputs with a default failure message when execution fails. - - This method is called when `async_execute()` fails. It sets all output - keys in `output_dict` to a failure message string. - - Args: - e: The exception that was raised during execution (if any) - **kwargs: Additional keyword arguments (not used by default) - """ - if isinstance(self.output_keys, str): - self.output_dict[self.output_keys] = f"{self.name} execution failed!" - else: - for output_key in self.output_keys: - self.output_dict[output_key] = f"{self.name} execution failed!" diff --git a/flowllm/core/op/base_mcp_op.py b/flowllm/core/op/base_mcp_op.py deleted file mode 100644 index 5fbb6e8..0000000 --- a/flowllm/core/op/base_mcp_op.py +++ /dev/null @@ -1,102 +0,0 @@ -"""Base operator for calling external MCP tools asynchronously. - -This module defines `BaseMcpOp`, which leverages `BaseAsyncToolOp` to fetch the -tool schema from a configured MCP provider and execute the tool via -`FastMcpClient`. It supports overriding input parameter requirements and -propagates inputs/outputs through the shared context managed by the base class. -""" - -from typing import List, Optional - -from .base_async_tool_op import BaseAsyncToolOp -from ..context import C -from ..schema import ToolCall -from ..utils import FastMcpClient - - -@C.register_op() -class BaseMcpOp(BaseAsyncToolOp): - """Async op that invokes an external MCP tool. - - The tool's `ToolCall` schema is loaded from `C.external_mcp_tool_call_dict` - using `mcp_name` and `tool_name`. Inputs are prepared by the base class and - passed to the MCP server; results are written back to the context. - """ - - def __init__( - self, - mcp_name: str = "", - tool_name: str = "", - save_answer: bool = True, - input_schema_required: List[str] = None, - input_schema_optional: List[str] = None, - input_schema_deleted: List[str] = None, - max_retries: int = 3, - timeout: Optional[float] = None, - raise_exception: bool = False, - **kwargs, - ): - """Initialize the MCP operator. - - Args: - mcp_name: The MCP server name configured in service config. - tool_name: The MCP tool name to call on the server. - save_answer: Whether to save primary output to `response.answer`. - input_schema_required: Parameter names to force as required. - input_schema_optional: Parameter names to force as optional. - input_schema_deleted: Parameter names to remove from input schema. - max_retries: Max retries for MCP client calls. - timeout: Optional timeout (seconds) for the MCP request. - raise_exception: Whether to raise when execution fails. - **kwargs: Forwarded to `BaseAsyncToolOp` / `BaseAsyncOp`. - """ - self.mcp_name: str = mcp_name - self.tool_name: str = tool_name - self.input_schema_required: List[str] = input_schema_required - self.input_schema_optional: List[str] = input_schema_optional - self.input_schema_deleted: List[str] = input_schema_deleted - self.timeout: Optional[float] = timeout - super().__init__(save_answer=save_answer, max_retries=max_retries, raise_exception=raise_exception, **kwargs) - # Example MCP marketplace: https://bailian.console.aliyun.com/?tab=mcp#/mcp-market - - def build_tool_call(self) -> ToolCall: - """Build the `ToolCall` for this MCP tool, applying overrides. - - Returns: - ToolCall: A deep-copied tool schema with requirement overrides and - deletions applied as configured. - """ - tool_call_dict = C.external_mcp_tool_call_dict[self.mcp_name] - tool_call: ToolCall = tool_call_dict[self.tool_name].model_copy(deep=True) - - # Override parameter requirements if specified - if self.input_schema_required: - for name in self.input_schema_required: - tool_call.input_schema[name].required = True - - if self.input_schema_optional: - for name in self.input_schema_optional: - tool_call.input_schema[name].required = False - - if self.input_schema_deleted: - for name in self.input_schema_deleted: - tool_call.input_schema.pop(name, None) - - return tool_call - - async def async_execute(self): - """Execute the MCP tool call and store the result. - - Uses `FastMcpClient` with configured retries and timeout, invokes the - tool with `self.input_dict`, then writes the returned value via - `set_result` so that post-execution hooks can persist it to context. - """ - mcp_server_config = C.service_config.external_mcp[self.mcp_name] - async with FastMcpClient( - name=self.mcp_name, - config=mcp_server_config, - max_retries=self.max_retries, - timeout=self.timeout, - ) as client: - result: str = await client.call_tool(self.tool_name, arguments=self.input_dict) - self.set_output(result) diff --git a/flowllm/core/op/base_op.py b/flowllm/core/op/base_op.py deleted file mode 100644 index 0d6f40b..0000000 --- a/flowllm/core/op/base_op.py +++ /dev/null @@ -1,649 +0,0 @@ -"""Base operation class for flowllm operations. - -This module provides the BaseOp class, which serves as the foundation for all -operations in the flowllm framework. It includes support for synchronous and -asynchronous execution, caching, retries, and operation composition. -""" - -import copy -from abc import ABC, abstractmethod -from pathlib import Path -from typing import Callable, List, Union - -from loguru import logger -from tqdm import tqdm - -from ..context import FlowContext, PromptHandler, C, BaseContext -from ..embedding_model import BaseEmbeddingModel -from ..llm import BaseLLM -from ..schema import LLMConfig, EmbeddingModelConfig, Message, ToolCall, LLMTokenCountConfig -from ..storage import CacheHandler -from ..token import BaseToken -from ..utils import Timer, camel_to_snake -from ..vector_store import BaseVectorStore - - -class BaseOp(ABC): - """Base class for all operations in the flowllm framework. - - This class provides the core functionality for operations, including: - - Synchronous and asynchronous execution modes - - Retry mechanism with configurable retry counts - - Caching support with expiration - - Operation composition (sequential, parallel) - - LLM, embedding model, and vector store integration - - Prompt handling and formatting - - Prompt file convention: - Each Op can be paired with a prompt configuration file for templates. - By default, `prompt_path` is inferred from the Op's source file path by - replacing the suffix "op.py" with "prompt.yaml". For example: - - - file: my_feature_op.py -> prompt: my_feature_prompt.yaml - - file: chat_op.py -> prompt: chat_prompt.yaml - - This means the prompt file should be located in the same directory as - the Op file, and named by replacing "op.py" with "prompt.yaml". - - You can override this behavior by explicitly passing `prompt_path` when - initializing the Op. - - Attributes: - name: Operation name, auto-generated from class name if not provided - async_mode: Whether to run in async mode - max_retries: Maximum number of retry attempts - raise_exception: Whether to raise exceptions on failure - enable_multithread: Whether to enable multithreading for task submission - language: Language for prompt handling - prompt_path: Path to prompt configuration file - llm: LLM instance or configuration name - embedding_model: Embedding model instance or configuration name - vector_store: Vector store instance or configuration name - enable_cache: Whether to enable caching - cache_path: Cache storage path template - cache_expire_hours: Cache expiration time in hours - ops: List or dict of sub-operations - op_params: Additional operation parameters - context: Current execution context - timer: Timer for tracking execution time - - Example: - ```python - class MyOp(BaseOp): - def execute(self): - return "result" - - op = MyOp() - result = op.call() - ``` - """ - - file_path: str = __file__ - - def __new__(cls, *args, **kwargs): - """Create a new instance and save initialization arguments for copying. - - This method saves the initialization arguments so that `copy()` can - recreate the instance with the same parameters. This is necessary for - deep copying operations. - - Args: - *args: Positional arguments passed to __init__ - **kwargs: Keyword arguments passed to __init__ - - Returns: - New instance with saved initialization arguments stored in - `_init_args` and `_init_kwargs` attributes - """ - instance = super().__new__(cls) - instance._init_args = copy.copy(args) - instance._init_kwargs = copy.copy(kwargs) - return instance - - def __init__( - self, - name: str = "", - async_mode: bool = False, - max_retries: int = 1, - raise_exception: bool = False, - enable_multithread: bool = True, - language: str = "", - prompt_path: str = "", - llm: str = "default", - embedding_model: str = "default", - vector_store: str = "default", - enable_cache: bool = False, - cache_path: str = "cache/{op_name}", - cache_expire_hours: float = 0.1, - ops: List["BaseOp"] = None, - **kwargs, - ): - """Initialize the base operation. - - Args: - name: Operation name, auto-generated if empty - async_mode: Whether to run in async mode - max_retries: Maximum number of retry attempts - raise_exception: Whether to raise exceptions on failure - enable_multithread: Whether to enable multithreading - language: Language for prompt handling - prompt_path: Path to prompt configuration file. If not provided, it - will be inferred by replacing the Op file's suffix "op.py" with - "prompt.yaml" (same directory). For example: - - `xxx_op.py` -> `xxx_prompt.yaml` - llm: LLM configuration name or instance - embedding_model: Embedding model configuration name or instance - vector_store: Vector store configuration name or instance - enable_cache: Whether to enable caching - cache_path: Cache storage path template with {op_name} placeholder - cache_expire_hours: Cache expiration time in hours - ops: List of sub-operations - **kwargs: Additional operation parameters - """ - super().__init__() - assert self.__class__.__name__.endswith("Op"), f"class_name={self.__class__.__name__} must end with 'Op'" - - self.name: str = name or camel_to_snake(self.__class__.__name__) - self.async_mode: bool = async_mode - self.max_retries: int = max_retries - self.raise_exception: bool = raise_exception - self.enable_multithread: bool = enable_multithread - self.language: str = language or C.language - default_prompt_path: str = self.file_path.replace("op.py", "prompt.yaml") - self.prompt_path: Path = Path(prompt_path if prompt_path else default_prompt_path) - self.prompt = PromptHandler(language=self.language).load_prompt_by_file(self.prompt_path) - self._llm: BaseLLM | str = llm - self._embedding_model: BaseEmbeddingModel | str = embedding_model - self._vector_store: BaseVectorStore | str = vector_store - self.enable_cache: bool = enable_cache - self.cache_path: str = cache_path - self.cache_expire_hours: float = cache_expire_hours - self.ops: List["BaseOp"] | BaseContext = ops if ops else BaseContext() - self.op_params: dict = kwargs - - self.task_list: list = [] - self.timer = Timer(name=self.name) - self.context: FlowContext | None = None - self._cache: CacheHandler | None = None - self.llm_config: LLMConfig | None = None - self.embedding_model_config: EmbeddingModelConfig | None = None - - @property - def short_name(self) -> str: - """Get the short name of the operation (without '_op' suffix). - - Returns: - Short name without '_op' suffix - """ - return self.name.removesuffix("_op") - - @property - def cache(self): - """Get the cache handler instance (lazy initialization). - - Returns: - CacheHandler instance if caching is enabled, None otherwise - """ - if self.enable_cache and self._cache is None: - self._cache = CacheHandler(self.cache_path.format(op_name=self.short_name)) - return self._cache - - def save_load_cache(self, key: str, fn: Callable, **kwargs): - """Save or load from cache. - - If caching is enabled, checks cache first. If not found, executes the - function and saves the result. Otherwise, executes the function directly. - - Args: - key: Cache key for storing/retrieving the result - fn: Function to execute if cache miss - **kwargs: Additional arguments for cache load operation - - Returns: - Cached result if available, otherwise result from function execution - """ - if self.enable_cache: - result = self.cache.load(key, **kwargs) - if result is None: - result = fn() - self.cache.save(key, result, expire_hours=self.cache_expire_hours) - else: - logger.info(f"load {key} from cache") - else: - result = fn() - - return result - - def before_execute(self): - """Hook method called before execute(). Override in subclasses. - - This method is called automatically by `call()` before executing - the main `execute()` method. Use this to perform any setup, - validation, or preprocessing needed before execution. - - Example: - ```python - def before_execute(self): - # Validate inputs - if not self.context.get("input"): - raise ValueError("Input is required") - ``` - """ - - def after_execute(self): - """Hook method called after execute(). Override in subclasses. - - This method is called automatically by `call()` after successfully - executing the main `execute()` method. Use this to perform any - cleanup, post-processing, or result transformation. - - Example: - ```python - def after_execute(self): - # Post-process results - if self.context.response: - self.context.response.answer = self.context.response.answer.upper() - ``` - """ - - @abstractmethod - def execute(self): - """Main execution method. Must be implemented in subclasses. - - Returns: - Execution result - """ - - def default_execute(self, e: Exception = None, **kwargs): - """Default execution method when main execution fails. Override in subclasses. - - This method is called when `execute()` fails and `raise_exception` - is False. It provides a fallback mechanism to return a default result - instead of raising an exception. - - Args: - e: The exception that was raised during execution (if any) - **kwargs: Additional keyword arguments - - Returns: - Default execution result - - Example: - ```python - def default_execute(self, e: Exception = None, **kwargs): - logger.warning(f"Execution failed: {e}, returning default result") - return {"status": "error", "message": str(e)} - ``` - """ - - @staticmethod - def build_context(context: FlowContext = None, **kwargs): - """Build or update a flow context. - - Args: - context: Existing flow context, creates new one if None - **kwargs: Additional context updates - - Returns: - FlowContext instance with updates applied - """ - if not context: - context = FlowContext() - if kwargs: - context.update(kwargs) - return context - - def call(self, context: FlowContext = None, **kwargs): - """Execute the operation synchronously. - - This method handles the full execution lifecycle including retries, - error handling, and context management. It automatically calls - `before_execute()`, `execute()`, and `after_execute()` in sequence. - - Args: - context: Flow context for this execution. If None, a new context - will be created. - **kwargs: Additional context updates to merge into the context - - Returns: - Execution result from `execute()`, context response if result - is None, or None if both are None - - Raises: - Exception: If execution fails and `raise_exception` is True and - `max_retries` is exhausted - """ - self.context = self.build_context(context, **kwargs) - with self.timer: - result = None - if self.max_retries == 1 and self.raise_exception: - self.before_execute() - result = self.execute() - self.after_execute() - - else: - for i in range(self.max_retries): - try: - self.before_execute() - result = self.execute() - self.after_execute() - break - - except Exception as e: - logger.exception(f"op={self.name} execute failed, error={e.args}") - - if self.raise_exception and i == self.max_retries - 1: - raise e - - result = self.default_execute(e) # pylint: disable=E1111 - - if result is not None: - return result - - elif self.context is not None and self.context.response is not None: - return self.context.response - - else: - return None - - def submit_task(self, fn, *args, **kwargs) -> "BaseOp": - """Submit a task for execution (multithreaded or synchronous). - - If `enable_multithread` is True, the task is submitted to a thread pool - for concurrent execution. Otherwise, it is executed immediately and - the result is stored. Tasks can be collected using `join_task()`. - - Args: - fn: Function to execute - *args: Positional arguments for the function - **kwargs: Keyword arguments for the function - - Returns: - Self for method chaining - - Example: - ```python - def my_task(x): - return x * 2 - - self.submit_task(my_task, 5) - results = self.join_task() - ``` - """ - if self.enable_multithread: - task = C.thread_pool.submit(fn, *args, **kwargs) - self.task_list.append(task) - - else: - result = fn(*args, **kwargs) - if result: - if isinstance(result, list): - result.extend(result) - else: - result.append(result) - - return self - - def join_task(self, task_desc: str = None) -> list: - """Wait for all submitted tasks to complete and collect results. - - This method waits for all tasks submitted via `submit_task()` to - complete and collects their results. If multithreading is enabled, - a progress bar is displayed. The task list is cleared after collection. - - Args: - task_desc: Description for progress bar display. If None, uses - the operation name. - - Returns: - List of task results. If a task returns a list, its elements are - extended into the result list; otherwise, the result is appended. - """ - result = [] - if self.enable_multithread: - for task in tqdm(self.task_list, desc=task_desc or self.name): - t_result = task.result() - if t_result: - if isinstance(t_result, list): - result.extend(t_result) - else: - result.append(t_result) - - else: - result.extend(self.task_list) - - self.task_list.clear() - return result - - def check_async(self, op: "BaseOp"): - """Check if the given operation has the same async mode as self. - - Args: - op: Operation to check - - Raises: - AssertionError: If async modes don't match - """ - assert self.async_mode == op.async_mode, f"async_mode must be the same. {self.async_mode}!={op.async_mode}" - - def add_op(self, op: "BaseOp", name: str = ""): - """Add a sub-operation to this operation. - - The operation is added to `self.ops`. If `ops` is a list, the operation - is appended. If `ops` is a dict, the operation is stored with the given - name (or the operation's own name if not provided). - - Args: - op: Operation to add - name: Name for the operation (only used if ops is a dict). If empty - and ops is a dict, uses `op.name`. - """ - if isinstance(self.ops, list): - self.ops.append(op) - else: - if not name: - name = op.name - self.ops[name] = op - - def __lshift__(self, ops: Union["BaseOp", dict, list]): - """Left shift operator for adding operations (op << other_op, op << {"search": op1, "find": op2}). - - Args: - ops: Operation(s) to add - can be BaseOp, dict of ops, or list of ops - - Returns: - Self for method chaining - - Raises: - ValueError: If op type is invalid - """ - if isinstance(ops, BaseOp): - self.check_async(ops) - self.add_op(ops) - - elif isinstance(ops, dict): - for name, op in ops.items(): - self.add_op(op, name) - - elif isinstance(ops, list): - for op in ops: - self.add_op(op) - - else: - raise ValueError(f"Invalid op type: {type(ops)}") - - return self - - def __rshift__(self, op: "BaseOp"): - """Right shift operator for creating sequential operations (op >> other_op). - - Args: - op: Operation to chain sequentially - - Returns: - SequentialOp containing self and op - """ - self.check_async(op) - from .sequential_op import SequentialOp - - sequential_op = SequentialOp(ops=[self], async_mode=self.async_mode) - - if isinstance(op, SequentialOp): - sequential_op.ops.extend(op.ops) - else: - sequential_op.ops.append(op) - return sequential_op - - def __or__(self, op: "BaseOp"): - """Bitwise OR operator for creating parallel operations (op | other_op). - - Args: - op: Operation to run in parallel - - Returns: - ParallelOp containing self and op - """ - self.check_async(op) - from .parallel_op import ParallelOp - - parallel_op = ParallelOp(ops=[self], async_mode=self.async_mode) - - if isinstance(op, ParallelOp): - parallel_op.ops.extend(op.ops) - else: - parallel_op.ops.append(op) - - return parallel_op - - def copy(self, **kwargs): - """Create a deep copy of this operation. - - Creates a new instance using the saved initialization arguments and - recursively copies all sub-operations. Any additional kwargs override - the original initialization parameters. - - Args: - **kwargs: Additional parameters to override in the copy - - Returns: - New operation instance with copied configuration and sub-operations - - Raises: - NotImplementedError: If the ops type is not supported (not a list - or BaseContext) - """ - copy_op = self.__class__(*self._init_args, **self._init_kwargs, **kwargs) - if self.ops: - if isinstance(self.ops, list): - copy_op.ops.clear() - for op in self.ops: - copy_op.ops.append(op.copy()) - - elif isinstance(self.ops, BaseContext): - copy_op.ops.clear() - for name, op in self.ops.items(): - copy_op.ops[name] = op.copy() - - else: - raise NotImplementedError("ops type not supported") - - return copy_op - - @property - def llm(self) -> BaseLLM: - """Get the LLM instance (lazy initialization). - - Returns: - BaseLLM instance configured from service config - """ - if isinstance(self._llm, str): - self.llm_config: LLMConfig = C.service_config.llm[self._llm] - llm_cls = C.get_llm_class(self.llm_config.backend) - self._llm = llm_cls(model_name=self.llm_config.model_name, **self.llm_config.params) - - return self._llm - - @property - def embedding_model(self) -> BaseEmbeddingModel: - """Get the embedding model instance (lazy initialization). - - Returns: - BaseEmbeddingModel instance configured from service config - """ - if isinstance(self._embedding_model, str): - self.embedding_model_config: EmbeddingModelConfig = C.service_config.embedding_model[self._embedding_model] - embedding_model_cls = C.get_embedding_model_class(self.embedding_model_config.backend) - self._embedding_model = embedding_model_cls( - model_name=self.embedding_model_config.model_name, - **self.embedding_model_config.params, - ) - - return self._embedding_model - - @property - def vector_store(self) -> BaseVectorStore: - """Get the vector store instance (lazy initialization). - - Returns: - BaseVectorStore instance configured from service config - """ - if isinstance(self._vector_store, str): - self._vector_store = C.get_vector_store(self._vector_store) - return self._vector_store - - @property - def service_config_metadata(self) -> dict: - """Get the service config metadata for this operation. - - Returns: - Service config metadata - """ - return C.service_config.metadata - - def prompt_format(self, prompt_name: str, **kwargs) -> str: - """Format a prompt template with provided variables. - - Args: - prompt_name: Name of the prompt template - **kwargs: Variables to substitute in the template - - Returns: - Formatted prompt string - """ - return self.prompt.prompt_format(prompt_name=prompt_name, **kwargs) - - def get_prompt(self, prompt_name: str) -> str: - """Get a prompt template by name. - - Args: - prompt_name: Name of the prompt template - - Returns: - Prompt template string - """ - return self.prompt.get_prompt(prompt_name=prompt_name) - - def token_count(self, messages: List[Message], tools: List[ToolCall] | None = None, **kwargs) -> int: - """Count tokens for the given messages and optional tools. - - Args: - messages: List of messages to count tokens for - tools: Optional list of tool calls to include in token count - **kwargs: Additional keyword arguments passed to the token counter - - Returns: - Total token count as an integer - """ - if self.llm_config is None: - llm_config: LLMConfig = C.service_config.llm[self._llm] - else: - llm_config = self.llm_config - - token_count_config: LLMTokenCountConfig = llm_config.token_count - - token_count_cls = C.get_token_counter_class(token_count_config.backend) - model_name = token_count_config.model_name - params = token_count_config.params - token_counter: BaseToken = token_count_cls(model_name=model_name, **params) - - return token_counter.token_count(messages, tools, **kwargs) diff --git a/flowllm/core/op/base_ray_op.py b/flowllm/core/op/base_ray_op.py deleted file mode 100644 index 539240d..0000000 --- a/flowllm/core/op/base_ray_op.py +++ /dev/null @@ -1,251 +0,0 @@ -"""Base Ray operation class for distributed parallel execution. - -This module provides the BaseRayOp class, which extends BaseOp with Ray-based -distributed parallelization capabilities. It allows operations to be executed -in parallel across multiple workers using Ray for efficient distributed computing. -""" - -from abc import ABC - -import pandas as pd -from loguru import logger -from tqdm import tqdm - -from .base_op import BaseOp -from ..context import BaseContext, C - - -class BaseRayOp(BaseOp, ABC): - """Base class for operations with Ray-based distributed parallelization. - - This class extends BaseOp to provide distributed parallel execution using Ray. - It automatically detects list parameters for parallelization and distributes - work across multiple Ray workers. Large data structures (DataFrames, Series, - dicts, lists, Contexts) are automatically converted to Ray objects for - efficient sharing across workers. - - Attributes: - ray_task_list: List of Ray task futures being tracked - - Example: - ```python - class MyRayOp(BaseRayOp): - def execute(self): - # Automatically parallelizes over 'items' list - results = self.submit_and_join_ray_task( - fn=process_item, - items=[1, 2, 3, 4, 5], - other_param="value" - ) - return results - ``` - """ - - def __init__(self, **kwargs): - """Initialize the Ray operation. - - Args: - **kwargs: Arguments passed to BaseOp.__init__ - """ - super().__init__(**kwargs) - self.ray_task_list = [] - - def submit_and_join_parallel_op(self, op: BaseOp, **kwargs): - """Submit a parallel operation and wait for all results. - - Automatically detects the first list parameter in kwargs as the parallel key - and distributes the operation's execution across Ray workers. - - Args: - op: BaseOp instance to execute in parallel - **kwargs: Arguments to pass to op.call(), must contain at least one list - - Returns: - List of results from all parallel executions - - Raises: - AssertionError: If no list parameter is found in kwargs - """ - parallel_key = None - for key, value in kwargs.items(): - if isinstance(value, list): - parallel_key = key - logger.info(f"using first list parallel_key={parallel_key}") - break - assert parallel_key is not None - - return self.submit_and_join_ray_task( - fn=op.call, - parallel_key=parallel_key, - task_desc=f"{op.short_name}.parallel", - context=self.context, - **kwargs, - ) - - def submit_and_join_ray_task( - self, - fn, - parallel_key: str = "", - task_desc: str = "", - **kwargs, - ): - """Submit Ray tasks for parallel execution and wait for all results. - - This method distributes work across multiple Ray workers by splitting a - list parameter. Large data structures are automatically converted to Ray - objects for efficient sharing. - - Args: - fn: Function to execute in parallel. Will be called with each item - from the parallel_list as the parallel_key parameter. - parallel_key: Name of the keyword argument containing the list to - parallelize over. If empty, auto-detects the first list parameter. - task_desc: Description for progress bar display - **kwargs: Additional arguments to pass to fn. Large data structures - (DataFrame, Series, dict, list, BaseContext) are converted to - Ray objects automatically. - - Returns: - List of results from all parallel executions, flattened if individual - results are lists - - Raises: - AssertionError: If no list parameter is found when parallel_key is empty - RuntimeError: If Ray is not configured (ray_max_workers <= 1) - """ - import ray - - max_workers = C.service_config.ray_max_workers - self.ray_task_list.clear() - - # Auto-detect parallel key if not provided - if not parallel_key: - for key, value in kwargs.items(): - if isinstance(value, list): - parallel_key = key - logger.info(f"using first list parallel_key={parallel_key}") - break - - # Extract the list to parallelize over - parallel_list = kwargs.pop(parallel_key) - assert isinstance(parallel_list, list) - - # Convert pandas DataFrames to Ray objects for efficient sharing - for key in sorted(kwargs.keys()): - value = kwargs[key] - if isinstance(value, pd.DataFrame | pd.Series | dict | list | BaseContext): - kwargs[key] = ray.put(value) - - # Create and submit tasks for each worker - for i in range(max_workers): - self.submit_ray_task( - fn=self.ray_task_loop, - parallel_key=parallel_key, - parallel_list=parallel_list, - actor_index=i, - max_workers=max_workers, - internal_fn=fn, - **kwargs, - ) - logger.info(f"ray.submit task_desc={task_desc} id={i}") - - # Wait for all tasks to complete and collect results - result = self.join_ray_task(task_desc=task_desc) - logger.info(f"{task_desc} complete. result_size={len(result)} resources={ray.available_resources()}") - return result - - @staticmethod - def ray_task_loop( - parallel_key: str, - parallel_list: list, - actor_index: int, - max_workers: int, - internal_fn, - **kwargs, - ): - """Worker loop that processes a subset of items from parallel_list. - - Each worker processes items assigned to it based on actor_index and - max_workers using a round-robin distribution (items[actor_index::max_workers]). - - Args: - parallel_key: Name of the keyword argument to set with parallel_value - parallel_list: Full list of items to process - actor_index: Index of this worker (0 to max_workers-1) - max_workers: Total number of workers - internal_fn: Function to call for each item - **kwargs: Additional arguments to pass to internal_fn - - Returns: - List of results from processing assigned items - """ - result = [] - for parallel_value in parallel_list[actor_index::max_workers]: - kwargs.update({"actor_index": actor_index, parallel_key: parallel_value}) - t_result = internal_fn(**kwargs) - if t_result: - if isinstance(t_result, list): - result.extend(t_result) - else: - result.append(t_result) - return result - - def submit_ray_task(self, fn, *args, **kwargs): - """Submit a single Ray task for asynchronous execution. - - Initializes Ray if not already initialized and creates a remote function - to execute the task. The task is added to ray_task_list for later joining. - - Args: - fn: Function to execute remotely - *args: Positional arguments for fn - **kwargs: Keyword arguments for fn - - Returns: - Self for method chaining - - Raises: - RuntimeError: If Ray is not configured (ray_max_workers <= 1) - """ - import ray - - if C.service_config.ray_max_workers <= 1: - raise RuntimeError("Ray is not configured. Please set ray_max_workers > 1 in service config.") - - # Initialize Ray if not already done - if not ray.is_initialized(): - logger.warning(f"Ray is not initialized. Initializing Ray with {C.service_config.ray_max_workers} workers.") - ray.init(num_cpus=C.service_config.ray_max_workers) - - # Create remote function and submit task - remote_fn = ray.remote(fn) - task = remote_fn.remote(*args, **kwargs) - self.ray_task_list.append(task) - return self - - def join_ray_task(self, task_desc: str = None) -> list: - """Wait for all submitted Ray tasks to complete and collect results. - - Processes tasks with a progress bar and collects results. If individual - results are lists, they are flattened into the final result list. - - Args: - task_desc: Description for progress bar display. If None, uses - "{self.name}_ray" - - Returns: - List of all task results, flattened if individual results are lists - """ - result = [] - # Process each task and collect results with progress bar - import ray - - for task in tqdm(self.ray_task_list, desc=task_desc or f"{self.name}_ray"): - t_result = ray.get(task) - if t_result: - if isinstance(t_result, list): - result.extend(t_result) - else: - result.append(t_result) - self.ray_task_list.clear() - return result diff --git a/flowllm/core/op/parallel_op.py b/flowllm/core/op/parallel_op.py deleted file mode 100644 index 61c3ad5..0000000 --- a/flowllm/core/op/parallel_op.py +++ /dev/null @@ -1,81 +0,0 @@ -"""Parallel operation execution. - -This module provides the ParallelOp class, which executes operations -in parallel, collecting all results at the end. -""" - -from typing import Union - -from .base_async_op import BaseAsyncOp -from .base_op import BaseOp - - -class ParallelOp(BaseAsyncOp): - """Operation that executes sub-operations in parallel. - - This operation executes all sub-operations concurrently and collects - their results. All operations receive the same context. - - Example: - ```python - op1 = MyOp1() - op2 = MyOp2() - parallel = op1 | op2 # Creates ParallelOp - results = parallel.call(context=FlowContext()) - ``` - """ - - def execute(self): - """Execute operations in parallel in sync mode. - - Uses multithreading to execute operations concurrently. - - Returns: - List of results from all operations - """ - for op in self.ops: - assert not op.async_mode - self.submit_task(op.call, context=self.context) - return self.join_task(task_desc="parallel execution") - - async def async_execute(self): - """Execute operations in parallel in async mode. - - Uses asyncio tasks to execute operations concurrently. - - Returns: - List of results from all operations - """ - for op in self.ops: - assert op.async_mode - assert isinstance(op, BaseAsyncOp) - self.submit_async_task(op.async_call, context=self.context) - return await self.join_async_task() - - def __lshift__(self, op: Union["BaseOp", dict, list]): - """Left shift operator is not supported in ParallelOp. - - Args: - op: Operation (not used) - - Raises: - RuntimeError: Always raises, as `<<` is not supported - """ - raise RuntimeError(f"`<<` is not supported in {self.name}") - - def __or__(self, op: BaseOp): - """Bitwise OR operator for adding operations to parallel execution. - - Args: - op: Operation to run in parallel - - Returns: - Self with the operation added to parallel execution - """ - self.check_async(op) - - if isinstance(op, ParallelOp): - self.ops.extend(op.ops) - else: - self.ops.append(op) - return self diff --git a/flowllm/core/op/sequential_op.py b/flowllm/core/op/sequential_op.py deleted file mode 100644 index 2578c1c..0000000 --- a/flowllm/core/op/sequential_op.py +++ /dev/null @@ -1,78 +0,0 @@ -"""Sequential operation execution. - -This module provides the SequentialOp class, which executes operations -sequentially one after another, passing the context through the chain. -""" - -from typing import Union - -from .base_async_op import BaseAsyncOp -from .base_op import BaseOp - - -class SequentialOp(BaseAsyncOp): - """Operation that executes sub-operations sequentially. - - This operation executes its sub-operations in order, passing the context - from one operation to the next. Each operation receives the output context - from the previous operation. - - Example: - ```python - op1 = MyOp1() - op2 = MyOp2() - sequential = op1 >> op2 # Creates SequentialOp - result = sequential.call(context=FlowContext()) - ``` - """ - - def execute(self): - """Execute operations sequentially in sync mode. - - Returns: - Result from the last operation in the sequence - """ - result = None - for op in self.ops: - assert op.async_mode is False - result = op.call(context=self.context) - return result - - async def async_execute(self): - """Execute operations sequentially in async mode. - - Returns: - Result from the last operation in the sequence - """ - result = None - for op in self.ops: - assert op.async_mode is True - assert isinstance(op, BaseAsyncOp) - result = await op.async_call(context=self.context) - return result - - def __lshift__(self, op: Union["BaseOp", dict, list]): - """Left shift operator is not supported in SequentialOp. - - Args: - op: Operation (not used) - - Raises: - RuntimeError: Always raises, as `<<` is not supported - """ - raise RuntimeError(f"`<<` is not supported in {self.name}") - - def __rshift__(self, op: BaseOp): - """Right shift operator for chaining operations sequentially. - - Args: - op: Operation to chain sequentially - - Returns: - Self with the operation added to the sequence - """ - if isinstance(op, SequentialOp): - self.ops.extend(op.ops) - else: - self.ops.append(op) - return self diff --git a/flowllm/core/schema/__init__.py b/flowllm/core/schema/__init__.py deleted file mode 100644 index b45f575..0000000 --- a/flowllm/core/schema/__init__.py +++ /dev/null @@ -1,39 +0,0 @@ -"""Schema definitions for flowllm core components.""" - -from .flow_request import FlowRequest -from .flow_response import FlowResponse -from .flow_stream_chunk import FlowStreamChunk -from .message import Message, Trajectory -from .service_config import ( - CmdConfig, - EmbeddingModelConfig, - FlowConfig, - HttpConfig, - LLMConfig, - LLMTokenCountConfig, - MCPConfig, - ServiceConfig, - VectorStoreConfig, -) -from .tool_call import ToolAttr, ToolCall -from .vector_node import VectorNode - -__all__ = [ - "FlowRequest", - "FlowResponse", - "FlowStreamChunk", - "Message", - "Trajectory", - "CmdConfig", - "EmbeddingModelConfig", - "FlowConfig", - "HttpConfig", - "LLMConfig", - "LLMTokenCountConfig", - "MCPConfig", - "ServiceConfig", - "VectorStoreConfig", - "ToolAttr", - "ToolCall", - "VectorNode", -] diff --git a/flowllm/core/schema/flow_request.py b/flowllm/core/schema/flow_request.py deleted file mode 100644 index b0512ac..0000000 --- a/flowllm/core/schema/flow_request.py +++ /dev/null @@ -1,18 +0,0 @@ -"""Flow request schema for API requests.""" - -from typing import List - -from pydantic import Field, BaseModel, ConfigDict - -from .message import Message - - -class FlowRequest(BaseModel): - """Represents a flow execution request with query, messages, and workspace context.""" - - query: str = Field(default="") - messages: List[Message] = Field(default_factory=list) - workspace_id: str = Field(default="") - metadata: dict = Field(default_factory=dict) - - model_config = ConfigDict(extra="allow") diff --git a/flowllm/core/schema/flow_response.py b/flowllm/core/schema/flow_response.py deleted file mode 100644 index 9d0fda5..0000000 --- a/flowllm/core/schema/flow_response.py +++ /dev/null @@ -1,11 +0,0 @@ -"""Flow response schema for API responses.""" - -from pydantic import Field, BaseModel - - -class FlowResponse(BaseModel): - """Represents a complete flow execution response with answer and messages.""" - - answer: str | dict | list = Field(default="") - success: bool = Field(default=True) - metadata: dict = Field(default_factory=dict) diff --git a/flowllm/core/schema/flow_stream_chunk.py b/flowllm/core/schema/flow_stream_chunk.py deleted file mode 100644 index 29bc4b0..0000000 --- a/flowllm/core/schema/flow_stream_chunk.py +++ /dev/null @@ -1,15 +0,0 @@ -"""Flow stream chunk schema for streaming responses.""" - -from pydantic import Field, BaseModel - -from ..enumeration import ChunkEnum - - -class FlowStreamChunk(BaseModel): - """Represents a chunk of streaming flow output.""" - - flow_id: str = Field(default="") - chunk_type: ChunkEnum = Field(default=ChunkEnum.ANSWER) - chunk: str | dict | list = Field(default="") - done: bool = Field(default=False) - metadata: dict = Field(default_factory=dict) diff --git a/flowllm/core/schema/message.py b/flowllm/core/schema/message.py deleted file mode 100644 index b62b42b..0000000 --- a/flowllm/core/schema/message.py +++ /dev/null @@ -1,105 +0,0 @@ -"""Message and trajectory schema definitions for conversation management.""" - -import datetime -from typing import List - -from pydantic import BaseModel, ConfigDict, Field, model_validator - -from .tool_call import ToolCall -from ..enumeration import Role - - -class ContentBlock(BaseModel): - """ - examples: - { - "type": "image_url", - "image_url": { - "url": "https://img.alicdn.com/imgextra/i1/O1CN01gDEY8M1W114Hi3XcN_!!6000000002727-0-tps-1024-406.jpg" - }, - } - { - "type": "video", - "video": [ - "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20241108/xzsgiz/football1.jpg", - "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20241108/tdescd/football2.jpg", - "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20241108/zefdja/football3.jpg", - "https://help-static-aliyun-doc.aliyuncs.com/file-manage-files/zh-CN/20241108/aedbqh/football4.jpg", - ], - } - { - "type": "text", - "text": "这道题怎么解答？" - } - """ - - model_config = ConfigDict(extra="allow") - - type: str = Field(default="") - content: str | dict | list = Field(default="") - - @model_validator(mode="before") - @classmethod - def init_block(cls, data: dict): - """Initialize content block by extracting content based on type field.""" - result = data.copy() - content_type = data.get("type", "") - if content_type and content_type in data: - result["content"] = data[content_type] - return result - - def simple_dump(self) -> dict: - """Convert ContentBlock to a simple dictionary format.""" - result = { - "type": self.type, - self.type: self.content, - **self.model_extra, - } - - return result - - -class Message(BaseModel): - """Represents a message in a conversation with role, content, and optional tool calls.""" - - name: str | None = Field(default=None) - role: Role = Field(default=Role.USER) - content: str | List[ContentBlock] = Field(default="") - reasoning_content: str = Field(default="") - tool_calls: List[ToolCall] = Field(default_factory=list) - tool_call_id: str = Field(default="") - time_created: str = Field(default_factory=lambda: datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")) - metadata: dict = Field(default_factory=dict) - - def dump_content(self) -> str | list[dict]: - """Dump message content to string or list of dicts.""" - if isinstance(self.content, str): - return self.content - elif isinstance(self.content, list): - return [x.simple_dump() for x in self.content] - else: - raise ValueError(f"Invalid content type: {type(self.content)}") - - def simple_dump(self, add_reasoning: bool = True) -> dict: - """Convert Message to a simple dictionary format for API serialization.""" - result: dict = {"role": self.role.value, "content": self.dump_content()} - - if add_reasoning: - result["reasoning_content"] = self.reasoning_content - - if self.tool_calls: - result["tool_calls"] = [x.simple_output_dump() for x in self.tool_calls] - - if self.tool_call_id: - result["tool_call_id"] = self.tool_call_id - - return result - - -class Trajectory(BaseModel): - """Represents a conversation trajectory with messages and optional scoring.""" - - task_id: str = Field(default="") - messages: List[Message] = Field(default_factory=list) - score: float = Field(default=0.0) - metadata: dict = Field(default_factory=dict) diff --git a/flowllm/core/schema/service_config.py b/flowllm/core/schema/service_config.py deleted file mode 100644 index e5b9107..0000000 --- a/flowllm/core/schema/service_config.py +++ /dev/null @@ -1,102 +0,0 @@ -"""Service configuration schema definitions for various service components.""" - -from typing import Dict, List - -from pydantic import BaseModel, Field - -from .tool_call import ToolCall - - -class MCPConfig(BaseModel): - """Configuration for MCP (Model Context Protocol) transport settings.""" - - transport: str = Field(default="") - host: str = Field(default="0.0.0.0") - port: int = Field(default=8001) - - -class HttpConfig(BaseModel): - """Configuration for HTTP server settings.""" - - host: str = Field(default="0.0.0.0") - port: int = Field(default=8001) - timeout_keep_alive: int = Field(default=3600) - limit_concurrency: int = Field(default=1000) - - -class CmdConfig(BaseModel): - """Configuration for command-line execution.""" - - flow: str = Field(default="") - params: dict = Field(default_factory=dict) - - -class FlowConfig(ToolCall): - """Configuration for flow execution with optional streaming and caching. - - In addition to tool-call metadata inherited from `ToolCall`, this config - controls execution behavior such as streaming output and response caching. - Caching is only applied to non-streaming invocations. - """ - - flow_content: str = Field(default="") - stream: bool = Field(default=False) - enable_cache: bool = Field(default=False, description="Enable non-stream response caching") - cache_path: str = Field(default="cache/{flow_name}", description="Cache path template; supports {flow_name}") - cache_expire_hours: float = Field(default=0.1, description="Cache TTL (hours)") - - -class LLMTokenCountConfig(BaseModel): - """Configuration for LLM token count estimation.""" - - backend: str = Field(default="base") - model_name: str = Field(default="") - params: dict = Field(default_factory=dict) - - -class LLMConfig(BaseModel): - """Configuration for LLM (Large Language Model) backend settings.""" - - backend: str = Field(default="") - model_name: str = Field(default="") - token_count: LLMTokenCountConfig = Field(default_factory=LLMTokenCountConfig) - params: dict = Field(default_factory=dict) - - -class EmbeddingModelConfig(BaseModel): - """Configuration for embedding model backend settings.""" - - backend: str = Field(default="") - model_name: str = Field(default="") - params: dict = Field(default_factory=dict) - - -class VectorStoreConfig(BaseModel): - """Configuration for vector store backend settings.""" - - backend: str = Field(default="") - embedding_model: str = Field(default="") - params: dict = Field(default_factory=dict) - - -class ServiceConfig(BaseModel): - """Main service configuration aggregating all component configurations.""" - - backend: str = Field(default="") - enable_logo: bool = Field(default=True) - language: str = Field(default="") - thread_pool_max_workers: int = Field(default=16) - ray_max_workers: int = Field(default=-1) - init_logger: bool = Field(default=True) - disabled_flows: List[str] = Field(default_factory=list) - enabled_flows: List[str] = Field(default_factory=list) - - cmd: CmdConfig = Field(default_factory=CmdConfig) - mcp: MCPConfig = Field(default_factory=MCPConfig) - external_mcp: Dict[str, dict] = Field(default_factory=dict, description="External MCP Server config") - http: HttpConfig = Field(default_factory=HttpConfig) - flow: Dict[str, FlowConfig] = Field(default_factory=dict) - llm: Dict[str, LLMConfig] = Field(default_factory=dict) - embedding_model: Dict[str, EmbeddingModelConfig] = Field(default_factory=dict) - vector_store: Dict[str, VectorStoreConfig] = Field(default_factory=dict) - metadata: dict = Field(default_factory=dict) diff --git a/flowllm/core/schema/tool_call.py b/flowllm/core/schema/tool_call.py deleted file mode 100644 index 39dc6d4..0000000 --- a/flowllm/core/schema/tool_call.py +++ /dev/null @@ -1,205 +0,0 @@ -"""Tool call schema definitions for representing and managing tool calls.""" - -import json -from typing import Dict, List, Literal - -from mcp.types import Tool -from pydantic import BaseModel, Field, model_validator, ConfigDict - -TOOL_ATTR_TYPE = Literal["string", "array", "integer", "number", "boolean", "object"] - - -class ToolAttr(BaseModel): - """Attributes for tool parameters.""" - - type: TOOL_ATTR_TYPE = Field(default="string", description="tool attribute type") - description: str = Field(default="", description="tool attribute description") - required: bool = Field(default=True, description="tool attribute required") - enum: List[str] | None = Field(default=None, description="tool attribute enum") - items: dict = Field(default_factory=dict, description="tool attribute items") - - model_config = ConfigDict(extra="allow") - - def simple_input_dump(self, version: str = "default") -> dict: - """Convert ToolAttr to input format dictionary for API requests.""" - if version == "default": - result: dict = { - "type": self.type, - "description": self.description, - } - - if self.enum: - result["enum"] = self.enum - - if self.items: - result["items"] = self.items - - return result - - else: - raise NotImplementedError(f"version {version} not supported") - - -class ToolCall(BaseModel): - """ - input: - { - "type": "function", - "function": { - "name": "get_current_weather", - "description": "It is very useful when you want to check the weather of a specified city.", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "Cities or counties, such as Beijing, Hangzhou, Yuhang District, etc.", - } - }, - "required": ["location"] - } - } - } - output: - { - "index": 0, - "id": "call_6596dafa2a6a46f7a217da", - "function": { - "arguments": "{\"location\": \"Beijing\"}", - "name": "get_current_weather" - }, - "type": "function", - } - """ - - index: int = Field(default=0) - id: str = Field(default="") - type: str = Field(default="function") - name: str = Field(default="") - - arguments: str = Field(default="", description="tool execution arguments") - - description: str = Field(default="") - input_schema: Dict[str, ToolAttr] = Field(default_factory=dict) - output_schema: Dict[str, ToolAttr] = Field(default_factory=dict) - - @model_validator(mode="before") - @classmethod - def init_tool_call(cls, data: dict): - """Initialize ToolCall from raw data dictionary, extracting function info.""" - tool_type = data.get("type", "") - tool_type_dict = data.get(tool_type, {}) - - # Create a new dict to avoid modifying the original data - result = data.copy() - - if "name" in tool_type_dict: - result["name"] = tool_type_dict["name"] - - if "arguments" in tool_type_dict: - result["arguments"] = tool_type_dict["arguments"] - - if "description" in tool_type_dict: - result["description"] = tool_type_dict["description"] - - if "parameters" in tool_type_dict: - parameters = tool_type_dict["parameters"] - properties: dict = parameters.get("properties", {}) - required: list = parameters.get("required", []) - result["input_schema"] = {} - for name, param_attrs in properties.items(): - tool_attr = ToolAttr(**param_attrs) - tool_attr.required = name in required - result["input_schema"][name] = tool_attr - - return result - - @property - def argument_dict(self) -> dict: - """Parse and return arguments as a dictionary.""" - return json.loads(self.arguments) - - def check_argument(self) -> bool: - """Check if arguments can be parsed as valid JSON.""" - try: - _ = self.argument_dict - return True - except Exception: - return False - - @staticmethod - def _build_schema_dict(schema: Dict[str, ToolAttr]) -> dict: - required_list = [] - properties = {} - for name, tool_attr in schema.items(): - if tool_attr.required: - required_list.append(name) - properties[name] = tool_attr.simple_input_dump() - - return { - "type": "object", - "properties": properties, - "required": required_list, - } - - def simple_input_dump(self, version: str = "default") -> dict: - """Convert ToolCall to input format dictionary for API requests.""" - if version == "default": - return { - "type": self.type, - self.type: { - "name": self.name, - "description": self.description, - "parameters": self._build_schema_dict(self.input_schema), - }, - } - - else: - raise NotImplementedError(f"version {version} not supported") - - def simple_output_dump(self, version: str = "default") -> dict: - """Convert ToolCall to output format dictionary for API responses.""" - if version == "default": - return { - "index": self.index, - "id": self.id, - self.type: { - "arguments": self.arguments, - "name": self.name, - }, - "type": self.type, - } - else: - raise NotImplementedError(f"version {version} not supported") - - @classmethod - def from_mcp_tool(cls, tool: Tool) -> "ToolCall": - """Create a ToolCall instance from an MCP Tool object.""" - input_schema = {} - properties = tool.inputSchema["properties"] - required = tool.inputSchema.get("required", []) - for name, attr_dict in properties.items(): - tool_attr = ToolAttr(**attr_dict) - if name in required: - tool_attr.required = True - input_schema[name] = tool_attr - - return cls( - name=tool.name, - description=tool.description, - input_schema=input_schema, - ) - - def to_mcp_tool(self) -> Tool: - """Convert this ToolCall to an MCP Tool object.""" - tool = Tool( - name=self.name, - description=self.description, - inputSchema=self._build_schema_dict(self.input_schema), - ) - - # Build outputSchema from output_schema if present - if self.output_schema: - tool.outputSchema = self._build_schema_dict(self.output_schema) - - return tool diff --git a/flowllm/core/schema/vector_node.py b/flowllm/core/schema/vector_node.py deleted file mode 100644 index 480aefb..0000000 --- a/flowllm/core/schema/vector_node.py +++ /dev/null @@ -1,16 +0,0 @@ -"""Vector node schema for representing nodes with embeddings.""" - -from typing import List -from uuid import uuid4 - -from pydantic import BaseModel, Field - - -class VectorNode(BaseModel): - """Represents a node with content, vector embeddings, and metadata.""" - - unique_id: str = Field(default_factory=lambda: uuid4().hex) - workspace_id: str = Field(default="") - content: str = Field(default="") - vector: List[float] | None = Field(default=None) - metadata: dict = Field(default_factory=dict) diff --git a/flowllm/core/service/__init__.py b/flowllm/core/service/__init__.py deleted file mode 100644 index aa2c5b0..0000000 --- a/flowllm/core/service/__init__.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Service interfaces and concrete implementations for FlowLLM. - -This subpackage defines the standard service abstraction used to connect flows -to external runtimes and tools, and provides reference implementations: - -- ``BaseService``: Defines the abstract lifecycle and request/response - contract that all services follow. -- ``CmdService``: Executes local command-line subprocesses as a service. -- ``HttpService``: Exposes a simple HTTP/JSON service implementation. -- ``MCPService``: Integrates with Model Context Protocol (MCP) servers. - -Typical usage: - from flowllm.core.service import BaseService, HttpService - -The module-level ``__all__`` re-exports the primary service classes for -convenience. -""" - -from .base_service import BaseService -from .cmd_service import CmdService -from .http_service import HttpService -from .mcp_service import MCPService - -__all__ = [ - "BaseService", - "CmdService", - "HttpService", - "MCPService", -] diff --git a/flowllm/core/service/base_service.py b/flowllm/core/service/base_service.py deleted file mode 100644 index f71747a..0000000 --- a/flowllm/core/service/base_service.py +++ /dev/null @@ -1,78 +0,0 @@ -"""Service base definitions for FlowLLM. - -This module defines the abstract `BaseService` that concrete services (CLI, HTTP, -MCP, etc.) extend to integrate registered flows and run the application. -""" - -from abc import ABC - -from loguru import logger - -from ..context import C -from ..flow import BaseFlow, BaseToolFlow -from ..schema import ServiceConfig - - -class BaseService(ABC): - """Abstract base class for all services. - - Services are responsible for integrating registered flows into a runtime - (e.g., CLI, HTTP server, MCP) and then starting that runtime. - """ - - def __init__(self, service_config: ServiceConfig): - """Initialize the service. - - - service_config: The configuration object for the current service mode. - """ - self.service_config: ServiceConfig = service_config - - def integrate_flow(self, _flow: BaseFlow) -> bool: - """Integrate a standard flow into the service. - - Return True if the flow was integrated and should be logged; False - otherwise. Default implementation does nothing. - """ - return False - - def integrate_tool_flow(self, _flow: BaseToolFlow) -> bool: - """Integrate a tool-callable flow into the service. - - Return True if the tool flow was integrated; False otherwise. - """ - return False - - def integrate_stream_flow(self, _flow: BaseFlow) -> bool: - """Integrate a streaming flow into the service. - - Return True if the streaming flow was integrated; False otherwise. - """ - return False - - def run(self): - """Integrate all registered flows and start the service runtime. - - Iterates over registered flows, integrates them according to type, then - prints the logo (optional) and suppresses deprecation warnings. Concrete - services should call super().run() before their own startup logic. - """ - flow_names = [] - for _, flow in C.flow_dict.items(): - assert isinstance(flow, BaseFlow) - if flow.stream: - if self.integrate_stream_flow(flow): - flow_names.append(flow.name) - - elif isinstance(flow, BaseToolFlow): - if self.integrate_tool_flow(flow): - flow_names.append(flow.name) - - else: - if self.integrate_flow(flow): - flow_names.append(flow.name) - - logger.info(f"integrate {','.join(flow_names)}") - - import warnings - - warnings.filterwarnings("ignore", category=DeprecationWarning) diff --git a/flowllm/core/service/cmd_service.py b/flowllm/core/service/cmd_service.py deleted file mode 100644 index 801c172..0000000 --- a/flowllm/core/service/cmd_service.py +++ /dev/null @@ -1,31 +0,0 @@ -"""Command-line service for executing a single configured flow. - -This service runs a specified flow once, synchronously or asynchronously, -logging the final answer if present. -""" - -import asyncio - -from loguru import logger - -from .base_service import BaseService -from ..context import C -from ..flow import CmdFlow - - -@C.register_service("cmd") -class CmdService(BaseService): - """Service that executes a configured command flow and exits.""" - - def run(self): - """Execute the configured flow and print the result if available.""" - super().run() - cmd_config = self.service_config.cmd - flow = CmdFlow(flow=cmd_config.flow) - if flow.async_mode: - response = asyncio.run(flow.async_call(**self.service_config.cmd.params)) - else: - response = flow.call(**self.service_config.cmd.params) - - if response.answer: - logger.info(f"response.answer={response.answer}") diff --git a/flowllm/core/service/http_service.py b/flowllm/core/service/http_service.py deleted file mode 100644 index 0cb3a67..0000000 --- a/flowllm/core/service/http_service.py +++ /dev/null @@ -1,129 +0,0 @@ -"""HTTP service for serving flows via FastAPI. - -This service exposes registered flows as HTTP endpoints, including regular, -tool-callable, and streaming flows. It also provides a health check endpoint. -""" - -import asyncio -import os -from typing import AsyncGenerator - -import uvicorn -from fastapi import FastAPI -from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import StreamingResponse -from loguru import logger - -from .base_service import BaseService -from ..context import C -from ..enumeration import ChunkEnum -from ..flow import BaseFlow, BaseToolFlow -from ..schema import FlowResponse, FlowStreamChunk -from ..utils.pydantic_utils import create_pydantic_model - - -@C.register_service("http") -class HttpService(BaseService): - """FastAPI-based HTTP server that exposes registered flows as endpoints.""" - - def __init__(self, **kwargs): - """Initialize the FastAPI application and middleware.""" - super().__init__(**kwargs) - self.app = FastAPI(title=os.getenv("FLOW_APP_NAME")) - self.app.add_middleware( - CORSMiddleware, - allow_origins=["*"], - allow_credentials=True, - allow_methods=["*"], - allow_headers=["*"], - ) - - def health_check(): - return {"status": "healthy"} - - self.app.get("/health")(health_check) - - def integrate_flow(self, flow: BaseFlow) -> bool: - """Expose a non-tool, non-stream flow as a POST endpoint returning JSON.""" - request_model = create_pydantic_model(flow.name) - - async def execute_endpoint(request: request_model) -> FlowResponse: - return await flow.async_call(**request.model_dump(exclude_none=True)) - - self.app.post(f"/{flow.name}", response_model=FlowResponse)(execute_endpoint) - return True - - def integrate_tool_flow(self, flow: BaseToolFlow) -> bool: - """Expose a tool-callable flow as a POST endpoint with OpenAPI details.""" - request_model = create_pydantic_model(flow.name, input_schema=flow.tool_call.input_schema) - - async def execute_endpoint(request: request_model) -> FlowResponse: - return await flow.async_call(**request.model_dump(exclude_none=True)) - - # include tool description in OpenAPI, parameters are described via request_model fields - self.app.post( - f"/{flow.name}", - response_model=FlowResponse, - description=flow.tool_call.description, - )(execute_endpoint) - return True - - def integrate_stream_flow(self, flow: BaseFlow) -> bool: - """Expose a streaming flow as a Server-Sent Events endpoint.""" - request_model = create_pydantic_model(flow.name) - - async def execute_stream_endpoint(request: request_model) -> StreamingResponse: - stream_queue = asyncio.Queue() - task = asyncio.create_task( - flow.async_call(stream_queue=stream_queue, **request.model_dump(exclude_none=True)), - ) - - async def generate_stream() -> AsyncGenerator[bytes, None]: - while True: - try: - stream_chunk: FlowStreamChunk = await asyncio.wait_for(stream_queue.get(), timeout=1.0) - if stream_chunk.done: - yield "data:[DONE]\n\n".encode("utf-8") - await task # Ensure task completes - break - yield f"data:{stream_chunk.model_dump_json()}\n\n".encode("utf-8") - - except asyncio.TimeoutError: - # Timeout: check if task has completed or failed - if task.done(): - try: - await task # This will raise exception if task failed - except Exception as e: - logger.exception(f"flow={flow.name} encounter error with args={e.args}") - - # Task failed, send error chunk - error_chunk = FlowStreamChunk(chunk_type=ChunkEnum.ERROR, chunk=str(e), done=True) - yield f"data:{error_chunk.model_dump_json()}\n\n".encode("utf-8") - yield "data:[DONE]\n\n".encode("utf-8") - break - - else: - # Task completed successfully but no done chunk received - # This shouldn't happen normally, but handle gracefully - yield "data:[DONE]\n\n".encode("utf-8") - break - - # Task still running, continue waiting for chunks - continue - - return StreamingResponse(generate_stream(), media_type="text/event-stream") - - self.app.post(f"/{flow.name}")(execute_stream_endpoint) - return True - - def run(self): - """Run the FastAPI app with the configured Uvicorn settings.""" - super().run() - http_config = self.service_config.http - uvicorn.run( - self.app, - host=http_config.host, - port=http_config.port, - timeout_keep_alive=http_config.timeout_keep_alive, - limit_concurrency=http_config.limit_concurrency, - ) diff --git a/flowllm/core/service/mcp_service.py b/flowllm/core/service/mcp_service.py deleted file mode 100644 index 6881c2e..0000000 --- a/flowllm/core/service/mcp_service.py +++ /dev/null @@ -1,62 +0,0 @@ -"""MCP service for exposing flows as Model Context Protocol tools. - -This service registers tool-callable flows with FastMCP and runs the selected -transport (SSE, HTTP, or STDIO). -""" - -import os - -from fastmcp import FastMCP -from fastmcp.tools import FunctionTool - -from .base_service import BaseService -from ..context import C -from ..flow import BaseToolFlow -from ..utils.pydantic_utils import create_pydantic_model - - -@C.register_service("mcp") -class MCPService(BaseService): - """Service that exposes tool flows via FastMCP transports.""" - - def __init__(self, **kwargs): - """Initialize the MCP server instance.""" - super().__init__(**kwargs) - self.mcp = FastMCP(name=os.getenv("FLOW_APP_NAME")) - - def integrate_tool_flow(self, flow: BaseToolFlow) -> bool: - """Register a tool-callable flow as an MCP tool.""" - request_model = create_pydantic_model(flow.name, flow.tool_call.input_schema) - - async def execute_tool(**kwargs): - response = await flow.async_call(**request_model(**kwargs).model_dump(exclude_none=True)) - return response.answer - - # add tool - tool_call_schema = flow.tool_call.simple_input_dump() - parameters = tool_call_schema[tool_call_schema["type"]]["parameters"] - tool = FunctionTool( - name=flow.name, - description=flow.tool_call.description, - fn=execute_tool, - parameters=parameters, - ) - - self.mcp.add_tool(tool) - return True - - def run(self): - """Run the MCP server using the configured transport and options.""" - super().run() - mcp_config = self.service_config.mcp - - if mcp_config.transport == "stdio": - self.mcp.run(transport="stdio", show_banner=False) - else: - assert mcp_config.transport in ["http", "sse", "streamable-http"] - self.mcp.run( - transport=mcp_config.transport, - host=mcp_config.host, - port=mcp_config.port, - show_banner=False, - ) diff --git a/flowllm/core/storage/__init__.py b/flowllm/core/storage/__init__.py deleted file mode 100644 index 4cf4ba3..0000000 --- a/flowllm/core/storage/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -"""Storage module for managing data caching and persistence. - -This module provides utilities for caching and storing various data types with -support for expiration, metadata tracking, and automatic cleanup. It includes: - -- CacheHandler: Cache utility class supporting DataFrame, list, string, and dict - with local storage and data expiration functionality -""" - -from .cache_handler import CacheHandler - -__all__ = [ - "CacheHandler", -] diff --git a/flowllm/core/storage/cache_handler.py b/flowllm/core/storage/cache_handler.py deleted file mode 100644 index aad6b3e..0000000 --- a/flowllm/core/storage/cache_handler.py +++ /dev/null @@ -1,463 +0,0 @@ -""" -CacheHandler utility that supports DataFrame, list, string, -and dict with local storage and data expiration functionality -""" - -import json -from datetime import datetime, timedelta -from pathlib import Path -from typing import Optional, Dict, Any, Union - -import pandas as pd -from loguru import logger - - -class CacheHandler: - """ - Cache utility class supporting DataFrame, list, string, and dict - - Features: - - Support for DataFrame (CSV), list (JSON), string (TXT), and dict (JSON) - - Support for data expiration time settings - - Automatic cleanup of expired data - - Recording and managing update timestamps - """ - - def __init__(self, cache_dir: str = "cache"): - self.cache_dir = Path(cache_dir) - self.cache_dir.mkdir(parents=True, exist_ok=True) - self.metadata_file = self.cache_dir / "metadata.json" - self.metadata = {} - self._load_metadata() - - def _load_metadata(self): - """Load metadata""" - if self.metadata_file.exists(): - try: - with open(self.metadata_file, "r", encoding="utf-8") as f: - self.metadata = json.load(f) - except Exception as e: - logger.warning(f"Failed to load metadata: {e}") - self.metadata = {} - - def _save_metadata(self): - """Save metadata""" - try: - with open(self.metadata_file, "w", encoding="utf-8") as f: - json.dump(self.metadata, f, ensure_ascii=False, indent=2) - except Exception as e: - logger.error(f"Failed to save metadata: {e}") - - @staticmethod - def _get_file_extension(data_type: type) -> str: - """Get file extension for data type""" - if data_type is pd.DataFrame: - return ".csv" - elif data_type == dict: - return ".json" - elif data_type == list: - return ".json" - elif data_type == str: - return ".txt" - else: - return ".dat" - - def _get_file_path(self, key: str, data_type: type = None) -> Path: - """Get data file path with appropriate extension""" - if data_type is None: - # Try to get extension from metadata - if key in self.metadata and "data_type" in self.metadata[key]: - stored_type_name = self.metadata[key]["data_type"] - if stored_type_name == "DataFrame": - extension = ".csv" - elif stored_type_name == "dict": - extension = ".json" - elif stored_type_name == "list": - extension = ".json" - elif stored_type_name == "str": - extension = ".txt" - else: - extension = ".dat" - else: - extension = ".dat" - else: - extension = self._get_file_extension(data_type) - - return self.cache_dir / f"{key}{extension}" - - @staticmethod - def _save_dataframe(data: pd.DataFrame, file_path: Path, **kwargs) -> Dict[str, Any]: - """Save DataFrame as CSV""" - csv_params = { - "index": False, - "encoding": "utf-8", - } - csv_params.update(kwargs) - data.to_csv(file_path, **csv_params) - return { - "row_count": len(data), - "column_count": len(data.columns), - "file_size": file_path.stat().st_size, - } - - @staticmethod - def _load_dataframe(file_path: Path, **kwargs) -> pd.DataFrame: - """Load DataFrame from CSV""" - csv_params = { - "encoding": "utf-8", - } - csv_params.update(kwargs) - return pd.read_csv(file_path, **csv_params) - - @staticmethod - def _save_dict(data: dict, file_path: Path, **kwargs) -> Dict[str, Any]: - """Save dict as JSON""" - json_params = { - "ensure_ascii": False, - "indent": 2, - **kwargs, - } - json_params.update(kwargs) - with open(file_path, "w", encoding="utf-8") as f: - json.dump(data, f, **json_params) - return { - "key_count": len(data), - "file_size": file_path.stat().st_size, - } - - @staticmethod - def _load_dict(file_path: Path, **kwargs) -> dict: - """Load dict from JSON""" - with open(file_path, "r", encoding="utf-8") as f: - return json.load(f, **kwargs) - - @staticmethod - def _save_list(data: list, file_path: Path, **kwargs) -> Dict[str, Any]: - """Save list as JSON""" - json_params = { - "ensure_ascii": False, - "indent": 2, - } - json_params.update(kwargs) - with open(file_path, "w", encoding="utf-8") as f: - json.dump(data, f, **json_params) - return { - "item_count": len(data), - "file_size": file_path.stat().st_size, - } - - @staticmethod - def _load_list(file_path: Path, **kwargs) -> list: - """Load list from JSON""" - with open(file_path, "r", encoding="utf-8") as f: - return json.load(f, **kwargs) - - @staticmethod - def _save_string(data: str, file_path: Path, **kwargs) -> Dict[str, Any]: - """Save string as TXT""" - encoding = kwargs.get("encoding", "utf-8") - with open(file_path, "w", encoding=encoding) as f: - f.write(data) - return { - "char_count": len(data), - "file_size": file_path.stat().st_size, - } - - @staticmethod - def _load_string(file_path: Path, **kwargs) -> str: - """Load string from TXT""" - encoding = kwargs.get("encoding", "utf-8") - with open(file_path, "r", encoding=encoding) as f: - return f.read() - - def _save_data(self, data: Any, file_path: Path, data_type: type, **kwargs) -> Dict[str, Any]: - """Save data based on type""" - if data_type is pd.DataFrame: - return self._save_dataframe(data, file_path, **kwargs) - elif data_type == dict: - return self._save_dict(data, file_path, **kwargs) - elif data_type == list: - return self._save_list(data, file_path, **kwargs) - elif data_type == str: - return self._save_string(data, file_path, **kwargs) - else: - raise ValueError(f"Unsupported data type: {data_type}") - - def _load_data(self, file_path: Path, data_type_name: str, **kwargs) -> Any: - """Load data based on type name""" - if data_type_name == "DataFrame": - return self._load_dataframe(file_path, **kwargs) - elif data_type_name == "dict": - return self._load_dict(file_path, **kwargs) - elif data_type_name == "list": - return self._load_list(file_path, **kwargs) - elif data_type_name == "str": - return self._load_string(file_path, **kwargs) - else: - raise ValueError(f"Unsupported data type: {data_type_name}") - - def _is_expired(self, key: str) -> bool: - """Check if data is expired""" - if key not in self.metadata: - return True - - expire_time_str = self.metadata[key].get("expire_time") - if not expire_time_str: - return False # No expiration time set, never expires - - expire_time = datetime.fromisoformat(expire_time_str) - return datetime.now() > expire_time - - def save( - self, - key: str, - data: Union[pd.DataFrame, dict, list, str], - expire_hours: Optional[float] = None, - **kwargs, - ) -> bool: - """ - Save data to cache - - Args: - key: Cache key name - data: Data to save (DataFrame, dict, list, or str) - expire_hours: Expiration time in hours, None means never expires - **kwargs: Additional parameters for save operations (e.g., encoding for string) - - Returns: - bool: Whether save was successful - """ - try: - data_type = type(data) - - # Validate data type - if data_type not in [pd.DataFrame, dict, list, str]: - logger.error(f"Unsupported data type: {data_type}") - return False - - file_path = self._get_file_path(key, data_type) - - # Save data - handler_metadata = self._save_data(data, file_path, data_type, **kwargs) - - # Update metadata - current_time = datetime.now() - self.metadata[key] = { - "created_time": current_time.isoformat(), - "updated_time": current_time.isoformat(), - "expire_time": ((current_time + timedelta(hours=expire_hours)).isoformat() if expire_hours else None), - "data_type": data_type.__name__, - **handler_metadata, - } - - self._save_metadata() - return True - - except Exception as e: - logger.exception(f"Failed to save data: {e}") - return False - - def load( - self, - key: str, - auto_clean_expired: bool = True, - **kwargs, - ) -> Optional[Any]: - """ - Load data from cache - - Args: - key: Cache key name - auto_clean_expired: Whether to automatically clean expired data - **kwargs: Additional parameters for load operations (e.g., encoding for string) - - Returns: - Optional[Any]: Loaded data, returns None if not exists or expired - """ - try: - # Check if expired - if self._is_expired(key): - if auto_clean_expired: - self.delete(key) - return None - - file_path = self._get_file_path(key) - if not file_path.exists(): - return None - - # Get data type from metadata - if key not in self.metadata or "data_type" not in self.metadata[key]: - logger.warning(f"No data type information found for key '{key}'") - return None - - data_type_name = self.metadata[key]["data_type"] - - # Load data - data = self._load_data(file_path, data_type_name, **kwargs) - - # Update last access time - if key in self.metadata: - self.metadata[key]["last_accessed"] = datetime.now().isoformat() - self._save_metadata() - - return data - - except Exception as e: - logger.exception(f"Failed to load data: {e}") - return None - - def exists(self, key: str, check_expired: bool = True) -> bool: - """ - Check if cache exists - - Args: - key: Cache key name - check_expired: Whether to check expiration status - - Returns: - bool: Whether cache exists and is not expired - """ - if check_expired and self._is_expired(key): - return False - - file_path = self._get_file_path(key) - return file_path.exists() and key in self.metadata - - def delete(self, key: str) -> bool: - """ - Delete cache - - Args: - key: Cache key name - - Returns: - bool: Whether deletion was successful - """ - try: - file_path = self._get_file_path(key) - - # Delete data file - if file_path.exists(): - file_path.unlink() - - # Delete metadata - if key in self.metadata: - del self.metadata[key] - self._save_metadata() - - return True - - except Exception as e: - logger.exception(f"Failed to delete cache: {e}") - return False - - def clean_expired(self) -> int: - """ - Clean all expired caches - - Returns: - int: Number of cleaned caches - """ - expired_keys = [] - - for key in list(self.metadata.keys()): - if self._is_expired(key): - expired_keys.append(key) - - cleaned_count = 0 - for key in expired_keys: - if self.delete(key): - cleaned_count += 1 - - return cleaned_count - - def get_info(self, key: str) -> Optional[Dict[str, Any]]: - """ - Get cache information - - Args: - key: Cache key name - - Returns: - Optional[Dict]: Cache information including creation time, update time, expiration time, etc. - """ - if key not in self.metadata: - return None - - info = self.metadata[key].copy() - info["key"] = key - info["is_expired"] = self._is_expired(key) - info["file_path"] = str(self._get_file_path(key)) - - return info - - def list_all(self, include_expired: bool = False) -> Dict[str, Dict[str, Any]]: - """ - List all caches - - Args: - include_expired: Whether to include expired caches - - Returns: - Dict: Information of all caches - """ - result = {} - - for key in self.metadata: - if not include_expired and self._is_expired(key): - continue - - info = self.get_info(key) - if info: - result[key] = info - - return result - - def get_cache_stats(self) -> Dict[str, Any]: - """ - Get cache statistics - - Returns: - Dict: Cache statistics information - """ - total_count = len(self.metadata) - expired_count = sum(1 for key in self.metadata if self._is_expired(key)) - active_count = total_count - expired_count - - total_size = 0 - for key in self.metadata: - file_path = self._get_file_path(key) - if file_path.exists(): - total_size += file_path.stat().st_size - - return { - "total_count": total_count, - "active_count": active_count, - "expired_count": expired_count, - "total_size_bytes": total_size, - "total_size_mb": round(total_size / (1024 * 1024), 2), - "cache_dir": str(self.cache_dir), - } - - def clear_all(self) -> bool: - """ - Clear all caches - - Returns: - bool: Whether clearing was successful - """ - try: - # Delete all data files (CSV, JSON, TXT, and other supported formats) - for data_file in self.cache_dir.glob("*"): - if data_file.is_file() and data_file.name != "metadata.json": - data_file.unlink() - - # Clear metadata - self.metadata = {} - self._save_metadata() - - return True - - except Exception as e: - logger.exception(f"Failed to clear cache: {e}") - return False diff --git a/flowllm/core/token/__init__.py b/flowllm/core/token/__init__.py deleted file mode 100644 index 45fbc54..0000000 --- a/flowllm/core/token/__init__.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Token counting adapters for FlowLLM runtimes. - -This subpackage centralizes reusable helpers for estimating token usage across -different model providers and SDKs: - -- ``BaseToken``: Provides a character-length heuristic, defaulting to len(content)/4. -- ``OpenAIToken``: Implements counting logic compatible with OpenAI/gpt-* APIs. -- ``HuggingFaceToken``: Uses HuggingFace tokenizers and chat templates. - -Typical usage: - from flowllm.core.token import OpenAIToken - -The module-level ``__all__`` surfaces the primary token helpers for convenience. -""" - -from .base_token import BaseToken -from .huggingface_token import HuggingFaceToken -from .openai_token import OpenAIToken - -__all__ = [ - "BaseToken", - "HuggingFaceToken", - "OpenAIToken", -] diff --git a/flowllm/core/token/base_token.py b/flowllm/core/token/base_token.py deleted file mode 100644 index ffe3124..0000000 --- a/flowllm/core/token/base_token.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Base token counting helper implementation. - -This module defines the common base class for token counting utilities used -by language models. It provides a simple, model-agnostic implementation that -estimates token usage from message content and optional tools. -""" - -import math -from typing import List - -from loguru import logger - -from ..context import C -from ..schema import Message, ToolCall - - -@C.register_token_counter("base") -class BaseToken: - """Common base class for adapters that count tokens for a model. - - This class provides a lightweight, heuristic token counter that can be - used as a default implementation when a model-specific tokenizer is not - available. It estimates token usage based on the total character length of - messages and tools. - """ - - def __init__(self, model_name: str, **kwargs): - """Initialize the base token counter. - - Args: - model_name: Name of the LLM, or the name / path of the - pretrained model this counter is associated with. - **kwargs: Additional model-specific configuration parameters. - They are stored for potential use by subclasses but are not - interpreted in this base implementation. - """ - self.model_name: str = model_name - self.kwargs: dict = kwargs - - def token_count( - self, - messages: List[Message], - tools: List[ToolCall] | None = None, - **_kwargs, - ) -> int: - """Estimate the token count for the provided messages and tools. - - This implementation uses a simple heuristic that divides the total - number of characters by 4 to approximate token usage, which aligns - roughly with many BPE-based tokenizers for English text. - - Args: - messages: Conversation history to be sent to the model. - tools: Optional list of tool call definitions associated with - the request. Tool names, descriptions, and arguments are - included in the character-based estimation. - **_kwargs: Extra keyword arguments ignored by this base - implementation but accepted for API compatibility with - more advanced token counters. - - Returns: - Estimated token count as a non-negative integer. Returns ``0`` - when there is no content to count. - """ - total_chars = 0 - logger.info("token count: using rule") - - for message in messages: - content = message.content - if isinstance(content, bytes): - content = content.decode("utf-8", errors="ignore") - total_chars += len(content) - - if message.reasoning_content: - total_chars += len(message.reasoning_content) - - if tools: - for tool in tools: - total_chars += len(tool.name) - total_chars += len(tool.description) - total_chars += len(tool.arguments) - - return math.ceil(total_chars / 4) if total_chars else 0 diff --git a/flowllm/core/token/huggingface_token.py b/flowllm/core/token/huggingface_token.py deleted file mode 100644 index eca301f..0000000 --- a/flowllm/core/token/huggingface_token.py +++ /dev/null @@ -1,99 +0,0 @@ -"""Token counting utilities for HuggingFace chat models.""" - -import os -from typing import Any, List - -from loguru import logger - -from .base_token import BaseToken -from ..context import C -from ..schema import Message, ToolCall - - -@C.register_token_counter("hf") -class HuggingFaceToken(BaseToken): - """Estimate token usage with `transformers.AutoTokenizer.apply_chat_template`.""" - - def __init__( - self, - model_name: str, - *, - use_fast: bool = False, - trust_remote_code: bool = False, - use_mirror: bool = True, - **kwargs, - ): - """Initialize HuggingFace token counter. - - Args: - model_name: Name or path of the HuggingFace model. - use_fast: Pass-through flag for ``AutoTokenizer.from_pretrained``. - trust_remote_code: Whether to trust remote code when loading tokenizers. - use_mirror: If ``True``, set ``HF_ENDPOINT`` to the HuggingFace mirror. - **kwargs: Extra keyword arguments forwarded to ``BaseToken``. - """ - super().__init__(model_name=model_name, **kwargs) - - self.use_fast: bool = use_fast - self.trust_remote_code: bool = trust_remote_code - self.use_mirror: bool = use_mirror - - # Lazily initialized tokenizer instance - self._tokenizer: Any | None = None - - def _ensure_tokenizer(self) -> Any: - """Lazily create and cache the tokenizer instance.""" - if self._tokenizer is not None: - return self._tokenizer - - if self.use_mirror: - os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com") - - from transformers import AutoTokenizer - - logger.info(f"token count: using model={self.model_name}") - tokenizer = AutoTokenizer.from_pretrained( - self.model_name, - use_fast=self.use_fast, - trust_remote_code=self.trust_remote_code, - **self.kwargs, - ) - - if tokenizer.chat_template is None: - raise ValueError(f"The tokenizer for model {self.model_name} is missing a chat template.") - - self._tokenizer = tokenizer - return tokenizer - - @staticmethod - def _serialize_messages(messages: List[Message]) -> list[dict]: - """Convert internal `Message` objects to simple dictionaries.""" - return [msg.simple_dump() for msg in messages] - - @staticmethod - def _serialize_tools(tools: List[ToolCall] | None) -> list[dict] | None: - """Convert tool schemas into JSON-serializable dictionaries.""" - if not tools: - return None - return [tool.simple_input_dump() for tool in tools] - - def token_count( - self, - messages: List[Message], - tools: List[ToolCall] | None = None, - **kwargs, - ) -> int: - """Return the token length computed by the HuggingFace tokenizer.""" - tokenizer = self._ensure_tokenizer() - serialized_messages = self._serialize_messages(messages) - serialized_tools = self._serialize_tools(tools) - - tokenized = tokenizer.apply_chat_template( - serialized_messages, - add_generation_prompt=kwargs.pop("add_generation_prompt", False), - tokenize=True, - return_tensors="np", - tools=serialized_tools, - **kwargs, - )[0] - return int(tokenized.size if hasattr(tokenized, "size") else len(tokenized)) diff --git a/flowllm/core/token/openai_token.py b/flowllm/core/token/openai_token.py deleted file mode 100644 index 48fd02f..0000000 --- a/flowllm/core/token/openai_token.py +++ /dev/null @@ -1,51 +0,0 @@ -"""Token counting implementation for OpenAI-compatible models.""" - -import json -from typing import List - -from loguru import logger - -from .base_token import BaseToken -from ..context import C -from ..schema import Message, ToolCall - - -@C.register_token_counter("openai") -class OpenAIToken(BaseToken): - """The OpenAI token counting class.""" - - def token_count( - self, - messages: List[Message], - tools: List[ToolCall] | None = None, - **_kwargs, - ) -> int: - """Estimate token usage for messages and tool payloads using tiktoken.""" - import tiktoken - - try: - encoding = tiktoken.encoding_for_model(self.model_name) - logger.info(f"token count: using model={self.model_name}") - - except KeyError: - encoding = tiktoken.get_encoding("o200k_base") - logger.info("token count: using model=o200k_base") - - num_tokens = 0 - # <|im_start|>system\n...<|im_end|> - for message in messages: - msg_tokens = 4 - if message.content: - msg_tokens += len(encoding.encode(message.content)) - if message.tool_calls: - for tool_call in message.tool_calls: - tool_call_content = json.dumps(tool_call.simple_output_dump(), ensure_ascii=False) - msg_tokens += len(encoding.encode(tool_call_content)) - - num_tokens += msg_tokens - - if tools: - for tool in tools: - tool_content = json.dumps(tool.simple_input_dump(), ensure_ascii=False) - num_tokens += len(encoding.encode(tool_content)) - return num_tokens diff --git a/flowllm/core/utils/__init__.py b/flowllm/core/utils/__init__.py deleted file mode 100644 index 8d2c92d..0000000 --- a/flowllm/core/utils/__init__.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Utility modules for the flowllm core package. - -This package provides utility functions and classes for: -- HTTP client operations for executing flows -- Common utility functions (string conversion, env loading, content extraction) -- MCP (Model Context Protocol) client operations -- LLM message formatting and processing utilities -- Timer utilities for measuring execution time -- Logger utilities for initializing and configuring loguru logger -- Logo and banner printing utilities -- Pydantic configuration parser for loading and merging configurations -- Tushare API client for financial market data -""" - -from .common_utils import ( - camel_to_snake, - load_env, - parse_flow_expression, - singleton, - snake_to_camel, -) -from .fastmcp_client import FastMcpClient -from .http_client import HttpClient -from .llm_utils import ( - extract_content, - format_messages, - merge_messages_content, - parse_message_by_keys, -) -from .logger_utils import init_logger -from .logo_utils import print_logo -from .pydantic_config_parser import PydanticConfigParser -from .pydantic_utils import create_pydantic_model -from .timer import Timer, timer - -__all__ = [ - "HttpClient", - "FastMcpClient", - "camel_to_snake", - "snake_to_camel", - "load_env", - "extract_content", - "parse_flow_expression", - "singleton", - "Timer", - "timer", - "init_logger", - "print_logo", - "create_pydantic_model", - "PydanticConfigParser", - "format_messages", - "merge_messages_content", - "parse_message_by_keys", -] diff --git a/flowllm/core/utils/common_utils.py b/flowllm/core/utils/common_utils.py deleted file mode 100644 index 80fac2c..0000000 --- a/flowllm/core/utils/common_utils.py +++ /dev/null @@ -1,223 +0,0 @@ -"""Common utility functions for string conversion, environment loading, content extraction, -and flow expression parsing. - -This module provides utility functions for: -- String format conversion (camelCase to snake_case and vice versa) -- Environment variable loading from .env files -- Extracting and parsing content from Markdown code blocks -- Parsing flow expressions into composed `BaseOp` trees -""" - -import os -import re -from pathlib import Path -from typing import TYPE_CHECKING - -from loguru import logger - -if TYPE_CHECKING: - from ..op import BaseOp - -ENV_LOADED = False - - -def camel_to_snake(content: str) -> str: - """Convert a camelCase or PascalCase string to snake_case. - - Args: - content: The camelCase or PascalCase string to convert. - - Returns: - The converted snake_case string. - - Example: - ```python - camel_to_snake("BaseWorker") - # 'base_worker' - camel_to_snake("MyLLMClass") - # 'my_llm_class' - ``` - """ - content = content.replace("LLM", "Llm") - snake_str = re.sub(r"(? str: - """Convert a snake_case string to PascalCase (camelCase with first letter capitalized). - - Args: - content: The snake_case string to convert. - - Returns: - The converted PascalCase string. - - Example: - ```python - snake_to_camel("base_worker") - # 'BaseWorker' - snake_to_camel("my_llm_class") - # 'MyLLMClass' - ``` - """ - camel_str = "".join(x.capitalize() for x in content.split("_")) - camel_str = camel_str.replace("Llm", "LLM") - return camel_str - - -def _load_env(path: Path): - """Load environment variables from a .env file. - - Reads the specified .env file line by line, parses key-value pairs, - and sets them as environment variables. Lines starting with '#' are - treated as comments and skipped. - - Args: - path: Path to the .env file to load. - """ - with path.open() as f: - for line in f: - line = line.strip() - if line.startswith("#"): - continue - - line_split = line.strip().split("=", 1) - if len(line_split) >= 2: - key = line_split[0].strip() - value = line_split[1].strip().strip('"') - os.environ[key] = value - - -def load_env(path: str | Path = None, enable_log: bool = True): - """Load environment variables from a .env file. - - This function ensures that environment variables are loaded only once - (controlled by the ENV_LOADED flag). If a path is provided, it loads - from that specific path. Otherwise, it searches for a .env file in the - current directory and up to 4 parent directories. - - Args: - path: Optional path to the .env file. If None, searches for .env - in current and parent directories. - enable_log: Whether to log when the .env file is found and loaded. - - Note: - The function uses a global flag to ensure it only runs once per - execution. Subsequent calls will be ignored. - """ - global ENV_LOADED - if ENV_LOADED: - return - - if path is not None: - path = Path(path) - if path.exists(): - _load_env(path) - ENV_LOADED = True - - else: - for i in range(5): - path = Path("../" * i + ".env") - if path.exists(): - if enable_log: - logger.info(f"load env_path={path}") - _load_env(path) - ENV_LOADED = True - return - - logger.warning(".env not found") - - -def singleton(cls): - """Decorator to create a singleton class. - - Ensures that only one instance of the decorated class is created. - Subsequent instantiations will return the same instance. - - Args: - cls: The class to be decorated as a singleton. - - Returns: - A wrapper function that returns the singleton instance of the class. - - Example: - ```python - @singleton - class Config: - def __init__(self): - self.value = 42 - - c1 = Config() - c2 = Config() - c1 is c2 # True - ``` - """ - _instance = {} - - def _singleton(*args, **kwargs): - if cls not in _instance: - _instance[cls] = cls(*args, **kwargs) - return _instance[cls] - - return _singleton - - -def parse_flow_expression(flow_content: str) -> "BaseOp": - """Parse a textual flow script into a composed `BaseOp` operation tree. - - The flow content must explicitly instantiate op classes that have been - registered in `C.registry_dict["op"]`. Multi-line content is supported: - all preceding non-empty lines are executed in a restricted environment to - prepare context (e.g., variable assignments), and the last non-empty line - is evaluated as an expression that must return a `BaseOp`. - - Rules for the final line: - - The last non-empty line MUST be an expression that evaluates to a - single `BaseOp` instance or a composed `BaseOp` expression (e.g., - `OpA() >> OpB()` or `OpA() | OpB()`). - - It CANNOT be an assignment statement (e.g., `opx = OpA() >> OpB()`). - Assignments are only allowed in preceding lines. - - The evaluated result of this final expression is the flow to run. - - Examples: - OpA() >> OpB() - - op = ContainerOp() - op.ops.search = OpA() - op.ops.find = OpB() - (op | OpB()) >> (OpA() | OpC()) >> op - - Args: - flow_content: Raw flow script text. - - Returns: - BaseOp: The parsed flow as an executable operation tree. - - Raises: - ValueError: If the provided content is empty. - AssertionError: If the last line does not evaluate to a `BaseOp`. - """ - from ..context import C - from ..enumeration import RegistryEnum - from ..op import BaseOp - - flow_content = flow_content.strip() - if not flow_content: - raise ValueError("flow content is empty") - - # Prepare environment with registered op classes available by name. - # When an op is registered without an explicit name, fall back to its class name. - op_registry = C.registry_dict[RegistryEnum.OP] - env: dict = {name or cls.__name__: cls for name, cls in op_registry.items()} - - # Execute all but the last non-empty line to set up context - lines = [x.strip() for x in flow_content.splitlines() if x.strip()] - if len(lines) > 1: - exec_content = "\n".join(lines[:-1]) - exec(exec_content, {"__builtins__": {}}, env) - - # Evaluate the final line; it must return a BaseOp - last_line_expr = lines[-1] - result = eval(last_line_expr, {"__builtins__": {}}, env) - assert isinstance(result, BaseOp), f"Expression '{last_line_expr}' did not evaluate to a BaseOp instance" - return result diff --git a/flowllm/core/utils/fastmcp_client.py b/flowllm/core/utils/fastmcp_client.py deleted file mode 100644 index e21e7ad..0000000 --- a/flowllm/core/utils/fastmcp_client.py +++ /dev/null @@ -1,335 +0,0 @@ -"""Asynchronous MCP client using FastMCP. - -This module provides an async MCP client wrapper around FastMCP's Client class, -supporting both stdio (command) and HTTP (URL) connection methods similar to -the original mcp_client.py implementation. -""" - -import asyncio -import os -import shutil -from typing import List, Optional, Union - -import mcp.types -from fastmcp import Client -from fastmcp.client.client import CallToolResult -from fastmcp.client.transports import StdioTransport, SSETransport, StreamableHttpTransport -from loguru import logger - -from ..schema.tool_call import ToolCall - - -class FastMcpClient: - """Asynchronous MCP client using FastMCP, supporting stdio and HTTP connections. - - This client wraps FastMCP's Client class and provides a similar interface to - the original McpClient, supporting: - - Stdio connections via command - - HTTP connections via URL (SSE or Streamable HTTP) - - Retry mechanisms - - Timeout configuration - - Environment variable injection - - Example: - ```python - # Using stdio connection - config = { - "command": "npx", - "args": ["-y", "@modelcontextprotocol/server-filesystem"], - } - async with FastMcpClient("mcp", config) as client: - tools = await client.list_tools() - result = await client.call_tool("read_file", {"path": "/tmp/test.txt"}) - - # Using HTTP connection - config = { - "type": "sse", - "url": "https://example.com/sse", - "headers": {"Authorization": "Bearer {API_KEY}"}, - } - async with FastMcpClient("mcp", config) as client: - tools = await client.list_tools() - ``` - """ - - def __init__( - self, - name: str, - config: dict, - append_env: bool = False, - max_retries: int = 3, - timeout: Optional[float] = None, - ): - """Initialize the FastMCP client. - - Args: - name: Name identifier for the MCP server - config: Configuration dictionary with either: - - For stdio: `command` (and optionally `args`, `env`) - - For HTTP: `url` (and optionally `type`, `headers`, `timeout`, `sse_read_timeout`) - append_env: Whether to append environment variables to the config's env - max_retries: Maximum number of retries for connection and operations - timeout: Optional timeout in seconds for operations - """ - self.name: str = name - self.config: dict = config - self.append_env: bool = append_env - self.max_retries: int = max_retries - self.timeout: Optional[float] = timeout - - self.client: Optional[Client] = None - self._transport = self._create_transport() - - def _create_transport(self): - """Create the appropriate transport based on config.""" - command = self.config.get("command") - - if command: - # Stdio transport - if command == "npx": - command = shutil.which("npx") or command - - env_params: dict = {} - if self.append_env: - env_params.update(os.environ) - if self.config.get("env"): - env_params.update(self.config["env"]) - - return StdioTransport( - command=command, - args=self.config.get("args", []), - env=env_params if env_params else None, - cwd=self.config.get("cwd"), - ) - else: - # HTTP transport (SSE or Streamable HTTP) - url = self.config["url"] - transport_type = self.config.get("type", "").lower() - kwargs: dict = {"url": url} - - # Handle headers with environment variable substitution - if self.config.get("headers"): - headers = self.config.get("headers").copy() - if headers.get("Authorization"): - assert isinstance(headers["Authorization"], str) - headers["Authorization"] = headers["Authorization"].format(**os.environ) - kwargs["headers"] = headers - - # Handle timeout - if "timeout" in self.config: - kwargs["sse_read_timeout"] = self.config["timeout"] - - # Determine transport type based on config or URL - if transport_type in ["streamable_http", "streamablehttp"]: - return StreamableHttpTransport(**kwargs) - else: - # Default to SSE for /sse endpoints, otherwise StreamableHttp - if url.endswith("/sse"): - return SSETransport(**kwargs) - else: - # For URLs without /sse, use StreamableHttpTransport - return StreamableHttpTransport(**kwargs) - - async def __aenter__(self) -> "FastMcpClient": - """Async context manager entry.""" - for i in range(self.max_retries): - try: - self.client = Client( - transport=self._transport, - name=self.name, - timeout=self.timeout, - ) - - # Use FastMCP Client as async context manager - if self.timeout is not None: - await asyncio.wait_for( - self.client.__aenter__(), - timeout=self.timeout, - ) - else: - await self.client.__aenter__() - - break - - except asyncio.TimeoutError as exc: - logger.exception(f"{self.name} start timeout after {self.timeout}s") - - # Clean up before retrying - await self._cleanup_client() - - if i == self.max_retries - 1: - raise TimeoutError(f"{self.name} start timeout after {self.timeout}s") from exc - - await asyncio.sleep(1 + i) - - except Exception as e: - logger.exception( - f"{self.name} start failed with {e}. " f"Retry {i + 1}/{self.max_retries} in {1 + i}s...", - ) - - # Clean up before retrying - await self._cleanup_client() - - await asyncio.sleep(1 + i) - - if i == self.max_retries - 1: - raise e - - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: - """Async context manager exit.""" - for i in range(self.max_retries): - try: - await self._cleanup_client() - break - - except Exception as e: - logger.exception( - f"{self.name} close failed with {e}. " f"Retry {i + 1}/{self.max_retries} in {1 + i}s...", - ) - await asyncio.sleep(1 + i) - - if i == self.max_retries - 1: - break - - self.client = None - - async def _cleanup_client(self): - """Clean up the client connection.""" - if self.client: - try: - await self.client.__aexit__(None, None, None) - except Exception: - pass - try: - await self.client.close() - except Exception: - pass - self.client = None - - async def list_tools(self) -> List[mcp.types.Tool]: - """List available tools from the MCP server. - - Returns: - List of MCP Tool objects - - Raises: - RuntimeError: If the client is not connected - TimeoutError: If the operation times out - """ - if not self.client: - raise RuntimeError(f"Server {self.name} not initialized") - - tools = [] - for i in range(self.max_retries): - try: - if self.timeout is not None: - tools = await asyncio.wait_for( - self.client.list_tools(), - timeout=self.timeout, - ) - else: - tools = await self.client.list_tools() - break - - except asyncio.TimeoutError as exc: - logger.exception(f"{self.name} list tools timeout after {self.timeout}s") - - if i == self.max_retries - 1: - raise TimeoutError(f"{self.name} list tools timeout after {self.timeout}s") from exc - - await asyncio.sleep(1 + i) - - except Exception as e: - logger.exception( - f"{self.name} list tools failed with {e}. " f"Retry {i + 1}/{self.max_retries} in {1 + i}s...", - ) - await asyncio.sleep(1 + i) - - if i == self.max_retries - 1: - raise e - - return tools - - async def list_tool_calls(self) -> List[ToolCall]: - """List available tools as ToolCall objects. - - Returns: - List of ToolCall objects converted from MCP Tools - - Raises: - RuntimeError: If the client is not connected - """ - if not self.client: - raise RuntimeError(f"Server {self.name} not initialized") - - tools = await self.list_tools() - return [ToolCall.from_mcp_tool(t) for t in tools] - - async def call_tool( - self, - tool_name: str, - arguments: dict, - parse_result: bool = False, - ) -> Union[str, CallToolResult]: - """Call a tool on the MCP server. - - Args: - tool_name: Name of the tool to call - arguments: Dictionary of arguments to pass to the tool - parse_result: parse CallToolResult to str - - Returns: - The result from the tool call - - Raises: - RuntimeError: If the client is not connected - TimeoutError: If the operation times out - """ - if not self.client: - raise RuntimeError(f"Server {self.name} not initialized") - - result = None - for i in range(self.max_retries): - try: - if self.timeout is not None: - result = await asyncio.wait_for( - self.client.call_tool(tool_name, arguments), - timeout=self.timeout, - ) - else: - result = await self.client.call_tool(tool_name, arguments) - - if parse_result: - if len(result.content) == 1: - return result.content[0].text - - else: - text_content = [] - for block in result.content: - if hasattr(block, "text"): - text_content.append(block.text) - return "\n".join(text_content) if text_content else result - else: - return result - - except asyncio.TimeoutError as exc: - logger.exception(f"{self.name}.{tool_name} call_tool timeout after {self.timeout}s") - - if i == self.max_retries - 1: - raise TimeoutError(f"{self.name}.{tool_name} call_tool timeout after {self.timeout}s") from exc - - await asyncio.sleep(1 + i) - - except Exception as e: - logger.exception( - f"{self.name}.{tool_name} call_tool failed with {e}. " - f"Retry {i + 1}/{self.max_retries} in {1 + i}s...", - ) - await asyncio.sleep(1 + i) - - if i == self.max_retries - 1: - raise e - - return result diff --git a/flowllm/core/utils/http_client.py b/flowllm/core/utils/http_client.py deleted file mode 100644 index baa71bc..0000000 --- a/flowllm/core/utils/http_client.py +++ /dev/null @@ -1,214 +0,0 @@ -"""Asynchronous HTTP client for executing flows with retry mechanism. - -This module provides an HttpClient class for making asynchronous HTTP requests -to flow endpoints with built-in retry logic and error handling. - -Key features: -- Async/await support for non-blocking HTTP requests -- Automatic retry mechanism with configurable retry attempts -- Flexible error handling (raise exceptions or return None) -- Context manager support for resource cleanup -- Health check and endpoint listing capabilities -""" - -import json -from typing import Dict, Optional, AsyncIterator - -import httpx -from loguru import logger - -from ..schema import FlowResponse - - -class HttpClient: - """Asynchronous HTTP client for executing flows with retry mechanism. - - This class provides an async HTTP client for interacting with flow endpoints. - It includes built-in retry logic, error handling, and context manager support. - - Features: - - Automatic retry on failure with configurable retry attempts - - Flexible error handling (raise exceptions or return None) - - Async/await support for non-blocking operations - - Context manager support for automatic resource cleanup - - Health check and endpoint discovery capabilities - """ - - def __init__( - self, - base_url: str = "http://localhost:8001", - timeout: float = 3600, - max_retries: int = 3, - raise_exception: bool = True, - ): - """Initialize the async HTTP client. - - Args: - base_url: Base URL for the flow service endpoint. Defaults to - "http://localhost:8001". - timeout: Request timeout in seconds. Defaults to 3600 (1 hour). - max_retries: Maximum number of retry attempts on failure. Defaults to 3. - raise_exception: Whether to raise exceptions on final failure. If False, - returns None after all retries fail. Defaults to True. - """ - self.base_url = base_url.rstrip("/") - self.timeout = timeout - self.max_retries = max_retries - self.raise_exception = raise_exception - self.client = httpx.AsyncClient(timeout=timeout) - - async def __aenter__(self): - """Async context manager entry point. - - Returns: - The HttpClient instance. - """ - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - """Async context manager exit point. - - Closes the HTTP client connection and cleans up resources. - - Args: - exc_type: Exception type if an exception occurred. - exc_val: Exception value if an exception occurred. - exc_tb: Exception traceback if an exception occurred. - """ - await self.client.aclose() - - async def close(self): - """Close the HTTP client connection and clean up resources. - - This method should be called when done using the client if not using - the context manager (async with statement). - """ - await self.client.aclose() - - async def health_check(self) -> Dict[str, str]: - """Check the health status of the flow service. - - Makes a GET request to the /health endpoint to verify the service - is running and responding. - - Returns: - A dictionary containing health check information. - - Raises: - httpx.HTTPStatusError: If the health check request fails. - """ - response = await self.client.get(f"{self.base_url}/health") - response.raise_for_status() - return response.json() - - async def execute_flow(self, flow_name: str, **kwargs) -> Optional[FlowResponse]: - """ - Execute the flow asynchronously with retry logic and error handling. - - This method: - 1. Makes HTTP POST request to the flow endpoint - 2. Retries on failure according to max_retries setting - 3. Returns FlowResponse or None based on raise_exception setting - - Args: - flow_name: Name of the flow to execute - **kwargs: Additional parameters passed to the flow - - Returns: - FlowResponse on success, None on failure if raise_exception=False - - Raises: - httpx.HTTPStatusError: If HTTP request fails and raise_exception=True - - curl example: - curl -X POST http://localhost:8002/demo_http_flow \ - -H "Content-Type: application/json" \ - -d '{ - "query": "what is ai?" - }' - """ - endpoint = f"{self.base_url}/{flow_name}" - result = None - - if self.max_retries == 1 and self.raise_exception: - response = await self.client.post(endpoint, json=kwargs) - response.raise_for_status() - result = FlowResponse(**response.json()) - - else: - for i in range(self.max_retries): - try: - response = await self.client.post(endpoint, json=kwargs) - response.raise_for_status() - result = FlowResponse(**response.json()) - break - except Exception as e: - logger.exception( - f"execute_flow failed for flow={flow_name}, attempt={i + 1}/{self.max_retries}, error={e.args}", - ) - - if i == self.max_retries - 1: - if self.raise_exception: - raise e - - return result - - async def list_endpoints(self) -> dict: - """List all available endpoints from the flow service. - - Retrieves the OpenAPI specification which contains information about - all available endpoints and their schemas. - - Returns: - A dictionary containing the OpenAPI specification with endpoint - information. - - Raises: - httpx.HTTPStatusError: If the request to retrieve endpoints fails. - """ - response = await self.client.get(f"{self.base_url}/openapi.json") - response.raise_for_status() - return response.json() - - async def execute_stream_flow(self, flow_name: str, **kwargs) -> AsyncIterator[Dict[str, str]]: - """Execute the flow and stream responses as they arrive. - - This streams Server-Sent Events style lines from the endpoint, parsing - entries that begin with "data:". Each valid JSON payload is expected to - include keys like "chunk_type" and "chunk" and will be yielded as a - dictionary: {"type": str, "content": str}. The stream terminates when a - line with payload "[DONE]" is received. - - Args: - flow_name: Name of the flow endpoint to call. - **kwargs: JSON payload sent to the endpoint. - - Yields: - Dict with keys "type" and "content" for each streamed chunk. If - an error occurs and raise_exception=False, an error chunk will be yielded. - """ - endpoint = f"{self.base_url}/{flow_name}" - - async with self.client.stream("POST", endpoint, json=kwargs) as response: - response.raise_for_status() - - async for line in response.aiter_lines(): - if not line: - continue - - if line.startswith("data:"): - data_content = line[5:].strip() - - if data_content == "[DONE]": - break - - try: - json_data = json.loads(data_content) - chunk_type = json_data.get("chunk_type", "answer") - chunk_content = json_data.get("chunk", "") - - if chunk_content: - yield {"type": chunk_type, "content": chunk_content} - except Exception: - # Skip lines that aren't valid JSON payloads - continue diff --git a/flowllm/core/utils/llm_utils.py b/flowllm/core/utils/llm_utils.py deleted file mode 100644 index bbf9738..0000000 --- a/flowllm/core/utils/llm_utils.py +++ /dev/null @@ -1,240 +0,0 @@ -"""Utility functions for LLM-related operations. - -This module provides utility functions for formatting and processing messages -for LLM interactions. -""" - -import json -import re -from typing import List - -from ..enumeration import Role -from ..schema import Message - - -def format_messages(messages: List[Message]) -> str: - """Format messages into a readable string representation. - - Converts a list of Message objects into a simple text format where each - message is represented as "role: content" on a separate line. This is - useful for logging, debugging, or displaying conversation history in - a human-readable format. - - Args: - messages: List of Message objects to format. Each message should have - a role and content attribute. - - Returns: - Formatted string representation of messages, with each message on a - new line in the format "role: content". Messages are separated by - newline characters. - - Example: - ```python - messages = [ - Message(role=Role.USER, content="Hello"), - Message(role=Role.ASSISTANT, content="Hi there!") - ] - formatted = format_messages(messages) - # Returns: "user: Hello\nassistant: Hi there!" - ``` - """ - formatted_lines = [] - for msg in messages: - formatted_lines.append(f"{msg.role.value}: {msg.content}") - return "\n".join(formatted_lines) - - -def merge_messages_content(messages: List[Message | dict]) -> str: - """Merge messages content into a formatted string representation. - - This function processes a list of messages (either Message objects or dicts) - and formats them into a structured string. Different message roles are - formatted differently: - - ASSISTANT: Includes reasoning content, main content, and tool calls - - USER: Includes the user content - - TOOL: Includes tool call results - - Each message is prefixed with a step number (starting from 0) to indicate - its position in the conversation sequence. - - Args: - messages: List of Message objects or dictionaries to merge. If a dict - is provided, it will be converted to a Message object. - - Returns: - Formatted string representation of all messages with step numbers. - Each message is separated by newlines and includes role information. - - Example: - ```python - messages = [ - Message(role=Role.USER, content="What's the weather?"), - Message(role=Role.ASSISTANT, content="Let me check", - tool_calls=[ToolCall(name="get_weather", arguments={})]) - ] - result = merge_messages_content(messages) - # Returns formatted string with step numbers and role information - ``` - """ - content_collector = [] - for i, message in enumerate(messages): - if isinstance(message, dict): - message = Message(**message) - - if message.role is Role.ASSISTANT: - line = ( - f"### step.{i} role={message.role.value} content=\n{message.reasoning_content}\n\n{message.content}\n" - ) - if message.tool_calls: - for tool_call in message.tool_calls: - line += f" - tool call={tool_call.name}\n params={tool_call.arguments}\n" - content_collector.append(line) - - elif message.role is Role.USER: - line = f"### step.{i} role={message.role.value} content=\n{message.content}\n" - content_collector.append(line) - - elif message.role is Role.TOOL: - line = f"### step.{i} role={message.role.value} tool call result=\n{message.content}\n" - content_collector.append(line) - - return "\n".join(content_collector) - - -def extract_content(text: str, language_tag: str = "json", greedy: bool = False): - """Extract and parse content from Markdown code blocks. - - Searches for content within Markdown code blocks (triple backticks) - and extracts it. If the language tag is "json", attempts to parse - the extracted content as JSON. If no code block is found, returns - the original text (or parsed JSON if applicable). - - Args: - text: The text to search for code blocks. - language_tag: The language tag of the code block (e.g., "json", - "python"). Defaults to "json". The tag matching is case-sensitive - and allows optional whitespace around the tag. - greedy: If False (default), uses non-greedy matching to find the - nearest closing ```. If True, uses greedy matching to find - the farthest closing ```. This is useful when there are multiple - code blocks with the same language tag. - - Returns: - If language_tag is "json": - - Parsed JSON object/dict/list if valid JSON is found - - None if JSON parsing fails - Otherwise: - - Extracted content string from code block (stripped of whitespace) - - Original text if no code block is found - - Example: - ```python - extract_content("```json\\n{\"key\": \"value\"}\\n```") # noqa - # {'key': 'value'} - extract_content("``` json\\n{\"key\": \"value\"}\\n```") # noqa - # {'key': 'value'} - extract_content("```python\\nprint('hello')\\n```", "python") - # "print('hello')" - extract_content("Some text without code blocks") - # "Some text without code blocks" - ``` - """ - # Use non-greedy (.*?) or greedy (.*) matching based on the greedy parameter - quantifier = ".*" if greedy else ".*?" - pattern = rf"```\s*{re.escape(language_tag)}\s*({quantifier})\s*```" - match = re.search(pattern, text, re.DOTALL) - - if match: - result = match.group(1).strip() - else: - result = text - - if language_tag == "json": - try: - result = json.loads(result) - - except json.JSONDecodeError: - result = None - - return result - - -def parse_message_by_keys(content: str, keys: List[str]) -> dict: - """Parse content by splitting it using a sequence of keys as delimiters. - - This function sequentially splits the content string using each key in the - provided list as a delimiter. It extracts the content segments between - consecutive keys and returns them as a dictionary where each key maps to - its corresponding content segment. - - The parsing process works as follows: - 1. For each key in the sequence, split the remaining content at the first - occurrence of that key - 2. Store the content before the key (for previous key) and continue with - the content after the key - 3. For the last key, store both the content before and after it - - Args: - content: The string content to parse - keys: List of keys (delimiters) to use for splitting, in order - - Returns: - Dictionary mapping each key to its corresponding content segment. - Keys that don't appear in the content will have empty string values. - Content before the first key is discarded. - - Example: - ```python - content = "prefixkey1middlekey2suffix" - keys = ["key1", "key2"] - result = parse_message_by_keys(content, keys) - # Returns: {"key1": "middle", "key2": "suffix"} - ``` - - Note: - - Only the first occurrence of each key is used for splitting - - Content before the first key is not included in the result - - If a key is not found, the remaining content is assigned to the - previous key (or discarded if it's the first key) - """ - origin_content = content - result = {} - # Iterate through pairs of (previous_key, current_key) - for pre_key, key in zip([None] + keys[:-1], keys): - # Split content at the first occurrence of current key (max 1 split) - origin_split = origin_content.strip().split(key, 1) - - if pre_key is None: - # First key: discard content before it, keep content after - if len(origin_split) >= 2: - origin_content = origin_split[1] - # Special case: if this is the only key, store the result - if len(keys) == 1: - result[key] = origin_split[1] - else: - # Key not found, keep original content for next iteration - # Special case: if this is the only key and not found, store empty - if len(keys) == 1: - result[key] = "" - - elif key == keys[-1]: - # Last key: store content before it (for previous key) and after it - if len(origin_split) >= 2: - result[pre_key] = origin_split[0] - result[key] = origin_split[1] - else: - # Key not found, assign remaining content to previous key - result[pre_key] = origin_content - result[key] = "" - - else: - # Middle keys: store content before current key, continue with content after - if len(origin_split) >= 2: - result[pre_key] = origin_split[0] - origin_content = origin_split[1] - else: - # Key not found, assign remaining content to previous key and continue - result[pre_key] = origin_content - - return result diff --git a/flowllm/core/utils/logger_utils.py b/flowllm/core/utils/logger_utils.py deleted file mode 100644 index d290947..0000000 --- a/flowllm/core/utils/logger_utils.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Logger utilities for initializing and configuring loguru logger. - -This module provides utilities for setting up loguru logger with file and console outputs, -including automatic log rotation, retention, and compression. -""" - -import os -import sys -from datetime import datetime - - -def init_logger(log_dir: str = "logs", level: str = "INFO"): - """Initialize and configure the loguru logger. - - Sets up loguru logger with two sinks: - 1. File sink: Logs are written to files in the specified directory with rotation, - retention, and compression settings. - 2. Console sink: Logs are also output to stdout with colorization. - - Args: - log_dir: Directory path where log files will be stored. Defaults to "logs". - level: Logging level (e.g., "INFO", "DEBUG", "WARNING", "ERROR"). Defaults to "INFO". - - Note: - The log filename is automatically generated using the current timestamp - in the format "YYYY-MM-DD_HH-MM-SS.log". - Files are rotated at midnight, retained for 7 days, and compressed as zip. - """ - from loguru import logger - - logger.remove() - os.makedirs(log_dir, exist_ok=True) - - current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") - - log_filename = f"{current_ts}.log" - log_filepath = os.path.join(log_dir, log_filename) - - logger.add( - log_filepath, - level=level, - rotation="00:00", - retention="7 days", - compression="zip", - encoding="utf-8", - format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {message}", - ) - - logger.add( - sink=sys.stdout, - level=level, - format="{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {message}", - colorize=True, - ) diff --git a/flowllm/core/utils/logo_utils.py b/flowllm/core/utils/logo_utils.py deleted file mode 100644 index 656a1f5..0000000 --- a/flowllm/core/utils/logo_utils.py +++ /dev/null @@ -1,81 +0,0 @@ -"""Logo and banner printing utilities. - -This module provides utilities for printing styled logos and service information -using pyfiglet for ASCII art and rich for terminal formatting. -""" - -from typing import TYPE_CHECKING - -from pyfiglet import Figlet -from rich.console import Console, Group -from rich.panel import Panel -from rich.table import Table -from rich.text import Text - -if TYPE_CHECKING: - from ..schema import ServiceConfig - - -def print_logo(service_config: "ServiceConfig", app_name: str, width: int = 400): - """Print a styled logo and service information banner. - - Creates and prints a formatted banner containing: - - ASCII art logo generated from the app name - - Service configuration information (backend, URL, transport) - - Version information for FlowLLM and related dependencies - - Args: - service_config: Service configuration object containing backend settings. - app_name: Name of the application to display in the logo. - width: Width for the ASCII art logo generation. Defaults to 400. - - The output is styled with colors and formatted as a panel using the rich library. - """ - f = Figlet(font="slant", width=width) - logo: str = f.renderText(app_name) - logo_text = Text(logo, style="bold green") - - info_table = Table.grid(padding=(0, 1)) - info_table.add_column(style="bold", justify="center") # Emoji column - info_table.add_column(style="bold cyan", justify="left") # Label column - info_table.add_column(style="white", justify="left") # Value column - - info_table.add_row("📦", "Backend:", service_config.backend) - - if service_config.backend == "http": - info_table.add_row("🔗", "URL:", f"http://{service_config.http.host}:{service_config.http.port}") - elif service_config.backend == "mcp": - info_table.add_row("📚", "Transport:", service_config.mcp.transport) - if service_config.mcp.transport == "sse": - info_table.add_row( - "🔗", - "URL:", - f"http://{service_config.mcp.host}:{service_config.mcp.port}/sse", - ) - - info_table.add_row("", "", "") - import flowllm - - info_table.add_row("🚀", "FlowLLM version:", Text(flowllm.__version__, style="dim white", no_wrap=True)) - - if service_config.backend == "http": - import fastapi - - info_table.add_row("📚", "FastAPI version:", Text(fastapi.__version__, style="dim white", no_wrap=True)) - elif service_config.backend == "mcp": - import fastmcp - - info_table.add_row("📚", "FastMCP version:", Text(fastmcp.__version__, style="dim white", no_wrap=True)) - panel_content = Group(logo_text, "", info_table) - - panel = Panel( - panel_content, - title=app_name, - title_align="left", - border_style="dim", - padding=(1, 4), - expand=False, - ) - - console = Console(stderr=False) - console.print(Group("\n", panel, "\n")) diff --git a/flowllm/core/utils/pydantic_config_parser.py b/flowllm/core/utils/pydantic_config_parser.py deleted file mode 100644 index c418d98..0000000 --- a/flowllm/core/utils/pydantic_config_parser.py +++ /dev/null @@ -1,397 +0,0 @@ -"""Pydantic configuration parser for loading and merging configurations from multiple sources. - -This module provides a generic parser that can load configurations from YAML files, -command-line arguments, and programmatic updates, with support for: -- Dot notation configuration keys (e.g., 'a.b.c=value') -- Deep merging of multiple configuration sources -- Automatic type conversion from strings -""" - -import copy -import json -from pathlib import Path -from typing import Any, Generic, List, Type, TypeVar - -import yaml -from loguru import logger -from pydantic import BaseModel - -T = TypeVar("T", bound=BaseModel) - - -class PydanticConfigParser(Generic[T]): - """Generic parser for Pydantic-based configuration management. - - This parser supports loading configurations from multiple sources with a priority order: - 1. Default values from Pydantic model - 2. YAML configuration files - 3. Command-line arguments (dot notation format) - 4. Programmatic updates via update_config() - - The parser automatically handles: - - Type conversion (string to int, float, bool, None, JSON) - - Deep merging of nested dictionaries - - Dot notation for nested configuration keys - - Attributes: - current_file: Path to the current file, used for relative config file resolution - default_config_name: Default configuration file name (without .yaml extension) - config_class: The Pydantic model class used for validation - config_dict: The current merged configuration dictionary - - Example: - ```python - from pydantic import BaseModel - - class AppConfig(BaseModel): - host: str = "localhost" - port: int = 8080 - - parser = PydanticConfigParser(AppConfig) - config = parser.parse_args("config=app.yaml", "server.port=9000") - print(config.port) # 9000 - ``` - """ - - current_file: str = __file__ - default_config_name: str = "default" - - def __init__(self, config_class: Type[T]): - """Initialize configuration parser. - - Args: - config_class: Pydantic configuration model class that defines the schema - and default values for the configuration. - - Raises: - TypeError: If config_class is not a subclass of BaseModel. - """ - self.config_class = config_class - self.config_dict: dict = {} - - def parse_dot_notation(self, dot_list: List[str]) -> dict: - """Parse dot notation format configuration list into nested dictionary. - - Converts a list of dot-notation key-value pairs into a nested dictionary structure. - Automatically performs type conversion on values. - - Args: - dot_list: Configuration list in format ['a.b.c=value', 'x.y=123', 'z=true']. - Items without '=' are ignored. - - Returns: - Parsed nested dictionary. For example, ['a.b=1', 'a.c=2'] becomes - {'a': {'b': 1, 'c': 2}}. - - Example: - ```python - parser = PydanticConfigParser(SomeConfig) - result = parser.parse_dot_notation(['server.host=localhost', 'server.port=8080']) - print(result) - # {'server': {'host': 'localhost', 'port': 8080}} - ``` - """ - config_dict = {} - - for item in dot_list: - if "=" not in item: - continue - - key_path, value_str = item.split("=", 1) - keys = key_path.split(".") - - # Automatic type conversion - value = self._convert_value(value_str) - - # Build nested dictionary - current_dict = config_dict - for key in keys[:-1]: - if key not in current_dict: - current_dict[key] = {} - current_dict = current_dict[key] - - current_dict[keys[-1]] = value - - return config_dict - - @staticmethod - def _convert_value(value_str: str) -> Any: - """Automatically convert string values to appropriate Python types. - - Attempts to convert string values to their appropriate Python types: - - Boolean: 'true'/'false' -> bool - - None: 'none'/'null' -> None - - Integer: Numeric strings without decimal point -> int - - Float: Numeric strings with decimal point or scientific notation -> float - - JSON: Valid JSON strings -> parsed JSON object - - String: Everything else -> str - - Args: - value_str: String value to convert. - - Returns: - Converted value with appropriate Python type. If conversion fails, - returns the original string. - - Example: - ```python - PydanticConfigParser._convert_value("123") # 123 - PydanticConfigParser._convert_value("true") # True - PydanticConfigParser._convert_value('{"key": "value"}') # {'key': 'value'} - ``` - """ - value_str = value_str.strip() - - if value_str.lower() in ("true", "false"): - return value_str.lower() == "true" - - if value_str.lower() in ("none", "null"): - return None - - try: - if "." not in value_str and "e" not in value_str.lower(): - return int(value_str) - - return float(value_str) - - except ValueError: - pass - - try: - return json.loads(value_str) - except (json.JSONDecodeError, ValueError): - pass - - return value_str - - @staticmethod - def load_from_yaml(yaml_path: str | Path) -> dict: - """Load configuration from YAML file. - - Reads and parses a YAML configuration file into a Python dictionary. - - Args: - yaml_path: Path to the YAML configuration file. Can be a string or Path object. - - Returns: - Configuration dictionary parsed from the YAML file. - - Raises: - FileNotFoundError: If the YAML file does not exist at the specified path. - yaml.YAMLError: If the YAML file contains invalid YAML syntax. - - Example: - ```python - config = PydanticConfigParser.load_from_yaml("config.yaml") - print(config) - # {'host': 'localhost', 'port': 8080} - ``` - """ - if isinstance(yaml_path, str): - yaml_path = Path(yaml_path) - - if not yaml_path.exists(): - raise FileNotFoundError(f"Configuration file does not exist: {yaml_path}") - - with yaml_path.open(encoding="utf-8") as f: - return yaml.safe_load(f) - - def merge_configs(self, *config_dicts: dict) -> dict: - """Deep merge multiple configuration dictionaries. - - Merges multiple configuration dictionaries in order, with later dictionaries - overriding values from earlier ones. Performs deep merging for nested dictionaries, - meaning nested values are merged rather than replaced entirely. - - Args: - *config_dicts: Multiple configuration dictionaries to merge. Merging order - is from first to last, with later configs taking precedence. - - Returns: - Deeply merged configuration dictionary. All provided dictionaries are merged - into a single result. - - Example: - ```python - parser = PydanticConfigParser(SomeConfig) - base = {'a': {'x': 1, 'y': 2}} - update = {'a': {'y': 3, 'z': 4}, 'b': 5} - result = parser.merge_configs(base, update) - print(result) - # {'a': {'x': 1, 'y': 3, 'z': 4}, 'b': 5} - ``` - """ - result = {} - - for config_dict in config_dicts: - result = self._deep_merge(result, config_dict) - - return result - - def _deep_merge(self, base_dict: dict, update_dict: dict) -> dict: - """Deep merge two dictionaries recursively. - - Recursively merges update_dict into base_dict. When both dictionaries contain - the same key and both values are dictionaries, they are merged recursively. - Otherwise, the value from update_dict overwrites the value in base_dict. - - Args: - base_dict: Base dictionary to merge into. This dictionary is not modified. - update_dict: Dictionary containing updates to merge into base_dict. - - Returns: - New merged dictionary containing the combined result of both dictionaries. - The base_dict and update_dict are not modified. - - Example: - ```python - parser = PydanticConfigParser(SomeConfig) - base = {'a': 1, 'nested': {'x': 1}} - update = {'b': 2, 'nested': {'y': 2}} - result = parser._deep_merge(base, update) - print(result) - # {'a': 1, 'b': 2, 'nested': {'x': 1, 'y': 2}} - ``` - """ - result = base_dict.copy() - - for key, value in update_dict.items(): - if key in result and isinstance(result[key], dict) and isinstance(value, dict): - result[key] = self._deep_merge(result[key], value) - else: - result[key] = value - - return result - - def parse_args(self, *args) -> T: - """Parse command line arguments and return validated configuration object. - - Parses configuration from multiple sources in priority order: - 1. Default values from Pydantic model - 2. YAML configuration file(s) (specified via config= or default_config_name) - - Supports multiple configs separated by commas: config=default,search - - Later configs override earlier ones - 3. Command-line overrides (dot notation format) - - Command-line arguments should be in format: - - 'config=' or 'c=' to specify YAML config file - - 'config=,' to specify multiple config files (file2 overrides file1) - - 'key=value' or 'nested.key=value' for configuration overrides - - Args: - *args: Command line arguments. Can include: - - 'config=' or 'c=' to specify YAML config file - - Dot notation overrides like 'server.host=localhost' - - Returns: - Validated configuration object of type T (the Pydantic model class). - The configuration is validated against the Pydantic schema. - - Raises: - FileNotFoundError: If the specified config file does not exist. - AssertionError: If no config file is specified and default_config_name is empty. - ValidationError: If the merged configuration does not match the Pydantic schema. - - Example: - ```python - parser = PydanticConfigParser(AppConfig) - config = parser.parse_args("config=app.yaml", "server.port=9000") - print(config.port) # 9000 (overridden from YAML default) - - # Multiple config files: search.yaml overrides default.yaml - config = parser.parse_args("config=default,search") - ``` - """ - configs_to_merge = [] - - # 1. Default configuration (from Pydantic model) - default_config = self.config_class().model_dump() - configs_to_merge.append(default_config) - - # 2. YAML configuration file(s) - config = "" - filter_args = [] - for arg in args: - if "=" not in arg: - continue - - arg = arg.lstrip("--").lstrip("-") - - if "c=" in arg or "config=" in arg: - config = arg.split("=")[-1] - else: - filter_args.append(arg) - - if not config: - if self.default_config_name: - config = self.default_config_name - assert config, "add `config=` in cmd!" - - # Support multiple config files separated by commas - # Example: config=default,search (search will override default) - config_list = [c.strip() for c in config.split(",") if c.strip()] - - # Load all config files in order (later configs override earlier ones) - for single_config in config_list: - if not single_config.endswith(".yaml"): - single_config += ".yaml" - - # Resolve config file path - config_path = Path(self.current_file).parent / single_config - if config_path.exists(): - logger.info(f"load config={config_path}") - else: - logger.warning(f"config={config_path} not found, try {single_config}") - config_path = Path(single_config) - assert config_path.exists(), f"config={config_path} not found" - - yaml_config = self.load_from_yaml(config_path) - configs_to_merge.append(yaml_config) - - # 3. Command line override configuration - if args: - cli_config = self.parse_dot_notation(filter_args) - configs_to_merge.append(cli_config) - - # Merge all configurations - self.config_dict = self.merge_configs(*configs_to_merge) - - # Create and validate final configuration object - return self.config_class.model_validate(self.config_dict) - - def update_config(self, **kwargs) -> T: - """Update configuration object using keyword arguments. - - Updates the current configuration with new values provided as keyword arguments. - Supports dot notation via double underscores (e.g., server__host becomes server.host). - - Args: - **kwargs: Configuration items to update. Keys can use double underscores - (__) to represent nested paths (e.g., server__host='localhost' becomes - server.host='localhost'). Values are automatically type-converted. - - Returns: - Updated and validated configuration object of type T. The original - config_dict is not modified. - - Example: - ```python - parser = PydanticConfigParser(AppConfig) - config = parser.parse_args("config=app.yaml") - updated = parser.update_config(server__port=9000, debug=True) - print(updated.port) # 9000 - print(updated.debug) # True - ``` - """ - # Convert kwargs to dot notation format - dot_list = [] - for key, value in kwargs.items(): - # support double underscore as dot replacement (server__host -> server.host) - dot_key = key.replace("__", ".") - dot_list.append(f"{dot_key}={value}") - - # Parse and merge configuration - override_config = self.parse_dot_notation(dot_list) - final_config = self.merge_configs(copy.deepcopy(self.config_dict), override_config) - - return self.config_class.model_validate(final_config) diff --git a/flowllm/core/utils/pydantic_utils.py b/flowllm/core/utils/pydantic_utils.py deleted file mode 100644 index 4b01568..0000000 --- a/flowllm/core/utils/pydantic_utils.py +++ /dev/null @@ -1,68 +0,0 @@ -"""Pydantic helpers for dynamically creating request models. - -This module exposes utilities to construct Pydantic models at runtime based on -Flow input schemas. It maps simple schema types to Python types and builds a -`FlowRequest`-derived model with properly typed, optional/required fields and -helpful descriptions. - -Exports: -- `create_pydantic_model`: Build a `pydantic.BaseModel` subclass for a flow's - input schema, using `FlowRequest` as a base. -""" - -from typing import Dict, Optional - -from pydantic import create_model, Field - -from . import snake_to_camel -from ..schema import ToolAttr, FlowRequest - -TYPE_MAPPING = { - "string": str, - "array": list, - "integer": int, - "number": float, - "boolean": bool, - "object": dict, -} - - -def create_pydantic_model(flow_name: str, input_schema: Dict[str, ToolAttr] = None): - """Create a Pydantic model class for a flow's input schema. - - The returned model subclasses `FlowRequest` and includes one field per - parameter in ``input_schema``. Parameter types are inferred from a fixed - mapping of schema type strings to Python types, and optionality is respected - according to each parameter's ``required`` flag. Field descriptions are - preserved via ``pydantic.Field``. - - Args: - flow_name: The name of the flow. Used to derive the generated model - class name in CamelCase. - input_schema: A mapping from parameter name to its attributes. - If ``None`` or empty, only the fields from the `FlowRequest` base are present. - - Returns: - A dynamically created `pydantic.BaseModel` subclass named - ``{SnakeToCamel(flow_name)}Model`` that inherits from `FlowRequest` and - contains fields described by ``input_schema``. - - Raises: - AssertionError: If a parameter's type is not one of the supported keys - in ``TYPE_MAPPING``. - """ - fields = {} - - if input_schema: - for param_name, param_config in input_schema.items(): - assert ( - param_config.type in TYPE_MAPPING - ), f"flow_name={flow_name} had invalid type: {param_config.type}! supported={sorted(TYPE_MAPPING.keys())}" - field_type = TYPE_MAPPING[param_config.type] - - if not param_config.required: - fields[param_name] = (Optional[field_type], Field(default=None, description=param_config.description)) - else: - fields[param_name] = (field_type, Field(default=..., description=param_config.description)) - - return create_model(f"{snake_to_camel(flow_name)}Model", __base__=FlowRequest, **fields) diff --git a/flowllm/core/utils/timer.py b/flowllm/core/utils/timer.py deleted file mode 100644 index 5d94c32..0000000 --- a/flowllm/core/utils/timer.py +++ /dev/null @@ -1,153 +0,0 @@ -"""Timer utilities for measuring execution time. - -This module provides a Timer class and a timer decorator for measuring execution time -of code blocks and functions. Supports both synchronous and asynchronous operations. -""" - -import functools -import inspect -import time -from typing import Optional, TypeVar, Callable, Union, Awaitable - -from loguru import logger - -T = TypeVar("T") - - -class Timer: - """A context manager and timer utility for measuring execution time. - - Supports both synchronous and asynchronous context managers. - - Usage as synchronous context manager: - with Timer("my_operation"): - # your code here - - Usage as asynchronous context manager: - async with Timer("my_operation"): - # your code here - await some_async_function() - - Usage with decorator: - @timer("my_function") - def my_function(): - # your code here - - @timer("my_async_function") - async def my_async_function(): - # your async code here - await some_async_function() - """ - - def __init__(self, name: str, use_ms: bool = False, stack_level: int = 2): - """Initialize a Timer instance. - - Args: - name: Name of the timer, used in log messages - use_ms: If True, display time in milliseconds, otherwise in seconds - stack_level: Stack level for loguru logger (default 2) - """ - self.name: str = name - self.use_ms: bool = use_ms - self.stack_level: int = stack_level - - self.time_start: float = 0.0 - self.time_cost: float = 0.0 - - def __enter__(self) -> "Timer": - """Start the timer.""" - self.time_start = time.time() - logger.info(f"========== timer.{self.name} start ==========", stacklevel=self.stack_level) - return self - - def __exit__(self, exc_type: Optional[type], exc_val: Optional[Exception], exc_tb: Optional[object]) -> None: - """Stop the timer and log the elapsed time. - - Args: - exc_type: Exception type if an exception occurred, None otherwise - exc_val: Exception value if an exception occurred, None otherwise - exc_tb: Exception traceback if an exception occurred, None otherwise - """ - time_end = time.time() - self.time_cost = time_end - self.time_start - if self.use_ms: - time_str = f"{self.time_cost * 1000:.2f}ms" - else: - time_str = f"{self.time_cost:.3f}s" - - logger.info(f"========== timer.{self.name} end, time_cost={time_str} ==========", stacklevel=self.stack_level) - - async def __aenter__(self) -> "Timer": - """Start the timer (async version).""" - self.time_start = time.time() - logger.info(f"========== timer.{self.name} start ==========", stacklevel=self.stack_level) - return self - - async def __aexit__(self, exc_type: Optional[type], exc_val: Optional[Exception], exc_tb: Optional[object]) -> None: - """Stop the timer and log the elapsed time (async version). - - Args: - exc_type: Exception type if an exception occurred, None otherwise - exc_val: Exception value if an exception occurred, None otherwise - exc_tb: Exception traceback if an exception occurred, None otherwise - """ - time_end = time.time() - self.time_cost = time_end - self.time_start - if self.use_ms: - time_str = f"{self.time_cost * 1000:.2f}ms" - else: - time_str = f"{self.time_cost:.3f}s" - - logger.info(f"========== timer.{self.name} end, time_cost={time_str} ==========", stacklevel=self.stack_level) - - -def timer(name: Optional[str] = None, use_ms: bool = False, stack_level: int = 2) -> Callable[ - [Callable[..., Union[T, Awaitable[T]]]], - Callable[..., Union[T, Awaitable[T]]], -]: - """A decorator factory for timing function execution. - - Automatically detects if the decorated function is async and handles it accordingly. - Supports both synchronous and asynchronous functions. - - Args: - name: Optional name for the timer. If None, uses the function's name - use_ms: If True, display time in milliseconds, otherwise in seconds - stack_level: Stack level for loguru logger (default 2, +1 is added for the decorator wrapper) - - Returns: - A decorator function that can be applied to any function (sync or async) - - Usage: - @timer("my_function") - def my_function(): - pass - - @timer("my_async_function") - async def my_async_function(): - await some_async_function() - - @timer(use_ms=True) - def another_function(): - pass - """ - - def decorator(func: Callable[..., Union[T, Awaitable[T]]]) -> Callable[..., Union[T, Awaitable[T]]]: - if inspect.iscoroutinefunction(func): - # Async function - @functools.wraps(func) - async def async_wrapper(*args, **kwargs) -> T: - async with Timer(name=name or func.__name__, use_ms=use_ms, stack_level=stack_level + 1): - return await func(*args, **kwargs) - - return async_wrapper - else: - # Sync function - @functools.wraps(func) - def sync_wrapper(*args, **kwargs) -> T: - with Timer(name=name or func.__name__, use_ms=use_ms, stack_level=stack_level + 1): - return func(*args, **kwargs) - - return sync_wrapper - - return decorator diff --git a/flowllm/core/vector_store/__init__.py b/flowllm/core/vector_store/__init__.py deleted file mode 100644 index d2de675..0000000 --- a/flowllm/core/vector_store/__init__.py +++ /dev/null @@ -1,34 +0,0 @@ -"""Vector store module for managing vector embeddings. - -This module provides implementations of vector stores for storing, searching, -and managing vector embeddings along with their associated metadata. It includes: - -- BaseVectorStore: Abstract base class defining the vector store interface -- LocalVectorStore: File-based local vector store implementation -- ChromaVectorStore: ChromaDB-based vector store -- EsVectorStore: Elasticsearch-based vector store -- QdrantVectorStore: Qdrant-based vector store -- MemoryVectorStore: In-memory vector store for fast access -- PgVectorStore: PostgreSQL pgvector-based vector store - -All vector stores support workspace-based organization and provide both synchronous -and asynchronous interfaces for operations such as insertion, search, and deletion. -""" - -from .base_vector_store import BaseVectorStore -from .chroma_vector_store import ChromaVectorStore -from .es_vector_store import EsVectorStore -from .local_vector_store import LocalVectorStore -from .memory_vector_store import MemoryVectorStore -from .pgvector_vector_store import PgVectorStore -from .qdrant_vector_store import QdrantVectorStore - -__all__ = [ - "BaseVectorStore", - "ChromaVectorStore", - "EsVectorStore", - "LocalVectorStore", - "MemoryVectorStore", - "PgVectorStore", - "QdrantVectorStore", -] diff --git a/flowllm/core/vector_store/base_vector_store.py b/flowllm/core/vector_store/base_vector_store.py deleted file mode 100644 index 14d2af4..0000000 --- a/flowllm/core/vector_store/base_vector_store.py +++ /dev/null @@ -1,297 +0,0 @@ -"""Base class for vector store implementations. - -This module provides the abstract base class for vector stores, which are used -to store, search, and manage vector embeddings along with their associated metadata. -Vector stores support workspace-based organization and provide both synchronous -and asynchronous interfaces for all operations. -""" - -import asyncio -from abc import ABC, abstractmethod -from functools import partial -from pathlib import Path -from typing import List, Dict, Any, Optional -from types import TracebackType - -from ..context.service_context import C -from ..embedding_model import BaseEmbeddingModel -from ..schema import VectorNode - - -class BaseVectorStore(ABC): - """Abstract base class for vector store implementations. - - This class defines the interface that all vector store implementations must - follow. It provides methods for managing workspaces, inserting and searching - vector nodes, and exporting/importing data. Both synchronous and asynchronous - versions of all operations are available. - - Attributes: - embedding_model: Optional embedding model used for generating embeddings - from text queries. If None, nodes must be inserted with pre-computed - embeddings. - - Subclasses must implement all abstract methods to provide the actual vector - storage functionality. - """ - - def __init__(self, embedding_model: BaseEmbeddingModel | None = None, **kwargs): - # Initialize the vector store with an optional embedding model - self.embedding_model: BaseEmbeddingModel | None = embedding_model - self.kwargs: dict = kwargs - - @staticmethod - async def _run_sync_in_executor(sync_func, *args, **kwargs): - """Run a synchronous function in a thread pool executor. - - This utility method is useful for wrapping synchronous I/O operations - (like file reads/writes) in async methods to avoid blocking the event loop. - - Args: - sync_func: The synchronous function to execute. - *args: Positional arguments to pass to the function. - **kwargs: Keyword arguments to pass to the function. - - Returns: - The return value of the synchronous function. - """ - loop = asyncio.get_running_loop() - return await loop.run_in_executor(C.thread_pool, partial(sync_func, *args, **kwargs)) - - def get_node_embeddings(self, nodes: VectorNode | List[VectorNode]) -> List[VectorNode]: - """Generate embeddings for nodes using the embedding model. - - Args: - nodes: Single VectorNode or list of VectorNodes to generate embeddings for. - - Returns: - List[VectorNode]: List of VectorNodes with generated embeddings. - - Raises: - ValueError: If embedding_model is None. - """ - if self.embedding_model is None: - raise ValueError("embedding_model is None. Cannot generate embeddings without an embedding model.") - return self.embedding_model.get_node_embeddings(nodes) - - async def async_get_node_embeddings(self, nodes: VectorNode | List[VectorNode]) -> List[VectorNode]: - """Asynchronously generate embeddings for nodes using the embedding model. - - Args: - nodes: Single VectorNode or list of VectorNodes to generate embeddings for. - - Returns: - List[VectorNode]: List of VectorNodes with generated embeddings. - - Raises: - ValueError: If embedding_model is None. - """ - if self.embedding_model is None: - raise ValueError("embedding_model is None. Cannot generate embeddings without an embedding model.") - return await self.embedding_model.async_get_node_embeddings(nodes) - - def get_embeddings(self, query: str | List[str]): - """Generate embeddings for text queries using the embedding model. - - Args: - query: Single string or list of strings to generate embeddings for. - - Returns: - Embeddings for the input query/queries. - - Raises: - ValueError: If embedding_model is None. - """ - if self.embedding_model is None: - raise ValueError("embedding_model is None. Cannot generate embeddings without an embedding model.") - return self.embedding_model.get_embeddings(query) - - async def async_get_embeddings(self, query: str | List[str]): - """Asynchronously generate embeddings for text queries using the embedding model. - - Args: - query: Single string or list of strings to generate embeddings for. - - Returns: - Embeddings for the input query/queries. - - Raises: - ValueError: If embedding_model is None. - """ - if self.embedding_model is None: - raise ValueError("embedding_model is None. Cannot generate embeddings without an embedding model.") - return await self.embedding_model.async_get_embeddings(query) - - @abstractmethod - def exist_workspace(self, workspace_id: str, **kwargs) -> bool: - """Check if a workspace exists in the vector store.""" - - @abstractmethod - async def async_exist_workspace(self, workspace_id: str, **kwargs) -> bool: - """Asynchronously check if a workspace exists in the vector store.""" - - @abstractmethod - def delete_workspace(self, workspace_id: str, **kwargs): - """Delete a workspace from the vector store.""" - - @abstractmethod - async def async_delete_workspace(self, workspace_id: str, **kwargs): - """Asynchronously delete a workspace from the vector store.""" - - @abstractmethod - def create_workspace(self, workspace_id: str, **kwargs): - """Create a new workspace in the vector store.""" - - @abstractmethod - async def async_create_workspace(self, workspace_id: str, **kwargs): - """Asynchronously create a new workspace in the vector store.""" - - @abstractmethod - def list_workspace_nodes(self, workspace_id: str, **kwargs) -> List[VectorNode]: - """List all nodes in a workspace.""" - - @abstractmethod - async def async_list_workspace_nodes(self, workspace_id: str, **kwargs) -> List[VectorNode]: - """Asynchronously list all nodes in a workspace.""" - - @abstractmethod - def dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs): - """Dump workspace data to a file or path.""" - - @abstractmethod - async def async_dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs): - """Asynchronously dump workspace data to a file or path.""" - - @abstractmethod - def load_workspace( - self, - workspace_id: str, - path: str | Path = "", - nodes: List[VectorNode] | None = None, - callback_fn=None, - **kwargs, - ): - """Load workspace data from a file or path, or from provided nodes.""" - - @abstractmethod - async def async_load_workspace( - self, - workspace_id: str, - path: str | Path = "", - nodes: List[VectorNode] | None = None, - callback_fn=None, - **kwargs, - ): - """Asynchronously load workspace data from a file or path, or from provided nodes.""" - - @abstractmethod - def copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs): - """Copy one workspace to another.""" - - @abstractmethod - async def async_copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs): - """Asynchronously copy one workspace to another.""" - - @abstractmethod - def list_workspace(self, **kwargs) -> List[str]: - """List all existing workspaces.""" - - @abstractmethod - async def async_list_workspace(self, **kwargs) -> List[str]: - """Asynchronously list all existing workspaces.""" - - @abstractmethod - def search( - self, - query: str, - workspace_id: str, - top_k: int = 1, - filter_dict: Optional[Dict[str, Any]] = None, - **kwargs, - ) -> List[VectorNode]: - """Search for similar vectors in the workspace. - - When query is empty (empty string or None), the search degrades to a - filter-only search without vector similarity ranking. In this case, - results are returned based solely on filter_dict criteria, up to top_k - items, without similarity scores. - - Args: - query: Text query to search for. If empty/None, performs filter-only - search without vector similarity. - workspace_id: Identifier of the workspace to search in. - top_k: Number of top results to return. Defaults to 1. - filter_dict: Optional dictionary of filters to apply to nodes. - **kwargs: Additional keyword arguments for implementation-specific options. - - Returns: - List[VectorNode]: List of matching nodes. When query is provided, - nodes are sorted by similarity score (highest first) with a - "score" key in metadata. When query is empty, nodes are returned - in storage order without scores. - """ - - @abstractmethod - async def async_search( - self, - query: str, - workspace_id: str, - top_k: int = 1, - filter_dict: Optional[Dict[str, Any]] = None, - **kwargs, - ) -> List[VectorNode]: - """Search for similar vectors in the workspace. - - Async version of search(). When query is empty, degrades to filter-only - search without vector similarity ranking. - - See search() for full documentation. - """ - - @abstractmethod - def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs): - """Insert nodes into the workspace.""" - - @abstractmethod - async def async_insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs): - """Asynchronously insert nodes into the workspace.""" - - @abstractmethod - def delete(self, node_ids: str | List[str], workspace_id: str, **kwargs): - """Delete nodes from the workspace by their IDs.""" - - @abstractmethod - async def async_delete(self, node_ids: str | List[str], workspace_id: str, **kwargs): - """Asynchronously delete nodes from the workspace by their IDs.""" - - def close(self): - """Close the vector store and clean up resources.""" - - async def async_close(self): - """Asynchronously close the vector store and clean up resources.""" - - def __enter__(self) -> "BaseVectorStore": - """Allow usage as a synchronous context manager.""" - return self - - async def __aenter__(self) -> "BaseVectorStore": - """Allow usage as an asynchronous context manager.""" - return self - - def __exit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - traceback: TracebackType | None, - ) -> bool | None: - """Ensure resources close when exiting a sync context.""" - self.close() - - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - traceback: TracebackType | None, - ) -> bool | None: - """Ensure resources close when exiting an async context.""" - await self.async_close() diff --git a/flowllm/core/vector_store/chroma_vector_store.py b/flowllm/core/vector_store/chroma_vector_store.py deleted file mode 100644 index a5696db..0000000 --- a/flowllm/core/vector_store/chroma_vector_store.py +++ /dev/null @@ -1,818 +0,0 @@ -""" -ChromaDB vector store implementation. - -This module provides a ChromaDB-based vector store that stores vector nodes -in ChromaDB collections. It supports workspace management, vector similarity search, -metadata filtering, and provides both synchronous and asynchronous operations. -""" - -import os -from pathlib import Path -from typing import List, Dict, Any, Optional - -from loguru import logger - -from .memory_vector_store import MemoryVectorStore -from ..context import C -from ..schema import VectorNode - -# Disable ChromaDB telemetry to avoid PostHog warnings -os.environ.setdefault("ANONYMIZED_TELEMETRY", "false") - - -@C.register_vector_store("chroma") -class ChromaVectorStore(MemoryVectorStore): - """ChromaDB vector store implementation. - - This class provides a ChromaDB-based vector store that uses ChromaDB collections - for workspace management. Each workspace corresponds to a ChromaDB collection, - and vector nodes are stored with their embeddings, documents, and metadata. - - Attributes: - store_dir: Directory path where ChromaDB data is persisted (default: "./chroma_vector_store"). - collections: Dictionary mapping workspace_id to ChromaDB Collection objects. - _client: Private ChromaDB client instance. - - The store supports both synchronous and asynchronous operations, with async methods - using thread pools to execute ChromaDB operations without blocking the event loop. - """ - - # ==================== Initialization ==================== - - def __init__(self, store_dir: str = "./chroma_vector_store", **kwargs): - """Initialize the ChromaDB client with telemetry disabled. - - Args: - store_dir: Directory path where ChromaDB data is persisted. - **kwargs: Additional keyword arguments passed to MemoryVectorStore. - """ - super().__init__(**kwargs) - self.store_dir = store_dir - self.collections: dict = {} - # Disable telemetry to avoid PostHog warnings - # Use PersistentClient explicitly to avoid singleton conflicts - from chromadb import PersistentClient, ClientAPI - from chromadb.config import Settings - - self._client: ClientAPI = PersistentClient( - path=self.store_dir, - settings=Settings(anonymized_telemetry=False), - ) - logger.info(f"ChromaDB client initialized with store_dir={self.store_dir}") - - @property - def store_path(self) -> Path: - """ - Get the storage directory path. - - Returns: - Path object representing the storage directory. - """ - return Path(self.store_dir) - - # ==================== Static Helper Methods ==================== - - @staticmethod - def _build_chroma_filters( - filter_dict: Optional[Dict[str, Any]] = None, - ) -> tuple[Optional[Dict], Optional[List[str]]]: - """Build ChromaDB where clause from filter_dict. - - Converts a filter dictionary into ChromaDB's where clause format. - Supports both term filters (exact match) and range filters (gte, lte, gt, lt). - Handles unique_id specially since it's stored as ChromaDB's document ID. - - Args: - filter_dict: Dictionary of filters to apply. Can contain: - - Term filters: {"key": "value"} -> exact match - - Range filters: {"key": {"gte": 1, "lte": 10}} -> range query - - unique_id: Filters by document ID (stored separately from metadata) - - Keys can use "metadata." prefix (will be stripped for ChromaDB) - - Returns: - Tuple of (where_clause, ids_filter): - - where_clause: ChromaDB where clause dictionary for metadata filters, or None - - ids_filter: List of document IDs to filter by, or None - """ - if not filter_dict: - return None, None - - where_conditions = {} - ids_filter = None - - for key, filter_value in filter_dict.items(): - # Handle unique_id specially - it's stored as ChromaDB's document ID - if key == "unique_id": - if isinstance(filter_value, list): - ids_filter = filter_value - else: - ids_filter = [filter_value] - continue - - # Strip "metadata." prefix if present (ChromaDB stores metadata fields directly) - if key.startswith("metadata."): - chroma_key = key[len("metadata.") :] - else: - chroma_key = key - - if isinstance(filter_value, dict): - # Range filter: {"gte": 1, "lte": 10} - range_conditions = {} - if "gte" in filter_value: - range_conditions["$gte"] = filter_value["gte"] - if "lte" in filter_value: - range_conditions["$lte"] = filter_value["lte"] - if "gt" in filter_value: - range_conditions["$gt"] = filter_value["gt"] - if "lt" in filter_value: - range_conditions["$lt"] = filter_value["lt"] - if range_conditions: - where_conditions[chroma_key] = range_conditions - elif isinstance(filter_value, list): - # List filter: use $in operator for OR logic - where_conditions[chroma_key] = {"$in": filter_value} - else: - # Term filter: direct value comparison - where_conditions[chroma_key] = filter_value - - return (where_conditions if where_conditions else None, ids_filter) - - # ==================== Private Helper Methods ==================== - - def _get_collection(self, workspace_id: str): - """Get or create a ChromaDB collection for the given workspace. - - Args: - workspace_id: The workspace identifier. - - Returns: - ChromaDB Collection object for the workspace. - """ - if workspace_id not in self.collections: - self.collections[workspace_id] = self._client.get_or_create_collection(workspace_id) - return self.collections[workspace_id] - - # ==================== Workspace Management Methods ==================== - - def exist_workspace(self, workspace_id: str, **kwargs) -> bool: - """Check if a workspace exists in the vector store. - - Args: - workspace_id: The workspace identifier. - **kwargs: Additional keyword arguments (unused). - - Returns: - True if the workspace exists, False otherwise. - """ - # Check cache first for better performance - if workspace_id in self.collections: - return True - # Fall back to checking the client - try: - self._client.get_collection(workspace_id) - return True - except Exception: - return False - - async def async_exist_workspace(self, workspace_id: str, **kwargs) -> bool: - """Async version of exist_workspace. - - Args: - workspace_id: The workspace identifier. - **kwargs: Additional keyword arguments (unused). - - Returns: - True if the workspace exists, False otherwise. - """ - return await self._run_sync_in_executor(self.exist_workspace, workspace_id, **kwargs) - - def delete_workspace(self, workspace_id: str, **kwargs): - """Delete a workspace from the vector store. - - Args: - workspace_id: The workspace identifier to delete. - **kwargs: Additional keyword arguments (unused). - """ - if self.exist_workspace(workspace_id): - self._client.delete_collection(workspace_id) - if workspace_id in self.collections: - del self.collections[workspace_id] - logger.info(f"Deleted workspace_id={workspace_id} from ChromaDB") - - async def async_delete_workspace(self, workspace_id: str, **kwargs): - """Async version of delete_workspace. - - Args: - workspace_id: The workspace identifier to delete. - **kwargs: Additional keyword arguments (unused). - """ - return await self._run_sync_in_executor(self.delete_workspace, workspace_id, **kwargs) - - def create_workspace(self, workspace_id: str, **kwargs): - """Create a new workspace in the vector store. - - Args: - workspace_id: The workspace identifier to create. - **kwargs: Additional keyword arguments (unused). - """ - self.collections[workspace_id] = self._client.get_or_create_collection(workspace_id) - logger.info(f"Created workspace_id={workspace_id} in ChromaDB") - - async def async_create_workspace(self, workspace_id: str, **kwargs): - """Async version of create_workspace. - - Args: - workspace_id: The workspace identifier to create. - **kwargs: Additional keyword arguments (unused). - """ - return await self._run_sync_in_executor(self.create_workspace, workspace_id, **kwargs) - - def list_workspace_nodes(self, workspace_id: str, **kwargs) -> List[VectorNode]: - """List all nodes in a workspace. - - Args: - workspace_id: The workspace identifier. - **kwargs: Additional keyword arguments (unused). - - Returns: - List of VectorNode objects from the workspace. - """ - if not self.exist_workspace(workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return [] - - collection = self._get_collection(workspace_id) - results = collection.get(include=["documents", "metadatas", "embeddings"]) - - nodes = [] - documents = results.get("documents") - metadatas = results.get("metadatas") - embeddings = results.get("embeddings") - for i in range(len(results["ids"])): - node = VectorNode( - workspace_id=workspace_id, - unique_id=results["ids"][i], - content=documents[i] if documents is not None else "", - metadata=metadatas[i] if metadatas is not None else {}, - vector=embeddings[i] if embeddings is not None else None, - ) - nodes.append(node) - return nodes - - async def async_list_workspace_nodes(self, workspace_id: str, **kwargs) -> List[VectorNode]: - """Async version of list_workspace_nodes. - - Args: - workspace_id: The workspace identifier. - **kwargs: Additional keyword arguments (unused). - - Returns: - List of VectorNode objects from the workspace. - """ - return await self._run_sync_in_executor(self.list_workspace_nodes, workspace_id, **kwargs) - - def dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs): - """Export a workspace from ChromaDB to disk at the specified path. - - Args: - workspace_id: Identifier of the workspace to export. - path: Directory path where to write the exported workspace file. - callback_fn: Optional callback function to transform nodes during export. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - - Returns: - dict: Dictionary with "size" key indicating number of nodes exported, - or empty dict if workspace doesn't exist or path is empty. - """ - if not self.exist_workspace(workspace_id): - logger.warning(f"workspace_id={workspace_id} not found in ChromaDB!") - return {} - - if not path: - logger.warning("path is empty, cannot dump workspace!") - return {} - - nodes = self.list_workspace_nodes(workspace_id, **kwargs) - - return self._dump_to_path( - nodes=nodes, - workspace_id=workspace_id, - path=path, - callback_fn=callback_fn, - **kwargs, - ) - - async def async_dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs): - """Async version of dump_workspace. - - Args: - workspace_id: Identifier of the workspace to export. - path: Directory path where to write the exported workspace file. - callback_fn: Optional callback function to transform nodes during export. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - - Returns: - dict: Dictionary with "size" key indicating number of nodes exported, - or empty dict if workspace doesn't exist or path is empty. - """ - return await self._run_sync_in_executor( - self.dump_workspace, - workspace_id, - path, - callback_fn, - **kwargs, - ) - - def load_workspace( - self, - workspace_id: str, - path: str | Path = "", - nodes: Optional[List[VectorNode]] = None, - callback_fn=None, - **kwargs, - ): - """Load a workspace into ChromaDB from disk, optionally merging with provided nodes. - - This method replaces any existing workspace with the same ID, - then loads nodes from the specified path and/or the provided nodes list. - - Args: - workspace_id: Identifier for the workspace to create/load. - path: Directory path containing the workspace file to load. - If empty, only loads from nodes parameter. - nodes: Optional list of VectorNode instances to merge with loaded nodes. - callback_fn: Optional callback function to transform node dictionaries. - **kwargs: Additional keyword arguments to pass to load operations. - - Returns: - dict: Dictionary with "size" key indicating total number of nodes loaded. - """ - if self.exist_workspace(workspace_id): - self.delete_workspace(workspace_id=workspace_id, **kwargs) - logger.info(f"Cleared existing workspace_id={workspace_id} from ChromaDB") - - self.create_workspace(workspace_id=workspace_id, **kwargs) - - all_nodes: List[VectorNode] = [] - - if nodes: - all_nodes.extend(nodes) - - if path: - all_nodes.extend( - self._load_from_path( - path=path, - workspace_id=workspace_id, - callback_fn=callback_fn, - **kwargs, - ), - ) - - if all_nodes: - self.insert(nodes=all_nodes, workspace_id=workspace_id, **kwargs) - - logger.info(f"Loaded workspace_id={workspace_id} with {len(all_nodes)} nodes into ChromaDB") - return {"size": len(all_nodes)} - - async def async_load_workspace( - self, - workspace_id: str, - path: str | Path = "", - nodes: Optional[List[VectorNode]] = None, - callback_fn=None, - **kwargs, - ): - """Async version of load_workspace. - - Args: - workspace_id: Identifier for the workspace to create/load. - path: Directory path containing the workspace file to load. - If empty, only loads from nodes parameter. - nodes: Optional list of VectorNode instances to merge with loaded nodes. - callback_fn: Optional callback function to transform node dictionaries. - **kwargs: Additional keyword arguments to pass to load operations. - - Returns: - dict: Dictionary with "size" key indicating total number of nodes loaded. - """ - return await self._run_sync_in_executor( - self.load_workspace, - workspace_id, - path, - nodes, - callback_fn, - **kwargs, - ) - - def copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs): - """Copy all nodes from one workspace to another in ChromaDB. - - Args: - src_workspace_id: Identifier of the source workspace. - dest_workspace_id: Identifier of the destination workspace. - Created if it doesn't exist. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - dict: Dictionary with "size" key indicating number of nodes copied, - or empty dict if source workspace doesn't exist. - """ - if not self.exist_workspace(src_workspace_id): - logger.warning(f"src_workspace_id={src_workspace_id} not found in ChromaDB!") - return {} - - if not self.exist_workspace(dest_workspace_id): - self.create_workspace(workspace_id=dest_workspace_id, **kwargs) - - src_nodes = self.list_workspace_nodes(src_workspace_id, **kwargs) - node_size = len(src_nodes) - - new_nodes = [] - for node in src_nodes: - new_node = VectorNode(**node.model_dump()) - new_node.workspace_id = dest_workspace_id - new_nodes.append(new_node) - - if new_nodes: - self.insert(nodes=new_nodes, workspace_id=dest_workspace_id, **kwargs) - - logger.info(f"Copied {node_size} nodes from {src_workspace_id} to {dest_workspace_id}") - return {"size": node_size} - - async def async_copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs): - """Async version of copy_workspace. - - Args: - src_workspace_id: Identifier of the source workspace. - dest_workspace_id: Identifier of the destination workspace. - Created if it doesn't exist. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - dict: Dictionary with "size" key indicating number of nodes copied, - or empty dict if source workspace doesn't exist. - """ - return await self._run_sync_in_executor( - self.copy_workspace, - src_workspace_id, - dest_workspace_id, - **kwargs, - ) - - def list_workspace(self, **kwargs) -> List[str]: - """List all existing workspaces (collections) in ChromaDB. - - Returns: - List[str]: Workspace identifiers (collection names). - """ - return [c.name for c in self._client.list_collections()] - - async def async_list_workspace(self, **kwargs) -> List[str]: - """Async version of list_workspace. - - Returns: - List[str]: Workspace identifiers (collection names). - """ - return await self._run_sync_in_executor(self.list_workspace, **kwargs) - - # ==================== Search Methods ==================== - - def search( - self, - query: str, - workspace_id: str, - top_k: int = 1, - filter_dict: Optional[Dict[str, Any]] = None, - **kwargs, - ) -> List[VectorNode]: - """Search for similar vector nodes in the workspace. - - When query is empty (empty string or None), the search degrades to a - filter-only search without vector similarity ranking. - - Args: - query: Text query to search for. If empty/None, performs filter-only - search without vector similarity. - workspace_id: The workspace identifier. - top_k: Number of top results to return (default: 1). - filter_dict: Optional metadata filters to apply. - **kwargs: Additional keyword arguments (unused). - - Returns: - List of VectorNode objects matching the query, ordered by similarity. - Returns empty list if workspace doesn't exist. - """ - if not self.exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return [] - - collection = self._get_collection(workspace_id) - - # Build where clause and ids filter from filter_dict - where_clause, ids_filter = self._build_chroma_filters(filter_dict) - - # When query is empty, degrade to filter-only search - use_vector_search = bool(query) - - if use_vector_search: - query_vector = self.get_embeddings(query) - query_kwargs = { - "query_embeddings": [query_vector], - "n_results": top_k, - } - if where_clause: - query_kwargs["where"] = where_clause - if ids_filter: - query_kwargs["ids"] = ids_filter - results = collection.query(**query_kwargs) - else: - # Filter-only search without vector similarity - get_kwargs = { - "limit": top_k, - "include": ["documents", "metadatas"], - } - if where_clause: - get_kwargs["where"] = where_clause - if ids_filter: - get_kwargs["ids"] = ids_filter - results = collection.get(**get_kwargs) - # Normalize results format to match query results - results = { - "ids": [results["ids"]], - "documents": [results["documents"]], - "metadatas": [results["metadatas"]], - "distances": None, - } - - nodes = [] - for i in range(len(results["ids"][0])): - node = VectorNode( - workspace_id=workspace_id, - unique_id=results["ids"][0][i], - content=results["documents"][0][i], - metadata=results["metadatas"][0][i], - ) - # ChromaDB returns distances, convert to similarity score - # Note: ChromaDB uses L2 (Euclidean) distance by default, but also supports - # cosine and inner product. Using 1 / (1 + distance) normalizes to [0, 1] - # and works for both L2 and cosine distance metrics. - if use_vector_search and results.get("distances") and len(results["distances"][0]) > i: - distance = results["distances"][0][i] - node.metadata["score"] = 1.0 / (1.0 + distance) - nodes.append(node) - - return nodes - - async def async_search( - self, - query: str, - workspace_id: str, - top_k: int = 1, - filter_dict: Optional[Dict[str, Any]] = None, - **kwargs, - ) -> List[VectorNode]: - """Async version of search using async embedding and run_in_executor for ChromaDB operations. - - When query is empty (empty string or None), the search degrades to a - filter-only search without vector similarity ranking. - - Args: - query: Text query to search for. If empty/None, performs filter-only - search without vector similarity. - workspace_id: The workspace identifier. - top_k: Number of top results to return (default: 1). - filter_dict: Optional metadata filters to apply. - **kwargs: Additional keyword arguments (unused). - - Returns: - List of VectorNode objects matching the query, ordered by similarity. - Returns empty list if workspace doesn't exist. - """ - if not await self.async_exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return [] - - # Build where clause and ids filter from filter_dict - where_clause, ids_filter = self._build_chroma_filters(filter_dict) - - # When query is empty, degrade to filter-only search - use_vector_search = bool(query) - - if use_vector_search: - # Use async embedding - query_vector = await self.async_get_embeddings(query) - - # Execute ChromaDB query in thread pool - def _do_query(): - collection = self._get_collection(workspace_id) - query_kwargs = { - "query_embeddings": [query_vector], - "n_results": top_k, - } - if where_clause: - query_kwargs["where"] = where_clause - if ids_filter: - query_kwargs["ids"] = ids_filter - return collection.query(**query_kwargs) - - results = await self._run_sync_in_executor(_do_query) - else: - # Filter-only search without vector similarity - - def _do_get(): - collection = self._get_collection(workspace_id) - get_kwargs = { - "limit": top_k, - "include": ["documents", "metadatas"], - } - if where_clause: - get_kwargs["where"] = where_clause - if ids_filter: - get_kwargs["ids"] = ids_filter - return collection.get(**get_kwargs) - - get_results = await self._run_sync_in_executor(_do_get) - # Normalize results format to match query results - results = { - "ids": [get_results["ids"]], - "documents": [get_results["documents"]], - "metadatas": [get_results["metadatas"]], - "distances": None, - } - - nodes = [] - for i in range(len(results["ids"][0])): - node = VectorNode( - workspace_id=workspace_id, - unique_id=results["ids"][0][i], - content=results["documents"][0][i], - metadata=results["metadatas"][0][i], - ) - # ChromaDB returns distances, convert to similarity score - # Note: ChromaDB uses L2 (Euclidean) distance by default, but also supports - # cosine and inner product. Using 1 / (1 + distance) normalizes to [0, 1] - # and works for both L2 and cosine distance metrics. - if use_vector_search and results.get("distances") and len(results["distances"][0]) > i: - distance = results["distances"][0][i] - node.metadata["score"] = 1.0 / (1.0 + distance) - nodes.append(node) - - return nodes - - # ==================== Insert Methods ==================== - - def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs): - """Insert vector nodes into the workspace. - - If nodes don't have embeddings, they will be generated using the embedding model. - Creates the workspace if it doesn't exist. - - Args: - nodes: Single VectorNode or list of VectorNode objects to insert. - workspace_id: The workspace identifier. - **kwargs: Additional keyword arguments (unused). - """ - if not self.exist_workspace(workspace_id=workspace_id): - self.create_workspace(workspace_id=workspace_id) - - if isinstance(nodes, VectorNode): - nodes = [nodes] - - # Generate embeddings for nodes that don't have them - all_nodes = self.get_node_embeddings(nodes) - - collection = self._get_collection(workspace_id) - - # Check for existing nodes and use upsert logic - existing_ids = set(collection.get()["ids"]) - update_cnt = sum(1 for n in all_nodes if n.unique_id in existing_ids) - - collection.upsert( - ids=[n.unique_id for n in all_nodes], - embeddings=[n.vector for n in all_nodes], - documents=[n.content for n in all_nodes], - metadatas=[n.metadata for n in all_nodes], - ) - - total_nodes = collection.count() - logger.info( - f"Inserted into workspace_id={workspace_id} nodes.size={len(all_nodes)} " - f"total.size={total_nodes} update_cnt={update_cnt}", - ) - - async def async_insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs): - """Async version of insert using async embedding and run_in_executor for ChromaDB operations. - - If nodes don't have embeddings, they will be generated using the async embedding model. - Creates the workspace if it doesn't exist. - - Args: - nodes: Single VectorNode or list of VectorNode objects to insert. - workspace_id: The workspace identifier. - **kwargs: Additional keyword arguments (unused). - """ - if not await self.async_exist_workspace(workspace_id=workspace_id): - await self.async_create_workspace(workspace_id=workspace_id) - - if isinstance(nodes, VectorNode): - nodes = [nodes] - - # Use async embedding - all_nodes = await self.async_get_node_embeddings(nodes) - - # Execute ChromaDB operations in thread pool - def _do_insert(): - collection = self._get_collection(workspace_id) - existing_ids = set(collection.get()["ids"]) - update_cnt = sum(1 for n in all_nodes if n.unique_id in existing_ids) - - collection.upsert( - ids=[n.unique_id for n in all_nodes], - embeddings=[n.vector for n in all_nodes], - documents=[n.content for n in all_nodes], - metadatas=[n.metadata for n in all_nodes], - ) - return collection.count(), update_cnt - - total_nodes, update_cnt = await self._run_sync_in_executor(_do_insert) - - logger.info( - f"Async inserted into workspace_id={workspace_id} nodes.size={len(all_nodes)} " - f"total.size={total_nodes} update_cnt={update_cnt}", - ) - - # ==================== Delete Methods ==================== - - def delete(self, node_ids: str | List[str], workspace_id: str, **kwargs): - """Delete vector nodes from the workspace. - - Args: - node_ids: Single node ID or list of node IDs to delete. - workspace_id: The workspace identifier. - **kwargs: Additional keyword arguments (unused). - """ - if not self.exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return - - if isinstance(node_ids, str): - node_ids = [node_ids] - - collection = self._get_collection(workspace_id) - before_size = collection.count() - collection.delete(ids=node_ids) - after_size = collection.count() - - logger.info( - f"Deleted from workspace_id={workspace_id} before_size={before_size} " - f"after_size={after_size} deleted_cnt={before_size - after_size}", - ) - - async def async_delete(self, node_ids: str | List[str], workspace_id: str, **kwargs): - """Async version of delete using run_in_executor for ChromaDB operations. - - Args: - node_ids: Single node ID or list of node IDs to delete. - workspace_id: The workspace identifier. - **kwargs: Additional keyword arguments (unused). - """ - if not await self.async_exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return - - if isinstance(node_ids, str): - node_ids = [node_ids] - - # Execute ChromaDB operations in thread pool - def _do_delete(): - collection = self._get_collection(workspace_id) - before_size = collection.count() - collection.delete(ids=node_ids) - after_size = collection.count() - return before_size, after_size - - before_size, after_size = await self._run_sync_in_executor(_do_delete) - - logger.info( - f"Deleted from workspace_id={workspace_id} before_size={before_size} " - f"after_size={after_size} deleted_cnt={before_size - after_size}", - ) - - # ==================== Close Methods ==================== - - def close(self): - """ - Close the vector store and clean up resources. - - For ChromaVectorStore, this clears the collections cache. - ChromaDB handles its own persistence. - """ - self.collections.clear() - - async def async_close(self): - """ - Async version of close. - - For ChromaVectorStore, this clears the collections cache. - """ - self.close() diff --git a/flowllm/core/vector_store/es_vector_store.py b/flowllm/core/vector_store/es_vector_store.py deleted file mode 100644 index e2fb9b3..0000000 --- a/flowllm/core/vector_store/es_vector_store.py +++ /dev/null @@ -1,866 +0,0 @@ -"""Elasticsearch vector store implementation. - -This module provides an Elasticsearch-based vector store that stores vector nodes -in Elasticsearch indices. It supports workspace management, vector similarity search, -metadata filtering using Elasticsearch query DSL, and provides both synchronous and -asynchronous operations using native Elasticsearch clients. -""" - -import os -from pathlib import Path -from typing import List, Tuple, Dict, Any, Optional - -from loguru import logger - -from .memory_vector_store import MemoryVectorStore -from ..context import C -from ..schema import VectorNode - - -@C.register_vector_store("elasticsearch") -class EsVectorStore(MemoryVectorStore): - """Elasticsearch vector store implementation. - - This class provides a vector store backend using Elasticsearch for storing and - searching vector embeddings. It supports both synchronous and asynchronous operations - using native Elasticsearch clients, and includes metadata filtering capabilities - using Elasticsearch query DSL. - - Inherits from MemoryVectorStore for utility methods like _load_from_path, - _dump_to_path, and _matches_filters. - - Attributes: - hosts: Elasticsearch host(s) as a string or list of strings. Defaults to - the FLOW_ES_HOSTS environment variable or "http://localhost:9200". - basic_auth: Optional basic authentication credentials as a string or - tuple of (username, password). - batch_size: Batch size for bulk operations. Defaults to 1024. - """ - - # ==================== Initialization ==================== - - def __init__( - self, - hosts: str | List[str] = "http://localhost:9200", - basic_auth: str | Tuple[str, str] | None = None, - batch_size: int = 1024, - **kwargs, - ): - """Initialize Elasticsearch clients. - - Args: - hosts: Elasticsearch host(s) as a string or list of strings. - basic_auth: Optional basic authentication credentials. - batch_size: Batch size for bulk operations. - **kwargs: Additional keyword arguments passed to MemoryVectorStore. - """ - super().__init__(**kwargs) - self.hosts = hosts or os.getenv("FLOW_ES_HOSTS", "http://localhost:9200") - self.basic_auth = basic_auth - self.batch_size = batch_size - - if isinstance(self.hosts, str): - self.hosts = [self.hosts] - - from elasticsearch import Elasticsearch, AsyncElasticsearch - - self._client = Elasticsearch(hosts=self.hosts, basic_auth=self.basic_auth) - self._async_client = AsyncElasticsearch(hosts=self.hosts, basic_auth=self.basic_auth) - logger.info(f"Elasticsearch client initialized with hosts={self.hosts} basic_auth={self.basic_auth}") - - # ==================== Static Helper Methods ==================== - - @staticmethod - def doc2node(doc, workspace_id: str, is_vector_search: bool = False) -> VectorNode: - """Convert an Elasticsearch document to a VectorNode. - - Args: - doc: The Elasticsearch document hit from a search response. - workspace_id: The workspace identifier to assign to the node. - is_vector_search: Whether this is a vector similarity search result. - If True, the score is adjusted by subtracting 1.0 (since we add 1.0 - in script_score to avoid negative scores). - - Returns: - VectorNode: A VectorNode instance created from the document data. - """ - node = VectorNode(**doc["_source"]) - node.workspace_id = workspace_id - node.unique_id = doc["_id"] - # Only adjust score for vector search (where we added 1.0 in script_score) - if is_vector_search and "_score" in doc: - node.metadata["score"] = doc["_score"] - 1 - return node - - # Top-level VectorNode fields that are stored directly in Elasticsearch - # (not nested under metadata) - _TOP_LEVEL_FIELDS = {"unique_id", "workspace_id", "content"} - - @staticmethod - def _build_es_filters(filter_dict: Optional[Dict[str, Any]] = None) -> List[Dict]: - """Build Elasticsearch filter clauses from filter_dict. - - Converts a filter dictionary into Elasticsearch query filter clauses. - Supports both term filters (exact match) and range filters (gte, lte, gt, lt). - - Handles top-level VectorNode fields (unique_id, workspace_id, content) - separately from metadata fields. The unique_id field is stored as Elasticsearch's - _id, so it uses the ids query. Other top-level fields are stored directly in - _source without the metadata prefix. - - Args: - filter_dict: Dictionary of filter conditions. Keys can be: - - Top-level VectorNode fields: unique_id, workspace_id, content - - Metadata fields: node_type, category, etc. (auto-prefixed with "metadata.") - - Already-prefixed metadata fields: metadata.node_type, etc. - Values can be exact match values or range dictionaries like - {"gte": 1, "lte": 10}. - - Returns: - List[Dict]: List of Elasticsearch filter clauses. - """ - if not filter_dict: - return [] - - filters = [] - for key, filter_value in filter_dict.items(): - # Handle unique_id specially - it's stored as Elasticsearch's _id - if key == "unique_id": - # Use ids query to filter by document ID - # Support both single value and list of values - if isinstance(filter_value, list): - filters.append({"ids": {"values": filter_value}}) - else: - filters.append({"ids": {"values": [filter_value]}}) - continue - - # Determine the Elasticsearch field key - if key in EsVectorStore._TOP_LEVEL_FIELDS: - # Top-level fields (workspace_id, content) are stored directly in _source - es_key = key - elif key.startswith("metadata."): - # Already prefixed with metadata. - es_key = key - else: - # Assume it's a metadata field and prefix it - es_key = f"metadata.{key}" - - if isinstance(filter_value, dict): - # Range filter: {"gte": 1, "lte": 10} - range_conditions = {} - if "gte" in filter_value: - range_conditions["gte"] = filter_value["gte"] - if "lte" in filter_value: - range_conditions["lte"] = filter_value["lte"] - if "gt" in filter_value: - range_conditions["gt"] = filter_value["gt"] - if "lt" in filter_value: - range_conditions["lt"] = filter_value["lt"] - if range_conditions: - filters.append({"range": {es_key: range_conditions}}) - elif isinstance(filter_value, list): - # List filter: use terms query for OR logic - filters.append({"terms": {es_key: filter_value}}) - else: - # Term filter: direct value comparison - filters.append({"term": {es_key: filter_value}}) - - return filters - - # ==================== ES-specific Helper Methods ==================== - - def refresh(self, workspace_id: str): - """Refresh an Elasticsearch index to make recent changes visible for search. - - Args: - workspace_id: The identifier of the workspace/index to refresh. - """ - self._client.indices.refresh(index=workspace_id) - - async def async_refresh(self, workspace_id: str): - """Refresh an Elasticsearch index to make recent changes visible for search (async). - - Args: - workspace_id: The identifier of the workspace/index to refresh. - """ - await self._async_client.indices.refresh(index=workspace_id) - - # ==================== Workspace Management Methods ==================== - - def exist_workspace(self, workspace_id: str, **kwargs) -> bool: - """Check if an Elasticsearch index (workspace) exists. - - Args: - workspace_id: The identifier of the workspace/index to check. - **kwargs: Additional keyword arguments passed to Elasticsearch API. - - Returns: - bool: True if the index exists, False otherwise. - """ - return self._client.indices.exists(index=workspace_id) - - async def async_exist_workspace(self, workspace_id: str, **kwargs) -> bool: - """Check if an Elasticsearch index (workspace) exists (async). - - Args: - workspace_id: The identifier of the workspace/index to check. - **kwargs: Additional keyword arguments passed to Elasticsearch API. - - Returns: - bool: True if the index exists, False otherwise. - """ - return await self._async_client.indices.exists(index=workspace_id) - - def delete_workspace(self, workspace_id: str, **kwargs): - """Delete an Elasticsearch index (workspace). - - Args: - workspace_id: The identifier of the workspace/index to delete. - **kwargs: Additional keyword arguments passed to Elasticsearch API. - """ - return self._client.indices.delete(index=workspace_id, **kwargs) - - async def async_delete_workspace(self, workspace_id: str, **kwargs): - """Delete an Elasticsearch index (workspace) (async). - - Args: - workspace_id: The identifier of the workspace/index to delete. - **kwargs: Additional keyword arguments passed to Elasticsearch API. - """ - return await self._async_client.indices.delete(index=workspace_id, **kwargs) - - def create_workspace(self, workspace_id: str, **kwargs): - """Create a new Elasticsearch index (workspace) with vector field mappings. - - Args: - workspace_id: The identifier of the workspace/index to create. - **kwargs: Additional keyword arguments passed to Elasticsearch API. - - Returns: - The response from Elasticsearch create index API. - """ - body = { - "mappings": { - "properties": { - "workspace_id": {"type": "keyword"}, - "content": {"type": "text"}, - "metadata": {"type": "object"}, - "vector": { - "type": "dense_vector", - "dims": self.embedding_model.dimensions, - }, - }, - }, - } - return self._client.indices.create(index=workspace_id, body=body) - - async def async_create_workspace(self, workspace_id: str, **kwargs): - """Create a new Elasticsearch index (workspace) with vector field mappings (async). - - Args: - workspace_id: The identifier of the workspace/index to create. - **kwargs: Additional keyword arguments passed to Elasticsearch API. - - Returns: - The response from Elasticsearch create index API. - """ - body = { - "mappings": { - "properties": { - "workspace_id": {"type": "keyword"}, - "content": {"type": "text"}, - "metadata": {"type": "object"}, - "vector": { - "type": "dense_vector", - "dims": self.embedding_model.dimensions, - }, - }, - }, - } - return await self._async_client.indices.create(index=workspace_id, body=body) - - def list_workspace_nodes( - self, - workspace_id: str, - max_size: int = 10000, - **kwargs, - ) -> List[VectorNode]: - """List all nodes in a workspace. - - Args: - workspace_id: The identifier of the workspace to list. - max_size: Maximum number of nodes to retrieve (default: 10000). - **kwargs: Additional keyword arguments passed to Elasticsearch API. - - Returns: - List[VectorNode]: Vector nodes from the workspace. - """ - if not self.exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return [] - - response = self._client.search(index=workspace_id, body={"query": {"match_all": {}}, "size": max_size}) - nodes = [] - for doc in response["hits"]["hits"]: - node = self.doc2node(doc, workspace_id, is_vector_search=False) - nodes.append(node) - return nodes - - async def async_list_workspace_nodes( - self, - workspace_id: str, - max_size: int = 10000, - **kwargs, - ) -> List[VectorNode]: - """List all nodes in a workspace (async). - - Args: - workspace_id: The identifier of the workspace to list. - max_size: Maximum number of nodes to retrieve (default: 10000). - **kwargs: Additional keyword arguments passed to Elasticsearch API. - - Returns: - List[VectorNode]: Vector nodes from the workspace. - """ - if not await self.async_exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return [] - - response = await self._async_client.search( - index=workspace_id, - body={"query": {"match_all": {}}, "size": max_size}, - ) - nodes = [] - for doc in response["hits"]["hits"]: - node = self.doc2node(doc, workspace_id, is_vector_search=False) - nodes.append(node) - return nodes - - def dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs): - """Export a workspace to disk at the specified path. - - Args: - workspace_id: Identifier of the workspace to export. - path: Directory path where to write the exported workspace file. - callback_fn: Optional callback function to transform nodes during export. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - - Returns: - dict: Dictionary with "size" key indicating number of nodes exported, - or empty dict if workspace doesn't exist. - """ - if not self.exist_workspace(workspace_id=workspace_id, **kwargs): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return {} - - return self._dump_to_path( - nodes=self.list_workspace_nodes(workspace_id=workspace_id, **kwargs), - workspace_id=workspace_id, - path=path, - callback_fn=callback_fn, - **kwargs, - ) - - async def async_dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs): - """Export a workspace to disk at the specified path (async). - - Args: - workspace_id: Identifier of the workspace to export. - path: Directory path where to write the exported workspace file. - callback_fn: Optional callback function to transform nodes during export. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - - Returns: - dict: Dictionary with "size" key indicating number of nodes exported, - or empty dict if workspace doesn't exist. - """ - if not await self.async_exist_workspace(workspace_id=workspace_id, **kwargs): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return {} - - nodes = await self.async_list_workspace_nodes(workspace_id=workspace_id, **kwargs) - return await self._run_sync_in_executor( - self._dump_to_path, - nodes=nodes, - workspace_id=workspace_id, - path=path, - callback_fn=callback_fn, - **kwargs, - ) - - def load_workspace( - self, - workspace_id: str, - path: str | Path = "", - nodes: List[VectorNode] | None = None, - callback_fn=None, - **kwargs, - ): - """Load a workspace from disk, optionally merging with provided nodes. - - This method replaces any existing workspace with the same ID, then loads - nodes from the specified path and/or the provided nodes list. - - Args: - workspace_id: Identifier for the workspace to create/load. - path: Directory path containing the workspace file to load. - nodes: Optional list of VectorNode instances to merge with loaded nodes. - callback_fn: Optional callback function to transform node dictionaries. - **kwargs: Additional keyword arguments to pass to load operations. - - Returns: - dict: Dictionary with "size" key indicating total number of nodes loaded. - """ - if self.exist_workspace(workspace_id, **kwargs): - self.delete_workspace(workspace_id=workspace_id, **kwargs) - logger.info(f"delete workspace_id={workspace_id}") - - self.create_workspace(workspace_id=workspace_id, **kwargs) - - all_nodes: List[VectorNode] = [] - - if nodes: - all_nodes.extend(nodes) - - all_nodes.extend(self._load_from_path(path=path, workspace_id=workspace_id, callback_fn=callback_fn, **kwargs)) - - if all_nodes: - self.insert(nodes=all_nodes, workspace_id=workspace_id, **kwargs) - return {"size": len(all_nodes)} - - async def async_load_workspace( - self, - workspace_id: str, - path: str | Path = "", - nodes: List[VectorNode] | None = None, - callback_fn=None, - **kwargs, - ): - """Load a workspace from disk, optionally merging with provided nodes (async). - - This method replaces any existing workspace with the same ID, then loads - nodes from the specified path and/or the provided nodes list. - - Args: - workspace_id: Identifier for the workspace to create/load. - path: Directory path containing the workspace file to load. - nodes: Optional list of VectorNode instances to merge with loaded nodes. - callback_fn: Optional callback function to transform node dictionaries. - **kwargs: Additional keyword arguments to pass to load operations. - - Returns: - dict: Dictionary with "size" key indicating total number of nodes loaded. - """ - if await self.async_exist_workspace(workspace_id, **kwargs): - await self.async_delete_workspace(workspace_id=workspace_id, **kwargs) - logger.info(f"delete workspace_id={workspace_id}") - - await self.async_create_workspace(workspace_id=workspace_id, **kwargs) - - all_nodes: List[VectorNode] = [] - - if nodes: - all_nodes.extend(nodes) - - if path: - loaded_nodes = await self._run_sync_in_executor( - lambda: list( - self._load_from_path(path=path, workspace_id=workspace_id, callback_fn=callback_fn, **kwargs), - ), - ) - all_nodes.extend(loaded_nodes) - - if all_nodes: - await self.async_insert(nodes=all_nodes, workspace_id=workspace_id, **kwargs) - return {"size": len(all_nodes)} - - def copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs): - """Copy all nodes from one workspace to another. - - Args: - src_workspace_id: Identifier of the source workspace. - dest_workspace_id: Identifier of the destination workspace. - Created if it doesn't exist. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - dict: Dictionary with "size" key indicating number of nodes copied, - or empty dict if source workspace doesn't exist. - """ - if not self.exist_workspace(workspace_id=src_workspace_id, **kwargs): - logger.warning(f"src_workspace_id={src_workspace_id} does not exist!") - return {} - - if not self.exist_workspace(dest_workspace_id, **kwargs): - self.create_workspace(workspace_id=dest_workspace_id, **kwargs) - - nodes = [] - node_size = 0 - for node in self.list_workspace_nodes(workspace_id=src_workspace_id, **kwargs): - nodes.append(node) - node_size += 1 - if len(nodes) >= self.batch_size: - self.insert(nodes=nodes, workspace_id=dest_workspace_id, **kwargs) - nodes.clear() - - if nodes: - self.insert(nodes=nodes, workspace_id=dest_workspace_id, **kwargs) - return {"size": node_size} - - async def async_copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs): - """Copy all nodes from one workspace to another (async). - - Args: - src_workspace_id: Identifier of the source workspace. - dest_workspace_id: Identifier of the destination workspace. - Created if it doesn't exist. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - dict: Dictionary with "size" key indicating number of nodes copied, - or empty dict if source workspace doesn't exist. - """ - if not await self.async_exist_workspace(workspace_id=src_workspace_id, **kwargs): - logger.warning(f"src_workspace_id={src_workspace_id} does not exist!") - return {} - - if not await self.async_exist_workspace(dest_workspace_id, **kwargs): - await self.async_create_workspace(workspace_id=dest_workspace_id, **kwargs) - - src_nodes = await self.async_list_workspace_nodes(workspace_id=src_workspace_id, **kwargs) - nodes = [] - node_size = 0 - for node in src_nodes: - nodes.append(node) - node_size += 1 - if len(nodes) >= self.batch_size: - await self.async_insert(nodes=nodes, workspace_id=dest_workspace_id, **kwargs) - nodes.clear() - - if nodes: - await self.async_insert(nodes=nodes, workspace_id=dest_workspace_id, **kwargs) - return {"size": node_size} - - def list_workspace(self, **kwargs) -> List[str]: - """List all existing workspaces (indices) in Elasticsearch. - - Returns: - List[str]: Workspace identifiers (index names). - """ - return list(self._client.indices.get(index="*").keys()) - - async def async_list_workspace(self, **kwargs) -> List[str]: - """List all existing workspaces (indices) in Elasticsearch (async). - - Returns: - List[str]: Workspace identifiers (index names). - """ - result = await self._async_client.indices.get(index="*") - return list(result.keys()) - - # ==================== Search Methods ==================== - - def search( - self, - query: str, - workspace_id: str, - top_k: int = 1, - filter_dict: Optional[Dict[str, Any]] = None, - **kwargs, - ) -> List[VectorNode]: - """Search for similar vector nodes using cosine similarity. - - When query is empty (empty string or None), the search degrades to a - filter-only search without vector similarity ranking. - - Args: - query: The text query to search for. Will be embedded using the - embedding model. If empty/None, performs filter-only search. - workspace_id: The identifier of the workspace to search in. - top_k: Maximum number of results to return (default: 1). - filter_dict: Optional dictionary of metadata filters to apply. - **kwargs: Additional keyword arguments passed to Elasticsearch API. - - Returns: - List[VectorNode]: List of matching vector nodes. When query is provided, - nodes are sorted by similarity score (highest first) with a "score" - key in metadata. When query is empty, nodes are returned based on - filter criteria without scores. - """ - if not self.exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return [] - - # Build filters from filter_dict - es_filters = self._build_es_filters(filter_dict) - - # When query is empty, degrade to filter-only search - use_vector_search = bool(query) - - if use_vector_search: - query_vector = self.get_embeddings(query) - body = { - "query": { - "script_score": { - "query": {"bool": {"must": es_filters}} if es_filters else {"match_all": {}}, - "script": { - "source": "cosineSimilarity(params.query_vector, 'vector') + 1.0", - "params": {"query_vector": query_vector}, - }, - }, - }, - "size": top_k, - } - else: - # Filter-only search without vector similarity - body = { - "query": {"bool": {"must": es_filters}} if es_filters else {"match_all": {}}, - "size": top_k, - } - - response = self._client.search(index=workspace_id, body=body, **kwargs) - - nodes: List[VectorNode] = [] - for doc in response["hits"]["hits"]: - node = self.doc2node(doc, workspace_id, is_vector_search=use_vector_search) - nodes.append(node) - - return nodes - - async def async_search( - self, - query: str, - workspace_id: str, - top_k: int = 1, - filter_dict: Optional[Dict[str, Any]] = None, - **kwargs, - ) -> List[VectorNode]: - """Search for similar vector nodes using cosine similarity (async). - - When query is empty (empty string or None), the search degrades to a - filter-only search without vector similarity ranking. - - Args: - query: The text query to search for. Will be embedded using the - embedding model's async method. If empty/None, performs filter-only search. - workspace_id: The identifier of the workspace to search in. - top_k: Maximum number of results to return (default: 1). - filter_dict: Optional dictionary of metadata filters to apply. - **kwargs: Additional keyword arguments passed to Elasticsearch API. - - Returns: - List[VectorNode]: List of matching vector nodes. When query is provided, - nodes are sorted by similarity score (highest first) with a "score" - key in metadata. When query is empty, nodes are returned based on - filter criteria without scores. - """ - if not await self.async_exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return [] - - # Build filters from filter_dict - es_filters = self._build_es_filters(filter_dict) - - # When query is empty, degrade to filter-only search - use_vector_search = bool(query) - - if use_vector_search: - # Use async embedding - query_vector = await self.async_get_embeddings(query) - body = { - "query": { - "script_score": { - "query": {"bool": {"must": es_filters}} if es_filters else {"match_all": {}}, - "script": { - "source": "cosineSimilarity(params.query_vector, 'vector') + 1.0", - "params": {"query_vector": query_vector}, - }, - }, - }, - "size": top_k, - } - else: - # Filter-only search without vector similarity - body = { - "query": {"bool": {"must": es_filters}} if es_filters else {"match_all": {}}, - "size": top_k, - } - - response = await self._async_client.search(index=workspace_id, body=body, **kwargs) - - nodes: List[VectorNode] = [] - for doc in response["hits"]["hits"]: - node = self.doc2node(doc, workspace_id, is_vector_search=use_vector_search) - nodes.append(node) - - return nodes - - # ==================== Insert Methods ==================== - - def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, refresh: bool = True, **kwargs): - """Insert vector nodes into the Elasticsearch index. - - Args: - nodes: A single VectorNode or list of VectorNodes to insert. - workspace_id: The identifier of the workspace to insert into. - refresh: Whether to refresh the index after insertion (default: True). - **kwargs: Additional keyword arguments passed to Elasticsearch bulk API. - """ - if not self.exist_workspace(workspace_id=workspace_id): - self.create_workspace(workspace_id=workspace_id) - - if isinstance(nodes, VectorNode): - nodes = [nodes] - - embedded_nodes = [node for node in nodes if node.vector] - not_embedded_nodes = [node for node in nodes if not node.vector] - now_embedded_nodes = self.get_node_embeddings(not_embedded_nodes) if not_embedded_nodes else [] - - docs = [ - { - "_op_type": "index", - "_index": workspace_id, - "_id": node.unique_id, - "_source": { - "workspace_id": workspace_id, - "content": node.content, - "metadata": node.metadata, - "vector": node.vector, - }, - } - for node in embedded_nodes + now_embedded_nodes - ] - from elasticsearch.helpers import bulk - - status, error = bulk(self._client, docs, chunk_size=self.batch_size, **kwargs) - logger.info(f"insert docs.size={len(docs)} status={status} error={error}") - - if refresh: - self.refresh(workspace_id=workspace_id) - - async def async_insert( - self, - nodes: VectorNode | List[VectorNode], - workspace_id: str, - refresh: bool = True, - **kwargs, - ): - """Insert vector nodes into the Elasticsearch index (async). - - Args: - nodes: A single VectorNode or list of VectorNodes to insert. - workspace_id: The identifier of the workspace to insert into. - refresh: Whether to refresh the index after insertion (default: True). - **kwargs: Additional keyword arguments passed to Elasticsearch bulk API. - """ - if not await self.async_exist_workspace(workspace_id=workspace_id): - await self.async_create_workspace(workspace_id=workspace_id) - - if isinstance(nodes, VectorNode): - nodes = [nodes] - - embedded_nodes = [node for node in nodes if node.vector] - not_embedded_nodes = [node for node in nodes if not node.vector] - - # Use async embedding - now_embedded_nodes = await self.async_get_node_embeddings(not_embedded_nodes) if not_embedded_nodes else [] - - docs = [ - { - "_op_type": "index", - "_index": workspace_id, - "_id": node.unique_id, - "_source": { - "workspace_id": workspace_id, - "content": node.content, - "metadata": node.metadata, - "vector": node.vector, - }, - } - for node in embedded_nodes + now_embedded_nodes - ] - - from elasticsearch.helpers import async_bulk - - status, error = await async_bulk(self._async_client, docs, chunk_size=self.batch_size, **kwargs) - logger.info(f"async insert docs.size={len(docs)} status={status} error={error}") - - if refresh: - await self.async_refresh(workspace_id=workspace_id) - - # ==================== Delete Methods ==================== - - def delete(self, node_ids: str | List[str], workspace_id: str, refresh: bool = True, **kwargs): - """Delete vector nodes from the Elasticsearch index. - - Args: - node_ids: A single node ID or list of node IDs to delete. - workspace_id: The identifier of the workspace to delete from. - refresh: Whether to refresh the index after deletion (default: True). - **kwargs: Additional keyword arguments passed to Elasticsearch bulk API. - """ - if not self.exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return - - if isinstance(node_ids, str): - node_ids = [node_ids] - - actions = [ - { - "_op_type": "delete", - "_index": workspace_id, - "_id": node_id, - } - for node_id in node_ids - ] - from elasticsearch.helpers import bulk - - status, error = bulk(self._client, actions, chunk_size=self.batch_size, **kwargs) - logger.info(f"delete actions.size={len(actions)} status={status} error={error}") - - if refresh: - self.refresh(workspace_id=workspace_id) - - async def async_delete(self, node_ids: str | List[str], workspace_id: str, refresh: bool = True, **kwargs): - """Delete vector nodes from the Elasticsearch index (async). - - Args: - node_ids: A single node ID or list of node IDs to delete. - workspace_id: The identifier of the workspace to delete from. - refresh: Whether to refresh the index after deletion (default: True). - **kwargs: Additional keyword arguments passed to Elasticsearch bulk API. - """ - if not await self.async_exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return - - if isinstance(node_ids, str): - node_ids = [node_ids] - - actions = [ - { - "_op_type": "delete", - "_index": workspace_id, - "_id": node_id, - } - for node_id in node_ids - ] - - from elasticsearch.helpers import async_bulk - - status, error = await async_bulk(self._async_client, actions, chunk_size=self.batch_size, **kwargs) - logger.info(f"async delete actions.size={len(actions)} status={status} error={error}") - - if refresh: - await self.async_refresh(workspace_id=workspace_id) - - # ==================== Close Methods ==================== - - def close(self): - """Close the synchronous Elasticsearch client connection.""" - self._client.close() - - async def async_close(self): - """Close the asynchronous Elasticsearch client connection.""" - await self._async_client.close() diff --git a/flowllm/core/vector_store/local_vector_store.py b/flowllm/core/vector_store/local_vector_store.py deleted file mode 100644 index d0851de..0000000 --- a/flowllm/core/vector_store/local_vector_store.py +++ /dev/null @@ -1,590 +0,0 @@ -""" -Local file-based vector store implementation. - -This module provides a local file-based vector store that stores vector nodes -in JSONL format on disk. It extends MemoryVectorStore to provide persistence -while maintaining fast in-memory operations. -""" - -from pathlib import Path -from typing import List, Optional, Dict, Any - -from loguru import logger - -from .memory_vector_store import MemoryVectorStore -from ..context.service_context import C -from ..schema.vector_node import VectorNode - - -@C.register_vector_store("local") -class LocalVectorStore(MemoryVectorStore): - """ - Local file-based vector store implementation. - - This vector store extends MemoryVectorStore to persist all data to JSONL files - on disk. Each workspace is stored as a separate file. It combines the fast - in-memory operations of MemoryVectorStore with automatic file persistence. - - Attributes: - store_dir: Directory path where workspace files are stored. - Defaults to "./local_vector_store". - """ - - # ==================== Initialization ==================== - - def __init__(self, store_dir: str = "./local_vector_store", **kwargs): - """ - Initialize the vector store by creating the storage directory if it doesn't exist. - - Args: - store_dir: Directory path where workspace files are stored. - **kwargs: Additional keyword arguments passed to MemoryVectorStore. - """ - super().__init__(**kwargs) - self.store_dir = store_dir - store_path = Path(self.store_dir) - store_path.mkdir(parents=True, exist_ok=True) - logger.info(f"LocalVectorStore initialized with store_dir={self.store_dir}") - - @property - def store_path(self) -> Path: - """ - Get the storage directory path. - - Returns: - Path object representing the storage directory. - """ - return Path(self.store_dir) - - # ==================== Workspace Management Methods ==================== - - def exist_workspace(self, workspace_id: str, **kwargs) -> bool: - """ - Check if a workspace exists (either in memory or on disk). - - Args: - workspace_id: Identifier of the workspace to check. - **kwargs: Additional keyword arguments (unused). - - Returns: - bool: True if the workspace exists in memory or on disk, False otherwise. - """ - # Check memory first, then disk - if super().exist_workspace(workspace_id, **kwargs): - return True - workspace_path = self.store_path / f"{workspace_id}.jsonl" - return workspace_path.exists() - - async def async_exist_workspace(self, workspace_id: str, **kwargs) -> bool: - """ - Async version of exist_workspace. - - Args: - workspace_id: Identifier of the workspace to check. - **kwargs: Additional keyword arguments (unused). - - Returns: - bool: True if the workspace exists in memory or on disk, False otherwise. - """ - return self.exist_workspace(workspace_id, **kwargs) - - def delete_workspace(self, workspace_id: str, **kwargs): - """ - Delete a workspace from both memory and disk. - - Args: - workspace_id: Identifier of the workspace to delete. - **kwargs: Additional keyword arguments (unused). - """ - # Delete from memory - super().delete_workspace(workspace_id, **kwargs) - # Delete from disk - workspace_path = self.store_path / f"{workspace_id}.jsonl" - if workspace_path.is_file(): - workspace_path.unlink() - logger.info(f"Deleted workspace file: {workspace_path}") - - async def async_delete_workspace(self, workspace_id: str, **kwargs): - """ - Async version of delete_workspace. - - Args: - workspace_id: Identifier of the workspace to delete. - **kwargs: Additional keyword arguments (unused). - """ - return await self._run_sync_in_executor(self.delete_workspace, workspace_id, **kwargs) - - def create_workspace(self, workspace_id: str, **kwargs): - """ - Create a new empty workspace in memory and on disk. - - Args: - workspace_id: Identifier for the new workspace. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - """ - # Create in memory - super().create_workspace(workspace_id, **kwargs) - # Create empty file on disk - self._dump_to_path(nodes=[], workspace_id=workspace_id, path=self.store_path, show_progress=False, **kwargs) - - async def async_create_workspace(self, workspace_id: str, **kwargs): - """ - Async version of create_workspace. - - Args: - workspace_id: Identifier for the new workspace. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - """ - return await self._run_sync_in_executor(self.create_workspace, workspace_id, **kwargs) - - def list_workspace_nodes(self, workspace_id: str, **kwargs) -> List[VectorNode]: - """ - List all nodes in a workspace. - - If the workspace is in memory, returns from memory. Otherwise, loads from disk. - - Args: - workspace_id: Identifier of the workspace. - **kwargs: Additional keyword arguments to pass to _load_from_path. - - Returns: - List[VectorNode]: All nodes from the workspace. - """ - # If in memory, return from memory - if workspace_id in self._memory_store: - return super().list_workspace_nodes(workspace_id, **kwargs) - # Otherwise, load from disk - return self._load_from_path( - path=self.store_path, - workspace_id=workspace_id, - **kwargs, - ) - - async def async_list_workspace_nodes(self, workspace_id: str, **kwargs) -> List[VectorNode]: - """ - Async version of list_workspace_nodes. - - Args: - workspace_id: Identifier of the workspace. - **kwargs: Additional keyword arguments to pass to _load_from_path. - - Returns: - List[VectorNode]: All nodes from the workspace. - """ - # If in memory, return from memory (fast path) - if workspace_id in self._memory_store: - return super().list_workspace_nodes(workspace_id, **kwargs) - # Otherwise, load from disk asynchronously - return await self._run_sync_in_executor( - self._load_from_path, - path=self.store_path, - workspace_id=workspace_id, - **kwargs, - ) - - def dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs): - """ - Export a workspace to disk at the specified path. - - Args: - workspace_id: Identifier of the workspace to export. - path: Directory path where to write the exported workspace file. - If empty, uses the current store_path. - callback_fn: Optional callback function to transform nodes during export. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - - Returns: - dict: Dictionary with "size" key indicating number of nodes exported, - or empty dict if workspace doesn't exist. - """ - if not self.exist_workspace(workspace_id=workspace_id, **kwargs): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return {} - - # Use provided path or default to store_path - dump_path = path if path else self.store_path - - return self._dump_to_path( - nodes=self.list_workspace_nodes(workspace_id=workspace_id, **kwargs), - workspace_id=workspace_id, - path=dump_path, - callback_fn=callback_fn, - **kwargs, - ) - - async def async_dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs): - """ - Async version of dump_workspace. - - Args: - workspace_id: Identifier of the workspace to export. - path: Directory path where to write the exported workspace file. - If empty, uses the current store_path. - callback_fn: Optional callback function to transform nodes during export. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - - Returns: - dict: Dictionary with "size" key indicating number of nodes exported, - or empty dict if workspace doesn't exist. - """ - return await self._run_sync_in_executor( - self.dump_workspace, - workspace_id, - path, - callback_fn, - **kwargs, - ) - - def load_workspace( - self, - workspace_id: str, - path: str | Path = "", - nodes: List[VectorNode] | None = None, - callback_fn=None, - **kwargs, - ): - """ - Load a workspace from disk into memory, optionally merging with provided nodes. - - This method replaces any existing workspace with the same ID, then loads - nodes from the specified path and/or the provided nodes list. - - Args: - workspace_id: Identifier for the workspace to create/load. - path: Directory path containing the workspace file to load. - If empty, loads from the current store_path. - nodes: Optional list of VectorNode instances to merge with loaded nodes. - callback_fn: Optional callback function to transform node dictionaries. - **kwargs: Additional keyword arguments to pass to load operations. - - Returns: - dict: Dictionary with "size" key indicating total number of nodes loaded. - """ - if self.exist_workspace(workspace_id, **kwargs): - self.delete_workspace(workspace_id=workspace_id, **kwargs) - logger.info(f"Deleted existing workspace_id={workspace_id}") - - # Create workspace in memory - super().create_workspace(workspace_id=workspace_id, **kwargs) - - all_nodes: List[VectorNode] = [] - - if nodes: - all_nodes.extend(nodes) - - # Load from path (use store_path if path is empty) - load_path = path if path else self.store_path - all_nodes.extend( - self._load_from_path(path=load_path, workspace_id=workspace_id, callback_fn=callback_fn, **kwargs), - ) - - if all_nodes: - self.insert(nodes=all_nodes, workspace_id=workspace_id, **kwargs) - - return {"size": len(all_nodes)} - - async def async_load_workspace( - self, - workspace_id: str, - path: str | Path = "", - nodes: List[VectorNode] | None = None, - callback_fn=None, - **kwargs, - ): - """ - Async version of load_workspace. - - Args: - workspace_id: Identifier for the workspace to create/load. - path: Directory path containing the workspace file to load. - If empty, loads from the current store_path. - nodes: Optional list of VectorNode instances to merge with loaded nodes. - callback_fn: Optional callback function to transform node dictionaries. - **kwargs: Additional keyword arguments to pass to load operations. - - Returns: - dict: Dictionary with "size" key indicating total number of nodes loaded. - """ - return await self._run_sync_in_executor( - self.load_workspace, - workspace_id, - path, - nodes, - callback_fn, - **kwargs, - ) - - def copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs): - """ - Copy all nodes from one workspace to another. - - Args: - src_workspace_id: Identifier of the source workspace. - dest_workspace_id: Identifier of the destination workspace. - Created if it doesn't exist. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - dict: Dictionary with "size" key indicating number of nodes copied, - or empty dict if source workspace doesn't exist. - """ - if not self.exist_workspace(workspace_id=src_workspace_id, **kwargs): - logger.warning(f"src_workspace_id={src_workspace_id} does not exist!") - return {} - - if not self.exist_workspace(dest_workspace_id, **kwargs): - self.create_workspace(workspace_id=dest_workspace_id, **kwargs) - - # Ensure source workspace is loaded into memory - self._ensure_workspace_loaded(src_workspace_id, **kwargs) - - # Use parent's copy_workspace for in-memory copy - result = super().copy_workspace(src_workspace_id, dest_workspace_id, **kwargs) - - # Persist destination workspace to disk - self._persist_workspace(dest_workspace_id, **kwargs) - - return result - - async def async_copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs): - """ - Async version of copy_workspace. - - Args: - src_workspace_id: Identifier of the source workspace. - dest_workspace_id: Identifier of the destination workspace. - Created if it doesn't exist. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - dict: Dictionary with "size" key indicating number of nodes copied, - or empty dict if source workspace doesn't exist. - """ - return await self._run_sync_in_executor( - self.copy_workspace, - src_workspace_id, - dest_workspace_id, - **kwargs, - ) - - def list_workspace(self, **kwargs) -> List[str]: - """ - List all existing workspaces (from disk). - - Returns: - List[str]: Workspace identifiers discovered in the storage directory. - """ - return [p.stem for p in self.store_path.glob("*.jsonl") if p.is_file()] - - async def async_list_workspace(self, **kwargs) -> List[str]: - """ - Async version of list_workspace. - - Returns: - List[str]: Workspace identifiers discovered in the storage directory. - """ - return await self._run_sync_in_executor(self.list_workspace, **kwargs) - - # ==================== Helper Methods ==================== - - def _ensure_workspace_loaded(self, workspace_id: str, **kwargs): - """ - Ensure a workspace is loaded into memory from disk if not already present. - - Args: - workspace_id: Identifier of the workspace to load. - **kwargs: Additional keyword arguments to pass to _load_from_path. - """ - if workspace_id not in self._memory_store: - nodes = self._load_from_path(path=self.store_path, workspace_id=workspace_id, **kwargs) - self._memory_store[workspace_id] = {node.unique_id: node for node in nodes} - - async def _async_ensure_workspace_loaded(self, workspace_id: str, **kwargs): - """ - Async version of _ensure_workspace_loaded. - - Args: - workspace_id: Identifier of the workspace to load. - **kwargs: Additional keyword arguments to pass to _load_from_path. - """ - if workspace_id not in self._memory_store: - nodes = await self._run_sync_in_executor( - self._load_from_path, - path=self.store_path, - workspace_id=workspace_id, - **kwargs, - ) - self._memory_store[workspace_id] = {node.unique_id: node for node in nodes} - - def _persist_workspace(self, workspace_id: str, **kwargs): - """ - Persist a workspace from memory to disk. - - Args: - workspace_id: Identifier of the workspace to persist. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - """ - if workspace_id in self._memory_store: - nodes = list(self._memory_store[workspace_id].values()) - self._dump_to_path( - nodes=nodes, - workspace_id=workspace_id, - path=self.store_path, - show_progress=False, - **kwargs, - ) - - async def _async_persist_workspace(self, workspace_id: str, **kwargs): - """ - Async version of _persist_workspace. - - Args: - workspace_id: Identifier of the workspace to persist. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - """ - if workspace_id in self._memory_store: - nodes = list(self._memory_store[workspace_id].values()) - await self._run_sync_in_executor( - self._dump_to_path, - nodes=nodes, - workspace_id=workspace_id, - path=self.store_path, - show_progress=False, - **kwargs, - ) - - # ==================== Search Methods ==================== - - def search( - self, - query: str, - workspace_id: str, - top_k: int = 1, - filter_dict: Optional[Dict[str, Any]] = None, - **kwargs, - ) -> List[VectorNode]: - """ - Search for similar nodes using vector similarity. - - Loads workspace from disk if not in memory, then performs search. - - Args: - query: Text query to search for. If empty/None, performs filter-only - search without vector similarity. - workspace_id: Identifier of the workspace to search in. - top_k: Number of top results to return. Defaults to 1. - filter_dict: Optional dictionary of filters to apply to nodes. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - List[VectorNode]: List of matching nodes sorted by similarity score. - """ - # Ensure workspace is loaded into memory - self._ensure_workspace_loaded(workspace_id, **kwargs) - # Use parent's search implementation - return super().search(query, workspace_id, top_k, filter_dict, **kwargs) - - async def async_search( - self, - query: str, - workspace_id: str, - top_k: int = 1, - filter_dict: Optional[Dict[str, Any]] = None, - **kwargs, - ) -> List[VectorNode]: - """ - Async version of search. - - Args: - query: Text query to search for. If empty/None, performs filter-only - search without vector similarity. - workspace_id: Identifier of the workspace to search in. - top_k: Number of top results to return. Defaults to 1. - filter_dict: Optional dictionary of filters to apply to nodes. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - List[VectorNode]: List of matching nodes sorted by similarity score. - """ - # Ensure workspace is loaded into memory - await self._async_ensure_workspace_loaded(workspace_id, **kwargs) - # Use parent's async search implementation - return await super().async_search(query, workspace_id, top_k, filter_dict, **kwargs) - - # ==================== Insert Methods ==================== - - def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs): - """ - Insert or update nodes in a workspace. - - Nodes are inserted into memory and then persisted to disk. - - Args: - nodes: Single VectorNode or list of VectorNode instances to insert/update. - workspace_id: Identifier of the workspace. - **kwargs: Additional keyword arguments to pass to operations. - """ - # Ensure workspace is loaded into memory first - self._ensure_workspace_loaded(workspace_id, **kwargs) - # Insert into memory - super().insert(nodes, workspace_id, **kwargs) - # Persist to disk - self._persist_workspace(workspace_id, **kwargs) - - async def async_insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs): - """ - Async version of insert. - - Args: - nodes: Single VectorNode or list of VectorNode instances to insert/update. - workspace_id: Identifier of the workspace. - **kwargs: Additional keyword arguments to pass to operations. - """ - # Ensure workspace is loaded into memory first - await self._async_ensure_workspace_loaded(workspace_id, **kwargs) - # Insert into memory using async embedding - await super().async_insert(nodes, workspace_id, **kwargs) - # Persist to disk - await self._async_persist_workspace(workspace_id, **kwargs) - - # ==================== Delete Methods ==================== - - def delete(self, node_ids: str | List[str], workspace_id: str, **kwargs): - """ - Delete nodes from a workspace by their unique IDs. - - Args: - node_ids: Single unique_id string or list of unique_id strings to delete. - workspace_id: Identifier of the workspace. - **kwargs: Additional keyword arguments to pass to operations. - """ - if not self.exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return - - # Ensure workspace is loaded into memory - self._ensure_workspace_loaded(workspace_id, **kwargs) - # Delete from memory - super().delete(node_ids, workspace_id, **kwargs) - # Persist to disk - self._persist_workspace(workspace_id, **kwargs) - - async def async_delete(self, node_ids: str | List[str], workspace_id: str, **kwargs): - """ - Async version of delete. - - Args: - node_ids: Single unique_id string or list of unique_id strings to delete. - workspace_id: Identifier of the workspace. - **kwargs: Additional keyword arguments to pass to operations. - """ - if not self.exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return - - # Ensure workspace is loaded into memory - await self._async_ensure_workspace_loaded(workspace_id, **kwargs) - # Delete from memory - super().delete(node_ids, workspace_id, **kwargs) - # Persist to disk - await self._async_persist_workspace(workspace_id, **kwargs) diff --git a/flowllm/core/vector_store/memory_vector_store.py b/flowllm/core/vector_store/memory_vector_store.py deleted file mode 100644 index 4a746c3..0000000 --- a/flowllm/core/vector_store/memory_vector_store.py +++ /dev/null @@ -1,833 +0,0 @@ -""" -In-memory vector store implementation. - -This module provides an in-memory vector store that keeps all data in memory -for fast access. It inherits from LocalVectorStore for utility methods and can -persist data to disk when needed via dump_workspace and load_workspace methods. -""" - -import json -import math -from pathlib import Path -from typing import List, Dict, Optional, Any, Iterable - -from loguru import logger -from tqdm import tqdm - -from .base_vector_store import BaseVectorStore -from ..context import C -from ..schema import VectorNode - -# fcntl is Unix/Linux specific, not available on Windows -try: - import fcntl - - HAS_FCNTL = True -except ImportError: - fcntl = None - HAS_FCNTL = False - logger.warning("fcntl module not available (Windows). File locking will be disabled.") - - -def _acquire_lock(file_obj, lock_type): - """Acquire file lock if fcntl is available.""" - if HAS_FCNTL: - fcntl.flock(file_obj, lock_type) - - -def _release_lock(file_obj): - """Release file lock if fcntl is available.""" - if HAS_FCNTL: - fcntl.flock(file_obj, fcntl.LOCK_UN) - - -@C.register_vector_store("memory") -class MemoryVectorStore(BaseVectorStore): - """ - In-memory vector store that keeps all data in memory for fast access. - - This vector store keeps all data in memory and only persists to disk when - dump_workspace is called. It can load previously saved data via load_workspace. - Provides both synchronous and asynchronous APIs for all operations. - """ - - # ==================== Initialization ==================== - - def __init__(self, **kwargs): - """ - Initialize the memory vector store. - - Args: - **kwargs: Keyword arguments passed to the parent BaseVectorStore class. - """ - super().__init__(**kwargs) - self._memory_store: Dict[str, Dict[str, VectorNode]] = {} - - # ==================== Static Helper Methods ==================== - - @staticmethod - def _load_from_path( - path: str | Path, - workspace_id: str, - callback_fn=None, - show_progress: bool = True, - **kwargs, - ) -> List[VectorNode]: - """ - Load vector nodes from a JSONL file on disk. - - Args: - path: Directory path containing the workspace file. - workspace_id: Identifier for the workspace to load. - callback_fn: Optional callback function to transform node dictionaries - before creating VectorNode instances. - show_progress: Whether to show tqdm progress bar. Defaults to True. - **kwargs: Additional keyword arguments to pass to VectorNode constructor. - - Returns: - List[VectorNode]: List of loaded vector nodes. - - Note: - This method uses file locking (shared lock) when available to ensure - thread-safe reads. On Windows, file locking is disabled. - """ - workspace_path = Path(path) / f"{workspace_id}.jsonl" - if not workspace_path.exists(): - logger.warning(f"workspace_path={workspace_path} does not exist!") - return [] - - nodes = [] - with workspace_path.open() as f: - _acquire_lock(f, fcntl.LOCK_SH if HAS_FCNTL else None) - try: - lines = tqdm(f, desc="load from path") if show_progress else f - for line in lines: - if line.strip(): - node_dict = json.loads(line.strip()) - if callback_fn: - node: VectorNode = callback_fn(node_dict) - assert isinstance(node, VectorNode) - else: - node: VectorNode = VectorNode(**node_dict, **kwargs) - node.workspace_id = workspace_id - nodes.append(node) - - finally: - _release_lock(f) - return nodes - - @staticmethod - def _dump_to_path( - nodes: Iterable[VectorNode], - workspace_id: str, - path: str | Path = "", - callback_fn=None, - ensure_ascii: bool = False, - show_progress: bool = True, - **kwargs, - ): - """ - Write vector nodes to a JSONL file on disk. - - Args: - nodes: Iterable of VectorNode instances to write. - workspace_id: Identifier for the workspace. - path: Directory path where the workspace file should be written. - callback_fn: Optional callback function to transform VectorNode instances - before serialization. - ensure_ascii: If True, ensure all non-ASCII characters are escaped. - Defaults to False. - show_progress: Whether to show tqdm progress bar. Defaults to True. - **kwargs: Additional keyword arguments to pass to json.dumps. - - Returns: - dict: Dictionary with "size" key indicating number of nodes written. - - Note: - This method uses file locking (exclusive lock) when available to ensure - thread-safe writes. On Windows, file locking is disabled. - """ - dump_path: Path = Path(path) - dump_path.mkdir(parents=True, exist_ok=True) - dump_file = dump_path / f"{workspace_id}.jsonl" - - count = 0 - with dump_file.open("w") as f: - _acquire_lock(f, fcntl.LOCK_EX if HAS_FCNTL else None) - try: - nodes_iter = tqdm(nodes, desc="dump to path") if show_progress else nodes - for node in nodes_iter: - node.workspace_id = workspace_id - if callback_fn: - node_dict = callback_fn(node) - assert isinstance(node_dict, dict) - else: - node_dict = node.model_dump() - - f.write(json.dumps(node_dict, ensure_ascii=ensure_ascii, **kwargs)) - f.write("\n") - count += 1 - - return {"size": count} - finally: - _release_lock(f) - - @staticmethod - def _matches_filters(node: VectorNode, filter_dict: dict = None) -> bool: - """ - Check if a node matches all filters in filter_dict. - - Supports both term filters (exact value match) and range filters - (gte, lte, gt, lt). Keys can access both top-level VectorNode fields - (unique_id, workspace_id, content) and nested metadata using dot notation - (e.g., "metadata.node_type"). - - Args: - node: VectorNode instance to check. - filter_dict: Dictionary of filters to apply. Can contain: - - Term filters: {"key": value} for exact matches - - Range filters: {"key": {"gte": min, "lte": max}} for ranges - - Top-level keys: {"unique_id": value, "content": value} - - Nested keys: {"metadata.node_type": value} - - Returns: - bool: True if node matches all filters, False otherwise. - Returns True if filter_dict is None or empty. - """ - if not filter_dict: - return True - - # Convert node to dict for unified access to all fields - node_dict = node.model_dump() - - for key, filter_value in filter_dict.items(): - # Navigate nested keys (e.g., "metadata.node_type") - value = node_dict - key_found = True - for key_part in key.split("."): - if isinstance(value, dict) and key_part in value: - value = value[key_part] - else: - key_found = False - break - - if not key_found: - return False - - # Handle different filter types - if isinstance(filter_value, dict): - # Range filter: {"gte": 1, "lte": 10} - range_match = True - if "gte" in filter_value and value < filter_value["gte"]: - range_match = False - elif "lte" in filter_value and value > filter_value["lte"]: - range_match = False - elif "gt" in filter_value and value <= filter_value["gt"]: - range_match = False - elif "lt" in filter_value and value >= filter_value["lt"]: - range_match = False - if not range_match: - return False - elif isinstance(filter_value, list): - # List filter: value must match any item in the list (OR logic) - if value not in filter_value: - return False - else: - # Term filter: direct value comparison - if value != filter_value: - return False - - return True - - @staticmethod - def calculate_similarity(query_vector: List[float], node_vector: List[float]): - """ - Calculate cosine similarity between two vectors. - - Args: - query_vector: Query embedding vector. - node_vector: Node embedding vector. - - Returns: - float: Cosine similarity score between -1 and 1 (typically 0-1 for normalized vectors). - - Raises: - AssertionError: If vectors are empty or have different dimensions. - """ - assert query_vector, "query_vector is empty!" - assert node_vector, "node_vector is empty!" - assert len(query_vector) == len( - node_vector, - ), f"query_vector.size={len(query_vector)} node_vector.size={len(node_vector)}" - - dot_product = sum(x * y for x, y in zip(query_vector, node_vector)) - norm_v1 = math.sqrt(sum(x**2 for x in query_vector)) - norm_v2 = math.sqrt(sum(y**2 for y in node_vector)) - return dot_product / (norm_v1 * norm_v2) - - # ==================== Workspace Management Methods ==================== - - def exist_workspace(self, workspace_id: str, **kwargs) -> bool: - """ - Check if a workspace exists in memory. - - Args: - workspace_id: Identifier of the workspace to check. - **kwargs: Additional keyword arguments (unused). - - Returns: - bool: True if the workspace exists in memory, False otherwise. - """ - return workspace_id in self._memory_store - - async def async_exist_workspace(self, workspace_id: str, **kwargs) -> bool: - """ - Async version of exist_workspace. - - Args: - workspace_id: Identifier of the workspace to check. - **kwargs: Additional keyword arguments (unused). - - Returns: - bool: True if the workspace exists in memory, False otherwise. - """ - return self.exist_workspace(workspace_id, **kwargs) - - def delete_workspace(self, workspace_id: str, **kwargs): - """ - Delete a workspace and all its nodes from memory. - - Args: - workspace_id: Identifier of the workspace to delete. - **kwargs: Additional keyword arguments (unused). - """ - if workspace_id in self._memory_store: - del self._memory_store[workspace_id] - logger.info(f"Deleted workspace_id={workspace_id} from memory") - - async def async_delete_workspace(self, workspace_id: str, **kwargs): - """ - Async version of delete_workspace. - - Args: - workspace_id: Identifier of the workspace to delete. - **kwargs: Additional keyword arguments (unused). - """ - return self.delete_workspace(workspace_id, **kwargs) - - def create_workspace(self, workspace_id: str, **kwargs): - """ - Create a new empty workspace in memory. - - Args: - workspace_id: Identifier for the new workspace. - **kwargs: Additional keyword arguments (unused). - """ - if workspace_id not in self._memory_store: - self._memory_store[workspace_id] = {} - logger.info(f"Created workspace_id={workspace_id} in memory") - - async def async_create_workspace(self, workspace_id: str, **kwargs): - """ - Async version of create_workspace. - - Args: - workspace_id: Identifier for the new workspace. - **kwargs: Additional keyword arguments (unused). - """ - return self.create_workspace(workspace_id, **kwargs) - - def list_workspace_nodes(self, workspace_id: str, **kwargs) -> List[VectorNode]: - """ - List all nodes in a workspace. - - Args: - workspace_id: Identifier of the workspace. - **kwargs: Additional keyword arguments (unused). - - Returns: - List[VectorNode]: All nodes from the workspace, or empty list if workspace doesn't exist. - """ - if workspace_id in self._memory_store: - return list(self._memory_store[workspace_id].values()) - return [] - - async def async_list_workspace_nodes(self, workspace_id: str, **kwargs) -> List[VectorNode]: - """ - Async version of list_workspace_nodes. - - Args: - workspace_id: Identifier of the workspace. - **kwargs: Additional keyword arguments (unused). - - Returns: - List[VectorNode]: All nodes from the workspace, or empty list if workspace doesn't exist. - """ - return self.list_workspace_nodes(workspace_id, **kwargs) - - def dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs): - """ - Export a workspace from memory to disk at the specified path. - - Args: - workspace_id: Identifier of the workspace to export. - path: Directory path where to write the exported workspace file. - callback_fn: Optional callback function to transform nodes during export. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - - Returns: - dict: Dictionary with "size" key indicating number of nodes exported, - or empty dict if workspace doesn't exist in memory or path is empty. - """ - if workspace_id not in self._memory_store: - logger.warning(f"workspace_id={workspace_id} not found in memory!") - return {} - - if not path: - logger.warning("path is empty, cannot dump workspace!") - return {} - - nodes = list(self._memory_store[workspace_id].values()) - - return self._dump_to_path( - nodes=nodes, - workspace_id=workspace_id, - path=path, - callback_fn=callback_fn, - **kwargs, - ) - - async def async_dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs): - """ - Async version of dump_workspace. - - This method performs the same dump operation as dump_workspace(), but runs - the file I/O operations in a thread pool for better performance in async contexts. - - Args: - workspace_id: Identifier of the workspace to export. - path: Directory path where to write the exported workspace file. - callback_fn: Optional callback function to transform nodes during export. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - - Returns: - dict: Dictionary with "size" key indicating number of nodes exported, - or empty dict if workspace doesn't exist in memory or path is empty. - """ - return await self._run_sync_in_executor( - self.dump_workspace, - workspace_id, - path, - callback_fn, - **kwargs, - ) - - def load_workspace( - self, - workspace_id: str, - path: str | Path = "", - nodes: Optional[List[VectorNode]] = None, - callback_fn=None, - **kwargs, - ): - """ - Load a workspace into memory from disk, optionally merging with provided nodes. - - This method replaces any existing workspace with the same ID in memory, - then loads nodes from the specified path and/or the provided nodes list. - - Args: - workspace_id: Identifier for the workspace to create/load. - path: Directory path containing the workspace file to load. - If empty, only loads from nodes parameter. - nodes: Optional list of VectorNode instances to merge with loaded nodes. - callback_fn: Optional callback function to transform node dictionaries. - **kwargs: Additional keyword arguments to pass to load operations. - - Returns: - dict: Dictionary with "size" key indicating total number of nodes loaded. - """ - if workspace_id in self._memory_store: - del self._memory_store[workspace_id] - logger.info(f"Cleared existing workspace_id={workspace_id} from memory") - - self.create_workspace(workspace_id=workspace_id, **kwargs) - - all_nodes: List[VectorNode] = [] - - if nodes: - all_nodes.extend(nodes) - - if path: - all_nodes.extend( - self._load_from_path(path=path, workspace_id=workspace_id, callback_fn=callback_fn, **kwargs), - ) - - if all_nodes: - self.insert(nodes=all_nodes, workspace_id=workspace_id, **kwargs) - - logger.info(f"Loaded workspace_id={workspace_id} with {len(all_nodes)} nodes into memory") - return {"size": len(all_nodes)} - - async def async_load_workspace( - self, - workspace_id: str, - path: str | Path = "", - nodes: Optional[List[VectorNode]] = None, - callback_fn=None, - **kwargs, - ): - """ - Async version of load_workspace. - - This method performs the same load operation as load_workspace(), but runs - the file I/O operations in a thread pool for better performance in async contexts. - - Args: - workspace_id: Identifier for the workspace to create/load. - path: Directory path containing the workspace file to load. - If empty, only loads from nodes parameter. - nodes: Optional list of VectorNode instances to merge with loaded nodes. - callback_fn: Optional callback function to transform node dictionaries. - **kwargs: Additional keyword arguments to pass to load operations. - - Returns: - dict: Dictionary with "size" key indicating total number of nodes loaded. - """ - return await self._run_sync_in_executor( - self.load_workspace, - workspace_id, - path, - nodes, - callback_fn, - **kwargs, - ) - - def copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs): - """ - Copy all nodes from one workspace to another in memory. - - Args: - src_workspace_id: Identifier of the source workspace. - dest_workspace_id: Identifier of the destination workspace. - Created if it doesn't exist. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - dict: Dictionary with "size" key indicating number of nodes copied, - or empty dict if source workspace doesn't exist in memory. - """ - if src_workspace_id not in self._memory_store: - logger.warning(f"src_workspace_id={src_workspace_id} not found in memory!") - return {} - - if dest_workspace_id not in self._memory_store: - self.create_workspace(workspace_id=dest_workspace_id, **kwargs) - - src_nodes = list(self._memory_store[src_workspace_id].values()) - node_size = len(src_nodes) - - new_nodes = [] - for node in src_nodes: - new_node = VectorNode(**node.model_dump()) - new_node.workspace_id = dest_workspace_id - new_nodes.append(new_node) - - self.insert(nodes=new_nodes, workspace_id=dest_workspace_id, **kwargs) - - logger.info(f"Copied {node_size} nodes from {src_workspace_id} to {dest_workspace_id}") - return {"size": node_size} - - async def async_copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs): - """ - Async version of copy_workspace. - - Args: - src_workspace_id: Identifier of the source workspace. - dest_workspace_id: Identifier of the destination workspace. - Created if it doesn't exist. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - dict: Dictionary with "size" key indicating number of nodes copied, - or empty dict if source workspace doesn't exist in memory. - """ - return self.copy_workspace(src_workspace_id, dest_workspace_id, **kwargs) - - def list_workspace(self, **kwargs) -> List[str]: - """ - List all existing workspaces in memory. - - Returns: - List[str]: Workspace identifiers currently present in memory. - """ - return list(self._memory_store.keys()) - - async def async_list_workspace(self, **kwargs) -> List[str]: - """ - Async version of list_workspace. - - Returns: - List[str]: Workspace identifiers currently present in memory. - """ - return self.list_workspace(**kwargs) - - # ==================== Search Methods ==================== - - def search( - self, - query: str, - workspace_id: str, - top_k: int = 1, - filter_dict: Optional[Dict[str, Any]] = None, - **kwargs, - ) -> List[VectorNode]: - """ - Search for similar nodes using vector similarity in memory. - - When query is empty (empty string or None), the search degrades to a - filter-only search without vector similarity ranking. - - Args: - query: Text query to search for. If empty/None, performs filter-only - search without vector similarity. - workspace_id: Identifier of the workspace to search in. - top_k: Number of top results to return. Defaults to 1. - filter_dict: Optional dictionary of filters to apply to nodes. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - List[VectorNode]: List of matching nodes. When query is provided, - nodes are sorted by similarity score (highest first) with a - "score" key in metadata. When query is empty, nodes are returned - in storage order without scores. - """ - if workspace_id not in self._memory_store: - logger.warning(f"workspace_id={workspace_id} not found in memory!") - return [] - - # When query is empty, degrade to filter-only search without vector similarity - use_vector_search = bool(query) - query_vector = self.get_embeddings(query) if use_vector_search else None - - nodes: List[VectorNode] = [] - - for node in self._memory_store[workspace_id].values(): - # Apply filters - if self._matches_filters(node, filter_dict): - if use_vector_search: - if node.vector: - score = self.calculate_similarity(query_vector, node.vector) - result_node = VectorNode(**node.model_dump()) - result_node.metadata["score"] = score - nodes.append(result_node) - else: - nodes.append(node) - # Early exit for filter-only search when we have enough results - if len(nodes) >= top_k: - break - - # Only sort by score when using vector search - if use_vector_search: - nodes = sorted(nodes, key=lambda x: x.metadata["score"], reverse=True) - - return nodes[:top_k] - - async def async_search( - self, - query: str, - workspace_id: str, - top_k: int = 1, - filter_dict: Optional[Dict[str, Any]] = None, - **kwargs, - ) -> List[VectorNode]: - """ - Async version of search using embedding model async capabilities. - - This method performs the same search operation as search(), but uses - async embedding generation for better performance in async contexts. - - When query is empty (empty string or None), the search degrades to a - filter-only search without vector similarity ranking. - - Args: - query: Text query to search for. If empty/None, performs filter-only - search without vector similarity. - workspace_id: Identifier of the workspace to search in. - top_k: Number of top results to return. Defaults to 1. - filter_dict: Optional dictionary of filters to apply to nodes. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - List[VectorNode]: List of matching nodes. When query is provided, - nodes are sorted by similarity score (highest first) with a - "score" key in metadata. When query is empty, nodes are returned - in storage order without scores. - """ - if workspace_id not in self._memory_store: - logger.warning(f"workspace_id={workspace_id} not found in memory!") - return [] - - # When query is empty, degrade to filter-only search without vector similarity - use_vector_search = bool(query) - query_vector = await self.async_get_embeddings(query) if use_vector_search else None - - nodes: List[VectorNode] = [] - - for node in self._memory_store[workspace_id].values(): - # Apply filters - if self._matches_filters(node, filter_dict): - if use_vector_search: - if node.vector: - score = self.calculate_similarity(query_vector, node.vector) - result_node = VectorNode(**node.model_dump()) - result_node.metadata["score"] = score - nodes.append(result_node) - else: - nodes.append(node) - # Early exit for filter-only search when we have enough results - if len(nodes) >= top_k: - break - - # Only sort by score when using vector search - if use_vector_search: - nodes = sorted(nodes, key=lambda x: x.metadata["score"], reverse=True) - - return nodes[:top_k] - - # ==================== Insert Methods ==================== - - def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs): - """ - Insert or update nodes in a workspace in memory. - - If a node with the same unique_id already exists, it will be updated. - New nodes will be added. All nodes are embedded before insertion. - Workspace is created automatically if it doesn't exist. - - Args: - nodes: Single VectorNode or list of VectorNode instances to insert/update. - workspace_id: Identifier of the workspace. - **kwargs: Additional keyword arguments to pass to operations. - """ - if isinstance(nodes, VectorNode): - nodes = [nodes] - - if workspace_id not in self._memory_store: - self.create_workspace(workspace_id=workspace_id, **kwargs) - - nodes: List[VectorNode] = self.get_node_embeddings(nodes) - - update_cnt = 0 - for node in nodes: - if node.unique_id in self._memory_store[workspace_id]: - update_cnt += 1 - - node.workspace_id = workspace_id - self._memory_store[workspace_id][node.unique_id] = node - - total_nodes = len(self._memory_store[workspace_id]) - logger.info( - f"Inserted into workspace_id={workspace_id} nodes.size={len(nodes)} " - f"total.size={total_nodes} update_cnt={update_cnt}", - ) - - async def async_insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs): - """ - Async version of insert using embedding model async capabilities. - - This method performs the same insert operation as insert(), but uses - async embedding generation for better performance in async contexts. - - Args: - nodes: Single VectorNode or list of VectorNode instances to insert/update. - workspace_id: Identifier of the workspace. - **kwargs: Additional keyword arguments to pass to operations. - """ - if isinstance(nodes, VectorNode): - nodes = [nodes] - - # Ensure workspace exists - if workspace_id not in self._memory_store: - self.create_workspace(workspace_id=workspace_id, **kwargs) - - # Use async embedding - nodes = await self.async_get_node_embeddings(nodes) - - update_cnt = 0 - for node in nodes: - if node.unique_id in self._memory_store[workspace_id]: - update_cnt += 1 - - node.workspace_id = workspace_id - self._memory_store[workspace_id][node.unique_id] = node - - total_nodes = len(self._memory_store[workspace_id]) - logger.info( - f"Async inserted into workspace_id={workspace_id} nodes.size={len(nodes)} " - f"total.size={total_nodes} update_cnt={update_cnt}", - ) - - # ==================== Delete Methods ==================== - - def delete(self, node_ids: str | List[str], workspace_id: str, **kwargs): - """ - Delete nodes from a workspace by their unique IDs in memory. - - Args: - node_ids: Single unique_id string or list of unique_id strings to delete. - workspace_id: Identifier of the workspace. - **kwargs: Additional keyword arguments to pass to operations. - """ - if workspace_id not in self._memory_store: - logger.warning(f"workspace_id={workspace_id} not found in memory!") - return - - if isinstance(node_ids, str): - node_ids = [node_ids] - - before_size = len(self._memory_store[workspace_id]) - deleted_cnt = 0 - - for node_id in node_ids: - if node_id in self._memory_store[workspace_id]: - del self._memory_store[workspace_id][node_id] - deleted_cnt += 1 - - after_size = len(self._memory_store[workspace_id]) - logger.info( - f"Deleted from workspace_id={workspace_id} before_size={before_size} " - f"after_size={after_size} deleted_cnt={deleted_cnt}", - ) - - async def async_delete(self, node_ids: str | List[str], workspace_id: str, **kwargs): - """ - Async version of delete. - - This method performs the same delete operation as delete() but is provided - for consistency in async contexts. - - Args: - node_ids: Single unique_id string or list of unique_id strings to delete. - workspace_id: Identifier of the workspace. - **kwargs: Additional keyword arguments to pass to operations. - """ - return self.delete(node_ids, workspace_id, **kwargs) - - # ==================== Close Methods ==================== - - def close(self): - """ - Close the vector store and clean up resources. - - For MemoryVectorStore, this clears all data from memory. - """ - self._memory_store.clear() - - async def async_close(self): - """ - Async version of close. - - For MemoryVectorStore, this clears all data from memory. - """ - self.close() diff --git a/flowllm/core/vector_store/pgvector_vector_store.py b/flowllm/core/vector_store/pgvector_vector_store.py deleted file mode 100644 index fa6b30b..0000000 --- a/flowllm/core/vector_store/pgvector_vector_store.py +++ /dev/null @@ -1,1246 +0,0 @@ -"""PostgreSQL pgvector vector store implementation. - -This module provides a PostgreSQL-based vector store that stores vector nodes -in PostgreSQL tables using the pgvector extension. It supports workspace management, -vector similarity search, metadata filtering using SQL WHERE clauses, and provides -both synchronous and asynchronous operations using psycopg and asyncpg. -""" - -# pylint: disable=too-many-lines - -import json -import os -import re -from pathlib import Path -from typing import List, Tuple, Dict, Any, Optional - -from loguru import logger - -from .memory_vector_store import MemoryVectorStore -from ..context import C -from ..schema import VectorNode - - -@C.register_vector_store("pgvector") -class PgVectorStore(MemoryVectorStore): - """PostgreSQL pgvector vector store implementation. - - This class provides a vector store backend using PostgreSQL with the pgvector - extension for storing and searching vector embeddings. It supports both synchronous - and asynchronous operations using psycopg and asyncpg, and includes metadata - filtering capabilities using SQL WHERE clauses. - - Attributes: - connection_string: PostgreSQL connection string for synchronous operations. - Defaults to the FLOW_PGVECTOR_CONNECTION_STRING environment variable - or "postgresql://localhost/postgres". - async_connection_string: PostgreSQL connection string for asynchronous operations. - Defaults to the FLOW_PGVECTOR_ASYNC_CONNECTION_STRING environment variable - or None (will use connection_string with asyncpg). - batch_size: Batch size for bulk operations. Defaults to 1024. - """ - - # ==================== Initialization ==================== - - def __init__( - self, - connection_string: str | None = None, - async_connection_string: str | None = None, - batch_size: int = 1024, - **kwargs, - ): - """Initialize PostgreSQL connections. - - Args: - connection_string: PostgreSQL connection string for synchronous operations. - async_connection_string: PostgreSQL connection string for asynchronous operations. - batch_size: Batch size for bulk operations. - **kwargs: Additional keyword arguments passed to MemoryVectorStore. - """ - super().__init__(**kwargs) - self.connection_string = connection_string or os.getenv( - "FLOW_PGVECTOR_CONNECTION_STRING", - "postgresql://localhost/postgres", - ) - self.async_connection_string = async_connection_string or os.getenv( - "FLOW_PGVECTOR_ASYNC_CONNECTION_STRING", - ) - self.batch_size = batch_size - - # Initialize synchronous connection - import psycopg - - self._conn = psycopg.connect(self.connection_string) - self._conn.autocommit = False - - # Initialize async connection (created lazily in async methods) - self._async_conn = None - # Use async_connection_string if provided, otherwise use connection_string - # asyncpg accepts standard postgresql:// connection strings directly - self._async_conn_string = self.async_connection_string or self.connection_string - - # Ensure pgvector extension exists - with self._conn.cursor() as cur: - cur.execute("CREATE EXTENSION IF NOT EXISTS vector") - self._conn.commit() - - logger.info( - f"PostgreSQL pgvector client initialized with connection_string={self.connection_string}", - ) - - # ==================== Static Helper Methods ==================== - - @staticmethod - def _get_table_name(workspace_id: str) -> str: - """Get the table name for a workspace. - - Args: - workspace_id: The workspace identifier. - - Returns: - str: The table name (sanitized workspace_id). - """ - # Sanitize workspace_id to be a valid PostgreSQL identifier - # Replace non-alphanumeric characters with underscores - sanitized = re.sub(r"[^a-zA-Z0-9_]", "_", workspace_id) - return f"workspace_{sanitized}" - - @staticmethod - def _row2node(row: Tuple, workspace_id: str) -> VectorNode: - """Convert a PostgreSQL row to a VectorNode. - - Args: - row: The PostgreSQL row tuple (unique_id, workspace_id, content, metadata, vector::text). - workspace_id: The workspace identifier to assign to the node. - - Returns: - VectorNode: A VectorNode instance created from the row data. - """ - unique_id, workspace_id_col, content, metadata, vector_str = row - # Parse vector string (format: [0.1,0.2,0.3] from pgvector) - # pgvector returns vector as string like '[0.1,0.2,0.3]' - vector = json.loads(vector_str) - - # Parse metadata if it's a string (psycopg may return JSONB as string in some cases) - if isinstance(metadata, str): - metadata = json.loads(metadata) - elif metadata is None: - metadata = {} - - node = VectorNode( - unique_id=unique_id, - workspace_id=workspace_id_col or workspace_id, - content=content, - metadata=metadata, - vector=vector, - ) - return node - - @staticmethod - def _build_sql_filters( - filter_dict: Optional[Dict[str, Any]] = None, - use_async: bool = False, - ) -> Tuple[str, List[Any]]: - """Build SQL WHERE clause from filter_dict. - - Converts a filter dictionary into SQL WHERE conditions. - Supports both term filters (exact match) and range filters (gte, lte, gt, lt). - - Args: - filter_dict: Dictionary of filter conditions. Keys are metadata field names, - values can be exact match values or range dictionaries like - {"gte": 1, "lte": 10}. - - Keys starting with "metadata." will have the prefix stripped - - "unique_id" is treated as a direct column reference - use_async: If True, use asyncpg-style placeholders ($1, $2), else use psycopg-style (%s). - - Returns: - Tuple[str, List[Any]]: SQL WHERE clause string and list of parameters. - """ - if not filter_dict: - return "", [] - - conditions = [] - params = [] - param_idx = 1 - - for key, filter_value in filter_dict.items(): - # Handle special keys that are stored as direct columns - if key == "unique_id": - # unique_id is a direct column, not in metadata JSONB - if use_async: - conditions.append(f"unique_id = ${param_idx}") - else: - conditions.append("unique_id = %s") - params.append(str(filter_value)) - param_idx += 1 - continue - - # Strip "metadata." prefix if present (since we're already accessing metadata column) - if key.startswith("metadata."): - metadata_key = key[len("metadata.") :] - else: - metadata_key = key - - # Handle nested keys by using JSONB path operators - jsonb_path = f"metadata->>'{metadata_key}'" - - if isinstance(filter_value, dict): - # Range filter: {"gte": 1, "lte": 10} - range_conditions = [] - if "gte" in filter_value: - if use_async: - range_conditions.append(f"({jsonb_path}::numeric) >= ${param_idx}") - else: - range_conditions.append(f"({jsonb_path}::numeric) >= %s") - params.append(filter_value["gte"]) - param_idx += 1 - if "lte" in filter_value: - if use_async: - range_conditions.append(f"({jsonb_path}::numeric) <= ${param_idx}") - else: - range_conditions.append(f"({jsonb_path}::numeric) <= %s") - params.append(filter_value["lte"]) - param_idx += 1 - if "gt" in filter_value: - if use_async: - range_conditions.append(f"({jsonb_path}::numeric) > ${param_idx}") - else: - range_conditions.append(f"({jsonb_path}::numeric) > %s") - params.append(filter_value["gt"]) - param_idx += 1 - if "lt" in filter_value: - if use_async: - range_conditions.append(f"({jsonb_path}::numeric) < ${param_idx}") - else: - range_conditions.append(f"({jsonb_path}::numeric) < %s") - params.append(filter_value["lt"]) - param_idx += 1 - if range_conditions: - conditions.append(f"({' AND '.join(range_conditions)})") - elif isinstance(filter_value, list): - # List filter: use IN clause for OR logic - if use_async: - placeholders = ", ".join(f"${param_idx + i}" for i in range(len(filter_value))) - conditions.append(f"{jsonb_path} IN ({placeholders})") - else: - placeholders = ", ".join(["%s"] * len(filter_value)) - conditions.append(f"{jsonb_path} IN ({placeholders})") - params.extend([str(v) for v in filter_value]) - param_idx += len(filter_value) - else: - # Term filter: direct value comparison - if use_async: - conditions.append(f"{jsonb_path} = ${param_idx}") - else: - conditions.append(f"{jsonb_path} = %s") - params.append(str(filter_value)) - param_idx += 1 - - where_clause = " AND ".join(conditions) - return where_clause, params - - # ==================== Workspace Management Methods ==================== - - def exist_workspace(self, workspace_id: str, **kwargs) -> bool: - """Check if a PostgreSQL table (workspace) exists. - - Args: - workspace_id: The identifier of the workspace/table to check. - **kwargs: Additional keyword arguments (unused). - - Returns: - bool: True if the table exists, False otherwise. - """ - table_name = self._get_table_name(workspace_id) - with self._conn.cursor() as cur: - cur.execute( - """ - SELECT EXISTS ( - SELECT FROM information_schema.tables - WHERE table_schema = 'public' - AND table_name = %s - ) - """, - (table_name,), - ) - return cur.fetchone()[0] - - async def async_exist_workspace(self, workspace_id: str, **kwargs) -> bool: - """Check if a PostgreSQL table (workspace) exists (async). - - Args: - workspace_id: The identifier of the workspace/table to check. - **kwargs: Additional keyword arguments (unused). - - Returns: - bool: True if the table exists, False otherwise. - """ - table_name = self._get_table_name(workspace_id) - conn = await self._get_async_conn() - result = await conn.fetchval( - """ - SELECT EXISTS ( - SELECT FROM information_schema.tables - WHERE table_schema = 'public' - AND table_name = $1 - ) - """, - table_name, - ) - return result - - def delete_workspace(self, workspace_id: str, **kwargs): - """Delete a PostgreSQL table (workspace). - - Args: - workspace_id: The identifier of the workspace/table to delete. - **kwargs: Additional keyword arguments (unused). - """ - table_name = self._get_table_name(workspace_id) - with self._conn.cursor() as cur: - cur.execute(f'DROP TABLE IF EXISTS "{table_name}" CASCADE') - self._conn.commit() - logger.info(f"Deleted workspace table: {table_name}") - - async def async_delete_workspace(self, workspace_id: str, **kwargs): - """Delete a PostgreSQL table (workspace) (async). - - Args: - workspace_id: The identifier of the workspace/table to delete. - **kwargs: Additional keyword arguments (unused). - """ - table_name = self._get_table_name(workspace_id) - logger.debug(f"Attempting to delete workspace table: {table_name}") - conn = await self._get_async_conn() - logger.debug(f"Connection established, executing DROP TABLE for {table_name}") - await conn.execute(f'DROP TABLE IF EXISTS "{table_name}" CASCADE') - logger.info(f"Successfully deleted workspace table: {table_name}") - - def create_workspace(self, workspace_id: str, **kwargs): - """Create a new PostgreSQL table (workspace) with vector field. - - Args: - workspace_id: The identifier of the workspace/table to create. - **kwargs: Additional keyword arguments (unused). - - Returns: - The response from PostgreSQL create table statement. - """ - table_name = self._get_table_name(workspace_id) - dimensions = self.embedding_model.dimensions - - with self._conn.cursor() as cur: - cur.execute( - f""" - CREATE TABLE IF NOT EXISTS "{table_name}" ( - unique_id TEXT PRIMARY KEY, - workspace_id TEXT NOT NULL, - content TEXT NOT NULL, - metadata JSONB NOT NULL, - vector vector({dimensions}) NOT NULL - ) - """, - ) - # Create index for vector similarity search using HNSW - # HNSW is preferred over IVFFlat as it works with any dataset size - # and provides better recall with similar performance - cur.execute( - f""" - CREATE INDEX IF NOT EXISTS "{table_name}_vector_idx" - ON "{table_name}" USING hnsw (vector vector_cosine_ops) - """, - ) - # Create index for metadata filtering - cur.execute( - f""" - CREATE INDEX IF NOT EXISTS "{table_name}_metadata_idx" - ON "{table_name}" USING gin (metadata) - """, - ) - self._conn.commit() - logger.info(f"Created workspace table: {table_name} with vector({dimensions})") - - async def async_create_workspace(self, workspace_id: str, **kwargs): - """Create a new PostgreSQL table (workspace) with vector field (async). - - Args: - workspace_id: The identifier of the workspace/table to create. - **kwargs: Additional keyword arguments (unused). - - Returns: - The response from PostgreSQL create table statement. - """ - table_name = self._get_table_name(workspace_id) - dimensions = self.embedding_model.dimensions - - conn = await self._get_async_conn() - await conn.execute( - f""" - CREATE TABLE IF NOT EXISTS "{table_name}" ( - unique_id TEXT PRIMARY KEY, - workspace_id TEXT NOT NULL, - content TEXT NOT NULL, - metadata JSONB NOT NULL, - vector vector({dimensions}) NOT NULL - ) - """, - ) - # Create index for vector similarity search using HNSW - # HNSW is preferred over IVFFlat as it works with any dataset size - # and provides better recall with similar performance - await conn.execute( - f""" - CREATE INDEX IF NOT EXISTS "{table_name}_vector_idx" - ON "{table_name}" USING hnsw (vector vector_cosine_ops) - """, - ) - # Create index for metadata filtering - await conn.execute( - f""" - CREATE INDEX IF NOT EXISTS "{table_name}_metadata_idx" - ON "{table_name}" USING gin (metadata) - """, - ) - logger.info(f"Created workspace table: {table_name} with vector({dimensions})") - - def list_workspace_nodes( - self, - workspace_id: str, - max_size: int = 10000, - **kwargs, - ) -> List[VectorNode]: - """List all nodes in a workspace. - - Args: - workspace_id: The identifier of the workspace to iterate over. - max_size: Maximum number of nodes to retrieve (default: 10000). - **kwargs: Additional keyword arguments (unused). - - Returns: - List[VectorNode]: Vector nodes from the workspace. - """ - if not self.exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return [] - - table_name = self._get_table_name(workspace_id) - nodes: List[VectorNode] = [] - with self._conn.cursor() as cur: - cur.execute( - f""" - SELECT unique_id, workspace_id, content, metadata, vector::text - FROM "{table_name}" - LIMIT %s - """, - (max_size,), - ) - for row in cur.fetchall(): - node = self._row2node(row, workspace_id) - nodes.append(node) - return nodes - - async def async_list_workspace_nodes( - self, - workspace_id: str, - max_size: int = 10000, - **kwargs, - ) -> List[VectorNode]: - """List all nodes in a workspace (async). - - Args: - workspace_id: The identifier of the workspace to iterate over. - max_size: Maximum number of nodes to retrieve (default: 10000). - **kwargs: Additional keyword arguments (unused). - - Returns: - List[VectorNode]: Vector nodes from the workspace. - """ - if not await self.async_exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return [] - - table_name = self._get_table_name(workspace_id) - conn = await self._get_async_conn() - rows = await conn.fetch( - f""" - SELECT unique_id, workspace_id, content, metadata, vector::text - FROM "{table_name}" - LIMIT $1 - """, - max_size, - ) - - nodes: List[VectorNode] = [] - for row in rows: - unique_id = row["unique_id"] - workspace_id_col = row["workspace_id"] - content = row["content"] - metadata = row["metadata"] - vector_str = row["vector"] - - # Parse vector string (format: [0.1,0.2,0.3] from pgvector) - vector = json.loads(vector_str) - - # Parse metadata if it's a string - if isinstance(metadata, str): - metadata = json.loads(metadata) - elif metadata is None: - metadata = {} - - node = VectorNode( - unique_id=unique_id, - workspace_id=workspace_id_col or workspace_id, - content=content, - metadata=metadata, - vector=vector, - ) - nodes.append(node) - - return nodes - - def dump_workspace( - self, - workspace_id: str, - path: str | Path = "", - callback_fn=None, - **kwargs, - ): - """Export a workspace from PostgreSQL to disk at the specified path. - - Args: - workspace_id: Identifier of the workspace to export. - path: Directory path where to write the exported workspace file. - callback_fn: Optional callback function to transform nodes during export. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - - Returns: - dict: Dictionary with "size" key indicating number of nodes exported, - or empty dict if workspace doesn't exist or path is empty. - """ - if not self.exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return {} - - if not path: - logger.warning("path is empty, cannot dump workspace!") - return {} - - nodes = self.list_workspace_nodes(workspace_id=workspace_id, **kwargs) - - return self._dump_to_path( - nodes=nodes, - workspace_id=workspace_id, - path=path, - callback_fn=callback_fn, - **kwargs, - ) - - async def async_dump_workspace( - self, - workspace_id: str, - path: str | Path = "", - callback_fn=None, - **kwargs, - ): - """Async version of dump_workspace. - - Args: - workspace_id: Identifier of the workspace to export. - path: Directory path where to write the exported workspace file. - callback_fn: Optional callback function to transform nodes during export. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - - Returns: - dict: Dictionary with "size" key indicating number of nodes exported, - or empty dict if workspace doesn't exist or path is empty. - """ - if not await self.async_exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return {} - - if not path: - logger.warning("path is empty, cannot dump workspace!") - return {} - - nodes = await self.async_list_workspace_nodes(workspace_id=workspace_id, **kwargs) - - return await self._run_sync_in_executor( - self._dump_to_path, - nodes=nodes, - workspace_id=workspace_id, - path=path, - callback_fn=callback_fn, - **kwargs, - ) - - def load_workspace( - self, - workspace_id: str, - path: str | Path = "", - nodes: Optional[List[VectorNode]] = None, - callback_fn=None, - **kwargs, - ): - """Load a workspace into PostgreSQL from disk, optionally merging with provided nodes. - - This method replaces any existing workspace with the same ID, - then loads nodes from the specified path and/or the provided nodes list. - - Args: - workspace_id: Identifier for the workspace to create/load. - path: Directory path containing the workspace file to load. - If empty, only loads from nodes parameter. - nodes: Optional list of VectorNode instances to merge with loaded nodes. - callback_fn: Optional callback function to transform node dictionaries. - **kwargs: Additional keyword arguments to pass to load operations. - - Returns: - dict: Dictionary with "size" key indicating total number of nodes loaded. - """ - if self.exist_workspace(workspace_id=workspace_id): - self.delete_workspace(workspace_id=workspace_id) - logger.info(f"Deleted existing workspace_id={workspace_id}") - - self.create_workspace(workspace_id=workspace_id, **kwargs) - - all_nodes: List[VectorNode] = [] - - if nodes: - all_nodes.extend(nodes) - - if path: - all_nodes.extend( - self._load_from_path( - path=path, - workspace_id=workspace_id, - callback_fn=callback_fn, - **kwargs, - ), - ) - - if all_nodes: - self.insert(nodes=all_nodes, workspace_id=workspace_id, **kwargs) - - logger.info(f"Loaded workspace_id={workspace_id} with {len(all_nodes)} nodes into PostgreSQL") - return {"size": len(all_nodes)} - - async def async_load_workspace( - self, - workspace_id: str, - path: str | Path = "", - nodes: Optional[List[VectorNode]] = None, - callback_fn=None, - **kwargs, - ): - """Async version of load_workspace. - - Args: - workspace_id: Identifier for the workspace to create/load. - path: Directory path containing the workspace file to load. - If empty, only loads from nodes parameter. - nodes: Optional list of VectorNode instances to merge with loaded nodes. - callback_fn: Optional callback function to transform node dictionaries. - **kwargs: Additional keyword arguments to pass to load operations. - - Returns: - dict: Dictionary with "size" key indicating total number of nodes loaded. - """ - if await self.async_exist_workspace(workspace_id=workspace_id): - await self.async_delete_workspace(workspace_id=workspace_id) - logger.info(f"Deleted existing workspace_id={workspace_id}") - - await self.async_create_workspace(workspace_id=workspace_id, **kwargs) - - all_nodes: List[VectorNode] = [] - - if nodes: - all_nodes.extend(nodes) - - if path: - loaded_nodes = await self._run_sync_in_executor( - lambda: list( - self._load_from_path(path=path, workspace_id=workspace_id, callback_fn=callback_fn, **kwargs), - ), - ) - all_nodes.extend(loaded_nodes) - - if all_nodes: - await self.async_insert(nodes=all_nodes, workspace_id=workspace_id, **kwargs) - - logger.info(f"Loaded workspace_id={workspace_id} with {len(all_nodes)} nodes into PostgreSQL") - return {"size": len(all_nodes)} - - def copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs): - """Copy all nodes from one workspace to another in PostgreSQL. - - Args: - src_workspace_id: Identifier of the source workspace. - dest_workspace_id: Identifier of the destination workspace. - Created if it doesn't exist. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - dict: Dictionary with "size" key indicating number of nodes copied, - or empty dict if source workspace doesn't exist. - """ - if not self.exist_workspace(workspace_id=src_workspace_id): - logger.warning(f"src_workspace_id={src_workspace_id} does not exist!") - return {} - - if not self.exist_workspace(workspace_id=dest_workspace_id): - self.create_workspace(workspace_id=dest_workspace_id, **kwargs) - - src_nodes = self.list_workspace_nodes(workspace_id=src_workspace_id, **kwargs) - node_size = len(src_nodes) - - new_nodes = [] - for node in src_nodes: - new_node = VectorNode(**node.model_dump()) - new_node.workspace_id = dest_workspace_id - new_nodes.append(new_node) - - if new_nodes: - self.insert(nodes=new_nodes, workspace_id=dest_workspace_id, **kwargs) - - logger.info(f"Copied {node_size} nodes from {src_workspace_id} to {dest_workspace_id}") - return {"size": node_size} - - async def async_copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs): - """Async version of copy_workspace. - - Args: - src_workspace_id: Identifier of the source workspace. - dest_workspace_id: Identifier of the destination workspace. - Created if it doesn't exist. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - dict: Dictionary with "size" key indicating number of nodes copied, - or empty dict if source workspace doesn't exist. - """ - if not await self.async_exist_workspace(workspace_id=src_workspace_id): - logger.warning(f"src_workspace_id={src_workspace_id} does not exist!") - return {} - - if not await self.async_exist_workspace(workspace_id=dest_workspace_id): - await self.async_create_workspace(workspace_id=dest_workspace_id, **kwargs) - - src_nodes = await self.async_list_workspace_nodes(workspace_id=src_workspace_id, **kwargs) - node_size = len(src_nodes) - - new_nodes = [] - for node in src_nodes: - new_node = VectorNode(**node.model_dump()) - new_node.workspace_id = dest_workspace_id - new_nodes.append(new_node) - - if new_nodes: - await self.async_insert(nodes=new_nodes, workspace_id=dest_workspace_id, **kwargs) - - logger.info(f"Copied {node_size} nodes from {src_workspace_id} to {dest_workspace_id}") - return {"size": node_size} - - def list_workspace(self, **kwargs) -> List[str]: - """List all existing workspaces (tables) in PostgreSQL. - - Returns: - List[str]: Workspace identifiers (table names without prefix). - """ - with self._conn.cursor() as cur: - cur.execute( - """ - SELECT table_name - FROM information_schema.tables - WHERE table_schema = 'public' - AND table_name LIKE 'workspace_%' - ORDER BY table_name - """, - ) - table_names = [row[0] for row in cur.fetchall()] - # Remove 'workspace_' prefix - workspace_ids = [name.replace("workspace_", "", 1) for name in table_names] - return workspace_ids - - async def async_list_workspace(self, **kwargs) -> List[str]: - """Async version of list_workspace. - - Returns: - List[str]: Workspace identifiers (table names without prefix). - """ - conn = await self._get_async_conn() - rows = await conn.fetch( - """ - SELECT table_name - FROM information_schema.tables - WHERE table_schema = 'public' - AND table_name LIKE 'workspace_%' - ORDER BY table_name - """, - ) - table_names = [row["table_name"] for row in rows] - # Remove 'workspace_' prefix - workspace_ids = [name.replace("workspace_", "", 1) for name in table_names] - return workspace_ids - - # ==================== Search Methods ==================== - - def search( - self, - query: str, - workspace_id: str, - top_k: int = 1, - filter_dict: Optional[Dict[str, Any]] = None, - **kwargs, - ) -> List[VectorNode]: - """Search for similar vector nodes using cosine similarity. - - When query is empty (empty string or None), the search degrades to a - filter-only search without vector similarity ranking. - - Args: - query: The text query to search for. Will be embedded using the - embedding model. If empty/None, performs filter-only search. - workspace_id: The identifier of the workspace to search in. - top_k: Maximum number of results to return (default: 1). - filter_dict: Optional dictionary of metadata filters to apply. - **kwargs: Additional keyword arguments (unused). - - Returns: - List[VectorNode]: List of matching vector nodes sorted by similarity score. - """ - if not self.exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return [] - - # Build filters from filter_dict - where_clause, filter_params = self._build_sql_filters(filter_dict) - table_name = self._get_table_name(workspace_id) - where_sql = f"WHERE {where_clause}" if where_clause else "" - - # When query is empty, degrade to filter-only search without vector similarity - use_vector_search = bool(query) - - if use_vector_search: - query_vector = self.get_embeddings(query) - # Convert query_vector to string format for pgvector - query_vector_str = "[" + ",".join(str(v) for v in query_vector) + "]" - - with self._conn.cursor() as cur: - # Build SQL with correct parameter order matching SQL clause order: - # 1. %s::vector in SELECT (for score calculation) - # 2. filter_params in WHERE clause - # 3. %s::vector in ORDER BY - # 4. %s for LIMIT - cur.execute( - f""" - SELECT unique_id, workspace_id, content, metadata, vector::text, - 1 - (vector <=> %s::vector) AS score - FROM "{table_name}" - {where_sql} - ORDER BY vector <=> %s::vector - LIMIT %s - """, - [query_vector_str] + filter_params + [query_vector_str, top_k], - ) - - nodes: List[VectorNode] = [] - for row in cur.fetchall(): - unique_id, workspace_id_col, content, metadata, vector_str, score = row - # Parse vector string (format: [0.1,0.2,0.3] from pgvector) - vector = json.loads(vector_str) - - # Parse metadata if it's a string - if isinstance(metadata, str): - metadata = json.loads(metadata) - elif metadata is None: - metadata = {} - - node = VectorNode( - unique_id=unique_id, - workspace_id=workspace_id_col or workspace_id, - content=content, - metadata=metadata, - vector=vector, - ) - node.metadata["score"] = float(score) - nodes.append(node) - - return nodes - else: - # Filter-only search without vector similarity - with self._conn.cursor() as cur: - # Parameters: filter_params first, then top_k - cur.execute( - f""" - SELECT unique_id, workspace_id, content, metadata, vector::text - FROM "{table_name}" - {where_sql} - LIMIT %s - """, - filter_params + [top_k], - ) - - nodes: List[VectorNode] = [] - for row in cur.fetchall(): - node = self._row2node(row, workspace_id) - nodes.append(node) - - return nodes - - async def async_search( - self, - query: str, - workspace_id: str, - top_k: int = 1, - filter_dict: Optional[Dict[str, Any]] = None, - **kwargs, - ) -> List[VectorNode]: - """Search for similar vector nodes using cosine similarity (async). - - When query is empty (empty string or None), the search degrades to a - filter-only search without vector similarity ranking. - - Args: - query: The text query to search for. Will be embedded using the - embedding model's async method. If empty/None, performs filter-only search. - workspace_id: The identifier of the workspace to search in. - top_k: Maximum number of results to return (default: 1). - filter_dict: Optional dictionary of metadata filters to apply. - **kwargs: Additional keyword arguments (unused). - - Returns: - List[VectorNode]: List of matching vector nodes sorted by similarity score. - """ - if not await self.async_exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return [] - - # Build filters from filter_dict (use asyncpg-style placeholders) - where_clause, filter_params = self._build_sql_filters(filter_dict, use_async=True) - table_name = self._get_table_name(workspace_id) - - conn = await self._get_async_conn() - - # When query is empty, degrade to filter-only search without vector similarity - use_vector_search = bool(query) - - if use_vector_search: - # Use async embedding - query_vector = await self.async_get_embeddings(query) - # Convert query_vector to string format for pgvector - query_vector_str = "[" + ",".join(str(v) for v in query_vector) + "]" - - # Build parameter list and adjust placeholder indices - # $1 is for query_vector, so filter_params start from $2 - if filter_params: - adjusted_where = where_clause - # Replace placeholders in reverse order to avoid $1 being replaced when replacing $10 - for i in range(len(filter_params) - 1, -1, -1): - old_placeholder = f"${i + 1}" - new_placeholder = f"${i + 2}" # Shift by 1 since $1 is query_vector - adjusted_where = adjusted_where.replace(old_placeholder, new_placeholder) - where_sql = f"WHERE {adjusted_where}" - else: - where_sql = "" - - # Calculate the last parameter index for top_k - top_k_param_idx = 1 + len(filter_params) + 1 - - query_sql = f""" - SELECT unique_id, workspace_id, content, metadata, vector::text, - 1 - (vector <=> $1::vector) AS score - FROM "{table_name}" - {where_sql} - ORDER BY vector <=> $1::vector - LIMIT ${top_k_param_idx} - """ - - rows = await conn.fetch(query_sql, query_vector_str, *filter_params, top_k) - - nodes: List[VectorNode] = [] - for row in rows: - unique_id = row["unique_id"] - workspace_id_col = row["workspace_id"] - content = row["content"] - metadata = row["metadata"] - vector_str = row["vector"] - score = row["score"] - - # Parse vector string (format: [0.1,0.2,0.3] from pgvector) - vector = json.loads(vector_str) - - # Parse metadata if it's a string - if isinstance(metadata, str): - metadata = json.loads(metadata) - elif metadata is None: - metadata = {} - - node = VectorNode( - unique_id=unique_id, - workspace_id=workspace_id_col or workspace_id, - content=content, - metadata=metadata, - vector=vector, - ) - node.metadata["score"] = float(score) - nodes.append(node) - - return nodes - else: - # Filter-only search without vector similarity - if filter_params: - where_sql = f"WHERE {where_clause}" - top_k_param_idx = len(filter_params) + 1 - else: - where_sql = "" - top_k_param_idx = 1 - - query_sql = f""" - SELECT unique_id, workspace_id, content, metadata, vector::text - FROM "{table_name}" - {where_sql} - LIMIT ${top_k_param_idx} - """ - - rows = await conn.fetch(query_sql, *filter_params, top_k) - - nodes: List[VectorNode] = [] - for row in rows: - unique_id = row["unique_id"] - workspace_id_col = row["workspace_id"] - content = row["content"] - metadata = row["metadata"] - vector_str = row["vector"] - - # Parse vector string - vector = json.loads(vector_str) - - # Parse metadata if it's a string - if isinstance(metadata, str): - metadata = json.loads(metadata) - elif metadata is None: - metadata = {} - - node = VectorNode( - unique_id=unique_id, - workspace_id=workspace_id_col or workspace_id, - content=content, - metadata=metadata, - vector=vector, - ) - nodes.append(node) - - return nodes - - # ==================== Insert Methods ==================== - - def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs): - """Insert vector nodes into the PostgreSQL table. - - Args: - nodes: A single VectorNode or list of VectorNodes to insert. - workspace_id: The identifier of the workspace to insert into. - **kwargs: Additional keyword arguments (unused). - """ - if not self.exist_workspace(workspace_id=workspace_id): - self.create_workspace(workspace_id=workspace_id) - - if isinstance(nodes, VectorNode): - nodes = [nodes] - - embedded_nodes = [node for node in nodes if node.vector] - not_embedded_nodes = [node for node in nodes if not node.vector] - now_embedded_nodes = self.get_node_embeddings(not_embedded_nodes) - - table_name = self._get_table_name(workspace_id) - all_nodes = embedded_nodes + now_embedded_nodes - - # Use batch insert with ON CONFLICT for upsert - with self._conn.cursor() as cur: - for i in range(0, len(all_nodes), self.batch_size): - batch = all_nodes[i : i + self.batch_size] - values = [] - for node in batch: - # pgvector accepts list directly or string format - vector_value = node.vector - if isinstance(vector_value, list): - # Convert list to string format for pgvector: '[0.1,0.2,0.3]' - vector_str = "[" + ",".join(str(v) for v in vector_value) + "]" - else: - vector_str = str(vector_value) - - values.append( - ( - node.unique_id, - workspace_id, - node.content, - json.dumps(node.metadata), - vector_str, - ), - ) - - # Use INSERT ... ON CONFLICT for upsert - cur.executemany( - f""" - INSERT INTO "{table_name}" (unique_id, workspace_id, content, metadata, vector) - VALUES (%s, %s, %s, %s, %s::vector) - ON CONFLICT (unique_id) DO UPDATE SET - workspace_id = EXCLUDED.workspace_id, - content = EXCLUDED.content, - metadata = EXCLUDED.metadata, - vector = EXCLUDED.vector - """, - values, - ) - - self._conn.commit() - logger.info(f"insert nodes.size={len(all_nodes)} into workspace_id={workspace_id}") - - async def async_insert( - self, - nodes: VectorNode | List[VectorNode], - workspace_id: str, - **kwargs, - ): - """Insert vector nodes into the PostgreSQL table (async). - - Args: - nodes: A single VectorNode or list of VectorNodes to insert. - workspace_id: The identifier of the workspace to insert into. - **kwargs: Additional keyword arguments (unused). - """ - if not await self.async_exist_workspace(workspace_id=workspace_id): - await self.async_create_workspace(workspace_id=workspace_id) - - if isinstance(nodes, VectorNode): - nodes = [nodes] - - embedded_nodes = [node for node in nodes if node.vector] - not_embedded_nodes = [node for node in nodes if not node.vector] - - # Use async embedding - now_embedded_nodes = await self.async_get_node_embeddings(not_embedded_nodes) - - table_name = self._get_table_name(workspace_id) - all_nodes = embedded_nodes + now_embedded_nodes - - conn = await self._get_async_conn() - - # Use batch insert with ON CONFLICT for upsert - for i in range(0, len(all_nodes), self.batch_size): - batch = all_nodes[i : i + self.batch_size] - - # Prepare batch data for executemany - batch_data = [] - for node in batch: - # pgvector accepts list directly or string format - vector_value = node.vector - if isinstance(vector_value, list): - # Convert list to string format for pgvector: '[0.1,0.2,0.3]' - vector_str = "[" + ",".join(str(v) for v in vector_value) + "]" - else: - vector_str = str(vector_value) - - batch_data.append( - ( - node.unique_id, - workspace_id, - node.content, - json.dumps(node.metadata), - vector_str, - ), - ) - - # Use executemany for batch insert - await conn.executemany( - f""" - INSERT INTO "{table_name}" (unique_id, workspace_id, content, metadata, vector) - VALUES ($1, $2, $3, $4, $5::vector) - ON CONFLICT (unique_id) DO UPDATE SET - workspace_id = EXCLUDED.workspace_id, - content = EXCLUDED.content, - metadata = EXCLUDED.metadata, - vector = EXCLUDED.vector - """, - batch_data, - ) - - logger.info(f"async insert nodes.size={len(all_nodes)} into workspace_id={workspace_id}") - - # ==================== Delete Methods ==================== - - def delete(self, node_ids: str | List[str], workspace_id: str, **kwargs): - """Delete vector nodes from the PostgreSQL table. - - Args: - node_ids: A single node ID or list of node IDs to delete. - workspace_id: The identifier of the workspace to delete from. - **kwargs: Additional keyword arguments (unused). - """ - if not self.exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return - - if isinstance(node_ids, str): - node_ids = [node_ids] - - table_name = self._get_table_name(workspace_id) - with self._conn.cursor() as cur: - cur.executemany( - f'DELETE FROM "{table_name}" WHERE unique_id = %s', - [(node_id,) for node_id in node_ids], - ) - self._conn.commit() - logger.info(f"delete node_ids.size={len(node_ids)} from workspace_id={workspace_id}") - - async def async_delete(self, node_ids: str | List[str], workspace_id: str, **kwargs): - """Delete vector nodes from the PostgreSQL table (async). - - Args: - node_ids: A single node ID or list of node IDs to delete. - workspace_id: The identifier of the workspace to delete from. - **kwargs: Additional keyword arguments (unused). - """ - if not await self.async_exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return - - if isinstance(node_ids, str): - node_ids = [node_ids] - - table_name = self._get_table_name(workspace_id) - conn = await self._get_async_conn() - - # Use batch delete with ANY for better performance - await conn.execute( - f'DELETE FROM "{table_name}" WHERE unique_id = ANY($1)', - node_ids, - ) - - logger.info(f"async delete node_ids.size={len(node_ids)} from workspace_id={workspace_id}") - - # ==================== Close Methods ==================== - - def close(self): - """Close the synchronous PostgreSQL connection.""" - if self._conn: - self._conn.close() - - async def async_close(self): - """Close the asynchronous PostgreSQL connection.""" - if self._async_conn: - await self._async_conn.close() - self._async_conn = None - - # ==================== Private Async Helper Methods ==================== - - async def _get_async_conn(self): - """Get or create async PostgreSQL connection.""" - if self._async_conn is None: - import asyncpg - - logger.debug(f"Establishing async PostgreSQL connection: {self._async_conn_string}") - self._async_conn = await asyncpg.connect(self._async_conn_string) - logger.debug("Async PostgreSQL connection established successfully") - return self._async_conn diff --git a/flowllm/core/vector_store/qdrant_vector_store.py b/flowllm/core/vector_store/qdrant_vector_store.py deleted file mode 100644 index 07126aa..0000000 --- a/flowllm/core/vector_store/qdrant_vector_store.py +++ /dev/null @@ -1,970 +0,0 @@ -"""Qdrant vector store implementation. - -This module provides a Qdrant-based vector store that uses Qdrant for storing -and searching vector embeddings. Qdrant is an open-source vector similarity -search engine that provides high-performance vector search capabilities. - -The implementation supports both synchronous and asynchronous operations, and -provides native Qdrant async client support for better async performance. -""" - -import os -import uuid -from pathlib import Path -from typing import List, Dict, Any, Optional - -from loguru import logger - -from .memory_vector_store import MemoryVectorStore -from ..context import C -from ..schema import VectorNode - - -@C.register_vector_store("qdrant") -class QdrantVectorStore(MemoryVectorStore): - """Qdrant-based vector store implementation. - - This class provides a vector store implementation using Qdrant as the - underlying storage engine. It supports workspace management, vector similarity - search with filtering, and provides both synchronous and asynchronous operations. - - The QdrantVectorStore can be configured to connect to: - - A local Qdrant instance (default: localhost:6333) - - A remote Qdrant instance via URL - - Qdrant Cloud via URL and API key - - Attributes: - url: Optional URL for connecting to Qdrant. If provided, host and port - are ignored. Useful for Qdrant Cloud or custom deployments. - host: Host address of the Qdrant server (default: localhost). - port: Port number of the Qdrant server (default: 6333). - api_key: Optional API key for authentication (required for Qdrant Cloud). - distance: Distance metric for vector similarity (default: COSINE). - Can be COSINE, EUCLIDEAN, or DOT. - _client: Private QdrantClient instance for synchronous operations. - _async_client: Private AsyncQdrantClient instance for asynchronous operations. - """ - - # ==================== Initialization ==================== - - def __init__( - self, - url: str | None = None, - host: str | None = None, - port: int | None = None, - api_key: str | None = None, - distance: str | None = None, - **kwargs, - ): - """Initialize Qdrant clients. - - Args: - url: Optional URL for connecting to Qdrant. If provided, host and port are ignored. - host: Host address of the Qdrant server (default: localhost from env or "localhost"). - port: Port number of the Qdrant server (default: 6333 from env or 6333). - api_key: Optional API key for authentication (required for Qdrant Cloud). - distance: Distance metric for vector similarity (default: COSINE). - **kwargs: Additional keyword arguments passed to MemoryVectorStore. - """ - super().__init__(**kwargs) - self.url = url - self.host = host or os.getenv("FLOW_QDRANT_HOST", "localhost") - self.port = port or int(os.getenv("FLOW_QDRANT_PORT", "6333")) - self.api_key = api_key - - # Set default distance to COSINE if not set - from qdrant_client.http.models import Distance - - if distance is None: - self.distance: Distance = Distance.COSINE - else: - self.distance: Distance = Distance(distance) - - # Build kwargs for QdrantClient initialization - client_kwargs = {} - if self.url is not None: - client_kwargs["url"] = self.url - else: - if self.host is not None: - client_kwargs["host"] = self.host - if self.port is not None: - client_kwargs["port"] = self.port - - if self.api_key is not None: - client_kwargs["api_key"] = self.api_key - - from qdrant_client import QdrantClient, AsyncQdrantClient - - self._client = QdrantClient(**client_kwargs) - self._async_client = AsyncQdrantClient(**client_kwargs) - - # Log connection info - if self.url: - logger.info(f"Qdrant client initialized with url={self.url}") - else: - logger.info(f"Qdrant client initialized with host={self.host} port={self.port}") - - # ==================== Static Helper Methods ==================== - - @staticmethod - def _convert_id_to_uuid(node_id: str) -> str: - """Convert a string ID to a UUID string. - - Uses UUID v5 with a fixed namespace to ensure deterministic conversion. - If the input is already a valid UUID, returns it as-is. - - Args: - node_id: The node ID to convert. - - Returns: - A valid UUID string. - """ - try: - # Check if it's already a valid UUID - uuid.UUID(node_id) - return node_id - except (ValueError, AttributeError): - # Convert string to UUID v5 using a fixed namespace - # Using DNS namespace as a standard base - return str(uuid.uuid5(uuid.NAMESPACE_DNS, str(node_id))) - - @staticmethod - def point2node(point, workspace_id: str) -> VectorNode: - """Convert Qdrant point to VectorNode. - - Converts a Qdrant Record object (returned from search or scroll operations) - into a VectorNode object, extracting the ID, vector, content, and metadata. - - Args: - point: The Qdrant Record object containing point data. - workspace_id: The workspace ID to assign to the VectorNode. - - Returns: - VectorNode: A VectorNode object created from the Qdrant point data. - """ - # Use original_id if available, otherwise use the point ID - original_id = point.payload.get("original_id", str(point.id)) - node = VectorNode( - unique_id=original_id, - workspace_id=workspace_id, - content=point.payload.get("content", ""), - metadata=point.payload.get("metadata", {}), - vector=point.vector, - ) - if hasattr(point, "score") and point.score is not None: - node.metadata["score"] = point.score - return node - - @staticmethod - def _build_qdrant_filters(filter_dict: Optional[Dict[str, Any]] = None): - """Build Qdrant filter from filter_dict. - - Converts a filter dictionary into a Qdrant Filter object. Supports: - - Term filters: Direct value matching (e.g., {"node_type": "n1"}) - - Range filters: Numeric range queries (e.g., {"age": {"gte": 18, "lte": 65}}) - - Args: - filter_dict: Optional dictionary of filters. Keys are field names (automatically - prefixed with "metadata." if not already present). Values can be: - - Direct values for term matching - - Dict with range operators: "gte", "lte", "gt", "lt" - - Returns: - Optional Filter object for Qdrant queries, or None if filter_dict is empty. - - Example: - ```python - # Term filter - filter_dict = {"node_type": "n1"} - - # Range filter - filter_dict = {"age": {"gte": 18, "lte": 65}} - ``` - """ - from qdrant_client.http.models import FieldCondition, MatchAny, MatchValue, Range - - if not filter_dict: - return None - - conditions = [] - for key, filter_value in filter_dict.items(): - # Handle special keys that are stored at payload root level - if key == "unique_id": - qdrant_key = "original_id" - # Handle nested keys by prefixing with metadata. - elif not key.startswith("metadata."): - qdrant_key = f"metadata.{key}" - else: - qdrant_key = key - - if isinstance(filter_value, dict): - # Range filter: {"gte": 1, "lte": 10} - range_conditions = {} - if "gte" in filter_value: - range_conditions["gte"] = filter_value["gte"] - if "lte" in filter_value: - range_conditions["lte"] = filter_value["lte"] - if "gt" in filter_value: - range_conditions["gt"] = filter_value["gt"] - if "lt" in filter_value: - range_conditions["lt"] = filter_value["lt"] - if range_conditions: - conditions.append( - FieldCondition( - key=qdrant_key, - range=Range(**range_conditions), - ), - ) - elif isinstance(filter_value, list): - # List filter: use MatchAny for OR logic - conditions.append( - FieldCondition( - key=qdrant_key, - match=MatchAny(any=filter_value), - ), - ) - else: - # Term filter: direct value comparison - conditions.append( - FieldCondition( - key=qdrant_key, - match=MatchValue(value=filter_value), - ), - ) - - if not conditions: - return None - - from qdrant_client.http.models import Filter - - return Filter(must=conditions) - - # ==================== Workspace Management Methods ==================== - - def exist_workspace(self, workspace_id: str, **kwargs) -> bool: - """Check if a collection exists in Qdrant. - - Args: - workspace_id: The ID of the workspace (collection) to check. - **kwargs: Additional keyword arguments (unused, kept for interface compatibility). - - Returns: - True if the collection exists, False otherwise. - """ - return self._client.collection_exists(collection_name=workspace_id) - - async def async_exist_workspace(self, workspace_id: str, **kwargs) -> bool: - """Async version of exist_workspace using native Qdrant async client. - - Args: - workspace_id: The ID of the workspace (collection) to check. - **kwargs: Additional keyword arguments (unused, kept for interface compatibility). - - Returns: - True if the collection exists, False otherwise. - """ - return await self._async_client.collection_exists(collection_name=workspace_id) - - def delete_workspace(self, workspace_id: str, **kwargs): - """Delete a collection from Qdrant. - - Args: - workspace_id: The ID of the workspace (collection) to delete. - **kwargs: Additional keyword arguments (unused, kept for interface compatibility). - - Returns: - The result of the delete_collection operation. - """ - return self._client.delete_collection(collection_name=workspace_id) - - async def async_delete_workspace(self, workspace_id: str, **kwargs): - """Async version of delete_workspace using native Qdrant async client. - - Args: - workspace_id: The ID of the workspace (collection) to delete. - **kwargs: Additional keyword arguments (unused, kept for interface compatibility). - - Returns: - The result of the delete_collection operation. - """ - return await self._async_client.delete_collection(collection_name=workspace_id) - - def create_workspace(self, workspace_id: str, **kwargs): - """Create a new collection in Qdrant. - - Creates a new collection with vector configuration based on the embedding - model's dimensions and the configured distance metric. - - Args: - workspace_id: The ID of the workspace (collection) to create. - **kwargs: Additional keyword arguments (unused, kept for interface compatibility). - - Returns: - The result of the create_collection operation. - """ - from qdrant_client.http.models import VectorParams - - return self._client.create_collection( - collection_name=workspace_id, - vectors_config=VectorParams( - size=self.embedding_model.dimensions, - distance=self.distance, - ), - ) - - async def async_create_workspace(self, workspace_id: str, **kwargs): - """Async version of create_workspace using native Qdrant async client. - - Creates a new collection with vector configuration based on the embedding - model's dimensions and the configured distance metric. - - Args: - workspace_id: The ID of the workspace (collection) to create. - **kwargs: Additional keyword arguments (unused, kept for interface compatibility). - - Returns: - The result of the create_collection operation. - """ - from qdrant_client.http.models import VectorParams - - return await self._async_client.create_collection( - collection_name=workspace_id, - vectors_config=VectorParams( - size=self.embedding_model.dimensions, - distance=self.distance, - ), - ) - - def list_workspace_nodes( - self, - workspace_id: str, - limit: int = 10000, - **kwargs, - ) -> List[VectorNode]: - """List all nodes in a workspace. - - Uses Qdrant's scroll API to paginate through all points in a collection, - converting them to VectorNode objects. - - Args: - workspace_id: The ID of the workspace to list nodes from. - limit: Maximum number of points to retrieve per scroll operation (default: 10000). - **kwargs: Additional keyword arguments (unused, kept for interface compatibility). - - Returns: - List[VectorNode]: All nodes in the workspace. - """ - nodes: List[VectorNode] = [] - offset = None - while True: - records, next_offset = self._client.scroll( - collection_name=workspace_id, - limit=limit, - offset=offset, - with_payload=True, - with_vectors=True, - ) - - if not records: - break - - for record in records: - node = self.point2node(record, workspace_id) - nodes.append(node) - - if next_offset is None: - break - offset = next_offset - - return nodes - - async def async_list_workspace_nodes( - self, - workspace_id: str, - limit: int = 10000, - **kwargs, - ) -> List[VectorNode]: - """Async version of list_workspace_nodes using native Qdrant async client. - - Uses Qdrant's async scroll API to paginate through all points in a collection, - converting them to VectorNode objects. - - Args: - workspace_id: The ID of the workspace to iterate nodes from. - limit: Maximum number of points to retrieve per scroll operation (default: 10000). - **kwargs: Additional keyword arguments (unused, kept for interface compatibility). - - Returns: - List[VectorNode]: All nodes in the workspace. - """ - nodes: List[VectorNode] = [] - offset = None - while True: - records, next_offset = await self._async_client.scroll( - collection_name=workspace_id, - limit=limit, - offset=offset, - with_payload=True, - with_vectors=True, - ) - - if not records: - break - - for record in records: - node = self.point2node(record, workspace_id) - nodes.append(node) - - if next_offset is None: - break - offset = next_offset - - return nodes - - def list_workspace(self, **kwargs) -> List[str]: - """ - List all existing workspaces (collections) in Qdrant. - - Returns: - List[str]: Workspace identifiers (collection names). - """ - return [c.name for c in self._client.get_collections().collections] - - async def async_list_workspace(self, **kwargs) -> List[str]: - """ - Async version of list_workspace using native Qdrant async client. - - Returns: - List[str]: Workspace identifiers (collection names). - """ - collections = await self._async_client.get_collections() - return [c.name for c in collections.collections] - - def dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs): - """Export a workspace to disk at the specified path. - - Args: - workspace_id: Identifier of the workspace to export. - path: Directory path where to write the exported workspace file. - callback_fn: Optional callback function to transform nodes during export. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - - Returns: - dict: Dictionary with "size" key indicating number of nodes exported, - or empty dict if workspace doesn't exist. - """ - if not self.exist_workspace(workspace_id=workspace_id, **kwargs): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return {} - - return self._dump_to_path( - nodes=self.list_workspace_nodes(workspace_id=workspace_id, **kwargs), - workspace_id=workspace_id, - path=path, - callback_fn=callback_fn, - **kwargs, - ) - - async def async_dump_workspace(self, workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs): - """Export a workspace to disk at the specified path (async). - - Args: - workspace_id: Identifier of the workspace to export. - path: Directory path where to write the exported workspace file. - callback_fn: Optional callback function to transform nodes during export. - **kwargs: Additional keyword arguments to pass to _dump_to_path. - - Returns: - dict: Dictionary with "size" key indicating number of nodes exported, - or empty dict if workspace doesn't exist. - """ - if not await self.async_exist_workspace(workspace_id=workspace_id, **kwargs): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return {} - - nodes = await self.async_list_workspace_nodes(workspace_id=workspace_id, **kwargs) - return await self._run_sync_in_executor( - self._dump_to_path, - nodes=nodes, - workspace_id=workspace_id, - path=path, - callback_fn=callback_fn, - **kwargs, - ) - - def load_workspace( - self, - workspace_id: str, - path: str | Path = "", - nodes: List[VectorNode] | None = None, - callback_fn=None, - **kwargs, - ): - """Load a workspace from disk, optionally merging with provided nodes. - - This method replaces any existing workspace with the same ID, then loads - nodes from the specified path and/or the provided nodes list. - - Args: - workspace_id: Identifier for the workspace to create/load. - path: Directory path containing the workspace file to load. - nodes: Optional list of VectorNode instances to merge with loaded nodes. - callback_fn: Optional callback function to transform node dictionaries. - **kwargs: Additional keyword arguments to pass to load operations. - - Returns: - dict: Dictionary with "size" key indicating total number of nodes loaded. - """ - if self.exist_workspace(workspace_id, **kwargs): - self.delete_workspace(workspace_id=workspace_id, **kwargs) - logger.info(f"delete workspace_id={workspace_id}") - - self.create_workspace(workspace_id=workspace_id, **kwargs) - - all_nodes: List[VectorNode] = [] - - if nodes: - all_nodes.extend(nodes) - - all_nodes.extend(self._load_from_path(path=path, workspace_id=workspace_id, callback_fn=callback_fn, **kwargs)) - - if all_nodes: - self.insert(nodes=all_nodes, workspace_id=workspace_id, **kwargs) - return {"size": len(all_nodes)} - - async def async_load_workspace( - self, - workspace_id: str, - path: str | Path = "", - nodes: List[VectorNode] | None = None, - callback_fn=None, - **kwargs, - ): - """Load a workspace from disk, optionally merging with provided nodes (async). - - This method replaces any existing workspace with the same ID, then loads - nodes from the specified path and/or the provided nodes list. - - Args: - workspace_id: Identifier for the workspace to create/load. - path: Directory path containing the workspace file to load. - nodes: Optional list of VectorNode instances to merge with loaded nodes. - callback_fn: Optional callback function to transform node dictionaries. - **kwargs: Additional keyword arguments to pass to load operations. - - Returns: - dict: Dictionary with "size" key indicating total number of nodes loaded. - """ - if await self.async_exist_workspace(workspace_id, **kwargs): - await self.async_delete_workspace(workspace_id=workspace_id, **kwargs) - logger.info(f"delete workspace_id={workspace_id}") - - await self.async_create_workspace(workspace_id=workspace_id, **kwargs) - - all_nodes: List[VectorNode] = [] - - if nodes: - all_nodes.extend(nodes) - - if path: - loaded_nodes = await self._run_sync_in_executor( - lambda: list( - self._load_from_path(path=path, workspace_id=workspace_id, callback_fn=callback_fn, **kwargs), - ), - ) - all_nodes.extend(loaded_nodes) - - if all_nodes: - await self.async_insert(nodes=all_nodes, workspace_id=workspace_id, **kwargs) - return {"size": len(all_nodes)} - - def copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs): - """Copy all nodes from one workspace to another. - - Args: - src_workspace_id: Identifier of the source workspace. - dest_workspace_id: Identifier of the destination workspace. - Created if it doesn't exist. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - dict: Dictionary with "size" key indicating number of nodes copied, - or empty dict if source workspace doesn't exist. - """ - if not self.exist_workspace(workspace_id=src_workspace_id, **kwargs): - logger.warning(f"src_workspace_id={src_workspace_id} does not exist!") - return {} - - if not self.exist_workspace(dest_workspace_id, **kwargs): - self.create_workspace(workspace_id=dest_workspace_id, **kwargs) - - src_nodes = self.list_workspace_nodes(workspace_id=src_workspace_id, **kwargs) - if src_nodes: - self.insert(nodes=src_nodes, workspace_id=dest_workspace_id, **kwargs) - - return {"size": len(src_nodes)} - - async def async_copy_workspace(self, src_workspace_id: str, dest_workspace_id: str, **kwargs): - """Copy all nodes from one workspace to another (async). - - Args: - src_workspace_id: Identifier of the source workspace. - dest_workspace_id: Identifier of the destination workspace. - Created if it doesn't exist. - **kwargs: Additional keyword arguments to pass to operations. - - Returns: - dict: Dictionary with "size" key indicating number of nodes copied, - or empty dict if source workspace doesn't exist. - """ - if not await self.async_exist_workspace(workspace_id=src_workspace_id, **kwargs): - logger.warning(f"src_workspace_id={src_workspace_id} does not exist!") - return {} - - if not await self.async_exist_workspace(dest_workspace_id, **kwargs): - await self.async_create_workspace(workspace_id=dest_workspace_id, **kwargs) - - src_nodes = await self.async_list_workspace_nodes(workspace_id=src_workspace_id, **kwargs) - if src_nodes: - await self.async_insert(nodes=src_nodes, workspace_id=dest_workspace_id, **kwargs) - - return {"size": len(src_nodes)} - - # ==================== Search Methods ==================== - - def search( - self, - query: str, - workspace_id: str, - top_k: int = 1, - filter_dict: Optional[Dict[str, Any]] = None, - **kwargs, - ) -> List[VectorNode]: - """Search for similar vectors in the workspace. - - When query is empty (empty string or None), the search degrades to a - filter-only search without vector similarity ranking. - - Args: - query: Text query to search for. If empty/None, performs filter-only - search without vector similarity. - workspace_id: The ID of the workspace to search in. - top_k: Number of most similar results to return (default: 1). - filter_dict: Optional dictionary of filters to apply to the search. - See _build_qdrant_filters for filter format details. - **kwargs: Additional keyword arguments passed to Qdrant's query_points method. - - Returns: - List[VectorNode]: List of matching nodes. When query is provided, - nodes are sorted by similarity score (highest first) with a - "score" key in metadata. When query is empty, nodes are returned - in storage order without scores. - """ - if not self.exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return [] - - # Build filters from filter_dict - qdrant_filter = self._build_qdrant_filters(filter_dict) - - # When query is empty, degrade to filter-only search without vector similarity - use_vector_search = bool(query) - - if use_vector_search: - query_vector = self.get_embeddings(query) - - response = self._client.query_points( - collection_name=workspace_id, - query=query_vector, - limit=top_k, - query_filter=qdrant_filter, - with_payload=True, - with_vectors=True, - **kwargs, - ) - - nodes: List[VectorNode] = [] - for scored_point in response.points: - node = self.point2node(scored_point, workspace_id) - node.metadata["score"] = scored_point.score - nodes.append(node) - - return nodes - else: - # Filter-only search using scroll API - records, _ = self._client.scroll( - collection_name=workspace_id, - limit=top_k, - scroll_filter=qdrant_filter, - with_payload=True, - with_vectors=True, - ) - - nodes: List[VectorNode] = [] - for record in records: - node = self.point2node(record, workspace_id) - nodes.append(node) - - return nodes - - async def async_search( - self, - query: str, - workspace_id: str, - top_k: int = 1, - filter_dict: Optional[Dict[str, Any]] = None, - **kwargs, - ) -> List[VectorNode]: - """Async version of search using native Qdrant async client and async embedding. - - When query is empty (empty string or None), the search degrades to a - filter-only search without vector similarity ranking. - - Args: - query: Text query to search for. If empty/None, performs filter-only - search without vector similarity. - workspace_id: The ID of the workspace to search in. - top_k: Number of most similar results to return (default: 1). - filter_dict: Optional dictionary of filters to apply to the search. - See _build_qdrant_filters for filter format details. - **kwargs: Additional keyword arguments passed to Qdrant's async query_points method. - - Returns: - List[VectorNode]: List of matching nodes. When query is provided, - nodes are sorted by similarity score (highest first) with a - "score" key in metadata. When query is empty, nodes are returned - in storage order without scores. - """ - if not await self.async_exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return [] - - # Build filters from filter_dict - qdrant_filter = self._build_qdrant_filters(filter_dict) - - # When query is empty, degrade to filter-only search without vector similarity - use_vector_search = bool(query) - - if use_vector_search: - # Use async embedding - query_vector = await self.async_get_embeddings(query) - - response = await self._async_client.query_points( - collection_name=workspace_id, - query=query_vector, - limit=top_k, - query_filter=qdrant_filter, - with_payload=True, - with_vectors=True, - **kwargs, - ) - - nodes: List[VectorNode] = [] - for scored_point in response.points: - node = self.point2node(scored_point, workspace_id) - node.metadata["score"] = scored_point.score - nodes.append(node) - - return nodes - else: - # Filter-only search using scroll API - records, _ = await self._async_client.scroll( - collection_name=workspace_id, - limit=top_k, - scroll_filter=qdrant_filter, - with_payload=True, - with_vectors=True, - ) - - nodes: List[VectorNode] = [] - for record in records: - node = self.point2node(record, workspace_id) - nodes.append(node) - - return nodes - - # ==================== Insert Methods ==================== - - def insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs): - """Insert nodes into the workspace. - - Inserts one or more nodes into the workspace. If nodes don't have embeddings, - they will be generated using the embedding_model. The workspace will be - created automatically if it doesn't exist. - - Args: - nodes: A single VectorNode or a list of VectorNode objects to insert. - Nodes without vectors will have embeddings generated automatically. - workspace_id: The ID of the workspace to insert nodes into. - **kwargs: Additional keyword arguments passed to Qdrant's upsert method. - - Note: - If a node with the same unique_id already exists, it will be updated - (upsert behavior). - """ - if not self.exist_workspace(workspace_id=workspace_id): - self.create_workspace(workspace_id=workspace_id) - - if isinstance(nodes, VectorNode): - nodes = [nodes] - - # Use parent's get_node_embeddings method - nodes = self.get_node_embeddings(nodes) - - from qdrant_client.http.models import PointStruct - - points = [ - PointStruct( - id=self._convert_id_to_uuid(node.unique_id), - vector=node.vector, - payload={ - "workspace_id": workspace_id, - "content": node.content, - "metadata": node.metadata, - "original_id": node.unique_id, # Store original ID in payload - }, - ) - for node in nodes - ] - - self._client.upsert( - collection_name=workspace_id, - points=points, - **kwargs, - ) - logger.info(f"insert points.size={len(points)} to workspace_id={workspace_id}") - - async def async_insert(self, nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs): - """Async version of insert using native Qdrant async client and async embedding. - - Inserts one or more nodes into the workspace asynchronously. If nodes don't - have embeddings, they will be generated using async embedding_model. The workspace - will be created automatically if it doesn't exist. - - Args: - nodes: A single VectorNode or a list of VectorNode objects to insert. - Nodes without vectors will have embeddings generated automatically. - workspace_id: The ID of the workspace to insert nodes into. - **kwargs: Additional keyword arguments passed to Qdrant's async upsert method. - - Note: - If a node with the same unique_id already exists, it will be updated - (upsert behavior). - """ - if not await self.async_exist_workspace(workspace_id=workspace_id): - await self.async_create_workspace(workspace_id=workspace_id) - - if isinstance(nodes, VectorNode): - nodes = [nodes] - - # Use parent's async_get_node_embeddings method - nodes = await self.async_get_node_embeddings(nodes) - - from qdrant_client.http.models import PointStruct - - points = [ - PointStruct( - id=self._convert_id_to_uuid(node.unique_id), - vector=node.vector, - payload={ - "workspace_id": workspace_id, - "content": node.content, - "metadata": node.metadata, - "original_id": node.unique_id, # Store original ID in payload - }, - ) - for node in nodes - ] - - await self._async_client.upsert( - collection_name=workspace_id, - points=points, - **kwargs, - ) - logger.info(f"async insert points.size={len(points)} to workspace_id={workspace_id}") - - # ==================== Delete Methods ==================== - - def delete(self, node_ids: str | List[str], workspace_id: str, **kwargs): - """Delete nodes from the workspace by their IDs. - - Args: - node_ids: A single node ID or a list of node IDs to delete. - workspace_id: The ID of the workspace containing the nodes to delete. - **kwargs: Additional keyword arguments passed to Qdrant's delete method. - - Note: - If the workspace doesn't exist, a warning is logged and the method returns - without performing any operation. - """ - if not self.exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return - - if isinstance(node_ids, str): - node_ids = [node_ids] - - # Convert node IDs to UUIDs - uuid_ids = [self._convert_id_to_uuid(node_id) for node_id in node_ids] - - from qdrant_client.http import models - - self._client.delete( - collection_name=workspace_id, - points_selector=models.PointIdsList( - points=uuid_ids, - ), - **kwargs, - ) - logger.info(f"delete node_ids.size={len(node_ids)} from workspace_id={workspace_id}") - - async def async_delete(self, node_ids: str | List[str], workspace_id: str, **kwargs): - """Async version of delete using native Qdrant async client. - - Args: - node_ids: A single node ID or a list of node IDs to delete. - workspace_id: The ID of the workspace containing the nodes to delete. - **kwargs: Additional keyword arguments passed to Qdrant's async delete method. - - Note: - If the workspace doesn't exist, a warning is logged and the method returns - without performing any operation. - """ - if not await self.async_exist_workspace(workspace_id=workspace_id): - logger.warning(f"workspace_id={workspace_id} does not exist!") - return - - if isinstance(node_ids, str): - node_ids = [node_ids] - - # Convert node IDs to UUIDs - uuid_ids = [self._convert_id_to_uuid(node_id) for node_id in node_ids] - - from qdrant_client.http import models - - await self._async_client.delete( - collection_name=workspace_id, - points_selector=models.PointIdsList( - points=uuid_ids, - ), - **kwargs, - ) - logger.info(f"async delete node_ids.size={len(node_ids)} from workspace_id={workspace_id}") - - # ==================== Close Methods ==================== - - def close(self): - """Close the Qdrant client. - - Closes the synchronous Qdrant client and releases any resources. - Should be called when the vector store is no longer needed. - """ - self._client.close() - - async def async_close(self): - """Async close the Qdrant client. - - Closes the asynchronous Qdrant client and releases any resources. - Should be called when the vector store is no longer needed in async contexts. - """ - await self._async_client.close() diff --git a/flowllm/enumeration/__init__.py b/flowllm/enumeration/__init__.py new file mode 100644 index 0000000..9887dde --- /dev/null +++ b/flowllm/enumeration/__init__.py @@ -0,0 +1,9 @@ +"""Enumeration""" + +from .chunk_enum import ChunkEnum +from .component_enum import ComponentEnum + +__all__ = [ + "ChunkEnum", + "ComponentEnum", +] diff --git a/flowllm/enumeration/chunk_enum.py b/flowllm/enumeration/chunk_enum.py new file mode 100644 index 0000000..af37157 --- /dev/null +++ b/flowllm/enumeration/chunk_enum.py @@ -0,0 +1,21 @@ +"""Chunk type enumeration for stream processing.""" + +from enum import Enum + + +class ChunkEnum(str, Enum): + """Chunk categories for streaming responses.""" + + REPLY_START = "reply_start" + REPLY_END = "reply_end" + + THINK = "think" + CONTENT = "content" + DATA = "data" + + TOOL_CALL = "tool_call" + TOOL_RESULT = "tool_result" + + USAGE = "usage" + ERROR = "error" + DONE = "done" diff --git a/flowllm/enumeration/component_enum.py b/flowllm/enumeration/component_enum.py new file mode 100644 index 0000000..e1a0784 --- /dev/null +++ b/flowllm/enumeration/component_enum.py @@ -0,0 +1,37 @@ +"""Component type enumeration.""" + +from enum import Enum + + +class ComponentEnum(str, Enum): + """Component types for dependency injection.""" + + BASE = "base" + + AS_LLM = "as_llm" + + AS_EMBEDDING = "as_embedding" + + EMBEDDING_STORE = "embedding_store" + + FILE_CHUNKER = "file_chunker" + + FILE_STORE = "file_store" + + FILE_GRAPH = "file_graph" + + FILE_CATALOG = "file_catalog" + + KEYWORD_INDEX = "keyword_index" + + SERVICE = "service" + + CLIENT = "client" + + STEP = "step" + + JOB = "job" + + TOKENIZER = "tokenizer" + + AGENT_WRAPPER = "agent_wrapper" diff --git a/flowllm/extensions/__init__.py b/flowllm/extensions/__init__.py deleted file mode 100644 index a61939f..0000000 --- a/flowllm/extensions/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -"""Extension packages for FlowLLM framework. - -This package provides extension modules that can be used in LLM-powered flows. -It includes ready-to-use extension packages for: - -- file_tool: File-related operations including editing and searching files -- data: Data-related operations including downloading stock data -- utils: Utility functions for date/time operations and other helpers -- skills: Skill-based operations for managing and executing specialized skills -""" - -from . import file_tool, skills, utils - -__all__ = [ - "file_tool", - "skills", - "utils", -] diff --git a/flowllm/extensions/file_tool/__init__.py b/flowllm/extensions/file_tool/__init__.py deleted file mode 100644 index e907314..0000000 --- a/flowllm/extensions/file_tool/__init__.py +++ /dev/null @@ -1,49 +0,0 @@ -"""File tool package for FlowLLM framework. - -This package provides file-related operations that can be used in LLM-powered flows. -It includes ready-to-use operations for: - -- EditOp: File editing operation for replacing text within files -- ExitPlanModeOp: Exit plan mode operation -- GlobOp: File search operation for finding files matching glob patterns -- GrepOp: Text search operation for finding patterns in files -- LSOp: List directory contents operation -- ReadFileOp: Read single file operation -- ReadManyFilesOp: Read multiple files operation matching glob patterns -- RipGrepOp: Advanced text search operation using ripgrep-like functionality -- ShellOp: Shell command execution operation -- SmartEditOp: Smart file editing operation with context awareness -- TaskOp: Task management operation -- WriteFileOp: Write file operation -- WriteTodosOp: To-do list management operation for tracking subtasks -""" - -from .edit_op import EditOp -from .exit_plan_mode_op import ExitPlanModeOp -from .glob_op import GlobOp -from .grep_op import GrepOp -from .ls_op import LSOp -from .read_file_op import ReadFileOp -from .read_many_files_op import ReadManyFilesOp -from .rip_grep_op import RipGrepOp -from .shell_op import ShellOp -from .smart_edit_op import SmartEditOp -from .task_op import TaskOp -from .write_file_op import WriteFileOp -from .write_todos_op import WriteTodosOp - -__all__ = [ - "EditOp", - "ExitPlanModeOp", - "GlobOp", - "GrepOp", - "LSOp", - "ReadFileOp", - "ReadManyFilesOp", - "RipGrepOp", - "ShellOp", - "SmartEditOp", - "TaskOp", - "WriteFileOp", - "WriteTodosOp", -] diff --git a/flowllm/extensions/file_tool/edit_op.py b/flowllm/extensions/file_tool/edit_op.py deleted file mode 100644 index 81bab3e..0000000 --- a/flowllm/extensions/file_tool/edit_op.py +++ /dev/null @@ -1,153 +0,0 @@ -"""File edit operation module. - -This module provides a tool operation for editing files by replacing text. -It supports creating new files, editing existing files, and replacing multiple occurrences. -""" - -from pathlib import Path -from typing import Optional - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - - -@C.register_op() -class EditOp(BaseAsyncToolOp): - """File edit operation. - - This operation replaces text within a file. By default, replaces a single - occurrence, but can replace multiple occurrences when expected_replacements - is specified. Supports creating new files when old_string is empty. - """ - - file_path = __file__ - - def __init__(self, **kwargs): - kwargs.setdefault("raise_exception", False) - super().__init__(**kwargs) - - def build_tool_call(self) -> ToolCall: - """Build and return the tool call schema for this operator.""" - tool_params = { - "name": "Edit", - "description": self.get_prompt("tool_desc"), - "input_schema": { - "file_path": { - "type": "string", - "description": self.get_prompt("file_path"), - "required": True, - }, - "old_string": { - "type": "string", - "description": self.get_prompt("old_string"), - "required": True, - }, - "new_string": { - "type": "string", - "description": self.get_prompt("new_string"), - "required": True, - }, - "expected_replacements": { - "type": "number", - "description": self.get_prompt("expected_replacements"), - "required": False, - }, - }, - } - - return ToolCall(**tool_params) - - async def async_execute(self): - """Execute the file edit operation.""" - file_path: str = self.input_dict.get("file_path", "").strip() - old_string: str = self.input_dict.get("old_string", "") - new_string: str = self.input_dict.get("new_string", "") - expected_replacements: Optional[int] = self.input_dict.get("expected_replacements") - - # Validate inputs - if not file_path: - raise ValueError("The 'file_path' parameter cannot be empty.") - - if expected_replacements is None: - expected_replacements = 1 - - if expected_replacements < 1: - raise ValueError("The 'expected_replacements' parameter must be at least 1.") - - # Resolve file path - file_path_obj = Path(file_path).expanduser().resolve() - - # Check if file exists - file_exists = file_path_obj.exists() and file_path_obj.is_file() - - # Handle new file creation - if not old_string and not file_exists: - # Create new file - file_path_obj.parent.mkdir(parents=True, exist_ok=True) - file_path_obj.write_text(new_string, encoding="utf-8") - self.set_output(f"Created new file: {file_path_obj}") - return - - # File must exist for editing - if not file_exists: - raise FileNotFoundError( - f"File not found: {file_path_obj}. Use an empty old_string to create a new file.", - ) - - # Cannot create file that already exists - if not old_string and file_exists: - raise ValueError( - f"Failed to edit. Attempted to create a file that already exists: {file_path_obj}", - ) - - # Read current content - current_content = file_path_obj.read_text(encoding="utf-8") - current_content = current_content.replace("\r\n", "\n") - - # Count occurrences - occurrences = current_content.count(old_string) - - # Validate occurrences - if occurrences == 0: - raise ValueError( - f"Failed to edit, could not find the string to replace. " - f"0 occurrences found for old_string in {file_path_obj}.", - ) - - if occurrences != expected_replacements: - occurrence_term = "occurrence" if expected_replacements == 1 else "occurrences" - raise ValueError( - f"Failed to edit, expected {expected_replacements} {occurrence_term} " - f"but found {occurrences} for old_string in {file_path_obj}.", - ) - - # Check if old_string and new_string are identical - if old_string == new_string: - raise ValueError( - f"No changes to apply. The old_string and new_string are identical in {file_path_obj}.", - ) - - # Perform replacement - new_content = current_content.replace(old_string, new_string, occurrences) - - # Check if content actually changed - if current_content == new_content: - raise ValueError( - f"No changes to apply. The new content is identical to the current content in {file_path_obj}.", - ) - - # Write new content - file_path_obj.write_text(new_content, encoding="utf-8") - - # Set output - result_msg = f"Successfully modified file: {file_path_obj} ({occurrences} replacements)." - self.set_output(result_msg) - - async def async_default_execute(self, e: Exception = None, **kwargs): - """Fill outputs with a default failure message when execution fails.""" - file_path: str = self.input_dict.get("file_path", "").strip() - error_msg = f'Failed to edit file "{file_path}"' - if e: - error_msg += f": {str(e)}" - self.set_output(error_msg) diff --git a/flowllm/extensions/file_tool/edit_prompt.yaml b/flowllm/extensions/file_tool/edit_prompt.yaml deleted file mode 100644 index 970363c..0000000 --- a/flowllm/extensions/file_tool/edit_prompt.yaml +++ /dev/null @@ -1,23 +0,0 @@ -tool_desc: | - Replaces text within a file. By default, replaces a single occurrence, but can replace multiple occurrences when \`expected_replacements\` is specified. This tool requires providing significant context around the change to ensure precise targeting. Always use the ${READ_FILE_TOOL_NAME} tool to examine the file's current content before attempting a text replacement. - The user has the ability to modify the \`new_string\` content. If modified, this will be stated in the response. - Expectation for required parameters: - 1. \`file_path\` is the path to the file to modify. - 2. \`old_string\` MUST be the exact literal text to replace (including all whitespace, indentation, newlines, and surrounding code etc.). - 3. \`new_string\` MUST be the exact literal text to replace \`old_string\` with (also including all whitespace, indentation, newlines, and surrounding code etc.). Ensure the resulting code is correct and idiomatic. - 4. NEVER escape \`old_string\` or \`new_string\`, that would break the exact literal text requirement. - **Important:** If ANY of the above are not satisfied, the tool will fail. CRITICAL for \`old_string\`: Must uniquely identify the single instance to change. Include at least 3 lines of context BEFORE and AFTER the target text, matching whitespace and indentation precisely. If this string matches multiple locations, or does not match exactly, the tool will fail. - **Multiple replacements:** Set \`expected_replacements\` to the number of occurrences you want to replace. The tool will replace ALL occurrences that match \`old_string\` exactly. Ensure the number of replacements matches your expectation. - -file_path: | - The path to the file to modify. - -old_string: | - The exact literal text to replace, preferably unescaped. For single replacements (default), include at least 3 lines of context BEFORE and AFTER the target text, matching whitespace and indentation precisely. For multiple replacements, specify expected_replacements parameter. If this string is not the exact literal text (i.e. you escaped it) or does not match exactly, the tool will fail. - -new_string: | - The exact literal text to replace `old_string` with, preferably unescaped. Provide the EXACT text. Ensure the resulting code is correct and idiomatic. - -expected_replacements: | - Optional: Number of replacements expected. Defaults to 1 if not specified. Use when you want to replace multiple occurrences. - diff --git a/flowllm/extensions/file_tool/exit_plan_mode_op.py b/flowllm/extensions/file_tool/exit_plan_mode_op.py deleted file mode 100644 index 30a46ba..0000000 --- a/flowllm/extensions/file_tool/exit_plan_mode_op.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Exit plan mode operation module. - -This module provides a tool operation for exiting plan mode after presenting -an implementation plan to the user for approval. -""" - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - - -@C.register_op() -class ExitPlanModeOp(BaseAsyncToolOp): - """Exit plan mode operation. - - This operation is used when in plan mode and ready to present the plan - to the user for approval before proceeding with implementation. - """ - - file_path = __file__ - - def __init__(self, **kwargs): - kwargs.setdefault("raise_exception", False) - super().__init__(**kwargs) - - def build_tool_call(self) -> ToolCall: - """Build and return the tool call schema for this operator.""" - tool_params = { - "name": "ExitPlanMode", - "description": self.get_prompt("tool_desc"), - "input_schema": { - "plan": { - "type": "string", - "description": self.get_prompt("plan"), - "required": True, - }, - }, - } - - return ToolCall(**tool_params) - - async def async_execute(self): - """Execute the exit plan mode operation.""" - plan: str = self.input_dict.get("plan", "").strip() - - if not plan: - raise ValueError("The 'plan' parameter cannot be empty.") - - result_msg = ( - f"Plan presented for approval:\n\n{plan}\n\nWaiting for user confirmation to proceed with implementation." - ) - self.set_output(result_msg) - - async def async_default_execute(self, e: Exception = None, **kwargs): - """Fill outputs with a default failure message when execution fails.""" - plan: str = self.input_dict.get("plan", "").strip() - error_msg = f'Failed to present plan "{plan[:50]}..."' if len(plan) > 50 else f'Failed to present plan "{plan}"' - if e: - error_msg += f": {str(e)}" - self.set_output(error_msg) diff --git a/flowllm/extensions/file_tool/exit_plan_mode_prompt.yaml b/flowllm/extensions/file_tool/exit_plan_mode_prompt.yaml deleted file mode 100644 index 75d4d4a..0000000 --- a/flowllm/extensions/file_tool/exit_plan_mode_prompt.yaml +++ /dev/null @@ -1,11 +0,0 @@ -tool_desc: | - Use this tool when you are in plan mode and have finished presenting your plan and are ready to code. This will prompt the user to exit plan mode. - IMPORTANT: Only use this tool when the task requires planning the implementation steps of a task that requires writing code. For research tasks where you're gathering information, searching files, reading files or in general trying to understand the codebase - do NOT use this tool. - - Examples: - 1. Initial task: "Search for and understand the implementation of vim mode in the codebase" - Do not use the exit plan mode tool because you are not planning the implementation steps of a task. - 2. Initial task: "Help me implement yank mode for vim" - Use the exit plan mode tool after you have finished planning the implementation steps of the task. - -plan: | - The plan you came up with, that you want to run by the user for approval. Supports markdown. The plan should be pretty concise. - diff --git a/flowllm/extensions/file_tool/glob_op.py b/flowllm/extensions/file_tool/glob_op.py deleted file mode 100644 index 6f1aeb5..0000000 --- a/flowllm/extensions/file_tool/glob_op.py +++ /dev/null @@ -1,296 +0,0 @@ -"""Glob file search operation module. - -This module provides a tool operation for finding files matching glob patterns. -It enables efficient file discovery based on name or path structure, especially -in large codebases. Files are sorted by modification time (newest first). -""" - -import fnmatch -import time -from pathlib import Path -from typing import List, Optional - -from loguru import logger -from pathspec import PathSpec -from pathspec.patterns.gitwildmatch import GitIgnorePattern - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - - -@C.register_op() -class GlobOp(BaseAsyncToolOp): - """Glob file search operation. - - This operation efficiently finds files matching specific glob patterns, - returning absolute paths sorted by modification time (newest first). - Supports gitignore patterns for filtering files. - """ - - file_path = __file__ - - def __init__(self, gitignore_patterns: List[str] = None, **kwargs): - kwargs.setdefault("raise_exception", False) - super().__init__(**kwargs) - self.gitignore_patterns = gitignore_patterns - - def build_tool_call(self) -> ToolCall: - """Build and return the tool call schema for this operator.""" - tool_params = { - "name": "FindFiles", - "description": self.get_prompt("tool_desc"), - "input_schema": { - "pattern": { - "type": "string", - "description": self.get_prompt("pattern"), - "required": True, - }, - "dir_path": { - "type": "string", - "description": self.get_prompt("dir_path"), - "required": False, - }, - "case_sensitive": { - "type": "boolean", - "description": self.get_prompt("case_sensitive"), - "required": False, - }, - }, - } - - return ToolCall(**tool_params) - - def _should_ignore_file( - self, - file_path: Path, - root_dir: Path, - ) -> bool: - """Check if a file should be ignored based on ignore patterns. - - Args: - file_path: Path to the file to check. - root_dir: Root directory for resolving relative paths. - - Returns: - True if the file should be ignored, False otherwise. - """ - if not self.gitignore_patterns: - return False - - ignore_spec = PathSpec.from_lines(GitIgnorePattern, self.gitignore_patterns) - - # Get relative path from root_dir - relative_path = file_path.relative_to(root_dir) - - # Check if file matches ignore patterns - return ignore_spec.match_file(str(relative_path)) - - async def async_execute(self): - """Execute the glob search operation.""" - pattern: str = self.input_dict.get("pattern", "").strip() - dir_path: Optional[str] = self.input_dict.get("dir_path") - case_sensitive: bool = self.input_dict.get("case_sensitive", False) - - # Validate pattern - if not pattern: - error_msg = "The 'pattern' parameter cannot be empty." - logger.error(f"{self.name}: {error_msg}") - self.set_output(error_msg) - return - - # Determine search directory - if dir_path: - search_dir = Path(dir_path).expanduser().resolve() - if not search_dir.exists(): - error_msg = f"Search path does not exist: {search_dir}" - logger.error(f"{self.name}: {error_msg}") - self.set_output(error_msg) - return - if not search_dir.is_dir(): - error_msg = f"Search path is not a directory: {search_dir}" - logger.error(f"{self.name}: {error_msg}") - self.set_output(error_msg) - return - else: - search_dir = Path.cwd() - - # Collect matching files - all_entries: List[Path] = [] - ignored_count = 0 - - # Check if pattern is an exact file path - full_path = search_dir / pattern - if full_path.exists() and full_path.is_file(): - # Use exact match - if not self._should_ignore_file( - full_path, - search_dir, - ): - all_entries.append(full_path) - else: - # Use glob pattern matching - matching_files = self._glob_match( - search_dir, - pattern, - case_sensitive=case_sensitive, - ) - - # Filter by ignore patterns - for file_path in matching_files: - if self._should_ignore_file( - file_path, - search_dir, - ): - ignored_count += 1 - else: - all_entries.append(file_path) - - # Check if any files found - if not all_entries: - message = f'No files found matching pattern "{pattern}"' - if dir_path: - message += f" within {search_dir}" - if ignored_count > 0: - message += f" ({ignored_count} files were ignored)" - self.set_output(message) - return - - # Sort files by modification time - now_timestamp = time.time() - # recency_threshold_days: Number of days to consider a file "recent" for sorting (default: 1). - recency_threshold_ms = self.op_params.get("recency_threshold_days", 1) * 24 * 60 * 60 * 1000 - - def get_sort_key(path: Path) -> tuple: - mtime = path.stat().st_mtime - mtime_ms = mtime * 1000 - is_recent = (now_timestamp * 1000) - mtime_ms < recency_threshold_ms - - if is_recent: - # Recent files: sort by mtime descending (newest first) - return 0, -mtime_ms - else: - # Old files: sort alphabetically - return 1, str(path) - - sorted_entries = sorted(all_entries, key=get_sort_key) - - # Format results - sorted_absolute_paths = [str(entry.resolve()) for entry in sorted_entries] - file_list_description = "\n".join(sorted_absolute_paths) - file_count = len(sorted_absolute_paths) - - result_message = f'Found {file_count} file(s) matching "{pattern}"' - if dir_path: - result_message += f" within {search_dir}" - if ignored_count > 0: - result_message += f" ({ignored_count} additional files were ignored)" - result_message += ", sorted by modification time (newest first):\n" - result_message += file_list_description - - self.set_output(result_message) - - async def async_default_execute(self, e: Exception = None, **kwargs): - """Fill outputs with a default failure message when execution fails.""" - pattern: str = self.input_dict.get("pattern", "").strip() - error_msg = f'Failed to search files matching pattern "{pattern}"' - if e: - error_msg += f": {str(e)}" - self.set_output(error_msg) - - def _glob_match( - self, - root_dir: Path, - pattern: str, - case_sensitive: bool = False, - ) -> List[Path]: - """Match files using glob pattern. - - Args: - root_dir: Root directory to search in. - pattern: Glob pattern to match. - case_sensitive: Whether matching should be case-sensitive. - - Returns: - List of matching file paths. - """ - matching_files: List[Path] = [] - - # Handle ** pattern (recursive match) - if "**" in pattern: - # Split pattern into parts - parts = pattern.split("**", 1) - prefix = parts[0].rstrip("/") - suffix = parts[1] if len(parts) > 1 else "" - - # Walk through directory tree - for path in root_dir.rglob("*"): - if not path.is_file(): - continue - - rel_path = path.relative_to(root_dir) - rel_str = str(rel_path).replace("\\", "/") - - # Check if path matches pattern - if self._matches_glob_pattern( - rel_str, - prefix, - suffix, - case_sensitive, - ): - matching_files.append(path) - else: - # For patterns without **, use fnmatch for better compatibility - # Walk through directory tree and match manually - pattern_normalized = pattern.replace("\\", "/") - for path in root_dir.rglob("*"): - if not path.is_file(): - continue - - rel_path = path.relative_to(root_dir) - rel_str = str(rel_path).replace("\\", "/") - - # Match using fnmatch - if case_sensitive: - if fnmatch.fnmatch(rel_str, pattern_normalized): - matching_files.append(path) - else: - if fnmatch.fnmatch(rel_str.lower(), pattern_normalized.lower()): - matching_files.append(path) - - return matching_files - - @staticmethod - def _matches_glob_pattern( - path_str: str, - prefix: str, - suffix: str, - case_sensitive: bool, - ) -> bool: - """Check if a path matches a glob pattern with **. - - Args: - path_str: Path string to check (relative to root). - prefix: Prefix pattern before **. - suffix: Suffix pattern after **. - case_sensitive: Whether matching should be case-sensitive. - - Returns: - True if path matches pattern, False otherwise. - """ - if not case_sensitive: - path_str = path_str.lower() - prefix = prefix.lower() - suffix = suffix.lower() - - # Check prefix - if prefix: - if not path_str.startswith(prefix): - return False - - # Check suffix - if suffix: - if not fnmatch.fnmatch(path_str, f"*{suffix}"): - return False - - return True diff --git a/flowllm/extensions/file_tool/glob_prompt.yaml b/flowllm/extensions/file_tool/glob_prompt.yaml deleted file mode 100644 index 5912bdc..0000000 --- a/flowllm/extensions/file_tool/glob_prompt.yaml +++ /dev/null @@ -1,11 +0,0 @@ -tool_desc: | - Efficiently finds files matching specific glob patterns (e.g., `src/**/*.ts`, `**/*.md`), returning absolute paths sorted by modification time (newest first). Ideal for quickly locating files based on their name or path structure, especially in large codebases. - -pattern: | - The glob pattern to match against (e.g., '**/*.py', 'docs/*.md'). - -dir_path: | - Optional: The absolute path to the directory to search within. If omitted, searches the current working directory. - -case_sensitive: | - Optional: Whether the search should be case-sensitive. Defaults to false. \ No newline at end of file diff --git a/flowllm/extensions/file_tool/grep_op.py b/flowllm/extensions/file_tool/grep_op.py deleted file mode 100644 index bcc0eb9..0000000 --- a/flowllm/extensions/file_tool/grep_op.py +++ /dev/null @@ -1,174 +0,0 @@ -"""Grep text search operation module. - -This module provides a tool operation for searching text patterns in files. -It enables efficient content-based search using regular expressions, with support -for glob pattern filtering and result limiting. -""" - -import fnmatch -import re -from pathlib import Path -from typing import List, Optional - -from loguru import logger - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - - -@C.register_op() -class GrepOp(BaseAsyncToolOp): - """Grep text search operation. - - This operation searches for text patterns in files using regular expressions. - Supports glob pattern filtering and result limiting. - """ - - file_path = __file__ - - def __init__(self, **kwargs): - kwargs.setdefault("raise_exception", False) - super().__init__(**kwargs) - - def build_tool_call(self) -> ToolCall: - """Build and return the tool call schema for this operator.""" - tool_params = { - "name": "Grep", - "description": self.get_prompt("tool_desc"), - "input_schema": { - "pattern": { - "type": "string", - "description": self.get_prompt("pattern"), - "required": True, - }, - "path": { - "type": "string", - "description": self.get_prompt("path"), - "required": False, - }, - "glob": { - "type": "string", - "description": self.get_prompt("glob"), - "required": False, - }, - "limit": { - "type": "number", - "description": self.get_prompt("limit"), - "required": False, - }, - }, - } - - return ToolCall(**tool_params) - - async def async_execute(self): - """Execute the grep search operation.""" - pattern: str = self.input_dict.get("pattern", "").strip() - path: Optional[str] = self.input_dict.get("path") - glob_pattern: Optional[str] = self.input_dict.get("glob") - limit: Optional[int] = self.input_dict.get("limit") - - # Validate pattern - if not pattern: - raise ValueError("The 'pattern' parameter cannot be empty.") - - # Validate regex pattern - regex = re.compile(pattern, re.IGNORECASE) - - # Determine search directory - if path: - search_dir = Path(path).expanduser().resolve() - if not search_dir.exists(): - raise ValueError(f"Search path does not exist: {search_dir}") - if not search_dir.is_dir(): - raise ValueError(f"Search path is not a directory: {search_dir}") - else: - search_dir = Path.cwd() - - # Collect matching files based on glob pattern - files_to_search: List[Path] = [] - if glob_pattern: - # Use glob pattern to filter files - glob_normalized = glob_pattern.replace("\\", "/") - for file_path in search_dir.rglob("*"): - if not file_path.is_file(): - continue - rel_path = file_path.relative_to(search_dir) - rel_str = str(rel_path).replace("\\", "/") - if fnmatch.fnmatch(rel_str.lower(), glob_normalized.lower()): - files_to_search.append(file_path) - else: - # Search all files recursively - files_to_search = [f for f in search_dir.rglob("*") if f.is_file()] - - # Search for matches - matches: List[dict] = [] - for file_path in files_to_search: - try: - content = file_path.read_text(encoding="utf-8", errors="ignore") - lines = content.split("\n") - for line_num, line in enumerate(lines, start=1): - if regex.search(line): - try: - relative_path = file_path.relative_to(search_dir) - except ValueError: - relative_path = file_path.name - matches.append( - { - "file_path": str(relative_path), - "line_number": line_num, - "line": line, - }, - ) - if limit and len(matches) >= limit: - break - if limit and len(matches) >= limit: - break - except Exception as e: - logger.debug(f"Could not read {file_path}: {str(e)}") - continue - - # Format results - if not matches: - search_location = f'in path "{path}"' if path else "in the workspace directory" - filter_desc = f' (filter: "{glob_pattern}")' if glob_pattern else "" - result_msg = f'No matches found for pattern "{pattern}" {search_location}{filter_desc}.' - self.set_output(result_msg) - return - - # Group matches by file - matches_by_file = {} - for match in matches: - file_key = match["file_path"] - if file_key not in matches_by_file: - matches_by_file[file_key] = [] - matches_by_file[file_key].append(match) - - # Build output - total_matches = len(matches) - match_term = "match" if total_matches == 1 else "matches" - search_location = f'in path "{path}"' if path else "in the workspace directory" - filter_desc = f' (filter: "{glob_pattern}")' if glob_pattern else "" - - output_lines = [ - f'Found {total_matches} {match_term} for pattern "{pattern}" {search_location}{filter_desc}:\n---', - ] - - for file_path in sorted(matches_by_file.keys()): - output_lines.append(f"File: {file_path}") - for match in sorted(matches_by_file[file_path], key=lambda x: x["line_number"]): - trimmed_line = match["line"].strip() - output_lines.append(f"L{match['line_number']}: {trimmed_line}") - output_lines.append("---") - - result_msg = "\n".join(output_lines) - self.set_output(result_msg) - - async def async_default_execute(self, e: Exception = None, **kwargs): - """Fill outputs with a default failure message when execution fails.""" - pattern: str = self.input_dict.get("pattern", "").strip() - error_msg = f'Failed to search for pattern "{pattern}"' - if e: - error_msg += f": {str(e)}" - self.set_output(error_msg) diff --git a/flowllm/extensions/file_tool/grep_prompt.yaml b/flowllm/extensions/file_tool/grep_prompt.yaml deleted file mode 100644 index 0c9a256..0000000 --- a/flowllm/extensions/file_tool/grep_prompt.yaml +++ /dev/null @@ -1,15 +0,0 @@ -tool_desc: | - A powerful search tool for finding patterns in files using regular expressions. Supports full regex syntax (e.g., "log.*Error", "function\\s+\\w+"), glob pattern filtering, and result limiting. Ideal for searching code or text content across multiple files. - -pattern: | - The regular expression pattern to search for in file contents. - -path: | - Optional: The directory to search in. Defaults to current working directory. - -glob: | - Optional: Glob pattern to filter files (e.g., "*.js", "*.{ts,tsx}"). - -limit: | - Optional: Maximum number of matching lines to return. Shows all matches if not specified. - diff --git a/flowllm/extensions/file_tool/ls_op.py b/flowllm/extensions/file_tool/ls_op.py deleted file mode 100644 index bee9898..0000000 --- a/flowllm/extensions/file_tool/ls_op.py +++ /dev/null @@ -1,129 +0,0 @@ -"""List directory operation module. - -This module provides a tool operation for listing files and subdirectories -in a specified directory path. -""" - -import fnmatch -from pathlib import Path -from typing import List, Optional - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - - -@C.register_op() -class LSOp(BaseAsyncToolOp): - """List directory operation. - - This operation lists the names of files and subdirectories directly - within a specified directory path. Can optionally ignore entries - matching provided glob patterns. - """ - - file_path = __file__ - - def __init__(self, **kwargs): - kwargs.setdefault("raise_exception", False) - super().__init__(**kwargs) - - def build_tool_call(self) -> ToolCall: - """Build and return the tool call schema for this operator.""" - tool_params = { - "name": "ListDirectory", - "description": self.get_prompt("tool_desc"), - "input_schema": { - "path": { - "type": "string", - "description": self.get_prompt("path"), - "required": True, - }, - "ignore": { - "type": "array", - "items": {"type": "string"}, - "description": self.get_prompt("ignore"), - "required": False, - }, - }, - } - - return ToolCall(**tool_params) - - def _should_ignore(self, name: str, patterns: Optional[List[str]]) -> bool: - """Check if a filename matches any of the ignore patterns. - - Args: - name: Filename to check. - patterns: Array of glob patterns to check against. - - Returns: - True if the filename should be ignored, False otherwise. - """ - if not patterns: - return False - - for pattern in patterns: - if fnmatch.fnmatch(name, pattern): - return True - return False - - async def async_execute(self): - """Execute the list directory operation.""" - path_str: str = self.input_dict.get("path", "").strip() - ignore: Optional[List[str]] = self.input_dict.get("ignore") - - if not path_str: - raise ValueError("The 'path' parameter cannot be empty.") - - dir_path = Path(path_str).expanduser().resolve() - - if not dir_path.exists(): - raise FileNotFoundError(f"Directory not found: {dir_path}") - - if not dir_path.is_dir(): - raise ValueError(f"Path is not a directory: {dir_path}") - - entries = [] - for item in dir_path.iterdir(): - if self._should_ignore(item.name, ignore): - continue - - try: - stat = item.stat() - is_dir = item.is_dir() - entries.append( - { - "name": item.name, - "path": str(item.resolve()), - "isDirectory": is_dir, - "size": 0 if is_dir else stat.st_size, - "modifiedTime": stat.st_mtime, - }, - ) - except OSError: - # Skip files that can't be accessed - continue - - # Sort entries (directories first, then alphabetically) - entries.sort(key=lambda x: (not x["isDirectory"], x["name"])) - - if not entries: - self.set_output(f"Directory {dir_path} is empty.") - return - - # Format output - directory_content = "\n".join( - f'[DIR] {entry["name"]}' if entry["isDirectory"] else entry["name"] for entry in entries - ) - - result_message = f"Directory listing for {dir_path}:\n{directory_content}" - self.set_output(result_message) - - async def async_default_execute(self, e: Exception = None, **kwargs): - """Fill outputs with a default failure message when execution fails.""" - path_str: str = self.input_dict.get("path", "").strip() - error_msg = f'Failed to list directory "{path_str}"' - if e: - error_msg += f": {str(e)}" - self.set_output(error_msg) diff --git a/flowllm/extensions/file_tool/ls_prompt.yaml b/flowllm/extensions/file_tool/ls_prompt.yaml deleted file mode 100644 index 6f72c03..0000000 --- a/flowllm/extensions/file_tool/ls_prompt.yaml +++ /dev/null @@ -1,9 +0,0 @@ -tool_desc: | - Lists the names of files and subdirectories directly within a specified directory path. Can optionally ignore entries matching provided glob patterns. - -path: | - The absolute path to the directory to list (e.g., '/home/user/project/src'). Must be an absolute path. - -ignore: | - Optional: Array of glob patterns to ignore (e.g., ['*.pyc', '__pycache__', '*.log']). - diff --git a/flowllm/extensions/file_tool/read_file_op.py b/flowllm/extensions/file_tool/read_file_op.py deleted file mode 100644 index 46c9d51..0000000 --- a/flowllm/extensions/file_tool/read_file_op.py +++ /dev/null @@ -1,120 +0,0 @@ -"""Read file operation module. - -This module provides a tool operation for reading file contents. -It supports reading entire files or specific line ranges for large files. -""" - -from pathlib import Path -from typing import Optional - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - - -@C.register_op() -class ReadFileOp(BaseAsyncToolOp): - """Read file operation. - - This operation reads and returns the content of a specified file. - For text files, it can read specific line ranges using offset and limit. - """ - - file_path = __file__ - - def __init__(self, **kwargs): - kwargs.setdefault("raise_exception", False) - super().__init__(**kwargs) - - def build_tool_call(self) -> ToolCall: - """Build and return the tool call schema for this operator.""" - tool_params = { - "name": "ReadFile", - "description": self.get_prompt("tool_desc"), - "input_schema": { - "absolute_path": { - "type": "string", - "description": self.get_prompt("absolute_path"), - "required": True, - }, - "offset": { - "type": "number", - "description": self.get_prompt("offset"), - "required": False, - }, - "limit": { - "type": "number", - "description": self.get_prompt("limit"), - "required": False, - }, - }, - } - - return ToolCall(**tool_params) - - async def async_execute(self): - """Execute the read file operation.""" - absolute_path: str = self.input_dict.get("absolute_path", "").strip() - offset: Optional[int] = self.input_dict.get("offset") - limit: Optional[int] = self.input_dict.get("limit") - - # Validate absolute_path - if not absolute_path: - raise ValueError("The 'absolute_path' parameter cannot be empty.") - - # Resolve file path - file_path_obj = Path(absolute_path).expanduser().resolve() - - # Check if file exists - if not file_path_obj.exists(): - raise FileNotFoundError(f"File not found: {file_path_obj}") - - if not file_path_obj.is_file(): - raise ValueError(f"Path is not a file: {file_path_obj}") - - # Read file content - content = file_path_obj.read_text(encoding="utf-8") - lines = content.split("\n") - - # Handle line range if specified - if offset is not None or limit is not None: - if offset is None: - offset = 0 - if limit is None: - limit = len(lines) - - # Validate offset and limit - if offset < 0: - raise ValueError("Offset must be a non-negative number") - if limit <= 0: - raise ValueError("Limit must be a positive number") - - total_lines = len(lines) - start = offset - end = min(offset + limit, total_lines) - - if start >= total_lines: - raise ValueError( - f"Offset {offset} is beyond file length ({total_lines} lines)", - ) - - selected_lines = lines[start:end] - result_content = "\n".join(selected_lines) - - # Format output with range information - if end < total_lines: - result = f"Showing lines {start}-{end - 1} of {total_lines} total lines.\n\n---\n\n{result_content}" - else: - result = result_content - else: - result = content - - self.set_output(result) - - async def async_default_execute(self, e: Exception = None, **kwargs): - """Fill outputs with a default failure message when execution fails.""" - absolute_path: str = self.input_dict.get("absolute_path", "").strip() - error_msg = f'Failed to read file "{absolute_path}"' - if e: - error_msg += f": {str(e)}" - self.set_output(error_msg) diff --git a/flowllm/extensions/file_tool/read_file_prompt.yaml b/flowllm/extensions/file_tool/read_file_prompt.yaml deleted file mode 100644 index 52f6a46..0000000 --- a/flowllm/extensions/file_tool/read_file_prompt.yaml +++ /dev/null @@ -1,12 +0,0 @@ -tool_desc: | - Reads and returns the content of a specified file. For text files, it can read specific line ranges using the 'offset' and 'limit' parameters. Use offset and limit to paginate through large files. - -absolute_path: | - The absolute path to the file to read (e.g., '/home/user/project/file.txt'). Relative paths are not supported. You must provide an absolute path. - -offset: | - Optional: For text files, the 0-based line number to start reading from. Requires 'limit' to be set. Use for paginating through large files. - -limit: | - Optional: For text files, maximum number of lines to read. Use with 'offset' to paginate through large files. If omitted when offset is provided, reads from offset to the end of the file. - diff --git a/flowllm/extensions/file_tool/read_many_files_op.py b/flowllm/extensions/file_tool/read_many_files_op.py deleted file mode 100644 index 8b998cd..0000000 --- a/flowllm/extensions/file_tool/read_many_files_op.py +++ /dev/null @@ -1,232 +0,0 @@ -"""Read many files operation module. - -This module provides a tool operation for reading content from multiple files -specified by glob patterns. It concatenates the content with separators. -""" - -import fnmatch -from pathlib import Path -from typing import List, Optional, Set - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - -DEFAULT_OUTPUT_SEPARATOR_FORMAT = "--- {filePath} ---" -DEFAULT_OUTPUT_TERMINATOR = "\n--- End of content ---" - - -@C.register_op() -class ReadManyFilesOp(BaseAsyncToolOp): - """Read many files operation. - - This operation reads content from multiple files matching glob patterns - and concatenates them with separators. - """ - - file_path = __file__ - - def __init__(self, **kwargs): - kwargs.setdefault("raise_exception", False) - super().__init__(**kwargs) - - def build_tool_call(self) -> ToolCall: - """Build and return the tool call schema for this operator.""" - tool_params = { - "name": "ReadManyFiles", - "description": self.get_prompt("tool_desc"), - "input_schema": { - "paths": { - "type": "array", - "items": {"type": "string"}, - "description": self.get_prompt("paths"), - "required": True, - }, - "include": { - "type": "array", - "items": {"type": "string"}, - "description": self.get_prompt("include"), - "required": False, - }, - "exclude": { - "type": "array", - "items": {"type": "string"}, - "description": self.get_prompt("exclude"), - "required": False, - }, - "dir_path": { - "type": "string", - "description": self.get_prompt("dir_path"), - "required": False, - }, - }, - } - - return ToolCall(**tool_params) - - def _glob_match( - self, - root_dir: Path, - pattern: str, - case_sensitive: bool = False, - ) -> List[Path]: - """Match files using glob pattern.""" - matching_files: List[Path] = [] - - # Check if pattern is an exact file path - full_path = root_dir / pattern - if full_path.exists() and full_path.is_file(): - matching_files.append(full_path) - return matching_files - - # Handle ** pattern (recursive match) - if "**" in pattern: - parts = pattern.split("**", 1) - prefix = parts[0].rstrip("/") - suffix = parts[1] if len(parts) > 1 else "" - - for path in root_dir.rglob("*"): - if not path.is_file(): - continue - - rel_path = path.relative_to(root_dir) - rel_str = str(rel_path).replace("\\", "/") - - if self._matches_glob_pattern(rel_str, prefix, suffix, case_sensitive): - matching_files.append(path) - else: - pattern_normalized = pattern.replace("\\", "/") - for path in root_dir.rglob("*"): - if not path.is_file(): - continue - - rel_path = path.relative_to(root_dir) - rel_str = str(rel_path).replace("\\", "/") - - if case_sensitive: - if fnmatch.fnmatch(rel_str, pattern_normalized): - matching_files.append(path) - else: - if fnmatch.fnmatch(rel_str.lower(), pattern_normalized.lower()): - matching_files.append(path) - - return matching_files - - @staticmethod - def _matches_glob_pattern( - path_str: str, - prefix: str, - suffix: str, - case_sensitive: bool, - ) -> bool: - """Check if a path matches a glob pattern with **.""" - if not case_sensitive: - path_str = path_str.lower() - prefix = prefix.lower() - suffix = suffix.lower() - - if prefix and not path_str.startswith(prefix): - return False - - if suffix and not fnmatch.fnmatch(path_str, f"*{suffix}"): - return False - - return True - - def _should_exclude(self, file_path: Path, exclude_patterns: List[str], search_dir: Path) -> bool: - """Check if a file should be excluded based on patterns.""" - if not exclude_patterns: - return False - - rel_path = file_path.relative_to(search_dir) - rel_str = str(rel_path).replace("\\", "/") - - for pattern in exclude_patterns: - if "**" in pattern: - parts = pattern.split("**", 1) - prefix = parts[0].rstrip("/") - suffix = parts[1] if len(parts) > 1 else "" - if self._matches_glob_pattern(rel_str, prefix, suffix, False): - return True - else: - if fnmatch.fnmatch(rel_str.lower(), pattern.lower()): - return True - - return False - - async def async_execute(self): - """Execute the read many files operation.""" - paths: List[str] = self.input_dict.get("paths", []) - include: List[str] = self.input_dict.get("include", []) - exclude: List[str] = self.input_dict.get("exclude", []) - dir_path: Optional[str] = self.input_dict.get("dir_path") - - if not paths: - raise ValueError("The 'paths' parameter cannot be empty.") - - # Determine search directory - if dir_path: - search_dir = Path(dir_path).expanduser().resolve() - if not search_dir.exists(): - raise FileNotFoundError(f"Search path does not exist: {search_dir}") - if not search_dir.is_dir(): - raise ValueError(f"Search path is not a directory: {search_dir}") - else: - search_dir = Path.cwd() - - # Collect all patterns - all_patterns = paths + include - all_files: Set[Path] = set() - - # Find matching files - for pattern in all_patterns: - matching_files = self._glob_match(search_dir, pattern, case_sensitive=False) - for file_path in matching_files: - if not self._should_exclude(file_path, exclude, search_dir): - all_files.add(file_path) - - if not all_files: - self.set_output("No files found matching the specified patterns.") - return - - # Sort files alphabetically - sorted_files = sorted(all_files) - - # Read and concatenate file contents - content_parts: List[str] = [] - skipped_files: List[tuple] = [] - - for file_path in sorted_files: - try: - content = file_path.read_text(encoding="utf-8") - separator = DEFAULT_OUTPUT_SEPARATOR_FORMAT.replace( - "{filePath}", - str(file_path), - ) - content_parts.append(f"{separator}\n\n{content}\n\n") - except Exception as e: - skipped_files.append((str(file_path), str(e))) - - if content_parts: - content_parts.append(DEFAULT_OUTPUT_TERMINATOR) - result = "".join(content_parts) - else: - result = "No files could be read successfully." - - # Add summary if there are skipped files - if skipped_files: - result += f"\n\nSkipped {len(skipped_files)} file(s):\n" - for file_path, reason in skipped_files[:5]: - result += f"- {file_path} (Reason: {reason})\n" - if len(skipped_files) > 5: - result += f"- ...and {len(skipped_files) - 5} more.\n" - - self.set_output(result) - - async def async_default_execute(self, e: Exception = None, **kwargs): - """Fill outputs with a default failure message when execution fails.""" - paths: List[str] = self.input_dict.get("paths", []) - error_msg = f"Failed to read files matching patterns {paths}" - if e: - error_msg += f": {str(e)}" - self.set_output(error_msg) diff --git a/flowllm/extensions/file_tool/read_many_files_prompt.yaml b/flowllm/extensions/file_tool/read_many_files_prompt.yaml deleted file mode 100644 index 339bf90..0000000 --- a/flowllm/extensions/file_tool/read_many_files_prompt.yaml +++ /dev/null @@ -1,15 +0,0 @@ -tool_desc: | - Reads content from multiple files specified by glob patterns or paths and concatenates them into a single output. Files are separated by "--- {filePath} ---" markers. Useful for reading multiple related files at once, such as all files in a directory, all files matching a pattern, or a collection of configuration files. - -paths: | - Required. An array of glob patterns or file paths to read. Examples: ['src/**/*.ts'], ['README.md', 'docs/'], ['*.py']. Paths are relative to the search directory (or current working directory if not specified). - -include: | - Optional. Additional glob patterns to include. These are merged with 'paths'. Example: ['*.test.ts'] to add test files. - -exclude: | - Optional. Glob patterns for files/directories to exclude. Example: ['**/*.log', 'temp/']. - -dir_path: | - Optional. The absolute path to the directory to search within. If omitted, searches the current working directory. - diff --git a/flowllm/extensions/file_tool/rip_grep_op.py b/flowllm/extensions/file_tool/rip_grep_op.py deleted file mode 100644 index 79592ae..0000000 --- a/flowllm/extensions/file_tool/rip_grep_op.py +++ /dev/null @@ -1,156 +0,0 @@ -"""Ripgrep text search operation module. - -This module provides a tool operation for searching text patterns in files -using ripgrep (rg). It enables efficient content-based search using regular -expressions with support for glob pattern filtering and result limiting. -""" - -import asyncio -from pathlib import Path -from typing import Optional - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - - -@C.register_op() -class RipGrepOp(BaseAsyncToolOp): - """Ripgrep text search operation. - - This operation searches for text patterns in files using ripgrep (rg). - Supports glob pattern filtering and result limiting. - """ - - file_path = __file__ - - def __init__(self, **kwargs): - kwargs.setdefault("raise_exception", False) - super().__init__(**kwargs) - - def build_tool_call(self) -> ToolCall: - """Build and return the tool call schema for this operator.""" - tool_params = { - "name": "RipGrep", - "description": self.get_prompt("tool_desc"), - "input_schema": { - "pattern": { - "type": "string", - "description": self.get_prompt("pattern"), - "required": True, - }, - "path": { - "type": "string", - "description": self.get_prompt("path"), - "required": False, - }, - "glob": { - "type": "string", - "description": self.get_prompt("glob"), - "required": False, - }, - "limit": { - "type": "number", - "description": self.get_prompt("limit"), - "required": False, - }, - }, - } - - return ToolCall(**tool_params) - - async def async_execute(self): - """Execute the ripgrep search operation.""" - pattern: str = self.input_dict.get("pattern", "").strip() - path: Optional[str] = self.input_dict.get("path") - glob_pattern: Optional[str] = self.input_dict.get("glob") - limit: Optional[int] = self.input_dict.get("limit") - - # Validate pattern - if not pattern: - raise ValueError("The 'pattern' parameter cannot be empty.") - - # Determine search path - if path: - search_path = Path(path).expanduser().resolve() - if not search_path.exists(): - raise ValueError(f"Search path does not exist: {search_path}") - else: - search_path = Path.cwd() - - # Build ripgrep command - rg_args = [ - "rg", - "--line-number", - "--no-heading", - "--with-filename", - "--ignore-case", - "--regexp", - pattern, - ] - - # Add glob pattern if provided - if glob_pattern: - rg_args.extend(["--glob", glob_pattern]) - - # Add search path - rg_args.append(str(search_path)) - - # Execute ripgrep - process = await asyncio.create_subprocess_exec( - *rg_args, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - stdout, stderr = await process.communicate() - - # Handle ripgrep exit codes - if process.returncode == 0: - raw_output = stdout.decode("utf-8").strip() - elif process.returncode == 1: - # No matches found - raw_output = "" - else: - error_msg = stderr.decode("utf-8").strip() - raise ValueError(f"ripgrep exited with code {process.returncode}: {error_msg}") - - # Build search description - search_location = f'in path "{path}"' if path else "in the workspace directory" - filter_desc = f' (filter: "{glob_pattern}")' if glob_pattern else "" - - # Check if we have any matches - if not raw_output: - result_msg = f'No matches found for pattern "{pattern}" {search_location}{filter_desc}.' - self.set_output(result_msg) - return - - # Split into lines and apply limit - all_lines = [line for line in raw_output.split("\n") if line.strip()] - total_matches = len(all_lines) - match_term = "match" if total_matches == 1 else "matches" - - # Apply limit if specified - lines_to_include = all_lines - truncated = False - if limit and len(all_lines) > limit: - lines_to_include = all_lines[:limit] - truncated = True - - # Build output - header = f'Found {total_matches} {match_term} for pattern "{pattern}" {search_location}{filter_desc}:\n---\n' - grep_output = "\n".join(lines_to_include) - - result_msg = header + grep_output - if truncated: - omitted = total_matches - len(lines_to_include) - result_msg += f"\n---\n[{omitted} {'line' if omitted == 1 else 'lines'} truncated] ..." - - self.set_output(result_msg) - - async def async_default_execute(self, e: Exception = None, **kwargs): - """Fill outputs with a default failure message when execution fails.""" - pattern: str = self.input_dict.get("pattern", "").strip() - error_msg = f'Failed to search for pattern "{pattern}"' - if e: - error_msg += f": {str(e)}" - self.set_output(error_msg) diff --git a/flowllm/extensions/file_tool/rip_grep_prompt.yaml b/flowllm/extensions/file_tool/rip_grep_prompt.yaml deleted file mode 100644 index 6faef04..0000000 --- a/flowllm/extensions/file_tool/rip_grep_prompt.yaml +++ /dev/null @@ -1,15 +0,0 @@ -tool_desc: | - A powerful search tool built on ripgrep for finding patterns in files using regular expressions. Supports full regex syntax (e.g., "log.*Error", "function\\s+\\w+"), glob pattern filtering, and result limiting. Ideal for efficient searching across large codebases. - -pattern: | - The regular expression pattern to search for in file contents. - -path: | - Optional: The file or directory to search in. Defaults to current working directory. - -glob: | - Optional: Glob pattern to filter files (e.g., "*.js", "*.{ts,tsx}"). - -limit: | - Optional: Maximum number of matching lines to return. Shows all matches if not specified. - diff --git a/flowllm/extensions/file_tool/shell_op.py b/flowllm/extensions/file_tool/shell_op.py deleted file mode 100644 index 5b71c8e..0000000 --- a/flowllm/extensions/file_tool/shell_op.py +++ /dev/null @@ -1,144 +0,0 @@ -"""Shell command execution operation module. - -This module provides a tool operation for executing shell commands. -It supports foreground and background execution, with optional directory -and description parameters. -""" - -import asyncio -from pathlib import Path -from typing import Optional - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - - -@C.register_op() -class ShellOp(BaseAsyncToolOp): - """Shell command execution operation. - - This operation executes shell commands and returns their output, - error messages, exit codes, and other execution details. - Supports both foreground and background execution. - """ - - file_path = __file__ - - def __init__(self, **kwargs): - kwargs.setdefault("raise_exception", False) - super().__init__(**kwargs) - - def build_tool_call(self) -> ToolCall: - """Build and return the tool call schema for this operator.""" - tool_params = { - "name": "ExecuteShell", - "description": self.get_prompt("tool_desc"), - "input_schema": { - "command": { - "type": "string", - "description": self.get_prompt("command"), - "required": True, - }, - "is_background": { - "type": "boolean", - "description": self.get_prompt("is_background"), - "required": True, - }, - "description": { - "type": "string", - "description": self.get_prompt("description"), - "required": False, - }, - "directory": { - "type": "string", - "description": self.get_prompt("directory"), - "required": False, - }, - }, - } - - return ToolCall(**tool_params) - - async def async_execute(self): - """Execute the shell command operation.""" - command: str = self.input_dict.get("command", "").strip() - is_background: bool = self.input_dict.get("is_background", False) - description: Optional[str] = self.input_dict.get("description") - directory: Optional[str] = self.input_dict.get("directory") - - # Validate command - if not command: - raise ValueError("The 'command' parameter cannot be empty.") - - # Determine working directory - if directory: - cwd = Path(directory).expanduser().resolve() - if not cwd.exists(): - raise FileNotFoundError(f"Directory does not exist: {cwd}") - if not cwd.is_dir(): - raise ValueError(f"Path is not a directory: {cwd}") - else: - cwd = Path.cwd() - - # Prepare command for execution - if is_background: - # For background execution, add & at the end if not present - if not command.rstrip().endswith("&"): - command = command.rstrip() + " &" - - # Execute command - if is_background: - # Background execution: start process and return immediately - process = await asyncio.create_subprocess_shell( - command, - cwd=str(cwd), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - shell=True, - ) - # Don't wait for background processes - pid = process.pid - result_message = f"Command started in background (PID: {pid})" - if description: - result_message += f": {description}" - result_message += f"\nCommand: {command}\nDirectory: {cwd}" - self.set_output(result_message) - else: - # Foreground execution: wait for completion - process = await asyncio.create_subprocess_shell( - command, - cwd=str(cwd), - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - shell=True, - ) - - stdout, stderr = await process.communicate() - - # Decode output - output = stdout.decode("utf-8", errors="replace") if stdout else "" - error_output = stderr.decode("utf-8", errors="replace") if stderr else "" - - # Build result message - result_parts = [ - f"Command: {command}", - f"Directory: {directory or '(current)'}", - f"Output: {output if output else '(empty)'}", - f"Error: {error_output if error_output else '(none)'}", - f"Exit Code: {process.returncode if process.returncode is not None else '(none)'}", - ] - - if description: - result_parts.insert(1, f"Description: {description}") - - result_message = "\n".join(result_parts) - self.set_output(result_message) - - async def async_default_execute(self, e: Exception = None, **kwargs): - """Fill outputs with a default failure message when execution fails.""" - command: str = self.input_dict.get("command", "").strip() - error_msg = f'Failed to execute shell command "{command}"' - if e: - error_msg += f": {str(e)}" - self.set_output(error_msg) diff --git a/flowllm/extensions/file_tool/shell_prompt.yaml b/flowllm/extensions/file_tool/shell_prompt.yaml deleted file mode 100644 index 3257c29..0000000 --- a/flowllm/extensions/file_tool/shell_prompt.yaml +++ /dev/null @@ -1,15 +0,0 @@ -tool_desc: | - Executes a shell command and returns its output, error messages, exit codes, and other execution details. Supports both foreground and background execution. For background execution, use is_background=true for long-running processes like development servers, watchers, or daemons. - -command: | - The exact shell command to execute. On Unix systems, commands are executed as `bash -c `. On Windows, commands are executed as `cmd.exe /c `. - -is_background: | - Whether to run the command in background. Default is false. Set to true for long-running processes like development servers, watchers, or daemons that should continue running without blocking further commands. - -description: | - Optional: Brief description of the command for the user. Be specific and concise. Ideally a single sentence. Can be up to 3 sentences for clarity. No line breaks. - -directory: | - Optional: The absolute path of the directory to run the command in. If not provided, the current working directory is used. Must be a directory within the workspace and must already exist. - diff --git a/flowllm/extensions/file_tool/smart_edit_op.py b/flowllm/extensions/file_tool/smart_edit_op.py deleted file mode 100644 index 5076fcf..0000000 --- a/flowllm/extensions/file_tool/smart_edit_op.py +++ /dev/null @@ -1,148 +0,0 @@ -"""Smart edit operation module. - -This module provides a tool operation for intelligently editing files by replacing text. -It supports exact matching, flexible matching (ignoring indentation), and regex-based matching. -""" - -from pathlib import Path - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall -from ..utils import ( - calculate_exact_replacement, - calculate_flexible_replacement, - calculate_regex_replacement, -) - - -@C.register_op() -class SmartEditOp(BaseAsyncToolOp): - """Smart edit operation. - - This operation intelligently replaces text within a file using multiple strategies: - exact matching, flexible matching (ignoring indentation), and regex-based matching. - """ - - file_path = __file__ - - def __init__(self, **kwargs): - kwargs.setdefault("raise_exception", False) - super().__init__(**kwargs) - - def build_tool_call(self) -> ToolCall: - """Build and return the tool call schema for this operator.""" - tool_params = { - "name": "SmartEdit", - "description": self.get_prompt("tool_desc"), - "input_schema": { - "file_path": { - "type": "string", - "description": self.get_prompt("file_path"), - "required": True, - }, - "old_string": { - "type": "string", - "description": self.get_prompt("old_string"), - "required": True, - }, - "new_string": { - "type": "string", - "description": self.get_prompt("new_string"), - "required": True, - }, - "instruction": { - "type": "string", - "description": self.get_prompt("instruction"), - "required": True, - }, - }, - } - - return ToolCall(**tool_params) - - async def async_execute(self): - """Execute the smart edit operation.""" - file_path: str = self.input_dict.get("file_path", "").strip() - old_string: str = self.input_dict.get("old_string", "") - new_string: str = self.input_dict.get("new_string", "") - instruction: str = self.input_dict.get("instruction", "").strip() - - if not file_path: - raise ValueError("The 'file_path' parameter cannot be empty.") - - if not instruction: - raise ValueError("The 'instruction' parameter cannot be empty.") - - file_path_obj = Path(file_path).expanduser().resolve() - file_exists = file_path_obj.exists() and file_path_obj.is_file() - - # Handle new file creation - if not old_string and not file_exists: - file_path_obj.parent.mkdir(parents=True, exist_ok=True) - file_path_obj.write_text(new_string, encoding="utf-8") - self.set_output(f"Created new file: {file_path_obj}") - return - - if not file_exists: - raise FileNotFoundError( - f"File not found: {file_path_obj}. Use an empty old_string to create a new file.", - ) - - if not old_string and file_exists: - raise ValueError( - f"Failed to edit. Attempted to create a file that already exists: {file_path_obj}", - ) - - # Read current content - current_content = file_path_obj.read_text(encoding="utf-8") - original_line_ending = "\r\n" if "\r\n" in current_content else "\n" - current_content = current_content.replace("\r\n", "\n") - - # Try replacement strategies - result = calculate_exact_replacement(current_content, old_string, new_string) - if not result: - result = calculate_flexible_replacement(current_content, old_string, new_string) - if not result: - result = calculate_regex_replacement(current_content, old_string, new_string) - - if not result: - raise ValueError( - f"Failed to edit, could not find the string to replace. " - f"0 occurrences found for old_string in {file_path_obj}. " - f"Ensure you're not escaping content incorrectly and check whitespace, indentation, and context.", - ) - - new_content, occurrences = result - - if occurrences == 0: - raise ValueError( - f"Failed to edit, 0 occurrences found for old_string in {file_path_obj}.", - ) - - if old_string == new_string: - raise ValueError( - f"No changes to apply. The old_string and new_string are identical in {file_path_obj}.", - ) - - # Restore original line endings - if original_line_ending == "\r\n": - new_content = new_content.replace("\n", "\r\n") - - # Write new content - file_path_obj.write_text(new_content, encoding="utf-8") - - if occurrences == 1: - result_msg = f"Successfully modified file: {file_path_obj} (1 replacement)." - else: - result_msg = f"Successfully modified file: {file_path_obj} ({occurrences} replacements)." - - self.set_output(result_msg) - - async def async_default_execute(self, e: Exception = None, **kwargs): - """Fill outputs with a default failure message when execution fails.""" - file_path: str = self.input_dict.get("file_path", "").strip() - error_msg = f'Failed to edit file "{file_path}"' - if e: - error_msg += f": {str(e)}" - self.set_output(error_msg) diff --git a/flowllm/extensions/file_tool/smart_edit_prompt.yaml b/flowllm/extensions/file_tool/smart_edit_prompt.yaml deleted file mode 100644 index 5b97b51..0000000 --- a/flowllm/extensions/file_tool/smart_edit_prompt.yaml +++ /dev/null @@ -1,39 +0,0 @@ -tool_desc: | - Intelligently replaces text within a file using multiple matching strategies (exact, flexible, regex-based). This tool requires providing significant context around the change to ensure precise targeting. Always use the ReadFile tool to examine the file's current content before attempting a text replacement. - - Expectation for required parameters: - 1. `file_path` MUST be an absolute path; otherwise an error will be thrown. - 2. `old_string` MUST be the exact literal text to replace (including all whitespace, indentation, newlines, and surrounding code etc.). - 3. `new_string` MUST be the exact literal text to replace `old_string` with (also including all whitespace, indentation, newlines, and surrounding code etc.). Ensure the resulting code is correct and idiomatic. - 4. `instruction` is the detailed instruction of what needs to be changed. Make it specific and detailed so developers or large language models can understand what needs to be changed. - 5. NEVER escape `old_string` or `new_string`, that would break the exact literal text requirement. - - **Important:** If ANY of the above are not satisfied, the tool will fail. CRITICAL for `old_string`: Must uniquely identify the single instance to change. Include at least 3 lines of context BEFORE and AFTER the target text, matching whitespace and indentation precisely. If this string matches multiple locations, or does not match exactly, the tool will fail. - - **Multiple replacements:** If there are multiple and ambiguous occurrences of the `old_string` in the file, the tool will also fail. - -file_path: | - The absolute path to the file to modify. Must start with '/'. - -old_string: | - The exact literal text to replace, preferably unescaped. Include at least 3 lines of context BEFORE and AFTER the target text, matching whitespace and indentation precisely. If this string is not the exact literal text (i.e. you escaped it) or does not match exactly, the tool will fail. - -new_string: | - The exact literal text to replace `old_string` with, preferably unescaped. Provide the EXACT text. Ensure the resulting code is correct and idiomatic. - -instruction: | - A clear, semantic instruction for the code change, acting as a high-quality prompt for an expert LLM assistant. It must be self-contained and explain the goal of the change. - - A good instruction should concisely answer: - 1. WHY is the change needed? (e.g., "To fix a bug where users can be null...") - 2. WHERE should the change happen? (e.g., "...in the 'renderUserProfile' function...") - 3. WHAT is the high-level change? (e.g., "...add a null check for the 'user' object...") - 4. WHAT is the desired outcome? (e.g., "...so that it displays a loading spinner instead of crashing.") - - **GOOD Example:** "In the 'calculateTotal' function, correct the sales tax calculation by updating the 'taxRate' constant from 0.05 to 0.075 to reflect the new regional tax laws." - - **BAD Examples:** - - "Change the text." (Too vague) - - "Fix the bug." (Doesn't explain the bug or the fix) - - "Replace the line with this new line." (Brittle, just repeats the other parameters) - diff --git a/flowllm/extensions/file_tool/task_op.py b/flowllm/extensions/file_tool/task_op.py deleted file mode 100644 index 434c160..0000000 --- a/flowllm/extensions/file_tool/task_op.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Task delegation operation module. - -This module provides a tool operation for delegating tasks to specialized subagents. -It enables primary agents to delegate complex, multi-step tasks to specialized agents. -""" - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - - -@C.register_op() -class TaskOp(BaseAsyncToolOp): - """Task delegation operation. - - This operation delegates tasks to specialized subagents for autonomous execution. - """ - - file_path = __file__ - - def __init__(self, **kwargs): - kwargs.setdefault("raise_exception", False) - super().__init__(**kwargs) - - def build_tool_call(self) -> ToolCall: - """Build and return the tool call schema for this operator.""" - tool_params = { - "name": "Task", - "description": self.get_prompt("tool_desc"), - "input_schema": { - "description": { - "type": "string", - "description": self.get_prompt("description"), - "required": True, - }, - "prompt": { - "type": "string", - "description": self.get_prompt("prompt"), - "required": True, - }, - "subagent_type": { - "type": "string", - "description": self.get_prompt("subagent_type"), - "required": True, - }, - }, - } - - return ToolCall(**tool_params) - - async def async_execute(self): - """Execute the task delegation operation.""" - description: str = self.input_dict.get("description", "").strip() - prompt: str = self.input_dict.get("prompt", "").strip() - subagent_type: str = self.input_dict.get("subagent_type", "").strip() - - # Validate parameters - if not description: - raise ValueError("The 'description' parameter cannot be empty.") - - if not prompt: - raise ValueError("The 'prompt' parameter cannot be empty.") - - if not subagent_type: - raise ValueError("The 'subagent_type' parameter cannot be empty.") - - # TODO: Implement actual subagent delegation logic - # This is a placeholder implementation that needs to be integrated - # with the actual subagent system when available - result = ( - f"Task delegated to {subagent_type} subagent: {description}\n\n" - f"Task prompt: {prompt}\n\n[Subagent execution result would appear here]" - ) - - self.set_output(result) - - async def async_default_execute(self, e: Exception = None, **kwargs): - """Fill outputs with a default failure message when execution fails.""" - description: str = self.input_dict.get("description", "").strip() - subagent_type: str = self.input_dict.get("subagent_type", "").strip() - error_msg = f'Failed to delegate task "{description}" to subagent "{subagent_type}"' - if e: - error_msg += f": {str(e)}" - self.set_output(error_msg) diff --git a/flowllm/extensions/file_tool/task_prompt.yaml b/flowllm/extensions/file_tool/task_prompt.yaml deleted file mode 100644 index 95d5475..0000000 --- a/flowllm/extensions/file_tool/task_prompt.yaml +++ /dev/null @@ -1,12 +0,0 @@ -tool_desc: | - Delegate tasks to specialized subagents for autonomous execution. This tool enables primary agents to delegate complex, multi-step tasks to specialized agents that can handle them independently. - -description: | - A short (3-5 word) description of the task to be delegated. - -prompt: | - The detailed task prompt for the subagent to perform. Should contain all necessary information for the subagent to execute the task autonomously. - -subagent_type: | - The type of specialized subagent to use for this task. Must match an available subagent type. - diff --git a/flowllm/extensions/file_tool/write_file_op.py b/flowllm/extensions/file_tool/write_file_op.py deleted file mode 100644 index 977cfdd..0000000 --- a/flowllm/extensions/file_tool/write_file_op.py +++ /dev/null @@ -1,89 +0,0 @@ -"""Write file operation module. - -This module provides a tool operation for writing content to files. -It supports creating new files or overwriting existing files, and automatically -creates parent directories if they don't exist. -""" - -from pathlib import Path - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - - -@C.register_op() -class WriteFileOp(BaseAsyncToolOp): - """Write file operation. - - This operation writes content to a specified file. If the file doesn't exist, - it will be created. If parent directories don't exist, they will be created automatically. - """ - - file_path = __file__ - - def __init__(self, **kwargs): - kwargs.setdefault("raise_exception", False) - super().__init__(**kwargs) - - def build_tool_call(self) -> ToolCall: - """Build and return the tool call schema for this operator.""" - tool_params = { - "name": "WriteFile", - "description": self.get_prompt("tool_desc"), - "input_schema": { - "file_path": { - "type": "string", - "description": self.get_prompt("file_path"), - "required": True, - }, - "content": { - "type": "string", - "description": self.get_prompt("content"), - "required": True, - }, - }, - } - - return ToolCall(**tool_params) - - async def async_execute(self): - """Execute the write file operation.""" - file_path: str = self.input_dict.get("file_path", "").strip() - content: str = self.input_dict.get("content", "") - - # Validate file_path - if not file_path: - raise ValueError("The 'file_path' parameter cannot be empty.") - - # Resolve file path - file_path_obj = Path(file_path).expanduser().resolve() - - # Check if path is a directory - if file_path_obj.exists() and file_path_obj.is_dir(): - raise ValueError(f"Path is a directory, not a file: {file_path_obj}") - - # Create parent directories if they don't exist - file_path_obj.parent.mkdir(parents=True, exist_ok=True) - - # Check if file exists - file_exists = file_path_obj.exists() and file_path_obj.is_file() - - # Write content to file - file_path_obj.write_text(content, encoding="utf-8") - - # Format success message - if file_exists: - result = f"Successfully overwrote file: {file_path_obj}" - else: - result = f"Successfully created and wrote to new file: {file_path_obj}" - - self.set_output(result) - - async def async_default_execute(self, e: Exception = None, **kwargs): - """Fill outputs with a default failure message when execution fails.""" - file_path: str = self.input_dict.get("file_path", "").strip() - error_msg = f'Failed to write file "{file_path}"' - if e: - error_msg += f": {str(e)}" - self.set_output(error_msg) diff --git a/flowllm/extensions/file_tool/write_file_prompt.yaml b/flowllm/extensions/file_tool/write_file_prompt.yaml deleted file mode 100644 index 02740d0..0000000 --- a/flowllm/extensions/file_tool/write_file_prompt.yaml +++ /dev/null @@ -1,9 +0,0 @@ -tool_desc: | - Writes content to a specified file in the local filesystem. If the file doesn't exist, it will be created. If parent directories don't exist, they will be created automatically. If the file already exists, it will be overwritten with the new content. - -file_path: | - The absolute path to the file to write to (e.g., '/home/user/project/file.txt'). Relative paths are not supported. You must provide an absolute path. - -content: | - The content to write to the file. - diff --git a/flowllm/extensions/file_tool/write_todos_op.py b/flowllm/extensions/file_tool/write_todos_op.py deleted file mode 100644 index d041fbc..0000000 --- a/flowllm/extensions/file_tool/write_todos_op.py +++ /dev/null @@ -1,106 +0,0 @@ -"""Write todos operation module. - -This module provides a tool operation for managing todo lists. -It enables tracking subtasks with status (pending, in_progress, completed, cancelled). -""" - -from typing import List - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - -TODO_STATUSES = ["pending", "in_progress", "completed", "cancelled"] - - -@C.register_op() -class WriteTodosOp(BaseAsyncToolOp): - """Write todos operation. - - This operation manages a todo list with subtasks that can be tracked - through different statuses: pending, in_progress, completed, cancelled. - """ - - file_path = __file__ - - def __init__(self, **kwargs): - kwargs.setdefault("raise_exception", False) - super().__init__(**kwargs) - - def build_tool_call(self) -> ToolCall: - """Build and return the tool call schema for this operator.""" - tool_params = { - "name": "WriteTodos", - "description": self.get_prompt("tool_desc"), - "input_schema": { - "todos": { - "type": "array", - "description": self.get_prompt("todos"), - "required": True, - "items": { - "type": "object", - "properties": { - "description": { - "type": "string", - "description": self.get_prompt("todo_description"), - }, - "status": { - "type": "string", - "description": self.get_prompt("todo_status"), - "enum": TODO_STATUSES, - }, - }, - "required": ["description", "status"], - }, - }, - }, - } - - return ToolCall(**tool_params) - - async def async_execute(self): - """Execute the write todos operation.""" - todos: List[dict] = self.input_dict.get("todos", []) - - # Validate todos - if not isinstance(todos, list): - raise ValueError("The 'todos' parameter must be an array") - - # Validate each todo item - for i, todo in enumerate(todos): - if not isinstance(todo, dict): - raise ValueError(f"Todo item at index {i} must be an object") - - if "description" not in todo or not isinstance(todo["description"], str): - raise ValueError(f"Todo item at index {i} must have a non-empty description string") - - if not todo["description"].strip(): - raise ValueError(f"Todo item at index {i} must have a non-empty description string") - - if "status" not in todo or todo["status"] not in TODO_STATUSES: - raise ValueError( - f"Todo item at index {i} must have a valid status ({', '.join(TODO_STATUSES)})", - ) - - # Validate only one in_progress task - in_progress_count = sum(1 for todo in todos if todo.get("status") == "in_progress") - if in_progress_count > 1: - raise ValueError("Only one task can be 'in_progress' at a time") - - # Format todo list - if not todos: - result_message = "Successfully cleared the todo list." - else: - todo_list_string = "\n".join( - f"{i + 1}. [{todo['status']}] {todo['description']}" for i, todo in enumerate(todos) - ) - result_message = f"Successfully updated the todo list. The current list is now:\n{todo_list_string}" - - self.set_output(result_message) - - async def async_default_execute(self, e: Exception = None, **kwargs): - """Fill outputs with a default failure message when execution fails.""" - error_msg = "Failed to update the todo list" - if e: - error_msg += f": {str(e)}" - self.set_output(error_msg) diff --git a/flowllm/extensions/file_tool/write_todos_prompt.yaml b/flowllm/extensions/file_tool/write_todos_prompt.yaml deleted file mode 100644 index b7f6ea5..0000000 --- a/flowllm/extensions/file_tool/write_todos_prompt.yaml +++ /dev/null @@ -1,35 +0,0 @@ -tool_desc: | - This tool can help you list out the current subtasks that are required to be completed for a given user request. The list of subtasks helps you keep track of the current task, organize complex queries and help ensure that you don't miss any steps. With this list, the user can also see the current progress you are making in executing a given task. - - Depending on the task complexity, you should first divide a given task into subtasks and then use this tool to list out the subtasks that are required to be completed for a given user request. - Each of the subtasks should be clear and distinct. - - Use this tool for complex queries that require multiple steps. If you find that the request is actually complex after you have started executing the user task, create a todo list and use it. If execution of the user task requires multiple steps, planning and generally is higher complexity than a simple Q&A, use this tool. - - DO NOT use this tool for simple tasks that can be completed in less than 2 steps. If the user query is simple and straightforward, do not use the tool. If you can respond with an answer in a single turn then this tool is not required. - - ## Task state definitions - - - pending: Work has not begun on a given subtask. - - in_progress: Marked just prior to beginning work on a given subtask. You should only have one subtask as in_progress at a time. - - completed: Subtask was successfully completed with no errors or issues. If the subtask required more steps to complete, update the todo list with the subtasks. All steps should be identified as completed only when they are completed. - - cancelled: As you update the todo list, some tasks are not required anymore due to the dynamic nature of the task. In this case, mark the subtasks as cancelled. - - ## Methodology for using this tool - 1. Use this todo list list as soon as you receive a user request based on the complexity of the task. - 2. Keep track of every subtask that you update the list with. - 3. Mark a subtask as in_progress before you begin working on it. You should only have one subtask as in_progress at a time. - 4. Update the subtask list as you proceed in executing the task. The subtask list is not static and should reflect your progress and current plans, which may evolve as you acquire new information. - 5. Mark a subtask as completed when you have completed it. - 6. Mark a subtask as cancelled if the subtask is no longer needed. - 7. You must update the todo list as soon as you start, stop or cancel a subtask. Don't batch or wait to update the todo list. - -todos: | - The complete list of todo items. This will replace the existing list. - -todo_description: | - The description of the task. - -todo_status: | - The current status of the task. Must be one of: pending, in_progress, completed, cancelled. - diff --git a/flowllm/extensions/skills/__init__.py b/flowllm/extensions/skills/__init__.py deleted file mode 100644 index 7095e70..0000000 --- a/flowllm/extensions/skills/__init__.py +++ /dev/null @@ -1,19 +0,0 @@ -"""Skills extension module for flowllm. - -This module provides operations for managing and executing skills, -including loading skill metadata, reading reference files, and running shell commands. -""" - -from .load_skill_metadata_op import LoadSkillMetadataOp -from .load_skill_op import LoadSkillOp -from .read_reference_file_op import ReadReferenceFileOp -from .run_shell_command_op import RunShellCommandOp -from .skill_agent_op import SkillAgentOp - -__all__ = [ - "LoadSkillMetadataOp", - "LoadSkillOp", - "ReadReferenceFileOp", - "RunShellCommandOp", - "SkillAgentOp", -] diff --git a/flowllm/extensions/skills/load_skill_metadata_op.py b/flowllm/extensions/skills/load_skill_metadata_op.py deleted file mode 100644 index f9c07f9..0000000 --- a/flowllm/extensions/skills/load_skill_metadata_op.py +++ /dev/null @@ -1,187 +0,0 @@ -"""Operation for loading skill metadata. - -This module provides the LoadSkillMetadataOp class which scans the skills -directory recursively and extracts metadata (name and description) from all -SKILL.md files. The metadata is parsed from YAML frontmatter in each SKILL.md -file and returned as a dictionary mapping skill names to their metadata. -""" - -from pathlib import Path - -from loguru import logger - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - - -@C.register_op() -class LoadSkillMetadataOp(BaseAsyncToolOp): - """Operation for loading metadata from all available skills. - - This tool scans the skills directory recursively for SKILL.md files and - extracts their metadata (name and description) from YAML frontmatter. - The metadata is returned as a dictionary where keys are skill names and - values contain the description and skill directory path. - - Returns: - dict: A dictionary mapping skill names to their metadata. Each entry - has the format: - { - "skill_name": { - "description": "Skill description text", - "skill_dir": "/path/to/skill/directory" - } - } - - Note: - The skills directory path is obtained from `self.context.skill_dir`. - Only SKILL.md files with valid YAML frontmatter containing both 'name' - and 'description' fields will be included in the results. - """ - - def build_tool_call(self) -> ToolCall: - """Build the tool call definition for load_skill_metadata. - - Creates and returns a ToolCall object that defines the load_skill_metadata - tool. This tool requires no input parameters and will scan the skills - directory to load all available skill metadata. - - Returns: - ToolCall: A ToolCall object defining the load_skill_metadata tool - with the following properties: - - name: "load_skill_metadata" - - description: Description of what the tool does - - input_schema: Empty dict (no input parameters required) - """ - tool_params = { - "name": "load_skill_metadata", - "description": "Load metadata (name and description) for all available skills from the skills directory.", - "input_schema": {}, - } - return ToolCall(**tool_params) - - @staticmethod - async def parse_skill_metadata(content: str, path: str) -> dict[str, str] | None: - """Extract skill metadata (name and description) from SKILL.md content. - - Parses YAML frontmatter from SKILL.md files to extract the skill name - and description. The frontmatter should be in the format: - ``` - --- - name: skill_name - description: skill description - --- - ``` - - The method splits the content by "---" delimiters to extract the - frontmatter section, then parses each line to find the 'name' and - 'description' fields. Values can be quoted or unquoted. - - Args: - content: The full content of the SKILL.md file as a string. - path: The file path (used for logging purposes when parsing fails). - - Returns: - dict[str, str] | None: A dictionary with 'name' and 'description' - keys containing the extracted values, or None if: - - No YAML frontmatter is found (less than 3 parts after splitting) - - The 'name' field is missing or empty - - The 'description' field is missing or empty - """ - # Split content by YAML frontmatter delimiters (---) - # Expected format: "---\n...frontmatter...\n---\n...content..." - # This should result in at least 3 parts: [before, frontmatter, after] - parts = content.split("---") - if len(parts) < 3: - logger.warning(f"No YAML frontmatter found in skill from {path}") - return None - - # Extract the frontmatter section (between the first two "---" delimiters) - frontmatter_text = parts[1].strip() - name = None - description = None - - # Parse each line in the frontmatter to find name and description - for line in frontmatter_text.split("\n"): - line = line.strip() - if line.startswith("name:"): - # Extract value after "name:", remove quotes if present - name = line.split(":", 1)[1].strip().strip("\"'") - elif line.startswith("description:"): - # Extract value after "description:", remove quotes if present - description = line.split(":", 1)[1].strip().strip("\"'") - - # Validate that both required fields are present - if not name or not description: - logger.warning(f"Missing name or description in skill from {path}") - return None - - return { - "name": name, - "description": description, - } - - async def async_execute(self): - """Execute the load skill metadata operation. - - Scans the skills directory recursively for all SKILL.md files, - extracts their metadata from YAML frontmatter, and builds a dictionary - mapping skill names to their metadata (description and directory path). - - The method: - 1. Gets the skills directory path from the context - 2. Recursively searches for all SKILL.md files - 3. Parses each file's frontmatter to extract metadata - 4. Builds a dictionary with skill names as keys - 5. Sets the output with the complete metadata dictionary - - Returns: - None: The result is set via `self.set_output()` with a dictionary - in the format: - { - "skill_name_1": { - "description": "Description of skill 1", - "skill_dir": "/path/to/skill1" - }, - "skill_name_2": { - "description": "Description of skill 2", - "skill_dir": "/path/to/skill2" - }, - ... - } - - Note: - Only skills with valid metadata (both name and description) are - included in the result. Invalid or missing metadata is logged as - a warning but does not stop the process. - """ - # Get the skills directory path from context - skill_dir = Path(self.context.skill_dir) - logger.info(f"🔧 Tool called: load_skill_metadata(path={skill_dir})") - - # Recursively find all SKILL.md files in the skills directory - skill_files = list(skill_dir.rglob("SKILL.md")) - assert skill_files, "No SKILL.md files found in skills directory" - - # Build dictionary mapping skill names to their metadata - skill_metadata_dict = {} - for skill_file in skill_files: - # Read the SKILL.md file content - content = skill_file.read_text(encoding="utf-8") - # Parse metadata from the file's frontmatter - metadata = await self.parse_skill_metadata(content, str(skill_file)) - - if metadata: - # Get the parent directory of the SKILL.md file as the skill directory - skill_dir = skill_file.parent.as_posix() - # Store metadata with skill name as key - skill_metadata_dict[metadata["name"]] = { - "description": metadata["description"], - "skill_dir": skill_dir, - } - logger.info(f"✅ Loaded skill {metadata['name']} metadata skill_dir={skill_dir}") - - logger.info(f"✅ Loaded {len(skill_metadata_dict)} skill metadata entries") - # Set the output with the complete metadata dictionary - self.set_output(skill_metadata_dict) diff --git a/flowllm/extensions/skills/load_skill_op.py b/flowllm/extensions/skills/load_skill_op.py deleted file mode 100644 index 77f9bdd..0000000 --- a/flowllm/extensions/skills/load_skill_op.py +++ /dev/null @@ -1,130 +0,0 @@ -"""Operation for loading a specific skill. - -This module provides the LoadSkillOp class which loads the content of a -SKILL.md file from a specified skill directory. The skill directory is -looked up from the skill metadata dictionary using the skill name. If the -SKILL.md file contains YAML frontmatter, only the content after the -frontmatter is returned; otherwise, the full file content is returned. -""" - -from pathlib import Path - -from loguru import logger - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - - -@C.register_op() -class LoadSkillOp(BaseAsyncToolOp): - """Operation for loading a specific skill's instructions. - - This tool loads the content of a SKILL.md file for a given skill name. - The skill directory is retrieved from the context's skill_metadata_dict, - which should be populated by LoadSkillMetadataOp beforehand. - - The tool: - 1. Takes a skill_name as input - 2. Looks up the skill directory from skill_metadata_dict - 3. Reads the SKILL.md file from that directory - 4. If YAML frontmatter is present, returns only the content after it - 5. If no frontmatter is found, returns the full file content - - Returns: - str: The skill instructions content. If the SKILL.md file has YAML - frontmatter (delimited by "---"), returns only the content after - the frontmatter. Otherwise, returns the complete file content. - If the skill is not found, returns an error message string. - - Note: - - The skill_name must exist in `self.context.skill_metadata_dict` - - The SKILL.md file must exist in the skill directory - - YAML frontmatter is detected by splitting on "---" delimiters - """ - - def build_tool_call(self) -> ToolCall: - """Build the tool call definition for load_skill. - - Creates and returns a ToolCall object that defines the load_skill - tool. This tool requires a skill_name parameter to identify which - skill's instructions to load. - - Returns: - ToolCall: A ToolCall object defining the load_skill tool with - the following properties: - - name: "load_skill" - - description: Description of what the tool does - - input_schema: A schema requiring a "skill_name" string - parameter that must be provided - """ - return ToolCall( - **{ - "name": "load_skill", - "description": "Load one skill's instructions from the SKILL.md.", - "input_schema": { - "skill_name": { - "type": "string", - "description": "skill name", - "required": True, - }, - }, - }, - ) - - async def async_execute(self): - """Execute the load skill operation. - - Loads the SKILL.md file content for the specified skill name. The - method retrieves the skill directory from the context's skill metadata - dictionary, reads the SKILL.md file, and extracts the instructions - content (excluding YAML frontmatter if present). - - The method: - 1. Extracts the skill_name from input_dict - 2. Looks up the skill directory from skill_metadata_dict - 3. Constructs the path to SKILL.md file - 4. Checks if the file exists - 5. Reads the file content - 6. Splits content by "---" to detect YAML frontmatter - 7. Returns content after frontmatter if present, otherwise full content - - Returns: - None: The result is set via `self.set_output()` with one of: - - The skill instructions (content after YAML frontmatter) - - The full file content (if no frontmatter is found) - - An error message string (if skill file is not found) - - Raises: - KeyError: If skill_name is not found in skill_metadata_dict. - This should be handled by ensuring LoadSkillMetadataOp is - called before LoadSkillOp. - - Note: - - If the file has YAML frontmatter (format: "---\n...\n---\n..."), - only the content after the second "---" is returned - - If no frontmatter is detected, the entire file content is returned - - File encoding is assumed to be UTF-8 - """ - # Extract skill name from input parameters - skill_name = self.input_dict["skill_name"] - # Look up the skill directory from the metadata dictionary - # This dictionary should be populated by LoadSkillMetadataOp - skill_dir = Path(self.context.skill_metadata_dict[skill_name]["skill_dir"]) - logger.info(f"🔧 Tool called: load_skill(skill_name='{skill_name}') with skill_dir={skill_dir}") - - # Construct the path to the SKILL.md file - skill_path = skill_dir / "SKILL.md" - - # Check if the SKILL.md file exists - if not skill_path.exists(): - content = f"❌ Skill '{skill_name}' not found" - logger.exception(content) - self.set_output(content) - return - - # Read the SKILL.md file content - content: str = skill_path.read_text(encoding="utf-8") - self.set_output(content) - - logger.info(f"✅ Loaded skill: {skill_name} size={len(content)}") diff --git a/flowllm/extensions/skills/read_reference_file_op.py b/flowllm/extensions/skills/read_reference_file_op.py deleted file mode 100644 index 58cc3e1..0000000 --- a/flowllm/extensions/skills/read_reference_file_op.py +++ /dev/null @@ -1,132 +0,0 @@ -"""Operation for reading reference files from skills. - -This module provides the ReadReferenceFileOp class which allows reading -reference files (e.g., forms.md, reference.md, ooxml.md) from skill directories. -The skill directory is looked up from the skill metadata dictionary using the -skill name. If the file is not found, an error message is returned. -""" - -from pathlib import Path - -from loguru import logger - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - - -@C.register_op() -class ReadReferenceFileOp(BaseAsyncToolOp): - """Operation for reading reference files from a skill directory. - - This tool allows reading reference files like forms.md, reference.md, - or ooxml.md from a specific skill's directory. The skill directory is - retrieved from the context's skill_metadata_dict, which should be - populated by LoadSkillMetadataOp beforehand. - - The tool: - 1. Takes a skill_name and file_name as input - 2. Looks up the skill directory from skill_metadata_dict - 3. Constructs the file path as {skill_dir}/{file_name} - 4. Reads the file content if it exists - 5. Returns the file content or an error message if not found - - Returns: - str: The content of the reference file read from the skill directory. - If the file is not found, returns an error message string indicating - that the file was not found in the specified skill. - - Note: - - The skill_name must exist in `self.context.skill_metadata_dict` - - The file_name can be a simple filename (e.g., "reference.md") or a - relative path within the skill directory - - File encoding is assumed to be UTF-8 - - The file path is constructed as: {skill_dir}/{file_name} - """ - - def build_tool_call(self) -> ToolCall: - """Build the tool call definition for read_reference_file. - - Creates and returns a ToolCall object that defines the read_reference_file - tool. This tool requires both skill_name and file_name parameters to - identify which reference file to read from which skill. - - Returns: - ToolCall: A ToolCall object defining the read_reference_file tool with - the following properties: - - name: "read_reference_file" - - description: Description of what the tool does - - input_schema: A schema requiring: - - "skill_name" (string, required): The name of the skill - - "file_name" (string, required): The reference file name or - file path relative to the skill directory - """ - return ToolCall( - **{ - "name": "read_reference_file", - "description": "Read a reference file from a skill (e.g., forms.md, reference.md, ooxml.md)", - "input_schema": { - "skill_name": { - "type": "string", - "description": "skill name", - "required": True, - }, - "file_name": { - "type": "string", - "description": "reference file name or file path", - "required": True, - }, - }, - }, - ) - - async def async_execute(self): - """Execute the read reference file operation. - - Reads a reference file from the specified skill directory. The method - retrieves the skill directory from the context's skill metadata - dictionary, constructs the file path, and reads the file content if it - exists. - - The method: - 1. Extracts the skill_name and file_name from input_dict - 2. Looks up the skill directory from skill_metadata_dict - 3. Constructs the file path as {skill_dir}/{file_name} - 4. Checks if the file exists - 5. Reads the file content if it exists - 6. Returns the file content or an error message if not found - - Returns: - None: The result is set via `self.set_output()` with one of: - - The file content (as a string) if the file exists - - An error message string if the file is not found - - Raises: - KeyError: If skill_name is not found in skill_metadata_dict. - This should be handled by ensuring LoadSkillMetadataOp is - called before ReadReferenceFileOp. - - Note: - - The file path is constructed as: {skill_dir}/{file_name} - - File encoding is assumed to be UTF-8 - - If the file does not exist, an error message is returned instead - of raising an exception - """ - skill_name = self.input_dict["skill_name"] - file_name = self.input_dict["file_name"] - skill_dir = Path(self.context.skill_metadata_dict[skill_name]["skill_dir"]) - logger.info( - f"🔧 Tool called: read_reference_file(skill_name='{skill_name}', file_name='{file_name}') " - f"with skill_dir={skill_dir}", - ) - - file_path = skill_dir / file_name - if not file_path.exists(): - content = f"File '{file_name}' not found in skill '{skill_name}'" - logger.exception(content) - self.set_output(content) - return - - result = file_path.read_text(encoding="utf-8") - logger.info(f"✅ Read file: {skill_name}/{file_name} size={len(result)}") - self.set_output(result) diff --git a/flowllm/extensions/skills/run_shell_command_op.py b/flowllm/extensions/skills/run_shell_command_op.py deleted file mode 100644 index d01f4db..0000000 --- a/flowllm/extensions/skills/run_shell_command_op.py +++ /dev/null @@ -1,183 +0,0 @@ -"""Operation for running shell commands. - -This module provides the RunShellCommandOp class which executes shell -commands in a subprocess, with automatic dependency detection and -installation for script files. The command is executed in the skill's -directory context, allowing scripts to access skill-specific files and -resources. -""" - -import asyncio -import os -import shutil -from pathlib import Path - -from loguru import logger - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - - -@C.register_op() -class RunShellCommandOp(BaseAsyncToolOp): - """Operation for running shell commands in a subprocess. - - This tool executes shell commands and can automatically detect and - install dependencies for script files (Python, JavaScript, Shell). - The command is executed in the skill's directory context, allowing - scripts to access skill-specific files and resources. - - The operation will: - 1. Extract skill_name and command from input - 2. Look up the skill directory from skill_metadata_dict - 3. Change to the skill directory before executing the command - 4. For Python commands, automatically detect and install dependencies - using pipreqs (if available and auto_install_deps parameter is enabled) - 5. Execute the command in a subprocess and capture stdout/stderr - 6. Return the combined output - - Returns: - str: The combined stdout and stderr output from the command execution. - The output is decoded as UTF-8 and stripped of leading/trailing - whitespace, with stdout and stderr concatenated with a newline. - - Note: - - The skill_name must exist in `self.context.skill_metadata_dict` - - The command is executed in the skill's directory using `cd {skill_dir} && {command}` - - For Python commands (containing "py"), the tool attempts to auto-install - dependencies using pipreqs if it's available in the system PATH and - the auto_install_deps parameter is enabled - - If pipreqs is not available or dependency installation fails, a warning - is logged but the command execution continues - - The subprocess uses the current environment variables (os.environ.copy()) - """ - - def __init__(self, auto_install_deps: bool = False, **kwargs): - """Initialize RunShellCommandOp. - - Args: - auto_install_deps: If True, enables automatic dependency installation for Python - commands. Defaults to False. - **kwargs: Additional keyword arguments passed to parent class. - """ - super().__init__(**kwargs) - self.auto_install_deps: bool = auto_install_deps - - def build_tool_call(self) -> ToolCall: - """Build the tool call definition for run_shell_command. - - Creates and returns a ToolCall object that defines the run_shell_command - tool. This tool requires both skill_name and command parameters to - identify which skill directory to use and what command to execute. - - Returns: - ToolCall: A ToolCall object defining the run_shell_command tool with - the following properties: - - name: "run_shell_command" - - description: Description of what the tool does - - input_schema: A schema requiring: - - "skill_name" (string, required): The name of the skill - - "command" (string, required): The shell command to execute - """ - return ToolCall( - **{ - "name": "run_shell_command", - "description": self.get_prompt("tool_desc"), - "input_schema": { - "skill_name": { - "type": "string", - "description": "skill name", - "required": True, - }, - "command": { - "type": "string", - "description": "shell command", - "required": True, - }, - }, - }, - ) - - async def async_execute(self): - """Execute the shell command operation. - - Executes a shell command in the specified skill's directory. For Python - commands, the method attempts to automatically detect and install - dependencies using pipreqs before executing the command. - - The method: - 1. Extracts skill_name and command from input_dict - 2. Looks up the skill directory from skill_metadata_dict - 3. For Python commands (containing "py"), checks if pipreqs is available - 4. If pipreqs is available, generates requirements.txt and installs dependencies - 5. Constructs the full command as `cd {skill_dir} && {command}` - 6. Executes the command in a subprocess with the current environment - 7. Captures stdout and stderr output - 8. Returns the combined output (stdout + stderr) - - Returns: - None: The result is set via `self.set_output()` with the combined - stdout and stderr output from the command execution. The output - is decoded as UTF-8 and formatted as: "{stdout}\n{stderr}" - - Raises: - KeyError: If skill_name is not found in skill_metadata_dict. - This should be handled by ensuring LoadSkillMetadataOp is - called before RunShellCommandOp. - - Note: - - Dependency auto-installation only occurs for commands containing "py" - and when the auto_install_deps parameter is enabled - - If pipreqs is not available, a warning is logged but execution continues - - If dependency installation fails, a warning is logged but the command - is still executed - - The command runs in the skill's directory, allowing access to - skill-specific files and resources - - Environment variables from the current process are passed to the subprocess - """ - # Extract skill name and command from input parameters - skill_name = self.input_dict["skill_name"] - command: str = self.input_dict["command"] - # Look up the skill directory from the metadata dictionary - # This dictionary should be populated by LoadSkillMetadataOp - skill_dir = Path(self.context.skill_metadata_dict[skill_name]["skill_dir"]) - logger.info(f"🔧 run shell command: skill_name={skill_name} skill_dir={skill_dir} command={command}") - - # Auto-install dependencies for Python scripts if pipreqs is available - # This helps ensure that Python scripts have their required dependencies - # Only install if auto_install_deps parameter is enabled - if self.auto_install_deps: - if "py" in command: - pipreqs_available = shutil.which("pipreqs") is not None - if pipreqs_available: - install_cmd = f"cd {skill_dir} && pipreqs . --force && pip install -r requirements.txt" - proc = await asyncio.create_subprocess_shell( - install_cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - stdout, stderr = await proc.communicate() - if proc.returncode != 0: - logger.warning(f"⚠️ Failed to install dependencies:\n{stdout.decode()}\n{stderr.decode()}") - else: - logger.info(f"✅ Dependencies installed successfully.\n{stdout.decode()}\n{stderr.decode()}") - else: - logger.info("ℹ️ pipreqs not found, skipping dependency auto-install.") - - # Construct the full command to execute in the skill directory - # This ensures the command runs in the correct context - full_command = f"cd {skill_dir} && {command}" - proc = await asyncio.create_subprocess_shell( - full_command, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - env=os.environ.copy(), - ) - - # Wait for the command to complete and capture output - stdout, stderr = await proc.communicate() - # Combine stdout and stderr output, decoded as UTF-8 - output = stdout.decode().strip() + "\n" + stderr.decode().strip() - logger.info(f"✅ Command executed: skill_name={skill_name} output={output}") - self.set_output(output) diff --git a/flowllm/extensions/skills/run_shell_command_prompt.yaml b/flowllm/extensions/skills/run_shell_command_prompt.yaml deleted file mode 100644 index 0741118..0000000 --- a/flowllm/extensions/skills/run_shell_command_prompt.yaml +++ /dev/null @@ -1,11 +0,0 @@ -tool_desc: | - run shell command in a subprocess. - Here you need to fill in skill_name. - This skill_name parameter allows you to navigate directly to the folder corresponding to skill_name, making it more convenient to use the scripts within that folder to execute commands. - If you want to know the exact path, you can use `pwd` to get the absolute path. - -tool_desc_zh: | - 在子进程中运行 shell 命令。 - 此处你需要填写 skill_name。 - 该 skill_name 参数可让你直接跳转到与 skill_name 对应的文件夹，从而更方便地使用该文件夹中的脚本来执行命令。 - 如果你想了解确切的路径，可以使用 `pwd` 命令获取绝对路径。 diff --git a/flowllm/extensions/skills/skill_agent_op.py b/flowllm/extensions/skills/skill_agent_op.py deleted file mode 100644 index 573b92d..0000000 --- a/flowllm/extensions/skills/skill_agent_op.py +++ /dev/null @@ -1,205 +0,0 @@ -"""Operation for a skill-enabled reactive agent. - -This module provides the SkillAgentOp class which extends ReactAgentOp to -automatically use pre-built skills relevant to the query. The agent loads -skill metadata from a specified skill directory and makes those skills -available as tools during the reasoning process. -""" - -import datetime -from typing import List - -from loguru import logger - -from ...core.context import C -from ...core.enumeration import Role -from ...core.op import BaseAsyncToolOp -from ...core.schema import Message, ToolCall -from ...gallery.agent import ReactAgentOp - - -@C.register_op() -class SkillAgentOp(ReactAgentOp): - """Reactive agent that automatically uses pre-built skills to answer queries. - - This agent extends ReactAgentOp to provide skill-based reasoning. It loads - skill metadata from a specified directory and makes those skills available - as tools during the agent's reasoning process. The agent can automatically - select and use relevant skills based on the user's query. - - The agent: - 1. Loads skill metadata from the specified skill directory - 2. Builds a system prompt that includes available skills - 3. Uses the React (Reasoning and Acting) pattern to iteratively reason - and call tools (including skills) - 4. Automatically passes skill_metadata_dict to skill operations when - they are invoked - - Attributes: - file_path (str): Path to the operation file, used for prompt loading. - Inherits all attributes from ReactAgentOp, including: - - llm: The language model to use (default: "qwen3_max_instruct") - - max_steps: Maximum number of reasoning steps (default: 5) - - tool_call_interval: Delay between tool calls in seconds (default: 1.0) - - add_think_tool: Whether to add a thinking tool (default: False) - - Note: - - The skill_dir must contain SKILL.md files with valid metadata - - Skills are loaded via LoadSkillMetadataOp before the agent starts - - The skill_metadata_dict is stored in context and passed to skill operations - """ - - file_path: str = __file__ - - def __init__( - self, - llm: str = "qwen3_30b_instruct", - max_steps: int = 50, - tool_call_interval: float = 1.0, - add_think_tool: bool = False, - **kwargs, - ): - """Initialize the skill agent with configuration. - - Args: - llm: The language model identifier to use for reasoning. - Default is "qwen3_max_instruct". - max_steps: Maximum number of reasoning steps the agent can take. - Default is 5. Note: This is passed as max_retries to the parent. - tool_call_interval: Delay in seconds between tool calls to avoid - rate limiting. Default is 1.0. - add_think_tool: Whether to add a thinking tool that allows the agent - to explicitly reason before taking actions. Default is False. - **kwargs: Additional keyword arguments passed to the parent class. - """ - super().__init__( - llm=llm, - max_steps=max_steps, - tool_call_interval=tool_call_interval, - add_think_tool=add_think_tool, - **kwargs, - ) - - def build_tool_call(self) -> ToolCall: - """Build the tool call definition for the skill agent. - - Creates and returns a ToolCall object that defines the skill agent tool. - This tool requires both query and skill_dir parameters to identify what - to answer and which skills are available. - - Returns: - ToolCall: A ToolCall object defining the skill agent tool with - the following properties: - - description: Description of what the tool does - - input_schema: A schema requiring: - - "query" (string, required): The user's query to answer - - "skill_dir" (string, required): The directory containing - skill definitions (SKILL.md files) - """ - return ToolCall( - **{ - "description": "Automatically uses pre-built Skills relevant to the query when needed.", - "input_schema": { - "query": { - "type": "string", - "description": "query", - "required": True, - }, - "skill_dir": { - "type": "string", - "description": "skill dir", - "required": True, - }, - }, - }, - ) - - async def build_messages(self) -> List[Message]: - """Build the initial messages for the agent conversation. - - Loads skill metadata from the specified skill directory and constructs - the initial conversation messages. The system prompt includes information - about available skills, allowing the agent to understand what tools - are at its disposal. - - The method: - 1. Extracts query and skill_dir from context - 2. Loads skill metadata using LoadSkillMetadataOp - 3. Stores the skill_metadata_dict in context for later use - 4. Formats skill metadata as a list for inclusion in the system prompt - 5. Builds system and user messages with the current time and skill info - - Returns: - List[Message]: A list containing: - - A SYSTEM message with the formatted system prompt including: - skill directory, current time, and list of available skills - - A USER message containing the user's query - - Note: - - The skill_metadata_dict is stored in context.skill_metadata_dict - for use by skill operations - - The system prompt is formatted using the "system_prompt" template - from the operation's prompt file - - Current time is formatted as "YYYY-MM-DD HH:MM:SS" - """ - # Extract query and skill directory from context - query: str = self.context.query - skill_dir: str = self.context.skill_dir - logger.info(f"SkillAgentOp processing query: {query} with access to skills in {skill_dir}") - - # Load skill metadata from the skill directory - # This populates the skill_metadata_dict with all available skills - from .load_skill_metadata_op import LoadSkillMetadataOp - - op = LoadSkillMetadataOp() - await op.async_call(skill_dir=skill_dir) - # Store the skill metadata dictionary in context for use by skill operations - self.context.skill_metadata_dict = skill_metadata_dict = op.output - - # Format skill metadata as a list for inclusion in the system prompt - skill_metadata_list = [f"- {k}: {v['description']}" for k, v in skill_metadata_dict.items()] - logger.info(f"SkillAgentOp loaded skill metadata: {skill_metadata_dict}") - - # Get current time for the system prompt - now_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") - # Build the initial conversation messages - messages = [ - Message( - role=Role.SYSTEM, - content=self.prompt_format( - "system_prompt", - time=now_time, - skill_metadata="\n".join(skill_metadata_list), - ), - ), - Message(role=Role.USER, content=query), - ] - - return messages - - async def execute_tool(self, op: BaseAsyncToolOp, tool_call: ToolCall): - """Execute a tool operation with skill metadata context. - - Overrides the parent class method to automatically pass the - skill_metadata_dict to skill operations when they are invoked. - This ensures that skill operations have access to the metadata - they need to function correctly. - - Args: - op: The tool operation to execute (e.g., LoadSkillOp, - ReadReferenceFileOp, RunShellCommandOp). - tool_call: The tool call object containing the arguments - for the operation. - - Note: - - The skill_metadata_dict is automatically passed to all tool - operations, allowing them to look up skill directories - - Additional arguments from tool_call.argument_dict are also - passed to the operation - - The operation is executed asynchronously via submit_async_task - """ - self.submit_async_task( - op.async_call, - skill_metadata_dict=self.context.skill_metadata_dict, - **tool_call.argument_dict, - ) diff --git a/flowllm/extensions/skills/skill_agent_prompt.yaml b/flowllm/extensions/skills/skill_agent_prompt.yaml deleted file mode 100644 index 487d549..0000000 --- a/flowllm/extensions/skills/skill_agent_prompt.yaml +++ /dev/null @@ -1,48 +0,0 @@ -system_prompt: | - You are a helpful AI assistant with access to specialized skills. - The current time is {time}. - When you encounter tasks involving specific domains or file formats, use the "load_skill" tool to gain expert knowledge. - Available skills (each line is "- : ") - {skill_metadata} - - Workflow: - 1. Identify if the task needs specialized knowledge. - 2. If specialized knowledge is needed, identify the most relevant skill from the available skills list. - 3. Use "load_skill" tool to get detailed instructions for the chosen skill. This will load the content of SKILL.md into your context. - 4. If the skill mentions reference files (e.g., forms.md), use "read_reference_file" tool to access their contents only when explicitly required for the task. - 5. If the skill includes executable scripts (e.g., fill_form.py), use "run_shell_command" tool with the appropriate shell commands to run them when necessary. Remember that only the script's output will be added to your context, not the script's code itself. - 6. Follow the instructions from the loaded skill - 7. Use available tools as needed - 8. After completing the task, call the "task_complete" tool to indicate that you are done with your task - - Important: - - Only load skills and additional resources when they are directly relevant to the current task - - Skill scripts are located in: scripts/ - - When running skill scripts, use the full path from current directory when creating the shell commands - - Example: scripts/check_fillable_fields.py - - If a task requires multiple skills, load and apply them sequentially as needed - - -system_prompt_zh: | - 你是一个具备专业技能访问权限的智能助手。 - 当前时间是 {time}。 - 当你遇到涉及特定领域或文件格式的任务时，请使用“load_skill”工具来获取专家级知识。 - 可用技能（每行格式为“- <技能名称>: <技能描述>”） - {skill_metadata} - - 工作流程： - 1. 判断当前任务是否需要专业知识。 - 2. 如果需要专业知识，请从可用技能列表中选择最相关的技能。 - 3. 使用“load_skill”工具加载所选技能的详细说明。这会将该技能目录下的 SKILL.md 文件内容载入你的上下文。 - 4. 如果技能说明中提到了参考文件（例如 forms.md），仅在任务明确需要时，才使用“read_reference_file”工具读取其内容。 - 5. 如果技能包含可执行脚本（例如 fill_form.py），在必要时使用“run_shell_command”工具运行相应的 shell 命令。请注意，只有脚本的输出结果会被添加到你的上下文中，而非脚本本身的代码。 - 6. 遵循已加载技能中的说明进行操作。 - 7. 根据需要使用其他可用工具。 - 8. 任务完成后，调用“task_complete”工具，表明你已完成当前任务。 - - 重要提示： - - 仅在当前任务直接相关时，才加载技能和额外资源。 - - 技能脚本位于：scripts/ - - 运行技能脚本时，请在 shell 命令中使用从当前目录出发的完整路径。 - - 示例：scripts/check_fillable_fields.py - - 如果一个任务需要多个技能，请按需依次加载并应用它们。 \ No newline at end of file diff --git a/flowllm/extensions/utils/__init__.py b/flowllm/extensions/utils/__init__.py deleted file mode 100644 index ba79ac7..0000000 --- a/flowllm/extensions/utils/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -"""Utility functions for FlowLLM framework extensions. - -This package provides utility functions that can be used across extension modules. -It includes file editing utilities for intelligent text replacement. -""" - -from .edit_utils import ( - calculate_exact_replacement, - calculate_flexible_replacement, - calculate_regex_replacement, - escape_regex, - restore_trailing_newline, -) - -__all__ = [ - "calculate_exact_replacement", - "calculate_flexible_replacement", - "calculate_regex_replacement", - "escape_regex", - "restore_trailing_newline", -] diff --git a/flowllm/extensions/utils/edit_utils.py b/flowllm/extensions/utils/edit_utils.py deleted file mode 100644 index 6f3eadf..0000000 --- a/flowllm/extensions/utils/edit_utils.py +++ /dev/null @@ -1,183 +0,0 @@ -"""File editing utility functions. - -This module provides utility functions for intelligently editing files by replacing text. -It supports exact matching, flexible matching (ignoring indentation), and regex-based matching. -""" - -import re - - -def escape_regex(s: str) -> str: - """Escape special regex characters in a string. - - This function escapes all special regex characters in the input string so that - it can be used as a literal string in a regular expression pattern. - - Args: - s: The string containing characters that may need escaping - - Returns: - A string with all special regex characters escaped - """ - return re.escape(s) - - -def restore_trailing_newline(original: str, modified: str) -> str: - """Restore trailing newline to match the original string. - - This function ensures that the modified string has the same trailing newline - behavior as the original string, preserving the file's newline format. - - Args: - original: The original string before modification - modified: The modified string that may have lost or gained a trailing newline - - Returns: - The modified string with trailing newline adjusted to match the original - """ - had_newline = original.endswith("\n") - if had_newline and not modified.endswith("\n"): - return modified + "\n" - elif not had_newline and modified.endswith("\n"): - return modified.rstrip("\n") - return modified - - -def calculate_exact_replacement( - content: str, - old_string: str, - new_string: str, -) -> tuple[str, int] | None: - """Perform exact string replacement in content. - - This function attempts to replace the old_string with new_string using exact - matching. It normalizes line endings (converts \\r\\n to \\n) before matching - and preserves the original trailing newline behavior. - - Args: - content: The original content string to modify - old_string: The exact string to find and replace - new_string: The replacement string - - Returns: - A tuple containing (modified_content, occurrence_count) if the old_string - is found, or None if no match is found. The occurrence_count indicates - how many times the old_string appears in the content. - """ - normalized_content = content - normalized_old = old_string.replace("\r\n", "\n") - normalized_new = new_string.replace("\r\n", "\n") - - occurrences = len(normalized_content.split(normalized_old)) - 1 - if occurrences > 0: - new_content = normalized_content.replace(normalized_old, normalized_new, 1) - new_content = restore_trailing_newline(content, new_content) - return new_content, occurrences - return None - - -def calculate_flexible_replacement( - content: str, - old_string: str, - new_string: str, -) -> tuple[str, int] | None: - """Perform flexible string replacement that ignores indentation differences. - - This function matches and replaces text by comparing line content while ignoring - leading whitespace (indentation). It preserves the indentation of the first matched - line when applying the replacement, making it useful for code editing where - indentation may vary. - - Args: - content: The original content string to modify - old_string: The string pattern to find (indentation is ignored during matching) - new_string: The replacement string (will be indented to match the original) - - Returns: - A tuple containing (modified_content, occurrence_count) if matches are found, - or None if no match is found. The occurrence_count indicates how many times - the pattern was found and replaced. - """ - normalized_content = content - normalized_old = old_string.replace("\r\n", "\n") - normalized_new = new_string.replace("\r\n", "\n") - - source_lines = normalized_content.split("\n") - search_lines_stripped = [line.strip() for line in normalized_old.split("\n") if line.strip()] - replace_lines = normalized_new.split("\n") - - if not search_lines_stripped: - return None - - occurrences = 0 - i = 0 - while i <= len(source_lines) - len(search_lines_stripped): - window = source_lines[i : i + len(search_lines_stripped)] - window_stripped = [line.strip() for line in window] - if all(window_stripped[j] == search_lines_stripped[j] for j in range(len(search_lines_stripped))): - occurrences += 1 - first_line = window[0] - indent_match = re.match(r"^(\s*)", first_line) - indent = indent_match.group(1) if indent_match else "" - new_block = [f"{indent}{line}" for line in replace_lines] - source_lines[i : i + len(search_lines_stripped)] = new_block - i += len(replace_lines) - else: - i += 1 - - if occurrences > 0: - new_content = "\n".join(source_lines) - new_content = restore_trailing_newline(content, new_content) - return new_content, occurrences - return None - - -def calculate_regex_replacement( - content: str, - old_string: str, - new_string: str, -) -> tuple[str, int] | None: - """Perform regex-based flexible replacement with whitespace tolerance. - - This function converts the old_string into a regex pattern by escaping special - characters and allowing flexible whitespace between tokens. It's useful for - matching code patterns where whitespace may vary. The replacement preserves - the indentation of the matched block. - - Args: - content: The original content string to modify - old_string: The string pattern to find (converted to regex with flexible whitespace) - new_string: The replacement string (will be indented to match the original) - - Returns: - A tuple containing (modified_content, 1) if a match is found, or None if - no match is found. Only the first match is replaced. - """ - normalized_old = old_string.replace("\r\n", "\n") - normalized_new = new_string.replace("\r\n", "\n") - - delimiters = ["(", ")", ":", "[", "]", "{", "}", ">", "<", "="] - processed = normalized_old - for delim in delimiters: - processed = processed.replace(delim, f" {delim} ") - - tokens = [t for t in processed.split() if t] - if not tokens: - return None - - escaped_tokens = [escape_regex(t) for t in tokens] - pattern = "\\s*".join(escaped_tokens) - final_pattern = f"^(\\s*){pattern}" - regex = re.compile(final_pattern, re.MULTILINE) - - match = regex.search(content) - if not match: - return None - - indent = match.group(1) or "" - new_lines = normalized_new.split("\n") - new_block = "\n".join(f"{indent}{line}" for line in new_lines) - - new_content = regex.sub(new_block, content, count=1) - new_content = restore_trailing_newline(content, new_content) - return new_content, 1 diff --git a/flowllm/gallery/__init__.py b/flowllm/gallery/__init__.py deleted file mode 100644 index e1b9cda..0000000 --- a/flowllm/gallery/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -"""Gallery package for FlowLLM framework.""" - -from . import agent -from . import search -from .chat_op import ChatOp -from .code_analyse_op import CodeAnalyseOp -from .execute_code_op import ExecuteCodeOp -from .gen_system_prompt_op import GenSystemPromptOp -from .stream_chat_op import StreamChatOp -from .token_count_op import TokenCountOp - -__all__ = [ - "agent", - "search", - "ChatOp", - "CodeAnalyseOp", - "ExecuteCodeOp", - "GenSystemPromptOp", - "StreamChatOp", - "TokenCountOp", -] diff --git a/flowllm/gallery/agent/__init__.py b/flowllm/gallery/agent/__init__.py deleted file mode 100644 index 7f9bab3..0000000 --- a/flowllm/gallery/agent/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -"""Convenience exports for agent operators.""" - -from .react_agent_op import ReactAgentOp -from .react_search_op import ReactSearchOp - -__all__ = [ - "ReactAgentOp", - "ReactSearchOp", -] diff --git a/flowllm/gallery/agent/react_agent_op.py b/flowllm/gallery/agent/react_agent_op.py deleted file mode 100644 index dcf3037..0000000 --- a/flowllm/gallery/agent/react_agent_op.py +++ /dev/null @@ -1,270 +0,0 @@ -"""Reactive agent operator that orchestrates tool-augmented LLM reasoning.""" - -import datetime -import time -from typing import List, Dict - -from loguru import logger - -from ...core.context import C, BaseContext -from ...core.enumeration import Role -from ...core.op import BaseAsyncToolOp -from ...core.schema import Message, ToolCall - - -@C.register_op() -class ReactAgentOp(BaseAsyncToolOp): - """React-style agent capable of iterative tool invocation.""" - - file_path: str = __file__ - - def __init__( - self, - llm: str = "qwen3_30b_instruct", - max_steps: int = 5, - tool_call_interval: float = 1.0, - add_think_tool: bool = False, - **kwargs, - ): - """Initialize the agent runtime configuration.""" - super().__init__(llm=llm, **kwargs) - self.max_steps: int = max_steps - self.tool_call_interval: float = tool_call_interval - self.add_think_tool: bool = add_think_tool - - def build_tool_call(self) -> ToolCall: - """Expose metadata describing how to invoke the agent.""" - return ToolCall( - **{ - "description": "A React agent that answers user queries.", - "input_schema": { - "query": { - "type": "string", - "description": "query", - "required": False, - }, - "messages": { - "type": "array", - "description": "messages", - "required": False, - }, - }, - }, - ) - - async def build_tool_op_dict(self) -> dict: - """Collect available tool operators from the execution context.""" - assert isinstance(self.ops, BaseContext), "self.ops must be BaseContext" - tool_op_dict: Dict[str, BaseAsyncToolOp] = { - op.tool_call.name: op for op in self.ops.values() if isinstance(op, BaseAsyncToolOp) - } - for op in tool_op_dict.values(): - op.language = self.language - return tool_op_dict - - async def build_messages(self) -> List[Message]: - """Build the initial message history for the LLM.""" - if "query" in self.input_dict: - query: str = self.input_dict["query"] - now_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") - messages = [ - Message(role=Role.SYSTEM, content=self.prompt_format(prompt_name="system_prompt", time=now_time)), - Message(role=Role.USER, content=query), - ] - logger.info(f"round0.system={messages[0].model_dump_json()}") - logger.info(f"round0.user={messages[1].model_dump_json()}") - - elif "messages" in self.input_dict: - messages = self.input_dict["messages"] - messages = [Message(**x) for x in messages] - - logger.info(f"round0.user={messages[-1].model_dump_json()}") - else: - raise ValueError("input_dict must contain either 'query' or 'messages'") - - return messages - - async def before_chat(self, messages: List[Message]): - """Prepare the message history for the LLM.""" - return messages - - async def execute_tool(self, op: BaseAsyncToolOp, tool_call: ToolCall): - """Execute a tool operation asynchronously using the provided tool call arguments.""" - self.submit_async_task(op.async_call, **tool_call.argument_dict) - - async def _reasoning_step( - self, - messages: List[Message], - tool_op_dict: Dict[str, BaseAsyncToolOp], - step: int, - ) -> tuple[Message, bool]: - """ - Perform the reasoning step: prepare messages and invoke LLM to generate response. - - This step represents the "think" phase of the ReAct loop, where the agent - analyzes the current context and decides what actions to take. - - Args: - messages: Current message history for the conversation. - tool_op_dict: Dictionary mapping tool names to their operator instances. - step: Current iteration step number (0-indexed). - - Returns: - A tuple containing: - - assistant_message: The LLM's response message, which may contain tool calls. - - should_continue: Boolean indicating whether to continue the loop - (False if no tool calls are needed, meaning the agent has finished). - """ - # Prepare messages for LLM input (e.g., formatting, filtering) - messages = await self.before_chat(messages) - - # Invoke LLM with current context and available tools - assistant_message: Message = await self.llm.achat( - messages=messages, - tools=[op.tool_call for op in tool_op_dict.values()], - ) - - # Append the assistant's response to message history - messages.append(assistant_message) - logger.info(f"round{step + 1}.assistant={assistant_message.model_dump_json()}") - - # Check if the agent wants to use tools or has finished reasoning - should_continue = bool(assistant_message.tool_calls) - - return assistant_message, should_continue - - async def _acting_step( - self, - assistant_message: Message, - tool_op_dict: Dict[str, BaseAsyncToolOp], - think_op: BaseAsyncToolOp, - step: int, - ) -> List[Message]: - """ - Perform the acting step: execute tool calls and collect results. - - This step represents the "act" phase of the ReAct loop, where the agent - executes the tools it decided to use in the reasoning step and incorporates - their results back into the conversation context. - - Args: - assistant_message: The LLM's response message containing tool calls. - tool_op_dict: Dictionary mapping tool names to their operator instances. - think_op: The think tool operator instance (for dynamic think tool management). - step: Current iteration step number (0-indexed). - - Returns: - A list of tool result messages to be appended to the conversation history. - """ - if not assistant_message.tool_calls: - return [] - - op_list: List[BaseAsyncToolOp] = [] - has_think_tool_flag: bool = False - tool_result_messages: List[Message] = [] - - # Phase 1: Submit all tool calls for parallel execution - for j, tool_call in enumerate(assistant_message.tool_calls): - # Track if think_tool was used (for dynamic tool management) - if tool_call.name == think_op.tool_call.name: - has_think_tool_flag = True - - # Validate tool exists in available tools - if tool_call.name not in tool_op_dict: - logger.exception(f"unknown tool_call.name={tool_call.name}") - continue - - logger.info( - f"round{step + 1}.{j} submit tool_calls={tool_call.name} " f"argument={tool_call.argument_dict}", - ) - - # Create a copy of the tool operator for this specific invocation - op_copy: BaseAsyncToolOp = tool_op_dict[tool_call.name].copy() - op_copy.tool_call.id = tool_call.id - op_list.append(op_copy) - - # Submit tool execution asynchronously - await self.execute_tool(op_copy, tool_call) - time.sleep(self.tool_call_interval) - - # Phase 2: Wait for all tool executions to complete - await self.join_async_task() - - # Phase 3: Collect tool results and format as tool messages - for j, op in enumerate(op_list): - tool_result = str(op.output) - tool_message = Message( - role=Role.TOOL, - content=tool_result, - tool_call_id=op.tool_call.id, - ) - tool_result_messages.append(tool_message) - logger.info( - f"round{step + 1}.{j} join tool_result={tool_result[:200]}...\n\n", - ) - - # Phase 4: Manage think_tool availability dynamically - # If think_tool was used, remove it to prevent repeated use in next step - # If it wasn't used, ensure it's available for the next step - if self.add_think_tool: - if not has_think_tool_flag: - tool_op_dict["think_tool"] = think_op - else: - tool_op_dict.pop("think_tool", None) - - return tool_result_messages - - async def async_execute(self): - """ - Main execution loop implementing the ReAct (Reasoning + Acting) pattern. - - The agent alternates between: - 1. Reasoning: Invoking the LLM to analyze the situation and decide on actions - 2. Acting: Executing the chosen tools and incorporating their results - - This loop continues until: - - The agent decides no more tools are needed (final answer reached) - - The maximum number of steps is reached - """ - from ..think_tool_op import ThinkToolOp - - # Initialize think tool operator if needed - think_op = ThinkToolOp(language=self.language) - - # Build dictionary of available tool operators from context - tool_op_dict = await self.build_tool_op_dict() - - # Optionally add think_tool to available tools - if self.add_think_tool: - tool_op_dict["think_tool"] = think_op - - # Initialize conversation message history - messages = await self.build_messages() - - # Main ReAct loop: alternate between reasoning and acting - for step in range(self.max_steps): - # Reasoning step: LLM analyzes context and decides on actions - assistant_message, should_continue = await self._reasoning_step( - messages, - tool_op_dict, - step, - ) - - # If no tool calls, the agent has reached a final answer - if not should_continue: - break - - # Acting step: execute tools and collect results - tool_result_messages = await self._acting_step( - assistant_message, - tool_op_dict, - think_op, - step, - ) - - # Append tool results to message history for next reasoning step - messages.extend(tool_result_messages) - - # Set final output and store full conversation history in metadata - self.set_output(messages[-1].content) - self.context.response.metadata["messages"] = messages diff --git a/flowllm/gallery/agent/react_agent_prompt.yaml b/flowllm/gallery/agent/react_agent_prompt.yaml deleted file mode 100644 index 84e65a6..0000000 --- a/flowllm/gallery/agent/react_agent_prompt.yaml +++ /dev/null @@ -1,5 +0,0 @@ -system_prompt: | - You are a helpful assistant. The current time is {time}. - -system_prompt_zh: | - 你是一个有用的助手。当前时间是 {time}。 diff --git a/flowllm/gallery/agent/react_search_op.py b/flowllm/gallery/agent/react_search_op.py deleted file mode 100644 index e1a654b..0000000 --- a/flowllm/gallery/agent/react_search_op.py +++ /dev/null @@ -1,23 +0,0 @@ -"""ReactAgentOp specialization that ensures search capability is available.""" - -from typing import Dict - -from .react_agent_op import ReactAgentOp -from ...core.context import C -from ...core.op import BaseAsyncToolOp - - -@C.register_op() -class ReactSearchOp(ReactAgentOp): - """Agent that guarantees a search tool fallback when none are configured.""" - - async def build_tool_op_dict(self) -> dict: - """Extend parent tools with a default search operator when needed.""" - tool_op_dict: Dict[str, BaseAsyncToolOp] = await super().build_tool_op_dict() - if not tool_op_dict: - from ..search.dashscope_search_op import DashscopeSearchOp - - search_op = DashscopeSearchOp() - tool_op_dict[search_op.tool_call.name] = search_op - - return tool_op_dict diff --git a/flowllm/gallery/chat_op.py b/flowllm/gallery/chat_op.py deleted file mode 100644 index ab309a6..0000000 --- a/flowllm/gallery/chat_op.py +++ /dev/null @@ -1,56 +0,0 @@ -"""LLM chat operation for interactive conversations with LLM.""" - -from loguru import logger - -from ..core.context import C -from ..core.enumeration import Role -from ..core.op import BaseAsyncOp -from ..core.schema import Message - - -@C.register_op() -class ChatOp(BaseAsyncOp): - """Operation for conducting chat conversations with LLM. - - This operation combines system prompt and user messages, sends them to the LLM, - and stores the response in context.response.answer. It requires both messages - and system_prompt to be present in the context. - """ - - file_path: str = __file__ - - def __init__( - self, - llm: str = "qwen3_30b_instruct", - **kwargs, - ): - """Initialize the ChatOp. - - Args: - llm: Name of the LLM to use for chat. - **kwargs: Additional arguments passed to BaseAsyncOp. - """ - super().__init__(llm=llm, **kwargs) - - async def async_execute(self): - """Execute the LLM chat operation. - - Combines system prompt and messages, sends to LLM, and stores the response - in context.response.answer. Requires context.messages (list of Message) and - context.system_prompt (str) to be set. - """ - messages = self.context.messages - assert isinstance(messages, list) and all( - isinstance(m, Message) for m in messages - ), "`messages` must be a list of Message objects!" - - system_prompt = self.context.system_prompt - assert system_prompt, "`system_prompt` is required!" - - messages = [Message(role=Role.SYSTEM, content=system_prompt)] + messages - logger.info(f"messages={messages}") - - response = await self.llm.achat(messages=messages, tools=None) - assert isinstance(response, Message), "Response must be a Message object!" - - self.context.response.answer = response.content.strip() diff --git a/flowllm/gallery/code_analyse_op.py b/flowllm/gallery/code_analyse_op.py deleted file mode 100644 index c96ba8d..0000000 --- a/flowllm/gallery/code_analyse_op.py +++ /dev/null @@ -1,244 +0,0 @@ -"""Code analysis operation for traversing and analyzing code files. - -This module provides the CodeAnalyseOp class which can traverse code files, -judge their relevance to a query, and produce explanations. - -Example usage: - flowllm backend=cmd \\ - cmd.flow="CodeAnalyseOp()" \\ - cmd.params.query="对话上下文(list of messages)超过一定长度后如何管理，保留关键代码" \\ - cmd.params.code_dir="/Users/yuli/workspace/gemini-cli/packages/core" \\ - cmd.params.file_suffix="ts,tsx" \\ - cmd.params.exclude_suffix="test.ts" \\ - cmd.params.output_dir="tmp" -""" - -import json -from pathlib import Path - -from loguru import logger -from tqdm import tqdm - -from ..core.context import C -from ..core.enumeration import Role -from ..core.op import BaseAsyncToolOp -from ..core.schema import Message, ToolCall -from ..core.utils.llm_utils import parse_message_by_keys - - -@C.register_op() -class CodeAnalyseOp(BaseAsyncToolOp): - """Operation for analyzing code files with relevance judgment and explanations. - - This operation traverses code files in a directory, judges their relevance - to a user query (if provided), and generates explanations. It supports - filtering by file suffixes and excluding specific patterns. - """ - - file_path: str = __file__ - - def __init__( - self, - llm: str = "qwen3_30b_instruct", - max_chars_per_file: int = 30000, - max_parallel_cnt: int = 4, - **kwargs, - ): - super().__init__(llm=llm, **kwargs) - self.max_chars_per_file = max_chars_per_file - self.max_parallel_cnt = max_parallel_cnt - - def build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "Traverse code files, judge relevance to a query, and produce explanations.", - "input_schema": { - "query": { - "type": "string", - "description": "User question or task the code should address. " - "If empty, the code will be directly analyzed without relevance judgment.", - "required": False, - }, - "code_dir": { - "type": "string", - "description": "Absolute path to the root directory of the code to inspect.", - "required": True, - }, - "file_suffix": { - "type": "string", - "description": "File suffix to include when scanning, e.g. 'py'.", - "required": True, - }, - "exclude_suffix": { - "type": "string", - "description": "File suffix to exclude when scanning, e.g. 'test.py'. " - "Multiple suffixes can be separated by commas.", - "required": False, - }, - "output_dir": { - "type": "string", - "description": "Absolute path to the directory where analysis files will be written.", - "required": True, - }, - }, - }, - ) - - async def async_execute(self): - query: str = self.input_dict.get("query", "").strip() - code_dir = Path(self.input_dict["code_dir"]).expanduser().resolve() - output_dir = Path(self.input_dict["output_dir"]).expanduser().resolve() - suffix_input = self.input_dict["file_suffix"].strip() - suffix_tokens = [token.strip() for token in suffix_input.split(",") if token.strip()] - normalized_suffixes = {token if token.startswith(".") else f".{token}" for token in suffix_tokens} - - if not normalized_suffixes: - error_msg = f"file_suffix is empty or invalid: {suffix_input}" - logger.error(error_msg) - self.set_output(error_msg) - return - - # 处理排除的后缀 - exclude_suffix_input = self.input_dict.get("exclude_suffix", "").strip() - exclude_suffix_tokens = [token.strip() for token in exclude_suffix_input.split(",") if token.strip()] - normalized_exclude_suffixes = { - token if token.startswith(".") else f".{token}" for token in exclude_suffix_tokens - } - - if not code_dir.exists() or not code_dir.is_dir(): - error_msg = f"code_dir does not exist or is not a directory: {code_dir}" - logger.error(error_msg) - self.set_output(error_msg) - return - - output_dir.mkdir(parents=True, exist_ok=True) - - matching_files = [ - path - for path in sorted(code_dir.rglob("*")) - if path.is_file() - and path.suffix in normalized_suffixes - and not any(path.name.endswith(exclude_suffix) for exclude_suffix in normalized_exclude_suffixes) - ] - - exclude_info = f" (excluding {sorted(normalized_exclude_suffixes)})" if normalized_exclude_suffixes else "" - logger.info( - f"{self.name}: Found {len(matching_files)} file(s) " - f"with suffixes {sorted(normalized_suffixes)}{exclude_info}", - ) - - # 根据 query 是否为空选择不同的 prompt - is_query_mode = bool(query) - if is_query_mode: - system_prompt = self.get_prompt("code_judge_system_prompt") - else: - system_prompt = self.get_prompt("code_analyse_system_prompt") - - max_parallel_cnt = max(1, self.max_parallel_cnt) - for file_path in tqdm(matching_files): - while len(self.task_list) >= max_parallel_cnt: - await self.join_async_task() - - self.submit_async_task( - self._process_single_file, - file_path, - code_dir, - output_dir, - query, - system_prompt, - is_query_mode, - ) - - await self.join_async_task() - - async def _process_single_file( - self, - file_path: Path, - code_dir: Path, - output_dir: Path, - query: str, - system_prompt: str, - is_query_mode: bool, - ): - raw_content = file_path.read_text(encoding="utf-8", errors="ignore") - - truncated = len(raw_content) > self.max_chars_per_file - snippet = raw_content[: self.max_chars_per_file] - - relative_path = str(file_path.relative_to(code_dir)) - - # 根据模式选择不同的 user prompt - if is_query_mode: - user_prompt = self.prompt_format( - "code_judge_user_prompt", - query=query, - file_path=str(file_path), - relative_path=relative_path, - truncated_flag="是" if truncated else "否", - code_language=file_path.suffix.lstrip(".") or "text", - code_content=snippet, - ) - else: - user_prompt = self.prompt_format( - "code_analyse_user_prompt", - file_path=str(file_path), - relative_path=relative_path, - truncated_flag="是" if truncated else "否", - code_language=file_path.suffix.lstrip(".") or "text", - code_content=snippet, - ) - - messages = [ - Message(role=Role.SYSTEM, content=system_prompt), - Message(role=Role.USER, content=user_prompt), - ] - - def parse_result(message: Message): - content = message.content.strip() - - if is_query_mode: - # 查询模式：需要判断相关性 - parsed = parse_message_by_keys(content, ["### think", "### result", "### explanation"]) - think_content = parsed.get("### think", "").strip() - result_content = parsed.get("### result", "").strip() - is_relevant: bool = "true" in result_content.lower() - if is_relevant: - explanation_content = parsed.get("### explanation", "").strip() - else: - explanation_content = "" - return { - "is_relevant": is_relevant, - "think_content": think_content, - "explanation_content": explanation_content, - } - else: - # 直接解读模式：不需要判断相关性，直接生成解释 - parsed = parse_message_by_keys(content, ["### think", "### explanation"]) - think_content = parsed.get("### think", "").strip() - explanation_content = parsed.get("### explanation", "").strip() - return { - "is_relevant": True, - "think_content": think_content, - "explanation_content": explanation_content, - } - - parsed_result = await self.llm.achat(messages=messages, callback_fn=parse_result) - logger.info( - f"file_path={file_path} parsed_result={json.dumps(parsed_result, ensure_ascii=False, indent=2)}", - ) - if not parsed_result["is_relevant"]: - logger.info(f"{self.name}: File {relative_path} deemed not relevant.") - return - - analysis_path = (output_dir / relative_path).with_suffix(".md") - analysis_path.parent.mkdir(parents=True, exist_ok=True) - - # think = parsed_result["think_content"] - explanation = parsed_result["explanation_content"] - analysis_text = "\n\n".join( - [ - str(file_path.absolute()), - explanation, - ], - ) - analysis_path.write_text(analysis_text, encoding="utf-8") diff --git a/flowllm/gallery/code_analyse_prompt.yaml b/flowllm/gallery/code_analyse_prompt.yaml deleted file mode 100644 index b1b74d4..0000000 --- a/flowllm/gallery/code_analyse_prompt.yaml +++ /dev/null @@ -1,64 +0,0 @@ -code_judge_system_prompt: | - 你是一名资深的软件工程与代码审阅专家，擅长判断源代码是否能回答用户问题，并在相关时给出面向实际需求的解释。请保持严谨、客观，并使用与用户问题相同的语言表达。请严格遵循指定的输出格式，使用 Markdown 标题格式输出。 - -code_judge_user_prompt: | - # 文件路径 - {file_path} - - # 相对路径 - {relative_path} - - # 代码片段 - - 是否截断: {truncated_flag} - ```{code_language} - {code_content} - ``` - - # 用户问题 - {query} - - # 任务 - 1. 在 ### think 标题下思考用户问题是否和代码片段有关联，并概述判断理由。 - 2. 在 ### result 标题下输出 true 或 false。仅当判断为相关（true）时，在 ### explanation 标题下总结代码如何回答用户问题，指出关键函数、数据结构或逻辑；如判断为不相关（false），不要输出 ### explanation。 - 3. 不要输出除此之外的任何内容。 - - # 输出格式 - 按顺序输出以下 Markdown 标题和内容： - ### think - 思考是否和问题相关 - - ### result - 是否相关的结果，只能输出 true 或 false - - 如判断为相关，请继续输出： - ### explanation - 解读与说明 - -code_analyse_system_prompt: | - 你是一名资深的软件工程与代码审阅专家，擅长解读和分析源代码。请对代码进行全面的解读，包括代码的功能、关键函数、数据结构、逻辑流程等。请保持严谨、客观，使用清晰的语言表达。请严格遵循指定的输出格式，使用 Markdown 标题格式输出。 - -code_analyse_user_prompt: | - # 文件路径 - {file_path} - - # 相对路径 - {relative_path} - - # 代码片段 - - 是否截断: {truncated_flag} - ```{code_language} - {code_content} - ``` - - # 任务 - 1. 在 ### think 标题下思考要如何进行解读。 - 2. 在 ### explanation 标题下详细解读代码 - 3. 不要输出除此之外的任何内容。 - - # 输出格式 - 按顺序输出以下 Markdown 标题和内容： - ### think - 思考要如何进行解读 - - ### explanation - 详细的代码解读与说明 diff --git a/flowllm/gallery/execute_code_op.py b/flowllm/gallery/execute_code_op.py deleted file mode 100644 index b2a8887..0000000 --- a/flowllm/gallery/execute_code_op.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Execute code operation for running Python code dynamically.""" - -import asyncio -import sys -from io import StringIO - -from loguru import logger - -from ..core.context import C -from ..core.op.base_async_tool_op import BaseAsyncToolOp -from ..core.schema import ToolCall - - -@C.register_op() -class ExecuteCodeOp(BaseAsyncToolOp): - """Operation for executing Python code dynamically. - - This operation allows for dynamic execution of Python code in scenarios such as - analysis or calculation. The code output is captured from stdout and returned - as the result. If execution fails, the exception message is returned. - """ - - def build_tool_call(self) -> ToolCall: - """Build the tool call definition for code execution. - - Returns: - ToolCall object defining the code execution tool with input schema. - """ - return ToolCall( - **{ - "description": "Execute python code can be used in scenarios such as analysis or calculation, " - "and the final result can be printed using the `print` function.", - "input_schema": { - "code": { - "type": "string", - "description": "code to be executed", - "required": True, - }, - }, - }, - ) - - def execute(self): - """Execute the Python code from input_dict. - - Captures stdout output and handles exceptions. The result is set using - set_result() method. - """ - old_stdout = sys.stdout - redirected_output = sys.stdout = StringIO() - - try: - code: str = self.input_dict["code"] - exec(code) - code_result = redirected_output.getvalue() - - except Exception as e: - logger.info(f"{self.name} encounter exception! error={e.args}") - code_result = str(e) - - sys.stdout = old_stdout - self.set_output(code_result) - - async def async_execute(self): - """Execute code asynchronously in a thread pool. - - Runs the synchronous execute() method in a thread pool to avoid blocking - the event loop. - """ - loop = asyncio.get_event_loop() - await loop.run_in_executor(C.thread_pool, self.execute) diff --git a/flowllm/gallery/gen_system_prompt_op.py b/flowllm/gallery/gen_system_prompt_op.py deleted file mode 100644 index 5ad4622..0000000 --- a/flowllm/gallery/gen_system_prompt_op.py +++ /dev/null @@ -1,85 +0,0 @@ -"""Generate system prompt operation for creating optimized system prompts.""" - -from loguru import logger - -from ..core.context import C -from ..core.enumeration import Role -from ..core.op import BaseAsyncOp -from ..core.schema import Message -from ..core.utils import extract_content, format_messages - - -@C.register_op() -class GenSystemPromptOp(BaseAsyncOp): - """Operation for generating optimized system prompts using LLM. - - This operation takes user query or messages and generates an optimized system - prompt using an LLM. The generated prompt is extracted from the LLM response, - which should contain think and prompt sections in markdown code blocks. - """ - - file_path: str = __file__ - - def __init__( - self, - llm: str = "qwen3_30b_instruct", - **kwargs, - ): - """Initialize the GenSystemPromptOp. - - Args: - llm: Name of the LLM to use for prompt generation. - **kwargs: Additional arguments passed to BaseAsyncOp. - """ - super().__init__(llm=llm, **kwargs) - - async def async_execute(self): - """Execute the system prompt generation. - - Reads query or messages from context, formats them, and uses LLM to generate - an optimized system prompt. The prompt is extracted from the response and - stored in context.system_prompt. - """ - query = self.context.get("query") - - if query: - messages = [Message(role=Role.USER, content=query)] - - else: - messages = self.context.get("messages", []) - if messages: - messages = [Message(**x) for x in messages] - - assert messages, "Both `query` and `messages` are not provided!" - self.context.messages = messages - - def callback_fn(message: Message): - """Extract system prompt from LLM response. - - The LLM response should contain think and prompt sections in markdown code blocks. - If extraction fails, fall back to the full message content. - """ - think_content = extract_content(message.content, "think") - prompt_content = extract_content(message.content, "prompt") - - logger.info(f"think_content={think_content}\nprompt_content={prompt_content}") - - # If prompt_content extraction failed, use the full message content as fallback - if prompt_content is None: - logger.warning("Failed to extract prompt from code block, using full message content") - # Get content, handling both str and bytes - content = message.content - if isinstance(content, bytes): - content = content.decode("utf-8", errors="ignore") - prompt_content = content.strip() - - return prompt_content - - user_prompt = self.prompt_format("gen_system_prompt_prompt", formated_messages=format_messages(messages)) - - system_prompt = await self.llm.achat( - messages=[Message(role=Role.USER, content=user_prompt)], - tools=None, - callback_fn=callback_fn, - ) - self.context.system_prompt = system_prompt diff --git a/flowllm/gallery/gen_system_prompt_prompt.yaml b/flowllm/gallery/gen_system_prompt_prompt.yaml deleted file mode 100644 index 9a5a2f4..0000000 --- a/flowllm/gallery/gen_system_prompt_prompt.yaml +++ /dev/null @@ -1,73 +0,0 @@ -gen_system_prompt_prompt: | - # Context - Below are the user's questions or conversation history: - {formated_messages} - - # Task - Based on the user's questions or conversation content, dynamically generate an appropriate system prompt that helps the LLM better understand and answer the user's questions. - - # Requirements - 1. **Emotion Recognition**: Analyze the emotional state of the user's question (such as anger, anxiety, happiness, confusion, etc.). If the user's emotion is negative, the system prompt should include empathy, patience, and understanding. - 2. **Scenario Adaptation**: Adjust the tone and professionalism of the system prompt according to the question type (such as technical questions, emotional counseling, daily conversations, etc.). - 3. **Role Positioning**: Clearly define the role the LLM should play (such as professional consultant, friend, assistant, etc.). - 4. **Response Style**: Determine the style that should be adopted in responses (such as concise, detailed, gentle, professional, etc.). - 5. **Special Needs**: Identify possible special needs of the user (such as needing encouragement, professional advice, comfort, etc.). - - # Examples - - If the user's question shows anger or frustration: the system prompt should emphasize empathy, patience, understanding, and reassurance - - If the user's question is a technical question: the system prompt should emphasize professionalism, accuracy, and clarity - - If the user's question is emotional counseling: the system prompt should emphasize listening, support, and advice - - If the user's question is a daily conversation: the system prompt should be natural, friendly, and casual - - # Output Format - Please output in the following format: - - ``` think - Analyze the user's question regarding emotions, scenarios, and role requirements, and think about what kind of system prompt should be generated to help the LLM better answer the user's question. - Include: - - User emotional state analysis - - Question type and scenario identification - - Traits and capabilities that need to be emphasized - - Response style suggestions - ``` - - ``` prompt - Output the generated system prompt here. It should be clear, specific, targeted, and able to guide the LLM to answer the user's question in an appropriate manner. - ``` - -gen_system_prompt_prompt_zh: | - # 上下文 - 以下是用户的问题或对话历史： - {formated_messages} - - # 任务 - 根据用户的问题或对话内容，动态生成一个适合的system prompt，帮助LLM更好地理解和回答用户的问题。 - - # 要求 - 1. **情绪识别**：分析用户问题的情绪状态（如生气、焦虑、开心、困惑等），如果用户情绪负面，system prompt应该包含共情能力、耐心和理解。 - 2. **场景适配**：根据问题类型（如技术问题、情感咨询、日常对话等）调整system prompt的语气和专业度。 - 3. **角色定位**：明确LLM应该扮演的角色（如专业顾问、朋友、助手等）。 - 4. **回答风格**：确定回答应该采用的风格（如简洁、详细、温和、专业等）。 - 5. **特殊需求**：识别用户可能的特殊需求（如需要鼓励、需要专业建议、需要安慰等）。 - - # 示例 - - 如果用户问题表现出愤怒或沮丧：system prompt应该强调共情、耐心、理解和安抚 - - 如果用户问题是技术问题：system prompt应该强调专业性、准确性和清晰度 - - 如果用户问题是情感咨询：system prompt应该强调倾听、支持和建议 - - 如果用户问题是日常对话：system prompt应该自然、友好、随和 - - # 输出格式 - 请按照以下格式输出： - - ``` think - 分析用户问题的情绪、场景、角色需求，思考应该生成什么样的system prompt可以让llm更好的回答用户问题。 - 包括： - - 用户情绪状态分析 - - 问题类型和场景识别 - - 需要强调的特质和能力 - - 回答风格建议 - ``` - - ``` prompt - 在这里输出生成的system prompt，应该清晰、具体、针对性强，能够指导LLM以合适的方式回答用户问题。 - ``` \ No newline at end of file diff --git a/flowllm/gallery/search/__init__.py b/flowllm/gallery/search/__init__.py deleted file mode 100644 index 7366b5f..0000000 --- a/flowllm/gallery/search/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -"""Search gallery.""" - -from .dashscope_search_op import DashscopeSearchOp -from .mcp_search_op import TongyiMcpSearchOp, BochaMcpSearchOp -from .mock_search_op import MockSearchOp -from .tavily_search_op import TavilySearchOp - -__all__ = [ - "DashscopeSearchOp", - "MockSearchOp", - "TongyiMcpSearchOp", - "BochaMcpSearchOp", - "TavilySearchOp", -] diff --git a/flowllm/gallery/search/dashscope_search_op.py b/flowllm/gallery/search/dashscope_search_op.py deleted file mode 100644 index 35b6dd7..0000000 --- a/flowllm/gallery/search/dashscope_search_op.py +++ /dev/null @@ -1,117 +0,0 @@ -"""Dashscope web search operation module. - -This module provides a tool operation for performing web searches using the Dashscope API. -It enables LLM models to retrieve relevant information from the internet by executing -search queries and returning formatted results. -""" - -import os - -import dashscope -from loguru import logger - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - - -@C.register_op() -class DashscopeSearchOp(BaseAsyncToolOp): - """A tool operation for performing web searches using Dashscope API. - - This operation enables LLM models to search the internet for information by - providing search keywords. It supports various search strategies and can - optionally use role prompts to enhance search queries. - - Attributes: - model: The Dashscope model to use for search (default: "qwen-plus"). - search_strategy: The search strategy to use (default: "max"). - enable_role_prompt: Whether to use role prompts for query enhancement. - api_key: Dashscope API key loaded from environment variable. - """ - - file_path: str = __file__ - - def __init__( - self, - model: str = "qwen-plus", - search_strategy: str = "max", - enable_role_prompt: bool = True, - **kwargs, - ): - super().__init__(**kwargs) - - self.model: str = model - self.search_strategy: str = search_strategy - self.enable_role_prompt: bool = enable_role_prompt - - self.api_key = os.getenv("DASHSCOPE_API_KEY", "") - # https://help.aliyun.com/zh/model-studio/web-search?spm=a2c4g.11186623.help-menu-2400256.d_0_7_0.670e253awctI43&scm=20140722.H_2867560._.OR_help-T_cn~zh-V_1 - - def build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "Use search keywords to retrieve relevant information from the internet. " - "If you have multiple keywords, please call this tool separately for each one.", - "input_schema": { - "query": { - "type": "string", - "description": "search keyword", - "required": True, - }, - }, - }, - ) - - async def async_execute(self): - query: str = self.input_dict["query"] - - if self.enable_cache: - cached_result = self.cache.load(query) - if cached_result: - self.set_output(cached_result["response_content"]) - return - - if self.enable_role_prompt: - user_query = self.prompt_format(prompt_name="role_prompt", query=query) - else: - user_query = query - logger.info(f"user_query={user_query}") - messages: list = [{"role": "user", "content": user_query}] - - response = await dashscope.AioGeneration.call( - api_key=self.api_key, - model=self.model, - messages=messages, - enable_search=True, # Enable web search - search_options={ - "forced_search": True, # Force web search - "enable_source": True, # Include search source information - "enable_citation": False, # Enable citation markers - "search_strategy": self.search_strategy, # Search strategy - }, - result_format="message", - ) - - search_results = [] - response_content = "" - - if hasattr(response, "output") and response.output: - if hasattr(response.output, "search_info") and response.output.search_info: - search_results = response.output.search_info.get("search_results", []) - - if hasattr(response.output, "choices") and response.output.choices and len(response.output.choices) > 0: - response_content = response.output.choices[0].message.content - - final_result = { - "query": query, - "search_results": search_results, - "response_content": response_content, - "model": self.model, - "search_strategy": self.search_strategy, - } - - if self.enable_cache: - self.cache.save(query, final_result, expire_hours=self.cache_expire_hours) - - self.set_output(final_result["response_content"]) diff --git a/flowllm/gallery/search/dashscope_search_prompt.yaml b/flowllm/gallery/search/dashscope_search_prompt.yaml deleted file mode 100644 index b8698ab..0000000 --- a/flowllm/gallery/search/dashscope_search_prompt.yaml +++ /dev/null @@ -1,13 +0,0 @@ -role_prompt: | - # user's question - {query} - - # task - Extract the original content related to the user's query directly from the context, maintain accuracy, and avoid excessive processing. - -role_prompt_zh: | - # 用户问题 - {query} - - # task - 直接从上下文中提取与用户问题相关的原始内容，保持准确性，避免过度处理。 \ No newline at end of file diff --git a/flowllm/gallery/search/mcp_search_op.py b/flowllm/gallery/search/mcp_search_op.py deleted file mode 100644 index 80f6ea9..0000000 --- a/flowllm/gallery/search/mcp_search_op.py +++ /dev/null @@ -1,68 +0,0 @@ -"""MCP web search operation module. - -This module provides tool operations for performing web searches using MCP (Model Context Protocol). -It enables LLM models to retrieve relevant information from the internet by executing -search queries through different MCP search providers. -""" - -from ...core.context import C -from ...core.op import BaseMcpOp - - -@C.register_op() -class TongyiMcpSearchOp(BaseMcpOp): - """A tool operation for performing web searches using Tongyi MCP. - - This operation enables LLM models to search the internet for information by - providing search keywords through the Tongyi search MCP server. It uses the - bailian_web_search tool to execute search queries. - - Attributes: - mcp_name: The MCP server name (default: "tongyi_search"). - tool_name: The tool name to call (default: "bailian_web_search"). - save_answer: Whether to save the search answer (default: True). - input_schema_optional: Optional input schema fields (default: ["count"]). - input_schema_deleted: Deleted input schema fields (default: ["ctx"]). - """ - - def __init__(self, **kwargs): - kwargs.update( - { - "mcp_name": "tongyi_search", - "tool_name": "bailian_web_search", - "save_answer": True, - "input_schema_optional": ["count"], - "input_schema_deleted": ["ctx"], - }, - ) - super().__init__(**kwargs) - - -@C.register_op() -class BochaMcpSearchOp(BaseMcpOp): - """A tool operation for performing web searches using BochaAI MCP. - - This operation enables LLM models to search the internet for information by - providing search keywords through the BochaAI search MCP server. It uses the - bocha_web_search tool to execute search queries with support for freshness - and count parameters. - - Attributes: - mcp_name: The MCP server name (default: "bochaai_search"). - tool_name: The tool name to call (default: "bocha_web_search"). - save_answer: Whether to save the search answer (default: True). - input_schema_optional: Optional input schema fields (default: ["freshness", "count"]). - input_schema_deleted: Deleted input schema fields (default: ["ctx"]). - """ - - def __init__(self, **kwargs): - kwargs.update( - { - "mcp_name": "bochaai_search", - "tool_name": "bocha_web_search", - "save_answer": True, - "input_schema_optional": ["freshness", "count"], - "input_schema_deleted": ["ctx"], - }, - ) - super().__init__(**kwargs) diff --git a/flowllm/gallery/search/mock_search_op.py b/flowllm/gallery/search/mock_search_op.py deleted file mode 100644 index 29aa140..0000000 --- a/flowllm/gallery/search/mock_search_op.py +++ /dev/null @@ -1,67 +0,0 @@ -"""Mock search operation that uses LLM to generate search results.""" - -import json -import random - -from loguru import logger - -from ...core.context import C -from ...core.enumeration import Role -from ...core.op import BaseAsyncOp -from ...core.schema import Message -from ...core.utils import extract_content - - -@C.register_op() -class MockSearchOp(BaseAsyncOp): - """Mock search operation that uses LLM to generate realistic search results. - - This operation takes a search query and uses an LLM to generate 0-5 mock - search results. Each result contains snippet, title, url, hostname, and hostlogo. - """ - - file_path: str = __file__ - - def __init__( - self, - llm: str = "qwen3_30b_instruct", - **kwargs, - ): - """Initialize the MockSearchOp. - - Args: - llm: Name of the LLM to use for generating search results. - **kwargs: Additional arguments passed to BaseAsyncToolOp. - """ - super().__init__(llm=llm, **kwargs) - - async def async_execute(self): - """Execute the mock search operation. - - Uses LLM to generate realistic search results based on the query. - Results are randomly generated between 0-5 items. - """ - query = self.context.query - if not query: - logger.warning(f"{self.name}: query is empty") - self.context.response.answer = "No results found." - return - - num_results = random.randint(0, 5) - user_prompt = self.prompt_format("mock_search_op_prompt", query=query, num_results=num_results) - - messages = [ - Message( - role=Role.SYSTEM, - content="You are a helpful assistant that generates realistic search results in JSON format.", - ), - Message(role=Role.USER, content=user_prompt), - ] - - logger.info(f"{self.name}: Generating {num_results} results for query: {query}") - - def callback_fn(message: Message): - return extract_content(message.content, "json") - - search_results: str = await self.llm.achat(messages=messages, callback_fn=callback_fn) - self.context.response.answer = json.dumps(search_results, ensure_ascii=False, indent=2) diff --git a/flowllm/gallery/search/mock_search_prompt.yaml b/flowllm/gallery/search/mock_search_prompt.yaml deleted file mode 100644 index a384252..0000000 --- a/flowllm/gallery/search/mock_search_prompt.yaml +++ /dev/null @@ -1,37 +0,0 @@ -mock_search_op_prompt: | - # Task - Generate {num_results} realistic search results for the query: "{query}". - - # Fields per item - Each result must be a JSON object with fields: - - snippet: 2-3 sentence summary - - title: page title - - url: realistic URL (e.g., https://example.com/article/title) - - hostname: domain (e.g., example.com) - - hostlogo: logo URL (e.g., https://example.com/logo.png) or empty string - - # Requirements - - Ensure relevance to the query - - Use diverse, realistic sources - - Ensure well-formed URLs - - If no relevant results, return an empty array - - # Output Format - First, think briefly about good sources and angles: - ``` think - your brief reasoning here - ``` - - Then output ONLY the JSON array wrapped in a json code block, nothing else: - ``` json - [ - {{ - "snippet": "核心内容", - "title": "...", - "url": "...", - "hostname": "...", - "hostlogo": "..." - }} - ] - ``` - diff --git a/flowllm/gallery/search/tavily_search_op.py b/flowllm/gallery/search/tavily_search_op.py deleted file mode 100644 index 9928b24..0000000 --- a/flowllm/gallery/search/tavily_search_op.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Tavily web search operation module. - -This module provides a tool operation for performing web searches using the Tavily API. -It enables LLM models to retrieve relevant information from the internet by executing -search queries and optionally extracting content from search results. -""" - -import json -import os -from typing import TYPE_CHECKING - -from loguru import logger - -from ...core.context import C -from ...core.op import BaseAsyncToolOp -from ...core.schema import ToolCall - -if TYPE_CHECKING: - from tavily import AsyncTavilyClient - - -@C.register_op() -class TavilySearchOp(BaseAsyncToolOp): - """A tool operation for performing web searches using Tavily API. - - This operation enables LLM models to search the internet for information by - providing search keywords. It supports optional content extraction from search - results with configurable character limits. - - Attributes: - enable_extract: Whether to extract raw content from search results (default: False). - item_max_char_count: Maximum character count per item when extracting (default: 20000). - all_max_char_count: Maximum total character count for all extracted items (default: 50000). - """ - - def __init__( - self, - enable_extract: bool = False, - item_max_char_count: int = 20000, - all_max_char_count: int = 50000, - **kwargs, - ): - super().__init__(**kwargs) - self.enable_extract: bool = enable_extract - self.item_max_char_count: int = item_max_char_count - self.all_max_char_count: int = all_max_char_count - - self._client: AsyncTavilyClient | None = None - - def build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "Use search keywords to retrieve relevant information from the internet.", - "input_schema": { - "query": { - "type": "string", - "description": "search keyword", - "required": True, - }, - }, - }, - ) - - @property - def client(self): - """Get or create the Tavily async client instance. - - Returns: - AsyncTavilyClient: The Tavily async client instance. - """ - if self._client is None: - from tavily import AsyncTavilyClient - - self._client = AsyncTavilyClient(api_key=os.environ["TAVILY_API_KEY"]) - return self._client - - async def async_execute(self): - query: str = self.input_dict["query"] - logger.info(f"tavily.query: {query}") - - if self.enable_cache: - cached_result = self.cache.load(query) - if cached_result: - self.set_output(json.dumps(cached_result, ensure_ascii=False, indent=2)) - return - - response = await self.client.search(query=query) - logger.info(f"tavily.response: {response}") - - if not self.enable_extract: - # 如果不需要 extract，直接返回 search 的结果 - if not response.get("results"): - raise RuntimeError("tavily return empty result") - - final_result = {item["url"]: item for item in response["results"]} - - if self.enable_cache and final_result: - self.cache.save(query, final_result, expire_hours=self.cache_expire_hours) - - self.set_output(json.dumps(final_result, ensure_ascii=False, indent=2)) - return - - # enable_extract=True 时的原有逻辑 - url_info_dict = {item["url"]: item for item in response["results"]} - response_extract = await self.client.extract(urls=[item["url"] for item in response["results"]]) - logger.info(f"tavily.response_extract: {response_extract}") - - final_result = {} - all_char_count = 0 - for item in response_extract["results"]: - url = item["url"] - raw_content: str = item["raw_content"] - if len(raw_content) > self.item_max_char_count: - raw_content = raw_content[: self.item_max_char_count] - if all_char_count + len(raw_content) > self.all_max_char_count: - raw_content = raw_content[: self.all_max_char_count - all_char_count] - - if raw_content: - final_result[url] = url_info_dict[url] - final_result[url]["raw_content"] = raw_content - all_char_count += len(raw_content) - - if not final_result: - raise RuntimeError("tavily return empty result") - - if self.enable_cache and final_result: - self.cache.save(query, final_result, expire_hours=self.cache_expire_hours) - - self.set_output(json.dumps(final_result, ensure_ascii=False, indent=2)) diff --git a/flowllm/gallery/stream_chat_op.py b/flowllm/gallery/stream_chat_op.py deleted file mode 100644 index 4847422..0000000 --- a/flowllm/gallery/stream_chat_op.py +++ /dev/null @@ -1,49 +0,0 @@ -"""Stream chat operation for flowllm. - -This module provides an async operation that streams chat responses from an LLM, -handling different chunk types including answers, thinking content, errors, and tool calls. -""" - -from loguru import logger - -from ..core.context import C -from ..core.enumeration import Role, ChunkEnum -from ..core.op import BaseAsyncOp -from ..core.schema import Message, FlowStreamChunk - - -@C.register_op() -class StreamChatOp(BaseAsyncOp): - """Async operation for streaming chat responses from an LLM. - - This operation processes messages and streams responses chunk by chunk, - handling different types of content including regular answers, reasoning - content, tool calls, and errors. It adds system prompts and processes - streaming chunks from the LLM. - """ - - file_path: str = __file__ - - def __init__( - self, - llm: str = "qwen3_30b_instruct", - **kwargs, - ): - super().__init__(llm=llm, **kwargs) - - async def async_execute(self): - messages = self.context.messages - assert isinstance(messages, list) and all( - isinstance(m, Message) for m in messages - ), "`messages` must be a list of Message objects!" - - system_prompt = self.context.system_prompt - assert system_prompt, "`system_prompt` is required!" - - messages = [Message(role=Role.SYSTEM, content=system_prompt)] + messages - logger.info(f"messages={messages}") - - async for stream_chunk in self.llm.astream_chat(messages): - assert isinstance(stream_chunk, FlowStreamChunk) - if stream_chunk.chunk_type in [ChunkEnum.ANSWER, ChunkEnum.THINK, ChunkEnum.ERROR, ChunkEnum.TOOL]: - await self.context.add_stream_chunk(stream_chunk) diff --git a/flowllm/gallery/think_tool_op.py b/flowllm/gallery/think_tool_op.py deleted file mode 100644 index b933e03..0000000 --- a/flowllm/gallery/think_tool_op.py +++ /dev/null @@ -1,37 +0,0 @@ -"""Lightweight tool op used to elicit intermediate thinking from the LLM.""" - -from ..core.context import C -from ..core.op import BaseAsyncToolOp -from ..core.schema import ToolCall - - -@C.register_op() -class ThinkToolOp(BaseAsyncToolOp): - """Utility operation that prompts the model for explicit reflection text.""" - - file_path = __file__ - - def __init__(self, add_output_reflection: bool = False, **kwargs): - super().__init__(**kwargs) - self.add_output_reflection: bool = add_output_reflection - - def build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "name": "think_tool", - "description": self.get_prompt("tool_desc"), - "input_schema": { - "reflection": { - "type": "string", - "description": self.get_prompt("reflection"), - "required": True, - }, - }, - }, - ) - - async def async_execute(self): - if self.add_output_reflection: - self.set_output(self.input_dict["reflection"]) - else: - self.set_output(self.get_prompt("reflection_output")) diff --git a/flowllm/gallery/think_tool_prompt.yaml b/flowllm/gallery/think_tool_prompt.yaml deleted file mode 100644 index 12a6c66..0000000 --- a/flowllm/gallery/think_tool_prompt.yaml +++ /dev/null @@ -1,32 +0,0 @@ -tool_desc: | - Before calling any external tool or when rethinking and planning is needed, you must invoke this tool for brief reflection. - The output must cover: - 1. Whether the current context is enough to answer the user directly, plus reasoning. - 2. If not, what information or validation is missing. - 3. A strategy to close the gap: which tool to call next, why, and key parameters or query terms. - Keep the reasoning tightly scoped to the current turn, avoid unrelated background, - and do not execute tools from here—only produce clear, actionable thoughts. - -reflection: | - 1) Can I answer now? Why? - 2) What is missing? - 3) Which tool + params next? - -tool_desc_zh: | - 每次准备调用任何外部工具之前或者需要重新思考规划，都必须先调用本工具进行简短思考。 - 输出需覆盖以下要点： - 1. 评估当前上下文是否足以直接回答用户问题，并解释理由。 - 2. 若不能回答，明确缺失的信息或验证步骤。 - 3. 针对缺口设计下一步策略：列出计划使用的工具、调用目的、关键参数或查询关键词。 - 思考要紧扣当前轮对话内容，避免复述无关背景，不要直接执行工具，只输出清晰推理。 - -reflection_zh: | - 1) 能直接回答吗？为什么？ - 2) 缺什么信息？ - 3) 下一步用哪个工具+参数？ - -reflection_output: | - Reflection has been recorded. - -reflection_output_zh: | - 已经记录反思 \ No newline at end of file diff --git a/flowllm/gallery/token_count_op.py b/flowllm/gallery/token_count_op.py deleted file mode 100644 index 3bcae25..0000000 --- a/flowllm/gallery/token_count_op.py +++ /dev/null @@ -1,38 +0,0 @@ -"""Token counting operation for calculating token usage in messages.""" - -from loguru import logger - -from ..core.context import C -from ..core.op import BaseAsyncOp -from ..core.schema import Message - - -@C.register_op() -class TokenCountOp(BaseAsyncOp): - """Operation for counting tokens in messages. - - This operation calculates the total number of tokens in the messages - from the context using the configured token counter. The token count - is logged and can be used for monitoring token usage in LLM operations. - - The operation requires messages to be present in the context. It uses - the token counting backend configured in the LLM service configuration - to accurately count tokens based on the model being used. - """ - - async def async_execute(self): - """Execute the token counting operation. - - Reads messages from context, converts them to Message objects, - calculates the token count using the configured token counter, - and logs the result. - - Requires: - context.messages: List of message dictionaries or Message objects - to count tokens for. - - The token count result is logged at INFO level. - """ - messages = [Message(**x) for x in self.context.messages] - cnt = self.token_count(messages) - logger.info(f"Token count result: {cnt}") diff --git a/flowllm/main.py b/flowllm/main.py deleted file mode 100644 index 5c827fe..0000000 --- a/flowllm/main.py +++ /dev/null @@ -1,91 +0,0 @@ -"""Main entry point for FlowLLM application.""" - -import sys - -from .config import ConfigParser -from .core.application import Application - - -class FlowLLMApp(Application): - """ - Convenience application class for FlowLLM framework. - - This class extends Application to provide a simplified interface with - default configurations. It automatically uses ConfigParser and loads - the default configuration file (default.yaml) if no custom config_path - is provided. - - This is the recommended entry point for most FlowLLM applications. - For advanced use cases requiring custom parsers or configurations, - use Application directly instead. - """ - - def __init__( - self, - *args, - llm_api_key: str = None, - llm_api_base: str = None, - embedding_api_key: str = None, - embedding_api_base: str = None, - config_path: str = None, - **kwargs, - ): - """ - Initialize FlowLLMApp with default configurations. - - Args: - *args: Additional arguments passed to ConfigParser. Examples: - - "llm.default.model_name=qwen3-30b-a3b-thinking-2507" - - "llm.default.backend=openai_compatible" - - "llm.default.params={'temperature': '0.6'}" - - "embedding_model.default.model_name=text-embedding-v4" - - "embedding_model.default.backend=openai_compatible" - - "embedding_model.default.params={'dimensions': 1024}" - - "vector_store.default.backend=memory" - - "vector_store.default.embedding_model=default" - - "vector_store.default.params={...}" - llm_api_key: API key for LLM service - llm_api_base: Base URL for LLM API - embedding_api_key: API key for embedding service - embedding_api_base: Base URL for embedding API - config_path: Path to custom configuration YAML file. If provided, loads - configuration from this file. If None, loads default.yaml automatically. - **kwargs: Additional keyword arguments passed to ConfigParser. - """ - super().__init__( - *args, - llm_api_key=llm_api_key, - llm_api_base=llm_api_base, - embedding_api_key=embedding_api_key, - embedding_api_base=embedding_api_base, - service_config=None, - parser=ConfigParser, - config_path=config_path, - load_default_config=True, - **kwargs, - ) - - -def main(): - """ - Main entry point for FlowLLM application. - - Initializes FlowLLMApp with command-line arguments and runs the service. - This function is typically called when running FlowLLM from the command line. - - Usage: - flowllm [config_path] [additional_config_overrides...] - - Example: - flowllm config=my_config.yaml llm.default.model_name=qwen3-max - - conda create -n fl_test python=3.10 - conda activate fl_test - conda env remove -n fl_test - """ - with FlowLLMApp(*sys.argv[1:]) as app: - app.run_service() - - -if __name__ == "__main__": - main() diff --git a/flowllm/schema/__init__.py b/flowllm/schema/__init__.py new file mode 100644 index 0000000..4425d8a --- /dev/null +++ b/flowllm/schema/__init__.py @@ -0,0 +1,17 @@ +"""Schema""" + +from .application_config import ApplicationConfig, ComponentConfig, JobConfig +from .emb_node import EmbNode +from .request import Request +from .response import Response +from .stream_chunk import StreamChunk + +__all__ = [ + "ApplicationConfig", + "ComponentConfig", + "EmbNode", + "JobConfig", + "Request", + "Response", + "StreamChunk", +] diff --git a/flowllm/schema/application_config.py b/flowllm/schema/application_config.py new file mode 100644 index 0000000..8350687 --- /dev/null +++ b/flowllm/schema/application_config.py @@ -0,0 +1,45 @@ +"""Application configuration schema.""" + +import os + +from pydantic import BaseModel, ConfigDict, Field + +from ..enumeration import ComponentEnum + + +class ComponentConfig(BaseModel): + """Base component config with extra fields for backend-specific options.""" + + model_config = ConfigDict(extra="allow") + + backend: str = Field(default="", description="Backend implementation class name") + + +class JobConfig(ComponentConfig): + """Config for a job: an ordered sequence of steps.""" + + description: str = Field(default="", description="Human-readable description") + parameters: dict = Field(default_factory=dict, description="Job-level parameters") + steps: list[ComponentConfig] = Field(default_factory=list, description="Ordered step configs") + enable_serve: bool = Field(default=True, description="Whether to expose this job through the service layer") + + +class ApplicationConfig(BaseModel): + """Root application config.""" + + app_name: str = Field(default=os.getenv("APP_NAME", "FlowLLM"), description="Application display name") + workspace_dir: str = Field(default=".flowllm", description="Workspace root directory for runtime files") + metadata_dir: str = Field(default="metadata", description="Subdirectory for FlowLLM persistent state") + session_dir: str = Field(default="session", description="Subdirectory for persisted agent sessions") + enable_logo: bool = Field(default=True, description="Show ASCII logo on startup") + timezone: str | None = Field(default="Asia/Shanghai", description="IANA timezone; None uses local time") + language: str = Field(default="", description="Default language for LLM interactions") + log_to_console: bool = Field(default=True, description="Log to console") + log_to_file: bool = Field(default=True, description="Log to file") + mcp_servers: dict[str, dict] = Field(default_factory=dict, description="MCP server configs by name") + service: ComponentConfig = Field(default_factory=ComponentConfig, description="Service endpoint config") + jobs: dict[str, JobConfig] = Field(default_factory=dict, description="Job definitions keyed by job name") + components: dict[ComponentEnum, dict[str, ComponentConfig]] = Field( + default_factory=dict, + description="Component registry keyed by type then name", + ) diff --git a/flowllm/schema/emb_node.py b/flowllm/schema/emb_node.py new file mode 100644 index 0000000..7e31b38 --- /dev/null +++ b/flowllm/schema/emb_node.py @@ -0,0 +1,32 @@ +"""Embedding node: text record with optional vector.""" + +from uuid import uuid4 + +import numpy as np +from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator + + +class EmbNode(BaseModel): + """Text record with optional embedding vector.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + id: str = Field(default_factory=lambda: uuid4().hex, description="Unique node id") + text: str = Field(default="", description="Text content") + embedding: np.ndarray | None = Field(default=None, description="Embedding vector (float16)") + metadata: dict = Field(default_factory=dict, description="Arbitrary metadata") + + @field_validator("embedding", mode="before") + @classmethod + def validate_embedding(cls, v): + """Coerce list/tuple to float16 ndarray.""" + if v is None: + return v + return np.array(v, dtype=np.float16) + + @field_serializer("embedding") + def serialize_embedding(self, v: np.ndarray | None, _info): + """Serialize ndarray to list for JSON.""" + if v is None: + return None + return v.tolist() diff --git a/flowllm/schema/request.py b/flowllm/schema/request.py new file mode 100644 index 0000000..38e7338 --- /dev/null +++ b/flowllm/schema/request.py @@ -0,0 +1,11 @@ +"""Service request schema.""" + +from pydantic import BaseModel, ConfigDict, Field + + +class Request(BaseModel): + """Incoming service request with extra fields for endpoint payloads.""" + + model_config = ConfigDict(extra="allow") + + metadata: dict | None = Field(default=None, description="Request metadata for context") diff --git a/flowllm/schema/response.py b/flowllm/schema/response.py new file mode 100644 index 0000000..3f1a596 --- /dev/null +++ b/flowllm/schema/response.py @@ -0,0 +1,15 @@ +"""Service response schema.""" + +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field + + +class Response(BaseModel): + """Standard response envelope with extra fields for endpoint output.""" + + model_config = ConfigDict(extra="allow") + + answer: str | Any = Field(default="", description="Response content or result data") + success: bool = Field(default=True, description="Whether the operation succeeded") + metadata: dict = Field(default_factory=dict, description="Additional context and diagnostics") diff --git a/flowllm/schema/stream_chunk.py b/flowllm/schema/stream_chunk.py new file mode 100644 index 0000000..b4fa8ff --- /dev/null +++ b/flowllm/schema/stream_chunk.py @@ -0,0 +1,25 @@ +"""Stream chunk schema for incremental LLM responses.""" + +from typing import Any + +from pydantic import BaseModel, Field + +from ..enumeration import ChunkEnum + + +class StreamChunk(BaseModel): + """A single chunk in a streaming response sequence.""" + + chunk_type: ChunkEnum = Field(default=ChunkEnum.CONTENT, description="Type of chunk content") + chunk: str | dict | list = Field(default="", description="Chunk payload") + done: bool = Field(default=False, description="Whether this is the final chunk") + + session_id: str | None = Field(default=None, description="Session identifier") + block_id: str | None = Field(default=None, description="Content block identifier") + tool_call_id: str | None = Field(default=None, description="Tool call identifier") + tool_call_name: str | None = Field(default=None, description="Tool call name") + media_type: str | None = Field(default=None, description="MIME type for data blocks") + input_tokens: int | None = Field(default=None, description="Prompt tokens consumed") + output_tokens: int | None = Field(default=None, description="Completion tokens generated") + + metadata: dict[str, Any] = Field(default_factory=dict, description="Chunk metadata") diff --git a/flowllm/steps/__init__.py b/flowllm/steps/__init__.py new file mode 100644 index 0000000..8eed3bd --- /dev/null +++ b/flowllm/steps/__init__.py @@ -0,0 +1,9 @@ +"""steps""" + +from . import common +from .base_step import BaseStep + +__all__ = [ + "BaseStep", + "common", +] diff --git a/flowllm/steps/base_step.py b/flowllm/steps/base_step.py new file mode 100644 index 0000000..9a09b13 --- /dev/null +++ b/flowllm/steps/base_step.py @@ -0,0 +1,189 @@ +"""Base step class for LLM workflow execution.""" + +import copy +from abc import abstractmethod, ABC +from typing import Any, TYPE_CHECKING + +from agentscope.model import ChatModelBase + +from ..components.agent_wrapper.base_agent_wrapper import BaseAgentWrapper +from ..components.base_component import ComponentMixin +from ..components.component_registry import R +from ..components.prompt_handler import PromptHandler +from ..components.runtime_context import RuntimeContext +from ..enumeration import ComponentEnum +from ..schema import ApplicationConfig, Response + +if TYPE_CHECKING: + from ..components import ApplicationContext + from ..components.job import BaseJob + +_UNSET = object() +_DispatchStep = str | dict[str, Any] + + +class Ref: + """Lazy-resolving descriptor for component dependencies (kwargs → context → registry, cached per instance).""" + + __slots__ = ("base_cls", "comp_enum", "attr", "optional", "key", "_cache_attr") + + def __init__(self, base_cls: type, comp_enum: ComponentEnum, attr: str | None = None, *, optional: bool = False): + self.base_cls = base_cls + self.comp_enum = comp_enum + self.attr = attr + self.optional = optional + self.key: str = "" + self._cache_attr: str = "" + + def __set_name__(self, owner: type, name: str) -> None: + self.key = name + self._cache_attr = f"_ref_{name}" + + def __get__(self, obj: "BaseStep | None", objtype: type | None = None): + if obj is None: + return self + cached = obj.__dict__.get(self._cache_attr, _UNSET) + if cached is not _UNSET: + return cached + value = self._resolve(obj) + obj.__dict__[self._cache_attr] = value + return value + + def __set__(self, obj: "BaseStep", value) -> None: + obj.__dict__[self._cache_attr] = value + + def __delete__(self, obj: "BaseStep") -> None: + obj.__dict__.pop(self._cache_attr, None) + + def _resolve(self, obj: "BaseStep"): + for source in (obj.kwargs, obj.context or {}): + value = source.get(self.key) + if isinstance(value, self.base_cls): + return value + + name = obj.kwargs.get(self.key, "default") + if obj.app_context is None: + if self.optional: + return None + raise RuntimeError(f"app_context is not set when resolving '{self.key}'") + comp = obj.app_context.components[self.comp_enum].get(name) + if comp is None: + if self.optional: + return None + raise KeyError(f"Component '{name}' not found in {self.comp_enum.value}") + return getattr(comp, self.attr) if self.attr else comp + + +class BaseStep(ComponentMixin, ABC): + """Composable unit of an LLM workflow.""" + + component_type = ComponentEnum.STEP + + as_llm: ChatModelBase = Ref(ChatModelBase, ComponentEnum.AS_LLM, "model") + agent_wrapper: BaseAgentWrapper = Ref(BaseAgentWrapper, ComponentEnum.AGENT_WRAPPER, optional=True) + + def __new__(cls, *args, **kwargs): + # Snapshot init args so copy() can rebuild an equivalent instance later. + instance = object.__new__(cls) + instance._init_args, instance._init_kwargs = copy.copy(args), copy.copy(kwargs) + return instance + + def __init__( + self, + name: str | None = None, + backend: str = "", + app_context: "ApplicationContext | None" = None, + language: str = "", + prompt_dict: dict[str, str] | None = None, + input_mapping: dict[str, str] | None = None, + output_mapping: dict[str, str] | None = None, + dispatch_steps: list[_DispatchStep] | None = None, + **kwargs, + ): + super().__init__(name=name, backend=backend, app_context=app_context, **kwargs) + self.language = language or (self.app_context.app_config.language if self.app_context is not None else "") + self.input_mapping = input_mapping + self.output_mapping = output_mapping + self.dispatch_step_specs = list(dispatch_steps or []) + self.context: RuntimeContext | None = None + # MRO reverse walk so derived-class prompts win, then overlay caller overrides. + self.prompt = PromptHandler(language=self.language) + for cls in reversed(self.__class__.__mro__): + self.prompt.load_prompt_by_class(cls) + self.prompt.load_prompt_dict(prompt_dict) + + @abstractmethod + async def execute(self): + """Run the step's logic against ``self.context``.""" + + async def __call__(self, context: RuntimeContext | None = None, **kwargs): + # Clear cached Ref values so context-supplied overrides take effect. + for key in [k for k in self.__dict__ if k.startswith("_ref_")]: + del self.__dict__[key] + self.context = RuntimeContext.from_context(context, **kwargs) + assert self.context is not None + if self.input_mapping: + self.context.apply_mapping(self.input_mapping) + result = await self.execute() + if self.output_mapping: + self.context.apply_mapping(self.output_mapping) + return result + + def prompt_format(self, prompt_name: str, **kwargs) -> str: + """Format a named prompt template with the given kwargs.""" + return self.prompt.prompt_format(prompt_name=prompt_name, **kwargs) + + def get_prompt(self, prompt_name: str) -> str: + """Return a named prompt template as-is.""" + return self.prompt.get_prompt(prompt_name=prompt_name) + + def config_value(self, key: str): + """Return an app config value, falling back to ApplicationConfig defaults.""" + defaults = ApplicationConfig() + cfg = self.app_context.app_config if self.app_context is not None else defaults + value = getattr(cfg, key) + return getattr(defaults, key) if value in (None, "") else value + + def copy(self, **kwargs) -> "BaseStep": + """Construct a new instance from the original init args, applying overrides.""" + return self.__class__(*self._init_args, **{**self._init_kwargs, **kwargs}) + + def get_job(self, name: str, /) -> "BaseJob | None": + """Return a job by name.""" + if self.app_context is None: + raise RuntimeError("Cannot get job without an app context") + return self.app_context.jobs.get(name) + + async def run_job(self, name: str, /, **kwargs) -> Response: + """Execute a job by name and kwargs, return the final response.""" + job: "BaseJob | None" = self.get_job(name) + if job is None: + raise RuntimeError(f"Job {name} not found") + return await job(**kwargs) + + def _resolve_dispatch_step(self, raw: _DispatchStep): + """Resolve a dispatch step spec to (step class, init params).""" + if isinstance(raw, str): + params: dict[str, Any] = {"backend": raw} + elif isinstance(raw, dict): + params = dict(raw) + else: + raise TypeError(f"Invalid dispatch step spec: {raw!r}") + backend = params.get("backend", "") + if not backend: + raise ValueError("Dispatch step is missing the required 'backend' field") + step_cls = R.get(ComponentEnum.STEP, backend) + if step_cls is None: + raise RuntimeError(f"Unregistered step '{backend}'") + params["app_context"] = self.app_context + return step_cls, params + + async def dispatch_steps(self, dispatch_steps: list[_DispatchStep], **kwargs) -> list[Response]: + """Run dispatch steps against the current context, preserving existing context data.""" + if self.context is None: + raise RuntimeError("Cannot dispatch steps without a runtime context") + responses: list[Response] = [] + for raw in dispatch_steps: + step_cls, params = self._resolve_dispatch_step(raw) + responses.append(await step_cls(**params)(self.context, **kwargs)) + return responses diff --git a/flowllm/steps/common/__init__.py b/flowllm/steps/common/__init__.py new file mode 100644 index 0000000..5d93893 --- /dev/null +++ b/flowllm/steps/common/__init__.py @@ -0,0 +1,23 @@ +"""common steps""" + +from .add import AddStep +from .demo import DemoEchoStep1, DemoEchoStep2 +from .health_check import HealthCheckStep +from .help import HelpStep +from .llm_demo import LLMDemoStep +from .stream_demo import StreamDemoStep1, StreamDemoStep2 +from .stream_llm_demo import StreamLLMDemoStep +from .version import VersionStep + +__all__ = [ + "AddStep", + "DemoEchoStep1", + "DemoEchoStep2", + "HealthCheckStep", + "HelpStep", + "LLMDemoStep", + "StreamDemoStep1", + "StreamDemoStep2", + "StreamLLMDemoStep", + "VersionStep", +] diff --git a/flowllm/steps/common/add.py b/flowllm/steps/common/add.py new file mode 100644 index 0000000..7d1558d --- /dev/null +++ b/flowllm/steps/common/add.py @@ -0,0 +1,24 @@ +"""Simple step that adds two numbers — used as a demo tool for agent wrapper.""" + +from ..base_step import BaseStep +from ...components import R + + +@R.register("add_step") +class AddStep(BaseStep): + """Add two numbers (a, b from context) and return the sum as response answer.""" + + async def execute(self): + assert self.context is not None + try: + a, b = float(self.context.get("a", 0.0)), float(self.context.get("b", 0.0)) + except (TypeError, ValueError) as exc: + self.context.response.success = False + self.context.response.answer = f"Invalid add arguments: {exc}" + return self.context.response + result = a + b + self.logger.info(f"[{self.name}] add({a}, {b}) = {result}") + self.context.response.success = True + self.context.response.answer = str(result) + self.context.response.metadata.update({"a": a, "b": b, "result": result}) + return self.context.response diff --git a/flowllm/steps/common/demo.py b/flowllm/steps/common/demo.py new file mode 100644 index 0000000..1da1c7b --- /dev/null +++ b/flowllm/steps/common/demo.py @@ -0,0 +1,47 @@ +"""Demo steps for integration-testing the application stack.""" + +from ..base_step import BaseStep +from ...components import R + + +@R.register("demo_echo_step1") +class DemoEchoStep1(BaseStep): + """Read query/min_score from context, normalize, and write back for Step2.""" + + async def execute(self): + assert self.context is not None + query = self.context.get("query", "") + min_score = self.context.get("min_score", 0.5) + self.logger.info(f"[{self.name}] query={query!r}, min_score={min_score}") + processed_query = query.strip().lower() + adjusted_min_score = float(min_score) * 0.9 + self.context["processed_query"] = processed_query + self.context["adjusted_min_score"] = adjusted_min_score + return self.context.response + + +@R.register("demo_echo_step2") +class DemoEchoStep2(BaseStep): + """Consume Step1's outputs from context and emit the final response.""" + + async def execute(self): + assert self.context is not None + query = self.context.get("query", "") + min_score = self.context.get("min_score", 0.5) + processed_query = self.context.get("processed_query", "") + adjusted_min_score = self.context.get("adjusted_min_score", min_score) + self.logger.info( + f"[{self.name}] query={query!r}, min_score={min_score}, " + f"processed_query={processed_query!r}, adjusted_min_score={adjusted_min_score}", + ) + self.context.response.answer = f"echo: {processed_query} (min_score={adjusted_min_score})" + self.context.response.metadata.update( + { + "step": self.name, + "query": query, + "min_score": min_score, + "processed_query": processed_query, + "adjusted_min_score": adjusted_min_score, + }, + ) + return self.context.response diff --git a/flowllm/steps/common/health_check.py b/flowllm/steps/common/health_check.py new file mode 100644 index 0000000..14c7bf5 --- /dev/null +++ b/flowllm/steps/common/health_check.py @@ -0,0 +1,179 @@ +"""Concise health snapshot of FlowLLM runtime components.""" + +import sys +from collections.abc import Mapping + +import numpy as np + +from ..base_step import BaseStep +from ... import __version__ +from ...components import R +from ...enumeration import ComponentEnum + + +def _deep_size(obj, _seen: set | None = None) -> int: + """Recursive sizeof. Uses ndarray.nbytes; walks Mappings, sequences, __dict__.""" + if _seen is None: + _seen = set() + if id(obj) in _seen: + return 0 + _seen.add(id(obj)) + if isinstance(obj, np.ndarray): + return int(obj.nbytes) + sys.getsizeof(obj) + size = sys.getsizeof(obj) + if isinstance(obj, (str, bytes, bytearray, int, float, bool, type(None))): + extra = 0 + elif isinstance(obj, Mapping): + extra = sum(_deep_size(k, _seen) + _deep_size(v, _seen) for k, v in obj.items()) + elif isinstance(obj, (list, tuple, set, frozenset)): + extra = sum(_deep_size(item, _seen) for item in obj) + elif hasattr(obj, "__dict__"): + extra = _deep_size(vars(obj), _seen) + elif hasattr(obj, "__slots__"): + extra = sum(_deep_size(getattr(obj, s), _seen) for s in obj.__slots__ if hasattr(obj, s)) + else: + extra = 0 + return size + extra + + +def _mb_str(*objs) -> str: + """Sum deep size of objs and format as 'X.XX MB'.""" + seen: set = set() + total = sum(_deep_size(o, seen) for o in objs) + return f"{total / (1024 * 1024):.2f} MB" + + +def _embedding_status(comp) -> dict: + cache = getattr(comp, "_embedding_cache", {}) or {} + model = getattr(comp, "model", None) + try: + dims = comp.dimensions + except Exception: + dims = None + return { + "is_started": comp.is_started, + "is_healthy": getattr(comp, "is_healthy", None), + "model_name": getattr(model, "model", None), + "dimensions": dims, + "cache_size": len(cache), + "memory": _mb_str(cache), + } + + +def _file_graph_nx_status(comp, graph) -> dict: + """Networkx backend: virtuals are nodes without a 'node' payload.""" + n_real = sum(1 for _, d in graph.nodes(data=True) if "node" in d) + return { + "is_started": comp.is_started, + "n_nodes": n_real, + "n_edges": graph.number_of_edges(), + "n_virtual": graph.number_of_nodes() - n_real, + "memory": _mb_str(graph), + } + + +def _file_graph_local_status(comp) -> dict: + """Local backend: nodes / inverse edges / pending edges as separate dicts.""" + nodes = getattr(comp, "_nodes", {}) or {} + inverse = getattr(comp, "_inverse", {}) or {} + pending = getattr(comp, "_pending", {}) or {} + return { + "is_started": comp.is_started, + "n_nodes": len(nodes), + "n_edges": sum(len(s) for s in inverse.values()), + "n_pending": sum(len(s) for s in pending.values()), + "memory": _mb_str(nodes, inverse, pending), + } + + +def _file_graph_neo4j_status(comp) -> dict: + """Neo4j backend: counts cached on async component for sync health checks.""" + attrs = {k: getattr(comp, f"_{k}", 0) for k in ("n_nodes", "n_edges", "n_virtual")} + mem = _mb_str(getattr(comp, "_uri", ""), getattr(comp, "_database", ""), *attrs.values()) + return {"is_started": comp.is_started, **attrs, "memory": mem} + + +def _file_graph_status(comp) -> dict: + graph = getattr(comp, "_graph", None) + if graph is not None: + return _file_graph_nx_status(comp, graph) + if hasattr(comp, "_driver"): + return _file_graph_neo4j_status(comp) + return _file_graph_local_status(comp) + + +def _file_store_status(comp) -> dict: + chunks = getattr(comp, "file_chunks", {}) or {} + n_emb = sum(1 for c in chunks.values() if getattr(c, "embedding", None) is not None) + return { + "is_started": comp.is_started, + "n_chunks": len(chunks), + "n_chunks_with_embedding": n_emb, + "memory": _mb_str(chunks), + } + + +def _keyword_index_status(comp) -> dict: + vocab = getattr(comp, "vocab", {}) or {} + mem = _mb_str( + vocab, + getattr(comp, "inverted_index", {}) or {}, + getattr(comp, "doc_meta", {}) or {}, + getattr(comp, "_idf_cache", {}) or {}, + ) + return { + "is_started": comp.is_started, + "n_docs": getattr(comp, "n_docs", None), + "vocab_size": len(vocab), + "memory": mem, + } + + +_HANDLERS = { + ComponentEnum.EMBEDDING_STORE: _embedding_status, + ComponentEnum.FILE_GRAPH: _file_graph_status, + ComponentEnum.FILE_STORE: _file_store_status, + ComponentEnum.KEYWORD_INDEX: _keyword_index_status, +} + + +def _is_healthy(ctype: ComponentEnum, status: dict) -> bool: + """Unstarted = unhealthy; embedding model also requires is_healthy != False.""" + if not status.get("is_started"): + return False + if ctype is ComponentEnum.EMBEDDING_STORE and status.get("is_healthy") is False: + return False + return True + + +def _collect_components(app_context) -> tuple[dict, bool]: + """Walk every registered component type and produce {type: {name: status}}, plus overall flag.""" + components: dict = {} + healthy = True + for ctype, handler in _HANDLERS.items(): + bucket = {} + for name, comp in app_context.components.get(ctype, {}).items(): + status = handler(comp) + bucket[name] = status + if not _is_healthy(ctype, status): + healthy = False + components[ctype.value] = bucket + return components, healthy + + +@R.register("health_check_step") +class HealthCheckStep(BaseStep): + """Collect a concise health snapshot of the relevant components.""" + + async def execute(self): + assert self.context is not None + if self.app_context is not None: + components, healthy = _collect_components(self.app_context) + else: + components, healthy = {}, True + health = {"version": __version__, "healthy": healthy, "components": components} + self.logger.info(f"[{self.name}] health collected: {health}") + emoji, label = ("✅", "healthy") if healthy else ("❌", "unhealthy") + self.context.response.answer = f"{emoji} FlowLLM v{__version__} - {label}" + self.context.response.metadata["health"] = health + return self.context.response diff --git a/flowllm/steps/common/help.py b/flowllm/steps/common/help.py new file mode 100644 index 0000000..d220456 --- /dev/null +++ b/flowllm/steps/common/help.py @@ -0,0 +1,35 @@ +"""Return a one-line summary of every registered job for LLM consumption.""" + +from ..base_step import BaseStep +from ...components import R + + +def _format_params(parameters: dict) -> str: + props = (parameters or {}).get("properties") or {} + if not props: + return "no args" + required = set((parameters or {}).get("required") or []) + parts = [] + for name, schema in props.items(): + ptype = schema.get("type", "any") + suffix = "*" if name in required else f"={schema['default']}" if "default" in schema else "" + parts.append(f"{name}:{ptype}{suffix}") + return ", ".join(parts) + + +@R.register("help_step") +class HelpStep(BaseStep): + """List all registered jobs (excluding self and non-servable) as compact one-liners for an LLM.""" + + async def execute(self): + assert self.context is not None + lines = [] + if self.app_context is not None: + for name, job in self.app_context.jobs.items(): + if name == "help" or not getattr(job, "enable_serve", True): + continue + lines.append(f"🛠️ `{name}` — {job.description} 📥 {_format_params(job.parameters)}") + self.logger.info(f"[{self.name}] returning {len(lines)} jobs") + self.context.response.answer = "\n".join(lines) + self.context.response.metadata["job_count"] = len(lines) + return self.context.response diff --git a/flowllm/steps/common/llm_demo.py b/flowllm/steps/common/llm_demo.py new file mode 100644 index 0000000..283edc0 --- /dev/null +++ b/flowllm/steps/common/llm_demo.py @@ -0,0 +1,36 @@ +"""Demo step that drives an Agent via the agent_wrapper component.""" + +from pydantic import BaseModel + +from ..base_step import BaseStep +from ...components import R + + +@R.register("llm_demo_step") +class LLMDemoStep(BaseStep): + """Drive an Agent via agent_wrapper; reads query/sys_prompt from context, writes reply to response.""" + + DEFAULT_SYS_PROMPT = "You are a helpful assistant. Provide clear and detailed responses." + + async def execute(self): + assert self.context is not None + query: str = self.context.get("query", "") + sys_prompt: str = self.context.get("sys_prompt") or self.DEFAULT_SYS_PROMPT + structured_model: type[BaseModel] | None = self.context.get("structured_model") + if not query: + self.context.response.success = False + self.context.response.answer = "Skipped: empty query" + return self.context.response + wrapper_kwargs = {"system_prompt": sys_prompt, "job_tools": ["add"]} + if structured_model is not None: + wrapper_kwargs["output_schema"] = structured_model + result = await self.agent_wrapper.reply(query, **wrapper_kwargs) + structured_content = result.get("structured_output") + text = (result.get("result") or "").strip() + self.logger.info(f"[{self.name}] response: {text!r}") + self.context.response.success = True + self.context.response.answer = text + self.context.response.metadata.update( + {"query": query, "sys_prompt": sys_prompt, "response": text, "structured_output": structured_content}, + ) + return self.context.response diff --git a/flowllm/steps/common/stream_demo.py b/flowllm/steps/common/stream_demo.py new file mode 100644 index 0000000..bd0703a --- /dev/null +++ b/flowllm/steps/common/stream_demo.py @@ -0,0 +1,37 @@ +"""Streaming demo steps: step1 prepares text, step2 streams it char-by-char.""" + +import asyncio + +from ..base_step import BaseStep +from ...components import R +from ...enumeration import ChunkEnum + + +@R.register("stream_demo_step1") +class StreamDemoStep1(BaseStep): + """Read query from context, repeat it 10x, write back for Step2 to stream.""" + + async def execute(self): + assert self.context is not None + query = self.context.get("query", "") + repeat = int(self.context.get("repeat", 10)) + stream_text = (query * repeat) if query else "" + self.logger.info(f"[{self.name}] query={query!r}, repeat={repeat}, len={len(stream_text)}") + self.context["stream_text"] = stream_text + return self.context.response + + +@R.register("stream_demo_step2") +class StreamDemoStep2(BaseStep): + """Stream stream_text char-by-char as CONTENT chunks with 0.1s pacing.""" + + async def execute(self): + assert self.context is not None + stream_text: str = self.context.get("stream_text", "") + interval = float(self.context.get("interval", 0.1)) + self.logger.info(f"[{self.name}] streaming {len(stream_text)} chars, interval={interval}s") + for ch in stream_text: + await self.context.add_stream_string(ch, ChunkEnum.CONTENT) + await asyncio.sleep(interval) + + return self.context.response diff --git a/flowllm/steps/common/stream_llm_demo.py b/flowllm/steps/common/stream_llm_demo.py new file mode 100644 index 0000000..81c35c1 --- /dev/null +++ b/flowllm/steps/common/stream_llm_demo.py @@ -0,0 +1,48 @@ +"""Demo step that drives an Agent via the agent_wrapper component with streaming output.""" + +from ..base_step import BaseStep +from ...components import R +from ...enumeration import ChunkEnum + + +@R.register("stream_llm_demo_step") +class StreamLLMDemoStep(BaseStep): + """Drive an Agent via agent_wrapper with streaming support; falls back to non-streaming when disabled.""" + + DEFAULT_SYS_PROMPT = "You are a helpful assistant. Provide clear and detailed responses." + + async def execute(self): + assert self.context is not None + query: str = self.context.get("query", "") + sys_prompt: str = self.context.get("sys_prompt") or self.DEFAULT_SYS_PROMPT + if not query: + self.context.response.success = False + self.context.response.answer = "Skipped: empty query" + return self.context.response + wrapper_kwargs = {"system_prompt": sys_prompt, "job_tools": ["add"]} + if self.context.stream: + text = await self._stream_reply(query, **wrapper_kwargs) + else: + result = await self.agent_wrapper.reply(query, **wrapper_kwargs) + text = (result.get("result") or "").strip() + self.logger.debug(f"[{self.name}] response: {text!r}") + self.context.response.success = True + self.context.response.answer = text + self.context.response.metadata.update({"query": query, "sys_prompt": sys_prompt, "response": text}) + return self.context.response + + async def _stream_reply(self, query: str, **wrapper_kwargs) -> str: + """Stream unified chunks to the context stream queue.""" + assert self.context is not None + text_parts: list[str] = [] + + async for chunk in self.agent_wrapper.reply_stream(query, **wrapper_kwargs): + await self.context.add_stream_string(chunk.chunk, chunk.chunk_type) + + if chunk.chunk_type == ChunkEnum.CONTENT and isinstance(chunk.chunk, str): + text_parts.append(chunk.chunk) + + if chunk.session_id: + self.context.response.metadata["session_id"] = chunk.session_id + + return "".join(text_parts).strip() diff --git a/flowllm/steps/common/version.py b/flowllm/steps/common/version.py new file mode 100644 index 0000000..23554ed --- /dev/null +++ b/flowllm/steps/common/version.py @@ -0,0 +1,18 @@ +"""Return the package version.""" + +from ..base_step import BaseStep +from ...components import R + + +@R.register("version_step") +class VersionStep(BaseStep): + """Emit flowllm.__version__ as the response answer.""" + + async def execute(self): + assert self.context is not None + from ... import __version__ + + self.logger.info(f"[{self.name}] version={__version__}") + self.context.response.answer = __version__ + self.context.response.metadata["version"] = __version__ + return self.context.response diff --git a/flowllm/utils/__init__.py b/flowllm/utils/__init__.py new file mode 100644 index 0000000..d82c7f0 --- /dev/null +++ b/flowllm/utils/__init__.py @@ -0,0 +1,18 @@ +"""Utility modules.""" + +from .agent_state_io import AsStateHandler +from .common_utils import execute_stream_task +from .env_utils import load_env +from .logger_utils import get_logger +from .logo_utils import print_logo +from .service_utils import precheck_start, cli_find_flowllm + +__all__ = [ + "AsStateHandler", + "execute_stream_task", + "load_env", + "get_logger", + "print_logo", + "precheck_start", + "cli_find_flowllm", +] diff --git a/flowllm/utils/agent_state_io.py b/flowllm/utils/agent_state_io.py new file mode 100644 index 0000000..e38c5ad --- /dev/null +++ b/flowllm/utils/agent_state_io.py @@ -0,0 +1,79 @@ +"""AgentState JSONL dump / load.""" + +import os +from pathlib import Path +from uuid import uuid4 + +import aiofiles +from agentscope.message import Msg, UserMsg +from agentscope.state import AgentState + +_META_KEYS = ("session_id", "reply_id", "cur_iter") + + +class AsStateHandler: + """Serialize / deserialize AgentState to a JSONL file.""" + + def __init__(self, path: str | Path): + self.path = Path(path) + + @classmethod + def for_session(cls, directory: str | Path, session_id: str) -> "AsStateHandler": + """Create a handler for a session file.""" + if not session_id or Path(session_id).name != session_id: + raise ValueError(f"Invalid session_id: {session_id!r}") + return cls(Path(directory) / f"{session_id}.jsonl") + + def exists(self) -> bool: + """Check if state file exists.""" + return self.path.is_file() + + async def load_or_none(self) -> AgentState | None: + """Load state or return None.""" + if not self.exists(): + return None + return await self.load() + + async def delete(self) -> bool: + """Delete the state file if present.""" + if not self.exists(): + return False + self.path.unlink() + return True + + async def dump(self, state: AgentState) -> Path: + """Write state as JSONL.""" + self.path.parent.mkdir(parents=True, exist_ok=True) + header = UserMsg( + name="__state__", + content=state.summary or "", + metadata={k: getattr(state, k) for k in _META_KEYS}, + ) + tmp_path = self.path.with_name(f".{self.path.name}.{uuid4().hex}.tmp") + async with aiofiles.open(tmp_path, "w", encoding="utf-8") as f: + await f.write(header.model_dump_json() + "\n") + for msg in state.context: + await f.write(msg.model_dump_json() + "\n") + os.replace(tmp_path, self.path) + return self.path + + async def load(self) -> AgentState: + """Read state from JSONL.""" + async with aiofiles.open(self.path, encoding="utf-8") as f: + lines = (await f.read()).splitlines() + if not lines: + return AgentState() + + header = Msg.model_validate_json(lines[0]) + summary: str | list = ( + list(header.content) + if any(getattr(b, "type", None) == "data" for b in header.content) + else header.get_text_content() or "" + ) + + metadata = header.metadata or {} + return AgentState( + **{k: metadata.get(k, d) for k, d in [("session_id", ""), ("reply_id", ""), ("cur_iter", 0)]}, + summary=summary, + context=[Msg.model_validate_json(line) for line in lines[1:] if line.strip()], + ) diff --git a/flowllm/utils/common_utils.py b/flowllm/utils/common_utils.py new file mode 100644 index 0000000..b071eb4 --- /dev/null +++ b/flowllm/utils/common_utils.py @@ -0,0 +1,94 @@ +"""Async stream task execution.""" + +import asyncio +from collections.abc import AsyncGenerator +from typing import Any, Literal + +from .logger_utils import get_logger +from ..enumeration import ChunkEnum +from ..schema import StreamChunk + + +def _format_chunk( + chunk: StreamChunk, + output_format: Literal["str", "bytes", "chunk"], +) -> str | bytes | StreamChunk: + """Render a StreamChunk for transport.""" + if output_format == "chunk": + return chunk + data = "data:[DONE]\n\n" if chunk.done else f"data:{chunk.model_dump_json()}\n\n" + return data.encode() if output_format == "bytes" else data + + +async def execute_stream_task( + stream_queue: asyncio.Queue[StreamChunk], + task: asyncio.Task[Any], + task_name: str | None = None, + output_format: Literal["str", "bytes", "chunk"] = "str", +) -> AsyncGenerator[str | bytes | StreamChunk, None]: + """Yield chunks from stream_queue while monitoring task.""" + logger = get_logger() + consumer: asyncio.Task[StreamChunk] | None = None + try: + while True: + consumer = get_chunk = asyncio.create_task(stream_queue.get()) + done, _pending = await asyncio.wait({get_chunk, task}, return_when=asyncio.FIRST_COMPLETED) + + # Relay next chunk. + if task not in done: + chunk = get_chunk.result() + yield _format_chunk(chunk, output_format) + if chunk.done: + return + continue + + # Producer finished — capture pending chunk. + pending_chunk: StreamChunk | None = None + if get_chunk in done: + pending_chunk = get_chunk.result() + else: + get_chunk.cancel() + try: + await get_chunk + except asyncio.CancelledError: + pass + + # Surface task failure. + if task.cancelled(): + msg = f"Task cancelled: {task_name}" if task_name else "Task cancelled" + raise asyncio.CancelledError(msg) + exc = task.exception() + if exc is not None: + log_msg = f"Task error in {task_name}: {exc}" if task_name else f"Task error: {exc}" + logger.error(log_msg, exc_info=exc) + raise exc + + # Flush remaining chunks and emit done sentinel. + if pending_chunk is not None: + yield _format_chunk(pending_chunk, output_format) + if pending_chunk.done: + return + while not stream_queue.empty(): + chunk = stream_queue.get_nowait() + yield _format_chunk(chunk, output_format) + if chunk.done: + return + + yield _format_chunk(StreamChunk(chunk_type=ChunkEnum.DONE, chunk="", done=True), output_format) + return + + finally: + # Cancel pending consumer. + if consumer is not None and not consumer.done(): + consumer.cancel() + try: + await consumer + except asyncio.CancelledError: + pass + # Cancel producer to avoid leaks. + if not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass diff --git a/flowllm/utils/env_utils.py b/flowllm/utils/env_utils.py new file mode 100644 index 0000000..4df3b95 --- /dev/null +++ b/flowllm/utils/env_utils.py @@ -0,0 +1,54 @@ +"""Load .env files into os.environ (idempotent).""" + +import os +from pathlib import Path + +_LOADED = False +_LOADED_VALUES: dict[str, str] = {} + + +def _parse_env_file(path: str | Path) -> dict[str, str]: + """Parse a KEY=VALUE env file.""" + path = Path(path) + values: dict[str, str] = {} + for line in path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, value = line.split("=", 1) + key = key.strip() + if not key: + continue + values[key] = value.strip().strip("'\"") + return values + + +def _load_values(values: dict[str, str], *, override: bool) -> dict[str, str]: + loaded: dict[str, str] = {} + for key, value in values.items(): + if override or key not in os.environ: + os.environ[key] = value + loaded[key] = value + return loaded + + +def load_env(path: str | Path | None = None, *, override: bool = True) -> dict[str, str]: + """Load .env from given path or search cwd and parents.""" + global _LOADED + global _LOADED_VALUES + if path is None and _LOADED: + return dict(_LOADED_VALUES) + + if path: + path = Path(path) + if path.exists(): + return _load_values(_parse_env_file(path), override=override) + return {} + + for directory in [Path.cwd(), *Path.cwd().parents[:5]]: + env_path = directory / ".env" + if env_path.exists(): + _LOADED_VALUES = _load_values(_parse_env_file(env_path), override=override) + _LOADED = True + return dict(_LOADED_VALUES) + return {} diff --git a/flowllm/utils/logger_utils.py b/flowllm/utils/logger_utils.py new file mode 100644 index 0000000..fec6075 --- /dev/null +++ b/flowllm/utils/logger_utils.py @@ -0,0 +1,108 @@ +"""Logger utilities (loguru / stdlib).""" + +import logging +import os +import sys +from datetime import datetime +from logging.handlers import TimedRotatingFileHandler + +_logger = None + +_LOGURU_FORMAT = "{time:YYYY-MM-DD HH:mm:ss} | {level} | {file}:{line} | {function} | {message}" +_STDLIB_FORMAT = "%(asctime)s | %(levelname)s | %(filename)s:%(lineno)d | %(funcName)s | %(message)s" +_STDLIB_DATEFMT = "%Y-%m-%d %H:%M:%S" + + +def _enable_loguru() -> bool: + return os.getenv("FLOWLLM_DISABLE_LOGURU", "").lower() != "true" + + +def _init_loguru(log_dir: str, level: str, log_to_console: bool, log_to_file: bool): + from loguru import logger + + logger.remove() + + if log_to_console: + logger.add( + sink=sys.stdout, + level=level, + format=_LOGURU_FORMAT, + colorize=True, + ) + + if log_to_file: + try: + os.makedirs(log_dir, exist_ok=True) + current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + log_filepath = os.path.join(log_dir, f"{current_ts}.log") + + logger.add( + log_filepath, + level=level, + rotation="00:00", + retention="7 days", + compression="zip", + encoding="utf-8", + format=_LOGURU_FORMAT, + ) + except Exception as e: + logger.error(f"Error configuring file logging: {e}") + + return logger + + +def _init_stdlib(log_dir: str, level: str, log_to_console: bool, log_to_file: bool): + logger = logging.getLogger("flowllm") + logger.setLevel(level) + logger.propagate = False + + for handler in list(logger.handlers): + logger.removeHandler(handler) + + formatter = logging.Formatter(_STDLIB_FORMAT, datefmt=_STDLIB_DATEFMT) + + if log_to_console: + console_handler = logging.StreamHandler(sys.stdout) + console_handler.setLevel(level) + console_handler.setFormatter(formatter) + logger.addHandler(console_handler) + + if log_to_file: + try: + os.makedirs(log_dir, exist_ok=True) + current_ts = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + log_filepath = os.path.join(log_dir, f"{current_ts}.log") + + file_handler = TimedRotatingFileHandler( + log_filepath, + when="midnight", + backupCount=7, + encoding="utf-8", + ) + file_handler.setLevel(level) + file_handler.setFormatter(formatter) + logger.addHandler(file_handler) + except Exception as e: + logger.error(f"Error configuring file logging: {e}") + + return logger + + +def get_logger( + log_dir: str = "logs", + level: str = "INFO", + log_to_console: bool = True, + log_to_file: bool = True, + force_init: bool = False, +): + """Return the global logger, initializing on first call.""" + global _logger + + if _logger is not None and not force_init: + return _logger + + if _enable_loguru(): + _logger = _init_loguru(log_dir, level, log_to_console, log_to_file) + else: + _logger = _init_stdlib(log_dir, level, log_to_console, log_to_file) + return _logger diff --git a/flowllm/utils/logo_utils.py b/flowllm/utils/logo_utils.py new file mode 100644 index 0000000..01a40ee --- /dev/null +++ b/flowllm/utils/logo_utils.py @@ -0,0 +1,97 @@ +"""Startup banner.""" + +import colorsys +import importlib.metadata +import random +from typing import TYPE_CHECKING + +from rich.console import Console, Group +from rich.panel import Panel +from rich.table import Table +from rich.text import Text + +if TYPE_CHECKING: + from ..schema import ApplicationConfig + + +def _get_version(package_name: str) -> str: + """Return installed package version or empty string.""" + try: + return importlib.metadata.version(package_name) + except importlib.metadata.PackageNotFoundError: + return "" + + +def _hsv_rgb(h: float, s: float = 0.85, v: float = 0.98) -> tuple[int, int, int]: + """HSV → 0-255 RGB tuple.""" + r, g, b = colorsys.hsv_to_rgb(h % 1.0, s, v) + return int(r * 255), int(g * 255), int(b * 255) + + +def print_logo(app_config: "ApplicationConfig"): + """Print rainbow ASCII logo and runtime config.""" + ascii_art = [ + r" ███████╗ ██╗ ██████╗ ██╗ ██╗ ██╗ ██╗ ███╗ ███╗ ", + r" ██╔════╝ ██║ ██╔═══██╗ ██║ ██║ ██║ ██║ ████╗ ████║ ", + r" █████╗ ██║ ██║ ██║ ██║ █╗ ██║ ██║ ██║ ██╔████╔██║ ", + r" ██╔══╝ ██║ ██║ ██║ ██║███╗██║ ██║ ██║ ██║╚██╔╝██║ ", + r" ██║ ███████╗╚██████╔╝ ╚███╔███╔╝ ███████╗ ███████╗ ██║ ╚═╝ ██║ ", + r" ╚═╝ ╚══════╝ ╚═════╝ ╚══╝╚══╝ ╚══════╝ ╚══════╝ ╚═╝ ╚═╝ ", + ] + + hue_base = random.random() # random starting hue per startup + horizontal_span = 0.5 # half the wheel left-to-right + vertical_shift = 0.08 # small per-line nudge for 2D rainbow + + logo_text = Text() + for line_idx, line in enumerate(ascii_art): + line_len = max(1, len(line) - 1) + line_hue_start = hue_base + line_idx * vertical_shift + for i, char in enumerate(line): + ratio = i / line_len + r, g, b = _hsv_rgb(line_hue_start + horizontal_span * ratio) + logo_text.append(char, style=f"bold rgb({r},{g},{b})") + logo_text.append("\n") + + info_table = Table.grid(padding=(0, 1)) + info_table.add_column(style="bold", justify="center") + info_table.add_column(style="bold cyan", justify="left") + info_table.add_column(style="white", justify="left") + + # Backend-specific fields live in model_extra. + service = app_config.service + backend = service.backend + extra = service.model_extra or {} + + info_table.add_row("📦", "Backend:", backend) + + match backend: + case "http": + host = extra.get("host", "localhost") + port = extra.get("port", 8000) + info_table.add_row("🔗", "URL:", f"http://{host}:{port}") + info_table.add_row("📚", "FastAPI:", Text(_get_version("fastapi"), style="dim")) + case "mcp": + transport = extra.get("transport", "stdio") + info_table.add_row("🚌", "Transport:", transport) + if transport != "stdio": + host = extra.get("host", "localhost") + port = extra.get("port", 8000) + url = f"http://{host}:{port}" + if transport == "sse": + url += "/sse" + info_table.add_row("🔗", "URL:", url) + info_table.add_row("📚", "FastMCP:", Text(_get_version("fastmcp"), style="dim")) + + info_table.add_row("🚀", "FlowLLM:", Text(_get_version("flowllm"), style="dim")) + + panel = Panel( + Group(logo_text, info_table), + title=app_config.app_name, + title_align="left", + border_style="dim", + padding=(1, 4), + expand=False, + ) + + Console().print(Group("\n", panel, "\n")) diff --git a/flowllm/utils/service_utils.py b/flowllm/utils/service_utils.py new file mode 100644 index 0000000..77b8210 --- /dev/null +++ b/flowllm/utils/service_utils.py @@ -0,0 +1,97 @@ +"""Service discovery utilities.""" + +import asyncio +import socket +import subprocess +import sys + +from ..constants import FLOWLLM_DEFAULT_HOST, FLOWLLM_DEFAULT_PORT + + +async def find_flowllm(host: str, port: int) -> str: + """Probe host:port. Returns 'flowllm', 'occupied', or 'free'.""" + from ..components.client.http_client import HttpClient + + try: + async with HttpClient(host=host, port=port, timeout=2.0) as client: + async for _ in client(action="health_check"): + break + return "flowllm" + except Exception: + pass + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + try: + s.bind((host, port)) + return "free" + except OSError: + return "occupied" + + +def _sh(cmd: list[str]) -> str: + """Run cmd; return stdout, or '' on failure.""" + try: + return subprocess.check_output(cmd, stderr=subprocess.DEVNULL, text=True) + except (subprocess.CalledProcessError, FileNotFoundError): + return "" + + +def _pid_on_port(port: int) -> int | None: + """PID listening on TCP port, or None.""" + out = _sh(["lsof", "-nP", f"-iTCP:{port}", "-sTCP:LISTEN", "-t"]).strip() + return int(out.splitlines()[0]) if out else None + + +def _scan_flowllm_procs() -> list[tuple[int, str, int]]: + """List running flowllm start processes.""" + procs: list[tuple[int, str, int]] = [] + for line in _sh(["pgrep", "-af", "flowllm.* start"]).splitlines(): + parts = line.split() + if not parts or not parts[0].isdigit(): + continue + host, port = FLOWLLM_DEFAULT_HOST, FLOWLLM_DEFAULT_PORT + for t in parts[1:]: + if t.startswith("service.host="): + host = t.split("=", 1)[1] + elif t.startswith("service.port=") and t.split("=", 1)[1].isdigit(): + port = int(t.split("=", 1)[1]) + procs.append((int(parts[0]), host, port)) + return procs + + +async def _locate_flowllm() -> tuple[str, int, int | None] | None: + """Find a running flowllm instance.""" + if await find_flowllm(FLOWLLM_DEFAULT_HOST, FLOWLLM_DEFAULT_PORT) == "flowllm": + return FLOWLLM_DEFAULT_HOST, FLOWLLM_DEFAULT_PORT, _pid_on_port(FLOWLLM_DEFAULT_PORT) + for pid, host, port in _scan_flowllm_procs(): + if await find_flowllm(host, port) == "flowllm": + return host, port, pid + return None + + +def precheck_start(svc_config: dict | None) -> bool: + """Pre-flight check before starting.""" + host = (svc_config or {}).get("host") or FLOWLLM_DEFAULT_HOST + port = (svc_config or {}).get("port") or FLOWLLM_DEFAULT_PORT + port = int(port) + status = asyncio.run(find_flowllm(host, port)) + if status == "flowllm": + print(f"flowllm already running at {host}:{port}") + return False + if status == "occupied": + print( + f"port {port} occupied. Start on another port: flowllm start service.port=", + file=sys.stderr, + ) + sys.exit(1) + return True + + +def cli_find_flowllm() -> None: + """Print running flowllm HOST/PORT/PID.""" + found = asyncio.run(_locate_flowllm()) + if not found: + print("flowllm not started. Try: flowllm start", file=sys.stderr) + sys.exit(1) + host, port, pid = found + print(f"HOST={host} PORT={port} PID={pid or 'unknown'}") diff --git a/mkdocs.yml b/mkdocs.yml deleted file mode 100644 index 968ec48..0000000 --- a/mkdocs.yml +++ /dev/null @@ -1,68 +0,0 @@ -site_name: FlowLLM -site_url: https://flowllm-ai.github.io/flowllm/ -site_description: "FlowLLM: Simplifying LLM-based HTTP/MCP Service Development" -repo_url: https://github.com/flowllm-ai/flowllm -repo_name: flowllm-ai/flowllm - -theme: - name: shadcn - image: figure/logo.png - show_stargazers: true - locale: zh_CN - -plugins: - - search - - excalidraw - -nav: - - 欢迎: index.md - - 快速开始: - - 安装: zh/guide/installation.md - - 配置: zh/guide/config_guide.md - - 学习资料分享: - - HaluMem解读: zh/reading/20251114-halumem.md - - Gemini-CLI上下文管理: zh/reading/20251113-gemini-cli-context-management.md - - 上下文管理指南: zh/reading/20251110-manus-context-report.md - - 上下文管理视频材料: zh/reading/20251110-manus-context-raw.md - - 操作开发: - - 操作介绍: zh/guide/op_introduction.md - - 简单操作指南: zh/guide/async_op_minimal_guide.md - - LLM 操作指南: zh/guide/async_op_llm_guide.md - - 高级操作指南: zh/guide/async_op_advance_guide.md - - 工具操作指南: zh/guide/async_tool_op_guide.md - - File Tool Op 指南: zh/guide/file_tool_op_guide.md - - 向量存储指南: zh/guide/vector_store_guide.md - - 流程编排: - - 流程指南: zh/guide/flow_guide.md - - 服务使用: - - HTTP 服务指南: zh/guide/http_service_guide.md - - HTTP 流指南: zh/guide/http_stream_guide.md - - MCP 服务指南: zh/guide/mcp_service_guide.md - - CMD 服务指南: zh/guide/cmd_service_guide.md - - 客户端指南: zh/guide/client_guide.md - - 贡献: - - 贡献指南: zh/guide/contribution.md - -markdown_extensions: - admonition: - codehilite: - fenced_code: - footnotes: - extra: - pymdownx.blocks.details: - pymdownx.tabbed: - pymdownx.blocks.tab: - combine_header_slug: true - separator: ___ - pymdownx.progressbar: - pymdownx.snippets: - pymdownx.arithmatex: - generic: true - shadcn.extensions.echarts.alpha: - shadcn.extensions.codexec: - shadcn.extensions.iconify: - -# pip install mkdocs-shadcn -# mkdocs build -# mkdocs serve -# mkdocs gh-deploy --force \ No newline at end of file diff --git a/old/INSTALLATION.md b/old/INSTALLATION.md deleted file mode 100644 index f072a6e..0000000 --- a/old/INSTALLATION.md +++ /dev/null @@ -1,103 +0,0 @@ -# 🛠️ Installation Guide - -## Install from PyPI (Recommended) - -```bash -pip install flowllm -``` - -## Install with Optional Dependencies - -```bash -# For financial data support -pip install flowllm[fin] - -# For distributed computing -pip install flowllm[dist] - -# For Qwen model support -pip install flowllm[qwen] - -# Install all features -pip install flowllm[all] -``` - -## Install from Source - -```bash -git clone https://github.com/your-org/flowllm.git -cd flowllm -pip install . -``` - -## Environment Configuration - -Copy `example.env` to `.env` and modify the corresponding parameters: - -```bash -FLOW_APP_NAME=FlowLLM - -# LLM Configuration -FLOW_LLM_API_KEY=sk-xxxx -FLOW_LLM_BASE_URL=https://xxxx/v1 -FLOW_LLM_MODEL=qwen3-max - -# Embedding Configuration -FLOW_EMBEDDING_API_KEY=sk-xxxx -FLOW_EMBEDDING_BASE_URL=https://xxxx/v1 - -# Optional: Elasticsearch -# FLOW_ES_HOSTS=http://0.0.0.0:9200 - -# Optional: Additional API Keys -# FLOW_DASHSCOPE_API_KEY=sk-xxxx -# FLOW_TAVILY_API_KEY=sk-xxxx -``` - -## Verify Installation - -After installation, you can verify that FlowLLM is installed correctly: - -```bash -# Check version -python -c "import flowllm; print(flowllm.__version__)" - -# Run help command -flowllm --help -``` - -## Troubleshooting - -### Common Issues - -#### 1. Import Error - -If you encounter import errors, make sure you have installed all required dependencies: - -```bash -pip install -r requirements.txt -``` - -#### 2. Environment Variables Not Found - -Ensure your `.env` file is in the correct location (project root) or export environment variables manually: - -```bash -export FLOW_LLM_API_KEY=your_api_key -export FLOW_LLM_BASE_URL=your_base_url -``` - -#### 3. Port Already in Use - -If the default port (8002) is already in use, specify a different port: - -```bash -flowllm backend=http http.port=8003 -``` - -## Next Steps - -- Read the [Quick Start Guide](QUICKSTART.md) to get started -- Check out the [README](README.md) for architecture overview -- Explore the [documentation](docs/) for detailed guides - diff --git a/old/QUICKSTART.md b/old/QUICKSTART.md deleted file mode 100644 index f0319dc..0000000 --- a/old/QUICKSTART.md +++ /dev/null @@ -1,232 +0,0 @@ -# 🚀 Quick Start Guide - -## 30-Second Setup - -```bash -# Install -pip install flowllm - -# Set environment variables -export FLOW_LLM_API_KEY=sk-xxx -export FLOW_LLM_BASE_URL=https://api.openai.com/v1 - -# Start HTTP service (includes pre-built flows) -flowllm backend=http http.port=8002 - -# Or start MCP service -flowllm backend=mcp mcp.transport=stdio -``` - -That's it! Pre-built flows are now available as HTTP endpoints or MCP tools. - -## Using Pre-built Services - -```bash -# Use built-in LLM flow -curl -X POST http://localhost:8002/llm_flow \ - -H "Content-Type: application/json" \ - -d '{"messages": [{"role": "user", "content": "Hello!"}]}' - -# With streaming -curl -X POST http://localhost:8002/llm_flow_stream \ - -H "Content-Type: application/json" \ - -d '{"messages": [{"role": "user", "content": "Tell me a story"}]}' -``` - -## Creating Your First Custom Service - -### Step 1: Write an Operation (Python) - -```python -# my_ops.py -from flowllm import BaseOp, C - -class SummarizeOp(BaseOp): - """Summarize text using LLM""" - def call(self, text: str, max_length: int = 100) -> dict: - llm = C.get_llm("default") - prompt = f"Summarize this in {max_length} words: {text}" - summary = llm.call(messages=[{"role": "user", "content": prompt}]) - return {"summary": summary} -``` - -### Step 2: Configure as Service (YAML) - -```yaml -# my_config.yaml -import_config: default - -backend: http -http: - port: 8002 - -op: - summarize_op: - backend: base_op - class: SummarizeOp - -flow: - summarize: - description: "Summarize long text into concise format" - flow_content: summarize_op - tool: - parameters: - text: - type: string - description: "Text to summarize" - required: true - max_length: - type: integer - description: "Maximum summary length in words" - required: false - default: 100 -``` - -### Step 3: Launch - -```bash -flowllm config=my_config -``` - -### Step 4: Use - -```bash -# HTTP endpoint automatically created -curl -X POST http://localhost:8002/summarize \ - -H "Content-Type: application/json" \ - -d '{ - "text": "Long article text here...", - "max_length": 50 - }' -``` - -**No FastAPI routes, no manual validation, no extra code - just configuration!** - -## Using HTTP Client - -### Python Requests - -```python -import requests - -# Execute a flow -response = requests.post("http://localhost:8002/my_custom_flow", json={ - "input_text": "Hello, FlowLLM!" -}) - -print(response.json()) -``` - -### curl - -```bash -curl -X POST http://localhost:8002/my_custom_flow \ - -H "Content-Type: application/json" \ - -d '{ - "input_text": "Hello, FlowLLM!" - }' -``` - -### Node.js - -```javascript -fetch("http://localhost:8002/my_custom_flow", { - method: "POST", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify({ - input_text: "Hello, FlowLLM!" - }) -}) -.then(response => response.json()) -.then(data => console.log(data)); -``` - -## Using Python Client - -### Synchronous Execution - -```python -from flowllm import FlowLLMApp - -# Initialize app -with FlowLLMApp(load_default_config=True) as app: - # Execute flow synchronously - result = app.execute_flow( - "my_custom_flow", - input_text="Hello, FlowLLM!" - ) - print(result) -``` - -### Async Usage - -```python -import asyncio -from flowllm import FlowLLMApp - -async def main(): - async with FlowLLMApp(load_default_config=True) as app: - # Execute flow asynchronously - result = await app.async_execute_flow( - "my_custom_flow", - input_text="Hello, FlowLLM!" - ) - print(result) - -asyncio.run(main()) -``` - -## Next Steps - -- **Learn More**: Check out the [Configuration-Driven Development](README.md#-configuration-driven-development) section -- **Explore Examples**: See the [Example Workflows](README.md#-example-workflows) section -- **Advanced Features**: Read about [Core Features](README.md#-core-features) -- **Deep Dive**: Browse the [documentation](docs/) directory for detailed guides -- **Build**: Follow the complete example in [Configuration-Driven Development](README.md#-configuration-driven-development) - -## Common Patterns - -### Simple LLM Chat - -```python -from flowllm import BaseOp, FlowLLMApp, C - -class ChatOp(BaseOp): - def call(self, message: str) -> dict: - llm = C.get_llm("default") - response = llm.call(messages=[ - {"role": "user", "content": message} - ]) - return {"response": response} - -with FlowLLMApp(load_default_config=True) as app: - result = app.execute_flow("chat", message="Hello!") - print(result) -``` - -### Sequential Workflow - -```yaml -flow: - research_flow: - description: "Multi-step research workflow" - flow_content: "search_op >> summarize_op >> validate_op" -``` - -### Parallel Processing - -```yaml -flow: - parallel_analysis: - description: "Analyze from multiple perspectives" - flow_content: "(sentiment_op | keywords_op | summary_op)" -``` - -## Getting Help - -- **Documentation**: [doc/](docs/) directory -- **Issues**: [GitHub Issues](https://github.com/your-org/flowllm/issues) -- **Email**: jinli.yl@alibaba-inc.com - diff --git a/old/README.md b/old/README.md deleted file mode 100644 index e2fd58b..0000000 --- a/old/README.md +++ /dev/null @@ -1,239 +0,0 @@ -

- FlowLLM Logo -

- -

- FlowLLM: A Flexible Framework for Building LLM-Powered Flows
- Flow with Intelligence, Build with Simplicity. -

- -

- Python Version - PyPI Version - License -

- ---- - -FlowLLM is a **configuration-driven** framework for building LLM-powered applications. Write operations once, compose them via YAML configuration, and automatically get HTTP APIs and MCP tools—no boilerplate code needed. - -## 📖 Table of Contents - -- [Why FlowLLM?](#-why-flowllm) -- [Getting Started](#-getting-started) -- [Core Workflow](#-core-workflow) -- [Architecture](#-architecture) -- [Features](#-features) -- [Resources](#-resources) - ---- - -## 💡 Why FlowLLM? - -**The Problem**: Building LLM services traditionally requires writing boilerplate routes, validation, documentation, and orchestration code for each endpoint. - -**The Solution**: FlowLLM's configuration-driven approach lets you: - -- ✅ **Write Operations Once** - Focus on business logic in reusable Python ops -- ✅ **Configure, Don't Code** - Compose workflows using YAML configuration -- ✅ **Auto-Generate Services** - HTTP and MCP endpoints created automatically -- ✅ **Built-in Orchestration** - Sequential (`>>`), parallel (`|`), and nested flows -- ✅ **Zero Boilerplate** - No routes, validators, or service code needed - -| Feature | Traditional Approach | FlowLLM Approach | -|---------|---------------------|------------------| -| **Service Creation** | Write FastAPI/Flask routes, handlers, validation | Write YAML config - auto-registers HTTP + MCP | -| **API Documentation** | Manually write OpenAPI specs | Auto-generated from config | -| **Workflow Changes** | Modify Python code, test, redeploy | Update config, restart service | -| **Orchestration** | Write custom coordination code | Use expressions: `>>`, `\|`, `()` | - -**Perfect For**: Rapid prototyping, microservices, AI agent tools, data pipelines, enterprise AI applications. - ---- - -## 🚀 Getting Started - -### Installation - -```bash -pip install flowllm -``` - -For detailed setup instructions, see the [Installation Guide](INSTALLATION.md). - -### Quick Start - -See the [Quick Start Guide](QUICKSTART.md) to build your first LLM service in 30 seconds. - ---- - -## 🎯 Core Workflow - -``` -┌─────────────────┐ ┌──────────────────┐ ┌──────────────────────────┐ -│ Build Ops │ │ Configure YAML │ │ Auto-Register │ -│ (Python) │ → │ (Workflows) │ → │ Services │ -│ │ │ │ │ │ -│ • BaseOp │ │ flow: │ │ ┌────────────────────┐ │ -│ • BaseAsyncOp │ │ workflow: │ │ │ HTTP Service │ │ -│ • BaseMcpOp │ │ description │ │ │ POST /workflow │ │ -│ • BaseRayOp │ │ flow_content│ │ │ OpenAPI docs │ │ -│ │ │ tool: │ │ └────────────────────┘ │ -│ │ │ parameters│ │ │ -│ │ │ │ │ ┌────────────────────┐ │ -│ │ │ backend: http │ │ │ MCP Service │ │ -│ │ │ or mcp │ │ │ Tool: workflow │ │ -│ │ │ │ │ │ Auto-discovered │ │ -│ │ │ │ │ └────────────────────┘ │ -└─────────────────┘ └──────────────────┘ └──────────────────────────┘ -``` - -**Three Simple Steps:** - -1. **Create an Op** - Write a Python class implementing your business logic -2. **Configure in YAML** - Define workflow and service endpoints -3. **Launch** - Run one command to start your HTTP or MCP service - -**No manual routing, no endpoint definitions, no service code - just configuration!** - ---- - -## ✨ Architecture - -FlowLLM adopts a **three-layer configuration-driven architecture**: - -``` -┌─────────────────────────────────────────────────────────────────────┐ -│ Service Layer (外层) │ -│ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │ -│ │ HTTP Service │ │ MCP Service │ │ CMD Service │ │ -│ │ FastAPI │ │ FastMCP │ │ Command Line│ │ -│ └──────────────┘ └──────────────┘ └──────────────┘ │ -│ Auto-Register from Configuration │ -└─────────────────────────────┬───────────────────────────────────────┘ - │ -┌─────────────────────────────┴───────────────────────────────────────┐ -│ Flow Layer (中层) │ -│ • Sequential: op1 >> op2 >> op3 │ -│ • Parallel: (op1 | op2 | op3) │ -│ • Nested: op1 >> (op2 | op3) >> op4 │ -└─────────────────────────────┬───────────────────────────────────────┘ - │ -┌─────────────────────────────┴───────────────────────────────────────┐ -│ Foundation Layer (底层) │ -│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ │ -│ │ Op Lib │ │ LLM Lib │ │Embedding │ │ Storage │ │ -│ ├──────────┤ ├──────────┤ ├──────────┤ ├──────────┤ │ -│ │ BaseOp │ │ OpenAI │ │ OpenAI │ │ElasticS. │ │ -│ │BaseAsync │ │ LiteLLM │ │Compatible│ │ChromaDB │ │ -│ │BaseTool │ │DashScope │ │ │ │ Local │ │ -│ │BaseMcpOp │ │ Custom │ │ Custom │ │ Cache │ │ -│ │BaseRayOp │ │ │ │ │ │ │ │ -│ └──────────┘ └──────────┘ └──────────┘ └──────────┘ │ -└──────────────────────────────────────────────────────────────────────┘ -``` - -### Key Design Principles - -1. **Separation of Concerns** - Ops (business logic), Flows (orchestration), Services (protocol handling) -2. **Configuration over Code** - Ops in Python, Flows in YAML, Services auto-generated -3. **Dependency Injection** - ServiceContext manages shared resources (LLM, VectorStore, etc.) -4. **Registry Pattern** - Dynamic loading and discovery based on configuration - -### Complete Data Flow - -``` -Request → Service Layer (HTTP/MCP) - ↓ -Flow Layer (Parse expression → Build DAG) - ↓ -Foundation Layer (Execute ops with context) - ↓ -Response (JSON/MCP result) -``` - ---- - -## 🎯 Features - -### 📦 Pre-built Operations - -**Gallery Ops**: `SimpleLLMOp`, `ReactLLMOp`, `ExecuteCodeOp`, `TranslateCodeOp` -**Search Ops**: `TavilySearchOp`, `DashScopeSearchOp`, `McpSearchOp` -**Research Ops**: `DashScopeDeepResearchOp`, `LangChainDeepResearchOp` -**Data Ops**: Various extraction and processing operations - -### 🔧 Advanced Capabilities - -- **Multi-LLM Support** - OpenAI, LiteLLM (100+ providers), DashScope, custom providers -- **Vector Storage** - Elasticsearch, ChromaDB, local file-based, in-memory -- **Async/Streaming** - Full async support with SSE streaming responses -- **Distributed Computing** - Ray integration for scaling operations -- **Caching** - Intelligent caching with TTL and automatic serialization -- **Web Crawling** - Integrated `crawl4ai` for content extraction - -### 🧪 Workflow Patterns - -- **Simple LLM Chat** - Direct model interaction -- **Multi-Step Research** - Sequential search, summarization, validation -- **Parallel Processing** - Concurrent sentiment analysis, keyword extraction -- **Complex Pipelines** - Nested sequential and parallel operations - ---- - -## 📚 Resources - -### Documentation - -- **[Installation Guide](INSTALLATION.md)** - Setup and environment configuration -- **[Quick Start Guide](QUICKSTART.md)** - Build your first service -- **Specialized Guides** in `doc/`: - - [Deep Research Guide](docs/deep_research.md) - - [Financial Supply Guide](docs/fin_supply_readme.md) - - [Vector Store Guide](docs/vector_store.md) - -### Examples & Configuration - -- **Examples**: `test/` directory for practical examples -- **Configuration**: `flowllm/config/` for sample configs - -### Latest Updates - -- **[2025-10]** FlowLLM v0.1.10 - Enhanced async support and stability -- **[2025-09]** Financial data modules with 26+ pre-built flows -- **[2025-09]** Deep research with multiple search backends -- **[2025-08]** MCP (Model Context Protocol) support -- **[2025-06]** Multi-backend vector storage - ---- - - -## ⚖️ License - -Apache License 2.0 - see [LICENSE](./LICENSE) file for details. - ---- - -## 🌟 Star History - -If you find FlowLLM useful, please consider giving it a star! - - -1. op简介 -2. 写简单的异步op，增加llm和prompt -4. 其他的能力,emb,vectorstore，并发 -5. 进阶版工具tool没了能够 -6. 配置pipeline -7. http服务例子 -8. mcp服务例子 -9. http stream例子 -10. readme例子 学习课堂 -11. contribute -12. installation -13. quick start - - -写一个op介绍，要求简洁，口语化： -1. 包括继承关系总览，他们是怎么继承的 -2. 表格对比不同的op，从各个维度，我可以想到的，比如能力，参数、适合使用的场景等等维度 -3. 诶个详细介绍每一个op，不需要代码哈 \ No newline at end of file diff --git a/old/__init__.py b/old/__init__.py deleted file mode 100644 index 92c2f0e..0000000 --- a/old/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ - - -from flowllm.app import FlowLLMApp -from flowllm.context.service_context import C -from flowllm.op import BaseOp, BaseAsyncOp, BaseAsyncToolOp, BaseMcpOp, BaseRayOp diff --git a/old/app_ui.py b/old/app_ui.py deleted file mode 100644 index f36e0b2..0000000 --- a/old/app_ui.py +++ /dev/null @@ -1,584 +0,0 @@ -import json -import time -import uuid -from datetime import datetime - -import requests -import streamlit as st - -from flowllm.context import C - -APP_NAME: str = os.getenv("FLOW_APP_NAME") -available_models = [ - "langchain+brief+bailian_search", - "langchain+bailian_search", - "dashscope_deep_research", - "llm_flow_stream", - # Add more models here in the future -] - -st.set_page_config( - page_title=f"{APP_NAME} Chat", - page_icon="🤖", - layout="wide", - initial_sidebar_state="expanded", -) - -st.markdown( - """ - - -""", - unsafe_allow_html=True, -) - -if "conversations" not in st.session_state: - st.session_state.conversations = {} - -if "current_conversation_id" not in st.session_state: - st.session_state.current_conversation_id = None - -if "is_streaming" not in st.session_state: - st.session_state.is_streaming = False - -if "selected_model" not in st.session_state: - st.session_state.selected_model = available_models[0] - - -def create_new_conversation(): - conversation_id = str(uuid.uuid4()) - st.session_state.conversations[conversation_id] = { - "title": "New Conversation", - "messages": [], - "created_at": datetime.now(), - } - st.session_state.current_conversation_id = conversation_id - return conversation_id - - -def get_conversation_title(messages): - for msg in messages: - if msg["role"] == "user": - content = msg["content"] - return content[:20] + "..." if len(content) > 20 else content - return "New Conversation" - - -def render_chunk_content(chunk_type, content, is_collapsible=True): - """Render different chunk types with appropriate styling using Streamlit components""" - - chunk_icons = { - "answer": "💬", - "think": "🤔", - "error": "❌", - "tool": "🔧", - } - - chunk_labels = { - "answer": "Answer", - "think": "Thinking", - "error": "Error", - "tool": "Tool Call", - } - - icon = chunk_icons.get(chunk_type, "💬") - label = chunk_labels.get(chunk_type, "Response") - - # Convert newlines to HTML line breaks for proper display - formatted_content = content.replace("\n", "
") - - # Use Streamlit expander for collapsible types (think, error, tool) - if chunk_type in ["think", "error", "tool"] and is_collapsible: - # Set different default expanded states - expanded_default = chunk_type == "error" # Error expanded by default, others collapsed - - with st.expander(f"{icon} {label}", expanded=expanded_default): - st.markdown(f'
{formatted_content}
', unsafe_allow_html=True) - else: - # For non-collapsible types (answer) or during streaming - render as HTML - st.markdown( - f""" -
-
- {icon} {label} -
-
- {formatted_content} -
-
- """, - unsafe_allow_html=True, - ) - - -def stream_llm_response(messages, model_name="llm_flow_stream"): - try: - url = f"http://localhost:8002/{model_name}" - headers = {"Content-Type": "application/json"} - data = {"messages": messages} - - response = requests.post(url, headers=headers, json=data, stream=True) - - if response.status_code == 200: - for line in response.iter_lines(): - if line: - line_str = line.decode("utf-8") - if line_str.startswith("data:"): - data_content = line_str[5:] - - if data_content.strip() == "[DONE]": - break - - try: - json_data = json.loads(data_content) - chunk_type = json_data.get("chunk_type", "answer") - chunk_content = json_data.get("chunk", "") - - if chunk_content: - yield { - "type": chunk_type, - "content": chunk_content, - } - except json.JSONDecodeError: - continue - else: - yield { - "type": "error", - "content": f"API request failed, status code: {response.status_code}", - } - - except Exception as e: - yield { - "type": "error", - "content": f"Error: {str(e)}", - } - - -with st.sidebar: - st.markdown('
💬 Chat History
', unsafe_allow_html=True) - - # Model selection dropdown - st.markdown("**🤖 Model Selection**") - - selected_model = st.selectbox( - "Choose model:", - available_models, - index=( - available_models.index(st.session_state.selected_model) - if st.session_state.selected_model in available_models - else 0 - ), - key="model_selector", - ) - - # Update session state when selection changes - if selected_model != st.session_state.selected_model: - st.session_state.selected_model = selected_model - - st.markdown("---") - - # New conversation button - if st.button("➕ New Chat", key="new_chat", use_container_width=True): - create_new_conversation() - st.rerun() - - st.markdown("---") - - # Display conversation history - if st.session_state.conversations: - for conv_id, conv_data in reversed(list(st.session_state.conversations.items())): - # Update conversation title - if conv_data["messages"] and conv_data["title"] == "New Conversation": - conv_data["title"] = get_conversation_title(conv_data["messages"]) - - # Conversation item style - is_active = conv_id == st.session_state.current_conversation_id - button_style = "conversation-active" if is_active else "conversation-item" - - col1, col2 = st.columns([4, 1]) - with col1: - if st.button( - conv_data["title"], - key=f"conv_{conv_id}", - use_container_width=True, - ): - st.session_state.current_conversation_id = conv_id - st.rerun() - - with col2: - if st.button("🗑️", key=f"del_{conv_id}", help="Delete conversation"): - del st.session_state.conversations[conv_id] - if st.session_state.current_conversation_id == conv_id: - st.session_state.current_conversation_id = None - st.rerun() - - # Display creation time - st.caption(f"Created: {conv_data['created_at'].strftime('%m-%d %H:%M')}") - st.markdown("---") - -st.markdown(f'

{APP_NAME} Chat

', unsafe_allow_html=True) - -if not st.session_state.current_conversation_id: - create_new_conversation() - -current_conv = st.session_state.conversations.get(st.session_state.current_conversation_id, {}) - -chat_container = st.container() - -with chat_container: - if current_conv.get("messages"): - for message in current_conv["messages"]: - if message["role"] == "user": - st.markdown( - f""" -
-
- {message['content']} -
{message.get('timestamp', '')}
-
-
- """, - unsafe_allow_html=True, - ) - else: - # Handle structured message with chunks - if isinstance(message.get("content"), list): - # Sort chunks to show THINK, ERROR, TOOL first, then ANSWER - chunks = message["content"] - sorted_chunks = [] - - # First add THINK, ERROR, TOOL chunks (before answer) - for chunk_type in ["think", "error", "tool"]: - type_chunks = [c for c in chunks if c.get("type") == chunk_type] - sorted_chunks.extend(type_chunks) - - # Then add ANSWER chunks last - answer_chunks = [c for c in chunks if c.get("type") == "answer"] - sorted_chunks.extend(answer_chunks) - - # Create a container with limited width for all chunks - col1, col2 = st.columns([7, 3]) # 70% width for content, 30% empty - with col1: - # Render chunks with proper spacing - for idx, chunk in enumerate(sorted_chunks): - chunk_id = f"{message.get('id', 'msg')}_{idx}" - render_chunk_content( - chunk.get("type", "answer"), - chunk.get("content", ""), - is_collapsible=True, - ) - - st.markdown( - f'
{message.get("timestamp", "")}
', - unsafe_allow_html=True, - ) - else: - # Handle legacy single content messages - st.markdown( - f""" -
-
- {message['content']} -
{message.get('timestamp', '')}
-
-
- """, - unsafe_allow_html=True, - ) - -# Use chat_input instead of text_input + button combination -user_input = st.chat_input( - "Please enter your question...", - disabled=st.session_state.is_streaming, -) - -if user_input and user_input.strip(): - st.session_state.is_streaming = True - - user_message = { - "role": "user", - "content": user_input.strip(), - "timestamp": datetime.now().strftime("%H:%M"), - } - - current_conv["messages"].append(user_message) - - # Extract content from messages for API call - api_messages = [] - for msg in current_conv["messages"]: - if isinstance(msg.get("content"), list): - # Combine all chunks into a single content for API - combined_content = "\n".join( - [chunk.get("content", "") for chunk in msg["content"] if chunk.get("type") == "answer"] - ) - if combined_content: - api_messages.append({"role": msg["role"], "content": combined_content}) - else: - api_messages.append({"role": msg["role"], "content": msg["content"]}) - - st.markdown( - f""" -
-
- {user_input.strip()} -
{datetime.now().strftime('%H:%M')}
-
-
- """, - unsafe_allow_html=True, - ) - - assistant_placeholder = st.empty() - - # Store chunks by type - response_chunks = [] - current_chunks = {"answer": "", "think": "", "error": "", "tool": ""} - - with assistant_placeholder.container(): - response_container = st.empty() - - for chunk_data in stream_llm_response(api_messages, st.session_state.selected_model): - chunk_type = chunk_data.get("type", "answer") - chunk_content = chunk_data.get("content", "") - - # Accumulate content by type - current_chunks[chunk_type] += chunk_content - - # Clear and rebuild the streaming display - response_container.empty() - - with response_container.container(): - # Create a container with limited width for streaming content - col1, col2 = st.columns([7, 3]) # 70% width for content, 30% empty - with col1: - # Display chunks in order: THINK, ERROR, TOOL first, then ANSWER - chunk_order = ["think", "error", "tool", "answer"] - - for c_type in chunk_order: - if current_chunks[c_type].strip(): - chunk_id = f"streaming_{len(current_conv['messages'])}_{c_type}" - render_chunk_content( - c_type, - current_chunks[c_type], - is_collapsible=False, # Don't make collapsible during streaming - ) - - st.markdown('
Typing...
', unsafe_allow_html=True) - - time.sleep(0.05) # Add small delay to simulate typing effect - - # Convert accumulated chunks to structured format, ordered properly - final_chunks = [] - # Order: THINK, ERROR, TOOL first, then ANSWER last - chunk_order = ["think", "error", "tool", "answer"] - for c_type in chunk_order: - if current_chunks[c_type].strip(): - final_chunks.append( - { - "type": c_type, - "content": current_chunks[c_type].strip(), - } - ) - - assistant_message = { - "role": "assistant", - "content": final_chunks if final_chunks else [{"type": "answer", "content": "No response received"}], - "timestamp": datetime.now().strftime("%H:%M"), - "id": str(uuid.uuid4()), - } - - current_conv["messages"].append(assistant_message) - - if current_conv["title"] == "New Conversation": - current_conv["title"] = get_conversation_title(current_conv["messages"]) - - st.session_state.is_streaming = False - - st.rerun() diff --git a/old/client/__init__.py b/old/client/__init__.py deleted file mode 100644 index e909eeb..0000000 --- a/old/client/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .async_http_client import AsyncHttpClient -from .fastmcp_client import FastmcpClient -from .http_client import HttpClient -from .mcp_client import McpClient diff --git a/old/client/http_client.py b/old/client/http_client.py deleted file mode 100644 index ba3a576..0000000 --- a/old/client/http_client.py +++ /dev/null @@ -1,34 +0,0 @@ -from typing import Dict - -import httpx - -from flowllm.schema.flow_response import FlowResponse - - -class HttpClient: - - def __init__(self, base_url: str = "http://localhost:8001", timeout: float = 3600): - self.base_url = base_url.rstrip("/") # Remove trailing slash for consistent URL formatting - self.timeout = timeout - self.client = httpx.Client(timeout=timeout) # Create synchronous HTTP client with timeout - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - self.client.close() - - def close(self): - self.client.close() - - def health_check(self) -> Dict[str, str]: - response = self.client.get(f"{self.base_url}/health") - response.raise_for_status() - return response.json() - - def execute_flow(self, flow_name: str, **kwargs) -> FlowResponse: - endpoint = f"{self.base_url}/{flow_name}" - response = self.client.post(endpoint, json=kwargs) - response.raise_for_status() - result_data = response.json() - return FlowResponse(**result_data) diff --git a/old/client/mcp_client.py b/old/client/mcp_client.py deleted file mode 100644 index ea6037c..0000000 --- a/old/client/mcp_client.py +++ /dev/null @@ -1,236 +0,0 @@ -import asyncio -import os -import shutil -from contextlib import AsyncExitStack -from typing import List, Optional - -import mcp.types -from loguru import logger -from mcp import ClientSession, StdioServerParameters -from mcp.client.sse import sse_client -from mcp.client.stdio import stdio_client -from mcp.client.streamable_http import streamablehttp_client - -from flowllm.schema.tool_call import ToolCall - - -class McpClient: - - def __init__( - self, - name: str, - config: dict, - append_env: bool = False, - max_retries: int = 3, - timeout: Optional[float] = None, - ): - - self.name: str = name - self.config: dict = config - self.append_env: bool = append_env - self.max_retries: int = max_retries - self.timeout: Optional[float] = timeout - - self.session: ClientSession | None = None - self._exit_stack: AsyncExitStack = AsyncExitStack() - - async def astart(self): - command = shutil.which("npx") if self.config.get("command") == "npx" else self.config.get("command") - - if command: - env_params: dict = {} - if self.append_env: - env_params.update(os.environ) - if self.config.get("env"): - env_params.update(self.config["env"]) - - server_params = StdioServerParameters(command=command, args=self.config.get("args", []), env=env_params) - streams = await self._exit_stack.enter_async_context(stdio_client(server_params)) - - else: - kwargs = {"url": self.config["url"]} - if self.config.get("headers"): - headers = self.config.get("headers") - if headers.get("Authorization"): - assert isinstance(headers["Authorization"], str) - headers["Authorization"] = headers["Authorization"].format(**os.environ) - kwargs["headers"] = headers - if "timeout" in self.config: - kwargs["timeout"] = self.config["timeout"] - if "sse_read_timeout" in self.config: - kwargs["sse_read_timeout"] = self.config["sse_read_timeout"] - - if self.config.get("type") in ["streamable_http", "streamableHttp"]: - streams = await self._exit_stack.enter_async_context(streamablehttp_client(**kwargs)) - streams = (streams[0], streams[1]) - else: - streams = await self._exit_stack.enter_async_context(sse_client(**kwargs)) - - session = await self._exit_stack.enter_async_context(ClientSession(*streams)) - await session.initialize() - self.session = session - - async def __aenter__(self) -> "McpClient": - for i in range(self.max_retries): - try: - if self.timeout is not None: - await asyncio.wait_for(self.astart(), timeout=self.timeout) - else: - await self.astart() - break - - except asyncio.TimeoutError: - logger.exception(f"{self.name} start timeout after {self.timeout}s") - - # Clean up the exit stack before retrying - try: - await self._exit_stack.aclose() - except Exception: - pass - self._exit_stack = AsyncExitStack() - - if i == self.max_retries - 1: - raise TimeoutError(f"{self.name} start timeout after {self.timeout}s") - - await asyncio.sleep(1 + i) - - except Exception as e: - logger.exception( - f"{self.name} start failed with {e}. " f"Retry {i + 1}/{self.max_retries} in {1 + i}s...", - ) - - # Clean up the exit stack before retrying - try: - await self._exit_stack.aclose() - except Exception: - pass - self._exit_stack = AsyncExitStack() - - await asyncio.sleep(1 + i) - - if i == self.max_retries - 1: - break - - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb) -> None: - for i in range(self.max_retries): - try: - await self._exit_stack.aclose() - break - - except Exception as e: - logger.exception( - f"{self.name} close failed with {e}. " f"Retry {i + 1}/{self.max_retries} in {1 + i}s...", - ) - await asyncio.sleep(1 + i) - - if i == self.max_retries - 1: - break - - self.session = None - - async def list_tools(self) -> List[mcp.types.Tool]: - if not self.session: - raise RuntimeError(f"Server {self.name} not initialized") - - tools = [] - for i in range(self.max_retries): - try: - if self.timeout is not None: - tools_response = await asyncio.wait_for( - self.session.list_tools(), - timeout=self.timeout, - ) - else: - tools_response = await self.session.list_tools() - - tools = [ - tool - for item in tools_response - if isinstance(item, tuple) and item[0] == "tools" - for tool in item[1] - ] - break - - except asyncio.TimeoutError: - logger.exception(f"{self.name} list tools timeout after {self.timeout}s") - - if i == self.max_retries - 1: - raise TimeoutError(f"{self.name} list tools timeout after {self.timeout}s") - - await asyncio.sleep(1 + i) - - except Exception as e: - logger.exception( - f"{self.name} list tools failed with {e}. " f"Retry {i + 1}/{self.max_retries} in {1 + i}s...", - ) - await asyncio.sleep(1 + i) - - if i == self.max_retries - 1: - raise e - - return tools - - async def list_tool_calls(self) -> List[ToolCall]: - if not self.session: - raise RuntimeError(f"Server {self.name} not initialized") - - tools = await self.list_tools() - return [ToolCall.from_mcp_tool(t) for t in tools] - - async def call_tool(self, tool_name: str, arguments: dict): - if not self.session: - raise RuntimeError(f"Server {self.name} not initialized") - - result = None - for i in range(self.max_retries): - try: - if self.timeout is not None: - result = await asyncio.wait_for( - self.session.call_tool(tool_name, arguments), - timeout=self.timeout, - ) - else: - result = await self.session.call_tool(tool_name, arguments) - break - - except asyncio.TimeoutError: - logger.exception(f"{self.name}.{tool_name} call_tool timeout after {self.timeout}s") - - if i == self.max_retries - 1: - raise TimeoutError(f"{self.name}.{tool_name} call_tool timeout after {self.timeout}s") - - await asyncio.sleep(1 + i) - - except Exception as e: - logger.exception( - f"{self.name}.{tool_name} call_tool failed with {e}. " - f"Retry {i + 1}/{self.max_retries} in {1 + i}s...", - ) - await asyncio.sleep(1 + i) - - if i == self.max_retries - 1: - raise e - - return result - - -async def main(): - config = { - "type": "sse", - "url": "http://11.160.132.45:8010/sse", - "headers": {}, - } - - async with McpClient("mcp", config) as client: - tool_calls = await client.list_tool_calls() - for tool_call in tool_calls: - print(tool_call.model_dump_json()) - - # result = await client.call_tool("search", arguments={"query": "半导体行业PE中位数", "entity": "半导体"}) - # print(result) - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/old/config/__init__.py b/old/config/__init__.py deleted file mode 100644 index a3cb3c3..0000000 --- a/old/config/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .pydantic_config_parser import PydanticConfigParser diff --git a/old/config/deep_search.yaml b/old/config/deep_search.yaml deleted file mode 100644 index 75ad84e..0000000 --- a/old/config/deep_search.yaml +++ /dev/null @@ -1,106 +0,0 @@ -import_config: default - -flow: - dashscope_deep_research: - flow_content: dashscope_deep_research_op - stream: true - - langchain+bailian_search: - flow_content: langchain_deep_research_op2 << (conduct_research_op << bailian_web_search_op << think_tool_op << research_complete_op) << think_tool_op << research_complete_op - stream: true - - # langchain+bocha_search: - # flow_content: langchain_deep_research_op2 << (conduct_research_op << bailian_web_search_op << think_tool_op << research_complete_op) << think_tool_op << research_complete_op - # stream: true - - langchain+brief+bailian_search: - flow_content: langchain_deep_research_op << (conduct_research_op << bailian_web_search_op << think_tool_op << research_complete_op) << think_tool_op << research_complete_op - stream: true - - # langchain+brief+bocha_search: - # flow_content: langchain_deep_research_op << (conduct_research_op << bocha_web_search_op << think_tool_op << research_complete_op) << think_tool_op << research_complete_op - # stream: true - - task_react: - flow_content: task_react_op - stream: true - -op: - langchain_deep_research_op2: - backend: langchain_deep_research_op - params: - enable_research_brief: false - - tavily_search_op: - backend: tavily_search_op - params: - save_answer: false - output_schema_mapping: - tavily_search_result: long_text - - dashscope_search_op: - backend: dashscope_search_op - params: - save_answer: true - - bailian_web_search_op: - backend: base_mcp_op - params: - mcp_name: tongyi_search - tool_name: bailian_web_search - save_answer: true - input_schema_optional: - - count - input_schema_deleted: - - ctx - - # bocha_web_search_op: - # backend: base_mcp_op - # params: - # mcp_name: bochaai_search - # tool_name: bocha_web_search - # save_answer: true - # input_schema_optional: - # - freshness - # - count - - brave_web_search_op: - backend: base_mcp_op - params: - mcp_name: brave_search - tool_name: brave_web_search - save_answer: true - input_schema_optional: - - offset - - count - - bailian_web_parser_op: - backend: base_mcp_op - params: - mcp_name: tongyi_web_parser - tool_name: bailian_web_parser - save_answer: true - input_schema_optional: - - format - - extract_long_text_op: - backend: extract_long_text_op - params: - save_answer: true - -external_mcp: - tongyi_search: - type: sse - url: "https://dashscope.aliyuncs.com/api/v1/mcps/WebSearch/sse" - headers: - Authorization: "Bearer {FLOW_BAILIAN_API_KEY}" - tongyi_web_parser: - type: sse - url: "https://dashscope.aliyuncs.com/api/v1/mcps/WebParser/sse" - headers: - Authorization: "Bearer {FLOW_BAILIAN_API_KEY}" -# bochaai_search: -# type: sse -# url: "https://dashscope.aliyuncs.com/api/v1/mcps/bochaai_search/sse" -# headers: -# Authorization: "Bearer {FLOW_BAILIAN_API_KEY}" diff --git a/old/config/default.yaml b/old/config/default.yaml deleted file mode 100644 index 874df2f..0000000 --- a/old/config/default.yaml +++ /dev/null @@ -1,76 +0,0 @@ -backend: http -thread_pool_max_workers: 128 - -mcp: - transport: sse - host: "0.0.0.0" - port: 8001 - -http: - host: "0.0.0.0" - port: 8002 - - -flow: - llm_flow: - flow_content: simple_llm_op - stream: false - - llm_flow_stream: - flow_content: stream_llm_op - stream: true - -llm: - default: - backend: openai_compatible - model_name: qwen3-30b-a3b-instruct-2507 - params: - temperature: 0.6 - - qwen3_30b_instruct: - backend: openai_compatible - model_name: qwen3-30b-a3b-instruct-2507 - - qwen3_30b_thinking: - backend: openai_compatible - model_name: qwen3-30b-a3b-thinking-2507 - - qwen3_235b_instruct: - backend: openai_compatible - model_name: qwen3-235b-a22b-instruct-2507 - - qwen3_235b_thinking: - backend: openai_compatible - model_name: qwen3-235b-a22b-thinking-2507 - - qwen3_80b_instruct: - backend: openai_compatible - model_name: qwen3-next-80b-a3b-instruct - - qwen3_80b_thinking: - backend: openai_compatible - model_name: qwen3-next-80b-a3b-thinking -# params: -# presence_penalty: 2 - - qwen3_max_instruct: - backend: openai_compatible - model_name: qwen3-max - - qwen25_max_instruct: - backend: openai_compatible - model_name: qwen-max-2025-01-25 - -embedding_model: - default: - backend: openai_compatible - model_name: text-embedding-v4 - params: - dimensions: 1024 - -vector_store: - default: - backend: elasticsearch - embedding_model: default -# params: -# hosts: "http://localhost:9200" diff --git a/old/config/fin_basic.yaml b/old/config/fin_basic.yaml deleted file mode 100644 index dec1eac..0000000 --- a/old/config/fin_basic.yaml +++ /dev/null @@ -1,60 +0,0 @@ -import_config: default -backend: mcp - -flow: -# extract_entities_code: -# flow_content: extract_entities_code_op << bailian_web_search_op -# description: | -# Extract financial entities from the query, including types such as "stock", "bond", "fund", "cryptocurrency", "index", "commodity", "etf", etc. -# For entities like stocks or ETF funds, search for their corresponding codes. Finally, return the financial entities appearing in the query, including their types and codes. -# input_schema: -# query: -# type: string -# description: "user query" -# required: true - - akshare_market: - flow_content: akshare_market_op - - akshare_calculate: - flow_content: akshare_calculate_op - - extract_query: - flow_content: extract_query_op << tongyi_mcp_search_op - - company_factor: - flow_content: company_factor_op - -op: - extract_entities_code_op: - backend: extract_entities_code_op - params: - save_answer: true - - bailian_web_search_op: - backend: base_mcp_op - params: - mcp_name: tongyi_search - tool_name: bailian_web_search - save_answer: true - input_schema_optional: - - count - input_schema_deleted: - - ctx - - akshare_market_op: - backend: akshare_market_op - params: - save_answer: true - - akshare_calculate_op: - backend: akshare_calculate_op - params: - save_answer: true - -external_mcp: - tongyi_search: - type: sse - url: "https://dashscope.aliyuncs.com/api/v1/mcps/WebSearch/sse" - headers: - Authorization: "Bearer {FLOW_BAILIAN_API_KEY}" \ No newline at end of file diff --git a/old/config/fin_mcp.yaml b/old/config/fin_mcp.yaml deleted file mode 100644 index 9b7fee7..0000000 --- a/old/config/fin_mcp.yaml +++ /dev/null @@ -1,43 +0,0 @@ -import_config: default -backend: mcp - -flow: - qtf_brief_mcp: - flow_content: qtf_brief_mcp_op - - qtf_medium_mcp: - flow_content: qtf_medium_mcp_op - - qtf_full_mcp: - flow_content: qtf_full_mcp_op - -op: - qtf_brief_mcp_op: - backend: base_mcp_op - params: - mcp_name: qtf_mcp - tool_name: brief - save_answer: true - - qtf_medium_mcp_op: - backend: base_mcp_op - params: - mcp_name: qtf_mcp - tool_name: medium - save_answer: true - - qtf_full_mcp_op: - backend: base_mcp_op - params: - mcp_name: qtf_mcp - tool_name: full - save_answer: true - -# https://bailian.console.aliyun.com/?tab=mcp#/mcp-market -external_mcp: - qtf_mcp: - # https://bailian.console.aliyun.com/?tab=mcp#/mcp-market/detail/qtf_mcp - type: sse - url: "https://dashscope.aliyuncs.com/api/v1/mcps/qtf_mcp/sse" - headers: - Authorization: "Bearer {FLOW_BAILIAN_API_KEY}" diff --git a/old/config/fin_research.yaml b/old/config/fin_research.yaml deleted file mode 100644 index 8ffc645..0000000 --- a/old/config/fin_research.yaml +++ /dev/null @@ -1,35 +0,0 @@ -import_config: default -backend: mcp - -#flow: - - -#op: -# bailian_web_search_op: -# backend: base_mcp_op -# params: -# mcp_name: tongyi_search -# tool_name: bailian_web_search -# save_answer: true -# input_schema_optional: -# - count -# input_schema_deleted: -# - ctx -# -# bocha_web_search_op: -# backend: base_mcp_op -# params: -# mcp_name: bochaai_search -# tool_name: bocha_web_search -# save_answer: true -# input_schema_optional: -# - freshness -# - count - -external_mcp: - tongyi_search: - type: sse - url: "https://dashscope.aliyuncs.com/api/v1/mcps/WebSearch/sse" - headers: - Authorization: "Bearer {FLOW_BAILIAN_API_KEY}" - diff --git a/old/config/fin_search.yaml b/old/config/fin_search.yaml deleted file mode 100644 index 2f105f9..0000000 --- a/old/config/fin_search.yaml +++ /dev/null @@ -1,107 +0,0 @@ -import_config: default -backend: mcp - -flow: - tavily_search: - flow_content: tavily_search_op >> extract_long_text_op - description: "Use search keywords to retrieve relevant information from the internet. If there are multiple search keywords, please use each keyword separately to call this tool." - input_schema: - query: - type: string - description: "user query" - required: true - - dashscope_search: - flow_content: dashscope_search_op - - bailian_web_search: - flow_content: bailian_web_search_op - -# bocha_web_search: -# flow_content: bocha_web_search_op - - brave_web_search: - flow_content: brave_web_search_op - - bailian_web_parser: - flow_content: bailian_web_parser_op - -op: - tavily_search_op: - backend: tavily_search_op - params: - save_answer: false - output_schema_mapping: - tavily_search_result: long_text - - dashscope_search_op: - backend: dashscope_search_op - params: - save_answer: true - - bailian_web_search_op: - backend: base_mcp_op - params: - mcp_name: tongyi_search - tool_name: bailian_web_search - save_answer: true - input_schema_optional: - - count - input_schema_deleted: - - ctx - -# bocha_web_search_op: -# backend: base_mcp_op -# params: -# mcp_name: bochaai_search -# tool_name: bocha_web_search -# save_answer: true -# input_schema_optional: -# - freshness -# - count - - brave_web_search_op: - backend: base_mcp_op - params: - mcp_name: brave_search - tool_name: brave_web_search - save_answer: true - input_schema_optional: - - offset - - count - - bailian_web_parser_op: - backend: base_mcp_op - params: - mcp_name: tongyi_web_parser - tool_name: bailian_web_parser - save_answer: true - input_schema_optional: - - format - - extract_long_text_op: - backend: extract_long_text_op - params: - save_answer: true - -external_mcp: - tongyi_search: - type: sse - url: "https://dashscope.aliyuncs.com/api/v1/mcps/WebSearch/sse" - headers: - Authorization: "Bearer {FLOW_BAILIAN_API_KEY}" - tongyi_web_parser: - type: sse - url: "https://dashscope.aliyuncs.com/api/v1/mcps/WebParser/sse" - headers: - Authorization: "Bearer {FLOW_BAILIAN_API_KEY}" -# bochaai_search: -# type: sse -# url: "https://dashscope.aliyuncs.com/api/v1/mcps/bochaai_search/sse" -# headers: -# Authorization: "Bearer {FLOW_BAILIAN_API_KEY}" - brave_search: - type: sse - url: "https://dashscope.aliyuncs.com/api/v1/mcps/Brave_Search/sse" - headers: - Authorization: "Bearer {FLOW_BAILIAN_API_KEY}" \ No newline at end of file diff --git a/old/config/fin_supply.yaml b/old/config/fin_supply.yaml deleted file mode 100644 index 7682caa..0000000 --- a/old/config/fin_supply.yaml +++ /dev/null @@ -1,7 +0,0 @@ -import_config: default, fin_basic, fin_mcp, fin_search, fin_ths -backend: mcp -disabled_flows: - - llm_flow - - llm_flow_stream - - mock_tool_flow - - mock_async_tool_flow diff --git a/old/config/fin_ths.yaml b/old/config/fin_ths.yaml deleted file mode 100644 index fe972db..0000000 --- a/old/config/fin_ths.yaml +++ /dev/null @@ -1,282 +0,0 @@ -import_config: default -backend: mcp - -flow: - crawl_ths_company: -# flow_content: ths_company_op >> bailian_web_parser_op >> extract_long_text_op - flow_content: ths_company_op >> crawl4ai_op >> extract_long_text_op - description: "通过股票代码获取公司资料信息，例如：详细情况，高管介绍，发行相关，参控股公司，最后返回和query相关的信息。" - input_schema: - query: - type: string - description: "user query" - required: true - code: - type: string - description: "stock code" - required: true - - crawl_ths_holder: -# flow_content: ths_holder_op >> bailian_web_parser_op >> extract_long_text_op - flow_content: ths_holder_op >> crawl4ai_op >> extract_long_text_op - description: "通过股票代码获取股东研究信息，例如：股东人数、十大流通股东、十大股东、十大债券持有人、控股层级关系，最后返回和query相关的信息。" - input_schema: - query: - type: string - description: "user query" - required: true - code: - type: string - description: "stock code" - required: true - - crawl_ths_operate: -# flow_content: ths_operate_op >> bailian_web_parser_op >> extract_long_text_op - flow_content: ths_operate_op >> crawl4ai_op >> extract_long_text_op - description: "通过股票代码获取经营分析信息，例如：主营介绍、运营业务数据、主营构成分析、主要客户及供应商、董事会经营评述、产品价格，最后返回和query相关的信息。" - input_schema: - query: - type: string - description: "user query" - required: true - code: - type: string - description: "stock code" - required: true - - crawl_ths_equity: -# flow_content: ths_equity_op >> bailian_web_parser_op >> extract_long_text_op - flow_content: ths_equity_op >> crawl4ai_op >> extract_long_text_op - description: "通过股票代码获取股本结构信息，例如：解禁时间表、总股本构成、A股结构图、历次股本变动，最后返回和query相关的信息。" - input_schema: - query: - type: string - description: "user query" - required: true - code: - type: string - description: "stock code" - required: true - - crawl_ths_capital: -# flow_content: ths_capital_op >> bailian_web_parser_op >> extract_long_text_op - flow_content: ths_capital_op >> crawl4ai_op >> extract_long_text_op - description: "通过股票代码获取资本运作信息，例如：募集资金来源、项目投资、收购兼并、股权投资、参股IPO、股权转让、关联交易、质押解冻，最后返回和query相关的信息。" - input_schema: - query: - type: string - description: "user query" - required: true - code: - type: string - description: "stock code" - required: true - - crawl_ths_worth: -# flow_content: ths_worth_op >> bailian_web_parser_op >> extract_long_text_op - flow_content: ths_worth_op >> crawl4ai_op >> extract_long_text_op - description: "通过股票代码获取盈利预测信息，例如：业绩预测、业绩预测详表、研报评级，最后返回和query相关的信息。" - input_schema: - query: - type: string - description: "user query" - required: true - code: - type: string - description: "stock code" - required: true - - crawl_ths_news: -# flow_content: ths_news_op >> bailian_web_parser_op >> extract_long_text_op - flow_content: ths_news_op >> crawl4ai_op >> extract_long_text_op - description: "通过股票代码获取新闻公告信息，例如：新闻与股价联动、公告列表、热点新闻列表、研报列表，最后返回和query相关的信息。" - input_schema: - query: - type: string - description: "user query" - required: true - code: - type: string - description: "stock code" - required: true - - crawl_ths_concept: -# flow_content: ths_concept_op >> bailian_web_parser_op >> extract_long_text_op - flow_content: ths_concept_op >> crawl4ai_op >> extract_long_text_op - description: "通过股票代码获取概念题材信息，例如：常规概念、其他概念、题材要点、概念对比，最后返回和query相关的信息。" - input_schema: - query: - type: string - description: "user query" - required: true - code: - type: string - description: "stock code" - required: true - - crawl_ths_position: -# flow_content: ths_position_op >> bailian_web_parser_op >> extract_long_text_op - flow_content: ths_position_op >> crawl4ai_op >> extract_long_text_op - description: "通过股票代码获取主力持仓信息，例如：机构持股汇总、机构持股明细、被举牌情况、IPO获配机构，最后返回和query相关的信息。" - input_schema: - query: - type: string - description: "user query" - required: true - code: - type: string - description: "stock code" - required: true - - crawl_ths_finance: -# flow_content: ths_finance_op >> bailian_web_parser_op >> extract_long_text_op - flow_content: ths_finance_op >> crawl4ai_op >> extract_long_text_op - description: "通过股票代码获取财务分析信息，例如：财务诊断、财务指标、指标变动说明、资产负债构成、财务报告、杜邦分析，最后返回和query相关的信息。" - input_schema: - query: - type: string - description: "user query" - required: true - code: - type: string - description: "stock code" - required: true - - crawl_ths_bonus: -# flow_content: ths_bonus_op >> bailian_web_parser_op >> extract_long_text_op - flow_content: ths_bonus_op >> crawl4ai_op >> extract_long_text_op - description: "通过股票代码获取分红融资信息，例如：分红诊断、分红情况、增发机构获配明细、增发概况、配股概况，最后返回和query相关的信息。" - input_schema: - query: - type: string - description: "user query" - required: true - code: - type: string - description: "stock code" - required: true - - crawl_ths_event: -# flow_content: ths_event_op >> bailian_web_parser_op >> extract_long_text_op - flow_content: ths_event_op >> crawl4ai_op >> extract_long_text_op - description: "通过股票代码获取公司大事信息，例如：高管持股变动、股东持股变动、担保明细、违规处理、机构调研、投资者互动，最后返回和query相关的信息。" - input_schema: - query: - type: string - description: "user query" - required: true - code: - type: string - description: "stock code" - required: true - - crawl_ths_field: -# flow_content: ths_field_op >> bailian_web_parser_op >> extract_long_text_op - flow_content: ths_field_op >> crawl4ai_op >> extract_long_text_op - description: "通过股票代码获取行业对比信息，例如：行业地位、行业新闻，最后返回和query相关的信息。" - input_schema: - query: - type: string - description: "user query" - required: true - code: - type: string - description: "stock code" - required: true - -op: - ths_company_op: - backend: ths_url_op - params: - url_template: "https://basic.10jqka.com.cn/{code}/company.html#stockpage" - - ths_holder_op: - backend: ths_url_op - params: - url_template: "https://basic.10jqka.com.cn/{code}/holder.html#stockpage" - - ths_operate_op: - backend: ths_url_op - params: - url_template: "https://basic.10jqka.com.cn/{code}/operate.html#stockpage" - - ths_equity_op: - backend: ths_url_op - params: - url_template: "https://basic.10jqka.com.cn/{code}/equity.html#stockpage" - - ths_capital_op: - backend: ths_url_op - params: - url_template: "https://basic.10jqka.com.cn/{code}/capital.html#stockpage" - - ths_worth_op: - backend: ths_url_op - params: - url_template: "https://basic.10jqka.com.cn/{code}/worth.html#stockpage" - - ths_news_op: - backend: ths_url_op - params: - url_template: "https://basic.10jqka.com.cn/{code}/news.html#stockpage" - - ths_concept_op: - backend: ths_url_op - params: - url_template: "https://basic.10jqka.com.cn/{code}/concept.html#stockpage" - - ths_position_op: - backend: ths_url_op - params: - url_template: "https://basic.10jqka.com.cn/{code}/position.html#stockpage" - - ths_finance_op: - backend: ths_url_op - params: - url_template: "https://basic.10jqka.com.cn/{code}/finance.html#stockpage" - - ths_bonus_op: - backend: ths_url_op - params: - url_template: "https://basic.10jqka.com.cn/{code}/bonus.html#stockpage" - - ths_event_op: - backend: ths_url_op - params: - url_template: "https://basic.10jqka.com.cn/{code}/event.html#stockpage" - - ths_field_op: - backend: ths_url_op - params: - url_template: "https://basic.10jqka.com.cn/{code}/field.html#stockpage" - - bailian_web_parser_op: - backend: base_mcp_op - params: - mcp_name: tongyi_web_parser - tool_name: bailian_web_parser - input_schema_optional: - - format - output_schema_mapping: - bailian_web_parser_result: long_text - enable_cache: true - - # 可以使用crawl4ai_op 代替 bailian_web_parser_op - crawl4ai_op: - backend: crawl4ai_op - params: - output_schema_mapping: - crawl4ai_result: long_text - enable_cache: true - - extract_long_text_op: - backend: extract_long_text_op - params: - save_answer: true - -external_mcp: - tongyi_web_parser: - type: sse - url: "https://dashscope.aliyuncs.com/api/v1/mcps/WebParser/sse" - headers: - Authorization: "Bearer {FLOW_BAILIAN_API_KEY}" \ No newline at end of file diff --git a/old/config/mock.yaml b/old/config/mock.yaml deleted file mode 100644 index f53d440..0000000 --- a/old/config/mock.yaml +++ /dev/null @@ -1,7 +0,0 @@ -import_config: default - -flow: - mock_exception: - flow_content: mock_exception_op - params: - raise_exception: true diff --git a/old/config/pydantic_config_parser.py b/old/config/pydantic_config_parser.py deleted file mode 100644 index 8670a45..0000000 --- a/old/config/pydantic_config_parser.py +++ /dev/null @@ -1,312 +0,0 @@ -import copy -import json -from pathlib import Path -from typing import Any, Generic, List, Type, TypeVar - -import yaml -from loguru import logger -from pydantic import BaseModel - -T = TypeVar("T", bound=BaseModel) - - -class PydanticConfigParser(Generic[T]): - current_file: str = __file__ - default_config_name: str = "default" - - """ - Pydantic Configuration Parser - - Supported configuration sources (priority from low to high): - 1. Default configuration (Pydantic model default values) - 2. YAML configuration file (supports import_config for importing multiple files) - 3. Command line arguments (dot notation format) - 4. Runtime parameters - - The import_config field supports importing multiple configuration files - by separating them with commas. Example: - import_config: "base.yaml, auth.yaml, database.yaml" - """ - - def __init__(self, config_class: Type[T]): - """ - Initialize configuration parser - - Args: - config_class: Pydantic configuration model class - """ - self.config_class = config_class - self.config_dict: dict = {} - - def parse_dot_notation(self, dot_list: List[str]) -> dict: - """ - Parse dot notation format configuration list - - Args: - dot_list: Configuration list in format ['a.b.c=value', 'x.y=123'] - - Returns: - Parsed nested dictionary - """ - config_dict = {} - - for item in dot_list: - if "=" not in item: - continue - - key_path, value_str = item.split("=", 1) - keys = key_path.split(".") - - # Automatic type conversion - value = self._convert_value(value_str) - - # Build nested dictionary - current_dict = config_dict - for key in keys[:-1]: - if key not in current_dict: - current_dict[key] = {} - current_dict = current_dict[key] - - current_dict[keys[-1]] = value - - return config_dict - - @staticmethod - def _convert_value(value_str: str) -> Any: - """ - Automatically convert string values to appropriate Python types - - Args: - value_str: String value - - Returns: - Converted value - """ - value_str = value_str.strip() - - if value_str.lower() in ("true", "false"): - return value_str.lower() == "true" - - if value_str.lower() in ("none", "null"): - return None - - try: - if "." not in value_str and "e" not in value_str.lower(): - return int(value_str) - - return float(value_str) - - except ValueError: - pass - - try: - return json.loads(value_str) - except (json.JSONDecodeError, ValueError): - pass - - return value_str - - @staticmethod - def load_from_yaml(yaml_path: str | Path) -> dict: - """ - Load configuration from YAML file - - Args: - yaml_path: YAML file path - - Returns: - Configuration dictionary - """ - if isinstance(yaml_path, str): - yaml_path = Path(yaml_path) - - if not yaml_path.exists(): - raise FileNotFoundError(f"Configuration file does not exist: {yaml_path}") - - with yaml_path.open() as f: - return yaml.safe_load(f) - - def merge_configs(self, *config_dicts: dict) -> dict: - """ - Deep merge multiple configuration dictionaries - - Args: - *config_dicts: Multiple configuration dictionaries - - Returns: - Merged configuration dictionary - """ - result = {} - - for config_dict in config_dicts: - result = self._deep_merge(result, config_dict) - - return result - - def _deep_merge(self, base_dict: dict, update_dict: dict) -> dict: - """ - Deep merge two dictionaries - - Args: - base_dict: Base dictionary - update_dict: Update dictionary - - Returns: - Merged dictionary - """ - result = base_dict.copy() - - for key, value in update_dict.items(): - if key in result and isinstance(result[key], dict) and isinstance(value, dict): - result[key] = self._deep_merge(result[key], value) - else: - result[key] = value - - return result - - def _load_import_configs_recursive(self, config_dict: dict, visited_configs: set) -> List[dict]: - """ - Recursively load import_config configurations with cycle detection - - Supports importing multiple config files by separating them with commas. - Example: import_config: "base.yaml, auth.yaml, database.yaml" - - Args: - config_dict: Current configuration dictionary - visited_configs: Set of already visited config file paths to detect cycles - - Returns: - List of imported configurations in correct merge order (base configs first) - """ - import_config = config_dict.get("import_config", "") - if not import_config: - return [] - - # Support comma-separated config files - import_config_list = [config.strip() for config in import_config.split(",") if config.strip()] - if not import_config_list: - return [] - - all_imported_configs = [] - - for single_import_config in import_config_list: - # Normalize config path - if not single_import_config.endswith(".yaml"): - single_import_config += ".yaml" - - # Resolve config file path - import_config_path = Path(self.current_file).parent / single_import_config - if not import_config_path.exists(): - import_config_path = Path(single_import_config) - - # Convert to absolute path for cycle detection - abs_config_path = import_config_path.resolve() - - # Check for circular dependency - if str(abs_config_path) in visited_configs: - # logger.warning(f"Circular import detected for config: {abs_config_path}") - continue - - logger.info(f"flowllm using import_config_path={import_config_path}") - - # Add current config to visited set - visited_configs.add(str(abs_config_path)) - - # Load the import config - import_yaml_config = self.load_from_yaml(import_config_path) - - # Recursively load imports from the imported config - nested_imports = self._load_import_configs_recursive(import_yaml_config, visited_configs.copy()) - - # Add configs in correct order: deeper imports first, then current import - # This ensures that configs closer to the root have higher priority - all_imported_configs.extend(nested_imports) - all_imported_configs.append(import_yaml_config) - - return all_imported_configs - - def parse_args(self, *args) -> T: - """ - Parse command line arguments and return configuration object - - Args: - args: Command line arguments. - - Returns: - Parsed configuration object - """ - configs_to_merge = [] - - # 1. Default configuration (from Pydantic model) - default_config = self.config_class().model_dump() - configs_to_merge.append(default_config) - - # 2. YAML configuration file - config = "" - filter_args = [] - for arg in args: - if "=" not in arg: - continue - - arg = arg.lstrip("--").lstrip("-") - - if "c=" in arg or "config=" in arg: - config = arg.split("=")[-1] - else: - filter_args.append(arg) - - if not config: - if self.default_config_name: - config = self.default_config_name - assert config, "add `config=` in cmd!" - - if not config.endswith(".yaml"): - config += ".yaml" - - # load pre-built configs - config_path = Path(self.current_file).parent / config - if not config_path.exists(): - config_path = Path(config) - logger.info(f"load config={config_path}") - - yaml_config = self.load_from_yaml(config_path) - - # load import configs recursively - imported_configs = self._load_import_configs_recursive(yaml_config, set()) - configs_to_merge.extend(imported_configs) - - configs_to_merge.append(yaml_config) - - # 3. Command line override configuration - if args: - cli_config = self.parse_dot_notation(filter_args) - configs_to_merge.append(cli_config) - - # Merge all configurations - self.config_dict = self.merge_configs(*configs_to_merge) - - # Create and validate final configuration object - return self.config_class.model_validate(self.config_dict) - - def update_config(self, **kwargs) -> T: - """ - Update configuration object using keyword arguments - - Args: - **kwargs: Configuration items to update, supports dot notation, e.g. a.b.c='xxx' - - Returns: - Updated configuration object - """ - # Convert kwargs to dot notation format - dot_list = [] - for key, value in kwargs.items(): - # support double underscore as dot replacement (server__host -> server.host) - dot_key = key.replace("__", ".") - dot_list.append(f"{dot_key}={value}") - - # Parse and merge configuration - override_config = self.parse_dot_notation(dot_list) - final_config = self.merge_configs(copy.deepcopy(self.config_dict), override_config) - - return self.config_class.model_validate(final_config) diff --git a/old/desktop_pet.py b/old/desktop_pet.py deleted file mode 100755 index 767005d..0000000 --- a/old/desktop_pet.py +++ /dev/null @@ -1,1123 +0,0 @@ -#!/usr/bin/env python3 -""" -macOS 桌面小宠物 - 超可爱版本 🐱✨ - -特性： -- 🎨 温暖奶茶色配色，更加可爱 -- ✨ 闪亮的琥珀色大眼睛，带渐变效果 -- 💕 粉嫩鼻子和小嘴，带可爱腮红 -- 🐾 渐变粉色肉垫，更多细节 -- 🌊 轻柔的呼吸动画（透明度变化） -- 👁️ 自动眨眼动画（每3.5秒） -- 🎯 尾巴自动摆动动画 -- 🎪 悬停时有反应（显示问号表情） -- 🦘 点击时弹跳动画 -- 💬 美化的渐变色对话气泡 -- 🏃 双击让猫咪奔跑 -- 🖱️ 拖拽移动、右键菜单 -- 🤖 AI智能对话 -""" -import sys -import os -import random -import asyncio -from typing import Optional -from PyQt6.QtWidgets import (QApplication, QLabel, QMenu, QWidget, - QTextEdit, QVBoxLayout, QPushButton, - QHBoxLayout, QScrollArea, QSizePolicy) -from PyQt6.QtCore import Qt, QTimer, QPoint, QPropertyAnimation, QEasingCurve, pyqtSignal, QThread, QSize -from PyQt6.QtGui import QPainter, QColor, QFont, QAction, QPainterPath, QRegion -from qasync import QEventLoop - -# 添加项目路径以便导入flowllm模块 -sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - -from flowllm.llm.openai_compatible_llm import OpenAICompatibleLLM -from flowllm.schema.message import Message -from loguru import logger - -class ResponseBubble(QWidget): - """AI回答气泡窗口""" - - closed = pyqtSignal() # 关闭信号 - - def __init__(self, pet_pos: QPoint, initial_text: str = ""): - super().__init__() - self.response_text = initial_text - self.init_ui(pet_pos) - - def init_ui(self, pet_pos: QPoint): - """初始化UI""" - # 窗口设置 - self.setWindowFlags( - Qt.WindowType.FramelessWindowHint | - Qt.WindowType.WindowStaysOnTopHint | - Qt.WindowType.Tool - ) - self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground) - - # 初始大小 - self.bubble_width = 350 - self.bubble_height = 150 - self.setFixedSize(self.bubble_width + 20, self.bubble_height + 20) - - # 位置在宠物右侧 - bubble_x = pet_pos.x() + 130 - bubble_y = pet_pos.y() - 50 - - # 确保不超出屏幕 - screen = QApplication.primaryScreen().geometry() - bubble_x = max(20, min(bubble_x, screen.width() - self.bubble_width - 40)) - bubble_y = max(20, min(bubble_y, screen.height() - self.bubble_height - 40)) - - self.move(bubble_x, bubble_y) - - # 创建主容器 - self.main_widget = QWidget(self) - self.main_widget.setGeometry(10, 10, self.bubble_width, self.bubble_height) - self.main_widget.setStyleSheet(""" - QWidget { - background: qlineargradient(x1:0, y1:0, x2:0, y2:1, - stop:0 rgba(245, 255, 245, 255), - stop:1 rgba(232, 245, 233, 250)); - border: 3px solid #66BB6A; - border-radius: 15px; - box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1); - } - """) - - layout = QVBoxLayout(self.main_widget) - layout.setContentsMargins(12, 10, 12, 10) - layout.setSpacing(5) - - # 标题栏（带关闭按钮） - title_layout = QHBoxLayout() - title_label = QLabel("✨ AI小助手") - title_label.setStyleSheet(""" - QLabel { - background-color: transparent; - color: #2E7D32; - font-size: 14px; - font-weight: bold; - border: none; - } - """) - title_layout.addWidget(title_label) - title_layout.addStretch() - - # 右上角关闭按钮 - close_button = QPushButton("✕") - close_button.setFixedSize(20, 20) - close_button.setStyleSheet(""" - QPushButton { - background-color: #81C784; - color: white; - border: none; - border-radius: 10px; - font-size: 14px; - font-weight: bold; - } - QPushButton:hover { - background-color: #66BB6A; - } - QPushButton:pressed { - background-color: #4CAF50; - } - """) - close_button.clicked.connect(self.close_bubble) - title_layout.addWidget(close_button) - - layout.addLayout(title_layout) - - # 创建滚动区域 - scroll_area = QScrollArea() - scroll_area.setWidgetResizable(True) - scroll_area.setStyleSheet(""" - QScrollArea { - background-color: transparent; - border: none; - } - QScrollBar:vertical { - border: none; - background: #E8F5E9; - width: 8px; - border-radius: 4px; - } - QScrollBar::handle:vertical { - background: #81C784; - border-radius: 4px; - } - """) - - # 创建内容标签（单个标签显示所有内容） - self.content_label = QLabel() - self.content_label.setWordWrap(True) - self.content_label.setTextFormat(Qt.TextFormat.RichText) - self.content_label.setAlignment(Qt.AlignmentFlag.AlignLeft | Qt.AlignmentFlag.AlignTop) - self.content_label.setStyleSheet(""" - QLabel { - background-color: transparent; - padding: 5px; - } - """) - - scroll_area.setWidget(self.content_label) - layout.addWidget(scroll_area, stretch=1) - - # 设置初始文本 - if self.response_text: - self.update_response('answer', self.response_text) - - def get_chunk_style(self, chunk_type: str) -> dict: - """获取不同chunk类型的样式配置""" - styles = { - 'answer': { - 'color': '#333333', - 'prefix': '' - }, - 'think': { - 'color': '#FF8C00', - 'prefix': '🤔 ' - }, - 'error': { - 'color': '#DC143C', - 'prefix': '❌ ' - }, - 'tool': { - 'color': '#1E90FF', - 'prefix': '🔧 ' - } - } - return styles.get(chunk_type, styles['answer']) - - def update_response(self, chunk_type: str, content: str): - """流式追加内容，根据类型使用不同颜色""" - # 只处理这四种类型，忽略其他类型（如usage） - if chunk_type not in ['answer', 'think', 'error', 'tool']: - return - - # 获取样式 - style = self.get_chunk_style(chunk_type) - - # 格式化新内容 - content_html = content.replace('\n', '
') - new_html = f'{style["prefix"]}{content_html}' - - # 追加到现有内容 - current_html = self.content_label.text() - self.content_label.setText(current_html + new_html) - - # 调整标签大小以适应内容 - self.content_label.adjustSize() - - # 计算新的气泡大小 - content_height = self.content_label.height() - new_height = min(max(150, content_height + 80), 500) # 最小150，最大500 - new_width = 350 - - if new_height != self.bubble_height or new_width != self.bubble_width: - self.bubble_width = new_width - self.bubble_height = new_height - - # 平滑调整大小 - self.setFixedSize(self.bubble_width + 20, self.bubble_height + 20) - self.main_widget.setGeometry(10, 10, self.bubble_width, self.bubble_height) - - def close_bubble(self): - """关闭气泡""" - self.closed.emit() - self.close() - - def paintEvent(self, event): - """绘制圆角阴影""" - painter = QPainter(self) - painter.setRenderHint(QPainter.RenderHint.Antialiasing) - - # 绘制阴影 - shadow_path = QPainterPath() - shadow_path.addRoundedRect(12, 12, self.bubble_width, self.bubble_height, 15, 15) - painter.fillPath(shadow_path, QColor(0, 0, 0, 30)) - - -class ChatBubble(QWidget): - """对话气泡窗口""" - - closed = pyqtSignal() # 关闭信号 - - def __init__(self, pet_pos: QPoint, llm_client: OpenAICompatibleLLM, pet_widget): - super().__init__() - self.llm_client = llm_client - self.pet_widget = pet_widget - self.init_ui(pet_pos) - self.is_streaming = False - self.response_bubble: Optional[ResponseBubble] = None - - def init_ui(self, pet_pos: QPoint): - """初始化UI""" - # 窗口设置 - self.setWindowFlags( - Qt.WindowType.FramelessWindowHint | - Qt.WindowType.WindowStaysOnTopHint | - Qt.WindowType.Tool - ) - self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground) - - # 设置大小和位置（缩小气泡） - self.setFixedSize(400, 200) - bubble_x = pet_pos.x() - 420 # 在宠物左侧 - bubble_y = pet_pos.y() - 150 # 在宠物上方 - - # 确保不超出屏幕 - screen = QApplication.primaryScreen().geometry() - bubble_x = max(20, min(bubble_x, screen.width() - 420)) - bubble_y = max(20, min(bubble_y, screen.height() - 220)) - - self.move(bubble_x, bubble_y) - - # 创建主容器 - main_widget = QWidget(self) - main_widget.setGeometry(10, 10, 380, 180) - main_widget.setStyleSheet(""" - QWidget { - background: qlineargradient(x1:0, y1:0, x2:0, y2:1, - stop:0 rgba(255, 250, 252, 255), - stop:1 rgba(255, 240, 245, 250)); - border: 3px solid #FF69B4; - border-radius: 20px; - box-shadow: 0 4px 8px rgba(255, 105, 180, 0.2); - } - """) - - layout = QVBoxLayout(main_widget) - layout.setContentsMargins(15, 15, 15, 15) - layout.setSpacing(10) - - # 标题栏（带关闭按钮） - title_layout = QHBoxLayout() - title_label = QLabel("💕 和我聊天吧~") - title_label.setStyleSheet(""" - QLabel { - background-color: transparent; - color: #E91E63; - font-size: 16px; - font-weight: bold; - border: none; - text-shadow: 1px 1px 2px rgba(233, 30, 99, 0.2); - } - """) - title_layout.addWidget(title_label) - title_layout.addStretch() - - # 右上角关闭按钮 - close_button = QPushButton("✕") - close_button.setFixedSize(24, 24) - close_button.setStyleSheet(""" - QPushButton { - background-color: #FFB6C1; - color: white; - border: none; - border-radius: 12px; - font-size: 16px; - font-weight: bold; - } - QPushButton:hover { - background-color: #FF69B4; - } - QPushButton:pressed { - background-color: #FF1493; - } - """) - close_button.clicked.connect(self.close_bubble) - title_layout.addWidget(close_button) - - layout.addLayout(title_layout) - - # 输入框 - self.input_text = QTextEdit() - self.input_text.setPlaceholderText("输入你想说的话...（按回车发送）") - self.input_text.setStyleSheet(""" - QTextEdit { - background-color: white; - border: 2px solid #FFB6C1; - border-radius: 10px; - padding: 8px; - font-size: 13px; - } - """) - layout.addWidget(self.input_text, stretch=1) - - # 支持回车发送 - self.input_text.installEventFilter(self) - - # 设置焦点到输入框 - self.input_text.setFocus() - - def eventFilter(self, obj, event): - """事件过滤器 - 处理回车发送""" - if obj == self.input_text and event.type() == event.Type.KeyPress: - if event.key() == Qt.Key.Key_Return or event.key() == Qt.Key.Key_Enter: - # Enter 发送 - self.send_message() - return True - return super().eventFilter(obj, event) - - def send_message(self): - """发送消息""" - message = self.input_text.toPlainText().strip() - if not message: - return - - if self.is_streaming: - return - - # 清空输入框并禁用 - self.input_text.clear() - self.input_text.setEnabled(False) - self.input_text.setPlaceholderText("思考中...") - - # 异步调用LLM - asyncio.create_task(self.call_llm(message)) - - async def call_llm(self, user_message: str): - """调用LLM并在新气泡中流式显示回答""" - self.is_streaming = True - - try: - # 小猫表情变为思考状态 - self.pet_widget.current_expression = "(=^･ω･^=)🤔" - self.pet_widget.update() - - # 创建回答气泡 - self.response_bubble = ResponseBubble(self.pet_widget.pos(), "") - self.response_bubble.closed.connect(self.on_response_bubble_closed) - self.response_bubble.show() - - # 创建消息列表 - messages = [ - Message(role="system", content="你是一个可爱的桌面宠物助手，用温暖友好的语气回答问题。回答要简洁但有帮助。"), - Message(role="user", content=user_message) - ] - - # 小猫表情变为回答状态 - self.pet_widget.current_expression = "(=^ω^=)💡" - self.pet_widget.update() - - # 流式接收并更新回答气泡 - async for chunk_content, chunk_type in self.llm_client.astream_chat(messages): - if chunk_content and self.response_bubble: - # 根据类型追加内容 - self.response_bubble.update_response(chunk_type, chunk_content) - # 让UI有时间更新 - await asyncio.sleep(0.01) - - # 回答完成，开心表情 - self.pet_widget.current_expression = "(=^ω^=)✨" - self.pet_widget.update() - - except Exception as e: - logger.exception(f"出错了: {str(e)}") - if self.response_bubble: - self.response_bubble.update_response('error', f"抱歉，出错了: {str(e)}") - # 错误表情 - self.pet_widget.current_expression = "(=ＴェＴ=)" - self.pet_widget.update() - - finally: - self.is_streaming = False - self.input_text.setEnabled(True) - self.input_text.setPlaceholderText("输入你想说的话...（按回车发送）") - self.input_text.setFocus() - - def on_response_bubble_closed(self): - """回答气泡关闭时的回调""" - self.response_bubble = None - - def close_bubble(self): - """关闭气泡""" - self.closed.emit() - self.close() - - def paintEvent(self, event): - """绘制圆角阴影""" - painter = QPainter(self) - painter.setRenderHint(QPainter.RenderHint.Antialiasing) - - # 绘制阴影 - shadow_path = QPainterPath() - shadow_path.addRoundedRect(12, 12, 380, 180, 20, 20) - painter.fillPath(shadow_path, QColor(0, 0, 0, 30)) - - -class DesktopPet(QLabel): - """桌面宠物类""" - - def __init__(self, llm_client: OpenAICompatibleLLM): - super().__init__() - self.llm_client = llm_client - self.chat_bubble: Optional[ChatBubble] = None - self.init_ui() - self.init_animation() - self.dragging = False - self.drag_position = QPoint() - self.drag_start_pos = QPoint() - - # 双击检测 - self.last_click_time = 0 - self.double_click_threshold = 300 # 毫秒 - - # 跑步状态 - self.is_running = False - self.run_timer = QTimer(self) - self.run_timer.timeout.connect(self.run_step) - self.run_direction = 1 # 1向右，-1向左 - - def init_ui(self): - """初始化UI""" - # 窗口设置 - self.setWindowFlags( - Qt.WindowType.FramelessWindowHint | # 无边框 - Qt.WindowType.WindowStaysOnTopHint | # 置顶 - Qt.WindowType.Tool # 不在任务栏显示 - ) - self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground) # 透明背景 - self.setAttribute(Qt.WidgetAttribute.WA_Hover) # 启用悬停事件 - - # 宠物大小 - self.pet_size = 150 - self.setFixedSize(self.pet_size, self.pet_size + 30) # 增加高度以容纳尾巴 - - # 初始位置（屏幕右下角） - screen = QApplication.primaryScreen().geometry() - self.move(screen.width() - 200, screen.height() - 200) - - # 宠物状态 - self.state = "idle" - self.expressions = ["(=^･ω･^=)", "(=^‥^=)", "(=^･ｪ･^=)", "(=ＴェＴ=)", "(=^ω^=)"] - self.current_expression = self.expressions[0] - self.is_hovered = False # 悬停状态 - - def init_animation(self): - """初始化动画""" - # 表情变化定时器（默认禁用） - self.expression_timer = QTimer(self) - self.expression_timer.timeout.connect(self.change_expression) - # self.expression_timer.start(3000) # 每3秒换表情 - 默认禁用 - - # 随机移动定时器（默认禁用） - self.move_timer = QTimer(self) - self.move_timer.timeout.connect(self.random_move) - # self.move_timer.start(8000) # 每8秒随机移动 - 默认禁用 - - # 动画效果 - self.animation = QPropertyAnimation(self, b"pos") - self.animation.setEasingCurve(QEasingCurve.Type.InOutCubic) - self.animation.setDuration(1500) - - # 自动行为开关 - self.auto_behavior_enabled = False - - # 呼吸动画（轻微缩放效果） - self.breath_animation = QPropertyAnimation(self, b"windowOpacity") - self.breath_animation.setDuration(2000) - self.breath_animation.setStartValue(0.95) - self.breath_animation.setEndValue(1.0) - self.breath_animation.setEasingCurve(QEasingCurve.Type.InOutSine) - self.breath_animation.setLoopCount(-1) # 无限循环 - self.breath_animation.start() - - # 眨眼动画计数器 - self.blink_timer = QTimer(self) - self.blink_timer.timeout.connect(self.blink) - self.blink_timer.start(3500) # 每3.5秒眨眼一次 - self.is_blinking = False - - # 尾巴摆动动画 - self.tail_angle = 0 - self.tail_timer = QTimer(self) - self.tail_timer.timeout.connect(self.wag_tail) - self.tail_timer.start(50) # 每50ms更新尾巴角度 - - def paintEvent(self, event): - """绘制宠物 - 超可爱长毛英短形象""" - painter = QPainter(self) - painter.setRenderHint(QPainter.RenderHint.Antialiasing) - from PyQt6.QtGui import QPolygon, QPainterPath, QPen, QRadialGradient - from PyQt6.QtCore import QPointF - - # 更温暖可爱的配色 - 奶茶色/米白色 - if not self.dragging: - body_base = QColor(220, 200, 180) # 奶茶色 - body_light = QColor(240, 230, 220) # 米白色 - body_dark = QColor(180, 160, 140) # 深奶茶色 - else: - body_base = QColor(230, 210, 190) - body_light = QColor(250, 240, 230) - body_dark = QColor(190, 170, 150) - - # === 绘制阴影 === - painter.setBrush(QColor(0, 0, 0, 40)) - painter.setPen(Qt.PenStyle.NoPen) - painter.drawEllipse(25, 152, 100, 25) - - # === 绘制蓬松的摆动尾巴（在身体后面，很粗） === - import math - # 根据尾巴角度计算摆动位置 - tail_offset = math.sin(self.tail_angle) * 8 - tail_path = QPainterPath() - tail_path.moveTo(QPointF(118, 125)) - tail_path.cubicTo( - QPointF(135 + tail_offset, 110), - QPointF(142 + tail_offset, 135), - QPointF(130 + tail_offset, 155) - ) - # 蓬松尾巴 - 画得很粗，带渐变 - tail_gradient = QRadialGradient(130 + tail_offset, 140, 15) - tail_gradient.setColorAt(0, body_light) - tail_gradient.setColorAt(1, body_base) - tail_pen = QPen(body_base, 20, Qt.PenStyle.SolidLine, Qt.PenCapStyle.RoundCap) - painter.setPen(tail_pen) - painter.drawPath(tail_path) - - # 尾巴尖端毛茸茸效果（随尾巴摆动） - painter.setPen(Qt.PenStyle.NoPen) - painter.setBrush(tail_gradient) - painter.drawEllipse(int(120 + tail_offset), 148, 18, 18) - painter.setBrush(body_light) - painter.drawEllipse(int(124 + tail_offset), 152, 14, 14) - - # === 绘制圆润的身体（长毛效果 - 多层） === - painter.setPen(Qt.PenStyle.NoPen) - - # 外层毛发（浅色） - painter.setBrush(body_light) - painter.drawEllipse(28, 92, 94, 72) - - # 中层身体 - painter.setBrush(body_base) - painter.drawEllipse(32, 95, 86, 66) - - # === 绘制圆圆的头部（长毛英短的大圆脸） === - # 外层毛茸茸效果 - painter.setBrush(body_light) - painter.drawEllipse(22, 28, 106, 95) - - # 脸部主体 - painter.setBrush(body_base) - painter.drawEllipse(28, 32, 94, 87) - - # 脸颊毛茸茸（英短特征） - painter.setBrush(body_light) - painter.drawEllipse(20, 65, 35, 28) # 左脸颊 - painter.drawEllipse(95, 65, 35, 28) # 右脸颊 - - # === 绘制圆耳朵（英短的小圆耳） === - painter.setBrush(body_base) - painter.setPen(Qt.PenStyle.NoPen) - - # 左耳 - 圆润的 - painter.drawEllipse(30, 25, 28, 30) - # 右耳 - painter.drawEllipse(92, 25, 28, 30) - - # 耳朵毛茸茸边缘 - painter.setBrush(body_light) - painter.drawEllipse(33, 28, 22, 24) - painter.drawEllipse(95, 28, 22, 24) - - # 耳朵内部（粉色） - painter.setBrush(QColor(255, 200, 210)) - painter.drawEllipse(37, 35, 14, 16) - painter.drawEllipse(99, 35, 14, 16) - - # === 绘制超萌大眼睛（闪亮亮的） === - # 眼窝阴影（更柔和） - painter.setBrush(body_dark) - painter.drawEllipse(44, 60, 26, 28) - painter.drawEllipse(80, 60, 26, 28) - - # 眼白（更大更圆）- 如果眨眼则缩小 - if self.is_blinking: - painter.setBrush(QColor(255, 255, 255)) - painter.drawEllipse(46, 73, 24, 6) # 眯眼效果 - painter.drawEllipse(80, 73, 24, 6) - else: - painter.setBrush(QColor(255, 255, 255)) - painter.drawEllipse(46, 62, 24, 26) - painter.drawEllipse(80, 62, 24, 26) - - # 眼珠（温暖的琥珀色 - 超可爱）- 只在睁眼时显示 - if not self.is_blinking: - if not self.dragging: - iris_gradient = QRadialGradient(58, 72, 10) - iris_gradient.setColorAt(0, QColor(255, 180, 100)) # 亮琥珀色 - iris_gradient.setColorAt(1, QColor(200, 130, 60)) # 深琥珀色 - painter.setBrush(iris_gradient) - else: - painter.setBrush(QColor(200, 180, 160)) # 温暖灰色 - - painter.drawEllipse(52, 68, 16, 18) - painter.drawEllipse(84, 68, 16, 18) - - # 瞳孔（竖线状 - 猫咪特征） - painter.setBrush(QColor(0, 0, 0)) - painter.drawEllipse(57, 72, 5, 12) - painter.drawEllipse(89, 72, 5, 12) - - # 超闪亮的眼睛高光（多层次） - painter.setBrush(QColor(255, 255, 255, 230)) - painter.drawEllipse(54, 70, 6, 7) # 主高光 - painter.drawEllipse(86, 70, 6, 7) - painter.setBrush(QColor(255, 255, 255, 150)) - painter.drawEllipse(60, 77, 3, 4) # 副高光 - painter.drawEllipse(92, 77, 3, 4) - - # === 绘制粉嫩小鼻子（超可爱） === - # 鼻子主体（渐变粉色） - nose_gradient = QRadialGradient(75, 96, 5) - nose_gradient.setColorAt(0, QColor(255, 180, 190)) # 浅粉 - nose_gradient.setColorAt(1, QColor(240, 140, 160)) # 深粉 - painter.setBrush(nose_gradient) - painter.setPen(Qt.PenStyle.NoPen) - nose_polygon = QPolygon([ - QPoint(75, 93), - QPoint(70, 98), - QPoint(80, 98) - ]) - painter.drawPolygon(nose_polygon) - - # 鼻子高光（更闪亮） - painter.setBrush(QColor(255, 220, 230, 200)) - painter.drawEllipse(73, 94, 3, 3) - - # === 绘制超萌微笑嘴巴（猫咪式微笑） === - painter.setPen(QPen(QColor(160, 100, 100), 2)) - painter.setBrush(Qt.BrushStyle.NoBrush) - mouth_path = QPainterPath() - mouth_path.moveTo(QPointF(75, 98)) - mouth_path.lineTo(QPointF(75, 103)) - # 可爱的微笑嘴角（更弯更可爱） - mouth_path.moveTo(QPointF(75, 103)) - mouth_path.quadTo(QPointF(68, 107), QPointF(63, 104)) - mouth_path.moveTo(QPointF(75, 103)) - mouth_path.quadTo(QPointF(82, 107), QPointF(87, 104)) - painter.drawPath(mouth_path) - - # 添加腮红（更可爱） - painter.setPen(Qt.PenStyle.NoPen) - painter.setBrush(QColor(255, 150, 150, 60)) - painter.drawEllipse(35, 85, 20, 15) # 左腮红 - painter.drawEllipse(95, 85, 20, 15) # 右腮红 - - # === 绘制胡须（细而长） === - whisker_pen = QPen(QColor(140, 140, 140), 1.5) - painter.setPen(whisker_pen) - # 左边胡须 - painter.drawLine(45, 88, 12, 82) - painter.drawLine(45, 92, 10, 92) - painter.drawLine(45, 96, 12, 102) - # 右边胡须 - painter.drawLine(105, 88, 138, 82) - painter.drawLine(105, 92, 140, 92) - painter.drawLine(105, 96, 138, 102) - - # === 绘制超可爱毛茸茸前爪 === - painter.setPen(Qt.PenStyle.NoPen) - - # 左爪（渐变效果） - left_paw_gradient = QRadialGradient(52, 159, 15) - left_paw_gradient.setColorAt(0, body_light) - left_paw_gradient.setColorAt(1, body_base) - painter.setBrush(left_paw_gradient) - painter.drawEllipse(38, 145, 28, 30) - - # 右爪（渐变效果） - right_paw_gradient = QRadialGradient(96, 159, 15) - right_paw_gradient.setColorAt(0, body_light) - right_paw_gradient.setColorAt(1, body_base) - painter.setBrush(right_paw_gradient) - painter.drawEllipse(84, 145, 28, 30) - - # 粉嫩肉垫（渐变粉色，更可爱） - paw_gradient = QRadialGradient(52, 162, 6) - paw_gradient.setColorAt(0, QColor(255, 200, 210)) - paw_gradient.setColorAt(1, QColor(255, 160, 180)) - painter.setBrush(paw_gradient) - painter.drawEllipse(47, 159, 10, 8) # 左爪主肉垫 - - paw_gradient2 = QRadialGradient(96, 162, 6) - paw_gradient2.setColorAt(0, QColor(255, 200, 210)) - paw_gradient2.setColorAt(1, QColor(255, 160, 180)) - painter.setBrush(paw_gradient2) - painter.drawEllipse(91, 159, 10, 8) # 右爪主肉垫 - - # 小肉垫（脚趾） - painter.setBrush(QColor(255, 180, 190)) - painter.drawEllipse(42, 157, 5, 5) # 左爪小肉垫 - painter.drawEllipse(58, 157, 5, 5) - painter.drawEllipse(86, 157, 5, 5) # 右爪小肉垫 - painter.drawEllipse(102, 157, 5, 5) - - # === 绘制胸前的白色毛（长毛英短特征） === - painter.setBrush(QColor(230, 235, 240, 180)) - painter.drawEllipse(58, 115, 34, 25) - - # 如果正在拖拽，显示提示 - if self.dragging: - painter.setPen(QColor(140, 150, 160)) - font = QFont("Arial", 11, QFont.Weight.Bold) - painter.setFont(font) - painter.drawText(self.rect().adjusted(0, -12, 0, 0), - Qt.AlignmentFlag.AlignCenter, "喵~ (=ＴェＴ=)") - - def blink(self): - """眨眼动画""" - if not self.dragging and not self.is_running: - self.is_blinking = True - self.update() - # 150毫秒后睁开眼睛 - QTimer.singleShot(150, self.open_eyes) - - def open_eyes(self): - """睁开眼睛""" - self.is_blinking = False - self.update() - - def wag_tail(self): - """摆动尾巴""" - import math - # 平滑的摆动效果 - self.tail_angle += 0.15 - if self.tail_angle > 2 * math.pi: - self.tail_angle = 0 - self.update() - - def change_expression(self): - """改变表情""" - if not self.dragging: - self.current_expression = random.choice(self.expressions) - self.update() - - def random_move(self): - """随机移动""" - if self.dragging or self.animation.state() == QPropertyAnimation.State.Running: - return - - screen = QApplication.primaryScreen().geometry() - - # 随机选择新位置（保持在屏幕内） - new_x = random.randint(50, screen.width() - self.width() - 50) - new_y = random.randint(50, screen.height() - self.height() - 50) - - # 执行动画移动 - self.animation.setStartValue(self.pos()) - self.animation.setEndValue(QPoint(new_x, new_y)) - self.animation.start() - - # 移动时换个开心的表情 - self.current_expression = "(=^ω^=)✨" - self.update() - - def mousePressEvent(self, event): - """鼠标按下事件""" - if event.button() == Qt.MouseButton.LeftButton: - self.drag_start_pos = event.globalPosition().toPoint() - self.drag_position = event.globalPosition().toPoint() - self.frameGeometry().topLeft() - event.accept() - elif event.button() == Qt.MouseButton.RightButton: - self.show_context_menu(event.globalPosition().toPoint()) - - def mouseMoveEvent(self, event): - """鼠标移动事件""" - if event.buttons() == Qt.MouseButton.LeftButton: - # 计算移动距离 - move_distance = (event.globalPosition().toPoint() - self.drag_start_pos).manhattanLength() - - # 如果移动超过5像素，认为是拖拽 - if move_distance > 5: - if not self.dragging: - self.dragging = True - self.current_expression = "(=ＴェＴ=)" - self.update() - - new_pos = event.globalPosition().toPoint() - self.drag_position - - # 限制在屏幕范围内 - screen = QApplication.primaryScreen().geometry() - new_pos.setX(max(0, min(new_pos.x(), screen.width() - self.width()))) - new_pos.setY(max(0, min(new_pos.y(), screen.height() - self.height()))) - - self.move(new_pos) - event.accept() - - def mouseReleaseEvent(self, event): - """鼠标释放事件""" - if event.button() == Qt.MouseButton.LeftButton: - # 如果没有拖拽（移动距离很小），认为是点击 - move_distance = (event.globalPosition().toPoint() - self.drag_start_pos).manhattanLength() - - if move_distance <= 5 and not self.dragging: - # 检测双击 - import time - current_time = time.time() * 1000 # 转换为毫秒 - if current_time - self.last_click_time < self.double_click_threshold: - # 双击 - 让猫跑起来 - self.handle_double_click() - else: - # 单击 - 打开聊天 - self.handle_click() - self.last_click_time = current_time - - if self.dragging: - self.dragging = False - self.current_expression = random.choice(["(=^･ω･^=)", "(=^ω^=)"]) - self.update() - - event.accept() - - def handle_click(self): - """处理点击事件""" - # 点击弹出对话框 - self.show_chat_bubble() - self.current_expression = "(=^･ω･^=)💬" - self.bounce_animation() # 添加弹跳动画 - self.update() - - def handle_double_click(self): - """处理双击事件 - 切换跑步状态""" - if self.is_running: - # 停止跑步 - self.stop_running() - else: - # 开始跑步 - self.start_running() - - def start_running(self): - """开始跑步""" - self.is_running = True - self.current_expression = "(=^ω^=)💨" - self.update() - - # 确定跑步方向（朝向最近的屏幕边缘相反方向） - screen = QApplication.primaryScreen().geometry() - center_x = self.pos().x() + self.width() / 2 - if center_x < screen.width() / 2: - self.run_direction = 1 # 向右跑 - else: - self.run_direction = -1 # 向左跑 - - # 开始跑步定时器（更快的更新频率） - self.run_timer.start(30) # 30毫秒更新一次 - - def stop_running(self): - """停止跑步""" - self.is_running = False - self.run_timer.stop() - self.current_expression = "(=^･ω･^=)" - self.update() - - def run_step(self): - """跑步的每一步""" - if not self.is_running: - return - - screen = QApplication.primaryScreen().geometry() - current_pos = self.pos() - - # 每步移动的距离 - step_size = 8 - new_x = current_pos.x() + (step_size * self.run_direction) - - # 检查是否到达屏幕边缘 - if new_x <= 0: - # 到达左边缘，转向右 - new_x = 0 - self.run_direction = 1 - elif new_x >= screen.width() - self.width(): - # 到达右边缘，转向左 - new_x = screen.width() - self.width() - self.run_direction = -1 - - # 移动猫咪 - self.move(new_x, current_pos.y()) - - # 跑步时的表情动画 - run_expressions = ["(=^ω^=)💨", "(=^･ω･^=)💨", "(=^ω^=)✨"] - self.current_expression = random.choice(run_expressions) - self.update() - - def show_chat_bubble(self): - """显示对话气泡""" - if self.chat_bubble is None or not self.chat_bubble.isVisible(): - self.chat_bubble = ChatBubble(self.pos(), self.llm_client, self) - self.chat_bubble.closed.connect(self.on_bubble_closed) - self.chat_bubble.show() - - def on_bubble_closed(self): - """对话气泡关闭时的回调""" - self.current_expression = random.choice(self.expressions) - self.update() - - def bounce_animation(self): - """可爱的弹跳动画""" - original_pos = self.pos() - - # 向上跳 - jump_height = 15 - self.move(original_pos.x(), original_pos.y() - jump_height) - - # 100毫秒后回到原位 - QTimer.singleShot(100, lambda: self.move(original_pos)) - - # 改变表情显得更可爱 - old_expr = self.current_expression - self.current_expression = "(=^ω^=)✨" - self.update() - QTimer.singleShot(200, lambda: setattr(self, 'current_expression', old_expr) or self.update()) - - def toggle_auto_behavior(self): - """切换自动行为开关""" - self.auto_behavior_enabled = not self.auto_behavior_enabled - - if self.auto_behavior_enabled: - # 启用自动行为 - self.expression_timer.start(3000) - self.move_timer.start(8000) - else: - # 禁用自动行为 - self.expression_timer.stop() - self.move_timer.stop() - - def show_context_menu(self, pos): - """显示右键菜单""" - menu = QMenu() - menu.setStyleSheet(""" - QMenu { - background-color: #FFF5F7; - border: 2px solid #FFB6C1; - border-radius: 8px; - padding: 5px; - } - QMenu::item { - padding: 5px 20px; - border-radius: 4px; - } - QMenu::item:selected { - background-color: #FFD6E0; - } - """) - - # 聊天选项 - chat_action = QAction("💬 和我聊天", self) - chat_action.triggered.connect(self.show_chat_bubble) - menu.addAction(chat_action) - - menu.addSeparator() - - # 自动行为开关 - auto_behavior_text = "🔕 禁用自动行为" if self.auto_behavior_enabled else "🔔 启用自动行为" - auto_behavior_action = QAction(auto_behavior_text, self) - auto_behavior_action.triggered.connect(self.toggle_auto_behavior) - menu.addAction(auto_behavior_action) - - menu.addSeparator() - - # 表情菜单 - expression_menu = menu.addMenu("切换表情 (=^･ω･^=)") - for expr in self.expressions: - action = QAction(expr, self) - action.triggered.connect(lambda checked, e=expr: self.set_expression(e)) - expression_menu.addAction(action) - - menu.addSeparator() - - # 回到角落 - corner_action = QAction("回到角落 🏠", self) - corner_action.triggered.connect(self.move_to_corner) - menu.addAction(corner_action) - - # 随机漫步（单次） - wander_action = QAction("随机漫步 🚶", self) - wander_action.triggered.connect(self.random_move) - menu.addAction(wander_action) - - # 跑步切换 - run_text = "停止跑步 🛑" if self.is_running else "开始跑步 🏃" - run_action = QAction(run_text, self) - run_action.triggered.connect(self.handle_double_click) - menu.addAction(run_action) - - menu.addSeparator() - - # 退出 - exit_action = QAction("再见~ 👋", self) - exit_action.triggered.connect(self.close_app) - menu.addAction(exit_action) - - menu.exec(pos) - - def set_expression(self, expression): - """设置表情""" - self.current_expression = expression - self.update() - - def move_to_corner(self): - """移动到角落""" - screen = QApplication.primaryScreen().geometry() - target_pos = QPoint(screen.width() - 200, screen.height() - 200) - - self.animation.setStartValue(self.pos()) - self.animation.setEndValue(target_pos) - self.animation.start() - - def close_app(self): - """关闭应用""" - self.current_expression = "(=ＴェＴ=)" - self.update() - QTimer.singleShot(500, QApplication.quit) - - def event(self, event): - """处理悬停事件""" - if event.type() == event.Type.HoverEnter: - self.is_hovered = True - if not self.dragging and not self.is_running: - self.current_expression = "(=^･ω･^=)?" - self.update() - elif event.type() == event.Type.HoverLeave: - self.is_hovered = False - if not self.dragging and not self.is_running: - self.current_expression = self.expressions[0] - self.update() - return super().event(event) - - -def main(): - """主函数""" - # 使用qasync支持异步 - app = QApplication(sys.argv) - loop = QEventLoop(app) - asyncio.set_event_loop(loop) - - app.setQuitOnLastWindowClosed(True) - - # 从环境变量读取LLM配置 - api_key = os.getenv("FLOW_LLM_API_KEY", "") - api_base = os.getenv("FLOW_LLM_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1") - model = os.getenv("FLOW_LLM_MODEL", "qwen3-max") - - # 初始化LLM客户端 - llm_client = OpenAICompatibleLLM( - api_key=api_key, - base_url=api_base, - model_name=model - ) - - pet = DesktopPet(llm_client) - pet.show() - - with loop: - loop.run_forever() - - -if __name__ == "__main__": - main() - diff --git a/old/docs/INSTALLATION.md b/old/docs/INSTALLATION.md deleted file mode 100644 index f072a6e..0000000 --- a/old/docs/INSTALLATION.md +++ /dev/null @@ -1,103 +0,0 @@ -# 🛠️ Installation Guide - -## Install from PyPI (Recommended) - -```bash -pip install flowllm -``` - -## Install with Optional Dependencies - -```bash -# For financial data support -pip install flowllm[fin] - -# For distributed computing -pip install flowllm[dist] - -# For Qwen model support -pip install flowllm[qwen] - -# Install all features -pip install flowllm[all] -``` - -## Install from Source - -```bash -git clone https://github.com/your-org/flowllm.git -cd flowllm -pip install . -``` - -## Environment Configuration - -Copy `example.env` to `.env` and modify the corresponding parameters: - -```bash -FLOW_APP_NAME=FlowLLM - -# LLM Configuration -FLOW_LLM_API_KEY=sk-xxxx -FLOW_LLM_BASE_URL=https://xxxx/v1 -FLOW_LLM_MODEL=qwen3-max - -# Embedding Configuration -FLOW_EMBEDDING_API_KEY=sk-xxxx -FLOW_EMBEDDING_BASE_URL=https://xxxx/v1 - -# Optional: Elasticsearch -# FLOW_ES_HOSTS=http://0.0.0.0:9200 - -# Optional: Additional API Keys -# FLOW_DASHSCOPE_API_KEY=sk-xxxx -# FLOW_TAVILY_API_KEY=sk-xxxx -``` - -## Verify Installation - -After installation, you can verify that FlowLLM is installed correctly: - -```bash -# Check version -python -c "import flowllm; print(flowllm.__version__)" - -# Run help command -flowllm --help -``` - -## Troubleshooting - -### Common Issues - -#### 1. Import Error - -If you encounter import errors, make sure you have installed all required dependencies: - -```bash -pip install -r requirements.txt -``` - -#### 2. Environment Variables Not Found - -Ensure your `.env` file is in the correct location (project root) or export environment variables manually: - -```bash -export FLOW_LLM_API_KEY=your_api_key -export FLOW_LLM_BASE_URL=your_base_url -``` - -#### 3. Port Already in Use - -If the default port (8002) is already in use, specify a different port: - -```bash -flowllm backend=http http.port=8003 -``` - -## Next Steps - -- Read the [Quick Start Guide](QUICKSTART.md) to get started -- Check out the [README](README.md) for architecture overview -- Explore the [documentation](docs/) for detailed guides - diff --git a/old/docs/QUICKSTART.md b/old/docs/QUICKSTART.md deleted file mode 100644 index f0319dc..0000000 --- a/old/docs/QUICKSTART.md +++ /dev/null @@ -1,232 +0,0 @@ -# 🚀 Quick Start Guide - -## 30-Second Setup - -```bash -# Install -pip install flowllm - -# Set environment variables -export FLOW_LLM_API_KEY=sk-xxx -export FLOW_LLM_BASE_URL=https://api.openai.com/v1 - -# Start HTTP service (includes pre-built flows) -flowllm backend=http http.port=8002 - -# Or start MCP service -flowllm backend=mcp mcp.transport=stdio -``` - -That's it! Pre-built flows are now available as HTTP endpoints or MCP tools. - -## Using Pre-built Services - -```bash -# Use built-in LLM flow -curl -X POST http://localhost:8002/llm_flow \ - -H "Content-Type: application/json" \ - -d '{"messages": [{"role": "user", "content": "Hello!"}]}' - -# With streaming -curl -X POST http://localhost:8002/llm_flow_stream \ - -H "Content-Type: application/json" \ - -d '{"messages": [{"role": "user", "content": "Tell me a story"}]}' -``` - -## Creating Your First Custom Service - -### Step 1: Write an Operation (Python) - -```python -# my_ops.py -from flowllm import BaseOp, C - -class SummarizeOp(BaseOp): - """Summarize text using LLM""" - def call(self, text: str, max_length: int = 100) -> dict: - llm = C.get_llm("default") - prompt = f"Summarize this in {max_length} words: {text}" - summary = llm.call(messages=[{"role": "user", "content": prompt}]) - return {"summary": summary} -``` - -### Step 2: Configure as Service (YAML) - -```yaml -# my_config.yaml -import_config: default - -backend: http -http: - port: 8002 - -op: - summarize_op: - backend: base_op - class: SummarizeOp - -flow: - summarize: - description: "Summarize long text into concise format" - flow_content: summarize_op - tool: - parameters: - text: - type: string - description: "Text to summarize" - required: true - max_length: - type: integer - description: "Maximum summary length in words" - required: false - default: 100 -``` - -### Step 3: Launch - -```bash -flowllm config=my_config -``` - -### Step 4: Use - -```bash -# HTTP endpoint automatically created -curl -X POST http://localhost:8002/summarize \ - -H "Content-Type: application/json" \ - -d '{ - "text": "Long article text here...", - "max_length": 50 - }' -``` - -**No FastAPI routes, no manual validation, no extra code - just configuration!** - -## Using HTTP Client - -### Python Requests - -```python -import requests - -# Execute a flow -response = requests.post("http://localhost:8002/my_custom_flow", json={ - "input_text": "Hello, FlowLLM!" -}) - -print(response.json()) -``` - -### curl - -```bash -curl -X POST http://localhost:8002/my_custom_flow \ - -H "Content-Type: application/json" \ - -d '{ - "input_text": "Hello, FlowLLM!" - }' -``` - -### Node.js - -```javascript -fetch("http://localhost:8002/my_custom_flow", { - method: "POST", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify({ - input_text: "Hello, FlowLLM!" - }) -}) -.then(response => response.json()) -.then(data => console.log(data)); -``` - -## Using Python Client - -### Synchronous Execution - -```python -from flowllm import FlowLLMApp - -# Initialize app -with FlowLLMApp(load_default_config=True) as app: - # Execute flow synchronously - result = app.execute_flow( - "my_custom_flow", - input_text="Hello, FlowLLM!" - ) - print(result) -``` - -### Async Usage - -```python -import asyncio -from flowllm import FlowLLMApp - -async def main(): - async with FlowLLMApp(load_default_config=True) as app: - # Execute flow asynchronously - result = await app.async_execute_flow( - "my_custom_flow", - input_text="Hello, FlowLLM!" - ) - print(result) - -asyncio.run(main()) -``` - -## Next Steps - -- **Learn More**: Check out the [Configuration-Driven Development](README.md#-configuration-driven-development) section -- **Explore Examples**: See the [Example Workflows](README.md#-example-workflows) section -- **Advanced Features**: Read about [Core Features](README.md#-core-features) -- **Deep Dive**: Browse the [documentation](docs/) directory for detailed guides -- **Build**: Follow the complete example in [Configuration-Driven Development](README.md#-configuration-driven-development) - -## Common Patterns - -### Simple LLM Chat - -```python -from flowllm import BaseOp, FlowLLMApp, C - -class ChatOp(BaseOp): - def call(self, message: str) -> dict: - llm = C.get_llm("default") - response = llm.call(messages=[ - {"role": "user", "content": message} - ]) - return {"response": response} - -with FlowLLMApp(load_default_config=True) as app: - result = app.execute_flow("chat", message="Hello!") - print(result) -``` - -### Sequential Workflow - -```yaml -flow: - research_flow: - description: "Multi-step research workflow" - flow_content: "search_op >> summarize_op >> validate_op" -``` - -### Parallel Processing - -```yaml -flow: - parallel_analysis: - description: "Analyze from multiple perspectives" - flow_content: "(sentiment_op | keywords_op | summary_op)" -``` - -## Getting Help - -- **Documentation**: [doc/](docs/) directory -- **Issues**: [GitHub Issues](https://github.com/your-org/flowllm/issues) -- **Email**: jinli.yl@alibaba-inc.com - diff --git a/old/docs/README.md b/old/docs/README.md deleted file mode 100644 index 9321344..0000000 --- a/old/docs/README.md +++ /dev/null @@ -1,215 +0,0 @@ -

- FlowLLM: A Flexible Framework for Building LLM-Powered Flows
- Flow with Intelligence, Build with Simplicity. -

- -

- Python Version - PyPI Version - License -

- ---- - -FlowLLM is a **configuration-driven** framework for building LLM-powered applications. Write operations once, compose them via YAML configuration, and automatically get HTTP APIs and MCP tools—no boilerplate code needed. - -## 📖 Table of Contents - -- [Why FlowLLM?](#-why-flowllm) -- [Getting Started](#-getting-started) -- [Core Workflow](#-core-workflow) -- [Architecture](#-architecture) -- [Features](#-features) -- [Resources](#-resources) - ---- - -## 💡 Why FlowLLM? - -**The Problem**: Building LLM services traditionally requires writing boilerplate routes, validation, documentation, and orchestration code for each endpoint. - -**The Solution**: FlowLLM's configuration-driven approach lets you: - -- ✅ **Write Operations Once** - Focus on business logic in reusable Python ops -- ✅ **Configure, Don't Code** - Compose workflows using YAML configuration -- ✅ **Auto-Generate Services** - HTTP and MCP endpoints created automatically -- ✅ **Built-in Orchestration** - Sequential (`>>`), parallel (`|`), and nested flows -- ✅ **Zero Boilerplate** - No routes, validators, or service code needed - -| Feature | Traditional Approach | FlowLLM Approach | -|---------|---------------------|------------------| -| **Service Creation** | Write FastAPI/Flask routes, handlers, validation | Write YAML config - auto-registers HTTP + MCP | -| **API Documentation** | Manually write OpenAPI specs | Auto-generated from config | -| **Workflow Changes** | Modify Python code, test, redeploy | Update config, restart service | -| **Orchestration** | Write custom coordination code | Use expressions: `>>`, `\|`, `()` | - -**Perfect For**: Rapid prototyping, microservices, AI agent tools, data pipelines, enterprise AI applications. - ---- - -## 🚀 Getting Started - -### Installation - -```bash -pip install flowllm -``` - -For detailed setup instructions, see the [Installation Guide](INSTALLATION.md). - -### Quick Start - -See the [Quick Start Guide](QUICKSTART.md) to build your first LLM service in 30 seconds. - ---- - -## 🎯 Core Workflow - -``` -┌─────────────────┐ ┌──────────────────┐ ┌──────────────────────────┐ -│ Build Ops │ │ Configure YAML │ │ Auto-Register │ -│ (Python) │ → │ (Workflows) │ → │ Services │ -│ │ │ │ │ │ -│ • BaseOp │ │ flow: │ │ ┌────────────────────┐ │ -│ • BaseAsyncOp │ │ workflow: │ │ │ HTTP Service │ │ -│ • BaseMcpOp │ │ description │ │ │ POST /workflow │ │ -│ • BaseRayOp │ │ flow_content│ │ │ OpenAPI docs │ │ -│ │ │ tool: │ │ └────────────────────┘ │ -│ │ │ parameters│ │ │ -│ │ │ │ │ ┌────────────────────┐ │ -│ │ │ backend: http │ │ │ MCP Service │ │ -│ │ │ or mcp │ │ │ Tool: workflow │ │ -│ │ │ │ │ │ Auto-discovered │ │ -│ │ │ │ │ └────────────────────┘ │ -└─────────────────┘ └──────────────────┘ └──────────────────────────┘ -``` - -**Three Simple Steps:** - -1. **Create an Op** - Write a Python class implementing your business logic -2. **Configure in YAML** - Define workflow and service endpoints -3. **Launch** - Run one command to start your HTTP or MCP service - -**No manual routing, no endpoint definitions, no service code - just configuration!** - ---- - -## ✨ Architecture - -FlowLLM adopts a **three-layer configuration-driven architecture**: - -``` -┌─────────────────────────────────────────────────────────────────────┐ -│ Service Layer (外层) │ -│ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │ -│ │ HTTP Service │ │ MCP Service │ │ CMD Service │ │ -│ │ FastAPI │ │ FastMCP │ │ Command Line│ │ -│ └──────────────┘ └──────────────┘ └──────────────┘ │ -│ Auto-Register from Configuration │ -└─────────────────────────────┬───────────────────────────────────────┘ - │ -┌─────────────────────────────┴───────────────────────────────────────┐ -│ Flow Layer (中层) │ -│ • Sequential: op1 >> op2 >> op3 │ -│ • Parallel: (op1 | op2 | op3) │ -│ • Nested: op1 >> (op2 | op3) >> op4 │ -└─────────────────────────────┬───────────────────────────────────────┘ - │ -┌─────────────────────────────┴───────────────────────────────────────┐ -│ Foundation Layer (底层) │ -│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ │ -│ │ Op Lib │ │ LLM Lib │ │Embedding │ │ Storage │ │ -│ ├──────────┤ ├──────────┤ ├──────────┤ ├──────────┤ │ -│ │ BaseOp │ │ OpenAI │ │ OpenAI │ │ElasticS. │ │ -│ │BaseAsync │ │ LiteLLM │ │Compatible│ │ChromaDB │ │ -│ │BaseTool │ │DashScope │ │ │ │ Local │ │ -│ │BaseMcpOp │ │ Custom │ │ Custom │ │ Cache │ │ -│ │BaseRayOp │ │ │ │ │ │ │ │ -│ └──────────┘ └──────────┘ └──────────┘ └──────────┘ │ -└──────────────────────────────────────────────────────────────────────┘ -``` - -### Key Design Principles - -1. **Separation of Concerns** - Ops (business logic), Flows (orchestration), Services (protocol handling) -2. **Configuration over Code** - Ops in Python, Flows in YAML, Services auto-generated -3. **Dependency Injection** - ServiceContext manages shared resources (LLM, VectorStore, etc.) -4. **Registry Pattern** - Dynamic loading and discovery based on configuration - -### Complete Data Flow - -``` -Request → Service Layer (HTTP/MCP) - ↓ -Flow Layer (Parse expression → Build DAG) - ↓ -Foundation Layer (Execute ops with context) - ↓ -Response (JSON/MCP result) -``` - ---- - -## 🎯 Features - -### 📦 Pre-built Operations - -**Gallery Ops**: `SimpleLLMOp`, `ReactLLMOp`, `ExecuteCodeOp`, `TranslateCodeOp` -**Search Ops**: `TavilySearchOp`, `DashScopeSearchOp`, `McpSearchOp` -**Research Ops**: `DashScopeDeepResearchOp`, `LangChainDeepResearchOp` -**Data Ops**: Various extraction and processing operations - -### 🔧 Advanced Capabilities - -- **Multi-LLM Support** - OpenAI, LiteLLM (100+ providers), DashScope, custom providers -- **Vector Storage** - Elasticsearch, ChromaDB, local file-based, in-memory -- **Async/Streaming** - Full async support with SSE streaming responses -- **Distributed Computing** - Ray integration for scaling operations -- **Caching** - Intelligent caching with TTL and automatic serialization -- **Web Crawling** - Integrated `crawl4ai` for content extraction - -### 🧪 Workflow Patterns - -- **Simple LLM Chat** - Direct model interaction -- **Multi-Step Research** - Sequential search, summarization, validation -- **Parallel Processing** - Concurrent sentiment analysis, keyword extraction -- **Complex Pipelines** - Nested sequential and parallel operations - ---- - -## 📚 Resources - -### Documentation - -- **[Installation Guide](INSTALLATION.md)** - Setup and environment configuration -- **[Quick Start Guide](QUICKSTART.md)** - Build your first service -- **Specialized Guides** in `doc/`: - - [Deep Research Guide](docs/deep_research.md) - - [Financial Supply Guide](docs/fin_supply_readme.md) - - [Vector Store Guide](docs/vector_store.md) - -### Examples & Configuration - -- **Examples**: `test/` directory for practical examples -- **Configuration**: `flowllm/config/` for sample configs - -### Latest Updates - -- **[2025-10]** FlowLLM v0.1.10 - Enhanced async support and stability -- **[2025-09]** Financial data modules with 26+ pre-built flows -- **[2025-09]** Deep research with multiple search backends -- **[2025-08]** MCP (Model Context Protocol) support -- **[2025-06]** Multi-backend vector storage - ---- - - -## ⚖️ License - -Apache License 2.0 - see [LICENSE](./LICENSE) file for details. - ---- - -## 🌟 Star History - -If you find FlowLLM useful, please consider giving it a star! diff --git a/old/docs/async_task_handling.md b/old/docs/async_task_handling.md deleted file mode 100644 index 2b6eb61..0000000 --- a/old/docs/async_task_handling.md +++ /dev/null @@ -1,244 +0,0 @@ -# 异步任务处理指南 - -## 问题背景 - -在使用 `BaseAsyncOp.submit_async_task()` 和 `join_async_task()` 时，如果提交的任务中有异常或网络超时，可能会导致程序卡死。 - -### 原因分析 - -1. **异常传播延迟**：通过 `create_task()` 创建的任务如果抛出异常，不会立即传播，而是在 `await asyncio.gather()` 时才抛出 -2. **任务无法取消**：默认的 `gather()` 在遇到异常时不会自动取消其他任务 -3. **网络超时问题**：如果任务中有网络请求卡住（如 SSE 连接错误），没有超时机制会导致永久等待 - -## 解决方案 - -### 1. 改进的 `join_async_task` 方法 - -新的 `join_async_task` 方法支持以下功能： - -```python -async def join_async_task(self, timeout: float = None, return_exceptions: bool = False): - """ - 等待所有异步任务完成 - - Args: - timeout: 超时时间（秒），None表示无限等待 - return_exceptions: 如果为True，异常会作为结果返回；如果为False，第一个异常会被抛出 - - Returns: - 所有任务的结果列表 - """ -``` - -#### 关键改进： - -1. **超时保护**：设置 `timeout` 参数可以避免任务永久等待 -2. **异常取消机制**：当任何任务失败或超时时，自动取消所有其他任务 -3. **灵活的异常处理**：通过 `return_exceptions` 参数控制异常处理方式 - -### 2. 使用示例 - -#### 场景1：严格模式（遇到异常立即停止） - -```python -class MyOp(BaseAsyncOp): - async def async_execute(self): - # 提交多个任务 - for query in search_queries: - search_op = self.ops[0].copy() - self.submit_async_task(search_op.async_call, query=query) - - # 如果任何一个任务失败，立即抛出异常并取消其他任务 - await self.join_async_task() -``` - -#### 场景2：容错模式（收集所有成功的结果） - -```python -class MyOp(BaseAsyncOp): - async def async_execute(self): - # 提交多个任务 - for query in search_queries: - search_op = self.ops[0].copy() - self.submit_async_task(search_op.async_call, query=query) - - # 即使部分任务失败，也继续处理成功的结果 - results = await self.join_async_task(return_exceptions=True) - # results 中会包含成功的结果和异常对象 -``` - -#### 场景3：带超时的任务（推荐） - -```python -class MyOp(BaseAsyncOp): - async def async_execute(self): - # 提交多个任务 - for query in search_queries: - search_op = self.ops[0].copy() - self.submit_async_task(search_op.async_call, query=query) - - try: - # 设置总超时时间为 60 秒 - results = await self.join_async_task(timeout=60.0) - except asyncio.TimeoutError: - logger.error("Tasks timeout, some queries may have failed") - # 所有任务已被自动取消 -``` - -#### 场景4：容错 + 超时（最灵活） - -```python -class MyOp(BaseAsyncOp): - async def async_execute(self): - # 提交多个任务 - for query in search_queries: - search_op = self.ops[0].copy() - self.submit_async_task(search_op.async_call, query=query) - - try: - # 最多等待60秒，部分失败也继续 - results = await self.join_async_task(timeout=60.0, return_exceptions=True) - - # 分离成功和失败的结果 - success_results = [r for r in results if not isinstance(r, Exception)] - failed_results = [r for r in results if isinstance(r, Exception)] - - # 异常会自动通过 logger.exception 打印，这里只统计数量 - logger.info(f"Success: {len(success_results)}, Failed: {len(failed_results)}") - - except asyncio.TimeoutError: - logger.error("Tasks timeout") -``` - -### 3. MCP 客户端的超时配置 - -对于 MCP 调用，建议在两个层面设置超时： - -#### 3.1 MCP 客户端级别 - -```python -class BaseMcpOp(BaseAsyncToolOp): - def __init__(self, - mcp_name: str = "", - timeout: float = 30.0, # MCP 调用超时 - **kwargs): - self.timeout = timeout - super().__init__(**kwargs) - - async def async_execute(self): - async with McpClient(name=self.mcp_name, - config=config, - timeout=self.timeout) as client: - result = await client.call_tool(self.tool_name, arguments=self.input_dict) -``` - -#### 3.2 任务汇总级别 - -```python -class CompanySegmentFactorOp(BaseAsyncToolOp): - async def _execute_searches(self, search_queries: List[str]) -> str: - for query in search_queries: - search_op = self.ops[0].copy() - self.submit_async_task(search_op.async_call, query=query) - - # 设置总超时，避免所有搜索任务卡死 - await self.join_async_task(timeout=120.0, return_exceptions=True) -``` - -## 日志输出 - -`join_async_task` 使用 `logger.exception` 记录异常信息，会自动打印完整的堆栈跟踪： - -### return_exceptions=False（默认） - -```python -# 任何任务失败时，会打印异常堆栈并抛出 -await self.join_async_task(timeout=60.0) -``` - -日志输出示例： -``` -ERROR | join_async_task failed with SSEError, cancelling remaining tasks... -Traceback (most recent call last): - File "/path/to/mcp/client/sse.py", line 72, in sse_reader - ... -httpx_sse._exceptions.SSEError: Expected response header Content-Type to contain 'text/event-stream', got '' -``` - -### return_exceptions=True（容错模式） - -```python -# 每个失败的任务都会单独记录 -results = await self.join_async_task(timeout=60.0, return_exceptions=True) -``` - -日志输出示例： -``` -ERROR | Task failed with exception -Traceback (most recent call last): - ... -httpx_sse._exceptions.SSEError: ... - -ERROR | Task failed with exception -Traceback (most recent call last): - ... -asyncio.TimeoutError: ... -``` - -## 最佳实践 - -1. **总是设置超时**：对于网络请求或外部服务调用，一定要设置合理的超时时间 -2. **考虑容错**：如果任务之间相对独立，使用 `return_exceptions=True` 提高鲁棒性 -3. **日志记录**：`logger.exception` 会自动记录完整的异常堆栈，便于排查问题 -4. **分层超时**： - - 单个 MCP 调用：30-60秒 - - 批量任务汇总：根据任务数量和单个超时计算，留有余地 -5. **优雅降级**：即使部分任务失败，也能返回部分结果 - -## 常见错误场景 - -### 错误1：SSE 连接错误导致卡死 - -```python -# 错误：没有超时保护 -await self.join_async_task() - -# 正确：添加超时 -await self.join_async_task(timeout=60.0) -``` - -### 错误2：一个任务失败导致所有结果丢失 - -```python -# 错误：任何一个失败就抛异常 -results = await self.join_async_task() - -# 正确：收集所有成功的结果 -results = await self.join_async_task(return_exceptions=True) -success_results = [r for r in results if not isinstance(r, Exception)] -``` - -### 错误3：没有清理失败的任务 - -改进后的 `join_async_task` 会自动处理，不需要手动清理。 - -## 调试建议 - -如果遇到任务卡死问题： - -1. **添加详细日志**：在 `submit_async_task` 前后添加日志 -2. **检查超时设置**：确保 MCP 客户端和任务汇总都设置了超时 -3. **使用 return_exceptions**：先让程序运行起来，收集异常信息 -4. **逐个测试**：将并发任务改为串行执行，定位问题任务 - -```python -# 调试模式：串行执行，定位问题 -for query in search_queries: - try: - search_op = self.ops[0].copy() - result = await search_op.async_call(query=query) - logger.info(f"Query '{query}' success") - except Exception as e: - logger.error(f"Query '{query}' failed: {e}") -``` - diff --git a/old/docs/company_factor_op_implementation.md b/old/docs/company_factor_op_implementation.md deleted file mode 100644 index 435def7..0000000 --- a/old/docs/company_factor_op_implementation.md +++ /dev/null @@ -1,172 +0,0 @@ -# CompanyFactorOp 实现说明 - -## 概述 -`CompanyFactorOp` 是一个用于构建公司整体估值金融因子逻辑图的操作类，它综合分析公司的多个业务板块，生成统一的因子传导路径图。 - -## 核心功能 - -### 输入参数 -- `name`: 公司名称 -- `code`: 股票代码 - -### 输出结果 -返回一个包含以下字段的 JSON 对象： -```json -{ - "name": "公司名称", - "code": "股票代码", - "segments": ["板块1", "板块2", "板块3"], - "meta_list": ["Meta信息1", "Meta信息2", ...], - "factor_graph": "mermaid图代码" -} -``` - -## 执行流程 - -### Step 1: 获取业务板块 -- 调用 `CompanyOperationOp` 获取公司的主营业务板块 -- 过滤规则： - - 营收或利润至少有一项 ≥ 5% - - 按重要性排序（revenue + profit 降序） - - 最多保留 3 个板块 - -### Step 2: 分析各业务板块 -- 对每个保留的业务板块，串行调用 `CompanySegmentFactorOp` -- 每个板块返回： - - `meta_list`: 该板块的Meta信息列表 - - `mermaid_graph`: 该板块的因子传导图 - -### Step 3: 合并Meta信息 -- 收集所有板块的Meta信息 -- 使用LLM进行智能合并： - - 去除重复信息 - - 合并同义表述 - - 解决信息冲突 - - 按重要性排序 - -### Step 4: 融合因子逻辑图 -- 收集所有板块的因子传导图 -- 使用LLM进行图融合： - - 统一收敛点：`A[公司名 估值增加]` - - 提取和合并相同/相似因子 - - 优化传导路径（最多5条，每条≤4个节点） - - 确保所有因子有Meta信息支撑 - -### Step 5: 验证并修正Mermaid代码 -- 检查语法正确性 -- 验证结构规范性 -- 修复常见错误 -- 确保可正常渲染 - -### Step 6: 输出最终结果 -- 构建包含所有信息的结果字典 -- 以JSON格式输出 - -## CompanySegmentFactorOp 的改进 - -为了支持 `CompanyFactorOp`，对 `CompanySegmentFactorOp` 进行了改进： - -### 返回值变更 -- **修改前**：只返回 `mermaid_graph` 字符串 -- **修改后**：返回包含 `meta_list` 和 `mermaid_graph` 的JSON对象 - -### 主要改动 -1. **初始化变更**：`meta_list` 从字符串改为空列表 `[]` -2. **类型调整**：所有内部方法的 `meta_list` 参数类型从 `str` 改为 `List` -3. **JSON序列化**：在prompt中使用时，通过 `json.dumps()` 将列表转为字符串 -4. **返回格式**：最终返回包含两个字段的字典，并序列化为JSON字符串 - -### 向后兼容性 -- 这个改动会影响现有的使用 `CompanySegmentFactorOp` 的代码 -- 如需继续单独使用该Op且只获取mermaid图，可通过 `json.loads(op.output)["mermaid_graph"]` 提取 - -## 提示词设计 - -### 复用的提示词（来自CompanySegmentFactorOp） -- `init_mermaid_graph`: 初始化mermaid图 -- `factor_step1_prompt`: 生成搜索查询 -- `factor_step2_prompt`: 更新Meta信息列表 -- `factor_step3_prompt`: 更新因子传导图 - -### 新增的提示词 - -#### merge_meta_prompt -- 功能：合并多个板块的Meta信息 -- 要求：去重、同义合并、冲突解决、按重要性排序 - -#### merge_graphs_prompt -- 功能：融合多个板块的因子传导图 -- 要求： - - 统一收敛点 - - 因子去重和归类 - - 路径优化（最多5条，每条≤4节点） - - 依据约束（所有因子必须有Meta支撑） - -#### validate_mermaid_prompt -- 功能：验证并修正Mermaid代码 -- 检查项： - - 语法正确性（graph LR; 开头，箭头格式等） - - 结构规范性（收敛点、无循环、无孤立节点） - - 节点命名规范 - - 常见错误修复 - -## 配置参数 - -### CompanyFactorOp 参数 -- `llm`: LLM模型（默认 "qwen3_30b_instruct"） -- `revenue_threshold`: 营收阈值（默认 0.05，即5%） -- `profit_threshold`: 利润阈值（默认 0.05，即5%） -- `max_segments`: 最多保留的板块数（默认 3） - -### CompanySegmentFactorOp 参数 -- `llm`: LLM模型（默认 "qwen3_30b_instruct"） -- `max_steps`: 最大迭代步数（默认 3） -- `max_search_cnt`: 每次最多搜索查询数（默认 3） -- `max_approach_cnt`: 最多传导路径数（默认 3） - -## 使用示例 - -```python -from flowllm.app import FlowLLMApp -from flowllm.op.fin_search.company_factor_op import CompanyFactorOp - -async with FlowLLMApp(args=["config=fin_research"]): - op = CompanyFactorOp() - await op.async_call(name="阿里巴巴", code="09988") - print(op.output) -``` - -## 输出示例 - -```json -{ - "name": "阿里巴巴", - "code": "09988", - "segments": ["云计算", "电商", "AI"], - "meta_list": [ - "政府推出AI产业扶持政策", - "云计算市场需求增长30%", - "电商平台GMV同比增长15%", - "阿里云Q3营收增长28%" - ], - "factor_graph": "graph LR;\nAI需求增 --> 云服务收入增\n政策支持 --> 云服务收入增\n云服务收入增 --> A[阿里巴巴 估值增加]\n电商增长 --> A[阿里巴巴 估值增加]" -} -``` - -## 技术亮点 - -1. **模块化设计**：通过组合现有Op（CompanyOperationOp、CompanySegmentFactorOpV2）实现复杂功能 -2. **智能过滤**：基于营收和利润指标自动筛选核心业务板块 -3. **LLM驱动**：利用LLM的理解能力进行Meta信息合并和因子图融合 -4. **串行处理**：确保每个板块分析完整且可追溯 -5. **自动验证**：内置Mermaid代码验证和修正机制 -6. **结构化输出**：统一的JSON格式便于后续处理 - -## 注意事项 - -1. **执行时间**：由于需要串行处理多个板块，整体执行时间较长 -2. **依赖关系**：需要确保 `CompanyOperationOp` 和 `TongyiMcpSearchOp` 正常工作 -3. **网络搜索**：依赖外部搜索服务，可能受网络状况影响 -4. **LLM质量**：最终结果质量依赖于LLM的推理能力 -5. **配置调优**：可根据实际需求调整阈值和最大数量参数 - diff --git a/old/docs/deep_research.md b/old/docs/deep_research.md deleted file mode 100644 index ccd20a5..0000000 --- a/old/docs/deep_research.md +++ /dev/null @@ -1,59 +0,0 @@ -# 🔬 Deep Research Guide - -FlowLLM 提供多种深度研究能力，支持不同的搜索后端和研究模式。 - -## 📊 可用的研究流程 -在flowllm的框架内迁移了olangchain的pen_deep_research，对prompt和实现细节做了一些优化，修复了qwen3-max下会出现复读机的问题。 - -| 流程名称 | 描述 | 特点 | -|----------------------------------|-------------------------|------------------| -| `dashscope_deep_research` | DashScope 原生深度研究 | 快速研究，内置搜索优化 | -| `langchain+brief+bailian_search` | LangChain + 百炼搜索 + 研究摘要 | 全面研究，话题转换，详细报告 | -| `langchain+brief+bocha_search` | LangChain + 博查搜索 + 研究摘要 | 全面研究，替代搜索引擎，详细报告 | -| `langchain+bailian_search` | LangChain + 百炼搜索 | 直接研究，流程简化，执行快速 | -| `langchain+bocha_search` | LangChain + 博查搜索 | 直接研究，替代搜索源，执行快速 | - -## 💻 使用方法 - -以下是各个研究流程的调用方式： - - -**1. DashScope 深度研究** - -```bash -curl -X POST http://11.164.204.33:8002/dashscope_deep_research -H "Content-Type: application/json" -d '{"query": "什么是人工智能？"}' -``` - -**2. LangChain + 研究摘要 + 百炼搜索** - -```bash -curl -X POST http://11.164.204.33:8002/langchain+brief+bailian_search -H "Content-Type: application/json" -d '{"query": "分析2025年电动汽车的竞争格局"}' -``` - -**3. LangChain + 研究摘要 + 博查搜索** - -```bash -curl -X POST http://11.164.204.33:8002/langchain+brief+bocha_search -H "Content-Type: application/json" -d '{"query": "量子计算的最新发展是什么？"}' -``` - -**4. LangChain + 百炼搜索** - -```bash -curl -X POST http://11.164.204.33:8002/langchain+bailian_search -H "Content-Type: application/json" -d '{"query": "比较全球可再生能源的采用率"}' -``` - -**5. LangChain + 博查搜索** - -```bash -curl -X POST http://11.164.204.33:8002/langchain+bocha_search -H "Content-Type: application/json" -d '{"query": "区块链技术的当前趋势是什么？"}' -``` - -**6. query合成** - -```bash -curl -X POST http://11.164.204.33:8002/task_react -H "Content-Type: application/json" -d '{"items": ["", ""]}' -``` - -items: list of input_topic(str), e.g. ["", ""] or ["行业研究","行业研究","个股分析"]，不写就是在这几个中随机 - -增加`exist_list`参数告诉模型之前产出了哪些query，避免重复，`{"items": ["", ""], "exist_list": ["xx1", "xx2"]}` \ No newline at end of file diff --git a/old/docs/fin_supply_deploy.md b/old/docs/fin_supply_deploy.md deleted file mode 100644 index edfed11..0000000 --- a/old/docs/fin_supply_deploy.md +++ /dev/null @@ -1,49 +0,0 @@ -## 快速开始 - -### 环境准备 - -#### 1. 克隆项目 -```shell -git clone git@gitlab.alibaba-inc.com:OpenRepo/flowllm.git -``` - -#### 2. 准备环境变量 -从 `example.env` 复制一份为 `.env`，填入自己的API密钥和URL： -```shell -cp example.env .env -# 编辑.env文件，填入相关配置 -``` - -#### 3. 安装依赖 -要求 Python >= 3.12 -```shell -pip install .[fin] -``` - -### 启动服务 - -#### SSE模式启动 -```shell -flowllm --config=fin_supply -``` - -#### MCP Server配置（stdio模式） - -配置文件示例： -```json -{ - "mcpServers": { - "asio-fin-supply-server": { - "command": "flowllm", - "args": [ - "--config=fin_supply", - "--mcp.transport=stdio" - ], - "env": { - "FLOW_EMBEDDING_API_KEY": "sk-xxxx", - "FLOW_EMBEDDING_BASE_URL": "xxxx" - } - } - } -} -``` diff --git a/old/docs/fin_supply_quick_start.md b/old/docs/fin_supply_quick_start.md deleted file mode 100644 index 3b941c2..0000000 --- a/old/docs/fin_supply_quick_start.md +++ /dev/null @@ -1,248 +0,0 @@ -### MCP API调用示例 - -标准的MCP服务使用方法，参考 [mcp_supply_test.py](../test/mcp_supply_test.py) - -## MCP工具的HTTP API调用示例 - -以下是每个MCP工具对应的curl命令示例，部署在8002端口，用于HTTP API测试： - -### 1. 搜索引擎工具 - -```bash -# tavily_search - Tavily搜索 -curl -X POST http://11.164.204.33:8002/tavily_search \ - -H "Content-Type: application/json" \ - -d '{ - "query": "紫金怎么样？雪球" - }' - -# dashscope_search - DashScope搜索 -curl -X POST http://11.164.204.33:8002/dashscope_search \ - -H "Content-Type: application/json" \ - -d '{ - "query": "紫金怎么样？雪球" - }' - -# brave_web_search - Brave网页搜索 -curl -X POST http://11.164.204.33:8002/brave_web_search \ - -H "Content-Type: application/json" \ - -d '{ - "query": "紫金怎么样？雪球" - }' - -# bailian_web_search - 百炼网页搜索 -curl -X POST http://11.164.204.33:8002/bailian_web_search \ - -H "Content-Type: application/json" \ - -d '{ - "query": "紫金怎么样？雪球" - }' - -# bocha_web_search - 博茶网页搜索 -curl -X POST http://11.164.204.33:8002/bocha_web_search \ - -H "Content-Type: application/json" \ - -d '{ - "query": "紫金怎么样？雪球" - }' -``` - -### 2. 网页解析工具 - -```bash -# bailian_web_parser - 百炼网页解析 -curl -X POST http://11.164.204.33:8002/bailian_web_parser \ - -H "Content-Type: application/json" \ - -d '{ - "url": "https://basic.10jqka.com.cn/601899/field.html" - }' - -# 解析公司信息页面 -curl -X POST http://11.164.204.33:8002/bailian_web_parser \ - -H "Content-Type: application/json" \ - -d '{ - "url": "https://basic.10jqka.com.cn/601899/company.html" - }' -``` - -### 3. 实体提取工具 - -```bash -# extract_entities_code - 实体代码提取 -curl -X POST http://11.164.204.33:8002/extract_entities_code \ - -H "Content-Type: application/json" \ - -d '{ - "query": "阿里和腾讯哪个更好？和茅台比呢？" - }' -``` - -### 4. AKShare市场数据工具 - -```bash -# akshare_market - AKShare市场数据 -curl -X POST http://11.164.204.33:8002/akshare_market \ - -H "Content-Type: application/json" \ - -d '{ - "code": "000001" - }' - -# akshare_calculate - AKShare计算分析 -curl -X POST http://11.164.204.33:8002/akshare_calculate \ - -H "Content-Type: application/json" \ - -d '{ - "query": "平安银行最近行情走势", - "code": "000001" - }' -``` - -### 5. QTF量化工具 - -```bash -# qtf_brief_mcp - QTF简要信息 -curl -X POST http://11.164.204.33:8002/qtf_brief_mcp \ - -H "Content-Type: application/json" \ - -d '{ - "symbol": "SZ000001" - }' - -# qtf_medium_mcp - QTF中等详细信息 -curl -X POST http://11.164.204.33:8002/qtf_medium_mcp \ - -H "Content-Type: application/json" \ - -d '{ - "symbol": "SZ000001" - }' - -# qtf_full_mcp - QTF完整信息 -curl -X POST http://11.164.204.33:8002/qtf_full_mcp \ - -H "Content-Type: application/json" \ - -d '{ - "symbol": "SZ000001" - }' -``` - -### 6. 同花顺(THS)数据爬取工具 - -```bash -# crawl_ths_company - 爬取公司基本信息 -curl -X POST http://11.164.204.33:8002/crawl_ths_company \ - -H "Content-Type: application/json" \ - -d '{ - "query": "平安银行的公司情况如何？", - "code": "000001" - }' - -# crawl_ths_holder - 爬取股东结构信息 -curl -X POST http://11.164.204.33:8002/crawl_ths_holder \ - -H "Content-Type: application/json" \ - -d '{ - "query": "平安银行的股东结构如何？", - "code": "000001" - }' - -# crawl_ths_operate - 爬取主营业务信息 -curl -X POST http://11.164.204.33:8002/crawl_ths_operate \ - -H "Content-Type: application/json" \ - -d '{ - "query": "平安银行的主营业务是什么？", - "code": "000001" - }' - -# crawl_ths_equity - 爬取股本结构信息 -curl -X POST http://11.164.204.33:8002/crawl_ths_equity \ - -H "Content-Type: application/json" \ - -d '{ - "query": "平安银行的股本结构如何？", - "code": "000001" - }' - -# crawl_ths_capital - 爬取资本运作信息 -curl -X POST http://11.164.204.33:8002/crawl_ths_capital \ - -H "Content-Type: application/json" \ - -d '{ - "query": "平安银行有哪些资本运作？", - "code": "000001" - }' - -# crawl_ths_worth - 爬取业绩预测信息 -curl -X POST http://11.164.204.33:8002/crawl_ths_worth \ - -H "Content-Type: application/json" \ - -d '{ - "query": "平安银行的业绩预测如何？", - "code": "000001" - }' - -# crawl_ths_news - 爬取最新新闻 -curl -X POST http://11.164.204.33:8002/crawl_ths_news \ - -H "Content-Type: application/json" \ - -d '{ - "query": "平安银行最近有什么新闻？", - "code": "000001" - }' - -# crawl_ths_concept - 爬取概念题材信息 -curl -X POST http://11.164.204.33:8002/crawl_ths_concept \ - -H "Content-Type: application/json" \ - -d '{ - "query": "平安银行涉及哪些概念题材？", - "code": "000001" - }' - -# crawl_ths_position - 爬取机构持仓信息 -curl -X POST http://11.164.204.33:8002/crawl_ths_position \ - -H "Content-Type: application/json" \ - -d '{ - "query": "平安银行的机构持仓情况如何？", - "code": "000001" - }' - -# crawl_ths_finance - 爬取财务状况信息 -curl -X POST http://11.164.204.33:8002/crawl_ths_finance \ - -H "Content-Type: application/json" \ - -d '{ - "query": "平安银行的财务状况如何？", - "code": "000001" - }' - -# crawl_ths_bonus - 爬取分红情况信息 -curl -X POST http://11.164.204.33:8002/crawl_ths_bonus \ - -H "Content-Type: application/json" \ - -d '{ - "query": "平安银行的分红情况如何？", - "code": "000001" - }' - -# crawl_ths_event - 爬取重大事件信息 -curl -X POST http://11.164.204.33:8002/crawl_ths_event \ - -H "Content-Type: application/json" \ - -d '{ - "query": "平安银行最近有什么重大事件？", - "code": "000001" - }' - -# crawl_ths_field - 爬取行业地位信息 -curl -X POST http://11.164.204.33:8002/crawl_ths_field \ - -H "Content-Type: application/json" \ - -d '{ - "query": "平安银行在行业中的地位如何？", - "code": "000001" - }' -``` - -### 7. 其他股票代码示例 - -```bash -# 紫金矿业公司信息 -curl -X POST http://11.164.204.33:8002/crawl_ths_company \ - -H "Content-Type: application/json" \ - -d '{ - "query": "紫金矿业的公司情况如何？", - "code": "601899" - }' - -# 紫金矿业财务指标 -curl -X POST http://11.164.204.33:8002/crawl_ths_finance \ - -H "Content-Type: application/json" \ - -d '{ - "query": "紫金矿业的财务指标怎么样？", - "code": "601899" - }' -``` - diff --git a/old/docs/fin_supply_readme.md b/old/docs/fin_supply_readme.md deleted file mode 100644 index ef52ec8..0000000 --- a/old/docs/fin_supply_readme.md +++ /dev/null @@ -1,412 +0,0 @@ -## 背景 - -### 系统架构 - -采用分层架构设计： - -1. **操作层（Operation Layer）**：提供原子化的数据操作能力，每个Op（操作）负责特定的数据获取或处理任务 -2. **流程层（Flow Layer）**：通过Pipeline将多个Op组合成复杂的业务流程，支持串行（>>）和并行（|）执行和子Op（<<） -3. **服务层（Service Layer）**：基于MCP（Model Context Protocol）协议提供标准化的API接口 - -### 核心特性 - -- **模块化设计**：每个供给模块独立可配置，支持按需组合和扩展 -- **统一接口**：基于MCP协议提供标准化的工具调用接口，便于集成 -- **智能缓存**：支持多级缓存机制，提高数据获取效率和降低成本 -- **灵活配置**：通过YAML配置文件管理流程参数，支持环境隔离和版本管理 - -### 供给模块 - -系统目前提供以下金融供给模块，共计26个Flow： - -| MCP Flow名称 | 核心能力 | -|---------------------------|--------------------| -| **extract_entities_code** | 金融实体识别和code提取 | -| **akshare_market** | A股实时行情数据 | -| **akshare_calculate** | 股票技术分析计算 | -| **qtf_brief_mcp** | 股票基础分析摘要 | -| **qtf_medium_mcp** | 股票中等深度分析 | -| **qtf_full_mcp** | 股票全面深度分析 | -| **tavily_search** | tavily_search | -| **dashscope_search** | dashscope大模型强制开启搜索 | -| **bailian_web_search** | bailian_search | -| **bocha_web_search** | bocha search | -| **brave_web_search** | brave search | -| **bailian_web_parser** | bailian 网页解析 | -| **crawl_ths_company** | 同花顺公司基本资料 | -| **crawl_ths_holder** | 同花顺股东研究信息 | -| **crawl_ths_operate** | 同花顺经营分析信息 | -| **crawl_ths_equity** | 同花顺股本结构信息 | -| **crawl_ths_capital** | 同花顺资本运作信息 | -| **crawl_ths_worth** | 同花顺盈利预测信息 | -| **crawl_ths_news** | 同花顺新闻公告信息 | -| **crawl_ths_concept** | 同花顺概念题材信息 | -| **crawl_ths_position** | 同花顺主力持仓信息 | -| **crawl_ths_finance** | 同花顺财务分析信息 | -| **crawl_ths_bonus** | 同花顺分红融资信息 | -| **crawl_ths_event** | 同花顺公司大事信息 | -| **crawl_ths_field** | 同花顺行业对比信息 | - -### 供给类型统计 - -| 分类 | 供给名称 | 输入示例 | 输出示例 | -|------------|--------------------------------------------|------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| **通用搜索解析** | tavily_search（Tavily搜索） | `{"query": "2024年中国GDP增长情况"}` | `根据最新数据，2024年中国GDP同比增长5.2%，经济运行总体平稳，消费和投资成为主要增长动力...` | -| | dashscope_search（dashscope调用enable_search） | `{"query": "人工智能最新发展趋势"}` | `人工智能领域最新发展包括大语言模型技术突破、多模态AI应用普及、AI+行业深度融合等趋势...` | -| | bailian_web_search（百炼搜索） ✅ | `{"query": "新能源汽车市场现状", "count": 10}` | `2024年新能源汽车市场保持高速增长，销量突破800万辆，市场渗透率达到35%，比亚迪、特斯拉领跑市场...` | -| | bocha_web_search（Bocha搜索） | `{"query": "区块链技术应用", "freshness": 7}` | `区块链技术在金融、供应链、数字身份等领域应用不断深化，央行数字货币DCEP试点扩大，DeFi生态持续发展...` | -| | brave_web_search（Brave搜索） | `{"query": "量子计算突破", "count": 5}` | `量子计算领域取得重大突破，IBM发布1000+量子比特处理器，Google实现量子纠错里程碑，产业化进程加速...` | -| | bailian_web_parser（网页解析） ✅ | `{"url": "https://example.com"}` | `解析网页内容：标题、正文、关键信息等结构化数据，支持多种格式输出...` | -| **股票提槽** | extract_entities_code（提取实体&代号）✅ | `{"query": "阿里和腾讯哪个更好？和茅台比呢？"}` | `[{"entity": "阿里巴巴", "type": "stock", "codes": ["BABA", "09988.HK"]}, {"entity": "腾讯", "type": "stock", "codes": ["00700.HK"]}, {"entity": "茅台", "type": "stock", "codes": ["600519"]}]` | -| **股票基本面** | akshare_market（实时行情） | `{"code": "000001"}` | `000001的实时行情: {"代码": "000001", "名称": "平安银行", "最新价": 15.23, "涨跌幅": 2.35, "涨跌额": 0.35, "成交量": 123456, "成交额": 1876543210, "振幅": 3.45, "最高": 15.45, "最低": 14.98, "今开": 15.12, "昨收": 14.88}` | -| | crawl_ths_company（公司基本资料） | `{"query": "详细情况", "code": "000001"}` | `平安银行详细情况: 公司成立于1987年，注册资本1000亿元，主营业务为商业银行业务...` | -| | crawl_ths_holder（股东研究） | `{"query": "十大股东", "code": "000001"}` | `平安银行十大股东: 1.中国平安保险(集团)股份有限公司 58.03%, 2.香港中央结算有限公司 8.12%...` | -| | crawl_ths_operate（经营分析） ✅ | `{"query": "主营业务构成", "code": "000001"}` | `平安银行主营业务构成: 利息净收入占比78.5%，手续费及佣金净收入占比15.2%...` | -| | crawl_ths_equity（股本结构） | `{"query": "股本构成", "code": "000001"}` | `平安银行股本构成: 总股本194.4亿股，其中A股194.4亿股，流通股194.4亿股...` | -| | crawl_ths_capital（资本运作） | `{"query": "募集资金", "code": "000001"}` | `平安银行募集资金情况: 近期通过定向增发募集资金260亿元，主要用于补充核心一级资本...` | -| | crawl_ths_finance（财务分析）✅ | `{"query": "财务指标", "code": "000001"}` | `平安银行财务指标: ROE 11.1%，ROA 0.89%，资本充足率13.85%，核心一级资本充足率10.15%...` | -| | crawl_ths_bonus（分红融资） | `{"query": "分红情况", "code": "000001"}` | `平安银行分红情况: 2023年每10股派现2.31元，股息率约1.5%，连续多年稳定分红...` | -| | crawl_ths_field（行业对比） ✅ | `{"query": "行业地位", "code": "000001"}` | `平安银行行业地位: 在股份制银行中排名第3位，资产规模4.9万亿元，市场份额约2.8%...` | -| **股票技术面** | akshare_calculate（技术分析） | `{"code": "000001", "query": "最近五日成交量有放量吗？最近五日macd有金叉吗？"}` | `根据分析，该股票最近五日成交量相比前期有明显放量，平均成交量增长35%。MACD指标在最近五日出现金叉信号，DIF线上穿DEA线，显示短期趋势向好。` | -| | qtf_brief_mcp（简要分析） | `{"symbol": "SZ000001"}` | `平安银行(000001)简要分析: 当前股价15.23元，涨幅2.35%。基本面良好，ROE 11.1%，市盈率8.5倍。技术面显示短期趋势向好，建议关注。` | -| | qtf_medium_mcp（中等分析） | `{"symbol": "SZ000001"}` | `平安银行(000001)中等分析: 基本面分析-资产质量稳定，不良率1.02%；盈利能力强，ROE持续改善。技术面分析-MACD金叉，成交量放大，短期支撑位14.8元。估值合理，目标价16.5元。风险提示：银行业监管政策变化。` | -| | qtf_full_mcp（全面分析） ✅ | `{"symbol": "SZ000001"}` | `平安银行(000001)全面分析报告: 1.基本面：资产规模4.9万亿，净利润增长6.8%，ROE 11.1%，资本充足率13.85%。2.技术面：多项指标显示买入信号，支撑阻力位明确。3.估值：PB 0.85倍，PE 8.5倍，相对估值偏低。4.投资建议：买入，目标价16.5-17元，持有期6-12个月。` | -| | crawl_ths_worth（盈利预测） ✅ | `{"query": "业绩预测", "code": "000001"}` | `平安银行业绩预测: 2024年预测净利润增长8.5%，ROE预计达到11.2%，机构一致目标价16.5元...` | -| | crawl_ths_position（主力持仓） ✅ | `{"query": "机构持股", "code": "000001"}` | `平安银行机构持股: 基金持股占比12.5%，QFII持股占比3.2%，社保基金持股占比1.8%...` | -| **股票消息面** | crawl_ths_news（新闻公告） ✅ | `{"query": "最新公告", "code": "000001"}` | `平安银行最新公告: 发布2024年第三季度业绩报告，净利润同比增长6.8%，不良贷款率1.02%...` | -| | crawl_ths_concept（概念题材） ✅ | `{"query": "概念板块", "code": "000001"}` | `平安银行概念板块: 银行概念、金融科技、数字货币、区块链、人工智能等概念...` | -| | crawl_ths_event（公司大事） | `{"query": "高管变动", "code": "000001"}` | `平安银行高管变动: 近期董事会决议任命新任副行长，具有20年银行从业经验...` | - -### 技术优势 - -- **高性能**：支持异步处理和并行执行，可处理大规模数据请求 -- **高可用**：内置重试机制和异常处理，确保服务稳定性 -- **易扩展**：基于插件化架构，可快速添加新的数据源和分析能力 -- **标准化**：遵循MCP协议标准，与主流AI平台和工具无缝集成 - -## basic supply - -基础金融供给提供了核心的金融数据获取和分析能力，主要包含以下流程： - -| MCP Flow名称 | 能力描述 | Pipeline | 输入示例 | 输出示例 | -|-----------------------|-------------------|---------------------------------------------------|------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| extract_entities_code | 从查询中提取金融实体并获取对应代码 | extract_entities_code_op << bailian_web_search_op | `{"query": "阿里和腾讯哪个更好？和茅台比呢？"}` | `[{"entity": "阿里巴巴", "type": "stock", "codes": ["BABA", "09988.HK"]}, {"entity": "腾讯", "type": "stock", "codes": ["00700.HK"]}, {"entity": "茅台", "type": "stock", "codes": ["600519"]}]` | -| akshare_market | 获取A股股票实时行情数据 | akshare_market_op | `{"code": "000001"}` | `000001的实时行情: {"代码": "000001", "名称": "平安银行", "最新价": 15.23, "涨跌幅": 2.35, "涨跌额": 0.35, "成交量": 123456, "成交额": 1876543210, "振幅": 3.45, "最高": 15.45, "最低": 14.98, "今开": 15.12, "昨收": 14.88}` | -| akshare_calculate | 基于历史数据进行股票技术分析计算 | akshare_calculate_op | `{"code": "000001", "query": "最近五日成交量有放量吗？最近五日macd有金叉吗？"}` | `根据分析，该股票最近五日成交量相比前期有明显放量，平均成交量增长35%。MACD指标在最近五日出现金叉信号，DIF线上穿DEA线，显示短期趋势向好。` | - -### 代码示例 - -```python -import asyncio -from flowllm.client.fastmcp_client import FastmcpClient - - -async def basic_supply_example(): - """Basic supply使用示例""" - async with FastmcpClient(transport="sse", host="11.164.204.33", port=8001) as client: - # 1. 提取金融实体代码 - print("=== 提取金融实体代码 ===") - result = await client.call_tool("extract_entities_code", { - "query": "阿里和腾讯哪个更好？和茅台比呢？" - }) - print(f"提取结果: {result.content[0].text}") - - # 2. 获取实时行情 - print("\n=== 获取实时行情 ===") - result = await client.call_tool("akshare_market", { - "code": "000001" - }) - print(f"行情数据: {result.content[0].text}") - - # 3. 技术分析计算 - print("\n=== 技术分析计算 ===") - result = await client.call_tool("akshare_calculate", { - "code": "000001", - "query": "最近五日成交量有放量吗？最近五日macd有金叉吗？" - }) - print(f"分析结果: {result.content[0].text}") - - -if __name__ == "__main__": - asyncio.run(basic_supply_example()) -``` - -### 主要特点 - -- **实体识别**：智能识别查询中的金融实体（股票、基金、ETF等）并自动获取交易代码 -- **实时行情**：提供A股市场实时价格、成交量等关键指标 -- **技术分析**：基于历史数据进行各类技术指标计算和趋势分析 -- **缓存机制**：支持数据缓存以提高查询效率 -- **多数据源**：集成akshare和网络搜索能力 - -### 参考地址 - -- [akshare](https://github.com/akfamily/akshare) - -### 配置地址 - -- [fin_basic.yaml](../flowllm/config/fin_basic.yaml) - -## ths supply - -同花顺供给提供了全面的股票深度分析能力，通过爬取同花顺网站获取详细的公司信息和财务数据： - -| MCP Flow名称 | 能力描述 | Pipeline | 输入示例 | 输出示例 | -|--------------------|------------|------------------------------------------------------------------|-----------------------------------------|----------------------------------------------------------------| -| crawl_ths_company | 获取公司基本资料信息 | ths_company_op >> bailian_web_parser_op >> extract_long_text_op | `{"query": "详细情况", "code": "000001"}` | `平安银行详细情况: 公司成立于1987年，注册资本1000亿元，主营业务为商业银行业务...` | -| crawl_ths_holder | 获取股东研究信息 | ths_holder_op >> bailian_web_parser_op >> extract_long_text_op | `{"query": "十大股东", "code": "000001"}` | `平安银行十大股东: 1.中国平安保险(集团)股份有限公司 58.03%, 2.香港中央结算有限公司 8.12%...` | -| crawl_ths_operate | 获取经营分析信息 | ths_operate_op >> bailian_web_parser_op >> extract_long_text_op | `{"query": "主营业务构成", "code": "000001"}` | `平安银行主营业务构成: 利息净收入占比78.5%，手续费及佣金净收入占比15.2%...` | -| crawl_ths_equity | 获取股本结构信息 | ths_equity_op >> bailian_web_parser_op >> extract_long_text_op | `{"query": "股本构成", "code": "000001"}` | `平安银行股本构成: 总股本194.4亿股，其中A股194.4亿股，流通股194.4亿股...` | -| crawl_ths_capital | 获取资本运作信息 | ths_capital_op >> bailian_web_parser_op >> extract_long_text_op | `{"query": "募集资金", "code": "000001"}` | `平安银行募集资金情况: 近期通过定向增发募集资金260亿元，主要用于补充核心一级资本...` | -| crawl_ths_worth | 获取盈利预测信息 | ths_worth_op >> bailian_web_parser_op >> extract_long_text_op | `{"query": "业绩预测", "code": "000001"}` | `平安银行业绩预测: 2024年预测净利润增长8.5%，ROE预计达到11.2%，机构一致目标价16.5元...` | -| crawl_ths_news | 获取新闻公告信息 | ths_news_op >> bailian_web_parser_op >> extract_long_text_op | `{"query": "最新公告", "code": "000001"}` | `平安银行最新公告: 发布2024年第三季度业绩报告，净利润同比增长6.8%，不良贷款率1.02%...` | -| crawl_ths_concept | 获取概念题材信息 | ths_concept_op >> bailian_web_parser_op >> extract_long_text_op | `{"query": "概念板块", "code": "000001"}` | `平安银行概念板块: 银行概念、金融科技、数字货币、区块链、人工智能等概念...` | -| crawl_ths_position | 获取主力持仓信息 | ths_position_op >> bailian_web_parser_op >> extract_long_text_op | `{"query": "机构持股", "code": "000001"}` | `平安银行机构持股: 基金持股占比12.5%，QFII持股占比3.2%，社保基金持股占比1.8%...` | -| crawl_ths_finance | 获取财务分析信息 | ths_finance_op >> bailian_web_parser_op >> extract_long_text_op | `{"query": "财务指标", "code": "000001"}` | `平安银行财务指标: ROE 11.1%，ROA 0.89%，资本充足率13.85%，核心一级资本充足率10.15%...` | -| crawl_ths_bonus | 获取分红融资信息 | ths_bonus_op >> bailian_web_parser_op >> extract_long_text_op | `{"query": "分红情况", "code": "000001"}` | `平安银行分红情况: 2023年每10股派现2.31元，股息率约1.5%，连续多年稳定分红...` | -| crawl_ths_event | 获取公司大事信息 | ths_event_op >> bailian_web_parser_op >> extract_long_text_op | `{"query": "高管变动", "code": "000001"}` | `平安银行高管变动: 近期董事会决议任命新任副行长，具有20年银行从业经验...` | -| crawl_ths_field | 获取行业对比信息 | ths_field_op >> bailian_web_parser_op >> extract_long_text_op | `{"query": "行业地位", "code": "000001"}` | `平安银行行业地位: 在股份制银行中排名第3位，资产规模4.9万亿元，市场份额约2.8%...` | - -### 数据来源 - -- [来源地址](https://basic.10jqka.com.cn/300033/) -- 同花顺基本面数据 -- 公司资料、股东信息、经营分析 -- 股本结构、资本运作、盈利预测 -- 新闻公告、概念题材、主力持仓 -- 财务分析、分红融资、公司大事 -- 行业对比等全方位数据 - -### 代码示例 - -```python -import asyncio -from flowllm.client.fastmcp_client import FastmcpClient - - -async def ths_supply_example(): - """THS supply使用示例""" - async with FastmcpClient(transport="sse", host="11.164.204.33", port=8001) as client: - - # 示例股票代码 - stock_code = "000001" # 平安银行 - - # 1. 获取公司基本资料信息 - print("=== 公司基本资料 ===") - result = await client.call_tool("crawl_ths_company", { - "query": "平安银行的公司情况如何？", - "code": stock_code - }) - print(f"公司信息: {result.content[0].text}") - - # 2. 获取股东研究信息 - print("\n=== 股东研究 ===") - result = await client.call_tool("crawl_ths_holder", { - "query": "平安银行的股东结构如何？", - "code": stock_code - }) - print(f"股东信息: {result.content[0].text}") - - # 3. 获取经营分析信息 - print("\n=== 经营分析 ===") - result = await client.call_tool("crawl_ths_operate", { - "query": "平安银行的主营业务是什么？", - "code": stock_code - }) - print(f"经营信息: {result.content[0].text}") - - # 4. 获取财务分析信息 - print("\n=== 财务分析 ===") - result = await client.call_tool("crawl_ths_finance", { - "query": "平安银行的财务状况如何？", - "code": stock_code - }) - print(f"财务信息: {result.content[0].text}") - - # 5. 获取新闻公告信息 - print("\n=== 新闻公告 ===") - result = await client.call_tool("crawl_ths_news", { - "query": "平安银行最近有什么新闻？", - "code": stock_code - }) - print(f"新闻信息: {result.content[0].text}") - - # 6. 获取概念题材信息 - print("\n=== 概念题材 ===") - result = await client.call_tool("crawl_ths_concept", { - "query": "平安银行涉及哪些概念题材？", - "code": stock_code - }) - print(f"概念信息: {result.content[0].text}") - - # 7. 获取主力持仓信息 - print("\n=== 主力持仓 ===") - result = await client.call_tool("crawl_ths_position", { - "query": "平安银行的机构持仓情况如何？", - "code": stock_code - }) - print(f"持仓信息: {result.content[0].text}") - - # 8. 其他维度信息示例 - other_examples = [ - ("crawl_ths_equity", {"query": "股本结构如何？", "code": stock_code}), - ("crawl_ths_capital", {"query": "有哪些资本运作？", "code": stock_code}), - ("crawl_ths_worth", {"query": "业绩预测如何？", "code": stock_code}), - ("crawl_ths_bonus", {"query": "分红情况如何？", "code": stock_code}), - ("crawl_ths_event", {"query": "最近有什么重大事件？", "code": stock_code}), - ("crawl_ths_field", {"query": "在行业中的地位如何？", "code": stock_code}) - ] - - for tool_name, params in other_examples: - print(f"\n=== {tool_name} ===") - try: - result = await client.call_tool(tool_name, params) - print(f"结果: {result.content[0].text[:200]}...") - except Exception as e: - print(f"调用失败: {e}") - - -if __name__ == "__main__": - asyncio.run(ths_supply_example()) -``` - -### 配置地址 - -- [fin_ths.yaml](../flowllm/config/fin_ths.yaml) - -## mcp supply - -MCP供给提供了基于阿里云百炼平台的量化投资分析能力，通过qtf_mcp工具提供不同深度的股票分析服务： - -| MCP Flow名称 | 能力描述 | Pipeline | 输入示例 | 输出示例 | -|----------------|------------|-------------------|--------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------| -| qtf_brief_mcp | 提供股票基础分析摘要 | qtf_brief_mcp_op | `{"symbol": "SZ000001"}` | `平安银行(000001)简要分析: 当前股价15.23元，涨幅2.35%。基本面良好，ROE 11.1%，市盈率8.5倍。技术面显示短期趋势向好，建议关注。` | -| qtf_medium_mcp | 提供股票中等深度分析 | qtf_medium_mcp_op | `{"symbol": "SZ000001"}` | `平安银行(000001)中等分析: 基本面分析-资产质量稳定，不良率1.02%；盈利能力强，ROE持续改善。技术面分析-MACD金叉，成交量放大，短期支撑位14.8元。估值合理，目标价16.5元。风险提示：银行业监管政策变化。` | -| qtf_full_mcp | 提供股票全面深度分析 | qtf_full_mcp_op | `{"symbol": "SZ000001"}` | `平安银行(000001)全面分析报告: 1.基本面：资产规模4.9万亿，净利润增长6.8%，ROE 11.1%，资本充足率13.85%。2.技术面：多项指标显示买入信号，支撑阻力位明确。3.估值：PB 0.85倍，PE 8.5倍，相对估值偏低。4.投资建议：买入，目标价16.5-17元，持有期6-12个月。` | - -### 代码示例 - -```python -import asyncio -from flowllm.client.fastmcp_client import FastmcpClient - - -async def mcp_supply_example(): - """MCP supply使用示例""" - async with FastmcpClient(transport="sse", host="11.164.204.33", port=8001) as client: - # 1. 简要分析 - print("=== 股票简要分析 ===") - result = await client.call_tool("qtf_brief_mcp", { - "symbol": "SZ000001" - }) - print(f"简要分析: {result.content[0].text}") - - # 2. 中等深度分析 - print("\n=== 股票中等深度分析 ===") - result = await client.call_tool("qtf_medium_mcp", { - "symbol": "SZ000001" - }) - print(f"中等分析: {result.content[0].text}") - - # 3. 全面深度分析 - print("\n=== 股票全面深度分析 ===") - result = await client.call_tool("qtf_full_mcp", { - "symbol": "SZ000001" - }) - print(f"全面分析: {result.content[0].text}") - - -if __name__ == "__main__": - asyncio.run(mcp_supply_example()) -``` - -### 主要特点 - -- **brief模式**：提供股票的基础信息和简要分析结论，适合快速了解 -- **medium模式**：包含基本面、技术面和估值的中等深度分析，平衡详细度和效率 -- **full模式**：提供全面深度的投资分析报告，包含详细的投资建议和风险提示 - -### 数据来源 - -- [阿里云百炼MCP市场](https://bailian.console.aliyun.com/pre-publish?tab=mcp#/mcp-market/detail/qtf_mcp) - -### 配置地址 - -- [fin_mcp.yaml](../flowllm/config/fin_mcp.yaml) - -## search supply - -搜索供给提供了多样化的网络搜索和内容解析能力，集成了多个优质搜索引擎和解析工具： - -| MCP Flow名称 | Pipeline | 输入示例 | 输出示例 | -|--------------------|------------------------------------------|----------------------------------------|-------------------------------------------------------------| -| tavily_search | tavily_search_op >> extract_long_text_op | `{"query": "2024年中国GDP增长情况"}` | `根据最新数据，2024年中国GDP同比增长5.2%，经济运行总体平稳，消费和投资成为主要增长动力...` | -| dashscope_search | dashscope_search_op | `{"query": "人工智能最新发展趋势"}` | `人工智能领域最新发展包括大语言模型技术突破、多模态AI应用普及、AI+行业深度融合等趋势...` | -| bailian_web_search | bailian_web_search_op | `{"query": "新能源汽车市场现状", "count": 10}` | `2024年新能源汽车市场保持高速增长，销量突破800万辆，市场渗透率达到35%，比亚迪、特斯拉领跑市场...` | -| bocha_web_search | bocha_web_search_op | `{"query": "区块链技术应用", "freshness": 7}` | `区块链技术在金融、供应链、数字身份等领域应用不断深化，央行数字货币DCEP试点扩大，DeFi生态持续发展...` | -| brave_web_search | brave_web_search_op | `{"query": "量子计算突破", "count": 5}` | `量子计算领域取得重大突破，IBM发布1000+量子比特处理器，Google实现量子纠错里程碑，产业化进程加速...` | -| bailian_web_parser | bailian_web_parser_op | `{"url": "https://example.com"}` | `解析网页内容：标题、正文、关键信息等结构化数据，支持多种格式输出...` | - -### 代码示例 - -```python -import asyncio -from flowllm.client.fastmcp_client import FastmcpClient - - -async def search_supply_example(): - """Search supply使用示例""" - async with FastmcpClient(transport="sse", host="11.164.204.33", port=8001) as client: - # 1. Tavily搜索（带长文本提取） - print("=== Tavily搜索 ===") - result = await client.call_tool("tavily_search", { - "query": "2024年中国GDP增长情况" - }) - print(f"搜索结果: {result.content[0].text}") - - # 2. 通义搜索 - print("\n=== 通义搜索 ===") - result = await client.call_tool("dashscope_search", { - "query": "人工智能最新发展趋势" - }) - print(f"搜索结果: {result.content[0].text}") - - # 3. 百炼网络搜索（带数量限制） - print("\n=== 百炼网络搜索 ===") - result = await client.call_tool("bailian_web_search", { - "query": "新能源汽车市场现状", - "count": 10 - }) - print(f"搜索结果: {result.content[0].text}") - - # 4. Bocha搜索（带时效性） - print("\n=== Bocha搜索 ===") - result = await client.call_tool("bocha_web_search", { - "query": "区块链技术应用", - "freshness": 7 - }) - print(f"搜索结果: {result.content[0].text}") - - # 5. Brave搜索 - print("\n=== Brave搜索 ===") - result = await client.call_tool("brave_web_search", { - "query": "量子计算突破", - "count": 5 - }) - print(f"搜索结果: {result.content[0].text}") - - # 6. 网页内容解析 - print("\n=== 网页内容解析 ===") - result = await client.call_tool("bailian_web_parser", { - "url": "https://basic.10jqka.com.cn/601899/field.html" - }) - print(f"解析结果: {result.content[0].text}") - - -if __name__ == "__main__": - asyncio.run(search_supply_example()) -``` - -### 配置地址 - -- [fin_search.yaml](../flowllm/config/fin_search.yaml) \ No newline at end of file diff --git a/old/docs/future_work.md b/old/docs/future_work.md deleted file mode 100644 index 2b3c410..0000000 --- a/old/docs/future_work.md +++ /dev/null @@ -1,28 +0,0 @@ -# 20251024 todo -1. optimize docs - 1. optimize README.md - 1. 基础能力：实现LLM+EMB+vectorstore驱动的工作流服务，不同http/stream_http/mcp/cmd模式 - 1. install，然后以MCP为例 - 2. 实现底层op 同步或者异步，注意输入，注意op之间交互，注意输出 - 3. 定义参数： - 1. http/mcp/cmd - 2. LLM+EMB+vectorstore参数 - 3. 定义flow，定义desc和input 、 等价代码定义 - 4. flowllm启动，curl命令测试/python/node.js/python import - 5. http的不同 - 6. stream_http的不同 - 7. cmd的不同 - 2. blog - 1. reading paper - 2. learning - 3. application - 1. reme_ai - 2. llm+search ui+demo - 3. deepresearch - 4. 金融供给 - 5. 金融因子图 - 6. 金融deepresearch ? - 7. desktop pet -2. optimize code - 1. fastmcp提供更加高级的支持 - 2. \ No newline at end of file diff --git a/old/docs/making_sop_memories.md b/old/docs/making_sop_memories.md deleted file mode 100644 index 3a3ca3f..0000000 --- a/old/docs/making_sop_memories.md +++ /dev/null @@ -1,296 +0,0 @@ -# SOP Memory: Combining Atomic Operations into Complex Workflows - -## 1. Background - -In LLM application development, we often need to combine multiple basic operations (atomic operations) into more complex workflows. These workflows can handle complex tasks such as data retrieval, code generation, multi-turn dialogues, and more. By combining these atomic operations into Standard Operating Procedures (SOPs), we can: - -- Improve code reusability -- Simplify implementation of complex tasks -- Standardize common workflows -- Reduce development and maintenance costs - -This document introduces how to combine atomic operations (Ops) to form new composite operation tools using the FlowLLM framework. - -## 2. Technical Solution - -### 2.1 Atomic Operation Definition - -Each operation (Op) needs to define the following core attributes: - -```python -class BaseOp: - description: str # Description of the operation - input_schema: Dict[str, ParamAttr] # Input parameter schema definition - output_schema: Dict[str, ParamAttr] # Output parameter schema definition -``` - -Where `ParamAttr` defines parameter type, whether it's required, and other attributes: - -```python -class ParamAttr: - type: Type # Parameter type, such as str, int, Dict, etc. - required: bool = True # Whether it must be provided - default: Any = None # Default value - description: str = "" # Parameter description -``` - -### 2.2 SOP Composition Process - -#### Step 1: Create Atomic Operation Instances - -First, instantiate the required atomic operations: - -```python -from flowllm.op.gallery.mock_op import MockOp -from flowllm.op.search.tavily_search_op import TavilySearchOp -from flowllm.op.llm.react_v2_op import ReactV2Op - -# Create atomic operation instances -search_op = TavilySearchOp() -react_op = ReactV2Op() -summary_op = MockOp( - description="Summarize search results", - input_schema={"search_results": ParamAttr(type=str, description="Search results to summarize")}, - output_schema={"summary": ParamAttr(type=str, description="Summarized content")} -) -``` - -#### Step 2: Define Data Flow Between Operations - -Set up input-output relationships between operations, defining how data flows between them: - -```python -# Set input parameter sources -react_op.set_input("context", "search_summary") # react_op's context parameter is retrieved from search_summary in memory - -# Set output parameter destinations -search_op.set_output("results", "search_results") # search_op's results output to search_results in memory -summary_op.set_output("summary", "search_summary") # summary_op's summary output to search_summary in memory -``` - -#### Step 3: Build Operation Flow Graph - -Use operators to build the operation flow graph, defining execution order and parallel relationships: - -```python -# Build operation flow graph -flow = search_op >> summary_op >> react_op - -# Or more complex flows -# Parallel operations use the | operator, sequential operations use the >> operator -complex_flow = (search_op >> summary_op) | (another_search_op >> another_summary_op) >> react_op -``` - -Operator explanation: -- `>>`: Sequential execution, execute the next operation after the previous one completes -- `|`: Parallel execution, execute multiple operations simultaneously - -#### Step 4: Create Composite Operation Class - -Encapsulate the built operation flow into a new composite operation class: - -```python -from flowllm.op.base_op import BaseOp - -class SearchAndReactOp(BaseOp): - description = "Search for information and generate a response based on search results" - input_schema = ... - output_schema = ... - - def build_flow(self): - search_op = TavilySearchOp() - summary_op = MockOp() - react_op = ReactV2Op() - - # Set data flow - search_op.set_output("results", "search_results") - summary_op.set_input("search_results", "search_results") - summary_op.set_output("summary", "search_summary") - react_op.set_input("context", "search_summary") - react_op.set_output("response", "response") - - # Build operation flow graph - self.flow = search_op >> summary_op >> react_op - - async def execute(self, inputs: Dict[str, Any]) -> Dict[str, Any]: - # Execute operation flow - return await self.flow.execute(inputs) -``` - -### 2.3 Memory Management - -Data transfer between operations is accomplished through memory space: - -```python -# Memory space is a key-value store -memory = {} - -# set_input reads data from memory -op.set_input("param_name", "memory_key") # op.param_name = memory["memory_key"] - -# set_output writes data to memory -op.set_output("result_name", "memory_key") # memory["memory_key"] = op.result_name -``` - -## 3. Example: Building a Search Q&A SOP - -Below is a complete example demonstrating how to build a search Q&A SOP: - -```python -from flowllm.op.base_op import BaseOp -from flowllm.op.search.tavily_search_op import TavilySearchOp -from flowllm.op.llm.react_v2_op import ReactV2Op -from flowllm.schema.vector_node import ParamAttr - - -class SearchQAOp(BaseOp): - description = "Search for information and answer questions based on search results" - input_schema = { - "question": ParamAttr(type=str, description="User question") - } - output_schema = { - "answer": ParamAttr(type=str, description="Answer based on search results") - } - - def __init__(self): - super().__init__() - - # 1. Create atomic operation instances - search_op = TavilySearchOp() - react_op = ReactV2Op() - - # 2. Set data flow - search_op.set_input("query", "question") # Get search query from input question - search_op.set_output("results", "search_results") # Store search results in search_results - - react_op.set_input("question", "question") # Get question from input question - react_op.set_input("context", "search_results") # Get context from search_results - react_op.set_output("answer", "answer") # Store answer in answer - - # 3. Build operation flow graph - self.flow = search_op >> react_op - - async def execute(self, inputs: Dict[str, Any]) -> Dict[str, Any]: - return await self.flow.execute(inputs) -``` - -## 4. Solution Validation - -### 4.1 Unit Testing - -To verify the correctness of SOPs, we can write unit tests: - -```python -import unittest -import asyncio -from flowllm.op.gallery.mock_op import MockOp -from flowllm.schema.vector_node import ParamAttr - -class TestSOPMemory(unittest.TestCase): - def test_simple_flow(self): - # Create mock operations - mock1 = MockOp( - description="Mock operation 1", - input_schema={"input1": ParamAttr(type=str)}, - output_schema={"output1": ParamAttr(type=str)} - ) - mock2 = MockOp( - description="Mock operation 2", - input_schema={"input2": ParamAttr(type=str)}, - output_schema={"output2": ParamAttr(type=str)} - ) - - # Set data flow - mock1.set_output("output1", "intermediate") - mock2.set_input("input2", "intermediate") - - # Build flow - flow = mock1 >> mock2 - - # Execute flow - result = asyncio.run(flow.execute({"input1": "test input"})) - - # Verify results - self.assertIn("output2", result) -``` - -### 4.2 Integration Testing - -Validate the effectiveness of SOPs through real-world scenario testing: - -```python -async def test_search_qa(): - # Create SearchQA operation - search_qa = SearchQAOp() - - # Execute operation - result = await search_qa.execute({ - "question": "What is the capital of France?" - }) - - # Print results - print(f"Question: What is the capital of France?") - print(f"Answer: {result['answer']}") - -# Run test -if __name__ == "__main__": - asyncio.run(test_search_qa()) -``` - -### 4.3 Performance Benchmarking - -Measure SOP performance metrics: - -```python -import time -import asyncio -from flowllm.utils.timer import Timer - -async def benchmark_sop(sop_op, inputs, num_runs=10): - total_time = 0 - results = [] - - for _ in range(num_runs): - with Timer() as timer: - result = await sop_op.execute(inputs) - - total_time += timer.elapsed_time - results.append(result) - - avg_time = total_time / num_runs - print(f"Average execution time: {avg_time:.2f}s") - - return results, avg_time - -# Run benchmark -search_qa = SearchQAOp() -asyncio.run(benchmark_sop(search_qa, {"question": "What is quantum computing?"})) -``` - -## 5. Best Practices - -### 5.1 Operation Design Principles - -- **Single Responsibility**: Each atomic operation should focus on a single functionality -- **Clear Interfaces**: Clearly define input and output schemas -- **Composability**: Design operations with combination with other operations in mind -- **Error Handling**: Properly handle exceptional cases - -### 5.2 SOP Design Patterns - -- **Pipeline Pattern**: `op1 >> op2 >> op3` -- **Branch Pattern**: `op1 >> (op2 | op3) >> op4` -- **Aggregation Pattern**: `(op1 | op2) >> op3` -- **Conditional Pattern**: Choose different operation paths based on conditions - -### 5.3 Debugging Tips - -- Use logging to record the input and output of each operation -- Create visualization charts for complex SOPs -- Use mock operations for isolated testing - -## 6. Conclusion - -Through the SOP Memory mechanism, we can flexibly combine atomic operations to build complex LLM application workflows. This approach not only improves code maintainability and reusability but also makes the implementation of complex tasks simpler and more standardized. - -As more atomic operations are developed and refined, we can build a richer and more powerful SOP library, further enhancing the development efficiency and quality of LLM applications. \ No newline at end of file diff --git a/old/docs/todo.md b/old/docs/todo.md deleted file mode 100644 index afa098b..0000000 --- a/old/docs/todo.md +++ /dev/null @@ -1,6 +0,0 @@ -# TODO 20251015 -- [ ] 在ReMe中增加tool memory，支持summary & retrieve的能力。 - - [ ] 支持浩然能够支持日志的打印 - - [ ] 融合到asio的memory服务中 -- [ ] 增加flowllm的logo -- [ ] 增加flowllm的 \ No newline at end of file diff --git a/old/docs/vector_store.md b/old/docs/vector_store.md deleted file mode 100644 index 2a753c5..0000000 --- a/old/docs/vector_store.md +++ /dev/null @@ -1,887 +0,0 @@ -# 🚀 Vector Store API Guide - -This guide covers the vector store implementations available in flowllm, their APIs, and how to use them effectively. - -## 📋 Overview - -flowllm provides multiple vector store backends for different use cases: - -- **LocalVectorStore** (`backend=local`) - 📁 Simple file-based storage for development and small datasets -- **ChromaVectorStore** (`backend=chroma`) - 🔮 Embedded vector database for moderate scale -- **EsVectorStore** (`backend=elasticsearch`) - 🔍 Elasticsearch-based storage for production and large scale -- **QdrantVectorStore** (`backend=qdrant`) - 🎯 High-performance vector database with advanced filtering -- **MemoryVectorStore** (`backend=memory`) - ⚡ In-memory storage for ultra-fast access and testing - -All vector stores implement the `BaseVectorStore` interface, providing a consistent API across implementations. - -## 📊 Comparison Table - -| Feature | LocalVectorStore | ChromaVectorStore | EsVectorStore | QdrantVectorStore | MemoryVectorStore | -|----------------------|------------------|-------------------|---------------|-------------------|-------------------| -| **Storage** | File (JSONL) | Embedded DB | Elasticsearch | Qdrant Server | In-Memory | -| **Performance** | Medium | Good | Excellent | Excellent | Ultra-Fast | -| **Scalability** | < 10K vectors | < 1M vectors | > 1M vectors | > 10M vectors | < 1M vectors | -| **Persistence** | ✅ Auto | ✅ Auto | ✅ Auto | ✅ Auto | ⚠️ Manual | -| **Setup Complexity** | 🟢 Simple | 🟡 Medium | 🔴 Complex | 🟡 Medium | 🟢 Simple | -| **Dependencies** | None | ChromaDB | Elasticsearch | Qdrant | None | -| **Filtering** | ❌ Basic | ✅ Metadata | ✅ Advanced | ✅ Advanced | ❌ Basic | -| **Concurrency** | ❌ Limited | ✅ Good | ✅ Excellent | ✅ Excellent | ❌ Single Process | -| **Async Support** | ❌ No | ❌ No | ❌ No | ✅ Native | ❌ No | -| **Best For** | Development | Local Apps | Production | Production/Cloud | Testing | - -## 🔄 Common API Methods - -All vector store implementations share these core methods: - -### 🔄 Async Support - -All vector stores provide both synchronous and asynchronous versions of every method: - -```python -# Synchronous methods -store.search(query="example", workspace_id="workspace", top_k=5) -store.insert(nodes, workspace_id="workspace") - -# Asynchronous methods (with async_ prefix) -await store.async_search(query="example", workspace_id="workspace", top_k=5) -await store.async_insert(nodes, workspace_id="workspace") -``` - -### Workspace Management - -```python -# Check if workspace exists -store.exist_workspace(workspace_id: str) -> bool - -# Create a new workspace -store.create_workspace(workspace_id: str, **kwargs) - -# Delete a workspace -store.delete_workspace(workspace_id: str, **kwargs) - -# Copy a workspace -store.copy_workspace(src_workspace_id: str, dest_workspace_id: str, **kwargs) -``` - -### Data Operations - -```python -# Insert nodes (single or list) -store.insert(nodes: VectorNode | List[VectorNode], workspace_id: str, **kwargs) - -# Delete nodes by ID -store.delete(node_ids: str | List[str], workspace_id: str, **kwargs) - -# Search for similar nodes -store.search(query: str, workspace_id: str, top_k: int = 1, **kwargs) -> List[VectorNode] - -# Iterate through workspace nodes -for node in store.list_workspace_nodes(workspace_id: str, **kwargs): - # Process each node -``` - -### Import/Export - -```python -# Export workspace to file -store.dump_workspace(workspace_id: str, path: str | Path = "", callback_fn=None, **kwargs) - -# Import workspace from file -store.load_workspace(workspace_id: str, path: str | Path = "", nodes: List[VectorNode] = None, - callback_fn=None, **kwargs) -``` - -## ⚡ Vector Store Implementations - -### 1. 📁 LocalVectorStore (`backend=local`) - -A simple file-based vector store that saves data to local JSONL files. - -#### 💡 When to Use -- **Development and testing** - No external dependencies required 🛠️ -- **Small datasets** - Suitable for datasets with < 10,000 vectors 📊 -- **Single-user applications** - Limited concurrent access support 👤 - -#### ⚙️ Configuration - -```python -from flowllm.storage.vector_store import LocalVectorStore -from flowllm.embedding_model import OpenAICompatibleEmbeddingModel -from flowllm.utils.common_utils import load_env - -# Load environment variables (for API keys) -load_env() - -# Initialize embedding model -embedding_model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4") - -# Initialize vector store -vector_store = LocalVectorStore( - embedding_model=embedding_model, - store_dir="./file_vector_store", # Directory to store JSONL files - batch_size=1024 # Batch size for operations -) -``` - -#### 💻 Example Usage - -```python -from flowllm.schema.vector_node import VectorNode - -# Create workspace -workspace_id = "my_workspace" -vector_store.create_workspace(workspace_id) - -# Create nodes -nodes = [ - VectorNode( - unique_id="node1", - workspace_id=workspace_id, - content="Artificial intelligence is revolutionizing technology", - metadata={"category": "tech", "source": "article1"} - ), - VectorNode( - unique_id="node2", - workspace_id=workspace_id, - content="Machine learning enables data-driven insights", - metadata={"category": "tech", "source": "article2"} - ) -] - -# Insert nodes -vector_store.insert(nodes, workspace_id) - -# Search -results = vector_store.search("What is AI?", workspace_id, top_k=2) -for result in results: - print(f"Content: {result.content}") - print(f"Metadata: {result.metadata}") - print(f"Score: {result.metadata.get('score', 'N/A')}") -``` - -### 2. 🔮 ChromaVectorStore (`backend=chroma`) - -An embedded vector database that provides persistent storage with advanced features. - -#### 💡 When to Use -- **Local development** with persistence requirements 🏠 -- **Medium-scale applications** (10K - 1M vectors) 📈 -- **Applications requiring metadata filtering** 🔍 - -#### ⚙️ Configuration - -```python -from flowllm.storage.vector_store import ChromaVectorStore -from flowllm.embedding_model import OpenAICompatibleEmbeddingModel -from flowllm.utils.common_utils import load_env - -# Load environment variables -load_env() - -# Initialize embedding model -embedding_model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4") - -# Initialize vector store -vector_store = ChromaVectorStore( - embedding_model=embedding_model, - store_dir="./chroma_vector_store", # Directory for Chroma database - batch_size=1024 # Batch size for operations -) -``` - -#### 💻 Example Usage - -```python -from flowllm.schema.vector_node import VectorNode - -workspace_id = "chroma_workspace" - -# Check if workspace exists and create if needed -if not vector_store.exist_workspace(workspace_id): - vector_store.create_workspace(workspace_id) - -# Create nodes with metadata -nodes = [ - VectorNode( - unique_id="node1", - workspace_id=workspace_id, - content="Deep learning models require large datasets", - metadata={ - "category": "AI", - "difficulty": "advanced", - "topic": "deep_learning" - } - ), - VectorNode( - unique_id="node2", - workspace_id=workspace_id, - content="Transformer architecture revolutionized NLP", - metadata={ - "category": "AI", - "difficulty": "intermediate", - "topic": "transformers" - } - ) -] - -# Insert nodes -vector_store.insert(nodes, workspace_id) - -# Search -results = vector_store.search("deep learning", workspace_id, top_k=5) -for result in results: - print(f"Content: {result.content}") - print(f"Metadata: {result.metadata}") -``` - -### 3. 🔍 EsVectorStore (`backend=elasticsearch`) - -Production-grade vector search using Elasticsearch with advanced filtering and scaling capabilities. - -#### 💡 When to Use -- **Production environments** requiring high availability 🏭 -- **Large-scale applications** (1M+ vectors) 🚀 -- **Complex filtering requirements** on metadata 🎯 - -#### 🛠️ Setup Elasticsearch - -Before using EsVectorStore, set up Elasticsearch: - -##### Option 1: Docker Run -```bash -# Pull the latest Elasticsearch image -docker pull docker.elastic.co/elasticsearch/elasticsearch-wolfi:9.0.0 - -# Run Elasticsearch container -docker run -p 9200:9200 \ - -e "discovery.type=single-node" \ - -e "xpack.security.enabled=false" \ - -e "xpack.license.self_generated.type=trial" \ - -e "http.host=0.0.0.0" \ - docker.elastic.co/elasticsearch/elasticsearch-wolfi:9.0.0 -``` - -##### Environment Configuration -```bash -export FLOW_ES_HOSTS=http://localhost:9200 -``` - -#### ⚙️ Configuration - -```python -from flowllm.storage.vector_store import EsVectorStore -from flowllm.embedding_model import OpenAICompatibleEmbeddingModel -from flowllm.utils.common_utils import load_env -import os - -# Load environment variables -load_env() - -# Initialize embedding model -embedding_model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4") - -# Initialize vector store -vector_store = EsVectorStore( - embedding_model=embedding_model, - hosts=os.getenv("FLOW_ES_HOSTS", "http://localhost:9200"), # Elasticsearch hosts - basic_auth=None, # ("username", "password") for auth - batch_size=1024 # Batch size for bulk operations -) -``` - -#### 🎯 Advanced Filtering - -EsVectorStore supports advanced filtering capabilities through the `filter_dict` parameter: - -```python -# Term filters (exact match) -term_filter = { - "category": "technology", - "author": "research_team" -} - -# Range filters (numeric and date ranges) -range_filter = { - "score": {"gte": 0.8}, # Score >= 0.8 - "confidence": {"gte": 0.5, "lte": 0.9}, # Between 0.5 and 0.9 - "timestamp": {"gte": "2024-01-01", "lte": "2024-12-31"} -} - -# Combined filters (filters are combined with AND logic) -combined_filter = { - "category": "AI", - "confidence": {"gte": 0.9} -} - -# Search with filters applied -results = vector_store.search("machine learning", workspace_id, top_k=10, filter_dict=combined_filter) -``` - -#### ⚡ Performance Optimization - -```python -# Refresh index for immediate availability (useful after bulk inserts) -vector_store.insert(nodes, workspace_id, refresh=True) # Auto-refresh -vector_store.refresh(workspace_id) # Manual refresh - -# Bulk operations with custom batch size -vector_store.insert(large_node_list, workspace_id, refresh=False) # Skip refresh for speed -vector_store.refresh(workspace_id) # Refresh once after all inserts -``` - -#### 💻 Example Usage - -```python -from flowllm.schema.vector_node import VectorNode - -# Define workspace -workspace_id = "production_workspace" - -# Create workspace if needed -if not vector_store.exist_workspace(workspace_id): - vector_store.create_workspace(workspace_id) - -# Create nodes with rich metadata -nodes = [ - VectorNode( - unique_id="doc1", - workspace_id=workspace_id, - content="Transformer architecture revolutionized NLP", - metadata={ - "category": "AI", - "subcategory": "NLP", - "author": "research_team", - "timestamp": "2024-01-15", - "confidence": 0.95, - "tags": ["transformer", "nlp", "attention"] - } - ) -] - -# Insert with refresh for immediate availability -vector_store.insert(nodes, workspace_id, refresh=True) - -# Advanced search with filters -filter_dict = { - "category": "AI", - "confidence": {"gte": 0.9} -} - -results = vector_store.search("transformer models", workspace_id, top_k=5, filter_dict=filter_dict) - -for result in results: - print(f"Score: {result.metadata.get('score', 'N/A')}") - print(f"Content: {result.content}") - print(f"Metadata: {result.metadata}") -``` - -### 4. 🎯 QdrantVectorStore (`backend=qdrant`) - -A high-performance vector database designed for production workloads with native async support and advanced filtering. - -#### 💡 When to Use -- **Production environments** requiring high performance and reliability 🏭 -- **Large-scale applications** (10M+ vectors) with excellent horizontal scaling 🚀 -- **Applications requiring native async operations** for better concurrency ⚡ -- **Complex filtering and metadata queries** on large datasets 🎯 -- **Cloud-native deployments** with Qdrant Cloud support ☁️ - -#### 🛠️ Setup Qdrant - -Before using QdrantVectorStore, set up Qdrant: - -##### Option 1: Docker Run (Recommended for Development) -```bash -# Pull the latest Qdrant image -docker pull qdrant/qdrant - -# Run Qdrant container -docker run -p 6333:6333 -p 6334:6334 \ - -v $(pwd)/qdrant_storage:/qdrant/storage:z \ - qdrant/qdrant -``` - -##### Option 2: Qdrant Cloud -For production, you can use [Qdrant Cloud](https://cloud.qdrant.io/) for managed hosting. - -##### Environment Configuration -```bash -# For local setup -export FLOW_QDRANT_HOST=localhost -export FLOW_QDRANT_PORT=6333 - -# For cloud setup (optional) -export FLOW_QDRANT_API_KEY=your-api-key -``` - -#### ⚙️ Configuration - -```python -from flowllm.storage.vector_store import QdrantVectorStore -from flowllm.embedding_model import OpenAICompatibleEmbeddingModel -from flowllm.utils.common_utils import load_env -import os - -# Load environment variables -load_env() - -# Initialize embedding model -embedding_model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4") - -# Option 1: Use localhost with environment variables -vector_store = QdrantVectorStore( - embedding_model=embedding_model, - host=os.getenv("FLOW_QDRANT_HOST", "localhost"), - port=int(os.getenv("FLOW_QDRANT_PORT", "6333")), - batch_size=1024 -) - -# Option 2: Use URL (for Qdrant Cloud or remote servers) -vector_store = QdrantVectorStore( - embedding_model=embedding_model, - url="http://your-qdrant-server:6333", - api_key="your-api-key", # Optional, for cloud - batch_size=1024 -) - -# Option 3: Specify custom distance metric -from qdrant_client.http.models import Distance - -vector_store = QdrantVectorStore( - embedding_model=embedding_model, - host="localhost", - port=6333, - distance=Distance.COSINE, # or Distance.EUCLIDEAN, Distance.DOT - batch_size=1024 -) -``` - -#### 🎯 Advanced Filtering - -QdrantVectorStore supports advanced filtering capabilities similar to Elasticsearch: - -```python -# Term filters (exact match) -term_filter = { - "category": "AI", - "node_type": "research" -} - -# Range filters (numeric) -range_filter = { - "confidence": {"gte": 0.8, "lte": 1.0}, # Between 0.8 and 1.0 - "score": {"gt": 0.5} # Greater than 0.5 -} - -# Combined filters (all conditions must match - AND logic) -combined_filter = { - "category": "AI", - "confidence": {"gte": 0.9}, - "node_type": "research" -} - -# Search with filters -results = vector_store.search( - query="machine learning", - workspace_id=workspace_id, - top_k=10, - filter_dict=combined_filter -) -``` - -##### Filter Operations Supported: -- **Exact match**: `{"field": "value"}` -- **Range queries**: - - `gte`: Greater than or equal - - `lte`: Less than or equal - - `gt`: Greater than - - `lt`: Less than - -#### ⚡ Async Operations - -QdrantVectorStore provides **native async support** for all operations: - -```python -import asyncio - -async def main(): - # All operations have async equivalents - - # Check if workspace exists - exists = await vector_store.async_exist_workspace(workspace_id) - - # Create workspace - if not exists: - await vector_store.async_create_workspace(workspace_id) - - # Insert nodes with async embedding - await vector_store.async_insert(nodes, workspace_id) - - # Search with async embedding - results = await vector_store.async_search( - query="AI research", - workspace_id=workspace_id, - top_k=5, - filter_dict={"category": "AI"} - ) - - # Delete nodes - await vector_store.async_delete(node_ids, workspace_id) - - # Delete workspace - await vector_store.async_delete_workspace(workspace_id) - - # Close client - await vector_store.async_close() - -# Run async operations -asyncio.run(main()) -``` - -#### 💻 Example Usage - -```python -from flowllm.schema.vector_node import VectorNode - -workspace_id = "qdrant_workspace" - -# Check and create workspace -if not vector_store.exist_workspace(workspace_id): - vector_store.create_workspace(workspace_id) - -# Create nodes with rich metadata -nodes = [ - VectorNode( - unique_id="node1", - workspace_id=workspace_id, - content="Artificial intelligence is revolutionizing technology", - metadata={ - "category": "AI", - "node_type": "research", - "confidence": 0.95, - "author": "research_team" - } - ), - VectorNode( - unique_id="node2", - workspace_id=workspace_id, - content="Machine learning models require large datasets", - metadata={ - "category": "AI", - "node_type": "tutorial", - "confidence": 0.85, - "author": "education_team" - } - ), - VectorNode( - unique_id="node3", - workspace_id=workspace_id, - content="Deep learning excels at image recognition", - metadata={ - "category": "AI", - "node_type": "research", - "confidence": 0.92, - "author": "research_team" - } - ) -] - -# Insert nodes (upsert - creates or updates) -vector_store.insert(nodes, workspace_id) - -# Simple search -results = vector_store.search("What is AI?", workspace_id, top_k=3) -for result in results: - print(f"Content: {result.content}") - print(f"Score: {result.metadata.get('score', 'N/A')}") - print(f"Metadata: {result.metadata}") - print("-" * 50) - -# Advanced search with filters -filter_dict = { - "node_type": "research", - "confidence": {"gte": 0.9} -} - -filtered_results = vector_store.search( - query="AI technology", - workspace_id=workspace_id, - top_k=5, - filter_dict=filter_dict -) - -print(f"Found {len(filtered_results)} filtered results") -for result in filtered_results: - print(f"Content: {result.content}") - print(f"Metadata: {result.metadata}") - -# Iterate through all nodes -print("\nAll nodes in workspace:") -for node in vector_store.list_workspace_nodes(workspace_id, limit=100): - print(f"ID: {node.unique_id}, Content: {node.content[:50]}...") - -# Update a node (delete + insert) -updated_node = VectorNode( - unique_id="node1", - workspace_id=workspace_id, - content="Artificial intelligence is transforming industries worldwide", - metadata={ - "category": "AI", - "node_type": "research", - "confidence": 0.98, - "author": "research_team", - "updated": True - } -) -vector_store.delete("node1", workspace_id) -vector_store.insert(updated_node, workspace_id) - -# Export workspace for backup -vector_store.dump_workspace(workspace_id, path="./qdrant_backup") - -# Clean up -vector_store.close() -``` - -#### 🔄 Async Example - -```python -import asyncio -from flowllm.schema.vector_node import VectorNode - -async def async_example(): - workspace_id = "async_qdrant_workspace" - - # Create workspace - if not await vector_store.async_exist_workspace(workspace_id): - await vector_store.async_create_workspace(workspace_id) - - # Create nodes - nodes = [ - VectorNode( - unique_id="async_node1", - workspace_id=workspace_id, - content="Async operations enable better performance", - metadata={"type": "performance", "async": True} - ), - VectorNode( - unique_id="async_node2", - workspace_id=workspace_id, - content="Concurrent requests improve throughput", - metadata={"type": "performance", "async": True} - ) - ] - - # Insert with async embedding - await vector_store.async_insert(nodes, workspace_id) - - # Search with async embedding - results = await vector_store.async_search( - query="performance optimization", - workspace_id=workspace_id, - top_k=2, - filter_dict={"async": True} - ) - - for result in results: - print(f"Score: {result.metadata['score']:.4f}") - print(f"Content: {result.content}") - - # Cleanup - await vector_store.async_delete_workspace(workspace_id) - await vector_store.async_close() - -# Run async example -asyncio.run(async_example()) -``` - -#### 🌟 Key Features - -- **Native Async Support** - All operations have async equivalents for better concurrency -- **Upsert Operations** - Insert automatically updates existing nodes with the same ID -- **Advanced Filtering** - Support for term and range filters on metadata -- **High Performance** - Optimized for large-scale vector similarity search -- **Horizontal Scaling** - Supports clustering for distributed deployments -- **Multiple Distance Metrics** - Cosine, Euclidean, and Dot Product similarity -- **Persistent Storage** - Data is automatically persisted to disk -- **Efficient Iteration** - Scroll through large collections with pagination - -#### 🚨 Important Notes - -- **Collection = Workspace** - Qdrant uses "collections" which map to workspace_id -- **Automatic Embedding** - Nodes without vectors are automatically embedded -- **ID-based Upsert** - Using the same unique_id will update existing nodes -- **Metadata Indexing** - All metadata fields are automatically indexed for filtering -- **Connection Management** - Call `close()` or `async_close()` to cleanup connections - -#### 📊 Performance Tips - -1. **Batch Operations** - Insert multiple nodes at once for better performance -2. **Use Async** - For high-concurrency scenarios, use async methods -3. **Optimize Filters** - Use indexed metadata fields for faster filtering -4. **Pagination** - Use `list_workspace_nodes()` with appropriate `limit` for large collections -5. **Distance Metric** - Choose appropriate distance metric for your use case (COSINE for normalized vectors) - -### 5. ⚡ MemoryVectorStore (`backend=memory`) - -An ultra-fast in-memory vector store that keeps all data in RAM for maximum performance. - -#### 💡 When to Use -- **Testing and development** - Fastest possible operations for unit tests 🧪 -- **Small to medium datasets** that fit in memory (< 1M vectors) 💾 -- **Applications requiring ultra-low latency** search operations ⚡ -- **Temporary workspaces** that don't need persistence 🚀 - -#### ⚙️ Configuration - -```python -from flowllm.storage.vector_store import MemoryVectorStore -from flowllm.embedding_model import OpenAICompatibleEmbeddingModel -from flowllm.utils.common_utils import load_env - -# Load environment variables -load_env() - -# Initialize embedding model -embedding_model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4") - -# Initialize vector store -vector_store = MemoryVectorStore( - embedding_model=embedding_model, - store_dir="./memory_vector_store", # Directory for backup/restore operations - batch_size=1024 # Batch size for operations -) -``` - -#### 💻 Example Usage - -```python -from flowllm.schema.vector_node import VectorNode - -workspace_id = "memory_workspace" - -# Create workspace in memory -vector_store.create_workspace(workspace_id) - -# Create nodes -nodes = [ - VectorNode( - unique_id="mem_node1", - workspace_id=workspace_id, - content="Memory stores provide ultra-fast access to data", - metadata={ - "category": "performance", - "type": "memory", - "speed": "ultra_fast" - } - ), - VectorNode( - unique_id="mem_node2", - workspace_id=workspace_id, - content="In-memory databases excel at low-latency operations", - metadata={ - "category": "performance", - "type": "database", - "latency": "low" - } - ) -] - -# Insert nodes (stored in memory) -vector_store.insert(nodes, workspace_id) - -# Ultra-fast search -results = vector_store.search("fast memory access", workspace_id, top_k=2) -for result in results: - print(f"Content: {result.content}") - print(f"Score: {result.metadata.get('score', 'N/A')}") - -# Optional: Save to disk for backup -vector_store.dump_workspace(workspace_id, path="./backup") - -# Optional: Load from disk to memory -vector_store.load_workspace(workspace_id, path="./backup") -``` - -#### ⚡ Performance Benefits - -- **Zero I/O latency** - All operations happen in RAM -- **Instant search results** - No disk or network overhead -- **Perfect for testing** - Fast setup and teardown -- **Memory efficient** - Only stores what you need - -#### 🚨 Important Notes - -- **Data is volatile** - Lost when process ends unless explicitly saved -- **Memory usage** - Entire dataset must fit in available RAM -- **No persistence** - Use `dump_workspace()` to save to disk -- **Single process** - Not suitable for distributed applications - -## 📝 Working with VectorNode - -The `VectorNode` class is the fundamental data unit for all vector stores: - -```python -from flowllm.schema.vector_node import VectorNode - -# Create a node -node = VectorNode( - unique_id="unique_identifier", # Unique ID for the node (required) - workspace_id="my_workspace", # Workspace ID (required) - content="Text content to embed", # Content to be embedded (required) - metadata={ # Optional metadata - "source": "document1", - "category": "technology", - "timestamp": "2024-08-29" - }, - vector=None # Vector will be generated automatically if None -) -``` - -## 🔄 Import/Export Example - -Export and import workspaces for backup or transfer: - -```python -# Export workspace to file -vector_store.dump_workspace( - workspace_id="my_workspace", - path="./backup_data" # Directory to store the exported data -) - -# Import workspace from file -vector_store.load_workspace( - workspace_id="new_workspace", - path="./backup_data" # Directory containing the exported data -) - -# Copy workspace within the same store -vector_store.copy_workspace( - src_workspace_id="original_workspace", - dest_workspace_id="copied_workspace" -) -``` - -## 🧩 Integration with Embedding Models - -All vector stores require an embedding model to function: - -```python -from flowllm.embedding_model import OpenAICompatibleEmbeddingModel - -# Initialize embedding model -embedding_model = OpenAICompatibleEmbeddingModel( - dimensions=64, # Embedding dimensions - model_name="text-embedding-v4", # Model name - batch_size=32 # Batch size for embedding generation -) - -# Pass to vector store (example with LocalVectorStore) -# You can also use: ChromaVectorStore, EsVectorStore, QdrantVectorStore, or MemoryVectorStore -vector_store = LocalVectorStore( - embedding_model=embedding_model, - store_dir="./vector_store" -) -``` - -🎉 This guide provides everything you need to work with vector stores in flowllm. Choose the implementation that best fits your use case and scale up as needed! ✨ \ No newline at end of file diff --git a/old/extract_long_text_prompt.yaml b/old/extract_long_text_prompt.yaml deleted file mode 100644 index 340482f..0000000 --- a/old/extract_long_text_prompt.yaml +++ /dev/null @@ -1,22 +0,0 @@ -extract_content_prompt: | - # Long Text - {long_text} - - # Task - From the provided long text (e.g., scraped web content or other documents), **extract and organize** all segments that are directly relevant to the following query: - - # Current Datetime - {datetime} - - # Query - {query} - - # Requirements: - - **Extract precisely**: Preserve original wording and factual details — do not paraphrase or summarize unless absolutely necessary for coherence. - - **Clean and organize**: Remove any content that is irrelevant to the query. Retain only segments that directly address or substantively relate to the query. - - **Time-filtered**: Keep only content dated within the most recent few years (typically 3–5 years from {datetime}), unless older content is explicitly critical for context or continuity. Clearly indicate if older content is retained and why. - - Maintain strict accuracy and neutrality — avoid assumptions, interpretations, or editorializing. - - Format output in Markdown, using only headings up to level three (###). - - Use the same language as the query (e.g., if query is in English, output in English; if in Chinese, output in Chinese). - - Clearly separate and label each relevant excerpt if multiple non-contiguous segments are extracted. - - If no relevant content exists (after filtering for recency and relevance), explicitly state: No relevant content found matching the query. diff --git a/old/flow/__init__.py b/old/flow/__init__.py deleted file mode 100644 index ef8d016..0000000 --- a/old/flow/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .base_flow import BaseFlow -from .base_tool_flow import BaseToolFlow -from .cmd_flow import CmdFlow -from .expression_tool_flow import ExpressionToolFlow -from .mock_tool_flow import MockToolFlow, MockAsyncToolFlow diff --git a/old/flow/mock_tool_flow.py b/old/flow/mock_tool_flow.py deleted file mode 100644 index 906849d..0000000 --- a/old/flow/mock_tool_flow.py +++ /dev/null @@ -1,72 +0,0 @@ -from flowllm.context.service_context import C -from flowllm.flow.base_tool_flow import BaseToolFlow -from flowllm.op.gallery.mock_op import Mock1Op, Mock2Op, Mock3Op, Mock4Op, Mock5Op, Mock6Op -from flowllm.schema.tool_call import ToolCall, ParamAttrs - - -@C.register_flow(register_app="FlowLLM") -class MockToolFlow(BaseToolFlow): - - def __init__(self, stream: bool = False, **kwargs): - super().__init__(stream=stream, **kwargs) - - def build_flow(self): - mock1_op = Mock1Op() - mock2_op = Mock2Op() - mock3_op = Mock3Op() - return mock1_op >> ((mock2_op >> mock3_op) | mock1_op) >> (mock2_op | mock3_op) - - def build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "type": "function", - "name": "mock_data", - "description": "A mock tool that processes data through multiple operations and returns structured results", - "input_schema": { - "a": ParamAttrs( - type="string", - description="The input data to be processed", - required=True, - ), - "b": ParamAttrs( - type="string", - description="Processing mode: basic, advanced, or expert", - required=False, - ), - }, - }, - ) - - -@C.register_flow(register_app="FlowLLM") -class MockAsyncToolFlow(BaseToolFlow): - - def __init__(self, stream: bool = False, **kwargs): - super().__init__(stream=stream, **kwargs) - - def build_flow(self): - mock4_op = Mock4Op() - mock5_op = Mock5Op() - mock6_op = Mock6Op() - return mock4_op >> ((mock5_op >> mock6_op) | mock4_op) >> (mock5_op | mock6_op) - - def build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "type": "function", - "name": "mock_data", - "description": "A mock tool that processes data through multiple operations and returns structured results", - "input_schema": { - "a": ParamAttrs( - type="string", - description="The input data to be processed", - required=True, - ), - "b": ParamAttrs( - type="string", - description="Processing mode: basic, advanced, or expert", - required=False, - ), - }, - }, - ) diff --git a/old/how_to_write_async_op_simple.md b/old/how_to_write_async_op_simple.md deleted file mode 100644 index 1b204e9..0000000 --- a/old/how_to_write_async_op_simple.md +++ /dev/null @@ -1,178 +0,0 @@ -## 编写异步 Op 快速入门（简单版） - -这篇文档会告诉你如何快速写一个最简单的异步 Op，包含注册、读取输入、使用 LLM 和 prompt、写输出。 - -想看更复杂的功能（缓存、并行、流式调用等），可以看[完整版文档](./how_to_write_async_op_advanced.md)。 - ---- - -### 1. 注册 Op - -首先，你需要用 `@C.register_op()` 装饰器注册你的 Op： - -```python -from flowllm.core.context import C -from flowllm.core.op import BaseAsyncOp - -@C.register_op() -class MyAsyncOp(BaseAsyncOp): - async def async_execute(self): - # 这里写你的业务逻辑 - return "ok" -``` - ---- - -### 2. 读取输入 - -Op 的输入都放在内部的 `context` 里。调用的时候直接传参数：`await op.async_call(a=1, b=2)`，这些会自动注入到 context。 - -在 `async_execute()` 里读取： - -```python -@C.register_op() -class SumOp(BaseAsyncOp): - async def async_execute(self): - # 方式1：直接用 key 访问（没有会抛 KeyError） - a = self.context["a"] - b = self.context["b"] - - # 方式2：用 get 方法（没有返回默认值） - a = self.context.get("a", 0) - b = self.context.get("b", 0) - - return a + b - -# 调用： -# result = await SumOp().async_call(a=1, b=2) -``` - ---- - -### 3. 使用 Prompt 和 LLM - -`BaseAsyncOp` 已经帮你把 Prompt 和 LLM 集成好了，直接用就行。 - -#### 3.1 Prompt 文件 - -系统会根据你的 Op 文件名自动找对应的 prompt 文件： -- 规则：把 Op 文件名的 `op.py` 后缀换成 `prompt.yaml`，放在同一个目录下 -- 比如：`my_feature_op.py` → `my_feature_prompt.yaml` - -用的时候： -- `self.prompt_format("prompt_name", **vars)` - 渲染 prompt 模板 - -#### 3.2 调用 LLM - -```python -from flowllm.core.context import C -from flowllm.core.op import BaseAsyncOp -from flowllm.core.schema import Message -from flowllm.core.enumeration import Role -from flowllm.main import FlowLLMApp - -@C.register_op() -class QAOp(BaseAsyncOp): - async def async_execute(self): - # 1. 读取输入 - question = self.context.get("question", "") - - # 2. 渲染 prompt 模板 - prompt = self.prompt_format("qa", question=question) - - # 3. 调用 LLM - messages = [Message(role=Role.USER, content=prompt)] - answer_message = await self.llm.achat(messages) - answer = answer_message.content - - # 4. 写回输出（见下一节） - self.context.response.answer = answer - return answer - -# 使用（需要在应用上下文里） -async def main(): - async with FlowLLMApp() as app: - result = await QAOp().async_call(question="你好？") - print(result) -``` - -注意：调用 LLM 需要在应用上下文 `FlowLLMApp()` 里，这样服务配置才能正确加载。 - ---- - -### 4. 写输出 - -输出有两种方式： - -```python -@C.register_op() -class QAOp(BaseAsyncOp): - async def async_execute(self): - question = self.context.get("question", "") - prompt = self.prompt_format("qa", question=question) - messages = [Message(role=Role.USER, content=prompt)] - answer_message = await self.llm.achat(messages) - answer = answer_message.content - - # 方式1：直接返回 - return answer - - # 方式2：写回 context（推荐，便于上层使用） - self.context.response.answer = answer - return answer -``` - -推荐同时使用两种方式：既返回结果，也写回 `context.response.answer`，这样上层可以直接用。 - ---- - -### 5. 完整示例 - -看一个完整的例子： - -```python -from flowllm.core.context import C -from flowllm.core.op import BaseAsyncOp -from flowllm.core.schema import Message -from flowllm.core.enumeration import Role -from flowllm.main import FlowLLMApp -import asyncio - -@C.register_op() -class SimpleChatOp(BaseAsyncOp): - async def async_execute(self): - # 1. 读取输入 - question = self.context.get("question", "") - - # 2. 渲染 prompt - prompt = self.prompt_format("chat", question=question) - - # 3. 调用 LLM - messages = [Message(role=Role.USER, content=prompt)] - answer_message = await self.llm.achat(messages) - answer = answer_message.content - - # 4. 写回输出 - self.context.response.answer = answer - return answer - -# 使用 -async def main(): - async with FlowLLMApp() as app: - result = await SimpleChatOp().async_call(question="你好？") - print(result) - -if __name__ == "__main__": - asyncio.run(main()) -``` - ---- - -### 6. 参考示例 - -想看实际例子的话，可以看看： -- `flowllm/flowllm/gallery/gen_system_prompt_op.py` - 简单的 LLM 调用示例 -- `flowllm/flowllm/gallery/mock_search_op.py` - 使用 callback_fn 的示例 - -想看更复杂的功能（缓存、并行、流式调用、工具调用等），请查看[完整版文档](./how_to_write_async_op_advanced.md)。 - diff --git a/old/op/__init__.py b/old/op/__init__.py deleted file mode 100644 index ea7be27..0000000 --- a/old/op/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -from . import crawl -from . import deep_research -from . import fin_research_v2 -from . import fin_supply -from . import gallery -from . import search -from . import fin_search -from .base_async_op import BaseAsyncOp -from .base_async_tool_op import BaseAsyncToolOp -from .base_mcp_op import BaseMcpOp -from .base_op import BaseOp -from .base_ray_op import BaseRayOp diff --git a/old/op/fin_ah/__init__.py b/old/op/fin_ah/__init__.py deleted file mode 100644 index 3f1acbe..0000000 --- a/old/op/fin_ah/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -""" -AH股策略模块 -提供完整的AH股数据处理、特征工程和回测功能 -""" - -from flowllm.op.fin_ah.ah_download_op import AhDownloadOp -from flowllm.op.fin_ah.ah_fix_op import AhFixOp -from flowllm.op.fin_ah.ah_feature_op import ( - AhFeatureTableOp, - AhDailyFeatureOp, - AhWeeklyFeatureOp, -) -from flowllm.op.fin_ah.ah_backtest_op import ( - AhBacktestTableOp, - AhBacktestOp, -) - -__all__ = [ - "AhDownloadOp", - "AhFixOp", - "AhFeatureTableOp", - "AhDailyFeatureOp", - "AhWeeklyFeatureOp", - "AhBacktestTableOp", - "AhBacktestOp", -] diff --git a/old/op/fin_ah/ah_backtest_op.py b/old/op/fin_ah/ah_backtest_op.py deleted file mode 100644 index 2a113ee..0000000 --- a/old/op/fin_ah/ah_backtest_op.py +++ /dev/null @@ -1,436 +0,0 @@ -""" -AH股回测Op -负责回测策略效果： -1. 使用Ridge回归训练模型 -2. 计算IC和RIC -3. 生成选股池（top5） -4. 保存回测中间结果和最终结果 -""" - -import os -from typing import Dict, List, Optional, Tuple -import ray -import numpy as np -import pandas as pd -from loguru import logger -from scipy.stats import pearsonr, spearmanr -from sklearn.linear_model import Ridge - -from flowllm.app import FlowLLMApp -from flowllm.context.service_context import C -from flowllm.op import BaseOp, BaseRayOp -from flowllm.utils.plot_utils import plot_figure - -TOP_5 = 5 # 选股池大小（top5） -TOP_10 = 10 # 选股池大小（top10） -BLOCK_SIZE = 10 # 分块数量（用于分析收益分布） - - -@C.register_op(register_app="FlowLLM") -class AhBacktestTableOp(BaseRayOp): - """回测策略效果（父Op）""" - - def __init__( - self, - input_dir: str = "data/feature", - output_dir: str = "data/backtest", - max_samples: int = 512, - use_weekly: bool = False, - start_date: int = 20200101, - feature_columns: List[str] = None, - label_column: str = "a_label", - label_normalization: Optional[str] = None, - **kwargs, - ): - super().__init__(**kwargs) - self.input_dir = input_dir - self.output_dir = output_dir - self.max_samples = max_samples - self.use_weekly = use_weekly - self.start_date = start_date - self.feature_columns = feature_columns - self.label_column = label_column - self.label_normalization = label_normalization - # 自动根据use_weekly设置文件后缀 - self.file_suffix = "weekly" if use_weekly else "daily" - - # 验证label_normalization参数 - if label_normalization and label_normalization not in ["mean", "median"]: - raise ValueError(f"label_normalization must be 'mean', 'median', or None, got '{label_normalization}'") - - def _ensure_output_dir(self) -> None: - """确保输出目录存在""" - os.makedirs(self.output_dir, exist_ok=True) - - def _load_feature_data(self) -> pd.DataFrame: - """加载特征数据""" - cache_name = "feature_weekly.csv" if self.use_weekly else "feature_daily.csv" - feature_path = os.path.join(self.input_dir, cache_name) - - if not os.path.exists(feature_path): - raise FileNotFoundError( - f"Feature file not found: {feature_path}\n" f"Please run AhFeatureTableOp first to generate features.", - ) - - df = pd.read_csv(feature_path) - - # 删除NaN值 - nan_count = df.isnull().sum().sum() - if nan_count > 0: - logger.warning(f"Dropping {nan_count} NaN values from feature data") - df = df.dropna() - - logger.info(f"Loaded features: {df.shape} from {feature_path}") - return df - - def _get_output_filename(self, base_filename: str) -> str: - """生成带后缀的输出文件名""" - if self.file_suffix: - name, ext = os.path.splitext(base_filename) - return f"{name}_{self.file_suffix}{ext}" - return base_filename - - def _save_dataframe(self, df: pd.DataFrame, filename: str) -> None: - """保存DataFrame到CSV""" - output_filename = self._get_output_filename(filename) - output_path = os.path.join(self.output_dir, output_filename) - df.to_csv(output_path, index=False) - logger.info(f"Saved {output_filename}: {len(df)} rows") - - def _plot_strategies(self, final_df: pd.DataFrame, dates: List[int]) -> None: - """绘制策略收益曲线""" - strategy_cols = [c for c in final_df.columns if "_uplift" in c] - if not strategy_cols: - logger.warning("No strategy columns found for plotting") - return - - logger.info(f"Plotting {len(strategy_cols)} strategies: {strategy_cols}") - plot_dict = {col: [v / 100 for v in final_df[col].tolist()] for col in strategy_cols} - plot_filename = self._get_output_filename("ah_strategy.pdf") - plot_output = os.path.join(self.output_dir, plot_filename) - plot_figure(plot_dict, output_path=plot_output, xs=[str(d) for d in dates], ticks_gap=90) - logger.info(f"Saved strategy plot to {plot_output}") - - def _normalize_labels(self, feature_df: pd.DataFrame) -> pd.DataFrame: - """对每天的label进行归一化处理（减去均值或中位数）""" - if not self.label_normalization: - return feature_df - - logger.info(f"Normalizing labels using {self.label_normalization}") - feature_df = feature_df.copy() - - # 获取所有label列（包括a_label和a_label_1d到a_label_5d） - label_cols = [col for col in feature_df.columns if col.startswith("a_label")] - - for dt in feature_df["dt"].unique(): - dt_mask = feature_df["dt"] == dt - - for label_col in label_cols: - if label_col not in feature_df.columns: - continue - - label_values = feature_df.loc[dt_mask, label_col] - - if self.label_normalization == "mean": - baseline = label_values.mean() - elif self.label_normalization == "median": - baseline = label_values.median() - else: - continue - - feature_df.loc[dt_mask, label_col] = label_values - baseline - - logger.info(f"Label normalization completed for {len(label_cols)} label columns") - return feature_df - - def execute(self) -> None: - """执行回测""" - self._ensure_output_dir() - - # 加载特征数据 - feature_df = self._load_feature_data() - - # 对label进行归一化处理 - feature_df = self._normalize_labels(feature_df) - - # 设置context参数 - self.context.max_samples = self.max_samples - self.context.feature_df = feature_df - self.context.feature_columns = self.feature_columns - self.context.label_column = self.label_column - - # 获取回测日期列表 - dt_list = sorted(feature_df.loc[feature_df.a_flag == 1, "dt"].unique()) - self.context.dt_a_list = dt_list - logger.info(f"Available dates: {dt_list[0]} to {dt_list[-1]} ({len(dt_list)} days)") - - # 筛选回测日期 - backtest_dates = [d for d in dt_list if d >= self.start_date] - logger.info(f"Backtest dates: {backtest_dates[0]} to {backtest_dates[-1]} ({len(backtest_dates)} days)") - - # 并行执行回测 - result = self.submit_and_join_parallel_op(op=self.ops[0], dt=backtest_dates) - - # 整理并保存最终结果 - final_df = pd.DataFrame([r["final"] for r in result]).sort_values("dt") - self._save_dataframe(final_df, "backtest_final.csv") - - # 整理并保存中间结果（选股池） - intermediate_records = [stock for r in result for stock in r["intermediate"]] - intermediate_df = pd.DataFrame(intermediate_records) - self._save_dataframe(intermediate_df, "backtest_pools.csv") - - # 打印IC/RIC统计 - for metric in ["model_ic", "model_ric", "rule_ic", "rule_ric"]: - mean_val = final_df[metric].mean() - std_val = final_df[metric].std() - logger.info(f"{metric}: mean={mean_val:.4f}, std={std_val:.4f}") - - # 绘制策略收益曲线 - self._plot_strategies(final_df, backtest_dates) - - logger.info(f"Backtest completed: {len(backtest_dates)} days, results in {self.output_dir}") - - -class AhBacktestOp(BaseOp): - """单日回测（子Op）""" - - @staticmethod - def _calculate_ic(pred: np.ndarray, actual: np.ndarray) -> Tuple[float, float]: - """计算IC和RIC""" - try: - ic, _ = pearsonr(pred, actual) - ric, _ = spearmanr(pred, actual) - return ic, ric - except Exception as e: - ray.logger.exception(f"Failed to calculate IC: {e}") - return 0.0, 0.0 - - @staticmethod - def _create_stock_pool_records( - df: pd.DataFrame, - dt: int, - strategy: str, - pred_col: Optional[str] = None, - ) -> List[Dict]: - """创建选股池记录""" - records = [] - for _, row in df.iterrows(): - records.append( - { - "dt": dt, - "strategy": strategy, - "name": row["name"], - "code": row["code"], - "ah_ratio": row["ah_ratio"], - "pred_value": row[pred_col] if pred_col else None, - "label": row["a_label"], - "actual_uplift": row["uplift"], - } - ) - return records - - def execute(self) -> Dict: - dt = self.context.dt - feature_df = self.context.feature_df - max_samples = self.context.max_samples - - # 准备训练集和测试集 - train_dates = [d for d in self.context.dt_a_list if d < dt][-max_samples:-1] - train_df = feature_df.loc[feature_df.dt.isin(train_dates)].copy() - test_df = feature_df.loc[feature_df.dt == dt].copy() - - if test_df.empty: - ray.logger.warning(f"No test data for dt={dt}") - return {"final": {}, "intermediate": []} - - # 获取特征列和标签列 - feature_cols = self.context.feature_columns - label_col = self.context.label_column - pred_col = "pred_y" - - # 准备数据 - train_x = train_df[feature_cols].values - train_y = train_df[label_col].values - test_x = test_df[feature_cols].values - test_y = test_df[label_col].values - rule_score = test_df["ah_ratio"].values - - # 训练Ridge回归模型 - model = Ridge(alpha=1.0) - model.fit(train_x, train_y) - - # 预测 - pred_y = model.predict(test_x) - test_df[pred_col] = pred_y - - # 计算IC和RIC - model_ic, model_ric = self._calculate_ic(pred_y, test_y) - rule_ic, rule_ric = self._calculate_ic(rule_score, test_y) - - # 生成选股池（Top 5 和 Top 10） - model_pool_5 = test_df.nlargest(TOP_5, pred_col) - model_pool_10 = test_df.nlargest(TOP_10, pred_col) - rule_pool_5 = test_df.nlargest(TOP_5, "ah_ratio") - rule_pool_10 = test_df.nlargest(TOP_10, "ah_ratio") - - # 构建最终结果（用uplift计算实际收益，用label计算IC） - final_result = { - "dt": dt, - "size": len(test_df), - "model_ic": model_ic, - "model_ric": model_ric, - "rule_ic": rule_ic, - "rule_ric": rule_ric, - # Top5 结果 - "model_names": ",".join(model_pool_5["name"].tolist()), - "model_uplift": model_pool_5["uplift"].mean(), - "rule_names": ",".join(rule_pool_5["name"].tolist()), - "rule_uplift": rule_pool_5["uplift"].mean(), - # Top10 结果 - "model_names_10": ",".join(model_pool_10["name"].tolist()), - "model_uplift_10": model_pool_10["uplift"].mean(), - "rule_names_10": ",".join(rule_pool_10["name"].tolist()), - "rule_uplift_10": rule_pool_10["uplift"].mean(), - # 全部股票平均 - "all_uplift": test_df["uplift"].mean(), - } - - # 计算分块收益（按预测值排序后分成N块，使用uplift计算实际收益） - block_count = max(1, round(len(test_df) / BLOCK_SIZE)) - test_df_sorted = test_df.sort_values(pred_col, ascending=False) - - for i in range(BLOCK_SIZE): - start_idx = i * block_count - end_idx = (i + 1) * block_count - block_data = test_df_sorted.iloc[start_idx:end_idx] - final_result[f"p{i}_uplift"] = block_data["uplift"].mean() if len(block_data) > 0 else 0.0 - - # 构建中间结果（选股池详情） - intermediate_result = [] - intermediate_result.extend(self._create_stock_pool_records(model_pool_5, dt, "model_top5", pred_col)) - intermediate_result.extend(self._create_stock_pool_records(rule_pool_5, dt, "rule_top5")) - intermediate_result.extend(self._create_stock_pool_records(model_pool_10, dt, "model_top10", pred_col)) - intermediate_result.extend(self._create_stock_pool_records(rule_pool_10, dt, "rule_top10")) - - return {"final": final_result, "intermediate": intermediate_result} - - -def main( - input_dir: str = "data/feature", - output_dir: str = "data/backtest", - max_samples: int = 512, - use_weekly: bool = False, - start_date: int = 20200101, - feature_columns: List[str] = None, - label_column: str = "a_label", - label_normalization: Optional[str] = None, - ray_workers: int = 8, -): - """运行AH股回测 - - Args: - input_dir: 特征数据目录 - output_dir: 回测结果输出目录 - max_samples: 训练样本数量 - use_weekly: 是否使用周频数据 - start_date: 回测开始日期 - feature_columns: 使用的特征列表（必须指定） - label_column: 使用的标签列（默认a_label，可选a_label_1d, a_label_2d, a_label_3d, a_label_4d, a_label_5d） - label_normalization: label归一化方式（None/mean/median），每天的label减去所有股票的均值或中位数 - ray_workers: Ray并行workers数量 - """ - with FlowLLMApp(load_default_config=True) as app: - app.service_config.ray_max_workers = ray_workers - - op = ( - AhBacktestTableOp( - input_dir=input_dir, - output_dir=output_dir, - max_samples=max_samples, - use_weekly=use_weekly, - start_date=start_date, - feature_columns=feature_columns, - label_column=label_column, - label_normalization=label_normalization, - ) - << AhBacktestOp() - ) - - op.call() - - -if __name__ == "__main__": - # 日频回测配置 - 默认使用5天标签 - daily_config = { - "input_dir": "data/feature", - "output_dir": "data/backtest", - "max_samples": 512, - "use_weekly": False, - "start_date": 20200101, - "feature_columns": [ - "ah_ratio", - # "ah_amount", - # "avg_1d_ah_pct_diff", - # "avg_5d_ah_pct_diff", - # "avg_1d_ah_amount_ratio", - # "avg_5d_ah_amount_ratio", - "avg_1d_a_pct", - "avg_1d_hk_pct", - "avg_5d_a_pct", - "avg_5d_hk_pct", - # "avg_20d_a_pct", - # "avg_20d_hk_pct", - "avg_1d_a_amount", - "avg_1d_hk_amount", - # "avg_5d_a_amount", - # "avg_5d_hk_amount", - ], - "label_column": "a_label_5d", # 可选: a_label_1d, a_label_2d, a_label_3d, a_label_4d, a_label_5d - "label_normalization": "mean", # None / "mean" / "median" - "ray_workers": 8, - } - - # 周频回测配置 - weekly_config = { - "input_dir": "data/feature", - "output_dir": "data/backtest", - "max_samples": 512, - "use_weekly": True, - "start_date": 20200101, - "feature_columns": [ - "ah_ratio", - "ah_amount5", - "avg_3d_a_pct", - "avg_3d_hk_pct", - "avg_10d_a_pct", - "avg_10d_hk_pct", - ], - "label_column": "a_label", # 周频使用默认label - "label_normalization": "mean", # None / "mean" / "median" - "ray_workers": 8, - } - - # 执行日频回测 - main(**daily_config) - - # 执行周频回测 - # main(**weekly_config) - - # 示例: 使用1天标签进行回测 - # daily_1d_config = daily_config.copy() - # daily_1d_config["label_column"] = "a_label_1d" - # daily_1d_config["output_dir"] = "data/backtest_1d" - # main(**daily_1d_config) - - # 示例: 使用均值归一化的回测 - # daily_mean_config = daily_config.copy() - # daily_mean_config["label_normalization"] = "mean" - # daily_mean_config["output_dir"] = "data/backtest_mean" - # main(**daily_mean_config) - - # 示例: 使用中位数归一化的回测 - # daily_median_config = daily_config.copy() - # daily_median_config["label_normalization"] = "median" - # daily_median_config["output_dir"] = "data/backtest_median" - # main(**daily_median_config) diff --git a/old/op/fin_ah/ah_feature_op.py b/old/op/fin_ah/ah_feature_op.py deleted file mode 100644 index 9e6c8f5..0000000 --- a/old/op/fin_ah/ah_feature_op.py +++ /dev/null @@ -1,537 +0,0 @@ -""" -AH股特征标签宽表Op -负责生成特征和标签的大宽表： -1. 计算AH比价 -2. 计算成交额比 -3. 计算历史涨跌幅 -4. 计算A/H涨跌幅的差值 -5. 计算A/H成交额的比值 -6. 生成未来收益标签 -支持日频和周频两种模式 -""" - -import os -from typing import Dict, List, Optional - -import pandas as pd -import ray -from loguru import logger - -from flowllm.app import FlowLLMApp -from flowllm.context.service_context import C -from flowllm.op import BaseOp, BaseRayOp -from flowllm.utils.common_utils import find_dt_less_index, get_monday_fridays, next_friday_or_same - -HISTORY_WINDOWS = [1, 3, 5, 10, 15, 20, 30, 60] # 历史涨跌幅和成交额窗口 - - -@C.register_op(register_app="FlowLLM") -class AhFeatureTableOp(BaseRayOp): - """生成AH股特征标签大宽表（父Op）""" - - def __init__( - self, - input_dir: str = "data/fixed", - output_dir: str = "data/feature", - use_weekly: bool = False, - **kwargs, - ): - - super().__init__(**kwargs) - self.input_dir = input_dir - self.output_dir = output_dir - self.use_weekly = use_weekly - - def _ensure_output_dir(self) -> None: - """确保输出目录存在""" - os.makedirs(self.output_dir, exist_ok=True) - - def _load_ah_comparison(self) -> pd.DataFrame: - """读取AH对比数据""" - # 优先从fixed目录读取，否则从origin目录 - for dir_name in [self.input_dir, "data/origin"]: - path = os.path.join(dir_name, "stk_ah_comparison.csv") - if os.path.exists(path): - df = pd.read_csv(path) - ah_df = df.loc[df.trade_date == df.trade_date.max(), ["hk_code", "ts_code", "name"]].copy() - logger.info(f"Loaded {len(ah_df)} AH pairs from {path}") - return ah_df - raise FileNotFoundError("stk_ah_comparison.csv not found") - - def _load_data_from_files(self) -> None: - """从文件加载数据（支持独立运行）""" - logger.info(f"Loading data from {self.input_dir}...") - - # 1. 读取AH对比数据 - ah_df = self._load_ah_comparison() - - # 2. 读取汇率比率 - forex_ratio_path = os.path.join(self.input_dir, "hk_forex_ratio.csv") - hk_forex_df = pd.read_csv(forex_ratio_path, index_col=0) - logger.info(f"Loaded forex ratio: {len(hk_forex_df)} rows") - - # 3. 读取股票数据 - stock_dict = {} - a_date_counter = {} - hk_date_counter = {} - - # 使用列表推导式而不是tqdm迭代器 - records = ah_df.to_dict(orient="records") - total = len(records) - logger.info(f"Loading {total} stock pairs...") - - for idx, record in enumerate(records): - hk_code, ts_code, name = record["hk_code"], record["ts_code"], record["name"] - - if (idx + 1) % 10 == 0: - logger.info(f"Loading progress: {idx + 1}/{total}") - - try: - # 读取A股和HK股数据 - a_df = pd.read_csv(os.path.join(self.input_dir, f"daily_{ts_code}.csv")) - hk_df = pd.read_csv(os.path.join(self.input_dir, f"hk_daily_{hk_code}.csv")) - - # 获取日期列表并对齐 - hk_dates = sorted(hk_df["trade_date"].unique()) - min_hk_date = min(hk_dates) - a_dates = sorted([d for d in a_df["trade_date"].unique() if d >= min_hk_date]) - - stock_dict[name] = { - "a_code": ts_code, - "hk_code": hk_code, - "a_org_stock_df": a_df, - "hk_org_stock_df": hk_df, - "hk_dt_list": hk_dates, - "a_dt_list": a_dates, - } - - # 统计日期覆盖 - for dt in a_dates: - a_date_counter[dt] = a_date_counter.get(dt, 0) + 1 - for dt in hk_dates: - hk_date_counter[dt] = hk_date_counter.get(dt, 0) + 1 - - except Exception as e: - logger.exception(f"Failed to load {name} ({ts_code}/{hk_code}): {e}") - continue - - logger.info(f"Loaded {len(stock_dict)} stock pairs") - - # 存入context - self.context.ah_df = ah_df - self.context.hk_forex_df = hk_forex_df - self.context.stock_dict = stock_dict - self.context.dt_a_list = sorted(a_date_counter.keys()) - self.context.dt_hk_list = sorted(hk_date_counter.keys()) - - def _prepare_weekly_dates(self) -> None: - """准备周频日期列表（每周最后一个交易日）""" - weekly_dates = [] - monday_fridays = get_monday_fridays(self.context.dt_a_list[0], self.context.dt_a_list[-1]) - - for monday, friday in monday_fridays: - week_dates = [d for d in self.context.dt_a_list if monday <= str(d) <= friday] - if week_dates: - week_dates = sorted(week_dates) - weekly_dates.append(week_dates[-1]) # 取本周最后一个交易日 - - logger.info(f"Generated {len(weekly_dates)} weekly dates") - self.context.dt_a_weekly_list = weekly_dates - - def execute(self) -> None: - """执行特征生成""" - self._ensure_output_dir() - - # 从文件加载数据 - self._load_data_from_files() - - # 设置参数到context - self.context.use_weekly = self.use_weekly - - # 选择日频或周频模式 - if self.use_weekly: - self._prepare_weekly_dates() - f_op = self.ops[1] # AhWeeklyFeatureOp - dt_list = self.context.dt_a_weekly_list - cache_name = "feature_weekly.csv" - else: - f_op = self.ops[0] # AhDailyFeatureOp - dt_list = self.context.dt_a_list - cache_name = "feature_daily.csv" - - mode = "weekly" if self.use_weekly else "daily" - logger.info(f"Generating {mode} features for {len(dt_list)} dates") - - # 并行执行特征生成 - result = self.submit_and_join_parallel_op(op=f_op, dt=dt_list) - - # 合并并排序结果 - df = pd.DataFrame(result).sort_values(["dt", "code"]) - - # 验证数据完整性 - nan_count = df.isnull().sum().sum() - if nan_count > 0: - logger.warning(f"Feature table has {nan_count} NaN values") - - # 保存特征表 - output_path = os.path.join(self.output_dir, cache_name) - df.to_csv(output_path, index=False) - logger.info(f"Saved {mode} features: {df.shape} to {output_path}") - - -class AhDailyFeatureOp(BaseOp): - """生成日频特征（子Op）""" - - @staticmethod - def get_forex_ratio(dt: int, hk_forex_df: pd.DataFrame) -> Optional[float]: - """获取指定日期的汇率比率""" - forex_dates = [d for d in hk_forex_df.index if d < dt] - if not forex_dates: - return None - - forex_dt = forex_dates[find_dt_less_index(dt, forex_dates)] - return hk_forex_df.loc[forex_dt, "close"] - - @staticmethod - def calculate_history_features( - df: pd.DataFrame, - dt_list: List[int], - dt_index: int, - windows: List[int] = None, - ) -> Dict[str, float]: - """计算历史累积涨跌幅特征（复合收益率）""" - if windows is None: - windows = HISTORY_WINDOWS - - features = {} - for window in windows: - start_idx = max(0, dt_index - (window - 1)) - date_slice = dt_list[start_idx : dt_index + 1] - matched = df.loc[df.index.isin(date_slice), "pct_chg"] - - # 计算累积收益率：(1 + r1) * (1 + r2) * ... - 1 - if len(matched) > 0: - cumulative_return = 1.0 - for pct in matched: - cumulative_return *= 1 + pct / 100 - features[f"avg_{window}d"] = (cumulative_return - 1) * 100 - else: - features[f"avg_{window}d"] = 0 - return features - - @staticmethod - def calculate_history_amount_features( - df: pd.DataFrame, - dt_list: List[int], - dt_index: int, - windows: List[int] = None, - ) -> Dict[str, float]: - """计算历史成交额特征""" - if windows is None: - windows = HISTORY_WINDOWS - - features = {} - for window in windows: - start_idx = max(0, dt_index - (window - 1)) - date_slice = dt_list[start_idx : dt_index + 1] - matched = df.loc[df.index.isin(date_slice), "amount"] - features[f"avg_{window}d_amount"] = matched.mean() if len(matched) > 0 else 0 - return features - - @staticmethod - def _calculate_single_day_uplift( - df: pd.DataFrame, - dt_list: List[int], - dt_index: int, - day_offset: int = 1, - ) -> float: - """计算单日uplift：从某天开盘到下一天开盘的收益 - - Args: - df: 股票数据 - dt_list: 日期列表 - dt_index: 当前日期索引 - day_offset: 从当前日期往后偏移几天（1表示明天开盘到后天开盘） - - Returns: - 收益率（百分比） - """ - start_idx = dt_index + day_offset - end_idx = dt_index + day_offset + 1 - - if start_idx >= len(dt_list) or end_idx >= len(dt_list): - return 0.0 - - start_open = df.loc[dt_list[start_idx], "open"] - end_open = df.loc[dt_list[end_idx], "open"] - - return (end_open / start_open - 1) * 100 - - @staticmethod - def _calculate_avg_future_uplift( - df: pd.DataFrame, - dt_list: List[int], - dt_index: int, - n_days: int = 5, - ) -> float: - """计算未来N天的平均uplift（累积收益率除以天数） - - Args: - df: 股票数据 - dt_list: 日期列表 - dt_index: 当前日期索引 - n_days: 未来N天 - - Returns: - 未来N天累积uplift除以天数 - """ - # 计算累积收益率：(1 + r1) * (1 + r2) * ... * (1 + rN) - 1 - cumulative_return = 1.0 - valid_days = 0 - - for i in range(1, n_days + 1): - if dt_index + i + 1 >= len(dt_list): - break - - uplift = AhDailyFeatureOp._calculate_single_day_uplift(df, dt_list, dt_index, day_offset=i) - cumulative_return *= 1 + uplift / 100 - valid_days += 1 - - if valid_days == 0: - return 0.0 - - # 累积收益率除以实际天数 - total_return = (cumulative_return - 1) * 100 - return total_return / valid_days - - def execute(self): - result = [] - dt = self.context.dt - hk_forex_df = self.context.hk_forex_df - - # 获取汇率 - hk_forex_ratio = self.get_forex_ratio(dt, hk_forex_df) - if hk_forex_ratio is None: - raise RuntimeError(f"No forex data for dt={dt}") - - for name, stock_info in self.context.stock_dict.items(): - a_dt_list = stock_info["a_dt_list"] - if dt not in a_dt_list: - continue - - hk_dt_list = stock_info["hk_dt_list"] - dt_a_index = a_dt_list.index(dt) - dt_hk_index = find_dt_less_index(dt, hk_dt_list) - - # 检查数据充足性 - if dt_a_index < 1 or dt_hk_index is None or dt_hk_index < 1: - continue - - a_df = stock_info["a_org_stock_df"].set_index("trade_date") - hk_df = stock_info["hk_org_stock_df"].set_index("trade_date") - - # 当前价格和成交额 - a_curr_dt = a_dt_list[dt_a_index] - hk_curr_dt = hk_dt_list[dt_hk_index] - - current_a_close = a_df.loc[a_curr_dt, "close"] - current_hk_close = hk_df.loc[hk_curr_dt, "close"] - current_a_amount = a_df.loc[a_curr_dt, "amount"] - current_hk_amount = hk_df.loc[hk_curr_dt, "amount"] - - # 计算uplift：明天开盘到后天开盘的涨跌幅 - next_a_uplift = self._calculate_single_day_uplift(a_df, a_dt_list, dt_a_index, day_offset=1) - - # 计算多个label：未来1/2/3/4/5天uplift的平均值 - labels = {} - for n_days in [1, 2, 3, 4, 5]: - labels[f"a_label_{n_days}d"] = self._calculate_avg_future_uplift( - a_df, - a_dt_list, - dt_a_index, - n_days=n_days, - ) - - # 构建特征字典 - feature = { - "dt": dt, - "name": name, - "code": f"{stock_info['a_code']}+{stock_info['hk_code']}", - "current_a_close": current_a_close, - "current_hk_close": current_hk_close, - "hk_forex_ratio": hk_forex_ratio, - "a_flag": 1, - "hk_flag": 1, - "a_uplift": a_df.loc[a_curr_dt, "pct_chg"], - "hk_uplift": hk_df.loc[hk_curr_dt, "pct_chg"], - "ah_ratio": current_hk_close * hk_forex_ratio / current_a_close, - "a_label": labels["a_label_5d"], # 保持向后兼容，默认使用5天 - "uplift": next_a_uplift, - } - - # 添加多个label列 - feature.update(labels) - - # 添加历史涨跌幅特征 - a_hist = self.calculate_history_features(a_df, a_dt_list, dt_a_index) - hk_hist = self.calculate_history_features(hk_df, hk_dt_list, dt_hk_index) - - for window in HISTORY_WINDOWS: - feature[f"avg_{window}d_a_pct"] = a_hist[f"avg_{window}d"] - feature[f"avg_{window}d_hk_pct"] = hk_hist[f"avg_{window}d"] - # 添加涨跌幅的差值（A股涨跌幅 - H股涨跌幅） - feature[f"avg_{window}d_ah_pct_diff"] = a_hist[f"avg_{window}d"] - hk_hist[f"avg_{window}d"] - - # 添加历史成交额特征 - a_amount_hist = self.calculate_history_amount_features(a_df, a_dt_list, dt_a_index) - hk_amount_hist = self.calculate_history_amount_features(hk_df, hk_dt_list, dt_hk_index) - - for window in HISTORY_WINDOWS: - feature[f"avg_{window}d_a_amount"] = a_amount_hist[f"avg_{window}d_amount"] - feature[f"avg_{window}d_hk_amount"] = hk_amount_hist[f"avg_{window}d_amount"] - # 添加成交额的比值（H股成交额 / A股成交额） - a_amount = a_amount_hist[f"avg_{window}d_amount"] - feature[f"avg_{window}d_ah_amount_ratio"] = ( - hk_amount_hist[f"avg_{window}d_amount"] / a_amount if a_amount != 0 else 0 - ) - - result.append(feature) - - return result - - -class AhWeeklyFeatureOp(BaseOp): - """生成周频特征（子Op）""" - - @staticmethod - def _calculate_weekly_return(df: pd.DataFrame, label_dates: List[int]) -> float: - """计算周收益：第一天用开盘价，后续天累乘""" - if not label_dates: - return 0.0 - - # 第一天：open -> close - ratio = df.loc[label_dates[0], "close"] / df.loc[label_dates[0], "open"] - - # 后续天：累乘pct_chg - for dt in label_dates[1:]: - ratio *= 1 + df.loc[dt, "pct_chg"] / 100 - - return (ratio - 1) * 100 - - def execute(self): - result = [] - dt = self.context.dt - hk_forex_df = self.context.hk_forex_df - - # 获取汇率 - hk_forex_ratio = AhDailyFeatureOp.get_forex_ratio(dt, hk_forex_df) - if hk_forex_ratio is None: - ray.logger.warning(f"No forex data for dt={dt}") - return result - - for name, stock_info in self.context.stock_dict.items(): - a_dt_list = stock_info["a_dt_list"] - if dt not in a_dt_list: - continue - - hk_dt_list = stock_info["hk_dt_list"] - dt_a_index = a_dt_list.index(dt) - dt_hk_index = find_dt_less_index(dt, hk_dt_list) - - # 检查数据充足性 - if dt_a_index < 1 or dt_hk_index is None or dt_hk_index < 1: - continue - - # 计算下周日期范围（用于标签） - start_dt = a_dt_list[dt_a_index + 1] if dt_a_index < len(a_dt_list) - 1 else dt - end_dt = int(next_friday_or_same(str(start_dt))) - label_dates = [d for d in a_dt_list if start_dt <= d <= end_dt] - - if not label_dates or start_dt > end_dt: - continue - - a_df = stock_info["a_org_stock_df"].set_index("trade_date") - hk_df = stock_info["hk_org_stock_df"].set_index("trade_date") - - # 当前价格和成交额 - a_curr_dt = a_dt_list[dt_a_index] - hk_curr_dt = hk_dt_list[dt_hk_index] - - current_a_close = a_df.loc[a_curr_dt, "close"] - current_hk_close = hk_df.loc[hk_curr_dt, "close"] - current_a_amount = a_df.loc[a_curr_dt, "amount"] - current_hk_amount = hk_df.loc[hk_curr_dt, "amount"] - - # 计算下周收益：从下周开盘到下周收盘的涨跌幅 - # label和uplift保持一致（用于训练和回测） - next_week_return = self._calculate_weekly_return(a_df, label_dates) - - # 构建特征字典 - feature = { - "dt": dt, - "name": name, - "code": f"{stock_info['a_code']}+{stock_info['hk_code']}", - "current_a_close": current_a_close, - "current_hk_close": current_hk_close, - "hk_forex_ratio": hk_forex_ratio, - "a_flag": 1, - "hk_flag": 1, - "a_uplift": a_df.loc[a_curr_dt, "pct_chg"], - "hk_uplift": hk_df.loc[hk_curr_dt, "pct_chg"], - "ah_ratio": current_hk_close * hk_forex_ratio / current_a_close, - "a_label": next_week_return, - "uplift": next_week_return, - } - - # 添加历史涨跌幅特征 - a_hist = AhDailyFeatureOp.calculate_history_features(a_df, a_dt_list, dt_a_index) - hk_hist = AhDailyFeatureOp.calculate_history_features(hk_df, hk_dt_list, dt_hk_index) - - for window in HISTORY_WINDOWS: - feature[f"avg_{window}d_a_pct"] = a_hist[f"avg_{window}d"] - feature[f"avg_{window}d_hk_pct"] = hk_hist[f"avg_{window}d"] - # 添加涨跌幅的差值（A股涨跌幅 - H股涨跌幅） - feature[f"avg_{window}d_ah_pct_diff"] = a_hist[f"avg_{window}d"] - hk_hist[f"avg_{window}d"] - - # 添加历史成交额特征 - a_amount_hist = AhDailyFeatureOp.calculate_history_amount_features(a_df, a_dt_list, dt_a_index) - hk_amount_hist = AhDailyFeatureOp.calculate_history_amount_features(hk_df, hk_dt_list, dt_hk_index) - - for window in HISTORY_WINDOWS: - feature[f"avg_{window}d_a_amount"] = a_amount_hist[f"avg_{window}d_amount"] - feature[f"avg_{window}d_hk_amount"] = hk_amount_hist[f"avg_{window}d_amount"] - # 添加成交额的比值（H股成交额 / A股成交额） - a_amount = a_amount_hist[f"avg_{window}d_amount"] - feature[f"avg_{window}d_ah_amount_ratio"] = ( - hk_amount_hist[f"avg_{window}d_amount"] / a_amount if a_amount != 0 else 0 - ) - - result.append(feature) - - return result - - -def main(use_weekly: bool): - """生成AH股特征标签""" - with FlowLLMApp(load_default_config=True) as app: - app.service_config.ray_max_workers = 8 - - op = ( - AhFeatureTableOp( - input_dir="data/fixed", - output_dir="data/feature", - use_weekly=use_weekly, - ) - << AhDailyFeatureOp() - << AhWeeklyFeatureOp() - ) - - op.call() - - -if __name__ == "__main__": - # 生成日频特征 - main(use_weekly=False) - - # 生成周频特征 - # main(use_weekly=True) diff --git a/old/op/fin_ah/ah_op_old.py b/old/op/fin_ah/ah_op_old.py deleted file mode 100644 index 6a9f35e..0000000 --- a/old/op/fin_ah/ah_op_old.py +++ /dev/null @@ -1,533 +0,0 @@ -from typing import List - -import numpy as np -import pandas as pd -from loguru import logger -from scipy.stats import spearmanr, pearsonr -from sklearn.linear_model import Ridge -from tqdm import tqdm - -from flowllm.context.service_context import C -from flowllm.op import BaseOp, BaseRayOp -from flowllm.utils.common_utils import find_dt_less_index, get_monday_fridays, next_friday_or_same -from flowllm.utils.plot_utils import plot_figure -from flowllm.utils.tushare_client import TushareClient - - -@C.register_op(register_app="FlowLLM") -class AhOp(BaseRayOp): - - def __init__( - self, - enable_cache: bool = True, - cache_expire_hours: float = 24, - max_samples: int = 512, - use_open: bool = True, - use_weekly: bool = True, - **kwargs, - ): - super().__init__(enable_cache=enable_cache, cache_expire_hours=cache_expire_hours, **kwargs) - self.max_samples: int = max_samples - self.use_open: bool = use_open - self.use_weekly: bool = use_weekly - self.ts_client = TushareClient() - - @staticmethod - def fix_df(df: pd.DataFrame) -> pd.DataFrame: - flag_index = (df["pre_close"].isna()) | (df["pre_close"] == 0.0) - df.loc[flag_index, "pre_close"] = df["close"].shift(-1) - df.loc[flag_index, "change"] = df["close"] - df["pre_close"] - df.loc[flag_index, "pct_chg"] = (df["close"] / df["pre_close"] - 1) * 100 - df = df[:-1].copy() - return df - - def prepare_raw_data(self): - logger.info("start") - df = self.ts_client.request(api_name="stk_ah_comparison") - ah_df = df.loc[df.trade_date == df.trade_date.max(), ["hk_code", "ts_code", "name"]] - self.context.ah_df = ah_df - logger.info(f"ah_df=\n{ah_df}") - - df_cnh: pd.DataFrame = self.ts_client.request(api_name="fx_daily", ts_code="USDCNH.FXCM") - df_hkd: pd.DataFrame = self.ts_client.request(api_name="fx_daily", ts_code="USDHKD.FXCM") - df_cnh = ( - df_cnh.loc[df_cnh.trade_date > 20150101, ["trade_date", "bid_close"]] - .set_index("trade_date") - .rename( - columns={"bid_close": "cnh_close"}, - ) - ) - df_hkd = ( - df_hkd.loc[df_hkd.trade_date > 20150101, ["trade_date", "bid_close"]] - .set_index("trade_date") - .rename( - columns={"bid_close": "hkd_close"}, - ) - ) - hk_forex_df = df_cnh.join(df_hkd).ffill() - nan_size = hk_forex_df.isnull().sum().sum() - assert nan_size == 0, f"hk_forex_df has null. size={nan_size}" - hk_forex_df.loc[:, "close"] = hk_forex_df.loc[:, "cnh_close"] / hk_forex_df.loc[:, "hkd_close"] - self.context.hk_forex_df = hk_forex_df - - self.context.stock_dict = {} - a_dt_size_dict = {} - hk_dt_size_dict = {} - for line in tqdm(ah_df.to_dict(orient="records")): - hk_code = line["hk_code"] - ts_code = line["ts_code"] - name = line["name"] - - # logger.info(f"hk_code={hk_code} ts_code={ts_code} name={name}") - a_org_stock_df = self.ts_client.request(api_name="daily", ts_code=ts_code) - hk_org_stock_df = self.ts_client.request(api_name="hk_daily", ts_code=hk_code) - hk_org_stock_df = self.fix_df(hk_org_stock_df) - - a_org_stock_df = a_org_stock_df.loc[a_org_stock_df.trade_date > 20160101, :] - hk_org_stock_df = hk_org_stock_df.loc[hk_org_stock_df.trade_date > 20160101, :] - self.cache.save(f"{hk_code}", hk_org_stock_df) - - nan_size = a_org_stock_df.isnull().sum().sum() - assert nan_size == 0, f"{name}.{ts_code}.a_org_stock_df has null. size={nan_size}" - - nan_size = hk_org_stock_df.isnull().sum().sum() - assert nan_size == 0, f"{name}.{hk_code}.hk_org_stock_df has null. size={nan_size}" - - hk_dt_list = sorted(hk_org_stock_df.loc[:, "trade_date"].unique()) - min_hk_dt = min(hk_dt_list) - a_dt_list = sorted(a_org_stock_df.loc[:, "trade_date"].unique()) - a_dt_list = [x for x in a_dt_list if x >= min_hk_dt] - - self.context.stock_dict[name] = { - "a_code": ts_code, - "hk_code": hk_code, - "a_org_stock_df": a_org_stock_df, - "hk_org_stock_df": hk_org_stock_df, - "hk_dt_list": hk_dt_list, - "a_dt_list": a_dt_list, - } - - for dt in a_dt_list: - if dt not in a_dt_size_dict: - a_dt_size_dict[dt] = 0 - a_dt_size_dict[dt] += 1 - - for dt in hk_dt_list: - if dt not in hk_dt_size_dict: - hk_dt_size_dict[dt] = 0 - hk_dt_size_dict[dt] += 1 - - for dt, cnt in sorted(a_dt_size_dict.items(), key=lambda x: x[0]): - logger.info(f"{dt}: {cnt}") - - for dt, cnt in sorted(hk_dt_size_dict.items(), key=lambda x: x[0]): - logger.info(f"{dt}: {cnt}") - - self.context.dt_a_list = sorted(a_dt_size_dict.keys()) - self.context.dt_hk_list = sorted(hk_dt_size_dict.keys()) - - if self.use_weekly: - dt_a_weekly_list = [] - monday_friday_list: List[List[str]] = get_monday_fridays( - self.context.dt_a_list[0], - self.context.dt_a_list[-1], - ) - for monday_friday in monday_friday_list: - monday = str(monday_friday[0]) - friday = str(monday_friday[1]) - dt_a_weekly = [x for x in self.context.dt_a_list if monday <= str(x) <= friday] - if dt_a_weekly: - dt_a_weekly_list.append(dt_a_weekly[-1]) - logger.info(f"dt_a_weekly_list={dt_a_weekly_list}") - self.context.dt_a_weekly_list = dt_a_weekly_list - - def prepare_feature_data(self): - self.context.use_open = self.use_open - self.context.use_weekly = self.use_weekly - - if self.use_weekly: - f_op = self.ops[1] - dt_list = self.context.dt_a_weekly_list - cache_name = "feature_w_df" - else: - f_op = self.ops[0] - dt_list = self.context.dt_a_list - cache_name = "feature_df" - - result = self.submit_and_join_parallel_op(op=f_op, dt=dt_list) - - df = pd.DataFrame(result).sort_values(["dt", "code"]) - self.cache.save(cache_name, df) - - def prepare_backtest(self): - logger.info("start backtest") - b_op = self.ops[2] - self.context.max_samples = self.max_samples - self.context.use_open = self.use_open - - if self.use_weekly: - cache_name = "feature_w_df" - else: - cache_name = "feature_df" - feature_df = self.cache.load(cache_name) - nan_size = feature_df.isnull().sum().sum() - assert nan_size == 0, f"feature_df has null. size={nan_size}" - - self.context.feature_df = feature_df - logger.info(f"feature_df.shape={feature_df.shape}") - dt_a_list = sorted(set(feature_df.loc[feature_df.a_flag == 1, "dt"])) - self.context.dt_a_list = dt_a_list - logger.info(f"dt_a_list={dt_a_list[0]}...{dt_a_list[-1]} size={len(dt_a_list)}") - - dts = [x for x in dt_a_list if x >= 20200101] - result = self.submit_and_join_parallel_op(op=b_op, dt=dts) - - df = pd.DataFrame(result).sort_values(["dt"]) - self.cache.save("backtest_df", df) - - for key in ["model_ic", "model_ric", "rule_ic", "rule_ric"]: - logger.info(f"{key} mean={df[key].mean()} std={df[key].std()}") - - strategy_list = [x for x in df.columns.tolist() if "_uplift" in x] - logger.info(f"find strategy_list={strategy_list}") - - plot_dict = {x: [x / 100 for x in df.loc[:, x].tolist()] for x in strategy_list} - plot_figure(plot_dict, output_path="ah_strategy.pdf", xs=[str(x) for x in dts], ticks_gap=90) - - def execute(self): - self.prepare_raw_data() - self.prepare_feature_data() - - self.prepare_backtest() - - -class AhFeatureOp(BaseOp): - - def execute(self): - # actor_index = self.context.actor_index - result = [] - dt = self.context.dt - use_open: bool = self.context.use_open - hk_forex_df = self.context.hk_forex_df - # 晚一天时间 - hk_forex_dt_list = [x for x in sorted(hk_forex_df.index.unique()) if x < dt] - hk_forex_dt = hk_forex_dt_list[find_dt_less_index(dt, hk_forex_dt_list)] - hk_forex_ratio = hk_forex_df.loc[hk_forex_dt, "close"] - - for name, stock_info in self.context.stock_dict.items(): - t_result = { - "dt": dt, - "name": name, - "code": stock_info["a_code"] + "+" + stock_info["hk_code"], - } - - a_org_stock_df = stock_info["a_org_stock_df"].set_index("trade_date") - hk_org_stock_df = stock_info["hk_org_stock_df"].set_index("trade_date") - - a_dt_list = stock_info["a_dt_list"] - if dt not in a_dt_list: - continue - - hk_dt_list = stock_info["hk_dt_list"] - - dt_a_index = a_dt_list.index(dt) - dt_hk_index = find_dt_less_index(dt, hk_dt_list) - - if dt_a_index is None or len(a_dt_list[: dt_a_index + 1]) < 2: - continue - - if dt_hk_index is None or len(hk_dt_list[: dt_hk_index + 1]) < 2: - continue - - current_a_close = a_org_stock_df.loc[a_dt_list[dt_a_index], "close"] - current_hk_close = hk_org_stock_df.loc[hk_dt_list[dt_hk_index], "close"] - - current_a_amount = a_org_stock_df.loc[a_dt_list[dt_a_index], "amount"] - current_hk_amount = hk_org_stock_df.loc[hk_dt_list[dt_hk_index], "amount"] - - current_a_uplift = a_org_stock_df.loc[a_dt_list[dt_a_index], "pct_chg"] - current_hk_uplift = hk_org_stock_df.loc[hk_dt_list[dt_hk_index], "pct_chg"] - - if use_open: - if dt_a_index < len(a_dt_list) - 2: - # next_a_index = dt_a_index + 2 - t1_open_close_ratio = ( - a_org_stock_df.loc[a_dt_list[dt_a_index + 1], "close"] - / a_org_stock_df.loc[ - a_dt_list[dt_a_index + 1], - "open", - ] - ) - t2_open_close_ratio = ( - a_org_stock_df.loc[a_dt_list[dt_a_index + 2], "close"] - / a_org_stock_df.loc[ - a_dt_list[dt_a_index + 2], - "open", - ] - ) - next_a_uplift = a_org_stock_df.loc[a_dt_list[dt_a_index + 2], "pct_chg"] - next_a_uplift = ((1 + next_a_uplift / 100) / t2_open_close_ratio * t1_open_close_ratio - 1) * 100 - - else: - # next_a_index = dt_a_index - next_a_uplift = 0 - - if dt_hk_index < len(hk_dt_list) - 2: - next_hk_index = dt_hk_index + 2 - t1_open_close_ratio = ( - hk_org_stock_df.loc[hk_dt_list[dt_hk_index + 1], "close"] - / hk_org_stock_df.loc[hk_dt_list[dt_hk_index + 1], "open"] - ) - t2_open_close_ratio = ( - hk_org_stock_df.loc[hk_dt_list[dt_hk_index + 2], "close"] - / hk_org_stock_df.loc[hk_dt_list[dt_hk_index + 2], "open"] - ) - next_hk_uplift = hk_org_stock_df.loc[hk_dt_list[next_hk_index], "pct_chg"] - next_hk_uplift = ((1 + next_hk_uplift / 100) / t2_open_close_ratio * t1_open_close_ratio - 1) * 100 - - else: - # next_hk_index = dt_hk_index - next_hk_uplift = 0 - - else: - if dt_a_index < len(a_dt_list) - 1: - next_a_index = dt_a_index + 1 - else: - next_a_index = dt_a_index - next_a_uplift = a_org_stock_df.loc[a_dt_list[next_a_index], "pct_chg"] - - if dt_hk_index < len(hk_dt_list) - 1: - next_hk_index = dt_hk_index + 1 - else: - next_hk_index = dt_hk_index - next_hk_uplift = hk_org_stock_df.loc[hk_dt_list[next_hk_index], "pct_chg"] - - t_result["current_hk_close"] = current_hk_close - t_result["current_a_close"] = current_a_close - t_result["hk_forex_ratio"] = hk_forex_ratio - - t_result["a_flag"] = 1 if dt in a_dt_list else 0 - t_result["hk_flag"] = 1 if dt in hk_dt_list else 0 - - t_result["a_uplift"] = current_a_uplift - t_result["hk_uplift"] = current_hk_uplift - - for i in [1, 3, 5, 10, 20]: - a_start_idx = max(0, dt_a_index - (i - 1)) - hk_start_idx = max(0, dt_hk_index - (i - 1)) - - a_date_slice = a_dt_list[a_start_idx : dt_a_index + 1] - hk_date_slice = hk_dt_list[hk_start_idx : dt_hk_index + 1] - - a_matched = hk_org_stock_df.loc[hk_org_stock_df.index.isin(a_date_slice), "pct_chg"] - hk_matched = hk_org_stock_df.loc[hk_org_stock_df.index.isin(hk_date_slice), "pct_chg"] - - t_result[f"avg_{i}d_a_pct"] = a_matched.mean() if len(a_matched) > 0 else 0 - t_result[f"avg_{i}d_hk_pct"] = hk_matched.mean() if len(hk_matched) > 0 else 0 - - t_result["ah_amount"] = (current_hk_amount / current_a_amount - 1) if current_a_amount != 0 else 0 - t_result["ah_ratio"] = current_hk_close * hk_forex_ratio / current_a_close - - t_result["a_label"] = next_a_uplift - t_result["hk_label"] = next_hk_uplift - - result.append(t_result) - - return result - - -class AhWeeklyFeatureOp(BaseOp): - - def execute(self): - # actor_index = self.context.actor_index - result = [] - dt = self.context.dt - - use_open: bool = self.context.use_open - hk_forex_df = self.context.hk_forex_df - # 晚一天时间 - hk_forex_dt_list = [x for x in sorted(hk_forex_df.index.unique()) if x < dt] - hk_forex_dt = hk_forex_dt_list[find_dt_less_index(dt, hk_forex_dt_list)] - hk_forex_ratio = hk_forex_df.loc[hk_forex_dt, "close"] - - for name, stock_info in self.context.stock_dict.items(): - t_result = { - "dt": dt, - "name": name, - "code": stock_info["a_code"] + "+" + stock_info["hk_code"], - } - - a_org_stock_df = stock_info["a_org_stock_df"].set_index("trade_date") - hk_org_stock_df = stock_info["hk_org_stock_df"].set_index("trade_date") - - a_dt_list = stock_info["a_dt_list"] - if dt not in a_dt_list: - continue - - ########## - dt_a_index: int = a_dt_list.index(dt) - if dt_a_index != len(a_dt_list) - 1: - start_dt = a_dt_list[dt_a_index + 1] - else: - start_dt = dt - - end_dt = next_friday_or_same(str(start_dt)) - label_dts = [x for x in a_dt_list if int(start_dt) <= x <= int(end_dt)] - end_dt = label_dts[-1] - assert start_dt <= end_dt, f"{name} dt={dt} start_dt={start_dt} > end_dt={end_dt}" - assert len(label_dts) >= 1 - - ########## - - hk_dt_list = stock_info["hk_dt_list"] - - dt_hk_index = find_dt_less_index(dt, hk_dt_list) - - if dt_a_index is None or len(a_dt_list[: dt_a_index + 1]) < 2: - continue - - if dt_hk_index is None or len(hk_dt_list[: dt_hk_index + 1]) < 2: - continue - - current_a_close = a_org_stock_df.loc[a_dt_list[dt_a_index], "close"] - current_hk_close = hk_org_stock_df.loc[hk_dt_list[dt_hk_index], "close"] - - current_a_amount = a_org_stock_df.loc[a_dt_list[dt_a_index], "amount"] - current_hk_amount = hk_org_stock_df.loc[hk_dt_list[dt_hk_index], "amount"] - - current5_a_amount = a_org_stock_df.loc[a_dt_list[max(dt_a_index - 4, 0) : dt_a_index + 1], "amount"].mean() - current5_hk_amount = hk_org_stock_df.loc[ - hk_dt_list[max(dt_hk_index - 4, 0) : dt_hk_index + 1], "amount" - ].mean() - - current_a_uplift = a_org_stock_df.loc[a_dt_list[dt_a_index], "pct_chg"] - current_hk_uplift = hk_org_stock_df.loc[hk_dt_list[dt_hk_index], "pct_chg"] - - assert use_open - next_a_uplift = ( - a_org_stock_df.loc[int(label_dts[0]), "close"] - / a_org_stock_df.loc[ - int(label_dts[0]), - "open", - ] - ) - for x_dt in label_dts[1:]: - next_a_uplift *= 1 + a_org_stock_df.loc[int(x_dt), "pct_chg"] / 100 - next_a_uplift = (next_a_uplift - 1) * 100 - - next_hk_uplift = 0 - - t_result["current_hk_close"] = current_hk_close - t_result["current_a_close"] = current_a_close - t_result["hk_forex_ratio"] = hk_forex_ratio - - t_result["a_flag"] = 1 if dt in a_dt_list else 0 - t_result["hk_flag"] = 1 if dt in hk_dt_list else 0 - - t_result["a_uplift"] = current_a_uplift - t_result["hk_uplift"] = current_hk_uplift - - for i in [1, 3, 5, 10, 20]: - a_start_idx = max(0, dt_a_index - (i - 1)) - hk_start_idx = max(0, dt_hk_index - (i - 1)) - - a_date_slice = a_dt_list[a_start_idx : dt_a_index + 1] - hk_date_slice = hk_dt_list[hk_start_idx : dt_hk_index + 1] - - a_matched = hk_org_stock_df.loc[hk_org_stock_df.index.isin(a_date_slice), "pct_chg"] - hk_matched = hk_org_stock_df.loc[hk_org_stock_df.index.isin(hk_date_slice), "pct_chg"] - - t_result[f"avg_{i}d_a_pct"] = a_matched.mean() if len(a_matched) > 0 else 0 - t_result[f"avg_{i}d_hk_pct"] = hk_matched.mean() if len(hk_matched) > 0 else 0 - - t_result["ah_amount"] = (current_hk_amount / current_a_amount - 1) if current_a_amount != 0 else 0 - t_result["ah_amount5"] = (current5_hk_amount / current5_a_amount - 1) if current5_a_amount != 0 else 0 - t_result["ah_ratio"] = current_hk_close * hk_forex_ratio / current_a_close - - t_result["a_label"] = next_a_uplift - t_result["hk_label"] = next_hk_uplift - - result.append(t_result) - - return result - - -class AhBacktestOp(BaseOp): - - def execute(self): - dt = self.context.dt - feature_df: pd.DataFrame = self.context.feature_df - max_samples = self.context.max_samples - train_dt_a_list: List[str] = [x for x in self.context.dt_a_list if x < dt][-max_samples:-1] - - train_df: pd.DataFrame = feature_df.loc[feature_df.dt.isin(train_dt_a_list)] - test_df: pd.DataFrame = feature_df.loc[feature_df.dt == dt].copy() # Fix: 使用.copy()避免SettingWithCopyWarning - - ah_ratio_key = "ah_ratio" - label_key = "a_label" - pred_key = "pred_y_nd" - # dt,name,code,current_hk_close,current_a_close,hk_forex_ratio,a_flag,hk_flag,a_uplift,hk_uplift,ah_amount,ah_ratio,a_label,hk_label - x_cols = [ah_ratio_key, "ah_amount5"] # , "ah_amount" , ah_ratio_key, "hk_flag" - - for i in [20]: # 3, 5, 10, 20 - x_cols.append(f"avg_{i}d_a_pct") - x_cols.append(f"avg_{i}d_hk_pct") - - train_x_nd: np.ndarray = train_df.loc[:, x_cols].values - train_y_nd: np.ndarray = train_df.loc[:, label_key].values - test_x_nd: np.ndarray = test_df.loc[:, x_cols].values - test_y_nd: np.ndarray = test_df.loc[:, label_key].values - rule_y_nd: np.ndarray = test_df.loc[:, ah_ratio_key].values - - model = Ridge(alpha=1) - model.fit(train_x_nd, train_y_nd) - import ray - - ray.logger.info(f"dt={dt} w={model.coef_} b={model.intercept_}") - pred_y_nd: np.ndarray = model.predict(test_x_nd) - test_df.loc[:, pred_key] = pred_y_nd - - model_ic, _ = pearsonr(pred_y_nd, test_y_nd) - model_ric, _ = spearmanr(pred_y_nd, test_y_nd) # noqa - rule_ic, _ = pearsonr(rule_y_nd, test_y_nd) - rule_ric, _ = spearmanr(rule_y_nd, test_y_nd) # noqa - - model_df = test_df.sort_values(by=pred_key, ascending=False)[:5] - rule_df = test_df.sort_values(by=ah_ratio_key, ascending=False)[:5] - - result = { - "dt": dt, - "size": len(test_df), - "model_ic": model_ic, - "model_ric": model_ric, - "rule_ic": rule_ic, - "rule_ric": rule_ric, - "rule_names": ",".join(rule_df.loc[:, "name"].tolist()), - "rule_uplift": rule_df.loc[:, label_key].mean(), - "model_names": ",".join(model_df.loc[:, "name"].tolist()), - "model_uplift": model_df.loc[:, label_key].mean(), - "all_uplift": test_y_nd.mean(), - } - - block_size = 10 - block_cnt = round(len(test_df) / block_size) - test_df_sorted = test_df.sort_values(by=pred_key, ascending=False) - # test_df_sorted = test_df.sort_values(by=ah_ratio_key, ascending=False) - for i in range(block_size): - result[f"p{i}_uplift"] = test_df_sorted[i * block_cnt : (i + 1) * block_cnt].loc[:, label_key].mean() - return result - - -def main(): - from flowllm.app import FlowLLMApp - - with FlowLLMApp(load_default_config=True) as app: - app.service_config.ray_max_workers = 8 - - op = AhOp() - op = op << AhFeatureOp() << AhWeeklyFeatureOp() << AhBacktestOp() - print(op.call()) - - -if __name__ == "__main__": - main() diff --git a/old/op/gallery/__init__.py b/old/op/gallery/__init__.py deleted file mode 100644 index 3dcfc82..0000000 --- a/old/op/gallery/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -from .execute_code_op import ExecuteCodeOp -from .mock_exception_op import MockExceptionOp -from .mock_op import Mock1Op, Mock2Op, Mock3Op, Mock4Op, Mock5Op, Mock6Op -from .react_llm_op import ReactLLMOp -from .research_complete_op import ResearchCompleteOp -from .simple_llm_op import SimpleLLMOp -from .stream_llm_op import StreamLLMOp -from .task_react_op import TaskReactOp -from .terminate_op import TerminateOp -from .think_op import ThinkToolOp -from .translate_code_op import TranslateCodeOp diff --git a/old/op/gallery/mock_exception_op.py b/old/op/gallery/mock_exception_op.py deleted file mode 100644 index 5364e0b..0000000 --- a/old/op/gallery/mock_exception_op.py +++ /dev/null @@ -1,41 +0,0 @@ -import asyncio - -from loguru import logger - -from flowllm.context.service_context import C -from flowllm.op.base_async_tool_op import BaseAsyncToolOp -from flowllm.schema.tool_call import ToolCall - - -@C.register_op(register_app="FlowLLM") -class MockExceptionOp(BaseAsyncToolOp): - - def build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "mock exception tool", - "input_schema": { - "query": { - "type": "string", - "description": "query", - "required": True, - }, - }, - }, - ) - - async def async_execute(self): - query = self.input_dict.get("query") - await asyncio.sleep(1) - logger.info(f"start mock exception with {query}") - - if query == "run_time": - raise RuntimeError("run_time error") - - elif query == "value": - raise ValueError("value error") - - elif query == "not_im": - raise NotImplementedError("not_im error") - - self.set_result("normal") diff --git a/old/op/gallery/mock_op.py b/old/op/gallery/mock_op.py deleted file mode 100644 index 9627e83..0000000 --- a/old/op/gallery/mock_op.py +++ /dev/null @@ -1,162 +0,0 @@ -import asyncio -import time -from typing import List - -from loguru import logger - -from flowllm.context import C -from flowllm.context.flow_context import FlowContext -from flowllm.op.base_async_op import BaseAsyncOp -from flowllm.op.base_op import BaseOp -from flowllm.schema.message import Message, Role - - -@C.register_op(register_app="FlowLLM") -class Mock1Op(BaseOp): - def execute(self): - time.sleep(1) - a = self.context.get("a", 1) - b = self.context.get("b", 2) - logger.info(f"enter class={self.name}. a={a} b={b}") - - self.context.response.answer = f"{self.name} a={a} b={b} answer=47" - - -@C.register_op(register_app="FlowLLM") -class Mock2Op(Mock1Op): ... - - -@C.register_op(register_app="FlowLLM") -class Mock3Op(Mock1Op): ... - - -@C.register_op(register_app="FlowLLM") -class Mock4Op(BaseAsyncOp): - """使用 LLM 进行简单问答的异步 Op""" - - def __init__(self, llm: str = "qwen3_30b_instruct", **kwargs): - super().__init__(llm=llm, **kwargs) - - async def async_execute(self): - await asyncio.sleep(1) - a = self.context.get("a", 1) - b = self.context.get("b", 2) - logger.info(f"enter class={self.name}. a={a} b={b}") - - # 使用 LLM 进行调用 - query = self.context.get("query", f"calculate {a} + {b}") - messages: List[Message] = [Message(role=Role.USER, content=query)] - - # 调用 LLM - response: Message = await self.llm.achat(messages) - logger.info(f"LLM response: {response.content}") - - # 将 LLM 的响应保存到 context 中，供下一个 op 使用 - self.context.response.answer = response.content - self.context.set("llm_response", response.content) - - -@C.register_op(register_app="FlowLLM") -class Mock5Op(BaseAsyncOp): - """串联使用前一个 Op 的结果作为输入""" - - def __init__(self, llm: str = "qwen3_30b_instruct", **kwargs): - super().__init__(llm=llm, **kwargs) - - async def async_execute(self): - # 获取上一个 op 的结果 - previous_response = self.context.get("llm_response", "no previous response") - logger.info(f"enter class={self.name}. previous_response={previous_response}") - - # 基于上一个结果继续对话 - query = self.context.get("query", f"Based on the previous answer: {previous_response}, please elaborate more.") - messages: List[Message] = [ - Message(role=Role.USER, content=f"Previous result: {previous_response}"), - Message(role=Role.USER, content=query), - ] - - response: Message = await self.llm.achat(messages) - logger.info(f"LLM response in {self.name}: {response.content}") - - # 更新 context - self.context.response.answer = response.content - self.context.set("llm_response", response.content) - - -@C.register_op(register_app="FlowLLM") -class Mock6Op(BaseAsyncOp): - """支持多轮对话的 Mock Op""" - - def __init__(self, llm: str = "qwen3_30b_instruct", **kwargs): - super().__init__(llm=llm, **kwargs) - - async def async_execute(self): - # 获取历史消息 - messages = self.context.get("messages", []) - if not messages: - query = self.context.get("query", "Hello, how are you?") - messages = [Message(role=Role.USER, content=query)] - else: - # 将 dict 转换为 Message 对象 - messages = [Message(**m) if isinstance(m, dict) else m for m in messages] - - logger.info(f"enter class={self.name}. messages count={len(messages)}") - - # 调用 LLM - response: Message = await self.llm.achat(messages) - logger.info(f"LLM response in {self.name}: {response.content}") - - # 添加助手的回复到消息历史 - messages.append(response) - - # 保存消息历史和答案 - self.context.set("messages", [m.model_dump() for m in messages]) - self.context.response.answer = response.content - self.context.response.metadata["messages"] = messages - - -async def main(): - """演示多个 Op 之间的串联调用""" - from flowllm.app import FlowLLMApp - - async with FlowLLMApp(load_default_config=True): - # 示例 1: Mock4Op - 简单的 LLM 调用 - logger.info("=== Example 1: Mock4Op ===") - context1 = FlowContext(query="What is 42 + 5?") - op1 = Mock4Op() - await op1.async_call(context=context1) - logger.info(f"Mock4Op result: {context1.response.answer}") - - # 示例 2: Mock4Op -> Mock5Op - 串联调用 - logger.info("\n=== Example 2: Mock4Op -> Mock5Op ===") - context2 = FlowContext(query="What is artificial intelligence?", a=10, b=20) - - op2 = Mock4Op() - await op2.async_call(context=context2) - logger.info(f"Mock4Op result: {context2.response.answer}") - - op3 = Mock5Op() - context2.set("query", "Can you give me a practical example?") - await op3.async_call(context=context2) - logger.info(f"Mock5Op result: {context2.response.answer}") - - # 示例 3: Mock6Op - 多轮对话 - logger.info("\n=== Example 3: Mock6Op - Multi-turn ===") - context3 = FlowContext(query="Tell me about Python programming") - - op4 = Mock6Op() - await op4.async_call(context=context3) - logger.info(f"Turn 1: {context3.response.answer}") - - # 第二轮对话 - messages = context3.get("messages", []) - messages.append({"role": "user", "content": "What are its main features?"}) - context3.set("messages", messages) - - op5 = Mock6Op() - await op5.async_call(context=context3) - logger.info(f"Turn 2: {context3.response.answer}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/old/op/gallery/react_llm_op.py b/old/op/gallery/react_llm_op.py deleted file mode 100644 index c56003a..0000000 --- a/old/op/gallery/react_llm_op.py +++ /dev/null @@ -1,110 +0,0 @@ -import asyncio -import datetime -import json -import time -from typing import List, Dict - -from loguru import logger - -from flowllm.context.flow_context import FlowContext -from flowllm.context.service_context import C -from flowllm.op.base_async_tool_op import BaseAsyncToolOp -from flowllm.schema.message import Message, Role -from flowllm.schema.tool_call import ToolCall - - -@C.register_op(register_app="FlowLLM") -class ReactLLMOp(BaseAsyncToolOp): - file_path: str = __file__ - - def __init__(self, llm: str = "qwen3_30b_instruct", save_answer: bool = True, **kwargs): - super().__init__(llm=llm, save_answer=save_answer, **kwargs) - - def build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "use this query to query an LLM", - "input_schema": { - "query": { - "type": "string", - "description": "query", - "required": True, - }, - }, - }, - ) - - async def async_execute(self): - query: str = self.input_dict["query"] - - max_steps: int = int(self.op_params.get("max_steps", 10)) - from flowllm.op.search import DashscopeSearchOp - - tools: List[BaseAsyncToolOp] = [DashscopeSearchOp()] - tool_dict: Dict[str, BaseAsyncToolOp] = {x.tool_call.name: x for x in tools} - for name, tool_call in tool_dict.items(): - logger.info( - f"name={name} " f"tool_call={json.dumps(tool_call.tool_call.simple_input_dump(), ensure_ascii=False)}", - ) - - now_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") - user_prompt = self.prompt_format( - prompt_name="role_prompt", - time=now_time, - tools=",".join(list(tool_dict.keys())), - query=query, - ) - messages: List[Message] = [Message(role=Role.USER, content=user_prompt)] - logger.info(f"step.0 user_prompt={user_prompt}") - - for i in range(max_steps): - assistant_message: Message = await self.llm.achat(messages, tools=[x.tool_call for x in tools]) - messages.append(assistant_message) - logger.info( - f"assistant.round{i}.reasoning_content={assistant_message.reasoning_content}\n" - f"content={assistant_message.content}\n" - f"tool.size={len(assistant_message.tool_calls)}", - ) - - if not assistant_message.tool_calls: - break - - op_list: List[BaseAsyncToolOp] = [] - for j, tool_call in enumerate(assistant_message.tool_calls): - logger.info(f"submit step={i} tool_calls.name={tool_call.name} argument_dict={tool_call.argument_dict}") - - if tool_call.name not in tool_dict: - logger.warning(f"step={i} no tool_call.name={tool_call.name}") - continue - - op_copy = tool_dict[tool_call.name].copy() - op_list.append(op_copy) - self.submit_async_task(op_copy.async_call, **tool_call.argument_dict) - time.sleep(1) - - await self.join_async_task() - - for j, op in enumerate(op_list): - logger.info(f"submit step.index={i}.{j} tool_result={op.output}") - tool_result = str(op.output) - tool_message = Message(role=Role.TOOL, content=tool_result, tool_call_id=op.tool_call.id) - messages.append(tool_message) - - self.set_result(messages[-1].content) - self.context.response.messages = messages - - -async def main(): - from flowllm.app import FlowLLMApp - - async with FlowLLMApp(load_default_config=True): - context = FlowContext(query="茅台和五粮现在股价多少？") - - op = ReactLLMOp() - result = await op.async_call(context=context) - print(result) - print(op.output) - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/old/op/gallery/react_llm_prompt.yaml b/old/op/gallery/react_llm_prompt.yaml deleted file mode 100644 index 4e6cfc2..0000000 --- a/old/op/gallery/react_llm_prompt.yaml +++ /dev/null @@ -1,19 +0,0 @@ -role_prompt: | - You are a helpful assistant. - The current time is {time}. - Please proactively choose the most suitable tool or combination of tools based on the user's question, including {tools} etc. - Please first think about how to break down the problem into subtasks, what tools and parameters should be used for each subtask, and finally provide the tool call name and parameters. - Try calling the same tool multiple times with different parameters to obtain information from various perspectives. - Please determine the response language based on the language of the user's question. - - {query} - -role_prompt_zh: | - 你是一个有用的助手。 - 当前时间是 {time}。 - 请根据用户的问题，主动选择最合适的工具或工具组合，包括 {tools} 等。 - 请先思考如何将问题分解为子任务，每个子任务应使用哪些工具和参数，最后提供工具调用名称和参数。 - 尝试多次使用相同的工具，但使用不同的参数，从多个角度获取信息。 - 请根据用户问题的语言来确定回复的语言。 - - {query} diff --git a/old/op/gallery/research_complete_op.py b/old/op/gallery/research_complete_op.py deleted file mode 100644 index d411d23..0000000 --- a/old/op/gallery/research_complete_op.py +++ /dev/null @@ -1,32 +0,0 @@ -import asyncio - -from flowllm.context.flow_context import FlowContext -from flowllm.context.service_context import C -from flowllm.op.base_async_tool_op import BaseAsyncToolOp -from flowllm.schema.tool_call import ToolCall - - -@C.register_op(register_app="FlowLLM") -class ResearchCompleteOp(BaseAsyncToolOp): - - def build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "name": "research_complete", - "description": "Call this tool to indicate that the research is complete.", - }, - ) - - async def async_execute(self): - self.set_result(f"The research is complete.") - - -async def main(): - op = ResearchCompleteOp() - context = FlowContext() - await op.async_call(context) - print(f"Result: {op.output}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/old/op/gallery/task_react_op.py b/old/op/gallery/task_react_op.py deleted file mode 100644 index 7407901..0000000 --- a/old/op/gallery/task_react_op.py +++ /dev/null @@ -1,637 +0,0 @@ -from __future__ import annotations - -import asyncio -import os -from contextlib import suppress -from dataclasses import dataclass -from datetime import datetime -from typing import List, Optional - -from loguru import logger - -from flowllm.client.fastmcp_client import FastmcpClient as McpClient -from flowllm.context.flow_context import FlowContext -from flowllm.context.service_context import C -from flowllm.enumeration.chunk_enum import ChunkEnum -from flowllm.op.base_async_tool_op import BaseAsyncToolOp -from flowllm.schema.message import Message -from flowllm.schema.message import Role -from flowllm.schema.tool_call import ToolCall - - -def tools_schema_to_qwen_prompt(tools_schema: List[ToolCall]): - """ - 将 tools_schema 转换为符合 Qwen 模型 chat_template 的工具描述 prompt。 - - Args: - tools_schema (list): 工具列表，格式如下： - [ - { - "name": "tool_name", - "description": "工具描述", - "parameters": { - "type": "object", - "properties": { ... }, - "required": [ ... ] - } - } - ] - - Returns: - str: 包含 标签的完整 system 工具描述 prompt - """ - if not tools_schema: - return "" - - lines = [] - # lines.append("\n\n# Tools\n") - lines.append("You may call one or more functions to assist with the user query.\n") - lines.append("You are provided with function signatures within XML tags:") - lines.append("") - # 逐个添加工具定义（JSON 格式，不转义） - for tool in tools_schema: - tool_json = tool.simple_input_dump() - tool_json = json.dumps( - tool_json, - ensure_ascii=False, - separators=(",", ":"), # 紧凑格式，不加空格 - ) - lines.append(tool_json) - lines.append("\n") - lines.append( - "For each function call, return a json object with function name and arguments within and XML tags:" - ) - lines.append("") - lines.append('{"name": , "arguments": }') - lines.append("") - - return "\n".join(lines) - - -from typing import Any, Dict, List - - -def parse_tool_calls(text: str): - """ - 从包含 ... 的文本中解析工具调用信息。 - - 返回: - calls: 按出现顺序的调用列表，元素形如 {"name": str, "arguments": Any} - """ - # 提取每个 块中包裹的 JSON（允许跨行） - pattern = re.compile(r"\s*({.*?})\s*", re.DOTALL) - json_blobs = pattern.findall(text) - - calls = [] - for blob in json_blobs: - try: - obj = json.loads(blob) - except json.JSONDecodeError: - # 无法解析的块直接跳过 - continue - - name = obj.get("name") - arguments = obj.get("arguments") - if name is None: - # 必须有 name 字段 - continue - tool_call_dict = {"name": name, "arguments": arguments} - - calls.append(tool_call_dict) - - return calls - - -import re -import json -from typing import Tuple - - -def parse_final_query_and_type(text: str) -> Tuple[str, str]: - """ - 从包含 ... 的字符串中解析出 final_query 与 type。 - 返回 (final_query, type)。解析失败会抛出 ValueError。 - """ - # 1) 提取 ...JSON... 内的内容 - m = re.search(r"\s*(\{.*?\})\s*", text, flags=re.S | re.I) - if not m: - raise ValueError("未找到 ... 块") - - json_str = m.group(1).strip() - - # 2) 解析 JSON - try: - obj = json.loads(json_str) - except json.JSONDecodeError as e: - raise ValueError(f"JSON 解析失败: {e.msg}") from e - - # 3) 取字段并做校验 - if not isinstance(obj, dict): - raise ValueError("answer 块不是一个 JSON 对象") - if "final_query" not in obj or "type" not in obj or "simplified_query" not in obj: - raise ValueError("JSON 中缺少 final_query, simplified_query 或 type 字段") - - return obj - - -def print_message_list(messages, *, max_content_chars: int = 800): - """Pretty-print a list of FlowLLM `Message` objects. - - - Safely handles Role enums and missing fields - - Displays tool calls with parsed JSON arguments (when possible) - - Truncates long content to `max_content_chars` - """ - import json as _json - - divider = "-" * 60 - print("=" * 60) - for i, m in enumerate(messages): - role = getattr(m, "role", None) - role_str = getattr(role, "value", str(role)) - time_created = getattr(m, "time_created", None) or "-" - print(f"[{i}] role={role_str} time={time_created}") - - def _emit(label: str, text: str): - if not text: - return - if max_content_chars and len(text) > max_content_chars: - text = text[:max_content_chars] + "…" - print(f" {label}: {text}") - - _emit("content", getattr(m, "content", "")) - _emit("reasoning", getattr(m, "reasoning_content", "")) - - tool_calls = getattr(m, "tool_calls", None) or [] - if tool_calls: - print(f" tool_calls ({len(tool_calls)}):") - for tc in tool_calls: - # Support both direct fields and tc.function.* shape - name = getattr(tc, "name", None) - _type = getattr(tc, "type", None) or "function" - tc_id = getattr(tc, "id", "") - args = getattr(tc, "arguments", {}) - - fn = getattr(tc, "function", None) - if fn and not name: - name = getattr(fn, "name", None) - args = getattr(fn, "arguments", args) - - if isinstance(args, str): - try: - args_fmt = _json.dumps(_json.loads(args), ensure_ascii=False) - except Exception: - args_fmt = args - else: - try: - args_fmt = _json.dumps(args, ensure_ascii=False) - except Exception: - args_fmt = str(args) - - print(f" - {_type} {name} id={tc_id}") - if args_fmt: - print(f" args: {args_fmt}") - - tool_call_id = getattr(m, "tool_call_id", "") - if tool_call_id: - print(f" tool_call_id: {tool_call_id}") - - metadata = getattr(m, "metadata", None) - if metadata: - try: - meta_str = _json.dumps(metadata, ensure_ascii=False) - except Exception: - meta_str = str(metadata) - print(f" metadata: {meta_str}") - - print(divider) - - -# -------------------- 序列化 & 写入工具 -------------------- -def serialize_message(msg: Message) -> Dict[str, Any]: - """把 SDK 的 Message 转成可 JSON 序列化的简洁结构。""" - return { - "role": getattr(msg, "role", None), - "content": getattr(msg, "content", None), - } - - -def serialize_trajectory(history: List[Message]) -> List[Dict[str, Any]]: - return [serialize_message(m) for m in history] - - -def append_jsonl(path: str, obj: Dict[str, Any]) -> None: - os.makedirs(os.path.dirname(path) or ".", exist_ok=True) - with open(path, "a", encoding="utf-8", newline="\n") as f: - f.write(json.dumps(obj, ensure_ascii=False) + "\n") - f.flush() - - -def load_existing_queries_from_jsonl( - path: str, - field: str = "final_query", -): - """ - 读取 jsonl 文件，顺序返回已存在的 `final_query` 列表（去掉 None/空字符串）。 - 若文件不存在，返回空列表。 - """ - exists: List[str] = [] - if not os.path.exists(path): - return exists - with open(path, "r", encoding="utf-8") as f: - for line in f: - line = line.strip() - if not line: - continue - try: - obj = json.loads(line) - except Exception: - continue - val = obj.get(field) - if isinstance(val, str) and val.strip(): - exists.append(val.strip()) - return exists - - -TASK_AGENT_PROMPT = """ -你是一个【任务生成智能体】，目标是在了解一个特定的智能体应用（下称“目标 Agent”）及其可用工具后，先探索再出题，最后仅输出一个高质量模拟用户问题（Task/Query），供训练该“目标 Agent”使用。 - -当前时间: {time_stamp} - -[输入] -1) 目标 Agent 应用描述： -{application_description} -2) 可用工具： -{tools} -3) 题目范围与类型（可选）： -{topic} -4) 已存在的题目（请勿生成重复问题）： -{exist_tasks} - -[你的任务] -A. 构思：你必须从构思开始，根据“应用描述+范围”，在草稿中列出 1-3 个潜在出题方向，也可以通过热点搜索来确定方向。 -B. 探索：**按需**调用工具（0–N 次均可），以提升问题质量与可回答性。你可以自由决定调用顺序与组合。 -C. 迭代：交替进行“构思 ⇄ 探索”，收敛并打磨问题，使其： - - 清晰、具体、可被目标 Agent 在其工具边界内解答； - - 明确必要上下文，但不要过多或过于具体（具体信息需要目标Agent在解答时自行收集）； - - 具训练价值（避免过宽或一搜即得；鼓励结构化信息需求与可检验标准）。 -D. 输出：仅产出一个问题（必要）；可附元数据（可选）. - -[工具使用——灵活指引（非强制，供你自由选择与组合）] -- 热点捕捉：检索最近一段时间（如近 7–30 天）与主题相关的新闻/事件/政策，识别可转化为问题的变化点或冲突点。 -- 实时脉冲：对涉时效主题，拉取最新关键数据点（如 CPI/利率/价格/产能），并在问题中固化时间点或区间。 -- 口径校准：当问题涉及统计口径/定义（如 TTM/年度、含税口径、地区边界）时，做一次轻量确认以避免歧义。 -- 范围界定：需要对比时，生成候选实体清单（公司/子行业/指标），并据此在问题里限定范围或点名对象。 -- 证据线索：为后续回答留出可验证的证据维度（来源类型/指标表/公告字段），但不要在输出中暴露任何工具细节。 -- 细化与降噪：若发现问题过宽或不可答，借助工具结果收缩到可操作的时间窗、对象清单、指标集合或假设边界。 -- 失败降级：工具不可用/受限时，改写为无需实时数据也能回答的版本（例如改为区间回溯、方法论对比或结构化框架题）。 - -[注意事项] -1. 你有完全的调用自主权：是否调用、调用多少次、调用哪个工具、以何种顺序，都由你决定。 -2. 优先以更好问题为目标，而非收集更多信息；在最少必要信息下形成可执行、可检验的问题陈述。 -3. 注意生成问题的多样性，不限于内容、问法等方面。 -4. 合规边界（必要）：不得诱导违法/不当请求，不得要求内幕信息或操纵市场；避免侵犯隐私；避免医疗/法律等高风险建议的结论性指令；如使用时效数据，仅作构题的合理事实校验，并在问题中固化时间点或区间。 - -[输出格式] -在**最终问题尚未确定之前**：请先以如下格式输出思考过程，并适当调用工具。 - 思考过程 工具调用 -**当且仅当**你已确定最终问题后，输出以下格式的JSON对象。 - -{{ - "final_query": "<给出的单个高质量问题，1–2句>", - "simplified_query": "", - "type": "<问题种类，具体选项参考目标应用描述>" -}} - - -[严格要求] -- 如果探索表明问题过于宽泛或不可答，请主动收敛并补充限定。 -- 若涉及时效性数据，请在问题中明确时间点或区间（示例：“截至 2025-06 的…”、“2023Q4–2025Q2”）。 -- 如工具不可用或超时，必须进行范围收缩或改写为无需实时数据亦可回答的问题。 -""" - - -FIN_DEEPRESEARCH_DESCRIPTION = """ -【应用名称】金融 deepresearch -【使命】 -以有洞察力、可复核、可落地的研究框架回答金融相关问题，覆盖行业研究、事件解读、个股分析、宏观分析与股票检索五大场景。 -【覆盖场景与示例】 -1) 行业研究 - - 中国电解铝是否具备中长期投资价值？核心标的与优劣对比 - - 磷化工/氟化工竞争格局与胜出因子，优选公司及理由 - - 半导体上行期，“铲子股”vs“耗材股”的取舍与公司推荐 - - 上市黄金企业的储量、成本与估值梳理；优选标的 - - 快递行业竞争态势与龙头护城河评估 -2) 事件解读 - - 最近一次美国CPI结构拆解及对资产定价/风格的影响 - - “反内卷”导向对高竞争行业盈利与估值的影响 - - 外卖补贴战对三家平台估值的分化路径 - - 关税政策变化的受益板块与A股映射 - - 宠物经济兴起的底层动因与可投资标的 - - 阶段性A股上涨的驱动因素剖析 -3) 个股分析 - - 紫金矿业增长的可持续性与边际驱动 - - 海尔估值合理性与同业对标 - - 申通快递相对同业的优势/劣势与催化 - - 藏格矿业、赤峰黄金的股价驱动因素 - - 中宠股份在当前价位的性价比评估 -4) 宏观分析 - - 利率下行对大宗商品的影响路径与机会 - - 社零与工业数据对风格与板块轮动的含义 - - 全球实际利率中枢对权益估值的作用 - - 名义增速缺口对沪深300的影响量化 - - 股债利差/倒挂对A股风险偏好的影响 - - 上证指数突破关键位所需条件复盘 -5) 股票检索（情报/线索） - - 海外油服设备供应商认证与相关企业清单 - - 在越南有产能布局的A股出口公司盘点 - - 关税重构背景下“铲子”赛道受益环节与公司列表 - - 电力/电气设备领域具寡头力的出海企业 - - 汽车产业链全球化龙头与品牌出海的产品型公司 -""" - - -# -------------------- Batch item -------------------- -@dataclass -class BatchItem: - topic: str - - -# -------------------- Op: TaskReActOP (Agent merged into Op) -------------------- -@C.register_op() -class TaskReactOp(BaseAsyncToolOp): - - def __init__( - self, - llm: str = "qwen3_max_instruct", - system_prompt: Optional[str] = None, - **kwargs, - ) -> None: - super().__init__(llm=llm, **kwargs) - self.system_prompt = system_prompt - - def build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "Run dedup batch of Task ReAct agents and stream results (Agent merged). Provide exist_list to dedup; no jsonl read.", - "input_schema": { - "items": { - "type": "array", - "description": "list of objects with field 'topic'", - "required": True, - }, - "exist_list": { - "type": "array", - "description": "existing simplified queries to deduplicate against", - "required": False, - }, - "jsonl_path": { - "type": "string", - "description": "path to append results jsonl", - "required": False, - }, - "mcp_transport": {"type": "string", "required": False}, - "mcp_host": {"type": "string", "required": False}, - "mcp_port": {"type": "integer", "required": False}, - "max_steps": {"type": "integer", "required": False}, - "dedup_field": {"type": "string", "required": False}, - "exist_task_prompt_limit": {"type": "integer", "required": False}, - "per_item_timeout": {"type": "number", "required": False}, - "include_trajectory": {"type": "boolean", "required": False}, - }, - }, - ) - - # -------------------- Agent logic (as a helper method) -------------------- - async def _agent_execute( - self, - input_topic: str, - *, - mcp: McpClient, - max_steps: int = 30, - stream: bool = False, - exist_tasks: Optional[List[str]] = None, - max_tool_result_char: int = int(5e4), - ) -> Dict[str, Any]: - """ - Agent logic merged into the Op as a coroutine method. LLM calls remain non-streaming per requirement. - Returns final_query_and_type with trajectory. - """ - # llm = OpenAICompatibleBaseLLM(model_name=self.model_name) - - # 1) 读取可用工具 - tools_schema = await mcp.list_tool_calls() or [] - tools_prompt = tools_schema_to_qwen_prompt(tools_schema) - - # 2) 最小对话历史 - task_prompt = TASK_AGENT_PROMPT.format( - time_stamp=datetime.now().strftime("%Y%m%d-%H%M%S"), - application_description=FIN_DEEPRESEARCH_DESCRIPTION, - tools=tools_prompt, - topic=input_topic, - exist_tasks=json.dumps(exist_tasks or [], ensure_ascii=False), - ) - - history: List[Message] = [] - if self.system_prompt: - history.append(Message(role=Role.SYSTEM, content=self.system_prompt)) - history.append(Message(role=Role.USER, content=task_prompt)) - - last_reply: Optional[Message] = None - final_query_and_type: Dict[str, Any] = {} - - # 3) 多步循环 - for step_idx in range(1, max_steps + 1): - last_reply = await self.llm.achat( - history, - [], - enable_stream_print=stream, - ) - history.append(last_reply) - - content = getattr(last_reply, "content", None) or [] - await self.context.add_stream_chunk_and_type( - chunk=json.dumps({"event": "llm", "content": content}, ensure_ascii=False), - chunk_type=ChunkEnum.THINK, - ) - tool_calls = parse_tool_calls(content) - - if not tool_calls: - final_query_and_type = parse_final_query_and_type(content) or {} - break - - for tc in tool_calls: - name = tc.get("name") - args = tc.get("arguments") - if not name: - continue - - try: - result = await mcp.call_tool(name, args) - - text = None - if getattr(result, "content", None): - try: - text = result.content[0].text - except Exception: - text = str(result.content) - if not text and getattr(result, "structured_content", None): - text = json.dumps(result.structured_content, ensure_ascii=False) - if not text: - text = str(result) - - if isinstance(text, str) and len(text) > max_tool_result_char: - text = text[:max_tool_result_char] + "...[truncated]" - - history.append( - Message( - role=Role.USER, - content=f"[tool:{name}]" + text, - ), - ) - except Exception as e: - history.append( - Message( - role=Role.USER, - content=f"[tool:{name}] 调用失败: {e}", - ), - ) - - if not final_query_and_type: - content = getattr(last_reply, "content", None) or [] - final_query_and_type = parse_final_query_and_type(content) or {} - - final_query_and_type.update({"trajectory": history}) - return final_query_and_type - - # -------------------- Main Op execution -------------------- - async def async_execute(self): - # parse inputs - items_raw = self.input_dict.get("items") - if not items_raw: - raise RuntimeError("'items' is required") - - items = [BatchItem(**x) if isinstance(x, dict) else BatchItem(topic=str(x)) for x in items_raw] - - # New: accept exist_list directly and do not read jsonl to load history - provided_exist_list = self.input_dict.get("exist_list") - if provided_exist_list is None: - existed_list: List[str] = [] - else: - # ensure strings - existed_list = [str(x) for x in provided_exist_list] - - mcp_transport = self.input_dict.get("mcp_transport", "sse") - mcp_host = self.input_dict.get("mcp_host", "0.0.0.0") - mcp_port = int(self.input_dict.get("mcp_port", 8001)) - max_steps = int(self.input_dict.get("max_steps", 30)) - exist_task_prompt_limit = int(self.input_dict.get("exist_task_prompt_limit", 500)) - per_item_timeout = self.input_dict.get("per_item_timeout", 300.0) - include_trajectory = bool(self.input_dict.get("include_trajectory", True)) - - # prepare dedup / history - logger.info(f"existed_list length={len(existed_list)}") - exist_tasks_for_prompt = existed_list[-exist_task_prompt_limit:] - - # Open MCP once - async with McpClient(transport=mcp_transport, host=mcp_host, port=mcp_port) as mcp: - for idx, item in enumerate(items, 1): - topic = item.topic - logger.info(f"[Op-Batch-Dedup] ({idx}/{len(items)}) topic={topic!r}") - - # stream: notify start of item - await self.context.add_stream_chunk_and_type( - chunk=json.dumps( - {"event": "start_item", "index": idx, "total": len(items), "topic": topic}, ensure_ascii=False - ), - chunk_type=ChunkEnum.THINK, - ) - - # run merged agent under timeout - task = asyncio.create_task( - self._agent_execute( - topic, - mcp=mcp, - max_steps=max_steps, - stream=False, - exist_tasks=exist_tasks_for_prompt, - ), - ) - - try: - if per_item_timeout is not None and per_item_timeout > 0: - result = await asyncio.wait_for(task, timeout=per_item_timeout) - else: - result = await task - - # dedup check - simplified = result.get("simplified_query") - - record: Dict[str, Any] = { - "event": "end_item", - "input_topic": topic, - "final_query": result.get("final_query"), - "simplified_query": result.get("simplified_query"), - "type": result.get("type"), - "ts": datetime.now().isoformat(timespec="seconds"), - } - if include_trajectory: - record["trajectory"] = serialize_trajectory(result.get("trajectory", [])) - - if simplified is not None: - existed_list.append(simplified) - exist_tasks_for_prompt.append(simplified) - if len(exist_tasks_for_prompt) > exist_task_prompt_limit: - exist_tasks_for_prompt = exist_tasks_for_prompt[-exist_task_prompt_limit:] - - # stream final_query as ANSWER chunk - await self.context.add_stream_chunk_and_type( - chunk=json.dumps(record, ensure_ascii=False), - chunk_type=ChunkEnum.ANSWER, - ) - - except asyncio.TimeoutError: - task.cancel() - with suppress(asyncio.CancelledError): - await task - - err_msg = f"timeout after {per_item_timeout}s" - await self.context.add_stream_chunk_and_type( - chunk=json.dumps({"event": "error", "topic": topic, "error": err_msg}, ensure_ascii=False), - chunk_type=ChunkEnum.ERROR, - ) - - except Exception as e: - await self.context.add_stream_chunk_and_type( - chunk=json.dumps({"event": "error", "topic": topic, "error": str(e)}, ensure_ascii=False), - chunk_type=ChunkEnum.ERROR, - ) - - -async def main(): - from flowllm.app import FlowLLMApp - - async with FlowLLMApp(load_default_config=True): - - context = FlowContext(items=["", ""], stream_queue=asyncio.Queue()) - op = TaskReactOp() - - async def async_call(): - await op.async_call(context=context) - await context.add_stream_done() - - task = asyncio.create_task(async_call()) - - while True: - stream_chunk = await context.stream_queue.get() - if stream_chunk.done: - print("\nend") - await task - break - else: - print(stream_chunk.chunk, end="") - - await task - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/old/op/gallery/translate_code_op.py b/old/op/gallery/translate_code_op.py deleted file mode 100644 index 67d3496..0000000 --- a/old/op/gallery/translate_code_op.py +++ /dev/null @@ -1,581 +0,0 @@ -import asyncio -import json -from pathlib import Path -from typing import List, Dict, Tuple - -from loguru import logger -from tqdm.asyncio import tqdm_asyncio - -from flowllm.context.service_context import C -from flowllm.op.base_async_tool_op import BaseAsyncToolOp -from flowllm.schema.message import Message, Role -from flowllm.schema.tool_call import ToolCall, ParamAttrs -from flowllm.utils.common_utils import extract_content - - -@C.register_op(register_app="FlowLLM") -class TranslateCodeOp(BaseAsyncToolOp): - """ - TranslateCodeOp - TypeScript/TSX to Python Translation/Interpretation Operator - - This operator recursively finds all TypeScript (.ts) and TSX (.tsx) files in a given directory and: - - mode=True: Translates them to Python (saved to *_python directory with .py extension) - - mode=False: Interprets/explains them as markdown (saved to *_markdown directory with .md extension) - - Uses LLM with concurrent processing (pool size: 6). - """ - - file_path: str = __file__ - - def __init__( - self, - # llm="qwen3_max_instruct", - llm="qwen3_30b_instruct", - max_concurrent: int = 6, - max_retries: int = 3, - skip_existing: bool = True, - submit_interval: float = 2.0, - enable_print_output: bool = False, - **kwargs, - ): - """ - Initialize TranslateCodeOp - - Args: - max_concurrent: Maximum number of concurrent LLM calls (default: 6) - max_retries: Maximum number of retries for failed translations (default: 3) - skip_existing: Skip if target file already exists and is not empty (default: True) - submit_interval: Interval in seconds between submitting concurrent tasks (default: 2.0) - """ - super().__init__(llm=llm, enable_print_output=enable_print_output, **kwargs) - self.max_concurrent = max_concurrent - self.max_retries = max_retries - self.skip_existing = skip_existing - self.submit_interval = submit_interval - self.semaphore = None # Will be initialized in async_execute - - def build_tool_call(self) -> ToolCall: - """Build tool call schema for TranslateCodeOp""" - return ToolCall( - **{ - "name": "translate_code", - "description": "Recursively find all TypeScript (.ts) and TSX (.tsx) files in a directory and translate/interpret them", - "input_schema": { - "file_path": ParamAttrs( - type="str", - description="The directory path(s) to search for TypeScript/TSX files. Multiple paths can be separated by semicolons (;)", - required=True, - ), - "mode": ParamAttrs( - type="bool", - description="True = translate to Python (saved to *_python/*.py), False = interpret to markdown (saved to *_markdown/*.md)", - required=False, - ), - }, - }, - ) - - async def async_execute(self): - """ - Main execution method - - Steps: - 1. Get the input file path(s) and mode from context (supports multiple paths separated by semicolons) - 2. Recursively find all .ts and .tsx files in all paths - 3. Translate/interpret each file concurrently with pool size limit - 4. Return translation results - """ - # Get input file path(s) - file_path: str = self.input_dict.get("file_path", "") - if not file_path: - raise ValueError("file_path is required") - - # Get mode from input_dict (default: True for translation) - mode = self.input_dict.get("mode", True) - - # Split paths by semicolon and strip whitespace - file_paths = [path.strip() for path in file_path.split(";") if path.strip()] - mode_str = "Translation" if mode else "Interpretation" - logger.info(f"Mode: {mode_str}") - logger.info(f"Processing {len(file_paths)} path(s): {file_paths}") - - # Initialize semaphore for concurrent control - self.semaphore = asyncio.Semaphore(self.max_concurrent) - - # Find all TypeScript and TSX files recursively from all paths - ts_files = [] - for path in file_paths: - files = self._find_ts_files(path) - ts_files.extend(files) - logger.info(f"Found {len(files)} TypeScript/TSX files in {path}") - - # Remove duplicates while preserving order - seen = set() - unique_ts_files = [] - for file in ts_files: - if file not in seen: - seen.add(file) - unique_ts_files.append(file) - ts_files = unique_ts_files - - logger.info(f"Total: {len(ts_files)} unique TypeScript/TSX files") - - if not ts_files: - result = { - "status": "success", - "message": f"No TypeScript/TSX files found in the specified path(s)", - "paths": file_paths, - "translations": [], - } - self.set_result(json.dumps(result, ensure_ascii=False, indent=2)) - return - - # Analyze file statistics before translation - file_stats = self._analyze_file_statistics(ts_files) - logger.info(f"\n{'='*60}\nTypeScript Files Statistics:\n{'='*60}") - logger.info(f"Total files: {file_stats['total_files']}") - logger.info(f"Total characters: {file_stats['total_chars']:,}") - logger.info(f"Average characters per file: {file_stats['avg_chars']:.2f}") - logger.info(f"Median characters: {file_stats['median_chars']:,}") - logger.info(f"\nTop 5 Largest Files:") - for i, (file_path, chars) in enumerate(file_stats["top_5_largest"], 1): - logger.info(f" {i}. {file_path} - {chars:,} chars") - logger.info(f"{'='*60}\n") - - # Create a mapping from ts_file to its base_path for output path generation - file_to_base_path = {} - for ts_file in ts_files: - for base_path in file_paths: - if ts_file.startswith(base_path) or ts_file.startswith(str(Path(base_path).resolve())): - file_to_base_path[ts_file] = base_path - break - - # Translate/interpret files concurrently with retry mechanism - translations = await self._translate_files_with_retry(ts_files, mode, file_to_base_path) - - # Prepare result - successful_count = len([t for t in translations if t.get("status") == "success"]) - skipped_count = len([t for t in translations if t.get("status") == "skipped"]) - failed_count = len([t for t in translations if t.get("status") == "failed"]) - - result = { - "status": "success", - "mode": mode_str, - "paths": file_paths, - "total_files": len(ts_files), - "successful_translations": successful_count, - "skipped_translations": skipped_count, - "failed_translations": failed_count, - "file_statistics": file_stats, - "translations": translations, - } - - logger.info( - f"{mode_str} completed: {successful_count} processed, {skipped_count} skipped, {failed_count} failed (total: {len(ts_files)})", - ) - self.set_result(json.dumps(result, ensure_ascii=False, indent=2)) - - @staticmethod - def _find_ts_files(directory_path: str) -> List[str]: - """ - Recursively find all TypeScript (.ts) and TSX (.tsx) files in the given directory - - Args: - directory_path: Root directory to search - - Returns: - List of TypeScript/TSX file paths - """ - path = Path(directory_path) - - if not path.exists(): - logger.warning(f"Path does not exist: {directory_path}") - return [] - - if path.is_file(): - # If it's a file, check if it's a .ts or .tsx file - if path.suffix in [".ts", ".tsx"]: - return [str(path)] - else: - logger.warning(f"Path is not a TypeScript/TSX file: {directory_path}") - return [] - - # Recursively find all .ts and .tsx files - ts_files = [] - for ts_file in path.rglob("*.ts"): - if ts_file.is_file(): - ts_files.append(str(ts_file)) - for tsx_file in path.rglob("*.tsx"): - if tsx_file.is_file(): - ts_files.append(str(tsx_file)) - - return sorted(ts_files) - - @staticmethod - def _analyze_file_statistics(ts_files: List[str]) -> Dict: - """ - Analyze statistics of TypeScript files (character count, top 5 largest, average, etc.) - - Args: - ts_files: List of TypeScript file paths - - Returns: - Dictionary containing file statistics: - - total_files: Total number of files - - total_chars: Total character count - - avg_chars: Average characters per file - - median_chars: Median character count - - top_5_largest: List of (file_path, char_count) tuples for top 5 largest files - """ - if not ts_files: - return { - "total_files": 0, - "total_chars": 0, - "avg_chars": 0, - "median_chars": 0, - "top_5_largest": [], - } - - # Collect file sizes - file_sizes: List[Tuple[str, int]] = [] - for ts_file in ts_files: - try: - content = Path(ts_file).read_text(encoding="utf-8") - char_count = len(content) - file_sizes.append((ts_file, char_count)) - except Exception as e: - logger.warning(f"Error reading {ts_file} for statistics: {e}") - file_sizes.append((ts_file, 0)) - - # Sort by size (descending) - file_sizes.sort(key=lambda x: x[1], reverse=True) - - # Calculate statistics - char_counts = [size for _, size in file_sizes] - total_files = len(ts_files) - total_chars = sum(char_counts) - avg_chars = total_chars / total_files if total_files > 0 else 0 - - # Calculate median - sorted_counts = sorted(char_counts) - n = len(sorted_counts) - if n % 2 == 0: - median_chars = (sorted_counts[n // 2 - 1] + sorted_counts[n // 2]) / 2 - else: - median_chars = sorted_counts[n // 2] - - # Get top 5 largest files - top_5_largest = file_sizes[:5] - - return { - "total_files": total_files, - "total_chars": total_chars, - "avg_chars": avg_chars, - "median_chars": int(median_chars), - "top_5_largest": top_5_largest, - } - - @staticmethod - def _get_output_file_path(ts_file_path: str, base_path: str, mode: bool) -> str: - """ - Generate corresponding output file path from TypeScript file path - - Args: - ts_file_path: Path to the TypeScript file (e.g., 'a/b/c/dir/file.ts') - base_path: Base path that was originally searched (e.g., 'a/b/c') - mode: True = Python translation, False = Markdown interpretation - - Returns: - Corresponding output file path: - - mode=True: 'a/b/c_python/dir/file.py' - - mode=False: 'a/b/c_markdown/dir/file.md' - """ - ts_path = Path(ts_file_path) - base = Path(base_path) - - # Get the relative path from base to the ts file - try: - relative_path = ts_path.relative_to(base) - except ValueError: - # If ts_file_path is not relative to base_path, just use the filename - relative_path = ts_path.name - - # Determine output directory and extension based on mode - if mode: - # Translation mode: a/b/c -> a/b/c_python - output_base = Path(str(base) + "_python") - extension = ".py" - else: - # Interpretation mode: a/b/c -> a/b/c_markdown - output_base = Path(str(base) + "_markdown") - extension = ".md" - - # Construct output path with new extension - output_path = output_base / relative_path.with_suffix(extension) - - return str(output_path) - - async def _translate_files_with_retry( - self, ts_files: List[str], mode: bool, file_to_base_path: Dict[str, str] - ) -> List[Dict]: - """ - Translate/interpret TypeScript files with retry mechanism for failed cases - - Args: - ts_files: List of TypeScript file paths - mode: True = translate to Python, False = interpret to markdown - file_to_base_path: Mapping from ts_file to its base_path - - Returns: - List of translation results - """ - all_results = {} # file_path -> result dict - files_to_translate = ts_files.copy() - retry_count = 0 - - while files_to_translate and retry_count <= self.max_retries: - if retry_count > 0: - logger.info( - f"Retry attempt {retry_count}/{self.max_retries} for {len(files_to_translate)} failed files...", - ) - - # Translate current batch with progress bar - translations = await self._translate_files_concurrently( - files_to_translate, mode, file_to_base_path, retry_round=retry_count - ) - - # Collect results and identify failed files - failed_files = [] - for translation in translations: - file_path = translation["file_path"] - all_results[file_path] = translation - - # Add retry info if this is a retry attempt - if retry_count > 0: - if "retry_attempts" not in translation: - translation["retry_attempts"] = [] - translation["retry_attempts"].append(retry_count) - - # Collect failed files for retry - if translation.get("status") == "failed" and retry_count < self.max_retries: - failed_files.append(file_path) - - # Update files to translate for next retry - files_to_translate = failed_files - retry_count += 1 - - # Convert results dict back to list, preserving original order - final_results = [all_results[f] for f in ts_files] - - # Add final retry statistics - retry_stats = { - "total_retries": retry_count - 1, - "files_retried": len([r for r in final_results if "retry_attempts" in r]), - } - logger.info(f"Retry statistics: {retry_stats}") - - return final_results - - async def _translate_single_file_safe(self, ts_file: str, mode: bool, base_path: str): - """Wrapper to safely execute translation/interpretation and catch exceptions""" - try: - return await self._translate_single_file(ts_file, mode, base_path) - except Exception as e: - logger.error(f"{'Translation' if mode else 'Interpretation'} failed for {ts_file}: {str(e)}") - return { - "file_path": ts_file, - "status": "failed", - "error": str(e), - } - - async def _translate_files_concurrently( - self, ts_files: List[str], mode: bool, file_to_base_path: Dict[str, str], retry_round: int = 0 - ) -> List[Dict]: - """ - Translate/interpret TypeScript files concurrently with semaphore control and submit interval - - Args: - ts_files: List of TypeScript file paths - mode: True = translate to Python, False = interpret to markdown - file_to_base_path: Mapping from ts_file to its base_path - retry_round: Current retry round number (0 for initial attempt) - - Returns: - List of translation results - """ - # Progress bar description - if retry_round > 0: - desc = f"Retry {retry_round}/{self.max_retries}" - else: - desc = "Translating" if mode else "Interpreting" - - # Create a wrapper task that includes delay - async def delayed_task(index: int, ts_file: str): - # Add delay before each task execution (except the first one) - if index > 0: - await asyncio.sleep(self.submit_interval) - logger.debug(f"Starting task {index} after {self.submit_interval}s delay") - return await self._translate_single_file_safe(ts_file, mode, file_to_base_path.get(ts_file, ts_file)) - - # Create all delayed tasks - tasks = [delayed_task(i, ts_file) for i, ts_file in enumerate(ts_files)] - - # Execute all tasks concurrently with progress bar - results = await tqdm_asyncio.gather( - *tasks, - desc=desc, - total=len(tasks), - unit="file", - ) - - return results - - async def _translate_single_file(self, ts_file_path: str, mode: bool, base_path: str) -> Dict: - """ - Translate/interpret a single TypeScript file using LLM with semaphore control - - Args: - ts_file_path: Path to the TypeScript file - mode: True = translate to Python, False = interpret to markdown - base_path: Base path for generating output path - - Returns: - Translation/interpretation result dictionary - """ - async with self.semaphore: - try: - # Read TypeScript file content - ts_content = Path(ts_file_path).read_text(encoding="utf-8") - - # Get output file path based on mode - output_file_path = self._get_output_file_path(ts_file_path, base_path, mode) - - # Check if target file already exists and is not empty - if self.skip_existing and Path(output_file_path).exists(): - try: - existing_content = Path(output_file_path).read_text(encoding="utf-8").strip() - if existing_content: - logger.info( - f"Skipping {ts_file_path} - target file {output_file_path} already exists with content", - ) - return { - "file_path": ts_file_path, - "status": "skipped", - "message": f"Target file {output_file_path} already exists", - "ts_content": ts_content, - "output_code": existing_content, - "output_file_path": output_file_path, - } - except Exception as e: - logger.warning(f"Error reading existing file {output_file_path}: {e}") - - # Create prompt based on mode - if mode: - # Translation mode: translate to Python - prompt = self.prompt_format( - prompt_name="translate_ts_to_python_prompt", - ts_content=ts_content, - ) - task_name = "Translating" - language_tag = "python" - else: - # Interpretation mode: explain in markdown - prompt = self.prompt_format( - prompt_name="interpret_ts_to_markdown_prompt", - ts_content=ts_content, - ) - task_name = "Interpreting" - language_tag = "markdown" - - # Call LLM with callback to extract code - logger.info(f"{task_name} {ts_file_path}...") - - # Callback function to extract code (only for translation mode) - if mode: - # Translation mode: extract Python code block - def extract_code(msg): - content = msg.content - - # Extract code block - extracted = extract_content(content, language_tag=language_tag) - return extracted if extracted else content - - output_code = await self.llm.achat( - messages=[Message(role=Role.USER, content=prompt)], - callback_fn=extract_code, - enable_stream_print=False, - ) - else: - # Interpretation mode: keep all content as is (no code extraction) - output_message = await self.llm.achat( - messages=[Message(role=Role.USER, content=prompt)], - enable_stream_print=False, - ) - output_code = output_message.content - - # Save output code to disk - output_path = Path(output_file_path) - output_path.parent.mkdir(parents=True, exist_ok=True) # Create parent directories if needed - output_path.write_text(output_code, encoding="utf-8") - logger.info(f"Saved {task_name.lower()} output to {output_file_path}") - - result = { - "file_path": ts_file_path, - "status": "success", - "ts_content": ts_content, - "output_code": output_code, - "output_file_path": output_file_path, - } - - return result - - except Exception as e: - logger.exception(f"Error processing {ts_file_path}: {e}") - return { - "file_path": ts_file_path, - "status": "failed", - "error": str(e), - } - - -async def main(): - """ - Test function demonstrating both translation and interpretation modes - - Example usage: - - Translation mode (mode=True): Translates TypeScript to Python, saves to *_python/*.py - - Interpretation mode (mode=False): Explains TypeScript in Markdown, saves to *_markdown/*.md - """ - from flowllm.app import FlowLLMApp - from flowllm.context.flow_context import FlowContext - - async with FlowLLMApp(load_default_config=True): - # Example: Translation mode (default) - print("=" * 80) - print("Example 1: Translation Mode (TypeScript -> Python)") - print("=" * 80) - op_translate = TranslateCodeOp(max_concurrent=2) - context1 = FlowContext() - await op_translate.async_call( - context1, - file_path="/Users/yuli/workspace/qwen-code", - mode=True, - ) - print(f"Translation Result Summary:\n{op_translate.output}\n") - - # Example: Interpretation mode - print("=" * 80) - print("Example 2: Interpretation Mode (TypeScript -> Markdown)") - print("=" * 80) - op_interpret = TranslateCodeOp(max_concurrent=2) - context2 = FlowContext() - await op_interpret.async_call( - context2, - file_path="/Users/yuli/workspace/qwen-code", - mode=False, - ) - print(f"Interpretation Result Summary:\n{op_interpret.output}\n") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/old/op/gallery/translate_code_prompt.yaml b/old/op/gallery/translate_code_prompt.yaml deleted file mode 100644 index 6a7cdec..0000000 --- a/old/op/gallery/translate_code_prompt.yaml +++ /dev/null @@ -1,14 +0,0 @@ -translate_ts_to_python_prompt: | - ```typescript - {ts_content} - ``` - 请将上面代码翻译成 Python 代码。 - ```python - ... - ``` - -interpret_ts_to_markdown_prompt: | - ```typescript - {ts_content} - ``` - 请用中文详细解读代码，要求举例。 diff --git a/old/op/search/__init__.py b/old/op/search/__init__.py deleted file mode 100644 index 381eccc..0000000 --- a/old/op/search/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .dashscope_search_op import DashscopeSearchOp -from .llm_mock_search_op import LLMMockSearchOp -from .mcp_search_op import TongyiMcpSearchOp, BochaMcpSearchOp -from .tavily_search_op import TavilySearchOp diff --git a/old/op/search/dashscope_search_op.py b/old/op/search/dashscope_search_op.py deleted file mode 100644 index 8153cf4..0000000 --- a/old/op/search/dashscope_search_op.py +++ /dev/null @@ -1,119 +0,0 @@ -import asyncio -import os -from typing import Dict, Any, List - -import dashscope -from loguru import logger - -from flowllm.context import FlowContext, C -from flowllm.op.base_async_tool_op import BaseAsyncToolOp -from flowllm.schema.tool_call import ToolCall - - -@C.register_op(register_app="FlowLLM") -class DashscopeSearchOp(BaseAsyncToolOp): - file_path: str = __file__ - - def __init__( - self, - model: str = "qwen-plus", - search_strategy: str = "max", - enable_role_prompt: bool = True, - **kwargs, - ): - super().__init__(**kwargs) - - self.model: str = model - self.search_strategy: str = search_strategy - self.enable_role_prompt: bool = enable_role_prompt - - self.api_key = os.getenv("FLOW_DASHSCOPE_API_KEY", "") - - def build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "Use search keywords to retrieve relevant information from the internet. If there are multiple search keywords, please use each keyword separately to call this tool.", - "input_schema": { - "query": { - "type": "string", - "description": "search keyword", - "required": True, - }, - }, - }, - ) - - @staticmethod - def format_search_results(search_results: List[Dict[str, Any]]) -> str: - """Format search results for display""" - formatted_results = ["=" * 20 + " Search Results " + "=" * 20] - for web in search_results: - formatted_results.append(f"[{web['index']}]: [{web['title']}]({web['url']})") - - return "\n".join(formatted_results) - - async def async_execute(self): - query: str = self.input_dict["query"] - - if self.enable_cache: - cached_result = self.cache.load(query) - if cached_result: - self.set_result(cached_result["response_content"]) - return - - if self.enable_role_prompt: - user_query = self.prompt_format(prompt_name="role_prompt", query=query) - else: - user_query = query - logger.info(f"user_query={user_query}") - messages: list = [{"role": "user", "content": user_query}] - - response = await dashscope.AioGeneration.call( - api_key=self.api_key, - model=self.model, - messages=messages, - enable_search=True, # Enable web search - search_options={ - "forced_search": True, # Force web search - "enable_source": True, # Include search source information - "enable_citation": False, # Enable citation markers - "search_strategy": self.search_strategy, # Search strategy - }, - result_format="message", - ) - - search_results = [] - response_content = "" - - if hasattr(response, "output") and response.output: - if hasattr(response.output, "search_info") and response.output.search_info: - search_results = response.output.search_info.get("search_results", []) - - if hasattr(response.output, "choices") and response.output.choices and len(response.output.choices) > 0: - response_content = response.output.choices[0].message.content - - final_result = { - "query": query, - "search_results": search_results, - "response_content": response_content, - "model": self.model, - "search_strategy": self.search_strategy, - } - - if self.enable_cache: - self.cache.save(query, final_result, expire_hours=self.cache_expire_hours) - - self.set_result(final_result["response_content"]) - - -async def async_main(): - # op = DashscopeSearchOp(enable_role_prompt=False) - op = DashscopeSearchOp(model="qwen3-max", enable_role_prompt=False) - # context = FlowContext(query="what is AI?") - context = FlowContext(query="藏格矿业的业务主要有哪几块？营收和利润的角度分析 雪球") - await op.async_call(context=context) - print(context.dashscope_search_result) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/old/op/search/dashscope_search_prompt.yaml b/old/op/search/dashscope_search_prompt.yaml deleted file mode 100644 index b8698ab..0000000 --- a/old/op/search/dashscope_search_prompt.yaml +++ /dev/null @@ -1,13 +0,0 @@ -role_prompt: | - # user's question - {query} - - # task - Extract the original content related to the user's query directly from the context, maintain accuracy, and avoid excessive processing. - -role_prompt_zh: | - # 用户问题 - {query} - - # task - 直接从上下文中提取与用户问题相关的原始内容，保持准确性，避免过度处理。 \ No newline at end of file diff --git a/old/op/search/llm_mock_search_prompt.yaml b/old/op/search/llm_mock_search_prompt.yaml deleted file mode 100644 index 4718340..0000000 --- a/old/op/search/llm_mock_search_prompt.yaml +++ /dev/null @@ -1,76 +0,0 @@ -classification_prompt: | - You are a query complexity classifier. Analyze the following search query and classify it into one of three categories: - - 1. **simple** - Simple factual queries that: - - Ask for a single, direct fact - - Have a clear, unambiguous answer - - Require minimal context or explanation - - Examples: "What is the capital of France?", "Who invented the telephone?", "When did World War 2 end?" - - 2. **medium** - Medium complexity queries that: - - Require some explanation or context - - May involve multiple related facts - - Need balanced depth without being exhaustive - - Examples: "How does photosynthesis work?", "What are the main causes of climate change?", "Explain blockchain technology" - - 3. **complex** - Complex research queries that: - - Require comprehensive, multi-dimensional analysis - - Involve multiple subtopics or perspectives - - Need in-depth exploration and connections - - Examples: "Analyze the impact of AI on the global economy", "Compare different renewable energy solutions", "What are the geopolitical implications of space exploration?" - - Query to classify: {query} - - Respond with ONLY one word: simple, medium, or complex. - -classification_prompt_zh: | - 你是一个查询复杂度分类器。分析以下搜索查询并将其分类为以下三类之一： - - 1. **simple** - 简单事实查询： - - 询问单一、直接的事实 - - 有明确、无歧义的答案 - - 需要最少的上下文或解释 - - 示例："法国的首都是什么？"、"谁发明了电话？"、"第二次世界大战何时结束？" - - 2. **medium** - 中等复杂度查询： - - 需要一些解释或上下文 - - 可能涉及多个相关事实 - - 需要平衡的深度但不需要详尽无遗 - - 示例："光合作用如何工作？"、"气候变化的主要原因是什么？"、"解释区块链技术" - - 3. **complex** - 复杂研究查询： - - 需要全面、多维度的分析 - - 涉及多个子主题或观点 - - 需要深入探索和联系 - - 示例："分析人工智能对全球经济的影响"、"比较不同的可再生能源解决方案"、"太空探索的地缘政治影响是什么？" - - 要分类的查询：{query} - - 只用一个词回答：simple、medium 或 complex。 - -generation_prompt: | - You are a search engine generating mock search results. Generate a {content_length} response for the following query. - - Query: {query} - Complexity Level: {complexity} - - Instructions based on content length: - - **short**: Provide a concise answer in 1-3 sentences. Be direct and factual. - - **medium**: Provide a balanced answer in 2-4 paragraphs. Include key details and some context. - - **long**: Provide a comprehensive answer in 4-6 paragraphs. Include multiple perspectives, detailed explanations, and relevant context. - - Generate the search result content now: - -generation_prompt_zh: | - 你是一个搜索引擎，正在生成模拟搜索结果。为以下查询生成一个 {content_length} 的响应。 - - 查询：{query} - 复杂度级别：{complexity} - - 根据内容长度的指示： - - **short**（短）：提供 1-3 句话的简洁答案。要直接和事实性。 - - **medium**（中）：提供 2-4 段的平衡答案。包括关键细节和一些上下文。 - - **long**（长）：提供 4-6 段的全面答案。包括多个角度、详细解释和相关上下文。 - - 现在生成搜索结果内容： - diff --git a/old/op/simple_llm_op.py b/old/op/simple_llm_op.py deleted file mode 100644 index 4485d81..0000000 --- a/old/op/simple_llm_op.py +++ /dev/null @@ -1,54 +0,0 @@ -import asyncio -from typing import List - -from loguru import logger - -from flowllm.context.flow_context import FlowContext -from flowllm.context.service_context import C -from flowllm.op.base_async_tool_op import BaseAsyncToolOp -from flowllm.schema.message import Message, Role -from flowllm.schema.tool_call import ToolCall - - -@C.register_op(register_app="FlowLLM") -class SimpleLLMOp(BaseAsyncToolOp): - - def __init__(self, llm: str = "qwen3_30b_thinking", save_answer: bool = True, **kwargs): - super().__init__(llm=llm, save_answer=save_answer, **kwargs) - - def build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "use this query to query an LLM", - "input_schema": { - "query": { - "type": "string", - "description": "search keyword", - "required": True, - }, - }, - }, - ) - - async def async_execute(self): - query: str = self.input_dict["query"] - logger.info(f"query={query}") - messages: List[Message] = [Message(role=Role.USER, content=query)] - assistant_message: Message = await self.llm.achat(messages) - self.set_result(assistant_message.content) - - -async def main(): - from flowllm.app import FlowLLMApp - - async with FlowLLMApp(load_default_config=True): - context = FlowContext(query="hello", stream_queue=asyncio.Queue()) - - op = SimpleLLMOp() - result = await op.async_call(context=context) - print(op.output) - print(result) - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/old/op/stream_llm_op.py b/old/op/stream_llm_op.py deleted file mode 100644 index 8c22d8c..0000000 --- a/old/op/stream_llm_op.py +++ /dev/null @@ -1,83 +0,0 @@ -import asyncio -import json -from typing import List - -from loguru import logger - -from flowllm.context.flow_context import FlowContext -from flowllm.context.service_context import C -from flowllm.enumeration.chunk_enum import ChunkEnum -from flowllm.op.base_async_tool_op import BaseAsyncToolOp -from flowllm.schema.message import Message, Role -from flowllm.schema.tool_call import ToolCall - - -@C.register_op(register_app="FlowLLM") -class StreamLLMOp(BaseAsyncToolOp): - - def __init__(self, llm: str = "qwen3_30b_thinking", save_answer: bool = True, **kwargs): - super().__init__(llm=llm, save_answer=save_answer, **kwargs) - - def build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "description": "use this query to query an LLM", - "input_schema": { - "query": { - "type": "string", - "description": "search keyword", - "required": False, - }, - "messages": { - "type": "array", - "description": "messages", - "required": False, - }, - }, - }, - ) - - async def async_execute(self): - if self.input_dict.get("query"): - query: str = self.input_dict.get("query") - messages: List[Message] = [Message(role=Role.USER, content=query)] - elif self.input_dict.get("messages"): - messages: list = self.input_dict.get("messages") - messages: List[Message] = [Message(**x) for x in messages] - else: - raise RuntimeError("query or messages is required") - - logger.info(f"messages={messages}") - - async for chunk, chunk_type in self.llm.astream_chat(messages): # noqa - if chunk_type in [ChunkEnum.ANSWER, ChunkEnum.THINK, ChunkEnum.ERROR]: - await self.context.add_stream_chunk_and_type(chunk, chunk_type) - elif chunk_type == ChunkEnum.TOOL: - await self.context.add_stream_chunk_and_type( - chunk=json.dumps([x.model_dump() for x in chunk], ensure_ascii=False), - chunk_type=ChunkEnum.TOOL, - ) - - -async def main(): - from flowllm.app import FlowLLMApp - - async with FlowLLMApp(load_default_config=True): - - context = FlowContext(query="what is ai?", stream_queue=asyncio.Queue()) - op = StreamLLMOp() - task = asyncio.create_task(op.async_call(context=context)) - - while True: - stream_chunk = await context.stream_queue.get() - if stream_chunk.done: - print("\nend") - break - else: - print(stream_chunk.chunk, end="") - - await task - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/old/op/terminate_op.py b/old/op/terminate_op.py deleted file mode 100644 index 1805311..0000000 --- a/old/op/terminate_op.py +++ /dev/null @@ -1,42 +0,0 @@ -import asyncio - -from flowllm.context.flow_context import FlowContext -from flowllm.context.service_context import C -from flowllm.op.base_async_tool_op import BaseAsyncToolOp -from flowllm.schema.tool_call import ToolCall - - -@C.register_op(register_app="FlowLLM") -class TerminateOp(BaseAsyncToolOp): - - def build_tool_call(self) -> ToolCall: - return ToolCall( - **{ - "name": "terminate", - "description": "If you can answer the question based on the context, be sure to use the **terminate** tool.", - "input_schema": { - "status": { - "type": "str", - "description": "If the user's question can be answered, return success, otherwise return failure.", - "required": True, - "enum": ["success", "failure"], - }, - }, - }, - ) - - async def async_execute(self): - status: str = self.input_dict["status"] - assert status in ["success", "failure"], f"Invalid status: {status}" - self.set_result(f"The interaction has been completed with status: {status}") - - -async def main(): - op = TerminateOp() - context = FlowContext(status="success") - await op.async_call(context) - print(f"Result: {op.output}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/old/pyproject.toml b/old/pyproject.toml deleted file mode 100644 index 88648f8..0000000 --- a/old/pyproject.toml +++ /dev/null @@ -1,91 +0,0 @@ -[build-system] -requires = ["setuptools>=61.0", "wheel"] -build-backend = "setuptools.build_meta" - -[project] -name = "flowllm" -version = "0.1.11.6" -description = "A flexible framework for building LLM-powered flows and mcp services" -authors = [ - { name = "jinli.yl", email = "jinli.yl@alibaba-inc.com" }, -] -license = {file = "LICENSE"} -readme = "README.md" -keywords = ["llm", "ai", "flow", "framework", "openai", "chatgpt", "language-model", "mcp", "http"] -requires-python = ">=3.10" - -classifiers = [ - "Development Status :: 4 - Beta", - "Intended Audience :: Developers", - "Intended Audience :: Science/Research", - "License :: OSI Approved :: Apache Software License", - "Operating System :: OS Independent", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.10", - "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries :: Python Modules", - "Topic :: Software Development :: Libraries :: Application Frameworks", - "Typing :: Typed", -] - -dependencies = [ - "dashscope>=1.13.0", - "fastapi>=0.104.0", - "fastmcp>=0.1.0", - "httpx>=0.25.0", - - "loguru>=0.7.0", - "mcp>=0.1.0", - "numpy>=1.24.0", - "pandas>=2.0.0", - "openai>=1.3.0", - "pydantic>=2.0.0", - "PyYAML>=6.0.0", - "requests>=2.31.0", - "tqdm>=4.66.0", - "urllib3>=2.0.0", - "uvicorn[standard]>=0.23.0", - "pyfiglet>=0.8.0", - "rich>=13.5.0", - "tavily-python>=0.2.0", -] - -[project.optional-dependencies] - -full = [ - "litellm>=1.0.0", - "pytest>=8.4.2", - "pytest_asyncio>=1.2.0", -] - -reme = ["elasticsearch>=8.10.0", "chromadb>=0.4.0", "qdrant_client>=1.15.1"] - -ray = ["ray>=2.6.0", "scikit-learn>=1.3.0"] - -fin = ["akshare>=1.9.0", "crawl4ai>=0.7.4"] - -token = ["modelscope>=1.10.0", "diffusers>=0.21.0", "peft>=0.6.0", "transformers>=4.35.0", "tiktoken>=0.5.0"] - -all = ["flowllm[reme]", "flowllm[ray]", "flowllm[fin]", "flowllm[token]"] - -[tool.setuptools.packages.find] -where = ["."] -include = ["flowllm*"] -exclude = ["cookbook*"] - -[tool.setuptools.package-data] -flowllm = [ - "**/*.yaml", -] - -[project.urls] -Homepage = "https://github.com/xxx-ai/xxxx" -Documentation = "https://doc.xxx.io/" -Repository = "https://github.com/xxx-ai/xxx" - -[project.scripts] -flowllm = "flowllm.app:main" - -# conda create -n fl_test python=3.12 -# conda activate fl_test -# conda env remove -n fl_test \ No newline at end of file diff --git a/old/test/__init__.py b/old/test/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/old/test/auto_screenshot_analyzer.py b/old/test/auto_screenshot_analyzer.py deleted file mode 100644 index c26bd65..0000000 --- a/old/test/auto_screenshot_analyzer.py +++ /dev/null @@ -1,459 +0,0 @@ -#!/usr/bin/env python3 -""" -自动截图分析脚本 -每5分钟对Mac电脑进行截图，使用VL模型分析用户正在做什么，并保存分析结果 -""" - -import os -import time -import json -import base64 -from datetime import datetime -from typing import List, Dict -import threading -import signal -import sys - -from loguru import logger - -# 检查必要的依赖 -try: - from openai import OpenAI -except ImportError: - print("❌ 错误: 未找到 openai 包") - print("请运行: pip install openai") - sys.exit(1) - -try: - from PIL import ImageGrab, Image -except ImportError: - print("❌ 错误: 未找到 Pillow 包") - print("请运行: pip install Pillow") - sys.exit(1) - -try: - from flowllm.utils.common_utils import load_env -except ImportError: - print("❌ 错误: 未找到 flowllm 包") - print("请确保已安装 flowllm: pip install -e .") - sys.exit(1) - -# 加载环境变量 -load_env() - - -class AutoScreenshotAnalyzer: - def __init__(self): - # 检查必要的环境变量 - api_key = os.getenv("FLOW_LLM_API_KEY") - if not api_key: - print("❌ 错误: 未找到 FLOW_LLM_API_KEY 环境变量") - print("请设置环境变量或在 .env 文件中配置 API Key") - sys.exit(1) - - self.client = OpenAI( - api_key=api_key, - base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", - ) - - # 配置参数 - self.model_name = "qwen3-vl-235b-a22b-instruct" - self.interval_minutes = 2 - self.max_history_context = 10 # 最多保留10条历史记录作为上下文 - - # 创建保存目录 - self.screenshot_dir = "/Users/yuli/workspace/flowllm/logs/screenshots" - self.analysis_file = "/Users/yuli/workspace/flowllm/logs/activity_analysis.json" - - os.makedirs(self.screenshot_dir, exist_ok=True) - os.makedirs(os.path.dirname(self.analysis_file), exist_ok=True) - - # 加载历史分析记录 - self.analysis_history = self.load_analysis_history() - - self.running = True - - def load_analysis_history(self) -> List[Dict]: - """加载历史分析记录""" - if os.path.exists(self.analysis_file): - try: - with open(self.analysis_file, 'r', encoding='utf-8') as f: - return json.load(f) - except (json.JSONDecodeError, FileNotFoundError): - print("无法加载历史记录，从空记录开始") - return [] - return [] - - def save_analysis_history(self): - """保存分析历史到文件""" - try: - with open(self.analysis_file, 'w', encoding='utf-8') as f: - json.dump(self.analysis_history, f, ensure_ascii=False, indent=2) - print(f"✅ 分析历史已保存到: {self.analysis_file}") - except Exception as e: - print(f"❌ 保存分析历史失败: {e}") - - def get_screen_info(self) -> List[Dict]: - """获取屏幕信息""" - try: - screens = [] - - # 使用macOS的system_profiler获取详细的显示器信息 - try: - import subprocess - import re - - result = subprocess.run(['system_profiler', 'SPDisplaysDataType'], - capture_output=True, text=True, timeout=10) - - if result.returncode == 0: - output = result.stdout - print(f"🖥️ system_profiler 输出:\n{output[:500]}...") - - # 解析显示器信息 - displays = [] - display_blocks = re.split(r'\n\s*([^:]+):\s*\n', output)[1:] # 跳过第一个空元素 - - for i in range(0, len(display_blocks), 2): - if i + 1 >= len(display_blocks): - break - - display_name = display_blocks[i].strip() - display_info = display_blocks[i + 1] - - # 寻找分辨率信息 - resolution_match = re.search(r'Resolution:\s*(\d+)\s*x\s*(\d+)', display_info) - if resolution_match: - width = int(resolution_match.group(1)) - height = int(resolution_match.group(2)) - displays.append({ - 'name': display_name, - 'width': width, - 'height': height, - 'info': display_info - }) - - print(f"🖥️ 解析到 {len(displays)} 个显示器") - for i, display in enumerate(displays): - print(f" 显示器 {i+1}: {display['name']} - {display['width']}x{display['height']}") - - # 根据解析结果构建屏幕信息 - if len(displays) >= 2: - # 多显示器情况，为每个显示器创建屏幕区域 - # 假设显示器是水平排列的（这是最常见的情况） - x_offset = 0 - for i, display in enumerate(displays[:2]): # 最多支持2个屏幕 - screens.append({ - 'id': i, - 'bbox': (x_offset, 0, x_offset + display['width'], display['height']), - 'name': f"显示器{i+1} ({display['name'][:20]})" - }) - x_offset += display['width'] - - elif len(displays) == 1: - # 单显示器 - display = displays[0] - screens.append({ - 'id': 0, - 'bbox': (0, 0, display['width'], display['height']), - 'name': f"主显示器 ({display['name'][:20]})" - }) - - except Exception as e: - print(f"❌ system_profiler 解析失败: {e}") - - # 如果 system_profiler 解析失败，fallback 到简单的全屏截图 - if not screens: - print("🖥️ 使用 fallback 方案：全屏截图") - screens.append({ - 'id': 0, - 'bbox': None, # None 表示全屏截图 - 'name': '全屏幕' - }) - - return screens - - except Exception as e: - print(f"❌ 获取屏幕信息失败: {e}") - # 返回默认单屏幕配置 - return [{'id': 0, 'bbox': None, 'name': '默认屏幕'}] - - def take_screenshots(self) -> List[str]: - """截取所有屏幕并保存到文件""" - screens = self.get_screen_info() - print(f"🖥️ 一共 {len(screens)} 个屏幕") - - screenshot_paths = [] - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - - for screen in screens: - print(f"🖥️ 处理 {screen['name']} 屏幕") - - # 生成文件名 - filename = f"screenshot_{timestamp}_screen{screen['id']}.png" - filepath = os.path.join(self.screenshot_dir, filename) - - # 使用 macOS 的 screencapture 命令截取特定显示器 - # -D 参数指定显示器编号 (1, 2, 3...) - display_num = screen['id'] + 1 # screencapture 的显示器编号从1开始 - - import subprocess - try: - # 使用 screencapture 命令截取特定显示器 - result = subprocess.run([ - 'screencapture', - '-D', str(display_num), # 指定显示器 - '-x', # 不播放快门声音 - filepath - ], capture_output=True, text=True, timeout=10) - - if result.returncode == 0 and os.path.exists(filepath): - screenshot_paths.append(filepath) - print(f"📸 {screen['name']} 截图已保存: {filepath}") - else: - print(f"❌ screencapture 命令失败: {result.stderr}") - # 尝试使用 PIL 作为备用方案 - if screen['bbox']: - screenshot = ImageGrab.grab(bbox=screen['bbox']) - screenshot.save(filepath) - screenshot_paths.append(filepath) - print(f"📸 {screen['name']} 使用PIL备用截图: {filepath}") - else: - raise Exception(f"screencapture failed and no bbox available") - - except subprocess.TimeoutExpired: - print(f"❌ screencapture 命令超时") - raise Exception("screencapture timeout") - except Exception as e: - print(f"❌ {screen['name']} 截图失败: {e}") - raise e - - return screenshot_paths - - def encode_image(self, image_path: str) -> str: - """将图像文件转换为Base64编码""" - try: - with open(image_path, "rb") as image_file: - return base64.b64encode(image_file.read()).decode("utf-8") - except Exception as e: - print(f"❌ 图像编码失败: {e}") - return None - - def build_context_prompt(self, num_screens: int = 1) -> str: - """构建包含历史上下文的提示词""" - if num_screens > 1: - base_prompt = f"""你是一个智能助手，正在分析用户的多屏幕活动。用户有{num_screens}个屏幕，请根据所有截图内容综合分析用户当前在做什么。 - -请提供： -1. 每个屏幕的内容简要描述 -2. 多屏幕之间的关联性分析 -3. 推理用户正在进行的整体活动 -4. 一句话总结用户的当前状态 - -""" - else: - base_prompt = """你是一个智能助手，正在分析用户的屏幕活动。请根据截图内容分析用户当前在做什么。 - -请提供： -1. 当前屏幕内容的简要描述 -2. 推理用户正在进行的活动 -3. 一句话总结用户的当前状态 - -""" - - # 添加历史上下文 - if self.analysis_history: - base_prompt += "\n📜 历史活动上下文（最近的活动）:\n" - recent_history = self.analysis_history[-self.max_history_context:] - - for i, record in enumerate(recent_history, 1): - time_str = record.get('timestamp', '未知时间') - activity = record.get('activity_summary', '未知活动') - base_prompt += f"{i}. [{time_str}] {activity}\n" - - base_prompt += "\n基于以上历史活动上下文和当前截图，分析用户现在的活动状态。\n" - - base_prompt += "\n请用中文回答，保持简洁明了。" - return base_prompt - - def analyze_screenshots(self, image_paths: List[str]) -> Dict: - """使用VL模型批量分析多张截图""" - try: - if not image_paths: - return None - - # 编码所有图像 - image_contents = [] - valid_paths = [] - - for i, image_path in enumerate(image_paths): - base64_image = self.encode_image(image_path) - if base64_image: - image_contents.append({ - "type": "image_url", - "image_url": {"url": f"data:image/png;base64,{base64_image}"} - }) - valid_paths.append(image_path) - print(f"📷 图像 {i+1} 编码成功: {os.path.basename(image_path)}") - else: - print(f"❌ 图像 {i+1} 编码失败: {image_path}") - - if not image_contents: - return None - - # 构建消息 - 支持多张图片 - context_prompt = self.build_context_prompt(num_screens=len(valid_paths)) - - # 构建用户消息内容：先放所有图片，再放文本提示 - user_content = image_contents + [ - { - "type": "text", - "text": context_prompt - } - ] - - messages = [ - { - "role": "system", - "content": [{"type": "text", "text": "你是一个专业的屏幕活动分析助手。"}] - }, - { - "role": "user", - "content": user_content - } - ] - - # 调用模型 - print(f"🤖 正在分析 {len(valid_paths)} 张截图...") - completion = self.client.chat.completions.create( - model=self.model_name, - messages=messages, - temperature=0.3 - ) - - analysis_result = completion.choices[0].message.content - print(f"✅ 批量分析完成: {analysis_result[:100]}...") - - return { - "timestamp": datetime.now().isoformat(), - "screenshot_paths": valid_paths, - "num_screens": len(valid_paths), - "analysis_result": analysis_result, - "activity_summary": self.extract_summary(analysis_result) - } - - except Exception as e: - logger.exception(e) - print(f"❌ 批量分析截图失败: {e}") - raise e - return None - - def extract_summary(self, analysis_text: str) -> str: - """从分析结果中提取一句话总结""" - # 简单的提取逻辑，寻找包含"总结"或"状态"的行 - lines = analysis_text.split('\n') - for line in lines: - if any(keyword in line for keyword in ['总结', '状态', '正在', '在做']): - # 移除可能的序号和标点 - summary = line.strip('123456789. ：:').strip() - if summary: - return summary - - # 如果没有找到特定格式，返回前50个字符作为总结 - return analysis_text[:50] + "..." if len(analysis_text) > 50 else analysis_text - - def process_cycle(self): - """执行一次完整的截图分析流程""" - print(f"\n🔄 开始新的分析周期 - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") - - # 1. 多屏幕截图 - screenshot_paths = self.take_screenshots() - if not screenshot_paths: - print("❌ 没有截图成功，跳过本次分析") - return - - print(f"📸 成功截取 {len(screenshot_paths)} 个屏幕") - - # 2. 批量分析所有截图 - analysis = self.analyze_screenshots(screenshot_paths) - if not analysis: - print("❌ 分析失败，跳过本次记录") - return - - # 3. 保存到历史记录 - self.analysis_history.append(analysis) - - # 4. 限制历史记录长度，避免文件过大 - max_total_history = 1000 # 最多保留1000条记录 - if len(self.analysis_history) > max_total_history: - self.analysis_history = self.analysis_history[-max_total_history:] - - # 5. 持久化保存 - self.save_analysis_history() - - # 6. 显示当前分析结果 - screen_info = f"({analysis['num_screens']} 屏幕)" if analysis['num_screens'] > 1 else "" - print(f"📊 当前活动总结 {screen_info}: {analysis['activity_summary']}") - - def start_monitoring(self): - """开始监控循环""" - print(f"🚀 开始自动截图分析监控") - print(f"⏰ 间隔时间: {self.interval_minutes} 分钟") - print(f"📁 截图保存目录: {self.screenshot_dir}") - print(f"📄 分析结果文件: {self.analysis_file}") - print(f"📜 历史记录数量: {len(self.analysis_history)}") - print("按 Ctrl+C 停止监控\n") - - # 首次执行 - self.process_cycle() - - # 开始定时循环 - while self.running: - try: - # 等待指定时间 - for _ in range(self.interval_minutes * 60): # 转换为秒 - if not self.running: - break - time.sleep(1) - - if self.running: - self.process_cycle() - - except KeyboardInterrupt: - self.stop_monitoring() - break - except Exception as e: - print(f"❌ 监控循环出错: {e}") - print("等待下一个周期...") - - def stop_monitoring(self): - """停止监控""" - print("\n🛑 正在停止监控...") - self.running = False - self.save_analysis_history() - print("✅ 监控已停止，数据已保存") - - def signal_handler(self, signum, frame): - """处理系统信号""" - print(f"\n接收到信号 {signum}") - self.stop_monitoring() - sys.exit(0) - - -def main(): - """主函数""" - analyzer = AutoScreenshotAnalyzer() - - # 注册信号处理器 - signal.signal(signal.SIGINT, analyzer.signal_handler) - signal.signal(signal.SIGTERM, analyzer.signal_handler) - - try: - analyzer.start_monitoring() - except Exception as e: - print(f"❌ 程序运行出错: {e}") - analyzer.stop_monitoring() - - -if __name__ == "__main__": - main() diff --git a/old/test/client_examples.py b/old/test/client_examples.py deleted file mode 100644 index 339ce72..0000000 --- a/old/test/client_examples.py +++ /dev/null @@ -1,80 +0,0 @@ -""" -Example usage of FlowLLM service clients -""" -import asyncio -import json - -from flowllm.client import HttpClient, AsyncHttpClient, MCPClient, SyncMCPClient - - -def http_client_example(): - """Example of using synchronous HTTP client""" - print("=== HTTP Client Example (Sync) ===") - - with HttpClient(base_url="http://localhost:8001") as client: - health = client.health_check() - print(f"Service health: {health}") - - result = client.execute_tool_flow(flow_name="get_a_stock_infos", query="茅台怎么样？") - print(f"Flow result: {result.answer}") - - result = client.list_tool_flows() - print(json.dumps(result, ensure_ascii=False)) - -async def async_http_client_example(): - """Example of using asynchronous HTTP client""" - print("\n=== HTTP Client Example (Async) ===") - - async with AsyncHttpClient(base_url="http://localhost:8001") as client: - health = await client.health_check() - print(f"Service health: {health}") - - result = await client.execute_tool_flow(flow_name="get_a_stock_infos", query="茅台怎么样？") - print(f"Flow result: {result.answer}") - - result = await client.list_tool_flows() - print(json.dumps(result, ensure_ascii=False)) - - - -async def mcp_client_example(): - print("\n=== MCP Client Example (Async) ===") - - async with MCPClient(transport="sse", host="0.0.0.0", port=8001) as client: - tools = await client.list_tools() - for tool in tools: - print(tool.model_dump_json()) - - if tools: - result = await client.call_tool(tool_name=tools[0].name, arguments={"query": "茅台怎么样？"}) - print(f"tool result: {result}") - - -def sync_mcp_client_example(): - """Example of using synchronous MCP client""" - print("\n=== MCP Client Example (Sync) ===") - - with SyncMCPClient(transport="sse", host="0.0.0.0", port=8001) as client: - # List available tools - tools = client.list_tools() - for tool in tools: - print(tool.model_dump_json()) - - if tools: - result = client.call_tool(tool_name=tools[1].name, arguments={"query": "茅台怎么样？"}) - print(f"tool result: {result}") - - -async def main(): - await async_http_client_example() - # await mcp_client_example() - - -def main_sync(): - # http_client_example() - sync_mcp_client_example() - - -if __name__ == "__main__": - # asyncio.run(main()) - main_sync() diff --git a/old/test/http_client_test.py b/old/test/http_client_test.py deleted file mode 100644 index 37af53b..0000000 --- a/old/test/http_client_test.py +++ /dev/null @@ -1,74 +0,0 @@ -import asyncio -import json - -from flowllm.client import AsyncHttpClient - - -async def main(): - async with AsyncHttpClient("http://0.0.0.0:8002") as client: - - # Test 1: Health check - print("=" * 50) - print("Testing health check endpoint...") - try: - health_result = await client.health_check() - print(f"Health check result: {json.dumps(health_result, indent=2)}") - except Exception as e: - print(f"Health check failed: {e}") - - # Test 2: Get available endpoints - print("=" * 50) - print("Getting available endpoints...") - try: - openapi_schema = await client.list_available_endpoints() - if openapi_schema and "paths" in openapi_schema: - print("Available endpoints:") - for path, methods in openapi_schema["paths"].items(): - for method, details in methods.items(): - if method.upper() == "POST": - print(f" {method.upper()} {path} - {details.get('summary', 'No description')}") - else: - print("Could not retrieve endpoint information") - except Exception as e: - print(f"Failed to get endpoints: {e}") - - query1 = "阿里巴巴怎么样？" - query2 = "寒武纪还可以买吗？" - - # response = await client.execute_flow("mock_tool_flow", a=query2) - # print(f"result: {response.answer}") - # - # response = await client.execute_flow("mock_async_tool_flow", a=query2) - # print(f"result: {response.answer}") - - # response = await client.execute_flow("ant_search", query="阿里巴巴怎么样？", entity="阿里巴巴") - # print(f"result: {response.answer}") - # - # response = await client.execute_flow("ant_investment", entity="阿里巴巴", analysis_category="股票") - # print(f"result: {response.answer}") - - # response = await client.execute_flow("dashscope_search_tool_flow", query="阿里巴巴怎么样？") - # print(f"result: {response.answer}") - - # response = await client.execute_flow("get_a_stock_infos", query=query2) - # print(f"result: {response.answer}") - - # response = await client.execute_flow("get_a_stock_news", query=query2) - # print(f"result: {response.answer}") - - # response = await client.execute_flow("llm_flow", query=query2) - # print(f"result: {response.answer}") - - # response = await client.execute_flow("react_llm_tool_flow", query=query2) - # print(f"result: {response.answer}") - - # response = await client.execute_flow("tavily_search_tool_flow", query=query2) - # print(f"result: {response.answer}") - - """ - curl -X POST http://0.0.0.0:8002/llm_flow_stream -H "Content-Type: application/json" -d '{"query": "what is AI?"}' - """ - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/old/test/http_proxy_server.py b/old/test/http_proxy_server.py deleted file mode 100644 index 4f3e78a..0000000 --- a/old/test/http_proxy_server.py +++ /dev/null @@ -1,116 +0,0 @@ -#!/usr/bin/env python3 -""" -HTTP Proxy Server - Forwards all HTTP requests to a target server -""" -import requests -from flask import Flask, request, Response -import logging - -# Configure logging -logging.basicConfig( - level=logging.INFO, - format='%(asctime)s - %(levelname)s - %(message)s' -) -logger = logging.getLogger(__name__) - -app = Flask(__name__) - -# Configure target server -TARGET_HOST = "8.130.105.202" # Change to target IP -TARGET_PORT = 8010 # Change to target port -TARGET_SCHEME = "http" # http or https - -TARGET_URL = f"{TARGET_SCHEME}://{TARGET_HOST}:{TARGET_PORT}" - - -@app.route('/', defaults={'path': ''}, methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS', 'HEAD']) -@app.route('/', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS', 'HEAD']) -def proxy(path): - """ - Proxy all requests to the target server - """ - # Build target URL - url = f"{TARGET_URL}/{path}" - - # Get request parameters - params = request.args.to_dict() - - # Get request headers (exclude host header) - headers = {key: value for key, value in request.headers if key.lower() != 'host'} - - # Get request body - data = request.get_data() - - # Log request info - logger.info(f"Forwarding {request.method} request to: {url}") - logger.info(f"Query params: {params}") - logger.info(f"Headers: {dict(headers)}") - - try: - # Forward the request - resp = requests.request( - method=request.method, - url=url, - params=params, - headers=headers, - data=data, - cookies=request.cookies, - allow_redirects=False, - stream=True - ) - - # Build response headers (exclude certain headers) - excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection'] - response_headers = [ - (name, value) for name, value in resp.raw.headers.items() - if name.lower() not in excluded_headers - ] - - # Create response - response = Response( - resp.content, - status=resp.status_code, - headers=response_headers - ) - - logger.info(f"Response status: {resp.status_code}") - - return response - - except requests.exceptions.RequestException as e: - logger.error(f"Error forwarding request: {str(e)}") - return Response(f"Proxy error: {str(e)}", status=502) - - -if __name__ == '__main__': - # Run the proxy server - """ - curl -X POST http://8.130.105.202:8010/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "model": "Qwen3-8B", - "messages": [ - { - "role": "system", - "content": "You are a helpful assistant." - }, - { - "role": "user", - "content": "你是谁？" - } - ] - }' - """ - PROXY_HOST = "0.0.0.0" # Listen on all interfaces - PROXY_PORT = 5010 # Proxy server port (must be different from target port!) - - logger.info(f"Starting HTTP Proxy Server on {PROXY_HOST}:{PROXY_PORT}") - logger.info(f"Forwarding requests to {TARGET_URL}") - - app.run( - host=PROXY_HOST, - port=PROXY_PORT, - debug=False, - threaded=True - ) - diff --git a/old/test/mcp_client_test.py b/old/test/mcp_client_test.py deleted file mode 100644 index e1c276f..0000000 --- a/old/test/mcp_client_test.py +++ /dev/null @@ -1,98 +0,0 @@ -import asyncio -import json - -from loguru import logger - -from flowllm.client.fastmcp_client import FastmcpClient - - -async def main(): - # New config-based interface matching McpClient - config = { - "type": "sse", - "url": "http://11.164.204.33:8001/sse", # or "http://0.0.0.0:8001/sse" - "headers": {}, - "timeout": 30.0 - } - - async with FastmcpClient("test_client", config, max_retries=3) as client: - - print("=" * 50) - print("Getting available MCP tools...") - try: - tool_calls = await client.list_tool_calls() - if tool_calls: - for tool_call in tool_calls: - print(tool_call.simple_input_dump()) - else: - print("No tools found or failed to retrieve tools") - except Exception as e: - print(f"Failed to get tools: {e}") - - query1 = "阿里巴巴怎么样？" - query2 = "分析下寒武纪" - query3 = "分析一下纽威股份" - - test_cases = [ - # ("ant_search", {"query": "阿里巴巴怎么样？", "entity": "阿里巴巴"}), - # ("ant_investment", {"entity": "阿里巴巴", "analysis_category": "股票"}), - # ("dashscope_search", {"query": "阿里巴巴怎么样？"}), - # ("get_a_stock_infos", {"query": query2}), - # ("get_a_stock_news", {"query": query2}), - # ("tavily_search", {"query": query2}), - - # ("mock_tool_flow", {"a": query2}), - # ("mock_async_tool_flow", {"a": query2}), - # ("tdx_wenda_quotes", {"question": "半导体行业PE中位数"}), - # ("tdx_market_quotes", {"code": "000001", "setcode": "1"}), - - # ("extract_entities_code", {"query": query2}), - # ("extract_entities_code", {"query": query3}), - - # ("bailian_web_search", {"query": "阿里巴巴怎么样？"}), - # ("bailian_web_parser", {"url": "https://xueqiu.com/1821992043/355703692"}), - # ("crawl_ths_company", {"query": query2, "code": "688256"}), - # ("crawl_ths_holder", {"query": query2, "code": "688256"}), - # ("crawl_ths_operate", {"query": query2, "code": "688256"}), - # ("crawl_ths_equity", {"query": query2, "code": "688256"}), - # ("crawl_ths_capital", {"query": query2, "code": "688256"}), - # ("crawl_ths_worth", {"query": query2, "code": "688256"}), - # ("crawl_ths_news", {"query": query2, "code": "688256"}), - # ("crawl_ths_concept", {"query": query2, "code": "688256"}), - # ("crawl_ths_position", {"query": query2, "code": "688256"}), - # ("crawl_ths_finance", {"query": query2, "code": "688256"}), - # ("crawl_ths_bonus", {"query": query2, "code": "688256"}), - # ("crawl_ths_event", {"query": query2, "code": "688256"}), - # ("crawl_ths_field", {"query": query2, "code": "688256"}), - - # ("crawl_ths_company", {"query": query3, "code": "603699"}), - # ("crawl_ths_bonus", {"query": query3, "code": "603699"}), - ("crawl_ths_operate", {"query": query3, "code": "603699"}), - # ("extract_query", {"query": "巴巴和小米怎么样"}), - # ("company_factor", {"name": "紫金矿业", "code": "601899"}), - ] - - for tool_name, arguments in test_cases: - print("=" * 50) - print(f"Testing tool: {tool_name} with arguments={arguments}") - try: - # New interface: call_tool(tool_name, arguments) returns CallToolResult - result = await client.call_tool(tool_name, arguments) - print(f"Tool result: {result}") - - if result.content: - print(f"Result content: {result.content[0].text}") - if result.structured_content: - print(f"Structured content: {json.dumps(result.structured_content, indent=2, ensure_ascii=False)}") - if result.is_error: - print(f"⚠ Tool returned error: {result.content[0].text if result.content else 'Unknown error'}") - else: - print("✓ Tool call successful") - - except Exception as e: - logger.exception(e) - print(f"✗ Tool call failed: {e}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/old/test/mcp_supply_test.py b/old/test/mcp_supply_test.py deleted file mode 100644 index a821896..0000000 --- a/old/test/mcp_supply_test.py +++ /dev/null @@ -1,78 +0,0 @@ -import asyncio -import json - -from flowllm.client.fastmcp_client import FastmcpClient - - -async def main(): - # async with FastmcpClient(transport="sse", host="0.0.0.0", port=8001) as client: - async with FastmcpClient(transport="sse", host="11.164.204.33", port=8001) as client: - print("=" * 50) - print("Getting available MCP tools...") - try: - tool_calls = await client.list_tool_calls() - if tool_calls: - for tool_call in tool_calls: - print(tool_call.model_dump_json()) - else: - print("No tools found or failed to retrieve tools") - except Exception as e: - print(f"Failed to get tools: {e}") - - query1 = "紫金怎么样？雪球" - - test_cases = [ - - # ("tavily_search", {"query": query1}), - # ("dashscope_search", {"query": query1}), - # ("brave_web_search", {"query": query1}), - # ("bailian_web_search", {"query": query1}), - # ("bocha_web_search", {"query": query1}), - ("bailian_web_parser", {"url": "https://basic.10jqka.com.cn/601899/field.html"}), - # ("bailian_web_parser", {"url": "https://basic.10jqka.com.cn/601899/company.html"}), - # - # ("extract_entities_code", {"query": "阿里和腾讯哪个更好？和茅台比呢？"}), - # ("akshare_market", {"code": "000001"}), - # ("akshare_calculate", {"query": "平安银行最近行情走势", "code": "000001"}), - # - # ("qtf_brief_mcp", {"symbol": "SZ000001"}), - # ("qtf_medium_mcp", {"symbol": "SZ000001"}), - # ("qtf_full_mcp", {"symbol": "SZ000001"}), - - # THS公司各类信息爬取测试 - # ("crawl_ths_company", {"query": "平安银行的公司情况如何？", "code": "000001"}), - # ("crawl_ths_holder", {"query": "平安银行的股东结构如何？", "code": "000001"}), - # ("crawl_ths_operate", {"query": "平安银行的主营业务是什么？", "code": "000001"}), - # ("crawl_ths_equity", {"query": "平安银行的股本结构如何？", "code": "000001"}), - # ("crawl_ths_capital", {"query": "平安银行有哪些资本运作？", "code": "000001"}), - # ("crawl_ths_worth", {"query": "平安银行的业绩预测如何？", "code": "000001"}), - # ("crawl_ths_news", {"query": "平安银行最近有什么新闻？", "code": "000001"}), - # ("crawl_ths_concept", {"query": "平安银行涉及哪些概念题材？", "code": "000001"}), - # ("crawl_ths_position", {"query": "平安银行的机构持仓情况如何？", "code": "000001"}), - # ("crawl_ths_finance", {"query": "平安银行的财务状况如何？", "code": "000001"}), - # ("crawl_ths_bonus", {"query": "平安银行的分红情况如何？", "code": "000001"}), - # ("crawl_ths_event", {"query": "平安银行最近有什么重大事件？", "code": "000001"}), - # ("crawl_ths_field", {"query": "平安银行在行业中的地位如何？", "code": "000001"}), - - # 测试其他股票代码 - # ("crawl_ths_company", {"query": "紫金矿业的公司情况如何？", "code": "601899"}), - # ("crawl_ths_finance", {"query": "紫金矿业的财务指标怎么样？", "code": "601899"}), - - ] - - for tool_name, arguments in test_cases: - print("=" * 100) - print(f"Testing tool: {tool_name}") - try: - result = await client.call_tool(tool_name, arguments) - if result.content: - print(f"Result content: {result.content[0].text}") - if hasattr(result, 'structured_content') and result.structured_content: - print(f"Structured content: {json.dumps(result.structured_content, indent=2, ensure_ascii=False)}") - print("✓ Tool call successful") - except Exception as e: - print(f"✗ Tool call failed: {e}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/old/test/mcp_test/__init__.py b/old/test/mcp_test/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/old/test/mcp_test/config.yaml b/old/test/mcp_test/config.yaml deleted file mode 100644 index b38ea2a..0000000 --- a/old/test/mcp_test/config.yaml +++ /dev/null @@ -1,35 +0,0 @@ -# config.yaml -functions: - user_service: - name: "create_user" - description: "创建用户" - parameters: - username: - type: "string" - description: "用户名" - required: true - email: - type: "string" - description: "邮箱" - required: true - age: - type: "integer" - description: "年龄" - required: false - - order_service: - name: "create_order" - description: "创建订单" - parameters: - product_id: - type: "string" - description: "产品ID" - required: true - quantity: - type: "integer" - description: "数量" - required: true - price: - type: "number" - description: "价格" - required: false diff --git a/old/test/mcp_test/mcp_service_generator.py b/old/test/mcp_test/mcp_service_generator.py deleted file mode 100644 index 2123c83..0000000 --- a/old/test/mcp_test/mcp_service_generator.py +++ /dev/null @@ -1,83 +0,0 @@ -from typing import Dict, Optional - -import yaml -from fastmcp import FastMCP -from pydantic import BaseModel, create_model - - -class MCPServiceGenerator: - def __init__(self, config_file: str): - self.config = self._load_config(config_file) - self.app = FastMCP("Dynamic Service") - - def _load_config(self, config_file: str) -> Dict: - """加载配置文件""" - with open(config_file, 'r', encoding='utf-8') as f: - return yaml.safe_load(f) - - def _create_pydantic_model(self, func_name: str, params: Dict) -> BaseModel: - """动态创建 Pydantic 模型""" - fields = {} - - for param_name, param_config in params.items(): - param_type = param_config.get('type', 'string') - required = param_config.get('required', True) - description = param_config.get('description', '') - - # 类型映射 - type_mapping = { - 'string': str, - 'integer': int, - 'number': float, - 'boolean': bool - } - - field_type = type_mapping.get(param_type, str) - - if not required: - field_type = Optional[field_type] - fields[param_name] = (field_type, None) - else: - fields[param_name] = (field_type, ...) - - return create_model(f"{func_name}Model", **fields) - - def register_function(self, func_name: str, actual_func): - """注册函数到 FastMCP""" - func_config = self.config['functions'][func_name] - - # 创建动态模型 - model_class = self._create_pydantic_model( - func_name, - func_config['parameters'] - ) - - # 注册到 FastMCP - @self.app.tool( - name=func_config['name'], - description=func_config['description'] - ) - def dynamic_handler(params: model_class) -> str: - # 将 Pydantic 模型转换为 kwargs - kwargs = params.dict(exclude_none=True) - - # 调用实际函数 - try: - result = actual_func(**kwargs) - return f"执行成功: {result}" - except Exception as e: - return f"执行失败: {str(e)}" - - return dynamic_handler - - def auto_register_functions(self, function_mapping: Dict[str, callable]): - """自动注册所有配置的函数""" - for func_name, actual_func in function_mapping.items(): - if func_name in self.config['functions']: - self.register_function(func_name, actual_func) - print(f"✅ 已注册函数: {func_name}") - - def run(self, host: str = "localhost", port: int = 8000): - """启动服务""" - print(f"🚀 FastMCP 服务启动在 {host}:{port}") - self.app.run(transport="sse", host=host, port=port) diff --git a/old/test/mcp_test/my_func.py b/old/test/mcp_test/my_func.py deleted file mode 100644 index a8a775c..0000000 --- a/old/test/mcp_test/my_func.py +++ /dev/null @@ -1,55 +0,0 @@ -# your_functions.py -from mcp_service_generator import MCPServiceGenerator - - -def create_user(**kwargs): - """创建用户的实际业务逻辑""" - username = kwargs.get('username') - email = kwargs.get('email') - age = kwargs.get('age', 18) - - # 模拟业务逻辑 - user_data = { - 'id': 12345, - 'username': username, - 'email': email, - 'age': age, - 'created_at': '2024-01-01' - } - - return user_data - - -def create_order(**kwargs): - """创建订单的实际业务逻辑""" - product_id = kwargs.get('product_id') - quantity = kwargs.get('quantity') - price = kwargs.get('price', 99.99) - - # 模拟业务逻辑 - order_data = { - 'order_id': 'ORD-001', - 'product_id': product_id, - 'quantity': quantity, - 'total_price': price * quantity - } - - return order_data - - -# main.py -if __name__ == "__main__": - # 创建服务生成器 - generator = MCPServiceGenerator("config.yaml") - - # 函数映射 - function_mapping = { - 'user_service': create_user, - 'order_service': create_order - } - - # 自动注册所有函数 - generator.auto_register_functions(function_mapping) - - # 启动服务 - generator.run(host="0.0.0.0", port=8000) diff --git a/old/test/mcp_test/test_server.py b/old/test/mcp_test/test_server.py deleted file mode 100644 index 4a804d3..0000000 --- a/old/test/mcp_test/test_server.py +++ /dev/null @@ -1,31 +0,0 @@ -from fastmcp import FastMCP -from pydantic import BaseModel, Field - - -# 定义参数模型 -class CalculateParams(BaseModel): - a: int = Field(description="第一个数字") - b: int = Field(description="第二个数字") - operation: str = Field(description="运算类型", enum=["add", "subtract", "multiply", "divide"]) - - -# 创建 FastMCP 实例 -mcp = FastMCP("Calculator") - - -# 注册工具函数 -@mcp.tool() -def calculate(params: CalculateParams) -> str: - """执行基本数学运算""" - if params.operation == "add": - result = params.a + params.b - elif params.operation == "subtract": - result = params.a - params.b - elif params.operation == "multiply": - result = params.a * params.b - elif params.operation == "divide": - result = params.a / params.b if params.b != 0 else "除数不能为零" - - return f"结果: {result}" - -mcp.run(transport="sse", host="0.0.0.0", port=8000) \ No newline at end of file diff --git a/old/test/test_cache.py b/old/test/test_cache.py deleted file mode 100644 index 1e611cf..0000000 --- a/old/test/test_cache.py +++ /dev/null @@ -1,456 +0,0 @@ -""" -Comprehensive tests for DataCache and CacheDataHandler classes -""" - -import json -import os -import sys -import tempfile -import time -from pathlib import Path -from typing import Dict, Any - -import pandas as pd -import pytest - -# Add project root directory to Python path -sys.path.append(os.path.join(os.path.dirname(__file__), '..')) - -from flowllm.storage.cache_handler.data_cache import DataCache -from flowllm.storage.cache_handler.cache_data_handler import CacheDataHandler, DataFrameHandler, DictHandler - - -class CustomHandler(CacheDataHandler): - """Custom handler for testing custom data types""" - - def save(self, data: str, file_path: Path, **kwargs) -> Dict[str, Any]: - """Save string data to text file""" - with open(file_path, 'w', encoding='utf-8') as f: - f.write(data) - - return { - 'length': len(data), - 'file_size': file_path.stat().st_size - } - - def load(self, file_path: Path, **kwargs) -> str: - """Load string data from text file""" - with open(file_path, 'r', encoding='utf-8') as f: - return f.read() - - def get_file_extension(self) -> str: - return ".txt" - - -class TestCacheDataHandlers: - """Test individual cache data handlers""" - - def test_dataframe_handler(self): - """Test DataFrameHandler""" - handler = DataFrameHandler() - - # Create test DataFrame - df = pd.DataFrame({ - 'id': [1, 2, 3], - 'name': ['Alice', 'Bob', 'Charlie'], - 'age': [25, 30, 35] - }) - - with tempfile.TemporaryDirectory() as temp_dir: - file_path = Path(temp_dir) / "test.csv" - - # Test save - metadata = handler.save(df, file_path) - assert file_path.exists() - assert metadata['row_count'] == 3 - assert metadata['column_count'] == 3 - assert 'file_size' in metadata - - # Test load - loaded_df = handler.load(file_path) - pd.testing.assert_frame_equal(df, loaded_df) - - # Test extension - assert handler.get_file_extension() == ".csv" - - def test_dict_handler(self): - """Test DictHandler""" - handler = DictHandler() - - # Create test dict - data = { - 'name': 'Test', - 'value': 42, - 'items': ['a', 'b', 'c'] - } - - with tempfile.TemporaryDirectory() as temp_dir: - file_path = Path(temp_dir) / "test.json" - - # Test save - metadata = handler.save(data, file_path) - assert file_path.exists() - assert metadata['key_count'] == 3 - assert 'file_size' in metadata - - # Test load - loaded_data = handler.load(file_path) - assert loaded_data == data - - # Test extension - assert handler.get_file_extension() == ".json" - - def test_custom_handler(self): - """Test custom handler""" - handler = CustomHandler() - - # Create test string data - data = "This is a test string with some content." - - with tempfile.TemporaryDirectory() as temp_dir: - file_path = Path(temp_dir) / "test.txt" - - # Test save - metadata = handler.save(data, file_path) - assert file_path.exists() - assert metadata['length'] == len(data) - assert 'file_size' in metadata - - # Test load - loaded_data = handler.load(file_path) - assert loaded_data == data - - # Test extension - assert handler.get_file_extension() == ".txt" - - -class TestDataCache: - """Test DataCache class""" - - def setup_method(self): - """Setup for each test method""" - self.temp_dir = tempfile.mkdtemp() - self.cache = DataCache(self.temp_dir) - - def teardown_method(self): - """Cleanup after each test method""" - # Clean up temp directory - import shutil - shutil.rmtree(self.temp_dir, ignore_errors=True) - - def test_basic_operations_dataframe(self): - """Test basic operations with DataFrame""" - # Create test DataFrame - df = pd.DataFrame({ - 'id': [1, 2, 3, 4, 5], - 'value': [10, 20, 30, 40, 50] - }) - - # Test save - success = self.cache.save("test_df", df, expire_hours=1) - assert success is True - - # Test exists - assert self.cache.exists("test_df") is True - assert self.cache.exists("nonexistent") is False - - # Test load - loaded_df = self.cache.load("test_df") - assert loaded_df is not None - pd.testing.assert_frame_equal(df, loaded_df) - - # Test info - info = self.cache.get_info("test_df") - assert info is not None - assert info['data_type'] == 'DataFrame' - assert info['row_count'] == 5 - assert info['column_count'] == 2 - assert 'created_time' in info - assert 'file_path' in info - - # Test delete - success = self.cache.delete("test_df") - assert success is True - assert self.cache.exists("test_df") is False - - def test_basic_operations_dict(self): - """Test basic operations with dict""" - # Create test dict - data = { - 'users': ['Alice', 'Bob', 'Charlie'], - 'count': 3, - 'active': True - } - - # Test save - success = self.cache.save("test_dict", data, expire_hours=2) - assert success is True - - # Test exists - assert self.cache.exists("test_dict") is True - - # Test load - loaded_data = self.cache.load("test_dict") - assert loaded_data is not None - assert loaded_data == data - - # Test info - info = self.cache.get_info("test_dict") - assert info is not None - assert info['data_type'] == 'dict' - assert info['key_count'] == 3 - - # Test delete - success = self.cache.delete("test_dict") - assert success is True - assert self.cache.exists("test_dict") is False - - def test_expiration_functionality(self): - """Test cache expiration""" - data = {'test': 'value'} - - # Save data that expires quickly (0.001 hours ≈ 3.6 seconds) - success = self.cache.save("temp_data", data, expire_hours=0.001) - assert success is True - - # Should exist immediately - assert self.cache.exists("temp_data") is True - - # Should be loadable immediately - loaded_data = self.cache.load("temp_data", auto_clean_expired=False) - assert loaded_data == data - - # Wait for expiration - time.sleep(4) - - # Should be expired now - assert self.cache.exists("temp_data") is False - - # Loading should return None and auto-clean - loaded_data = self.cache.load("temp_data", auto_clean_expired=True) - assert loaded_data is None - - def test_no_expiration(self): - """Test cache without expiration""" - data = {'persistent': 'data'} - - # Save without expiration - success = self.cache.save("persistent_data", data) - assert success is True - - # Should exist - assert self.cache.exists("persistent_data") is True - - # Info should show no expiration - info = self.cache.get_info("persistent_data") - assert info['expire_time'] is None - assert info['is_expired'] is False - - def test_custom_handler_registration(self): - """Test registering and using custom handlers""" - # Register custom handler for str type - custom_handler = CustomHandler() - self.cache.register_handler(str, custom_handler) - - # Test save/load with custom handler - test_string = "This is a test string for custom handler" - - success = self.cache.save("custom_data", test_string, expire_hours=1) - assert success is True - - loaded_string = self.cache.load("custom_data") - assert loaded_string == test_string - - # Check file extension is correct - info = self.cache.get_info("custom_data") - assert info['file_path'].endswith('.txt') - - def test_cache_management(self): - """Test cache management operations""" - # Create multiple caches - df1 = pd.DataFrame({'a': [1, 2, 3]}) - df2 = pd.DataFrame({'b': [4, 5, 6]}) - dict1 = {'key1': 'value1'} - dict2 = {'key2': 'value2'} - - # Save some data - self.cache.save("df1", df1, expire_hours=24) - self.cache.save("df2", df2, expire_hours=24) - self.cache.save("dict1", dict1, expire_hours=24) - self.cache.save("dict2", dict2, expire_hours=0.001) # This will expire - - # Wait for one to expire - time.sleep(4) - - # Test list_all - all_caches = self.cache.list_all(include_expired=True) - assert len(all_caches) == 4 - - active_caches = self.cache.list_all(include_expired=False) - assert len(active_caches) == 3 # One should be expired - - # Test cache stats - stats = self.cache.get_cache_stats() - assert stats['total_count'] == 4 - assert stats['active_count'] == 3 - assert stats['expired_count'] == 1 - assert stats['total_size_bytes'] > 0 - assert 'cache_dir' in stats - - # Test clean expired - cleaned_count = self.cache.clean_expired() - assert cleaned_count == 1 - - # After cleaning, should have 3 total - stats_after = self.cache.get_cache_stats() - assert stats_after['total_count'] == 3 - assert stats_after['expired_count'] == 0 - - def test_clear_all(self): - """Test clearing all caches""" - # Create some data - df = pd.DataFrame({'x': [1, 2, 3]}) - data = {'test': 'value'} - - self.cache.save("df_data", df) - self.cache.save("dict_data", data) - - # Verify they exist - assert self.cache.exists("df_data") is True - assert self.cache.exists("dict_data") is True - - # Clear all - success = self.cache.clear_all() - assert success is True - - # Verify they're gone - assert self.cache.exists("df_data") is False - assert self.cache.exists("dict_data") is False - - # Stats should show empty cache - stats = self.cache.get_cache_stats() - assert stats['total_count'] == 0 - - def test_handler_kwargs(self): - """Test passing kwargs to handlers""" - df = pd.DataFrame({ - 'name': ['Alice', 'Bob'], - 'age': [25, 30] - }) - - # Save with custom CSV parameters - success = self.cache.save( - "custom_csv", - df, - expire_hours=1, - index=True, # Save index - sep=';' # Use semicolon separator - ) - assert success is True - - # Load with corresponding parameters - loaded_df = self.cache.load( - "custom_csv", - sep=';', # Use semicolon separator - index_col=0 # First column as index - ) - - assert loaded_df is not None - # The loaded DataFrame will have the index from the CSV - assert len(loaded_df) == 2 - assert list(loaded_df.columns) == ['name', 'age'] - - def test_error_handling(self): - """Test error handling scenarios""" - # Test loading non-existent cache - result = self.cache.load("nonexistent") - assert result is None - - # Test getting info for non-existent cache - info = self.cache.get_info("nonexistent") - assert info is None - - # Test deleting non-existent cache (should not raise error) - success = self.cache.delete("nonexistent") - assert success is True # Should return True even if nothing to delete - - # Test unsupported data type - class UnsupportedType: - pass - - unsupported_data = UnsupportedType() - success = self.cache.save("unsupported", unsupported_data) - assert success is False # Should fail gracefully - - -def create_sample_dataframe(rows: int = 10) -> pd.DataFrame: - """Create a sample DataFrame for testing""" - import numpy as np - np.random.seed(42) - - data = { - 'id': range(1, rows + 1), - 'name': [f'User_{i}' for i in range(1, rows + 1)], - 'value': np.random.randint(1, 100, rows), - 'category': np.random.choice(['A', 'B', 'C'], rows) - } - return pd.DataFrame(data) - - -def test_integration_example(): - """Integration test showing typical usage""" - print("=== Integration Test Example ===") - - with tempfile.TemporaryDirectory() as temp_dir: - cache = DataCache(temp_dir) - - # Test with DataFrame - df = create_sample_dataframe(20) - print(f"Created DataFrame with {len(df)} rows") - - # Save DataFrame - success = cache.save("sample_data", df, expire_hours=1) - print(f"Save DataFrame: {'Success' if success else 'Failed'}") - - # Test with dict - config_data = { - 'database_url': 'postgresql://localhost:5432/mydb', - 'cache_timeout': 3600, - 'features': ['feature1', 'feature2', 'feature3'], - 'settings': { - 'debug': True, - 'log_level': 'INFO' - } - } - - # Save dict - success = cache.save("config", config_data, expire_hours=2) - print(f"Save dict: {'Success' if success else 'Failed'}") - - # Load and verify - loaded_df = cache.load("sample_data") - loaded_config = cache.load("config") - - print(f"Load DataFrame: {'Success' if loaded_df is not None else 'Failed'}") - print(f"Load dict: {'Success' if loaded_config is not None else 'Failed'}") - - # Show cache stats - stats = cache.get_cache_stats() - print(f"\nCache Statistics:") - print(f" Total items: {stats['total_count']}") - print(f" Total size: {stats['total_size_mb']} MB") - print(f" Cache directory: {stats['cache_dir']}") - - # List all caches - all_caches = cache.list_all() - print(f"\nCached items:") - for key, info in all_caches.items(): - print(f" {key}: {info['data_type']} (created: {info['created_time'][:19]})") - - -if __name__ == "__main__": - # Run integration test - test_integration_example() - - print("\n" + "=" * 50) - print("Run 'pytest test_cache.py -v' for full test suite") diff --git a/old/test/test_config.py b/old/test/test_config.py deleted file mode 100644 index 29ffc37..0000000 --- a/old/test/test_config.py +++ /dev/null @@ -1,66 +0,0 @@ -from pydantic import BaseModel - -from flowllm.config.pydantic_config_parser import PydanticConfigParser - - -class TestConfig(BaseModel): - name: str = "test" - port: int = 8080 - debug: bool = False - - -class ServerConfig(BaseModel): - host: str = "localhost" - port: int = 3000 - ssl: bool = False - - -def test_generic_config_parser(): - """Test that the generic config parser works correctly with type hints""" - - # Test with TestConfig - parser1 = PydanticConfigParser[TestConfig](TestConfig) - config1 = parser1.parse_args([]) - - assert isinstance(config1, TestConfig) - assert config1.name == "test" - assert config1.port == 8080 - assert config1.debug is False - - # Test with ServerConfig - parser2 = PydanticConfigParser[ServerConfig](ServerConfig) - config2 = parser2.parse_args([]) - - assert isinstance(config2, ServerConfig) - assert config2.host == "localhost" - assert config2.port == 3000 - assert config2.ssl is False - - # Test update_config returns correct type - updated_config1 = parser1.update_config(name="updated", port=9090) - assert isinstance(updated_config1, TestConfig) - assert updated_config1.name == "updated" - assert updated_config1.port == 9090 - - updated_config2 = parser2.update_config(host="example.com", ssl=True) - assert isinstance(updated_config2, ServerConfig) - assert updated_config2.host == "example.com" - assert updated_config2.ssl is True - - -def test_backward_compatibility(): - """Test that the parser still works without explicit generic type""" - - # Without explicit generic type (backward compatibility) - parser = PydanticConfigParser(TestConfig) - config = parser.parse_args([]) - - assert isinstance(config, TestConfig) - assert config.name == "test" - assert config.port == 8080 - - -if __name__ == "__main__": - test_generic_config_parser() - test_backward_compatibility() - print("All tests passed!") diff --git a/old/test/test_craw4ai.py b/old/test/test_craw4ai.py deleted file mode 100644 index 8ad5b37..0000000 --- a/old/test/test_craw4ai.py +++ /dev/null @@ -1,34 +0,0 @@ -import asyncio - -from crawl4ai import * - - -async def main(): - browser_config = BrowserConfig( - headless=True, - java_script_enabled=True, - user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36", - viewport={"width": 1280, "height": 800}, - verbose=True) - - crawler_config = CrawlerRunConfig(cache_mode=CacheMode.BYPASS, - page_timeout=9000, - verbose=True) - - async with AsyncWebCrawler(config=browser_config) as crawler: - result = await crawler.arun( - # url="https://stockpage.10jqka.com.cn/000807/operate/", - # url="https://basic.10jqka.com.cn/601899/operate.html#stockpage", - # url="https://basic.10jqka.com.cn/601899/worth.html#stockpage", - url="https://basic.10jqka.com.cn/601899/finance.html#stockpage", - config=crawler_config, - js_code="window.scrollTo(0, document.body.scrollHeight);", - wait_for="document.querySelector('.loaded')" - ) - print(result.markdown) - - with open("output.md", "w", encoding="utf-8") as f: - f.write(result.markdown) - -if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file diff --git a/old/test/test_dashscope_llm.py b/old/test/test_dashscope_llm.py deleted file mode 100644 index 81f59f5..0000000 --- a/old/test/test_dashscope_llm.py +++ /dev/null @@ -1,190 +0,0 @@ -""" -Test module for DashscopeLLM implementation. - -This module contains unit tests for the DashscopeLLM class, -testing initialization, message conversion, and basic functionality. -""" - -import os -import pytest -from unittest.mock import patch, MagicMock - -from flowllm.llm.dashscope_llm import DashscopeLLM -from flowllm.schema.message import Message -from flowllm.schema.tool_call import ToolCall -from flowllm.enumeration.role import Role - - -class TestDashscopeLLM: - """Test cases for DashscopeLLM class""" - - def test_initialization(self): - """Test DashscopeLLM initialization with default parameters""" - llm = DashscopeLLM(model_name="qwen-plus-2025-04-28") - - assert llm.model_name == "qwen-plus-2025-04-28" - assert llm.temperature == 0.0000001 - assert llm.enable_search is False - assert llm.enable_thinking is False - assert llm.max_retries == 5 - - def test_initialization_with_custom_params(self): - """Test DashscopeLLM initialization with custom parameters""" - llm = DashscopeLLM( - model_name="qwq-plus-latest", - temperature=0.7, - enable_search=True, - enable_thinking=True, - max_retries=3 - ) - - assert llm.model_name == "qwq-plus-latest" - assert llm.temperature == 0.7 - assert llm.enable_search is True - assert llm.enable_thinking is True - assert llm.max_retries == 3 - - def test_convert_messages_to_dashscope(self): - """Test message conversion to Dashscope format""" - llm = DashscopeLLM(model_name="qwen-plus-2025-04-28") - - messages = [ - Message(role=Role.USER, content="Hello"), - Message(role=Role.ASSISTANT, content="Hi there!"), - Message(role=Role.USER, reasoning_content="Think about this", content="") - ] - - dashscope_messages = llm._convert_messages_to_dashscope(messages) - - assert len(dashscope_messages) == 3 - assert dashscope_messages[0].role == "user" - assert dashscope_messages[0].content == "Hello" - assert dashscope_messages[1].role == "assistant" - assert dashscope_messages[1].content == "Hi there!" - # Should use reasoning_content when content is empty - assert dashscope_messages[2].content == "Think about this" - - def test_convert_tools_to_dashscope(self): - """Test tool conversion to Dashscope format""" - llm = DashscopeLLM(model_name="qwen-plus-2025-04-28") - - tools = [ - ToolCall( - name="search_web", - description="Search the web for information", - parameters={ - "type": "object", - "properties": { - "query": {"type": "string", "description": "Search query"} - }, - "required": ["query"] - } - ) - ] - - dashscope_tools = llm._convert_tools_to_dashscope(tools) - - assert len(dashscope_tools) == 1 - assert dashscope_tools[0]["type"] == "function" - assert dashscope_tools[0]["function"]["name"] == "search_web" - assert dashscope_tools[0]["function"]["description"] == "Search the web for information" - assert "query" in dashscope_tools[0]["function"]["parameters"]["properties"] - - def test_convert_tools_to_dashscope_empty(self): - """Test tool conversion with empty tool list""" - llm = DashscopeLLM(model_name="qwen-plus-2025-04-28") - - dashscope_tools = llm._convert_tools_to_dashscope([]) - assert dashscope_tools is None - - dashscope_tools = llm._convert_tools_to_dashscope(None) - assert dashscope_tools is None - - @patch.dict(os.environ, {'DASHSCOPE_API_KEY': 'test_key'}) - def test_api_key_from_env(self): - """Test API key loading from environment variable""" - llm = DashscopeLLM(model_name="qwen-plus-2025-04-28") - assert llm.api_key == "test_key" - - def test_model_name_validation(self): - """Test model name validation with supported models""" - supported_models = [ - "qwen-plus-2025-04-28", - "qwq-plus-latest", - "qwen-max-2025-01-25", - "qwen-turbo", - "qwen-plus", - "qwen-max" - ] - - for model in supported_models: - llm = DashscopeLLM(model_name=model) - assert llm.model_name == model - - def test_search_options_default(self): - """Test default search options configuration""" - llm = DashscopeLLM(model_name="qwen-plus-2025-04-28") - - expected_options = { - "forced_search": False, - "enable_source": True, - "enable_citation": False, - "search_strategy": "pro" - } - - assert llm.search_options == expected_options - - def test_search_options_custom(self): - """Test custom search options configuration""" - custom_options = { - "forced_search": True, - "enable_source": False, - "enable_citation": True, - "search_strategy": "basic" - } - - llm = DashscopeLLM( - model_name="qwen-plus-2025-04-28", - search_options=custom_options - ) - - assert llm.search_options == custom_options - - -if __name__ == "__main__": - # Run basic tests without pytest - test_instance = TestDashscopeLLM() - - print("Running DashscopeLLM tests...") - - try: - test_instance.test_initialization() - print("✓ Initialization test passed") - - test_instance.test_initialization_with_custom_params() - print("✓ Custom parameters test passed") - - test_instance.test_convert_messages_to_dashscope() - print("✓ Message conversion test passed") - - test_instance.test_convert_tools_to_dashscope() - print("✓ Tool conversion test passed") - - test_instance.test_convert_tools_to_dashscope_empty() - print("✓ Empty tool conversion test passed") - - test_instance.test_model_name_validation() - print("✓ Model name validation test passed") - - test_instance.test_search_options_default() - print("✓ Default search options test passed") - - test_instance.test_search_options_custom() - print("✓ Custom search options test passed") - - print("\nAll tests passed! 🎉") - - except Exception as e: - print(f"❌ Test failed: {e}") - import traceback - traceback.print_exc() diff --git a/old/test/test_dataframe_cache.py b/old/test/test_dataframe_cache.py deleted file mode 100644 index 49a94ac..0000000 --- a/old/test/test_dataframe_cache.py +++ /dev/null @@ -1,205 +0,0 @@ -""" -DataFrame cache utility tests and usage examples -""" - -import os -import sys -import time - -import numpy as np -import pandas as pd - -# Add project root directory to Python path -sys.path.append(os.path.join(os.path.dirname(__file__), '..')) - -from flowllm.storage.cache_handler.dataframe_cache import DataFrameCache, save_dataframe, load_dataframe - - -def create_sample_dataframe(rows: int = 100) -> pd.DataFrame: - """Create sample DataFrame""" - np.random.seed(42) - data = { - 'id': range(1, rows + 1), - 'name': [f'User_{i}' for i in range(1, rows + 1)], - 'age': np.random.randint(18, 80, rows), - 'salary': np.random.normal(50000, 15000, rows).round(2), - 'department': np.random.choice(['Engineering', 'Sales', 'Marketing', 'HR'], rows), - 'join_date': pd.date_range('2020-01-01', periods=rows, freq='D') - } - return pd.DataFrame(data) - - -def test_basic_operations(): - """Test basic operations""" - print("=== Test Basic Operations ===") - - # Create cache instance - cache = DataFrameCache("test_cache") - - # Create sample data - df = create_sample_dataframe(50) - print(f"Created DataFrame with {len(df)} rows of data") - - # Save data (expires after 1 hour) - success = cache.save("employee_data", df, expire_hours=1) - print(f"Save result: {'Success' if success else 'Failed'}") - - # Check if exists - exists = cache.exists("employee_data") - print(f"Data exists: {exists}") - - # Load data - loaded_df = cache.load("employee_data") - if loaded_df is not None: - print(f"Successfully loaded data with {len(loaded_df)} rows") - print("First 5 rows:") - print(loaded_df.head()) - else: - print("Failed to load") - - # Get cache info - info = cache.get_info("employee_data") - if info: - print(f"\nCache info:") - print(f" Created time: {info['created_time']}") - print(f" File size: {info['file_size']} bytes") - print(f" Row count: {info['row_count']}") - print(f" Column count: {info['column_count']}") - - -def test_expiry_functionality(): - """Test expiry functionality""" - print("\n=== Test Expiry Functionality ===") - - cache = DataFrameCache("test_cache") - df = create_sample_dataframe(20) - - # Save data that expires quickly (0.001 hours = 3.6 seconds) - cache.save("temp_data", df, expire_hours=0.001) - print("Saved data that expires after 3.6 seconds") - - # Check immediately - exists_before = cache.exists("temp_data") - print(f"Immediate existence check: {exists_before}") - - # Wait for expiration - print("Waiting 5 seconds...") - time.sleep(5) - - # Check again - exists_after = cache.exists("temp_data") - print(f"Existence check after 5 seconds: {exists_after}") - - # Try to load expired data - expired_df = cache.load("temp_data") - print(f"Try to load expired data: {'Success' if expired_df is not None else 'Failed (automatically cleaned)'}") - - -def test_cache_management(): - """Test cache management functionality""" - print("\n=== Test Cache Management Functionality ===") - - cache = DataFrameCache("test_cache") - - # Create multiple caches - for i in range(3): - df = create_sample_dataframe(10 + i * 10) - cache.save(f"dataset_{i}", df, expire_hours=24) - - # Create an expired cache - df_expired = create_sample_dataframe(5) - cache.save("expired_dataset", df_expired, expire_hours=0.001) - time.sleep(4) # Wait for expiration - - # List all caches - all_caches = cache.list_all(include_expired=True) - print(f"Total cache count: {len(all_caches)}") - - active_caches = cache.list_all(include_expired=False) - print(f"Active cache count: {len(active_caches)}") - - # Get statistics - stats = cache.get_cache_stats() - print(f"\nCache statistics:") - print(f" Total count: {stats['total_count']}") - print(f" Active count: {stats['active_count']}") - print(f" Expired count: {stats['expired_count']}") - print(f" Total size: {stats['total_size_mb']} MB") - - # Clean expired caches - cleaned_count = cache.clean_expired() - print(f"Cleaned {cleaned_count} expired caches") - - # Clear all caches - cache.clear_all() - print("Cleared all caches") - - -def test_convenience_functions(): - """Test convenience functions""" - print("\n=== Test Convenience Functions ===") - - df = create_sample_dataframe(30) - - # Use convenience function to save - success = save_dataframe("convenience_test", df, expire_hours=2) - print(f"Convenience save: {'Success' if success else 'Failed'}") - - # Use convenience function to load - loaded_df = load_dataframe("convenience_test") - print(f"Convenience load: {'Success' if loaded_df is not None else 'Failed'}") - - if loaded_df is not None: - print(f"Data shape: {loaded_df.shape}") - - -def test_custom_csv_params(): - """Test custom CSV parameters""" - print("\n=== Test Custom CSV Parameters ===") - - cache = DataFrameCache("test_cache") - df = create_sample_dataframe(20) - - # Save with custom parameters - success = cache.save( - "custom_csv", - df, - expire_hours=1, - index=True, # Save index - sep=';' # Use semicolon separator - ) - print(f"Save with custom CSV parameters: {'Success' if success else 'Failed'}") - - # Load with corresponding parameters - loaded_df = cache.load( - "custom_csv", - sep=';', # Use semicolon separator - index_col=0 # First column as index - ) - - if loaded_df is not None: - print(f"Successfully loaded custom format data, shape: {loaded_df.shape}") - print("Index name:", loaded_df.index.name) - - # Cleanup - cache.delete("custom_csv") - - -if __name__ == "__main__": - print("DataFrame Cache Utility Test") - print("=" * 50) - - try: - test_basic_operations() - test_expiry_functionality() - test_cache_management() - test_convenience_functions() - test_custom_csv_params() - - print("\n" + "=" * 50) - print("All tests completed!") - - except Exception as e: - print(f"Error occurred during testing: {e}") - import traceback - traceback.print_exc() diff --git a/old/test/test_ifund.py b/old/test/test_ifund.py deleted file mode 100644 index d6dd04d..0000000 --- a/old/test/test_ifund.py +++ /dev/null @@ -1,9 +0,0 @@ -import requests -import json - -requestURL = "https://quantapi.51ifind.com/api/v1/basic_data_service" -requestHeaders = {"Content-Type":"application/json","access_token":""} -formData = {"codes":"601899.SH","indipara":[{"indicator":"ths_main_businuess_stock","indiparams":["2025-09-17"]}]} - -thsResponse = requests.post(url=requestURL,json=formData,headers=requestHeaders) -print(json.dumps(thsResponse.json(), ensure_ascii=False, indent=2)) \ No newline at end of file diff --git a/old/test/test_playwright.py b/old/test/test_playwright.py deleted file mode 100644 index 0c4957d..0000000 --- a/old/test/test_playwright.py +++ /dev/null @@ -1,176 +0,0 @@ -import asyncio -from playwright.async_api import async_playwright -import time -import os -from datetime import datetime - - -async def crawl_with_playwright(): - """使用Playwright爬取股票页面""" - - async with async_playwright() as p: - # 启动浏览器 - browser = await p.chromium.launch( - headless=False, # 显示浏览器窗口，方便调试 - args=[ - '--no-sandbox', - '--disable-blink-features=AutomationControlled', - '--disable-web-security' - ] - ) - - # 创建新页面 - context = await browser.new_context( - user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36", - viewport={"width": 1280, "height": 800} - ) - - page = await context.new_page() - - try: - print("开始访问页面...") - - # 访问页面 - await page.goto( - "https://stockpage.10jqka.com.cn/000807/operate/", - wait_until="domcontentloaded", # 等待DOM加载完成 - timeout=15000 # 15秒超时 - ) - - print("页面加载完成，等待内容...") - - # 等待页面稳定 - await page.wait_for_timeout(3000) - - # 滚动到页面底部，触发懒加载 - print("滚动页面加载更多内容...") - await page.evaluate("window.scrollTo(0, document.body.scrollHeight)") - - # 等待滚动后的内容加载 - await page.wait_for_timeout(2000) - - # 再次滚动确保所有内容加载 - await page.evaluate("window.scrollTo(0, document.body.scrollHeight)") - await page.wait_for_timeout(2000) - - # 尝试等待特定元素（如果知道的话） - try: - # 等待可能的加载指示器消失或内容出现 - await page.wait_for_function( - "document.readyState === 'complete'", - timeout=5000 - ) - except: - print("等待页面完全加载超时，继续...") - - # 获取页面内容 - print("获取页面内容...") - - # 获取页面标题 - title = await page.title() - print(f"页面标题: {title}") - - # 获取页面文本内容 - text_content = await page.evaluate(""" - () => { - // 移除脚本和样式标签 - const scripts = document.querySelectorAll('script, style'); - scripts.forEach(el => el.remove()); - - // 获取body的文本内容 - return document.body.innerText || document.body.textContent || ''; - } - """) - - # 获取HTML内容（备用） - html_content = await page.content() - - print(f"文本内容长度: {len(text_content)} 字符") - print(f"HTML内容长度: {len(html_content)} 字符") - - # 保存内容到文件 - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - - # 创建输出目录 - output_dir = "/Users/yuli/workspace/flowllm/test/crawl_output" - os.makedirs(output_dir, exist_ok=True) - - # 保存文本内容 - text_filename = f"stock_content_{timestamp}.txt" - text_filepath = os.path.join(output_dir, text_filename) - - with open(text_filepath, 'w', encoding='utf-8') as f: - f.write(f"页面标题: {title}\n") - f.write(f"爬取时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n") - f.write(f"URL: https://stockpage.10jqka.com.cn/000807/operate/\n") - f.write("=" * 50 + "\n\n") - f.write(text_content) - - # 保存HTML内容 - html_filename = f"stock_html_{timestamp}.html" - html_filepath = os.path.join(output_dir, html_filename) - - with open(html_filepath, 'w', encoding='utf-8') as f: - f.write(html_content) - - print(f"\n✅ 内容已保存到文件:") - print(f"📄 文本内容: {text_filepath}") - print(f"🌐 HTML内容: {html_filepath}") - - if len(text_content) > 100: - print("\n--- 页面文本内容预览 ---") - print(text_content[:1000] + "..." if len(text_content) > 1000 else text_content) - else: - print("文本内容较少，显示HTML片段:") - print(html_content[:1000] + "..." if len(html_content) > 1000 else html_content) - - return { - 'success': True, - 'title': title, - 'text': text_content, - 'html': html_content, - 'text_length': len(text_content), - 'html_length': len(html_content), - 'text_file': text_filepath, - 'html_file': html_filepath - } - - except Exception as e: - print(f"爬取过程中发生错误: {str(e)}") - return { - 'success': False, - 'error': str(e) - } - - finally: - # 关闭浏览器 - await browser.close() - - -async def main(): - """主函数""" - print("使用Playwright爬取网页...") - print("=" * 50) - - start_time = time.time() - result = await crawl_with_playwright() - end_time = time.time() - - print("\n" + "=" * 50) - print(f"爬取完成，耗时: {end_time - start_time:.2f} 秒") - - if result['success']: - print("✅ 爬取成功!") - print(f"页面标题: {result['title']}") - print(f"文本长度: {result['text_length']} 字符") - print(f"HTML长度: {result['html_length']} 字符") - print(f"📁 文件保存位置:") - print(f" 文本文件: {result['text_file']}") - print(f" HTML文件: {result['html_file']}") - else: - print("❌ 爬取失败!") - print(f"错误信息: {result['error']}") - - -if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file diff --git a/old/test/test_simple_flow.py b/old/test/test_simple_flow.py deleted file mode 100644 index acc6da6..0000000 --- a/old/test/test_simple_flow.py +++ /dev/null @@ -1,148 +0,0 @@ -#!/usr/bin/env python3 -""" -Test script for SimpleFlow implementation -Demonstrates parsing and execution of flow expressions -""" - -import os -import sys -import time -from concurrent.futures import ThreadPoolExecutor - -# Add the project root to the Python path -sys.path.insert(0, os.path.abspath('.')) - -from flowllm.engine.simple_flow_engine import SimpleFlowEngine -from flowllm.context.flow_context import FlowContext -from flowllm.context.service_context import C -from flowllm.op.base_op import BaseOp -from flowllm.schema.service_config import ServiceConfig, OpConfig - - -class TestOp(BaseOp): - """Test operation for demonstration purposes""" - - def execute(self): - # Simulate some processing time - time.sleep(0.1) - result = f"{self.name}_result" - print(f" 🔧 Executing {self.name} -> {result}") - return result - - -def create_test_context(): - """Create test contexts with sample operations""" - - # Setup thread pool for parallel execution - C["thread_pool"] = ThreadPoolExecutor(max_workers=4) - - # Register TestOp in the service context - C.op_registry.register("test_op")(TestOp) - - # Create service config with op configurations - service_config = ServiceConfig( - op={ - "op1": OpConfig(backend="test_op"), - "op2": OpConfig(backend="test_op"), - "op3": OpConfig(backend="test_op"), - "op4": OpConfig(backend="test_op") - } - ) - - # Create flow context - flow_context = FlowContext(service_config=service_config) - - return flow_context - - -def test_simple_expression(): - """Test simple sequential expression""" - print("\n" + "=" * 60) - print("TEST 1: Simple sequential expression 'op1 >> op2'") - print("=" * 60) - - flow_context = create_test_context() - - flow = SimpleFlowEngine( - flow_name="simple_sequential", - flow_content="op1 >> op2", - flow_context=flow_context) - - result = flow() - print(f"Final result: {result}") - - -def test_parallel_expression(): - """Test parallel expression""" - print("\n" + "=" * 60) - print("TEST 2: Parallel expression 'op1 | op2'") - print("=" * 60) - - flow_context = create_test_context() - - flow = SimpleFlowEngine( - flow_name="simple_parallel", - flow_content="op1 | op2", - flow_context=flow_context) - - result = flow() - print(f"Final result: {result}") - - -def test_mixed_expression(): - """Test mixed expression with parentheses""" - print("\n" + "=" * 60) - print("TEST 3: Mixed expression 'op1 >> (op2 | op3) >> op4'") - print("=" * 60) - - flow_context = create_test_context() - - flow = SimpleFlowEngine( - flow_name="mixed_flow", - flow_content="op1 >> (op2 | op3) >> op4", - flow_context=flow_context) - - result = flow() - print(f"Final result: {result}") - - -def test_complex_expression(): - """Test complex nested expression""" - print("\n" + "=" * 60) - print("TEST 4: Complex expression 'op1 >> (op1 | (op2 >> op3)) >> op4'") - print("=" * 60) - - flow_context = create_test_context() - - flow = SimpleFlowEngine( - flow_name="complex_flow", - flow_content="op1 >> (op1 | (op2 >> op3)) >> op4", - flow_context=flow_context) - - result = flow() - print(f"Final result: {result}") - - -def main(): - """Run all tests""" - print("🧪 Testing SimpleFlow Implementation") - print("This demonstrates parsing and execution of flow expressions") - - try: - test_simple_expression() - test_parallel_expression() - test_mixed_expression() - test_complex_expression() - - print("\n" + "=" * 60) - print("✅ All tests completed successfully!") - print("=" * 60) - - except Exception as e: - print(f"\n❌ Test failed with error: {e}") - import traceback - traceback.print_exc() - - -if __name__ == "__main__": - main() diff --git a/old/test/test_streaming.py b/old/test/test_streaming.py deleted file mode 100644 index 721bf4c..0000000 --- a/old/test/test_streaming.py +++ /dev/null @@ -1,49 +0,0 @@ -import time -import requests -import json -import sseclient - -from flowllm.context.service_context import C -from flowllm.flow.gallery.streaming_mock_flow import StreamingMockFlow -from flowllm.service.base_service import BaseService -from flowllm.service.http_service import HttpService - - -def test_streaming_response(): - """Test streaming response from HTTP service""" - # Register streaming mock tool flow - mock_flow = StreamingMockFlow() - C.register_flow("streaming_mock")(mock_flow) - - # Start HTTP service in a separate thread - import threading - service = BaseService.get_service("config=default") - thread = threading.Thread(target=service, daemon=True) - thread.start() - - # Wait for service to start - time.sleep(2) - - # Test regular endpoint - response = requests.post("http://localhost:8000/streaming_mock", json={"input_data": "test data"}) - print(f"Regular response: {response.json()}") - - # Test streaming endpoint - url = "http://localhost:8000/streaming_mock/stream" - headers = {"Accept": "text/event-stream"} - response = requests.post(url, json={"input_data": "test data"}, headers=headers, stream=True) - - client = sseclient.SSEClient(response) - for event in client.events(): - if event.data: - try: - data = json.loads(event.data) - print(f"Streaming chunk: {data}") - if data.get("is_complete"): - break - except json.JSONDecodeError: - print(f"Non-JSON data: {event.data}") - - -if __name__ == "__main__": - test_streaming_response() diff --git a/old/test/test_streaming_example.py b/old/test/test_streaming_example.py deleted file mode 100644 index 08e1725..0000000 --- a/old/test/test_streaming_example.py +++ /dev/null @@ -1,71 +0,0 @@ -import time -import requests -import json -import sseclient -import threading - -from flowllm.context.service_context import C -from flowllm.service.base_service import BaseService -from flowllm.flow.gallery.streaming_example_flow import StreamingExampleFlow - - -def test_streaming_example(): - """测试流式输出示例""" - # 注册流式输出示例流程 - C.register_flow("streaming_example")(StreamingExampleFlow()) - - # 在单独的线程中启动HTTP服务 - service = BaseService.get_service("config=default") - thread = threading.Thread(target=service, daemon=True) - thread.start() - - # 等待服务启动 - time.sleep(2) - - # 测试常规端点 - print("\n测试常规端点:") - response = requests.post( - "http://localhost:8000/streaming_example_flow", - json={"text": "这是一个流式输出的测试示例，看看效果如何。"} - ) - print(f"常规响应: {response.json()}") - - # 测试流式端点 - print("\n测试流式端点:") - url = "http://localhost:8000/streaming_example_flow/stream" - headers = {"Accept": "text/event-stream"} - response = requests.post( - url, - json={"text": "这是一个流式输出的测试示例，看看效果如何。", "delay": 0.3}, - headers=headers, - stream=True - ) - - client = sseclient.SSEClient(response) - for event in client.events(): - if event.data: - if event.data == "[DONE]": - print("\n流式输出完成") - break - - try: - data = json.loads(event.data) - chunk_type = data.get("chunk_type", "") - chunk = data.get("chunk", "") - - if chunk_type == "think": - print(f"\n[思考] {chunk}") - elif chunk_type == "answer": - print(chunk, end="", flush=True) - else: - print(f"\n[{chunk_type}] {chunk}") - - if data.get("done"): - print("\n流式输出完成") - break - except json.JSONDecodeError: - print(f"\n非JSON数据: {event.data}") - - -if __name__ == "__main__": - test_streaming_example() diff --git a/old/test/test_vl.py b/old/test/test_vl.py deleted file mode 100644 index 6067cf0..0000000 --- a/old/test/test_vl.py +++ /dev/null @@ -1,38 +0,0 @@ -import os -from openai import OpenAI -import base64 -from flowllm.utils.common_utils import load_env -load_env() - -# 编码函数： 将本地文件转换为 Base64 编码的字符串 -def encode_image(image_path): - with open(image_path, "rb") as image_file: - return base64.b64encode(image_file.read()).decode("utf-8") - -# 将xxxx/eagle.png替换为你本地图像的绝对路径 -# base64_image = encode_image("/Users/yuli/Documents/20250924133208.jpg") -base64_image = encode_image("/Users/yuli/Documents/20250924133848.jpg") - -client = OpenAI( - # 新加坡和北京地域的API Key不同。获取API Key：https://help.aliyun.com/zh/model-studio/get-api-key - api_key=os.getenv("FLOW_LLM_API_KEY"), - # 以下为北京地域url，若使用新加坡地域的模型，需将url替换为：https://dashscopep-intl.aliyuncs.com/api/v1/services/aigc/text-generation/generation - base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", -) - -messages = [ - {"role":"system","content":[{"type": "text", "text": "You are a helpful assistant."}]}, - {"role": "user","content": [ - # JPEG图像： f"data:image/jpeg;base64,{base64_image}" - {"type": "image_url","image_url": {"url": f"data:image/jpeg;base64,{base64_image}"},}, - {"type": "text", "text": "这些图描绘了什么内容，简单介绍下，同时推理用户正在干什么？一句话总结"}, - ], - } - ] - -completion = client.chat.completions.create( - model="qwen3-vl-235b-a22b-instruct", # 此处以qwen-vl-max-latest为例，可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/models - messages=messages, -) - -print(completion.choices[0].message.content) \ No newline at end of file diff --git a/old/test/test_vllm.py b/old/test/test_vllm.py deleted file mode 100644 index 5c93fc8..0000000 --- a/old/test/test_vllm.py +++ /dev/null @@ -1,40 +0,0 @@ -from openai import OpenAI - -""" -export CUDA_VISIBLE_DEVICES=1,2 -python -m vllm.entrypoints.openai.api_server \ - --model /mnt/data_cpfs/xielipeng.xlp/models/Qwen3-8B \ - --served-model-name Qwen3-8B \ - --host 0.0.0.0 \ - --port 8010 \ - --tensor-parallel-size 2 \ - --max-model-len 40960 - -curl --location 'http://8.130.105.202:8010/v1/chat/completions' \ - --header 'Content-Type: application/json' \ - --data '{ - "model" : "Qwen3-8B", - "messages" : [ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "what is your name?"} - ] - }' -""" - -try: - client = OpenAI( - base_url="http://8.130.105.202:8010/v1", - api_key="", - ) - - completion = client.chat.completions.create( - model="Qwen3-8B", # 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models - messages=[ - {'role': 'system', 'content': 'You are a helpful assistant.'}, - {'role': 'user', 'content': '你是谁？'} - ] - ) - print(completion.choices[0].message.content) -except Exception as e: - print(f"错误信息：{e}") - print("请参考文档：https://help.aliyun.com/zh/model-studio/developer-reference/error-code") diff --git a/old/test/tt.py b/old/test/tt.py deleted file mode 100644 index ea2e6bf..0000000 --- a/old/test/tt.py +++ /dev/null @@ -1,17 +0,0 @@ -import asyncio - -def sync_function(x): - import time - time.sleep(2) # 模拟耗时操作（如文件读写、CPU 计算） - return x * x - -async def async_wrapper(x): - # 在线程中运行同步函数 - result = await asyncio.to_thread(sync_function, x) - return result - -async def main(): - result = await async_wrapper(5) - print(result) # 输出: 25 - -asyncio.run(main()) \ No newline at end of file diff --git a/old/todo_files/__init__.py b/old/todo_files/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/old/todo_files/akshare/__init__.py b/old/todo_files/akshare/__init__.py deleted file mode 100644 index f90d358..0000000 --- a/old/todo_files/akshare/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .get_ak_a_code_op import GetAkACodeOp -from .get_ak_a_info_op import GetAkAInfoOp, GetAkASpotOp, GetAkAMoneyFlowOp, GetAkAFinancialInfoOp, \ - GetAkANewsOp, MergeAkAInfoOp diff --git a/old/todo_files/akshare/get_ak_a_code_op.py b/old/todo_files/akshare/get_ak_a_code_op.py deleted file mode 100644 index dd251d9..0000000 --- a/old/todo_files/akshare/get_ak_a_code_op.py +++ /dev/null @@ -1,107 +0,0 @@ -import json -from typing import List - -import akshare as ak -import pandas as pd -from flowllm.op.base_llm_op import BaseLLMOp -from loguru import logger - -from flowllm.context.flow_context import FlowContext -from flowllm.context.service_context import C -from flowllm.enumeration.role import Role -from flowllm.schema.message import Message -from flowllm.storage.cache_handler.data_cache import DataCache -from flowllm.utils.timer import timer - - -@C.register_op() -class GetAkACodeOp(BaseLLMOp): - file_path: str = __file__ - - def __init__(self, language: str = "zh", llm="qwen3_30b_instruct", **kwargs): - super().__init__(language=language, llm=llm, **kwargs) - - @staticmethod - def download_a_stock_df(): - cache = DataCache() - save_df_key: str = "all_a_stock_name_code" - if not cache.exists(save_df_key): - stock_sh_a_spot_em_df = ak.stock_sh_a_spot_em() - stock_sz_a_spot_em_df = ak.stock_sz_a_spot_em() - stock_bj_a_spot_em_df = ak.stock_bj_a_spot_em() - - df: pd.DataFrame = pd.concat([stock_sh_a_spot_em_df, stock_sz_a_spot_em_df, stock_bj_a_spot_em_df], axis=0) - df = df.drop(columns=["序号"]) - df = df.reset_index(drop=True) - df = df.sort_values(by="代码") - cache.save(save_df_key, df, expire_hours=0.25) - - df = cache.load(save_df_key, dtype={"代码": str}) - return df - - def get_name_code_dict(self) -> dict: - df = self.download_a_stock_df() - - name_code_dict = {} - for line in df.to_dict(orient="records"): - name = line["名称"].replace(" ", "") - code = line["代码"] - name_code_dict[name] = code - logger.info(f"name_code_dict.size={len(name_code_dict)} content={str(name_code_dict)[:50]}...") - return name_code_dict - - @staticmethod - def split_list(array_list: list, n: int): - if n <= 0: - raise ValueError - - length = len(array_list) - base_size = length // n - remainder = length % n - - start = 0 - for i in range(n): - size = base_size + (1 if i < remainder else 0) - end = start + size - yield array_list[start:end] - start = end - - @timer() - def find_stock_codes(self, stock_names: List[str]): - stock_names = "\n".join([x.strip() for x in stock_names if x]) - prompt = self.prompt_format(prompt_name="find_stock_name", - stock_names=stock_names, - query=self.context.query) - logger.info(f"prompt={prompt}") - - def callback_fn(msg: Message): - content = msg.content - if "```" in content: - content = content.split("```")[1] - content = content.strip("json") - content = json.loads(content.strip()) - return content - - codes: List[str] = self.llm.chat(messages=[Message(role=Role.USER, content=prompt)], - enable_stream_print=False, - callback_fn=callback_fn) - return codes - - def execute(self): - name_code_dict = self.get_name_code_dict() - stock_names = list(name_code_dict.keys()) - for p_stock_names in self.split_list(stock_names, n=1): - self.submit_task(self.find_stock_codes, stock_names=p_stock_names) - # time.sleep(1) - - stock_names = sorted(set(self.join_task())) - self.context.code_infos = {name_code_dict[n]: {"股票名称": n} for n in stock_names if n in name_code_dict} - logger.info(f"code_infos={self.context.code_infos}") - - -if __name__ == "__main__": - C.set_service_config().init_by_service_config() - context = FlowContext(query="茅台和五粮现在价格多少？") - - op = GetAkACodeOp() - op(context=context) diff --git a/old/todo_files/akshare/get_ak_a_code_prompt.yaml b/old/todo_files/akshare/get_ak_a_code_prompt.yaml deleted file mode 100644 index 010b11e..0000000 --- a/old/todo_files/akshare/get_ak_a_code_prompt.yaml +++ /dev/null @@ -1,21 +0,0 @@ -find_stock_name_zh: | - # 股票标准名称 - {stock_names} - - # 用户问题 - {query} - - # 任务 - 请你提取出用户问题中提到的股票名称，并通过**股票标准名称**找到对应的标准名称，并以json格式返回。 - 如果用户问题中没有提及股票名称或者**股票标准名称**中没有找到，请返回空"[]"。 - 如果用户只是提到了某个行业，并且没有明确提及股票名称，请返回空"[]"。 - 请思考后输出你的答案。 - - # 答案格式 - ```json - [ - "股票标准名称1", - "股票标准名称2", - ... - ] - ``` diff --git a/old/todo_files/akshare/get_ak_a_info_op.py b/old/todo_files/akshare/get_ak_a_info_op.py deleted file mode 100644 index 91ef6dd..0000000 --- a/old/todo_files/akshare/get_ak_a_info_op.py +++ /dev/null @@ -1,140 +0,0 @@ -import json -import time - -import akshare as ak -import pandas as pd -from loguru import logger -from tqdm import tqdm - -from flowllm.context.flow_context import FlowContext -from flowllm.context.service_context import C -from flowllm.op.base_op import BaseOp -from flowllm.utils.fetch_url import fetch_webpage_text - - -@C.register_op() -class GetAkAInfoOp(BaseOp): - - def execute_code(self, code: str) -> dict: - df = ak.stock_individual_info_em(symbol=code) - result = {} - for line in df.to_dict(orient="records"): - result[line["item"].strip()] = line["value"] - return {"基本信息": result} - - def execute(self): - max_retries: int = self.op_params.get("max_retries", 3) - for code, info_dict in self.context.code_infos.items(): - result = {} - for i in range(max_retries): - try: - result = self.execute_code(code) - break - - except Exception as _: - if i != max_retries - 1: - time.sleep(i * 2 + 1) - - if result: - info_dict.update(result) - - time.sleep(1) - logger.info(f"code_infos={json.dumps(self.context.code_infos, ensure_ascii=False, indent=2)}") - - -@C.register_op() -class GetAkASpotOp(GetAkAInfoOp): - - def execute_code(self, code: str) -> dict: - from flowllm.op.akshare import GetAkACodeOp - - df: pd.DataFrame = GetAkACodeOp.download_a_stock_df() - df = df.loc[df["代码"] == code, :] - result = {} - if len(df) > 0: - result["实时行情"] = df.to_dict(orient="records")[-1] - - return result - - -@C.register_op() -class GetAkAMoneyFlowOp(GetAkAInfoOp): - - def execute_code(self, code: str) -> dict: - df = ak.stock_individual_fund_flow(stock=code) - result = {} - if len(df) > 0: - result["资金流入流出"] = {k: str(v) for k, v in df.to_dict(orient="records")[-1].items()} - return result - - -@C.register_op() -class GetAkAFinancialInfoOp(GetAkAInfoOp): - - def execute_code(self, code: str) -> dict: - df = ak.stock_financial_abstract_ths(symbol=code, indicator="按报告期") - result = {} - if len(df) > 0: - result["财务信息"] = {k: str(v) for k, v in df.to_dict(orient="records")[-1].items()} - return result - - -@C.register_op() -class GetAkANewsOp(GetAkAInfoOp): - - def execute_code(self, code: str) -> dict: - stock_news_em_df = ak.stock_news_em(symbol=code) - top_n_news: int = self.op_params.get("top_n_news", 1) - - news_content_list = [] - for i, line in enumerate(tqdm(stock_news_em_df.to_dict(orient="records")[:top_n_news])): - url = line["新闻链接"] - # http://finance.eastmoney.com/a/202508133482756869.html - ts = url.split("/")[-1].split(".")[0] - date = ts[:8] - content = fetch_webpage_text(url).strip() - content = f"新闻{i}\n时间{date}\n{content}" - news_content_list.append(content) - - return {"新闻": "\n\n".join(news_content_list)} - - -@C.register_op() -class MergeAkAInfoOp(BaseOp): - - def execute(self): - code_content = {} - for code, info_dict in self.context.code_infos.items(): - content_list = [f"\n\n### {code}"] - for key, value in info_dict.items(): - content_list.append(f"\n#### {code}-{key}") - if isinstance(value, str): - content_list.append(value) - elif isinstance(value, dict): - for attr_name, attr_value in value.items(): - content_list.append(f"{attr_name}: {attr_value}") - elif isinstance(value, list): - content_list.extend([x.strip() for x in value if x]) - - code_content[code] = "\n".join(content_list) - - answer = "\n".join(code_content.values()) - logger.info(f"answer=\n{answer}") - self.context.response.answer = answer.strip() - - -if __name__ == "__main__": - C.set_service_config().init_by_service_config() - - code_infos = {"000858": {}, "600519": {}} - context = FlowContext(code_infos=code_infos, query="茅台和五粮现在价格多少？") - - op1 = GetAkAInfoOp() - op2 = GetAkASpotOp() - op3 = GetAkAMoneyFlowOp() - op4 = GetAkAFinancialInfoOp() - op5 = GetAkANewsOp() - op6 = MergeAkAInfoOp() - - op = op1 >> op2 >> op3 >> op4 >> op5 >> op6 - op(context=context) diff --git a/old/todo_files/company_segment_factor_prompt2.yaml b/old/todo_files/company_segment_factor_prompt2.yaml deleted file mode 100644 index e217f0c..0000000 --- a/old/todo_files/company_segment_factor_prompt2.yaml +++ /dev/null @@ -1,256 +0,0 @@ -factor_step1_prompt: | - # 角色定义 - 你是一名金融分析专家，专注于金融因子传导路径的逻辑推导与分析。 - - # 背景信息 - ## 搜索内容 - {search_content} - - ## 分析对象 - 目标公司：{name} - 业务板块：{segment} - - ## 当前金融因子逻辑图 - {mermaid_graph} - - # 任务说明 - 基于**当前金融因子逻辑图**，评估逻辑链的完整性，识别信息缺口或可深化的分析环节。 - 若存在需要补充的信息，请生成搜索查询列表（query list），每条查询将独立调用一次搜索API。 - - ## 查询生成规则 - 1. **查询格式要求**： - - 每条查询必须简洁明确，仅包含关键词，使用空格分隔 - - 单次最多生成 {max_search_cnt} 条查询 - - 2. **避免冗余搜索**： - - 若逻辑路径起点已是明确的具体事件、政策文件或新闻事实，无需为该因子再生成上游查询 - - 例如："2025年欧盟碳关税正式实施"已足够明确，无需再搜索其成因 - - 3. **关键分析维度**： - - **上游维度**：原材料价格、生产设备、供应商状况、成本结构变化 - - **企业自身维度**：产能扩张、技术突破、资本开支、运营效率改善 - - **下游维度**：终端需求变化、客户行业景气度、订单量与订单结构 - - **政策维度**：行业监管政策、政府补贴、关税与贸易壁垒 - - **公司动态维度**：财务报告、官方公告、管理层指引、产能规划 - - **新闻事件维度**：行业重大新闻、公司重要事件、政府政策发布、国际市场动态 - - 4. **分析视角差异**： - - **未盈利公司**：重点关注营收驱动因子，包括市场份额提升、行业总需求增长、新市场开拓等 - - **盈利公司**：重点关注利润驱动因子，包括成本优化、产品结构升级、议价能力提升等 - - **估值直接影响**：关注宏观因素对估值的直接作用，例如： - - 无风险利率下降 → DCF模型折现率降低 → 未来现金流现值提升 → 估值提升 - - 风险偏好改善 → 估值溢价提升 - - 5. **终止条件**： - - 若当前逻辑链已完整且无需补充新信息，则返回空列表 `[]` - - # 输出格式 - ## 思考过程 - [在此进行逐步分析，包括： - - 当前逻辑图的完整性评估 - - 识别出的信息缺口或可展开环节 - - 需要补充的信息及其必要性说明] - - ## 搜索查询列表 - ```json - [ - "查询关键词1 关键词2", - "查询关键词3 关键词4", - "查询关键词5 关键词6" - ] - ``` - - 请严格按照上述格式输出，先进行思考分析，再给出搜索查询列表。 - -factor_step2_prompt: | - # 角色定义 - 你是一名金融分析专家，专注于金融因子传导路径的逻辑推导与图谱构建。 - - # 背景信息 - ## 搜索内容 - {search_content} - - ## 分析对象 - 目标公司：{name} - 业务板块：{segment} - - ## 当前时间 - {current_time} - - ## 当前金融因子逻辑图 - {mermaid_graph} - - # 任务说明 - 根据**搜索内容**中提供的信息，更新现有的**金融因子逻辑图**。 - 更新操作包括：新增逻辑路径、删除过时路径、修改现有路径、在路径中新增/删除/修改因子节点。 - - ## 更新规则与要求 - - ### 1. 因子定义规范 - - 每个因子必须采用 **[实体 + 变化方向]** 的格式，确保具体、可观测 - - 因子描述应简洁明确，例如： - - ✓ "铜价上涨"、"美联储降息"、"公司新建产能投产" - - ✗ "原材料"、"政策变化"、"市场环境"（过于抽象） - - ### 2. 图谱结构要求 - - **统一终点**：所有逻辑路径的终点必须统一为 `A[{name} {segment} 估值增加]` - - **路径长度限制**：每条逻辑路径最多包含 6 个因子节点（含起点和终点） - - **起点要求**：每条路径的起点应为不需进一步向上追溯的客观因子，例如： - - 政策事件："2025年新能源汽车补贴政策发布" - - 市场事件："全球半导体供应短缺" - - 公司事件："公司完成新产线建设" - - ### 3. 特殊情况处理 - - 若起点因子为"利润贡献提升"等中间结果性因子，必须继续向上追溯至根源性因子 - - 例如："利润贡献提升"应追溯至"成本降低"或"售价提升"等更上游因子 - - ### 4. 图谱维护原则 - - **合并重复因子**：相同含义的因子应合并为一个节点 - - **删除过时因子**：与当前时间或搜索内容不符的因子应删除 - - **保持原图**：若搜索内容中无有效新信息，则直接原样输出当前图谱 - - ### 5. 分析视角指引 - - **未盈利公司**：侧重营收驱动路径 - - 市场份额扩大 → 营收增长 → 估值提升 - - 行业需求爆发 → 公司订单增加 → 营收增长 → 估值提升 - - - **盈利公司**：侧重利润驱动路径 - - 原材料价格下降 → 成本降低 → 利润率提升 → 估值提升 - - 产品结构优化 → 毛利率提升 → 利润增长 → 估值提升 - - - **估值直接驱动路径**： - - 央行降息 → 无风险利率下降 → DCF折现率降低 → 估值提升 - - 行业政策利好 → 市场风险偏好提升 → 估值溢价扩大 → 估值提升 - - ## 逻辑路径示例 - ``` - 美联储加息 -> 市场利率上升 -> 银行净息差扩大 -> 利息收入增加 -> [xx银行 借贷业务 估值增加] - ``` - - # 输出格式 - ## 思考过程 - [在此进行逐步分析，包括： - - 搜索内容中的关键信息提取 - - 对现有图谱的影响评估 - - 拟进行的更新操作及理由] - - ## 更新后的金融因子逻辑图 - ```mermaid - graph LR; - B[因子1] --> C[因子2] - C --> A[{name} {segment} 估值增加] - D[因子3] --> E[因子4] - E --> A - ``` - - 请严格按照上述格式输出，先进行思考分析，再给出更新后的图谱。 - - -factor_transmission_prompt: | - # 因子传导路径分析任务 - - ## 分析对象 - 目标公司：{name} - 业务板块：{segment} - - [no_first]## 搜索内容 - [no_first]{search_content} - - ## 现有因子传导逻辑图 - {mermaid_graph} - - # 任务概述 - 围绕 **{name}** 的 **{segment}** 业务，构建并迭代树状结构的**因子传导路径图**。 - 每条路径应从**宏观、行业或公司层面的可观测因子**出发，通过清晰的因果逻辑链，最终传导至 `[{name} {segment} 估值增加]`。 - - **标准路径示例**： - ``` - 美联储加息 -> 市场利率上升 -> 银行净息差扩大 -> 利息收入增加 -> [xx银行 借贷业务 估值增加] - ``` - - --- - - ## 任务一：更新因子传导逻辑图 - - ### 任务目标 - 根据 **搜索内容** 中的信息，更新现有的 **因子传导逻辑图**。 - - ### 更新要求 - 1. **因子规范性**： - - 每个因子必须是 **具体、可观察的实体及其变化方向** - - 正确示例："铜价上涨"、"美联储降息"、"公司新建产能投产" - - 错误示例："市场环境"、"政策因素"（过于抽象） - - 2. **路径结构要求**： - - 所有路径必须以 `[{name} {segment} 估值增加]` 为唯一终点 - - 单条路径最多包含 6 个因子节点 - - 路径起点应为明确的事件、政策或新闻，无需进一步向上追溯 - - 3. **图谱维护原则**： - - 合并语义重复的因子节点 - - 删除与当前时间不符或已过时的因子 - - 若搜索内容无有效信息，则保持原图不变，直接原样输出 - - 4. **特殊情况处理**： - - 若路径起点为"利润贡献提升"等中间结果性因子，需继续追溯至根本驱动因素 - - --- - - ## 任务二：生成搜索查询列表 - - ### 任务目标 - 基于 **更新后的因子传导逻辑图**，识别信息缺口或可深化的分析环节，生成搜索查询列表。 - - ### 查询生成原则 - 1. **避免冗余搜索**： - - 若路径起点已是明确的具体事件、政策或新闻（如"2025年欧盟碳关税正式实施"），则无需为该因子生成上游查询 - - 2. **关键分析维度**： - - **上游维度**：原材料价格、生产设备、供应商状况、成本结构 - - **企业自身维度**：产能变化、技术进展、资本开支、运营效率 - - **下游维度**：终端需求、客户行业景气度、订单数据 - - **政策维度**：行业监管、政府补贴、贸易政策 - - **公司动态维度**：财报数据、官方公告、管理层指引、产能规划 - - 3. **查询格式要求**： - - 每条查询应简洁明确，仅包含关键词，空格分隔 - - 单次最多生成 {max_search_cnt} 条查询 - - 4. **终止条件**： - - 若当前逻辑链已完整且无需补充信息，则返回空列表 `[]` - - --- - - # 输出格式 - - ## 任务一：思考过程 - [分析搜索内容对现有图谱的影响，说明拟进行的更新操作] - - ## 任务一：更新后的因子传导逻辑图 - ```mermaid - graph LR; - B[因子1] --> C[因子2] - C --> A[{name} {segment} 估值增加] - D[因子3] --> A - ``` - - ## 任务二：思考过程 - [分析当前图谱的信息缺口，说明需要补充的搜索方向] - - ## 任务二：搜索查询列表 - ```json - [ - "关键词1 关键词2", - "关键词3 关键词4" - ] - ``` - - 请严格按照上述格式输出，逐步完成两个任务。 - - -init_mermaid_graph: | - ```mermaid - graph LR; - A[{name} {segment} 估值增加] - ``` - diff --git a/old/todo_files/company_segment_factor_prompt_old.yaml b/old/todo_files/company_segment_factor_prompt_old.yaml deleted file mode 100644 index c1d2bb1..0000000 --- a/old/todo_files/company_segment_factor_prompt_old.yaml +++ /dev/null @@ -1,163 +0,0 @@ -factor_step1_prompt: | - # 角色 - 你是一名金融专家，擅长金融逻辑推导。 - - {search_content} - - # 目标公司与业务板块 - {name}的{segment}业务 - - # 当前的金融因子逻辑图 - {mermaid_graph} - - # 任务 - 基于**当前的金融因子逻辑图**，判断当前因子逻辑图是否存在信息缺口或可进一步展开的环节。 - 如有需要，请生成用于补充信息的搜索query list，每条query将单独调用一次搜索API，最多{max_search_cnt}条。 - 生成query时请遵循以下原则： - 1. 每一条query要求简洁，只包含关键词，空格分隔。 - 2. **无需重复搜索已明确的起点**：如果当前路径起点已是具体事件、具体政策或明确新闻，则无需再为此因子生成搜索query。 - 3. **关键因子**： - - 上游：原材料、设备、供应商、成本结构 - - 自身：产能、技术、资本开支、运营效率 - - 下游：终端需求、客户行业景气度、订单情况 - - 政策：行业监管、补贴、贸易壁垒 - - 公司动态：财报、公告、管理层指引、产能规划 - - 新闻：行业新闻、公司新闻、政府新闻、国际新闻 - 4. **关键角度**： - - 未盈利的公司可以更多的关注因子对于营收的影响，可以是市场占有率的提升，也可以是行业整体营收在未来的提升。 - - 盈利的公司可以更关注因子对于利润的提升。 - - 也可以关注一些宏观角度对估值的直接影响，例如宏观利率的下降可以DCF公式中的无风险利率的下降，进而导致未来现金流的上升，以及估值的提升。 - 5. 若当前逻辑链已完整或无新信息可补充，则返回空列表`[]`。 - - # 输出示例 - ## 思考 - ... - - ## 搜索query list - ```json - [ - "xxx 需求预测", - "xxx 出货指引", - "xxx 2025 年装机量" - ] - ``` - - 请先一步一步思考，然后按照输出示例返回结果。 - -factor_step2_prompt: | - # 角色 - 你是一名金融专家，擅长金融逻辑推导。 - - # 搜索内容 - {search_content} - - # 目标公司与业务板块 - {name}的{segment}业务 - - # 当前时间 - {current_time} - - # 当前的金融因子逻辑图 - {mermaid_graph} - - # 任务 - 请根据**搜索内容**中提供的信息，更新现有的**当前的金融因子逻辑图**，可以是新增、删除、修改其中的逻辑路径。也可以是在其中一条逻辑路径中新增、删除、修改因子。 - - ## 要求如下： - - 合并重复的因子 - - 删除过时的因子 - - 每个因子必须是[实体+变化方向]，同时要求简洁。例如：“黄金上涨”、“美联储降息”、“公司新建产能投产”）。 - - 可以有多条逻辑路径，但所有的逻辑路径最终都会指向A[{name} {segment} 估值增加]这个因子。 - - 每条逻辑路径最多包含6个因子。每条逻辑路径的起点应该是事件、政策或新闻等一些不需要被扩展的客观因子。终点是A[{name} {segment} 估值增加]这个因子。 - - 如果每条逻辑路径的起点是"利润贡献提升"因子，那可以进一步增加上游因子，什么因子导致利润贡献提升的？否则就不要增加该因子。 - - 若搜索内容中无有效新信息，则直接原样输出。 - - ## 单条逻辑路径示例 - 美联储加息 -> 市场利率上升 -> 银行净息差扩大 -> [xx银行 借贷业务 估值增加] - - ## 思考角度示例： - - 未盈利的公司可以更多的关注因子对于营收的影响，可以是市场占有率的提升，也可以是行业整体营收在未来的提升。 - - 盈利的公司可以更关注因子对于利润的提升。 - - 也可以关注一些宏观角度对估值的直接影响，例如宏观利率的下降可以DCF公式中的无风险利率的下降，进而导致未来现金流的上升，以及估值的提升。 - - # 输出示例 - ## 思考 - ... - - ## 更新的金融因子逻辑图 - ```mermaid - graph LR; - ... - ``` - - - 请先一步一步思考，然后按照输出示例返回结果。 - - -factor_transmission_prompt: | - # 因子传导路径分析任务 - - ## 目标公司与业务板块 - {name} 的 {segment} 业务 - - [no_first]## 搜索内容 - [no_first]{search_content} - - ## 现有因子传导逻辑图 - {mermaid_graph} - - ## 任务目标 - 你需要围绕{name}的{segment}业务，根据不同维度的搜索来迭代树状结构的**因子传导路径**。 - 每条路径应从**宏观、行业或公司层面的可观测因子**（如政策变动、原材料价格、产能变化、需求波动等）出发，通过清晰的因果逻辑链，最终传导至[{name} {segment} 估值增加]。 - 例如：美联储加息 -> 市场利率上升 -> 银行净息差扩大 -> [xx银行 借贷业务 估值增加] - - ### 任务-更新因子传导逻辑图 - 请根据 **搜索内容** 中提供的信息，更新现有的 **因子传导逻辑图**。 - 更新时需满足以下要求： - - 每个因子必须是**具体、可观察的实体及其变化方向**（例如：“铜价上涨”、“美联储降息”、“公司新建产能投产”）。 - - 若搜索内容中**无有效新信息**，则**保持原图不变**，直接原样输出。 - - ### 任务-生成搜索查询 - 基于**更新后的因子传导逻辑图**，判断当前路径是否存在信息缺口或可进一步展开的环节。 - 如有需要，请生成用于补充信息的**搜索查询列表**，每条查询将单独调用一次搜索 API。 - - 生成查询时请遵循以下原则： - 1. **无需重复搜索已明确的起点**：如果当前路径起点已是具体事件、政策或新闻（如“2025年欧盟碳关税正式实施”），则无需再为此因子生成上游查询。 - 2. **聚焦关键维度**：可从以下角度展开： - - 上游：原材料、设备、供应商、成本结构 - - 自身：产能、技术、资本开支、运营效率 - - 下游：终端需求、客户行业景气度、订单情况 - - 政策：行业监管、补贴、贸易壁垒 - - 公司动态：财报、公告、管理层指引、产能规划 - 3. 若当前逻辑链已完整或无新信息可补充，则返回空列表`[]`。 - - ## 输出示例 - ### 任务一思考 - ... - - ### 任务一输出示例（更新后的 Mermaid 图） - ```mermaid - graph LR; - B[factor1 增加] -> A[xxx 估值增加] - ``` - - ### 任务二思考 - ... - - ### 任务二输出示例（搜索查询列表） - ```json - [ - "xxx 上游价格", - "xxx 产能", - "xxx 需求预测" - ] - ``` - - 请先一步一步思考，然后按照输出示例返回结果。 - - -init_mermaid_graph: | - ```mermaid - graph LR; - A[{name} {segment} 估值增加] - ``` \ No newline at end of file diff --git a/old/todo_files/crawl_utils.py b/old/todo_files/crawl_utils.py deleted file mode 100644 index 76d45e7..0000000 --- a/old/todo_files/crawl_utils.py +++ /dev/null @@ -1,117 +0,0 @@ -import random -import time - -import requests -import urllib3 -from bs4 import BeautifulSoup - -# Disable SSL warnings (optional, for handling insecure HTTPS) -urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) - - -def get_random_headers(): - """Generate random headers to avoid detection""" - user_agents = [ - 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36', - 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/92.0.4515.107 Safari/537.36', - 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/93.0.4577.82 Safari/537.36', - 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36', - 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/14.1.1 Safari/605.1.15', - 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36', - 'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:89.0) Gecko/20100101 Firefox/89.0', - 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:89.0) Gecko/20100101 Firefox/89.0', - 'Mozilla/5.0 (X11; Linux x86_64; rv:89.0) Gecko/20100101 Firefox/89.0', - 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Edge/91.0.864.59' - ] - - accept_languages = [ - 'en-US,en;q=0.9', - 'zh-CN,zh;q=0.9,en;q=0.8', - 'en-GB,en;q=0.9', - 'fr-FR,fr;q=0.9,en;q=0.8', - 'de-DE,de;q=0.9,en;q=0.8' - ] - - accept_encodings = [ - 'gzip, deflate, br', - 'gzip, deflate', - 'br, gzip, deflate' - ] - - headers = { - 'User-Agent': random.choice(user_agents), - 'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8', - 'Accept-Language': random.choice(accept_languages), - 'Accept-Encoding': random.choice(accept_encodings), - 'Connection': 'keep-alive', - 'Upgrade-Insecure-Requests': '1', - } - - # Randomly add some optional headers - if random.random() > 0.5: - headers['DNT'] = '1' - if random.random() > 0.7: - headers['Cache-Control'] = 'max-age=0' - if random.random() > 0.6: - headers['Sec-Fetch-Dest'] = 'document' - headers['Sec-Fetch-Mode'] = 'navigate' - headers['Sec-Fetch-Site'] = 'none' - - return headers - - -def fetch_webpage_text(url, min_delay=1, max_delay=3): - """ - Fetch and extract text content from a webpage with randomization - - Args: - url (str): The URL to fetch - min_delay (int): Minimum delay before request (seconds) - max_delay (int): Maximum delay before request (seconds) - - Returns: - str: Extracted text content or error message - """ - # Add random delay to avoid being detected as bot - delay = random.uniform(min_delay, max_delay) - time.sleep(delay) - - headers = get_random_headers() - - # Random timeout between 8-15 seconds - timeout = random.randint(8, 15) - - try: - # Send request with random headers and timeout - response = requests.get(url, headers=headers, timeout=timeout, verify=False) - response.raise_for_status() # Check if request was successful - response.encoding = response.apparent_encoding # Auto-detect encoding - - # Parse HTML using BeautifulSoup - soup = BeautifulSoup(response.text, 'html.parser') - - # Remove script, style and navigation elements to avoid interference - for script_or_style in soup(['script', 'style', 'nav', 'footer', 'header']): - script_or_style.decompose() - - # Extract text content - text = soup.get_text() - - # Clean whitespace: remove extra blank lines and spaces - lines = (line.strip() for line in text.splitlines()) - chunks = (phrase.strip() for line in lines for phrase in line.split(" ")) - text = ' '.join(chunk for chunk in chunks if chunk) - - return text - - except requests.exceptions.RequestException as e: - return f"Request failed: {e}" - except Exception as e: - return f"Parsing failed: {e}" - - -# Example usage -if __name__ == "__main__": - url = "http://finance.eastmoney.com/a/202508133482756869.html" - text = fetch_webpage_text(url) - print(text) diff --git a/old/todo_files/miner_u_pdf_processor.py b/old/todo_files/miner_u_pdf_processor.py deleted file mode 100644 index 71bcdfe..0000000 --- a/old/todo_files/miner_u_pdf_processor.py +++ /dev/null @@ -1,726 +0,0 @@ -#!/usr/bin/env python3 -""" -MinerU PDF Processor - -A comprehensive PDF processing utility that leverages MinerU for extracting structured content -from PDF documents. Returns both Markdown content and structured content lists for further processing. - -This processor provides a high-level interface to MinerU's command-line tools, handling -file I/O, error management, and result parsing automatically. -""" - -import json -import logging -import os -import platform -import re -import subprocess -from pathlib import Path -from typing import Dict, List, Any, Tuple, Optional, Union - - -class MinerUPDFProcessor: - """ - MinerU-based PDF Processing Engine - - A robust PDF processor that wraps MinerU functionality to extract structured content - from PDF documents. Inspired by RAGAnything's processing logic but operates independently - with MinerU as the core engine. - - Features: - - Automatic MinerU installation validation - - Multiple parsing methods (auto, txt, ocr) - - Language-specific OCR optimization - - Structured content extraction with metadata - - Image path resolution and management - - Comprehensive error handling and logging - - Example: - processor = MinerUPDFProcessor(log_level="INFO") - content_list, markdown = processor.process_pdf("document.pdf") - """ - - def __init__(self, log_level: str = "INFO"): - """ - Initialize the PDF processor with logging configuration. - - Args: - log_level (str): Logging level for the processor. - Options: "DEBUG", "INFO", "WARNING", "ERROR" - - Raises: - RuntimeError: If MinerU is not properly installed or accessible - """ - # Configure logging system - logging.basicConfig( - level=getattr(logging, log_level.upper()), - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' - ) - self.logger = logging.getLogger(__name__) - - # Validate MinerU installation before proceeding - if not self.check_mineru_installation(): - raise RuntimeError( - "MinerU is not properly installed. Please install using:\n" - "pip install -U 'mineru[core]' or uv pip install -U 'mineru[core]'" - ) - - @classmethod - def create_with_defaults(cls, log_level: str = "INFO") -> "MinerUPDFProcessor": - """ - Create a MinerUPDFProcessor instance with default settings. - - Convenience method for quick instantiation with standard configuration. - - Args: - log_level (str): Logging level (default: "INFO") - - Returns: - MinerUPDFProcessor: Configured processor instance - """ - return cls(log_level=log_level) - - def check_mineru_installation(self) -> bool: - """ - Verify that MinerU is properly installed and accessible. - - Attempts to run the MinerU command-line tool to check its availability - and version information. - - Returns: - bool: True if MinerU is properly installed, False otherwise - """ - try: - # Configure subprocess parameters for cross-platform compatibility - subprocess_kwargs = { - "capture_output": True, - "text": True, - "check": True, - "encoding": "utf-8", - "errors": "ignore", - } - - # Hide console window on Windows systems - if platform.system() == "Windows": - subprocess_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW - - # Execute version check command - result = subprocess.run(["mineru", "--version"], **subprocess_kwargs) - self.logger.debug(f"MinerU version detected: {result.stdout.strip()}") - return True - except (subprocess.CalledProcessError, FileNotFoundError): - return False - - def _run_mineru_command( - self, - input_path: Union[str, Path], - output_dir: Union[str, Path], - method: str = "auto", - lang: Optional[str] = None, - backend: str = "pipeline", - start_page: Optional[int] = None, - end_page: Optional[int] = None, - formula: bool = True, - table: bool = True, - device: Optional[str] = None, - source: str = "modelscope", - vlm_url: Optional[str] = None, - ) -> None: - """ - Execute MinerU command-line tool with specified parameters. - - This method constructs and executes the MinerU command with all provided - options, handling cross-platform subprocess execution and error management. - - Args: - input_path (Union[str, Path]): Path to the input PDF file - output_dir (Union[str, Path]): Directory path for output files - method (str): Parsing method - "auto", "txt", or "ocr" - lang (Optional[str]): Document language for OCR optimization (e.g., "en", "ch", "ja") - backend (str): Processing backend to use - start_page (Optional[int]): Starting page number (0-based indexing) - end_page (Optional[int]): Ending page number (0-based indexing) - formula (bool): Enable mathematical formula parsing - table (bool): Enable table structure parsing - device (Optional[str]): Computing device for inference (e.g., "cuda", "cpu") - source (str): Model source repository - vlm_url (Optional[str]): VLM server URL (required for vlm-sglang-client backend) - - Raises: - subprocess.CalledProcessError: If MinerU command execution fails - FileNotFoundError: If MinerU executable is not found - RuntimeError: If MinerU is not properly installed - """ - # Build base command with required parameters - cmd = [ - "mineru", - "-p", str(input_path), - "-o", str(output_dir), - "-m", method, - # Note: backend and source parameters are commented out as they may not be - # available in all MinerU versions or configurations - # "-b", backend, - # "--source", source, - ] - - # Add optional parameters if specified - if lang: - cmd.extend(["-l", lang]) - if start_page is not None: - cmd.extend(["-s", str(start_page)]) - if end_page is not None: - cmd.extend(["-e", str(end_page)]) - if not formula: - cmd.extend(["-f", "false"]) - if not table: - cmd.extend(["-t", "false"]) - if device: - cmd.extend(["-d", device]) - if vlm_url: - cmd.extend(["-u", vlm_url]) - - try: - # Configure subprocess execution parameters - subprocess_kwargs = { - "capture_output": True, - "text": True, - "check": True, - "encoding": "utf-8", - "errors": "ignore", - } - - # Hide console window on Windows systems - if platform.system() == "Windows": - subprocess_kwargs["creationflags"] = subprocess.CREATE_NO_WINDOW - - self.logger.info(f"Executing MinerU command: {' '.join(cmd)}") - result = subprocess.run(cmd, **subprocess_kwargs) - - self.logger.info("MinerU command executed successfully") - if result.stdout: - self.logger.debug(f"MinerU output: {result.stdout}") - - except subprocess.CalledProcessError as e: - self.logger.exception(f"MinerU command execution failed: {e}") - if e.stderr: - self.logger.exception(f"Error details: {e.stderr}") - raise - except FileNotFoundError: - raise RuntimeError( - "MinerU command not found. Please ensure MinerU 2.0 is properly installed:\n" - "pip install -U 'mineru[core]' or uv pip install -U 'mineru[core]'" - ) - - def _read_output_files( - self, - output_dir: Path, - file_stem: str, - method: str = "auto" - ) -> Tuple[List[Dict[str, Any]], str]: - """ - Read and parse MinerU-generated output files. - - This method locates and reads the Markdown and JSON files generated by MinerU, - handling different directory structures and resolving image paths to absolute paths. - - Args: - output_dir (Path): Directory containing the MinerU output files - file_stem (str): Base filename without extension - method (str): Parsing method used ("auto", "txt", "ocr", "vlm") - - Returns: - Tuple[List[Dict[str, Any]], str]: A tuple containing: - - content_list: Structured content list with metadata - - markdown_content: Raw Markdown text content - """ - # Locate generated output files - handle both flat and nested directory structures - md_file = output_dir / f"{file_stem}.md" - json_file = output_dir / f"{file_stem}_content_list.json" - images_base_dir = output_dir - - # Check for nested subdirectory structure (common with newer MinerU versions) - file_stem_subdir = output_dir / file_stem - if file_stem_subdir.exists(): - md_file = file_stem_subdir / method / f"{file_stem}.md" - json_file = file_stem_subdir / method / f"{file_stem}_content_list.json" - images_base_dir = file_stem_subdir / method - - # Read Markdown content - md_content = "" - if md_file.exists(): - try: - with open(md_file, "r", encoding="utf-8") as f: - md_content = f.read() - self.logger.info(f"Successfully read Markdown file: {md_file}") - except Exception as e: - self.logger.warning(f"Failed to read Markdown file {md_file}: {e}") - else: - self.logger.warning(f"Markdown file not found: {md_file}") - - # Read structured content list from JSON - content_list = [] - if json_file.exists(): - try: - with open(json_file, "r", encoding="utf-8") as f: - content_list = json.load(f) - - # Convert relative image paths to absolute paths for proper access - self.logger.info(f"Resolving image paths relative to: {images_base_dir}") - for item in content_list: - if isinstance(item, dict): - # Process various image path fields that may be present - for field_name in ["img_path", "table_img_path", "equation_img_path"]: - if field_name in item and item[field_name]: - img_path = item[field_name] - if not os.path.isabs(img_path): - absolute_img_path = (images_base_dir / img_path).resolve() - item[field_name] = str(absolute_img_path) - self.logger.debug(f"Updated {field_name}: {img_path} -> {item[field_name]}") - - self.logger.info( - f"Successfully read JSON file: {json_file}, containing {len(content_list)} content blocks") - - except Exception as e: - self.logger.warning(f"Failed to read JSON file {json_file}: {e}") - else: - self.logger.warning(f"JSON file not found: {json_file}") - - return content_list, md_content - - def process_pdf( - self, - pdf_path: Union[str, Path], - output_dir: Optional[Union[str, Path]] = None, - method: str = "auto", - lang: Optional[str] = None, - backend: str = "pipeline", - **kwargs - ) -> Tuple[List[Dict[str, Any]], str]: - """ - Process a PDF file and extract structured content using MinerU. - - This is the main entry point for PDF processing. It validates input, - executes MinerU processing, and returns both structured content and markdown. - - Args: - pdf_path (Union[str, Path]): Path to the input PDF file - output_dir (Optional[Union[str, Path]]): Output directory path. - If None, creates 'mineru_output' in PDF's directory - method (str): Parsing method - "auto" (recommended), "txt", or "ocr" - lang (Optional[str]): Document language for OCR optimization - (e.g., "ch" for Chinese, "en" for English, "ja" for Japanese) - backend (str): Processing backend - "pipeline", "vlm-transformers", - "vlm-sglang-engine", or "vlm-sglang-client" - **kwargs: Additional MinerU parameters (start_page, end_page, formula, table, etc.) - - Returns: - Tuple[List[Dict[str, Any]], str]: A tuple containing: - - content_list: Structured list of content blocks with metadata - - markdown_content: Complete document in Markdown format - - Raises: - FileNotFoundError: If the specified PDF file does not exist - ValueError: If the file is not a valid PDF format - RuntimeError: If MinerU processing fails or encounters errors - """ - # Convert to Path object and validate input - pdf_path = Path(pdf_path) - if not pdf_path.exists(): - raise FileNotFoundError(f"PDF file does not exist: {pdf_path}") - - if not pdf_path.suffix.lower() == '.pdf': - raise ValueError(f"File is not a PDF format: {pdf_path}") - - name_without_suffix = pdf_path.stem - - # Prepare output directory - if output_dir: - base_output_dir = Path(output_dir) - else: - base_output_dir = pdf_path.parent / "mineru_output" - - base_output_dir.mkdir(parents=True, exist_ok=True) - - try: - # Execute MinerU processing - self.logger.info(f"Starting PDF processing: {pdf_path}") - - self._run_mineru_command( - input_path=pdf_path, - output_dir=base_output_dir, - method=method, - lang=lang, - backend=backend, - **kwargs - ) - - # Read generated output files - backend_method = method - if backend.startswith("vlm-"): - backend_method = "vlm" - - content_list, markdown_content = self._read_output_files( - base_output_dir, name_without_suffix, method=backend_method - ) - - # Generate processing statistics - content_stats = {} - for item in content_list: - if isinstance(item, dict): - content_type = item.get("type", "unknown") - content_stats[content_type] = content_stats.get(content_type, 0) + 1 - - self.logger.info(f"PDF processing completed! Extracted {len(content_list)} content blocks") - self.logger.info("Content type statistics:") - for content_type, count in content_stats.items(): - self.logger.info(f" - {content_type}: {count}") - - return content_list, markdown_content - - except Exception as e: - self.logger.exception(f"Error occurred during PDF processing: {str(e)}") - raise - - def save_results( - self, - content_list: List[Dict[str, Any]], - markdown_content: str, - output_path: Union[str, Path], - save_markdown: bool = True, - save_json: bool = True, - indent: int = 2 - ) -> Dict[str, Path]: - """ - Save processing results to files. - - Saves the extracted content in both JSON (structured) and Markdown (text) formats - for different use cases and downstream processing needs. - - Args: - content_list (List[Dict[str, Any]]): Structured content list with metadata - markdown_content (str): Complete document in Markdown format - output_path (Union[str, Path]): Output file path (without extension) - save_markdown (bool): Whether to save Markdown file - save_json (bool): Whether to save JSON file with structured content - indent (int): JSON file indentation for readability - - Returns: - Dict[str, Path]: Dictionary mapping file types to their saved paths - Keys: 'markdown', 'json' (if respective files were saved) - - Raises: - Exception: If file writing operations fail - """ - output_path = Path(output_path) - saved_files = {} - - try: - # Ensure output directory exists - output_path.parent.mkdir(parents=True, exist_ok=True) - - # Save Markdown file - if save_markdown and markdown_content: - md_path = output_path.with_suffix('.md') - with open(md_path, 'w', encoding='utf-8') as f: - f.write(markdown_content) - saved_files['markdown'] = md_path - self.logger.info(f"Markdown file saved: {md_path}") - - # Save JSON file with structured content - if save_json and content_list: - json_path = output_path.with_suffix('.json') - with open(json_path, 'w', encoding='utf-8') as f: - json.dump(content_list, f, indent=indent, ensure_ascii=False) - saved_files['json'] = json_path - self.logger.info(f"JSON file saved: {json_path}") - - return saved_files - - except Exception as e: - self.logger.exception(f"Error occurred while saving files: {e}") - raise - - @staticmethod - def get_content_statistics(content_list: List[Dict[str, Any]]) -> Dict[str, Any]: - """ - Generate detailed statistics about the processed content. - - Analyzes the content list to provide insights into document structure, - content types, and processing results. - - Args: - content_list (List[Dict[str, Any]]): Structured content list from MinerU - - Returns: - Dict[str, Any]: Dictionary containing various statistics: - - total_blocks: Total number of content blocks - - content_types: Count of each content type - - text_stats: Text-specific statistics (characters, words, etc.) - - image_count: Number of images found - - table_count: Number of tables found - """ - stats = { - "total_blocks": len(content_list), - "content_types": {}, - "text_stats": {"total_characters": 0, "total_words": 0, "title_levels": {}}, - "image_count": 0, - "table_count": 0, - "has_formulas": False - } - - for item in content_list: - if not isinstance(item, dict): - continue - - content_type = item.get("type", "unknown") - stats["content_types"][content_type] = stats["content_types"].get(content_type, 0) + 1 - - if content_type == "text": - text = item.get("text", "") - stats["text_stats"]["total_characters"] += len(text) - stats["text_stats"]["total_words"] += len(text.split()) - - level = item.get("text_level", 0) - if level > 0: - stats["text_stats"]["title_levels"][level] = stats["text_stats"]["title_levels"].get(level, 0) + 1 - - elif content_type == "image": - stats["image_count"] += 1 - - elif content_type == "table": - stats["table_count"] += 1 - - elif content_type == "formula": - stats["has_formulas"] = True - - return stats - - def validate_output_quality(self, content_list: List[Dict[str, Any]], markdown_content: str) -> Dict[str, Any]: - """ - Validate the quality and completeness of the processing output. - - Performs various checks to ensure the processed content meets quality standards - and provides warnings or suggestions for improvement. - - Args: - content_list (List[Dict[str, Any]]): Structured content list - markdown_content (str): Markdown content string - - Returns: - Dict[str, Any]: Validation results containing: - - is_valid: Overall validation status - - warnings: List of warning messages - - suggestions: List of improvement suggestions - - quality_score: Numeric quality score (0-100) - """ - validation = { - "is_valid": True, - "warnings": [], - "suggestions": [], - "quality_score": 100 - } - - # Check if content was extracted - if not content_list and not markdown_content.strip(): - validation["is_valid"] = False - validation["warnings"].append("No content was extracted from the PDF") - validation["quality_score"] = 0 - return validation - - # Check content diversity - stats = self.get_content_statistics(content_list) - if stats["total_blocks"] < 5: - validation["warnings"].append("Very few content blocks extracted - document may be complex or image-heavy") - validation["quality_score"] -= 20 - - # Check text content ratio - text_blocks = stats["content_types"].get("text", 0) - if text_blocks == 0: - validation["warnings"].append("No text blocks found - consider using OCR method for image-based PDFs") - validation["quality_score"] -= 30 - elif text_blocks / stats["total_blocks"] < 0.3: - validation["suggestions"].append("Low text content ratio - document may benefit from OCR processing") - validation["quality_score"] -= 10 - - # Check for images without processing - if stats["image_count"] > 0 and stats["content_types"].get("text", 0) == 0: - validation["suggestions"].append( - "Images detected but no text extracted - consider using VLM backend for image analysis") - - # Check markdown length vs content blocks - if len(markdown_content.strip()) < 100 and stats["total_blocks"] > 10: - validation["warnings"].append("Markdown content seems unusually short for the number of content blocks") - validation["quality_score"] -= 15 - - return validation - - -def chunk_pdf_content(content_list: List[Dict[str, Any]], max_length: int = 4000) -> List[str]: - """ - Split MinerU-parsed content list into text chunks of specified length. - - This utility function converts structured content from MinerU into manageable - text chunks suitable for downstream processing like embedding generation or - language model input. - - Args: - content_list (List[Dict[str, Any]]): MinerU-parsed structured content list - max_length (int): Maximum character length per chunk (default: 4000) - - Returns: - List[str]: List of text chunks, each prefixed with chunk metadata - including chunk number, total chunks, and character count - """ - - def extract_text(item: Dict[str, Any]) -> str: - """ - Extract text content from a single content item. - - Handles different content types (text, table, image) and formats them - appropriately for text-based processing. - - Args: - item (Dict[str, Any]): Single content item from MinerU output - - Returns: - str: Extracted and formatted text content - """ - if item.get("type") == "text": - text = item.get("text", "").strip() - if not text: - return "" - # Add markdown header formatting for titles - level = item.get("text_level", 0) - if level > 0: - return f"{'#' * min(level, 6)} {text}" - return text - - elif item.get("type") == "table": - parts = [] - if item.get("table_caption"): - parts.append("Table: " + " | ".join(item["table_caption"])) - if item.get("table_body"): - # Simple HTML tag cleanup and formatting - table_text = re.sub(r'<[^>]+>', ' | ', item["table_body"]) - table_text = re.sub(r'\s+', ' ', table_text).strip() - parts.append(table_text) - return "\n".join(parts) if parts else "" - - elif item.get("type") == "image": - if item.get("image_caption"): - return "Image: " + " | ".join(item["image_caption"]) - return "" - - return "" - - # Extract all text content from the structured list - all_text = "" - for item in content_list: - text = extract_text(item) - if text.strip(): - all_text += text + "\n" - - if not all_text.strip(): - return [] - - # Split into chunks based on max_length - chunks = [] - current_chunk = "" - - for line in all_text.split('\n'): - # Check if adding this line would exceed max_length - if len(current_chunk) + len(line) + 1 > max_length and current_chunk: - chunks.append(current_chunk.strip()) - current_chunk = line - else: - current_chunk += line + "\n" if current_chunk else line - - # Add the final chunk if it contains content - if current_chunk.strip(): - chunks.append(current_chunk.strip()) - - # Add chunk metadata headers - total_chunks = len(chunks) - marked_chunks = [] - for i, chunk in enumerate(chunks): - header = f"=== CHUNK {i + 1}/{total_chunks} ({len(chunk)} characters) ===\n" - marked_chunks.append(header + chunk) - - return marked_chunks - - -# Example usage and demonstration -if __name__ == "__main__": - """ - Example usage of the MinerUPDFProcessor class. - - This example demonstrates the basic workflow for processing a PDF file - and working with the extracted content. - """ - import sys - - - # Example usage - def example_usage(): - """Demonstrate basic PDF processing workflow.""" - try: - # Initialize processor - processor = MinerUPDFProcessor.create_with_defaults(log_level="INFO") - - # Example PDF path (replace with actual PDF file) - pdf_path = "example_document.pdf" - - if not Path(pdf_path).exists(): - print(f"Example PDF file not found: {pdf_path}") - print("Please provide a valid PDF file path to test the processor.") - return - - # Process PDF with different methods - print("Processing PDF with auto method...") - content_list, markdown_content = processor.process_pdf( - pdf_path=pdf_path, - method="auto", - lang="en" # Specify language for better OCR results - ) - - # Generate statistics - stats = processor.get_content_statistics(content_list) - print(f"Processing Statistics:") - print(f" Total blocks: {stats['total_blocks']}") - print(f" Content types: {stats['content_types']}") - print(f" Text characters: {stats['text_stats']['total_characters']}") - print(f" Text words: {stats['text_stats']['total_words']}") - - # Validate output quality - validation = processor.validate_output_quality(content_list, markdown_content) - print(f"Quality Score: {validation['quality_score']}/100") - if validation['warnings']: - print("Warnings:", validation['warnings']) - if validation['suggestions']: - print("Suggestions:", validation['suggestions']) - - # Save results - output_path = Path(pdf_path).stem + "_processed" - saved_files = processor.save_results( - content_list=content_list, - markdown_content=markdown_content, - output_path=output_path - ) - print(f"Results saved to: {saved_files}") - - # Create text chunks for downstream processing - chunks = chunk_pdf_content(content_list, max_length=2000) - print(f"Created {len(chunks)} text chunks") - - # Display first chunk as example - if chunks: - print("First chunk preview:") - print(chunks[0][:200] + "..." if len(chunks[0]) > 200 else chunks[0]) - - except Exception as e: - print(f"Error during processing: {e}") - sys.exit(1) - - - # Run example if script is executed directly - example_usage() diff --git a/old/todo_files/open_compress_op.py b/old/todo_files/open_compress_op.py deleted file mode 100644 index f96d9ac..0000000 --- a/old/todo_files/open_compress_op.py +++ /dev/null @@ -1,53 +0,0 @@ -from typing import List - -from loguru import logger - -from flowllm.context import C -from flowllm.enumeration.chunk_enum import ChunkEnum -from flowllm.enumeration.role import Role -from flowllm.op.base_async_tool_op import BaseAsyncToolOp -from flowllm.schema.message import Message -from flowllm.schema.tool_call import ToolCall -from flowllm.utils.common_utils import get_datetime - - -@C.register_op(register_app="FlowLLM") -class OpenCompressOp(BaseAsyncToolOp): - file_path: str = __file__ - - def __init__(self, llm: str = "qwen3_80b_instruct", return_answer: bool = False, **kwargs): - super().__init__(llm=llm, **kwargs) - self.return_answer: bool = return_answer - - def build_tool_call(self) -> ToolCall: - return ToolCall(**{ - "description": "compress messages", - "input_schema": { - "messages": { - "type": "array", - "description": "messages", - "required": True - } - } - }) - - async def async_execute(self): - messages: List[Message] = self.input_dict["messages"] - messages = [Message(**m) if isinstance(m, dict) else m for m in messages] - - compress_system_prompt: str = self.prompt_format("compress_system_prompt", date=get_datetime()) - messages = [ - Message(role=Role.SYSTEM, content=compress_system_prompt), - *messages, - Message(role=Role.USER, content=self.get_prompt("compress_user_prompt")) - ] - - logger.info(f"messages={messages}") - assistant_message = await self.llm.achat(messages=messages) - chunk_type: ChunkEnum = ChunkEnum.ANSWER if self.return_answer else ChunkEnum.THINK - content = f"{self.name}.{self.tool_index} content={assistant_message.content}" - await self.context.add_stream_chunk_and_type(content, chunk_type) - self.set_result(assistant_message.content) - - - diff --git a/old/todo_files/open_compress_prompt.yaml b/old/todo_files/open_compress_prompt.yaml deleted file mode 100644 index 1463abc..0000000 --- a/old/todo_files/open_compress_prompt.yaml +++ /dev/null @@ -1,43 +0,0 @@ -compress_system_prompt: | - You are a research assistant that has conducted research on a topic by calling several tools and web searches. Your job is now to clean up the findings, but preserve all of the relevant statements and information that the researcher has gathered. For context, today's date is {date}. - - - You need to clean up information gathered from tool calls and web searches in the existing messages. - All relevant information should be repeated and rewritten verbatim, but in a cleaner format. - The purpose of this step is just to remove any obviously irrelevant or duplicative information. - For example, if three sources all say "X", you could say "These three sources all stated X". - Only these fully comprehensive cleaned findings are going to be returned to the user, so it's crucial that you don't lose any information from the raw messages. - - - - 1. Your output findings should be fully comprehensive and include ALL of the information and sources that the researcher has gathered from tool calls and web searches. It is expected that you repeat key information verbatim. - 2. This report can be as long as necessary to return ALL of the information that the researcher has gathered. - 3. In your report, you should return inline citations for each source that the researcher found. - 4. You should include a "Sources" section at the end of the report that lists all of the sources the researcher found with corresponding citations, cited against statements in the report. - 5. Make sure to include ALL of the sources that the researcher gathered in the report, and how they were used to answer the question! - 6. It's really important not to lose any sources. A later LLM will be used to merge this report with others, so having all of the sources is critical. - 7. Do not change the language of tool call or web search results; maintain the same language as the original content. - - - - The report should be structured like this: - **List of Queries and Tool Calls Made** - **Fully Comprehensive Findings** - **List of All Relevant Sources (with citations in the report)** - - - - - Assign each unique URL a single citation number in your text - - End with ### Sources that lists each source with corresponding numbers - - IMPORTANT: Number sources sequentially without gaps (1,2,3,4...) in the final list regardless of which sources you choose - - Example format: - [1] Source Title: URL - [2] Source Title: URL - - - Critical Reminder: It is extremely important that any information that is even remotely relevant to the user's research topic is preserved verbatim (e.g. don't rewrite it, don't summarize it, don't paraphrase it). - - -compress_user_prompt: | - All above messages are about research conducted by an AI Researcher. Please clean up these findings, be careful not to include outdated information. - DO NOT summarize the information. I want the raw information returned, just in a cleaner format. Make sure all relevant information is preserved - you can rewrite findings verbatim. diff --git a/old/todo_files/react_v1_op.py b/old/todo_files/react_v1_op.py deleted file mode 100644 index 4b2250e..0000000 --- a/old/todo_files/react_v1_op.py +++ /dev/null @@ -1,109 +0,0 @@ -import datetime -import json -import time -from typing import List, Dict - -from loguru import logger - -from flowllm.context.flow_context import FlowContext -from flowllm.context.service_context import C -from flowllm.op.base_llm_op import BaseLLMOp -from flowllm.schema.flow_response import FlowResponse -from flowllm.schema.message import Message, Role - - -@C.register_op() -class ReactV1Op(BaseLLMOp): - file_path: str = __file__ - - def execute(self): - query: str = self.context.query - - max_steps: int = int(self.op_params.get("max_steps", 10)) - from flowllm.flow.base_tool_flow import BaseToolFlow - from flowllm.flow.gallery import DashscopeSearchToolFlow, CodeToolFlow, TerminateToolFlow - - tools: List[BaseToolFlow] = [DashscopeSearchToolFlow(), CodeToolFlow(), TerminateToolFlow()] - - """ - NOTE : x.tool_call.name != x.name - `x.tool_call.name` is tool's namex.name is flow's name(unique service name) - """ - tool_dict: Dict[str, BaseToolFlow] = {x.tool_call.name: x for x in tools} - for name, tool_call in tool_dict.items(): - logger.info(f"name={name} " - f"tool_call={json.dumps(tool_call.tool_call.simple_input_dump(), ensure_ascii=False)}") - - now_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') - has_terminate_tool = False - - user_prompt = self.prompt_format(prompt_name="role_prompt", - time=now_time, - tools=",".join(list(tool_dict.keys())), - query=query) - messages: List[Message] = [Message(role=Role.USER, content=user_prompt)] - logger.info(f"step.0 user_prompt={user_prompt}") - - for i in range(max_steps): - if has_terminate_tool: - assistant_message: Message = self.llm.chat(messages) - else: - assistant_message: Message = self.llm.chat(messages, tools=[x.tool_call for x in tools]) - - messages.append(assistant_message) - logger.info(f"assistant.{i}.reasoning_content={assistant_message.reasoning_content}\n" - f"content={assistant_message.content}\n" - f"tool.size={len(assistant_message.tool_calls)}") - - if has_terminate_tool: - break - - for tool in assistant_message.tool_calls: - if tool.name == "terminate": - has_terminate_tool = True - logger.info(f"step={i} find terminate tool, break.") - break - - if not has_terminate_tool and not assistant_message.tool_calls: - logger.warning(f"【bugfix】step={i} no tools, break.") - has_terminate_tool = True - - for j, tool_call in enumerate(assistant_message.tool_calls): - logger.info(f"submit step={i} tool_calls.name={tool_call.name} argument_dict={tool_call.argument_dict}") - - if tool_call.name not in tool_dict: - logger.warning(f"step={i} no tool_call.name={tool_call.name}") - continue - - self.submit_task(tool_dict[tool_call.name].__call__, **tool_call.argument_dict) - time.sleep(1) - - if not has_terminate_tool: - user_content_list = [] - for tool_result, tool_call in zip(self.join_task(), assistant_message.tool_calls): - logger.info(f"submit step={i} tool_calls.name={tool_call.name} tool_result={tool_result}") - if isinstance(tool_result, FlowResponse): - tool_result = tool_result.answer - else: - tool_result = str(tool_result) - user_content_list.append(f"\n{tool_result}\n") - user_content_list.append(self.prompt_format(prompt_name="next_prompt")) - assistant_message.tool_calls.clear() - messages.append(Message(role=Role.USER, content="\n".join(user_content_list))) - - else: - assistant_message.tool_calls.clear() - query = self.prompt_format(prompt_name="final_prompt", query=query) - messages.append(Message(role=Role.USER, content=query)) - - # Store results in context instead of response - self.context.response.messages = messages - self.context.response.answer = messages[-1].content - - -if __name__ == "__main__": - C.set_service_config().init_by_service_config() - context = FlowContext(query="茅台和五粮现在股价多少？") - - op = ReactV1Op() - op(context=context) diff --git a/old/todo_files/react_v1_prompt.yaml b/old/todo_files/react_v1_prompt.yaml deleted file mode 100644 index b2f5208..0000000 --- a/old/todo_files/react_v1_prompt.yaml +++ /dev/null @@ -1,54 +0,0 @@ -role_prompt: | - You are a helpful assistant. - The current time is {time}. - Please proactively choose the most suitable tool or combination of tools based on the user's question, including {tools} etc. - Please first think about how to break down the problem into subtasks, what tools and parameters should be used for each subtask, and finally provide the tool call name and parameters. - Try calling the same tool multiple times with different parameters to obtain information from various perspectives. - Please determine the response language based on the language of the user's question. - - {query} - -# write a complete and rigorous report to answer user's questions based on the context. -next_prompt: | - Think based on the current content and the user's question: Is the current context sufficient to answer the user's question? - - - If the current context is not sufficient to answer the user's question, consider what information is missing. - Re-plan and think about how to break down the missing information into subtasks. - For each subtask, determine what tools and parameters should be used for the query. - Please first provide the reasoning process, then give the tool call name and parameters. - - - If the current context is sufficient to answer the user's question, use the **terminate** tool. - -# Please determine the response language based on the language of the user's question. -final_prompt: | - Please integrate the context and provide a complete answer to the user's question. - - # User's Question - {query} - - - - -role_prompt_zh: | - 你是一个有用的助手。 - 当前时间是 {time}。 - 请根据用户的问题，主动选择最合适的工具或工具组合，包括 {tools} 等。 - 请先思考如何将问题分解为子任务，每个子任务应使用哪些工具和参数，最后提供工具调用名称和参数。 - 尝试多次使用相同的工具，但使用不同的参数，从多个角度获取信息。 - 请根据用户问题的语言来确定回复的语言。 - - {query} - -next_prompt_zh: | - 根据当前内容和用户的问题进行思考：当前上下文是否足以回答用户的问题？ - - 如果当前上下文不足以回答用户的问题，请考虑缺少哪些信息。 - 重新规划并思考如何将缺失的信息分解为子任务。 - 对于每个子任务，确定应使用哪些工具和参数进行查询。 - 请先提供推理过程，然后给出工具调用名称和参数。 - - 如果当前上下文足以回答用户的问题，请使用 **terminate** 工具。 - -final_prompt_zh: | - 请整合上下文，为用户的问题提供一个完整的答案。 - - # 用户的问题 - {query} diff --git a/old/todo_files/ths_op.py b/old/todo_files/ths_op.py deleted file mode 100644 index 9366840..0000000 --- a/old/todo_files/ths_op.py +++ /dev/null @@ -1,243 +0,0 @@ -import asyncio - -from crawl4ai import BrowserConfig, CrawlerRunConfig, CacheMode, AsyncWebCrawler -from loguru import logger - -from flowllm.context.flow_context import FlowContext -from flowllm.context.service_context import C -from flowllm.enumeration.role import Role -from flowllm.op.base_async_tool_op import BaseAsyncToolOp -from flowllm.schema.message import Message -from flowllm.schema.tool_call import ToolCall -from flowllm.utils.common_utils import get_datetime -from flowllm.utils.web_utils import get_random_user_agent - - -class ThsBaseOp(BaseAsyncToolOp): - file_path: str = __file__ - - def __init__(self, - url_key: str = "", - tool_description: str = "", - max_content_len: int = 30000, - llm: str = "qwen3_80b_instruct", - language: str = "zh", - enable_cache: bool = True, - cache_expire_hours: float = 1, - **kwargs): - - super().__init__(llm=llm, - language=language, - enable_cache=enable_cache, - cache_expire_hours=cache_expire_hours, - **kwargs) - - self.browser_config = BrowserConfig( - headless=True, - java_script_enabled=True, - user_agent=get_random_user_agent(), - viewport={"width": 1280, "height": 800}, - verbose=True) - - self.crawler_config = CrawlerRunConfig( - cache_mode=CacheMode.DISABLED, - page_timeout=9000, - verbose=True) - - self.url: str = f"https://basic.10jqka.com.cn/{{code}}/{url_key}.html#stockpage" - self.tool_description: str = tool_description - self.max_content_len: int = max_content_len - - def build_tool_call(self) -> ToolCall: - return ToolCall(**{ - "description": self.tool_description, - "input_schema": { - "code": { - "type": "string", - "description": "stock code", - "required": True - }, - } - }) - - async def async_execute(self): - code: str = self.input_dict["code"] - - if self.enable_cache: - cached_result = self.cache.load(code) - if cached_result: - self.set_result(cached_result["response_content"]) - return - - url = self.url.format(code=code) - - async with AsyncWebCrawler(config=self.browser_config) as crawler: - result = await crawler.arun( - url=url, - config=self.crawler_config, - # js_code="window.scrollTo(0, document.body.scrollHeight);", - # wait_for="document.querySelector('.loaded')" - ) - content = result.markdown[:self.max_content_len] - extract_content_prompt = self.prompt_format(prompt_name="extract_content_prompt", - date=get_datetime(), - code=code, - content=content) - assistant_message = await self.llm.achat(messages=[Message(role=Role.USER, content=extract_content_prompt)]) - self.set_result(assistant_message.content) - - final_result = { - "code": code, - "raw": result.markdown, - "response_content": assistant_message.content, - "model": self.llm.model_name, - } - - if self.enable_cache: - self.cache.save(code, final_result, expire_hours=self.cache_expire_hours) - - -@C.register_op(register_app="FlowLLM") -class ThsCompanyOp(ThsBaseOp): - - def __init__(self, **kwargs): - url_key: str = "company" - tool_description: str = "通过股票代码获取公司资料信息，例如：详细情况，高管介绍，发行相关，参控股公司。" - super().__init__(url_key=url_key, tool_description=tool_description, **kwargs) - - -@C.register_op(register_app="FlowLLM") -class ThsHolderOp(ThsBaseOp): - - def __init__(self, **kwargs): - url_key: str = "holder" - tool_description: str = "通过股票代码获取股东研究信息，例如：股东人数、十大流通股东、十大股东、十大债券持有人、控股层级关系。" - super().__init__(url_key=url_key, tool_description=tool_description, **kwargs) - - -@C.register_op(register_app="FlowLLM") -class ThsOperateOp(ThsBaseOp): - - def __init__(self, **kwargs): - url_key: str = "operate" - tool_description: str = "通过股票代码获取经营分析信息，例如：主营介绍、运营业务数据、主营构成分析、主要客户及供应商、董事会经营评述、产品价格。" - super().__init__(url_key=url_key, tool_description=tool_description, **kwargs) - - -@C.register_op(register_app="FlowLLM") -class ThsEquityOp(ThsBaseOp): - - def __init__(self, **kwargs): - url_key: str = "equity" - tool_description: str = "通过股票代码获取股本结构信息，例如：解禁时间表、总股本构成、A股结构图、历次股本变动" - super().__init__(url_key=url_key, tool_description=tool_description, **kwargs) - - -@C.register_op(register_app="FlowLLM") -class ThsCapitalOp(ThsBaseOp): - - def __init__(self, **kwargs): - url_key: str = "capital" - tool_description: str = "通过股票代码获取资本运作信息，例如：募集资金来源、项目投资、收购兼并、股权投资、参股IPO、股权转让、关联交易、质押解冻。" - super().__init__(url_key=url_key, tool_description=tool_description, **kwargs) - - -@C.register_op(register_app="FlowLLM") -class ThsWorthOp(ThsBaseOp): - - def __init__(self, **kwargs): - url_key: str = "worth" - tool_description: str = "通过股票代码获取盈利预测信息，例如：业绩预测、业绩预测详表、研报评级" - super().__init__(url_key=url_key, tool_description=tool_description, **kwargs) - - -@C.register_op(register_app="FlowLLM") -class ThsNewsOp(ThsBaseOp): - - def __init__(self, **kwargs): - url_key: str = "news" - tool_description: str = "通过股票代码获取新闻公告信息，例如：新闻与股价联动、公告列表、热点新闻列表、研报列表" - super().__init__(url_key=url_key, tool_description=tool_description, **kwargs) - - -@C.register_op(register_app="FlowLLM") -class ThsConceptOp(ThsBaseOp): - - def __init__(self, **kwargs): - url_key: str = "concept" - tool_description: str = "通过股票代码获取概念题材信息，例如：常规概念、其他概念、题材要点、概念对比" - super().__init__(url_key=url_key, tool_description=tool_description, **kwargs) - - -@C.register_op(register_app="FlowLLM") -class ThsPositionOp(ThsBaseOp): - - def __init__(self, **kwargs): - url_key: str = "position" - tool_description: str = "通过股票代码获取主力持仓信息，例如：机构持股汇总、机构持股明细、被举牌情况、IPO获配机构" - super().__init__(url_key=url_key, tool_description=tool_description, **kwargs) - - -@C.register_op(register_app="FlowLLM") -class ThsFinanceOp(ThsBaseOp): - - def __init__(self, **kwargs): - url_key: str = "finance" - tool_description: str = "通过股票代码获取财务分析信息，例如：财务诊断、财务指标、指标变动说明、资产负债构成、财务报告、杜邦分析" - super().__init__(url_key=url_key, tool_description=tool_description, **kwargs) - - -@C.register_op(register_app="FlowLLM") -class ThsBonusOp(ThsBaseOp): - - def __init__(self, **kwargs): - url_key: str = "bonus" - tool_description: str = "通过股票代码获取分红融资信息，例如：分红诊断、分红情况、增发机构获配明细、增发概况、配股概况" - super().__init__(url_key=url_key, tool_description=tool_description, **kwargs) - - -@C.register_op(register_app="FlowLLM") -class ThsEventOp(ThsBaseOp): - - def __init__(self, **kwargs): - url_key: str = "event" - tool_description: str = "通过股票代码获取公司大事信息，例如：高管持股变动、股东持股变动、担保明细、违规处理、机构调研、投资者互动" - super().__init__(url_key=url_key, tool_description=tool_description, **kwargs) - - -@C.register_op(register_app="FlowLLM") -class ThsFieldOp(ThsBaseOp): - - def __init__(self, **kwargs): - url_key: str = "field" - tool_description: str = "通过股票代码获取行业对比信息，例如：行业地位、行业新闻" - super().__init__(url_key=url_key, tool_description=tool_description, **kwargs) - -async def main(): - from flowllm.app import FlowLLMApp - async with FlowLLMApp(load_default_config=True): - code = "601899" - context = FlowContext(code=code) - - ops = [ - ThsCompanyOp(), - ThsHolderOp(), - ThsOperateOp(), - ThsEquityOp(), - ThsCapitalOp(), - ThsWorthOp(), - ThsNewsOp(), - ThsConceptOp(), - ThsPositionOp(), - ThsFinanceOp(), - ThsBonusOp(), - ThsEventOp(), - ThsFieldOp() - ] - for op in ops: - await op.async_call(context=context) - logger.info(op.output) - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/old/todo_files/ths_prompt.yaml b/old/todo_files/ths_prompt.yaml deleted file mode 100644 index 15af757..0000000 --- a/old/todo_files/ths_prompt.yaml +++ /dev/null @@ -1,16 +0,0 @@ -extract_content_prompt_zh: | - 请根据以下信息，从提供的网络爬取内容中严谨、完整地整理出与指定股票代码相关的内容： - - 当前时间 - {date} - - 股票代码 - {code} - - 原始内容 - {content} - - 请严格遵守以下要求： - - 所有整理出的信息必须直接来源于原始的同花顺网页内容，不得添加、推测或虚构任何信息（包括数字、日期、百分比等），严禁“幻觉”内容。 - - 所有数据（如财务数字、涨跌幅、市值、时间点等）必须能在原文中找到对应出处，若原文未明确提及，则不得写入。 - - 最终输出格式必须为 Markdown，且最大标题层级为三级标题（###），结构清晰、条理分明。 diff --git a/old/todo_files/ths_supply_op.py b/old/todo_files/ths_supply_op.py deleted file mode 100644 index d78ab14..0000000 --- a/old/todo_files/ths_supply_op.py +++ /dev/null @@ -1,7 +0,0 @@ -from flowllm import BaseLLMOp - - -class ThsSupplyOp(BaseLLMOp): - - async def async_execute(self): - ... \ No newline at end of file diff --git a/old/todo_files/tool_op_flow.py b/old/todo_files/tool_op_flow.py deleted file mode 100644 index 738f3b9..0000000 --- a/old/todo_files/tool_op_flow.py +++ /dev/null @@ -1,110 +0,0 @@ -from flowllm.op.gallery import ExecuteCodeOp -from flowllm.op.llm import SimpleLLMOp, ReactLLMOp, StreamLLMOp -from flowllm.op.search import DashscopeSearchOp -from flowllm.op.search import TavilySearchOp - -from flowllm.context.service_context import C -from flowllm.flow.base_tool_flow import BaseToolFlow -from flowllm.schema.tool_call import ToolCall - - -@C.register_flow() -class TavilySearchToolFlow(BaseToolFlow): - - def __init__(self, - enable_async: bool = True, - stream: bool = False, - service_type: str = "http+mcp", - **kwargs): - super().__init__(use_async=enable_async, stream=stream, service_type=service_type, **kwargs) - - def build_flow(self): - return TavilySearchOp(save_answer=True) - - def build_tool_call(self) -> ToolCall: - return self.flow_op.tool_call - - -@C.register_flow() -class DashscopeSearchToolFlow(BaseToolFlow): - - def __init__(self, - enable_async: bool = True, - stream: bool = False, - service_type: str = "http+mcp", - **kwargs): - super().__init__(use_async=enable_async, stream=stream, service_type=service_type, **kwargs) - - def build_flow(self): - return DashscopeSearchOp(save_answer=True) - - def build_tool_call(self) -> ToolCall: - return self.flow_op.tool_call - - -@C.register_flow() -class StreamLLMToolFlow(BaseToolFlow): - - def __init__(self, - enable_async: bool = True, - stream: bool = True, - service_type: str = "http", - **kwargs): - super().__init__(use_async=enable_async, stream=stream, service_type=service_type, **kwargs) - - def build_flow(self): - return StreamLLMOp() - - def build_tool_call(self) -> ToolCall: - return self.flow_op.tool_call - - -@C.register_flow() -class SimpleLLMToolFlow(BaseToolFlow): - - def __init__(self, - enable_async: bool = True, - stream: bool = False, - service_type: str = "http", - **kwargs): - super().__init__(use_async=enable_async, stream=stream, service_type=service_type, **kwargs) - - def build_flow(self): - return SimpleLLMOp(save_answer=True) - - def build_tool_call(self) -> ToolCall: - return self.flow_op.tool_call - - -@C.register_flow() -class ReactLLMToolFlow(BaseToolFlow): - - def __init__(self, - enable_async: bool = True, - stream: bool = False, - service_type: str = "http", - **kwargs): - super().__init__(use_async=enable_async, stream=stream, service_type=service_type, **kwargs) - - def build_flow(self): - return ReactLLMOp() - - def build_tool_call(self) -> ToolCall: - return self.flow_op.tool_call - - -@C.register_flow() -class CodeExecutionFlow(BaseToolFlow): - - def __init__(self, - enable_async: bool = True, - stream: bool = False, - service_type: str = "http+mcp", - **kwargs): - super().__init__(use_async=enable_async, stream=stream, service_type=service_type, **kwargs) - - def build_flow(self): - return ExecuteCodeOp(save_answer=True) - - def build_tool_call(self) -> ToolCall: - return self.flow_op.tool_call diff --git a/old/utils/__init__.py b/old/utils/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/old/utils/common_utils.py b/old/utils/common_utils.py deleted file mode 100644 index 9f7b8e0..0000000 --- a/old/utils/common_utils.py +++ /dev/null @@ -1,193 +0,0 @@ -import json -import os -import re -from datetime import datetime, timedelta -from pathlib import Path -from typing import List - -from loguru import logger - -ENV_LOADED = False - - -def camel_to_snake(content: str) -> str: - """ - BaseWorker -> base_worker - """ - # FIXME - content = content.replace("LLM", "Llm") - - snake_str = re.sub(r"(? str: - """ - base_worker -> BaseWorker - """ - camel_str = "".join(x.capitalize() for x in content.split("_")) - - # FIXME - camel_str = camel_str.replace("Llm", "LLM") - return camel_str - - -def _load_env(path: Path): - with path.open() as f: - for line in f: - line = line.strip() - if line.startswith("#"): - continue - - line_split = line.strip().split("=", 1) - if len(line_split) >= 2: - key = line_split[0].strip() - value = line_split[1].strip().strip('"') - os.environ[key] = value - - -def load_env(path: str | Path = None, enable_log: bool = True): - global ENV_LOADED - if ENV_LOADED: - return - - if path is not None: - path = Path(path) - if path.exists(): - _load_env(path) - ENV_LOADED = True - - else: - for i in range(5): - path = Path("../" * i + ".env") - if path.exists(): - if enable_log: - logger.info(f"load env_path={path}") - _load_env(path) - ENV_LOADED = True - return - - logger.warning(".env not found") - - -def extract_content(text: str, language_tag: str = "json"): - pattern = rf"```{re.escape(language_tag)}\s*(.*?)\s*```" - match = re.search(pattern, text, re.DOTALL) - - if match: - result = match.group(1).strip() - else: - result = text - - if language_tag == "json": - try: - result = json.loads(result) - - except json.JSONDecodeError: - result = None - - return result - - -def get_datetime(time_ft: str = "%Y-%m-%d %H:%M:%S"): - now = datetime.now() - formatted_time = now.strftime(time_ft) - return formatted_time - - -def get_monday_fridays(start_str: str, end_str: str): - start = datetime.strptime(str(start_str), "%Y%m%d") - end = datetime.strptime(str(end_str), "%Y%m%d") - if start > end: - return [] - - current = start - while current.weekday() != 4: - current += timedelta(days=1) - if current > end: - return [] - - result = [] - while current <= end: - result.append( - [ - (current - timedelta(days=4)).strftime("%Y%m%d"), - current.strftime("%Y%m%d"), - ] - ) - current += timedelta(days=7) - - return result - - -def next_friday_or_same(date_str): - dt = datetime.strptime(date_str, "%Y%m%d") - days_ahead = (4 - dt.weekday()) % 7 - next_fri = dt + timedelta(days=days_ahead) - return next_fri.strftime("%Y%m%d") - - -def find_dt_less_index(dt: str | int, dt_list: List[str | int]): - """ - Use binary search to find the index of the date that is closest to and less than dt. - Time complexity: O(log n) - """ - if not dt_list: - return None - - left, right = 0, len(dt_list) - 1 - - if dt < dt_list[left]: - return None - - if dt >= dt_list[right]: - return right - - while left < right: - mid = (left + right + 1) // 2 - if dt_list[mid] <= dt: - left = mid - else: - right = mid - 1 - - return left - - -from typing import List, Optional - - -def find_dt_greater_index(dt: str, dt_list: List[str]) -> Optional[int]: - """ - Use binary search to find the index of the date that is closest to and greater than dt. - Time complexity: O(log n) - - Args: - dt: Target date string (e.g., '2023-05-15') - dt_list: Sorted list of date strings in ascending order - - Returns: - Index of the first date in dt_list that is strictly greater than dt, - or None if no such date exists. - """ - if not dt_list: - return None - - left, right = 0, len(dt_list) - 1 - - # If dt is >= the last element, no greater element exists - if dt >= dt_list[right]: - return None - - # If dt is < the first element, the first element is the answer - if dt < dt_list[left]: - return left - - # Binary search for the first element > dt - while left < right: - mid = (left + right) // 2 - if dt_list[mid] <= dt: - left = mid + 1 - else: - right = mid - - return left diff --git a/old/utils/eastmoney_utils.py b/old/utils/eastmoney_utils.py deleted file mode 100644 index 83a1f96..0000000 --- a/old/utils/eastmoney_utils.py +++ /dev/null @@ -1,308 +0,0 @@ -import math -import os - -import akshare as ak -import pandas as pd -import requests - - -def get_ah_mapping() -> pd.DataFrame: - EAST_MONEY_COOKIE = os.getenv("EAST_MONEY_COOKIE") - url = "https://push2.eastmoney.com/api/qt/clist/get" - params = { - "np": "1", - "fltt": "1", - "invt": "2", - "fs": "b:DLMK0101", - "fields": "f193,f191,f192,f12,f13,f14,f1,f2,f4,f3,f152,f186,f190,f187,f189,f188", - "fid": "f3", - "pn": "1", - "pz": "100", - "po": "1", - "dect": "1", - "wbp2u": "|0|0|0|web", - } - headers = {} - if EAST_MONEY_COOKIE: - headers["Cookie"] = EAST_MONEY_COOKIE - temp_df = fetch_paginated_data(url, base_params=params, headers=headers) - - columns_mapping = { - "f193": "name", - "f12": "hk_code", - "f2": "hk_price", - "f3": "hk_pct_chg", - "f191": "a_code", - "f186": "a_price", - "f187": "a_pct_chg", - "f189": "ah_price_ratio", - "f188": "ah_premium_ratio", - } - - temp_df = temp_df.rename(columns=columns_mapping).loc[:, list(columns_mapping.values())] - temp_df["hk_price"] = pd.to_numeric(temp_df["hk_price"], errors="coerce") / 1000 - temp_df["hk_pct_chg"] = pd.to_numeric(temp_df["hk_pct_chg"], errors="coerce") / 100 - temp_df["a_price"] = pd.to_numeric(temp_df["a_price"], errors="coerce") / 100 - temp_df["a_pct_chg"] = pd.to_numeric(temp_df["a_pct_chg"], errors="coerce") / 100 - temp_df["ah_price_ratio"] = pd.to_numeric(temp_df["ah_price_ratio"], errors="coerce") / 100 - temp_df["ah_premium_ratio"] = pd.to_numeric(temp_df["ah_premium_ratio"], errors="coerce") / 100 - return temp_df - - -def fetch_paginated_data(url: str, base_params: dict, headers: dict = None, timeout: int = 15) -> pd.DataFrame: - params = base_params.copy() - response = requests.get(url, params=params, headers=headers, timeout=timeout) - data_json = response.json() - items_per_page = len(data_json["data"]["diff"]) - total_pages = math.ceil(data_json["data"]["total"] / items_per_page) - page_dataframes = [pd.DataFrame(data_json["data"]["diff"])] - - # Get remaining page data - for page in range(2, total_pages + 1): - params.update({"pn": page}) - response = requests.get(url, params=params, headers=headers, timeout=timeout) - data_json = response.json() - page_dataframe = pd.DataFrame(data_json["data"]["diff"]) - page_dataframes.append(page_dataframe) - - # Merge all data - merged_dataframe = pd.concat(page_dataframes, ignore_index=True) - merged_dataframe["f3"] = pd.to_numeric(merged_dataframe["f3"], errors="coerce") - merged_dataframe.sort_values(by=["f3"], ascending=False, inplace=True, ignore_index=True) - return merged_dataframe - - -def get_a_stock_df( - code: str, - period: str = "daily", - start_date: str = "20160101", - end_date: str = "20300101", - adjust: str = "hfq", - timeout: float = None, -) -> pd.DataFrame: - EAST_MONEY_COOKIE = os.getenv("EAST_MONEY_COOKIE") - - market_code = 1 if code.startswith("6") else 0 - adjust_dict = {"qfq": "1", "hfq": "2", "": "0"} - period_dict = {"daily": "101", "weekly": "102", "monthly": "103"} - url = "https://push2his.eastmoney.com/api/qt/stock/kline/get" - params = { - "fields1": "f1,f2,f3,f4,f5,f6", - "fields2": "f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61,f116", - "ut": "7eea3edcaed734bea9cbfc24409ed989", - "klt": period_dict[period], - "fqt": adjust_dict[adjust], - "secid": f"{market_code}.{code}", - "beg": start_date, - "end": end_date, - } - headers = {} - if EAST_MONEY_COOKIE: - headers["Cookie"] = EAST_MONEY_COOKIE - r = requests.get(url, params=params, headers=headers, timeout=timeout) - data_json = r.json() - if not (data_json["data"] and data_json["data"]["klines"]): - return pd.DataFrame() - - a_hist_df = pd.DataFrame([item.split(",") for item in data_json["data"]["klines"]]) - a_hist_df["code"] = code - a_hist_df.columns = [ - "date", - "open", - "close", - "high", - "low", - "vol", - "amount", - "swing", - "chg_pct", - "change", - "turnover_ratio", - "code", - ] - a_hist_df["date"] = pd.to_datetime(a_hist_df["date"], errors="coerce").dt.date - a_hist_df["open"] = pd.to_numeric(a_hist_df["open"], errors="coerce") - a_hist_df["close"] = pd.to_numeric(a_hist_df["close"], errors="coerce") - a_hist_df["high"] = pd.to_numeric(a_hist_df["high"], errors="coerce") - a_hist_df["low"] = pd.to_numeric(a_hist_df["low"], errors="coerce") - a_hist_df["vol"] = pd.to_numeric(a_hist_df["vol"], errors="coerce") - a_hist_df["amount"] = pd.to_numeric(a_hist_df["amount"], errors="coerce") - a_hist_df["swing"] = pd.to_numeric(a_hist_df["swing"], errors="coerce") - a_hist_df["chg_pct"] = pd.to_numeric(a_hist_df["chg_pct"], errors="coerce") - a_hist_df["change"] = pd.to_numeric(a_hist_df["change"], errors="coerce") - a_hist_df["turnover_ratio"] = pd.to_numeric(a_hist_df["turnover_ratio"], errors="coerce") - return a_hist_df - - -def get_hk_stock_df( - code: str, - period: str = "daily", - start_date: str = "20160101", - end_date: str = "20300101", - adjust: str = "hfq", - timeout: float = None, -) -> pd.DataFrame: - EAST_MONEY_COOKIE = os.getenv("EAST_MONEY_COOKIE") - - adjust_dict = {"qfq": "1", "hfq": "2", "": "0"} - period_dict = {"daily": "101", "weekly": "102", "monthly": "103"} - url = "https://33.push2his.eastmoney.com/api/qt/stock/kline/get" - params = { - "fields1": "f1,f2,f3,f4,f5,f6", - "fields2": "f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61", - "klt": period_dict[period], - "fqt": adjust_dict[adjust], - "secid": f"116.{code}", - "beg": start_date, - "end": end_date, - } - headers = {} - if EAST_MONEY_COOKIE: - headers["Cookie"] = EAST_MONEY_COOKIE - r = requests.get(url, params=params, headers=headers, timeout=timeout) - data_json = r.json() - if not (data_json["data"] and data_json["data"]["klines"]): - return pd.DataFrame() - - hk_hist_df = pd.DataFrame([item.split(",") for item in data_json["data"]["klines"]]) - hk_hist_df["code"] = code - hk_hist_df.columns = [ - "date", - "open", - "close", - "high", - "low", - "vol", - "amount", - "swing", - "chg_pct", - "change", - "turnover_ratio", - "code", - ] - - hk_hist_df["date"] = pd.to_datetime(hk_hist_df["date"], errors="coerce").dt.date - hk_hist_df["open"] = pd.to_numeric(hk_hist_df["open"], errors="coerce") - hk_hist_df["close"] = pd.to_numeric(hk_hist_df["close"], errors="coerce") - hk_hist_df["high"] = pd.to_numeric(hk_hist_df["high"], errors="coerce") - hk_hist_df["low"] = pd.to_numeric(hk_hist_df["low"], errors="coerce") - hk_hist_df["vol"] = pd.to_numeric(hk_hist_df["vol"], errors="coerce") - hk_hist_df["amount"] = pd.to_numeric(hk_hist_df["amount"], errors="coerce") - hk_hist_df["swing"] = pd.to_numeric(hk_hist_df["swing"], errors="coerce") - hk_hist_df["chg_pct"] = pd.to_numeric(hk_hist_df["chg_pct"], errors="coerce") - hk_hist_df["change"] = pd.to_numeric(hk_hist_df["change"], errors="coerce") - hk_hist_df["turnover_ratio"] = pd.to_numeric(hk_hist_df["turnover_ratio"], errors="coerce") - return hk_hist_df - - -def get_forex_df( - code: str = "HKDCNYC", - start_date: str = "20160101", - end_date: str = "20300101", - timeout: float = None, -) -> pd.DataFrame: - EAST_MONEY_COOKIE = os.getenv("EAST_MONEY_COOKIE") - code_market_map = { - "HKDCNYC": 120, - "USDCNYC": 120, - } - - url = "https://push2his.eastmoney.com/api/qt/stock/kline/get" - market_code = code_market_map[code] - params = { - "secid": f"{market_code}.{code}", - "klt": "101", - "fqt": "1", - "beg": start_date, - "end": end_date, - "iscca": "1", - "fields1": "f1,f2,f3,f4,f5,f6,f7,f8", - "fields2": "f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61,f62,f63,f64", - "ut": "f057cbcbce2a86e2866ab8877db1d059", - "forcect": 1, - } - headers = {} - if EAST_MONEY_COOKIE: - headers["Cookie"] = EAST_MONEY_COOKIE - r = requests.get(url, params=params, headers=headers, timeout=timeout) - data_json = r.json() - if not (data_json["data"] and data_json["data"]["klines"]): - return pd.DataFrame() - - forex_df = pd.DataFrame([item.split(",") for item in data_json["data"]["klines"]]) - forex_df["code"] = data_json["data"]["code"] - forex_df["name"] = data_json["data"]["name"] - forex_df.columns = [ - "date", - "open", - "close", - "high", - "low", - "-", - "-", - "swing", - "-", - "-", - "-", - "-", - "-", - "-", - "code", - "name", - ] - forex_df = forex_df[ - [ - "date", - "code", - "name", - "open", - "close", - "high", - "low", - "swing", - ] - ] - forex_df["date"] = pd.to_datetime(forex_df["date"], errors="coerce").dt.date - forex_df["open"] = pd.to_numeric(forex_df["open"], errors="coerce") - forex_df["close"] = pd.to_numeric(forex_df["close"], errors="coerce") - forex_df["high"] = pd.to_numeric(forex_df["high"], errors="coerce") - forex_df["low"] = pd.to_numeric(forex_df["low"], errors="coerce") - forex_df["swing"] = pd.to_numeric(forex_df["swing"], errors="coerce") - return forex_df - - -def get_history_dividend_detail(code: str): - result = {} - df = ak.stock_history_dividend_detail(symbol=code, indicator="分红") - if df is not None and len(df) > 0: - result["fh_anno"] = [str(x) for x in df.loc[:, "公告日期"].tolist()] - result["fh_prog"] = [str(x) for x in df.loc[:, "除权除息日"].tolist()] - return result - - -def main(): - # print(os.getenv("EAST_MONEY_COOKIE")) - # - # ah_mapping_df = get_ah_mapping() - # print(ah_mapping_df) - # - # a_stock_df = get_a_stock_df("000001") - # print(a_stock_df) - # - # hk_stock_df = get_hk_stock_df("01810") - # print(hk_stock_df) - # - # forex_hist_em_df = get_forex_df(code="HKDCNYC") - # print(forex_hist_em_df) - - # df = get_history_dividend_detail("300750") - df = get_history_dividend_detail("000001") - print(df) - - -if __name__ == "__main__": - from flowllm.utils.common_utils import load_env - - load_env() - - main() diff --git a/old/utils/plot_utils.py b/old/utils/plot_utils.py deleted file mode 100644 index 510fd1b..0000000 --- a/old/utils/plot_utils.py +++ /dev/null @@ -1,62 +0,0 @@ -from pathlib import Path -from typing import Dict - -from matplotlib import pyplot as plt - - -def calculate_max_drawdown(uplift: list): - """ - A list of uplift values following a pure -1. - """ - uplift_list = [x + 1 for x in uplift] - cumulative_max = uplift_list[0] - max_drawdown = 0 - - for value in uplift_list: - cumulative_max = max(cumulative_max, value) - drawdown = value / cumulative_max - 1 - max_drawdown = min(max_drawdown, drawdown) - return max_drawdown * 100 - - -def plot_figure( - plot_dict: Dict[str, list], - output_path: str | Path, - xs: list = None, - flag: bool = False, - ticks_gap: int = None, - enable_drawdown: bool = True, - extra_content: str = "", -): - dt_len = len(list(plot_dict.values())[0]) - plt.figure(figsize=(12, 6)) - - if xs is None: - xs = [i for i in range(dt_len)] - - for key, ys in plot_dict.items(): - if flag: - new_ys = ys - else: - new_ys = [] - for i in range(len(ys)): - new_ys.append(sum(ys[: i + 1])) - - label = f"{key}@{new_ys[-1]:.2f}" - if enable_drawdown: - label += f"@{calculate_max_drawdown(new_ys):.2f}%" - plt.plot(xs, new_ys, label=label) - - if ticks_gap is not None: - plt.xticks(ticks=range(0, len(xs), ticks_gap), labels=xs[::ticks_gap], rotation=45) - - plt.xlabel("x") - plt.ylabel("y") - plt.grid() - plt.legend() - - if extra_content: - plt.text(x=-0.5, y=-1, s=extra_content, fontsize=8) - - plt.savefig(output_path) - plt.close() diff --git a/old/utils/ridge_v2.py b/old/utils/ridge_v2.py deleted file mode 100644 index ee11a91..0000000 --- a/old/utils/ridge_v2.py +++ /dev/null @@ -1,56 +0,0 @@ -import numpy as np -from sklearn.linear_model import Ridge -from sklearn.preprocessing import StandardScaler - - -class RidgeV2: - - def __init__(self, bound: int = 4.6, use_ridge_v2=True, **kwargs): - self.bound: int = bound - self.use_ridge_v2: bool = use_ridge_v2 - self.kwargs: dict = kwargs - self.model = Ridge(**self.kwargs) - - def clear(self): - self.model = Ridge(**self.kwargs) - return self - - def sigmoid(self, x): - x = np.asarray(x, dtype=float) - x = np.clip(x, -self.bound, self.bound) - return 1 / (1 + np.exp(-x)) - - def inv_sigmoid(self, p): - p = np.asarray(p, dtype=float) - p = np.clip(p, self.sigmoid(-self.bound), self.sigmoid(self.bound)) - return np.log(p / (1 - p)) - - def fit(self, x, y, sample_weight=None): - if self.use_ridge_v2: - return self.model.fit(x, self.inv_sigmoid(y), sample_weight=sample_weight) - else: - return self.model.fit(x, y, sample_weight=sample_weight) - - def predict(self, x): - if self.use_ridge_v2: - return self.sigmoid(self.model.predict(x)) - else: - return self.model.predict(x) - - def fit_and_predict( - self, - train_x_nd: np.ndarray, - train_y_nd: np.ndarray, - test_x_nd: np.ndarray, - check_y: bool = True, - ): - if check_y: - assert np.all((train_y_nd >= 0) & (train_y_nd <= 1)) - - scaler = StandardScaler() - scaler.fit(train_x_nd) - train_x_nd = scaler.transform(train_x_nd) - test_x_nd = scaler.transform(test_x_nd) - self.model.fit(train_x_nd, train_y_nd) - pred_y_nd = self.model.predict(test_x_nd) - return np.minimum(np.maximum(pred_y_nd, 0), 1) diff --git a/old/utils/singleton.py b/old/utils/singleton.py deleted file mode 100644 index df7a1b3..0000000 --- a/old/utils/singleton.py +++ /dev/null @@ -1,9 +0,0 @@ -def singleton(cls): - _instance = {} - - def _singleton(*args, **kwargs): - if cls not in _instance: - _instance[cls] = cls(*args, **kwargs) - return _instance[cls] - - return _singleton diff --git a/old/utils/token_utils.py b/old/utils/token_utils.py deleted file mode 100644 index fd77c9c..0000000 --- a/old/utils/token_utils.py +++ /dev/null @@ -1,182 +0,0 @@ -from typing import TYPE_CHECKING, Any - -from loguru import logger - -if TYPE_CHECKING: - import tiktoken - from modelscope import AutoTokenizer as ModelScopeTokenizer - -# Lazy import flags -_TIKTOKEN_CHECKED = False -_TIKTOKEN_AVAILABLE = False -_MODELSCOPE_CHECKED = False -_MODELSCOPE_AVAILABLE = False - - -def _check_tiktoken() -> bool: - """Lazy check for tiktoken availability.""" - global _TIKTOKEN_CHECKED, _TIKTOKEN_AVAILABLE - if not _TIKTOKEN_CHECKED: - try: - import tiktoken # noqa: F401 - - _TIKTOKEN_AVAILABLE = True - except ImportError: - _TIKTOKEN_AVAILABLE = False - logger.debug("tiktoken not installed. OpenAI models will use character-based estimation.") - _TIKTOKEN_CHECKED = True - return _TIKTOKEN_AVAILABLE - - -def _check_modelscope() -> bool: - """Lazy check for modelscope availability.""" - global _MODELSCOPE_CHECKED, _MODELSCOPE_AVAILABLE - if not _MODELSCOPE_CHECKED: - try: - from modelscope import AutoTokenizer # noqa: F401 - - _MODELSCOPE_AVAILABLE = True - except ImportError: - _MODELSCOPE_AVAILABLE = False - logger.debug("modelscope not installed. Qwen models will use character-based estimation.") - _MODELSCOPE_CHECKED = True - return _MODELSCOPE_AVAILABLE - - -class TokenCounter: - - def __init__(self, default_type: str = "qwen", default_model: str = None): - self._tokenizer_cache = {} - - # Validate and set default type - valid_types = ("auto", "openai", "qwen", "char") - if default_type not in valid_types: - raise ValueError(f"default_type must be one of {valid_types}, got: {default_type}") - self.default_type = default_type - - # Set default model based on type - if default_model is None: - if default_type == "openai": - self.default_model = "gpt-4o" # o200k_base encoding - elif default_type == "qwen": - self.default_model = "Qwen/Qwen3-32B" - else: - self.default_model = "Qwen/Qwen3-32B" # fallback default - else: - self.default_model = default_model - - def count(self, text: str, model_name: str = None) -> int: - if not text: - return 0 - - # Use default model if not specified - if model_name is None: - model_name = self.default_model - - # Force specific tokenizer type if set - if self.default_type == "openai": - return self._count_with_tiktoken(text, model_name) - elif self.default_type == "qwen": - return self._count_with_qwen_tokenizer(text, model_name) - elif self.default_type == "char": - return self._estimate_from_chars(text) - - # Auto mode: detect from model name - model_lower = model_name.lower() - - # Route 1: GPT models → tiktoken - if model_lower.startswith("gpt"): - return self._count_with_tiktoken(text, model_name) - - # Route 2: Qwen models → ModelScope - # Detect by "qwen" prefix or "/" (ModelScope path format) - if model_lower.startswith("qwen") or "/" in model_name: - return self._count_with_qwen_tokenizer(text, model_name) - - # Route 3: Fallback → character estimation - return self._estimate_from_chars(text) - - def _count_with_tiktoken(self, text: str, model_name: str) -> int: - """Count tokens using tiktoken (for GPT models).""" - if not _check_tiktoken(): - logger.debug(f"tiktoken not available, using fallback for {model_name}") - return self._estimate_from_chars(text) - - try: - import tiktoken - - # Get encoding for the model - if model_name.startswith("gpt-4o"): - encoding_name = "o200k_base" - elif model_name.startswith("gpt-4") or model_name.startswith("gpt-3.5"): - encoding_name = "cl100k_base" - else: - encoding_name = "cl100k_base" # default - - # Cache encoding - if encoding_name not in self._tokenizer_cache: - self._tokenizer_cache[encoding_name] = tiktoken.get_encoding(encoding_name) - - encoding = self._tokenizer_cache[encoding_name] - return len(encoding.encode(text)) - - except Exception as e: - logger.debug(f"tiktoken failed for {model_name}: {e}, using fallback") - return self._estimate_from_chars(text) - - def _count_with_qwen_tokenizer(self, text: str, model_name: str) -> int: - """ - Count tokens using ModelScope tokenizer (for Qwen models). - - Expects a direct ModelScope path: "qwen/Qwen2.5-72B-Instruct" - """ - if not _check_modelscope(): - logger.debug(f"modelscope not available, using fallback for {model_name}") - return self._estimate_from_chars(text) - - try: - from modelscope import AutoTokenizer as ModelScopeTokenizer - - # Use model_name directly as ModelScope path - ms_model = model_name - - # Load and cache tokenizer - if ms_model not in self._tokenizer_cache: - logger.info(f"Loading Qwen tokenizer: {ms_model} (only ~2-5MB)") - tokenizer = ModelScopeTokenizer.from_pretrained(ms_model, trust_remote_code=True) - self._tokenizer_cache[ms_model] = tokenizer - - tokenizer = self._tokenizer_cache[ms_model] - return len(tokenizer.encode(text)) - - except Exception as e: - logger.warning( - f"Failed to load tokenizer for '{model_name}': {e}\n" - f"Ensure the ModelScope path is correct.\n" - f"Example: 'qwen/Qwen2.5-72B-Instruct' or 'qwen/Qwen-7B-Chat'\n" - f"Falling back to character estimation.", - ) - return self._estimate_from_chars(text) - - @staticmethod - def _estimate_from_chars(text: str) -> int: - """ - Fallback: Estimate tokens from character count. - Simple heuristic: tokens ≈ characters / 4 - """ - if not text: - return 0 - - # Slightly better heuristic for Chinese - chinese_chars = sum(1 for c in text if "\u4e00" <= c <= "\u9fff") - other_chars = len(text) - chinese_chars - - # Chinese: ~2 chars/token, English: ~4 chars/token - estimated = (chinese_chars // 2) + (other_chars // 4) - return max(1, estimated) - - -if __name__ == "__main__": - counter = TokenCounter() - token_count = counter.count("我爱吃苹果") - print(f"Token count for qwen: {token_count}") diff --git a/old/utils/web_utils.py b/old/utils/web_utils.py deleted file mode 100644 index a0a1548..0000000 --- a/old/utils/web_utils.py +++ /dev/null @@ -1,119 +0,0 @@ -import random -from typing import List - - -def get_random_user_agent() -> str: - if random.choice([True, False]): - return _get_predefined_user_agent() - else: - return _generate_dynamic_user_agent() - - -def _get_predefined_user_agent() -> str: - """Get a user agent from predefined list.""" - user_agents: List[str] = [ - # Chrome on Windows - "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36", - "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36", - "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/118.0.0.0 Safari/537.36", - # Chrome on macOS - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36", - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_6) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36", - # Chrome on Linux - "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36", - "Mozilla/5.0 (X11; Ubuntu; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36", - # Firefox on Windows - "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:121.0) Gecko/20100101 Firefox/121.0", - "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:120.0) Gecko/20100101 Firefox/120.0", - # Firefox on macOS - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:121.0) Gecko/20100101 Firefox/121.0", - # Safari on macOS - "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.2 Safari/605.1.15", - # Edge on Windows - "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36 Edg/120.0.0.0", - ] - - return random.choice(user_agents) - - -def _generate_dynamic_user_agent() -> str: - """Generate a user agent with random component combinations.""" - - # 操作系统组合 - os_combinations = [ - "Windows NT 10.0; Win64; x64", - "Windows NT 11.0; Win64; x64", - "Macintosh; Intel Mac OS X 10_15_7", - "Macintosh; Intel Mac OS X 10_15_6", - "Macintosh; Intel Mac OS X 11_7_10", - "X11; Linux x86_64", - "X11; Ubuntu; Linux x86_64", - "X11; Linux i686", - ] - - # Chrome版本 - chrome_versions = [ - "120.0.0.0", - "119.0.0.0", - "118.0.0.0", - "117.0.0.0", - "121.0.0.0", - "116.0.0.0", - "115.0.0.0", - ] - - # WebKit版本 - webkit_versions = [ - "537.36", - "537.35", - "537.34", - ] - - # Firefox版本 - firefox_versions = [ - "121.0", - "120.0", - "119.0", - "118.0", - "117.0", - ] - - # Safari版本 - safari_versions = [ - ("605.1.15", "17.2"), - ("605.1.15", "17.1"), - ("605.1.15", "16.6"), - ] - - # 随机选择浏览器类型 - browser_type = random.choice(["chrome", "firefox", "safari", "edge"]) - os_string = random.choice(os_combinations) - - if browser_type == "chrome": - chrome_version = random.choice(chrome_versions) - webkit_version = random.choice(webkit_versions) - return f"Mozilla/5.0 ({os_string}) AppleWebKit/{webkit_version} (KHTML, like Gecko) Chrome/{chrome_version} Safari/{webkit_version}" - - elif browser_type == "firefox": - if "Windows" in os_string: - firefox_version = random.choice(firefox_versions) - return f"Mozilla/5.0 ({os_string}; rv:{firefox_version}) Gecko/20100101 Firefox/{firefox_version}" - else: - # macOS/Linux Firefox - firefox_version = random.choice(firefox_versions) - return f"Mozilla/5.0 ({os_string}; rv:{firefox_version}) Gecko/20100101 Firefox/{firefox_version}" - - elif browser_type == "safari" and "Macintosh" in os_string: - webkit_version, safari_version = random.choice(safari_versions) - return f"Mozilla/5.0 ({os_string}) AppleWebKit/{webkit_version} (KHTML, like Gecko) Version/{safari_version} Safari/{webkit_version}" - - elif browser_type == "edge" and "Windows" in os_string: - chrome_version = random.choice(chrome_versions) - webkit_version = random.choice(webkit_versions) - edge_version = chrome_version # Edge版本通常与Chrome版本一致 - return f"Mozilla/5.0 ({os_string}) AppleWebKit/{webkit_version} (KHTML, like Gecko) Chrome/{chrome_version} Safari/{webkit_version} Edg/{edge_version}" - - # 如果组合不匹配，回退到Chrome - chrome_version = random.choice(chrome_versions) - webkit_version = random.choice(webkit_versions) - return f"Mozilla/5.0 ({os_string}) AppleWebKit/{webkit_version} (KHTML, like Gecko) Chrome/{chrome_version} Safari/{webkit_version}" diff --git a/pyproject.toml b/pyproject.toml index a85d8cf..d9d756b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,110 +1,71 @@ -[build-system] -requires = ["setuptools>=61.0", "wheel"] -build-backend = "setuptools.build_meta" - [project] name = "flowllm" dynamic = ["version"] -description = "FlowLLM: Simplifying LLM-based HTTP/MCP Service Development" +description = "FlowLLM: Build LLM applications with ease." +readme = "README.md" authors = [ { name = "jinli.yl", email = "jinli.yl@alibaba-inc.com" }, ] -license = {file = "LICENSE"} -readme = "README.md" -keywords = ["llm", "ai", "flow", "framework", "openai", "chatgpt", "language-model", "mcp", "http"] -requires-python = ">=3.10" +license = "Apache-2.0" +requires-python = ">=3.11" +keywords = ["llm", "flow", "agent", "ai", "mcp", "flowllm"] classifiers = [ "Development Status :: 4 - Beta", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.11", + "Operating System :: OS Independent", "Intended Audience :: Developers", "Intended Audience :: Science/Research", - "License :: OSI Approved :: Apache Software License", - "Operating System :: OS Independent", - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.10", "Topic :: Scientific/Engineering :: Artificial Intelligence", - "Topic :: Software Development :: Libraries :: Python Modules", - "Topic :: Software Development :: Libraries :: Application Frameworks", - "Typing :: Typed", ] dependencies = [ - "dashscope>=1.13.0", - "fastapi>=0.104.0", - "fastmcp>=0.1.0", - "httpx>=0.25.0", - "loguru>=0.7.0", - "numpy>=1.24.0", - "pandas>=2.0.0", - "openai>=1.3.0", - "pydantic>=2.0.0", - "PyYAML>=6.0.0", - "requests>=2.31.0", - "tqdm>=4.66.0", - "urllib3>=2.0.0", - "uvicorn[standard]>=0.23.0", - "pyfiglet>=0.8.0", - "rich>=13.5.0", - "pathspec>=0.12.1", - "litellm>=1.0.0", - "tavily-python>=0.2.0", + "agentscope>=2.0.2", + "aiofiles>=24.1.0", + "croniter>=2.0", + "fastapi>=0.135.1", + "fastmcp>=3.1.0", + "httpx>=0.28.1", + "loguru>=0.7.3", + "numpy>=2.2.6", + "pydantic>=2.12.5", + "pyyaml>=6.0.3", + "rich>=14.3.3", + "uvicorn>=0.41.0", ] [project.optional-dependencies] - +claude-code = [ + "claude-agent-sdk>=0.2.91", +] dev = [ - "pytest>=8.4.2", - "pytest_asyncio>=1.2.0", "pre-commit", - "mkdocs-shadcn", - "pipreqs", -] - -reme = [ - "elasticsearch>=8.10.0", - "chromadb>=0.4.0", - "qdrant_client>=1.15.1", - "pgvector>=0.4.2" -] - -fin = [ - "akshare>=1.9.0", - "crawl4ai>=0.7.4", -] - -token = [ -# "modelscope>=1.10.0", -# "diffusers>=0.21.0", -# "peft>=0.6.0", - "transformers>=4.35.0", - "tiktoken>=0.5.0", + "pytest>=8.0", + "pytest-asyncio>=0.23", ] - -skill = [ - -] - full = [ - "flowllm[dev,reme,fin,token,skill]", + "flowllm[claude-code]", + "flowllm[dev]", ] -[tool.setuptools.packages.find] -where = ["."] -include = ["flowllm*"] - -[tool.setuptools.package-data] -flowllm = [ - "**/*.yaml", -] - -[tool.setuptools.dynamic] -version = {attr = "flowllm.__version__"} - [project.urls] Homepage = "https://github.com/flowllm-ai/flowllm" -Documentation = "https://flowllm-ai.github.io/flowllm/" Repository = "https://github.com/flowllm-ai/flowllm" [project.scripts] -flowllm = "flowllm.main:main" +flowllm = "flowllm.application:main" + +[tool.setuptools] +packages = { find = { where = ["."], include = ["flowllm*"] } } +include-package-data = true + +[tool.setuptools.package-data] +"*" = ["py.typed", "**/*.yaml", "**/*.json"] +[tool.setuptools.dynamic] +version = { attr = "flowllm.__version__" } + +[build-system] +requires = ["setuptools>=45", "wheel"] +build-backend = "setuptools.build_meta" diff --git a/skills/flowllm_dev/SKILL.md b/skills/flowllm_dev/SKILL.md new file mode 100644 index 0000000..a1f69ad --- /dev/null +++ b/skills/flowllm_dev/SKILL.md @@ -0,0 +1,179 @@ +--- +name: flowllm-dev +description: FlowLLM repository development guidance. Use when working in the flowllm codebase to implement, debug, test, review, or document FlowLLM runtime behavior, including CLI/client calls, services, application wiring, jobs, steps, components, configuration, streaming, registry registration, and tests. +--- + +# FlowLLM Development + +Use this skill when changing or reviewing the FlowLLM repository. Treat the project as a configuration-driven LLM application framework whose execution path is: + +```text +CLI / Client -> Service -> Application -> Job -> Step -> Component +``` + +Prefer small, framework-shaped changes over ad hoc shortcuts. Read the local files before editing; the repository may move faster than this skill. + +## Start Here + +Read these files first for most development tasks: + +- `docs/zh/framework.md` for the architecture and extension model. +- `docs/zh/contributing.md` for development, tests, and contribution conventions. +- `flowllm/config/default.yaml` for built-in jobs, components, defaults, and env-var placeholders. +- The target implementation and nearby tests before making edits. + +Use `rg` / `rg --files` for discovery. Do not assume a backend is available just because it appears in config; confirm it is registered and imported. + +## Repository Map + +- `flowllm/application.py`: application lifecycle, workspace setup, component/job instantiation, dependency ordering, CLI dispatch. +- `flowllm/config/config_parser.py`: config file loading, dot-notation overrides, env-var expansion, scalar conversion. +- `flowllm/config/default.yaml`: default HTTP service, jobs, and model/embedding/agent components. +- `flowllm/components/base_component.py`: `BaseComponent`, lifecycle hooks, `bind()` dependency declaration, workspace paths. +- `flowllm/components/component_registry.py`: global registry `R`, keyed by `(ComponentEnum, backend_name)`. +- `flowllm/components/runtime_context.py`: per-request `data`, `response`, stream queue, and mapping behavior. +- `flowllm/components/job/`: `BaseJob`, `StreamJob`, background jobs, cron jobs. +- `flowllm/components/service/`: HTTP and MCP service exposure. +- `flowllm/components/client/`: CLI-side clients. +- `flowllm/steps/base_step.py`: `BaseStep`, `Ref`, prompt loading, mapping, dispatch steps, job calls. +- `flowllm/steps/common/`: built-in example/status steps. +- `flowllm/schema/`: pydantic models for config, requests, responses, and stream chunks. +- `flowllm/enumeration/`: component and chunk enums. +- `tests/unit/` and `tests/integration/`: preferred examples for test style and coverage boundaries. + +## Development Rules + +- Implement externally callable behavior as a `Job` configured under `jobs:` and exposed by the active service. +- Implement workflow atoms as `BaseStep` subclasses under `flowllm/steps/`. +- Implement reusable long-lived infrastructure as `BaseComponent` subclasses under `flowllm/components/`. +- Register every backend with `@R.register("")`. +- Ensure the module containing a new registration is imported by the relevant package `__init__.py`; registration happens at import time. +- Put default user-visible behavior in `flowllm/config/default.yaml` when `flowllm start` should expose it. +- Keep service, application, job, step, and component responsibilities separate. +- Update docs when CLI behavior, config keys, service endpoints, public jobs, or user-visible defaults change. +- Add focused tests for bug fixes, new steps, new jobs, config parsing, registry behavior, lifecycle ordering, service exposure, and streaming behavior. + +## Common Implementation Patterns + +### Add a Step and Job + +Create a step: + +```python +from flowllm.components import R +from flowllm.steps import BaseStep + + +@R.register("reverse_step") +class ReverseStep(BaseStep): + async def execute(self): + text = self.context.get("text", "") + self.context.response.answer = text[::-1] + return self.context.response +``` + +Expose it in config: + +```yaml +jobs: + reverse: + backend: base + description: "reverse text" + parameters: + type: object + properties: + text: + type: string + required: + - text + steps: + - backend: reverse_step +``` + +Then verify the package import path registers the step. + +### Add a Component + +Subclass `BaseComponent`, set a non-`BASE` `component_type`, register it, and put startup/cleanup in `_start()` and `_close()`: + +```python +from flowllm.components import BaseComponent, R +from flowllm.enumeration import ComponentEnum + + +@R.register("my_backend") +class MyComponent(BaseComponent): + component_type = ComponentEnum.EMBEDDING_STORE + + async def _start(self) -> None: + ... + + async def _close(self) -> None: + ... +``` + +Declare component dependencies with `BaseComponent.bind()` so `Application` can start components in topological order: + +```python +self.embedding_store = self.bind("default", BaseEmbeddingStore, optional=False) +``` + +### Use Components From Steps + +Use `BaseStep.Ref` / existing descriptors when a step needs configured infrastructure. `BaseStep` already provides: + +- `self.as_llm`: resolved to the configured AS LLM model. +- `self.agent_wrapper`: optional agent wrapper. + +For tests, pass concrete objects through step kwargs when possible; `Ref` checks kwargs and context before resolving from the application context. + +### Streaming + +Use `StreamJob` for SSE output. Streaming steps should enqueue chunks through `RuntimeContext`: + +```python +await self.context.add_stream_string(text, ChunkEnum.CONTENT) +``` + +`StreamJob` always emits a final `DONE` chunk. On exceptions it emits an `ERROR` chunk before `DONE`. + +### Config and CLI + +FlowLLM config supports: + +- Default loading from `flowllm/config/default.yaml`. +- `config=/path/to/app.yaml` for YAML/JSON config. +- Dot-notation overrides such as `service.port=8181`. +- `${VAR}` and `${VAR:-default}` env-var expansion. +- Automatic conversion for booleans, numbers, JSON lists/dicts, and null, while preserving leading-zero strings. + +The CLI entry point is `flowllm.application:main`. `flowllm start` starts the service; other actions call a server-side job of the same name through the selected client backend. + +## Testing + +Run the narrowest useful checks first, then broaden when touching shared surfaces: + +```bash +pytest tests/unit +pytest tests/integration +pytest +pre-commit run --all-files +``` + +Use unit tests for pure behavior such as config parsing, registry behavior, steps, mapping, and lifecycle helpers. Use integration tests when startup, service/client behavior, application wiring, streaming, or external-facing CLI behavior is involved. + +If LLM, embedding, Claude Code, MCP, network, or external service credentials are required and unavailable, state exactly which checks were skipped and why. + +## Review Checklist + +Before finishing a FlowLLM change, confirm: + +- New backends are registered with `R` and imported at package import time. +- Config keys match constructor parameters and pydantic schema expectations. +- `enable_serve` and job type match the intended HTTP/MCP exposure. +- Async lifecycle hooks are idempotent enough for start/close/restart paths. +- Required dependencies use `optional=False`; optional dependencies handle `None`. +- Steps mutate `RuntimeContext` intentionally and return/update `context.response` for non-stream jobs. +- Streaming paths emit content/error chunks and allow the final `DONE` marker. +- User-visible behavior is documented. +- Tests cover both the main path and the most likely failure path. diff --git a/tests/http_client_test.py b/tests/http_client_test.py deleted file mode 100644 index cfdaa14..0000000 --- a/tests/http_client_test.py +++ /dev/null @@ -1,77 +0,0 @@ -"""Test module for HttpClient functionality. - -This module contains tests and examples demonstrating the usage of HttpClient -for interacting with flow execution endpoints, including health checks, -endpoint listing, synchronous flow execution, and streaming flow execution. -""" - -import asyncio -import json - -from flowllm.core.utils import HttpClient - - -async def main(): - """Test function for HttpClient. - - This function demonstrates how to use HttpClient to: - - Check health status - - List available endpoints - - Execute flows synchronously - - Execute flows with streaming - """ - async with HttpClient("http://0.0.0.0:8002") as client: - - # Test 1: Health check - print("=" * 50) - print("Testing health check endpoint...") - try: - health_result = await client.health_check() - print(f"Health check result: {json.dumps(health_result, indent=2)}") - except Exception as e: - print(f"Health check failed: {e}") - - # Test 2: Get available endpoints - print("=" * 50) - print("Getting available endpoints...") - try: - openapi_schema = await client.list_endpoints() - if openapi_schema and "paths" in openapi_schema: - print("Available endpoints:") - for path, methods in openapi_schema["paths"].items(): - for method, details in methods.items(): - if method.upper() == "POST": - print(f" {method.upper()} {path} - {details.get('summary', 'No description')}") - else: - print("Could not retrieve endpoint information") - except Exception as e: - print(f"Failed to get endpoints: {e}") - - query = "阿里巴巴前景如何？" - - response = await client.execute_flow("demo_http_flow", query=query) - print(f"result: {response.answer}") - - # Example curl command for streaming endpoint: - # curl -X POST http://localhost:8002/demo_stream_http_flow \ - # -H "Content-Type: application/json" \ - # -d '{ - # "query": "what is ai" - # }' - print("=" * 50) - print("Testing streaming endpoint...") - try: - async for chunk in client.execute_stream_flow( - "demo_stream_http_flow", - query="what is ai", - ): - chunk_type = chunk.get("type", "answer") - chunk_content = chunk.get("content", "") - if chunk_content: - print(f"[{chunk_type}] {chunk_content}") - except Exception as e: - print(f"Streaming test failed: {e}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/old/op/examples/__init__.py b/tests/integration/__init__.py similarity index 100% rename from old/op/examples/__init__.py rename to tests/integration/__init__.py diff --git a/tests/integration/_workspace_fixture.py b/tests/integration/_workspace_fixture.py new file mode 100644 index 0000000..97463a3 --- /dev/null +++ b/tests/integration/_workspace_fixture.py @@ -0,0 +1,86 @@ +"""Shared fixture: temporary workspace environment for integration tests.""" + +from __future__ import annotations + +import contextlib +import os +import tempfile +from pathlib import Path +from typing import Any, Iterator + + +@contextlib.contextmanager +def temp_chdir(path) -> Iterator[Path]: + """Temporarily chdir, restoring on exit.""" + old = os.getcwd() + os.chdir(path) + try: + yield Path(path) + finally: + os.chdir(old) + + +class WorkspaceEnv: + """Temp workspace + Application lifecycle helpers.""" + + def __init__(self, workspace_dir: Path): + self.workspace_dir = workspace_dir + self._apps: list[Any] = [] + + async def make_app(self, *, config: str | None = None, **overrides) -> Any: + """Build, start, and track an Application for cleanup.""" + from flowllm import Application + from flowllm.config import resolve_app_config + + kwargs: dict[str, Any] = { + "log_to_console": False, + "log_to_file": False, + "enable_logo": False, + "workspace_dir": str(self.workspace_dir), + } + if config: + kwargs["config"] = config + kwargs.update(overrides) + cfg = resolve_app_config(**kwargs) + app = Application(**cfg) + await app.start() + self._apps.append(app) + return app + + async def close_all(self) -> None: + """Close all tracked apps. Idempotent.""" + for app in self._apps: + await app.close() + self._apps.clear() + + def session_state_files(self, prefix: str = "session_state_") -> list[Path]: + """List session-state JSONL files under session/.""" + session_dir = self.workspace_dir / "session" + if not session_dir.exists(): + return [] + return sorted(session_dir.rglob(f"{prefix}*.jsonl")) + + +@contextlib.contextmanager +def workspace_env( + *, + chdir: bool = True, + workspace_name: str = ".flowllm", + load_env_file: bool = True, +) -> Iterator[WorkspaceEnv]: + """Yield a fresh WorkspaceEnv in a temp directory.""" + if load_env_file: + from flowllm.utils import load_env + + load_env() + + with tempfile.TemporaryDirectory() as tmp_dir: + workspace = Path(tmp_dir).resolve() / workspace_name + workspace.mkdir(parents=True, exist_ok=True) + env = WorkspaceEnv(workspace_dir=workspace) + + if chdir: + with temp_chdir(workspace): + yield env + else: + yield env diff --git a/tests/integration/test_agent_session.py b/tests/integration/test_agent_session.py new file mode 100644 index 0000000..b182c84 --- /dev/null +++ b/tests/integration/test_agent_session.py @@ -0,0 +1,128 @@ +"""Integration tests for session persistence and forking in AgentWrapper.""" + +import asyncio +import sys +from pathlib import Path + +INTEGRATION_DIR = Path(__file__).resolve().parent +sys.path.insert(0, str(INTEGRATION_DIR)) + +# pylint: disable=wrong-import-position +from _workspace_fixture import workspace_env # noqa: E402 + +from flowllm.enumeration import ComponentEnum # noqa: E402 + + +async def _run_session_persistence() -> None: + """Same session_id across replies should share context.""" + with workspace_env() as env: + app = await env.make_app() + try: + wrapper = app.context.components[ComponentEnum.AGENT_WRAPPER]["default"] + + sid = "test-persist-session" + + _, msg_1 = await wrapper.reply( + "My favorite color is blue. Remember that.", + session_id=sid, + system_prompt="You are a helpful assistant. Keep answers short.", + ) + text_1 = (msg_1.get_text_content() or "").strip() + print(f"\n[session_persist] reply 1: {text_1!r}") + assert text_1, "Empty first reply" + + files_after_1 = env.session_state_files() + print(f"[session_persist] session files after reply 1: {files_after_1}") + assert len(files_after_1) == 1, f"Expected 1 session file, got {len(files_after_1)}" + assert sid in files_after_1[0].name + + _, msg_2 = await wrapper.reply( + "What is my favorite color?", + session_id=sid, + system_prompt="You are a helpful assistant. Keep answers short.", + ) + text_2 = (msg_2.get_text_content() or "").strip() + print(f"[session_persist] reply 2: {text_2!r}") + assert "blue" in text_2.lower(), f"Agent should recall 'blue' from session context, got: {text_2!r}" + + print("✓ test_session_persistence passed") + finally: + await env.close_all() + + +async def _run_fork_session() -> None: + """fork_session=True creates a new session file with a different ID.""" + with workspace_env() as env: + app = await env.make_app() + try: + wrapper = app.context.components[ComponentEnum.AGENT_WRAPPER]["default"] + + sid = "test-fork-origin" + + await wrapper.reply( + "The secret number is 42.", + session_id=sid, + system_prompt="You are a helpful assistant. Keep answers short.", + ) + files_before_fork = env.session_state_files() + assert len(files_before_fork) == 1 + + forked_sid, msg_fork = await wrapper.reply( + "What is the secret number?", + session_id=sid, + fork_session=True, + system_prompt="You are a helpful assistant. Keep answers short.", + ) + text_fork = (msg_fork.get_text_content() or "").strip() + print(f"\n[fork_session] forked reply: {text_fork!r}") + assert "42" in text_fork, f"Forked session should recall '42', got: {text_fork!r}" + + files_after_fork = env.session_state_files() + print(f"[fork_session] session files after fork: {[f.name for f in files_after_fork]}") + assert ( + len(files_after_fork) == 2 + ), f"Expected 2 session files (original + fork), got {len(files_after_fork)}" + + assert forked_sid != sid, f"Forked session_id should differ from original, got {forked_sid!r}" + + original_file = files_before_fork[0] + assert original_file.exists(), "Original session file should still exist after fork" + + print("✓ test_fork_session passed") + finally: + await env.close_all() + + +async def _run_no_session_id() -> None: + """No session_id means no session file created.""" + with workspace_env() as env: + app = await env.make_app() + try: + wrapper = app.context.components[ComponentEnum.AGENT_WRAPPER]["default"] + + _, msg = await wrapper.reply( + "Say hello.", + system_prompt="You are a helpful assistant. Keep answers short.", + ) + text = (msg.get_text_content() or "").strip() + print(f"\n[no_session] reply: {text!r}") + assert text, "Empty reply" + + files = env.session_state_files() + assert len(files) == 0, f"No session files should be created without session_id, found {files}" + + print("✓ test_no_session_id passed") + finally: + await env.close_all() + + +async def _run_all() -> None: + await _run_no_session_id() + await _run_session_persistence() + await _run_fork_session() + + +if __name__ == "__main__": + print("=== Agent session state integration tests ===") + asyncio.run(_run_all()) + print("\nAll integration tests passed!") diff --git a/tests/integration/test_embedding.py b/tests/integration/test_embedding.py new file mode 100644 index 0000000..4460773 --- /dev/null +++ b/tests/integration/test_embedding.py @@ -0,0 +1,139 @@ +"""Integration tests for embedding store through Application wiring.""" + +import asyncio +import sys +from pathlib import Path + +import numpy as np + +INTEGRATION_DIR = Path(__file__).resolve().parent +sys.path.insert(0, str(INTEGRATION_DIR)) + +# pylint: disable=wrong-import-position +from _workspace_fixture import workspace_env # noqa: E402 + +from flowllm.enumeration import ComponentEnum # noqa: E402 +from flowllm.schema import EmbNode # noqa: E402 + + +def test_embedding_health_check(): + """health_check() returns True with a valid API key.""" + + async def run(): + with workspace_env() as env: + app = await env.make_app() + try: + store = app.context.components[ComponentEnum.EMBEDDING_STORE]["default"] + result = await store.health_check(timeout=10.0) + assert result is True, f"health_check returned {result}" + assert store.is_healthy is True + print("✓ test_embedding_health_check passed") + finally: + await env.close_all() + + asyncio.run(run()) + + +def test_embedding_single_text(): + """Single text returns a valid embedding vector.""" + + async def run(): + with workspace_env() as env: + app = await env.make_app() + try: + store = app.context.components[ComponentEnum.EMBEDDING_STORE]["default"] + emb = await store.get_embedding("Hello, world!") + assert emb is not None, "get_embedding returned None" + assert emb.shape == (store.dimensions,), f"shape {emb.shape} != ({store.dimensions},)" + assert emb.dtype == np.float16, f"dtype {emb.dtype} != float16" + assert np.linalg.norm(emb) > 0, "embedding is a zero vector" + print(f"\n [single] len={len(emb)}, first5={emb[:5].tolist()}") + print("✓ test_embedding_single_text passed") + finally: + await env.close_all() + + asyncio.run(run()) + + +def test_embedding_multiple_texts(): + """Batch embedding returns correct count and shape.""" + + async def run(): + with workspace_env() as env: + app = await env.make_app() + try: + store = app.context.components[ComponentEnum.EMBEDDING_STORE]["default"] + texts = ["cat", "dog", "house"] + results = await store.get_embeddings(texts) + assert len(results) == 3, f"expected 3 results, got {len(results)}" + for i, emb in enumerate(results): + assert emb is not None, f"result[{i}] is None" + assert emb.shape == (store.dimensions,), f"result[{i}] shape mismatch" + assert np.linalg.norm(emb) > 0, f"result[{i}] is a zero vector" + print(f"\n [{texts[i]}] len={len(emb)}, first5={emb[:5].tolist()}") + print("✓ test_embedding_multiple_texts passed") + finally: + await env.close_all() + + asyncio.run(run()) + + +def test_embedding_cache_hit(): + """Repeated text returns cached result.""" + + async def run(): + with workspace_env() as env: + app = await env.make_app() + try: + store = app.context.components[ComponentEnum.EMBEDDING_STORE]["default"] + text = "test caching behavior" + emb1 = await store.get_embedding(text) + cache_size_after_first = len(store._cache) # pylint: disable=protected-access + + emb2 = await store.get_embedding(text) + cache_size_after_second = len(store._cache) # pylint: disable=protected-access + + assert emb1 is not None and emb2 is not None + assert cache_size_after_second == cache_size_after_first, "cache grew on second call" + assert np.array_equal(emb1, emb2), "cached embedding differs from original" + print(f"\n [cache] len={len(emb1)}, first5={emb1[:5].tolist()}") + print("✓ test_embedding_cache_hit passed") + finally: + await env.close_all() + + asyncio.run(run()) + + +def test_embedding_node_embeddings(): + """get_node_embeddings populates embedding on EmbNode objects.""" + + async def run(): + with workspace_env() as env: + app = await env.make_app() + try: + store = app.context.components[ComponentEnum.EMBEDDING_STORE]["default"] + nodes = [ + EmbNode(text="first node text"), + EmbNode(text="second node text"), + ] + result = await store.get_node_embeddings(nodes) + assert result is nodes, "get_node_embeddings should return the same list" + for i, node in enumerate(nodes): + assert node.embedding is not None, f"node[{i}].embedding is None" + assert node.embedding.shape == (store.dimensions,), f"node[{i}] shape mismatch" + print(f"\n [node{i}] len={len(node.embedding)}, first5={node.embedding[:5].tolist()}") + print("✓ test_embedding_node_embeddings passed") + finally: + await env.close_all() + + asyncio.run(run()) + + +if __name__ == "__main__": + print("=== Embedding integration tests ===") + test_embedding_health_check() + test_embedding_single_text() + test_embedding_multiple_texts() + test_embedding_cache_hit() + test_embedding_node_embeddings() + print("\nAll embedding integration tests passed!") diff --git a/tests/integration/test_llm.py b/tests/integration/test_llm.py new file mode 100644 index 0000000..10de6f0 --- /dev/null +++ b/tests/integration/test_llm.py @@ -0,0 +1,108 @@ +"""Integration tests for LLMDemoStep + Agent through Application wiring.""" + +import asyncio +import sys +from pathlib import Path +from typing import Literal + +from pydantic import BaseModel, Field + +INTEGRATION_DIR = Path(__file__).resolve().parent +sys.path.insert(0, str(INTEGRATION_DIR)) + +# pylint: disable=wrong-import-position +from _workspace_fixture import workspace_env # noqa: E402 + +from flowllm.steps.common.llm_demo import LLMDemoStep # noqa: E402 + + +class MathResult(BaseModel): + """Schema for math evaluation results.""" + + expression: str = Field(description="The math expression evaluated") + result: float = Field(description="Numeric result") + explanation: str = Field(description="Brief explanation") + + +class SentimentAnalysis(BaseModel): + """Schema for sentiment analysis results.""" + + sentiment: Literal["positive", "negative", "neutral"] = Field(description="Overall sentiment") + confidence: float = Field(description="Confidence 0-1") + key_phrases: list[str] = Field(description="Key phrases indicating sentiment") + + +async def _run_basic_chat(app) -> None: + step = LLMDemoStep(app_context=app.context) + response = await step( + query="What is 1 + 1? Reply with just the number.", + ) + text = (response.answer or "").strip() + print(f"\n[basic_chat] response: {text!r}") + assert text, "Empty assistant response" + assert "2" in text, f"Expected '2' in response, got: {text!r}" + print("✓ test_llm_demo_step_basic_chat passed") + + +async def _run_with_tool(app) -> None: + step = LLMDemoStep(app_context=app.context) + response = await step( + query="Use the add tool to compute 21 + 21 and report the result.", + sys_prompt="Use the `add` tool whenever the user asks to add numbers.", + ) + text = (response.answer or "").strip() + print(f"\n[with_tool] response: {text!r}") + assert "42" in text, f"Expected '42' in response, got: {text!r}" + print("✓ test_llm_demo_step_with_tool passed") + + +async def _run_structured_output(app) -> None: + step = LLMDemoStep(app_context=app.context) + response = await step( + query="What is 15 multiplied by 7? Show your work.", + sys_prompt="You are a math tutor. Solve the problem step by step.", + structured_model=MathResult, + ) + structured = response.metadata.get("structured_output") + print(f"\n[structured_output] result: {structured}") + assert structured is not None, "structured_output should not be None" + assert "result" in structured, "structured_output should have 'result' field" + assert structured["result"] == 105, f"Expected result=105, got: {structured['result']}" + assert "expression" in structured, "structured_output should have 'expression' field" + assert "explanation" in structured, "structured_output should have 'explanation' field" + print("✓ test_llm_demo_step_structured_output passed") + + +async def _run_structured_output_enum(app) -> None: + step = LLMDemoStep(app_context=app.context) + response = await step( + query="Analyze the sentiment: 'I absolutely love this product! It exceeded all my expectations.'", + sys_prompt="You are a sentiment analysis expert. Analyze the given text.", + structured_model=SentimentAnalysis, + ) + structured = response.metadata.get("structured_output") + print(f"\n[structured_enum] result: {structured}") + assert structured is not None, "structured_output should not be None" + assert structured["sentiment"] == "positive", f"Expected sentiment='positive', got: {structured['sentiment']}" + assert 0 <= structured["confidence"] <= 1, f"Confidence should be 0-1, got: {structured['confidence']}" + assert isinstance(structured["key_phrases"], list), "key_phrases should be a list" + assert len(structured["key_phrases"]) > 0, "key_phrases should not be empty" + print("✓ test_llm_demo_step_structured_output_enum passed") + + +async def _run_all() -> None: + with workspace_env() as env: + app = await env.make_app() + try: + await _run_basic_chat(app) + await _run_with_tool(app) + await _run_structured_output(app) + await _run_structured_output_enum(app) + finally: + await env.close_all() + + +if __name__ == "__main__": + print("=== LLMDemoStep + Agent integration tests ===") + asyncio.run(_run_all()) + print("\nAll integration tests passed!") diff --git a/tests/integration/test_stream_llm.py b/tests/integration/test_stream_llm.py new file mode 100644 index 0000000..339300b --- /dev/null +++ b/tests/integration/test_stream_llm.py @@ -0,0 +1,233 @@ +"""Integration tests for streaming output through StreamLLMDemoStep.""" + +import asyncio +import sys +from pathlib import Path + +INTEGRATION_DIR = Path(__file__).resolve().parent +sys.path.insert(0, str(INTEGRATION_DIR)) + +# pylint: disable=wrong-import-position +from _workspace_fixture import workspace_env # noqa: E402 + +from flowllm.enumeration import ChunkEnum # noqa: E402 +from flowllm.schema import StreamChunk # noqa: E402 +from flowllm.steps.common.stream_llm_demo import StreamLLMDemoStep # noqa: E402 +from flowllm.utils.common_utils import execute_stream_task # noqa: E402 + +# Inject the ``add`` job that StreamLLMDemoStep expects. +ADD_JOB = { + "backend": "base", + "description": "add two numbers", + "parameters": { + "type": "object", + "properties": { + "a": {"type": "number", "description": "first addend"}, + "b": {"type": "number", "description": "second addend"}, + }, + "required": ["a", "b"], + }, + "steps": [{"backend": "add_step"}], +} + + +async def _test_stream_llm_basic_chat(): + """Streams text chunks for a basic chat query.""" + with workspace_env() as env: + app = await env.make_app(jobs={"add": ADD_JOB}) + try: + step = StreamLLMDemoStep(app_context=app.context) + queue: asyncio.Queue = asyncio.Queue() + chunks: list[StreamChunk] = [] + + task = asyncio.create_task( + step( + stream_queue=queue, + query="Explain step by step how to compute 1 + 1, and give the final answer.", + ), + ) + + print("\n[stream_basic] streaming output:") + async for raw in execute_stream_task(queue, task, output_format="chunk"): + chunk: StreamChunk = raw # type: ignore[assignment] + chunks.append(chunk) + if chunk.chunk_type == ChunkEnum.CONTENT: + sys.stdout.write(chunk.chunk) + sys.stdout.flush() + + response = task.result() + + content_chunks = [c for c in chunks if c.chunk_type == ChunkEnum.CONTENT] + print(f"\n\n[stream_basic] got {len(content_chunks)} CONTENT chunks") + assert len(content_chunks) > 1, "Expected multiple CONTENT chunks for streaming" + + # Final answer should be populated + text = (response.answer or "").strip() + assert text, "Empty assistant response" + assert "2" in text, f"Expected '2' in response, got: {text!r}" + + streamed_text = "".join(c.chunk for c in content_chunks) + assert streamed_text.strip() == text, f"Stream text mismatch: {streamed_text!r} vs {text!r}" + print("✓ test_stream_llm_basic_chat passed") + finally: + await env.close_all() + + +async def _test_stream_llm_with_tool(): + """Streams tool call and result events when tools are used.""" + with workspace_env() as env: + app = await env.make_app(jobs={"add": ADD_JOB}) + try: + step = StreamLLMDemoStep(app_context=app.context) + queue: asyncio.Queue = asyncio.Queue() + chunks: list[StreamChunk] = [] + + task = asyncio.create_task( + step( + stream_queue=queue, + query="Use the add tool to compute 21 + 21 and report the result.", + sys_prompt="Use the `add` tool whenever the user asks to add numbers.", + ), + ) + + print("\n[stream_tool] streaming output:") + async for raw in execute_stream_task(queue, task, output_format="chunk"): + chunk: StreamChunk = raw # type: ignore[assignment] + chunks.append(chunk) + if chunk.chunk_type == ChunkEnum.CONTENT: + sys.stdout.write(chunk.chunk) + sys.stdout.flush() + elif chunk.chunk_type == ChunkEnum.TOOL_CALL: + sys.stdout.write(f"\033[33m{chunk.chunk}\033[0m") + sys.stdout.flush() + elif chunk.chunk_type == ChunkEnum.TOOL_RESULT: + sys.stdout.write(f"\033[32m{chunk.chunk}\033[0m") + sys.stdout.flush() + + response = task.result() + + tool_call_chunks = [c for c in chunks if c.chunk_type == ChunkEnum.TOOL_CALL] + tool_result_chunks = [c for c in chunks if c.chunk_type == ChunkEnum.TOOL_RESULT] + content_chunks = [c for c in chunks if c.chunk_type == ChunkEnum.CONTENT] + + print(f"\n\n[stream_tool] TOOL_CALL chunks: {len(tool_call_chunks)}") + print(f"[stream_tool] TOOL_RESULT chunks: {len(tool_result_chunks)}") + print(f"[stream_tool] CONTENT chunks: {len(content_chunks)}") + + assert len(tool_call_chunks) > 0, "Expected TOOL_CALL chunks" + assert len(tool_result_chunks) > 0, "Expected TOOL_RESULT chunks" + + text = (response.answer or "").strip() + print(f"[stream_tool] final answer: {text!r}") + assert "42" in text, f"Expected '42' in response, got: {text!r}" + print("✓ test_stream_llm_with_tool passed") + finally: + await env.close_all() + + +async def _test_stream_llm_fallback_no_stream(): + """Streaming still works without explicit stream_queue setup.""" + with workspace_env() as env: + app = await env.make_app(jobs={"add": ADD_JOB}) + try: + step = StreamLLMDemoStep(app_context=app.context) + queue: asyncio.Queue = asyncio.Queue() + chunks: list[StreamChunk] = [] + + task = asyncio.create_task( + step( + stream_queue=queue, + query="Explain step by step how to compute 1 + 1, and give the final answer.", + ), + ) + + print("\n[fallback_stream] streaming output:") + async for raw in execute_stream_task(queue, task, output_format="chunk"): + chunk: StreamChunk = raw # type: ignore[assignment] + chunks.append(chunk) + if chunk.chunk_type == ChunkEnum.CONTENT: + sys.stdout.write(chunk.chunk) + sys.stdout.flush() + elif chunk.chunk_type == ChunkEnum.THINK: + sys.stdout.write(f"\033[2m{chunk.chunk}\033[0m") + sys.stdout.flush() + + response = task.result() + text = (response.answer or "").strip() + content_chunks = [c for c in chunks if c.chunk_type == ChunkEnum.CONTENT] + print(f"\n\n[fallback_stream] got {len(content_chunks)} CONTENT chunks") + assert text, "Empty assistant response" + assert "2" in text, f"Expected '2' in response, got: {text!r}" + print("✓ test_stream_llm_fallback_no_stream passed") + finally: + await env.close_all() + + +def test_stream_llm_basic_chat(): + """Streams text chunks for basic chat.""" + asyncio.run(_test_stream_llm_basic_chat()) + + +def test_stream_llm_with_tool(): + """Streams tool call events when tools are used.""" + asyncio.run(_test_stream_llm_with_tool()) + + +def test_stream_llm_fallback_no_stream(): + """Fallback behavior without explicit stream setup.""" + asyncio.run(_test_stream_llm_fallback_no_stream()) + + +async def _demo_stream_print(): + """Demo: real-time streaming print with a longer query.""" + with workspace_env() as env: + app = await env.make_app(jobs={"add": ADD_JOB}) + try: + step = StreamLLMDemoStep(app_context=app.context) + queue: asyncio.Queue = asyncio.Queue() + + query = ( + "Please explain in detail how neural networks learn through backpropagation. " + "Include the chain rule, gradient descent, and give a concrete example with numbers." + ) + + task = asyncio.create_task( + step( + stream_queue=queue, + query=query, + sys_prompt="You are a knowledgeable AI teacher. Explain concepts thoroughly.", + ), + ) + + async for raw in execute_stream_task(queue, task, output_format="chunk"): + chunk: StreamChunk = raw # type: ignore[assignment] + if chunk.chunk_type == ChunkEnum.CONTENT: + sys.stdout.write(chunk.chunk) + sys.stdout.flush() + elif chunk.chunk_type == ChunkEnum.THINK: + sys.stdout.write(f"\033[2m{chunk.chunk}\033[0m") + sys.stdout.flush() + elif chunk.chunk_type == ChunkEnum.TOOL_CALL: + sys.stdout.write(f"\n\033[33m[tool_call] {chunk.chunk}\033[0m") + sys.stdout.flush() + elif chunk.chunk_type == ChunkEnum.TOOL_RESULT: + sys.stdout.write(f"\033[32m{chunk.chunk}\033[0m") + sys.stdout.flush() + print() + finally: + await env.close_all() + + +async def _run_all(): + print("=== StreamLLMDemoStep integration tests ===") + await _test_stream_llm_basic_chat() + await _test_stream_llm_with_tool() + await _test_stream_llm_fallback_no_stream() + print("\nAll stream integration tests passed!") + + +if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "demo": + asyncio.run(_demo_stream_print()) + else: + asyncio.run(_run_all()) diff --git a/tests/mcp_client_test.py b/tests/mcp_client_test.py deleted file mode 100644 index f444a73..0000000 --- a/tests/mcp_client_test.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Test module for FastMcpClient functionality. - -This module contains tests and examples demonstrating the usage of FastMcpClient -for interacting with MCP (Model Context Protocol) servers, including listing -available tools and calling MCP tools with arguments. -""" - -import asyncio -import json - -from loguru import logger - -from flowllm.core.utils import FastMcpClient - - -async def main(): - """Test function for FastMcpClient. - - This function demonstrates how to use FastMcpClient to: - - List available MCP tools - - Call MCP tools with arguments - - Handle tool results and errors - """ - # New config-based interface matching McpClient - config = { - "type": "sse", - "url": "http://0.0.0.0:8001/sse", - "headers": {}, - "timeout": 30.0, - } - - async with FastMcpClient("test_client", config, max_retries=3) as client: - - print("=" * 50) - print("Getting available MCP tools...") - try: - tool_calls = await client.list_tool_calls() - if tool_calls: - for tool_call in tool_calls: - print(tool_call.simple_input_dump()) - else: - print("No tools found or failed to retrieve tools") - except Exception as e: - print(f"Failed to get tools: {e}") - - query = "阿里巴巴前景怎么样？" - - test_cases = [ - ("demo_mcp_flow", {"query": query}), - ] - - for tool_name, arguments in test_cases: - print("=" * 50) - print(f"Testing tool: {tool_name} with arguments={arguments}") - try: - # New interface: call_tool(tool_name, arguments) returns CallToolResult - result = await client.call_tool(tool_name, arguments) - print(f"Tool result: {result}") - - if result.content: - print(f"Result content: {result.content[0].text}") - if result.structured_content: - print(f"Structured content: {json.dumps(result.structured_content, indent=2, ensure_ascii=False)}") - if result.is_error: - print(f"⚠ Tool returned error: {result.content[0].text if result.content else 'Unknown error'}") - else: - print("✓ Tool call successful") - - except Exception as e: - logger.exception(e) - print(f"✗ Tool call failed: {e}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/tests/test_base_context.py b/tests/test_base_context.py deleted file mode 100644 index 9eaae7c..0000000 --- a/tests/test_base_context.py +++ /dev/null @@ -1,27 +0,0 @@ -"""Test script for BaseContext. - -This script provides test functions for BaseContext class. -It can be run directly with: python test_base_context.py -""" - -from flowllm.core.context.base_context import BaseContext - - -def main(): - """Test function for BaseContext.""" - ctx = BaseContext(**{"name": "Alice", "age": 30, "city": "New York"}) - - print(ctx.name) - print(ctx.age) - print(ctx.city) - - ctx.email = "alice@example.com" - ctx["email"] = "alice@example.com" - print(ctx.email) - - print(ctx.keys()) - print(ctx) - - -if __name__ == "__main__": - main() diff --git a/tests/test_cache_handler.py b/tests/test_cache_handler.py deleted file mode 100644 index 2c711c5..0000000 --- a/tests/test_cache_handler.py +++ /dev/null @@ -1,532 +0,0 @@ -"""Test script for CacheHandler. - -This script provides comprehensive test functions for CacheHandler class. -It can be run directly with: python test_cache_handler.py -""" - -import shutil -from pathlib import Path -from time import sleep - -import pandas as pd -from loguru import logger - -from flowllm.core.storage.cache_handler import CacheHandler - - -def test_dataframe(): - """Test DataFrame save and load.""" - logger.info("=" * 50) - logger.info("Testing DataFrame save and load") - logger.info("=" * 50) - - cache = CacheHandler(cache_dir="test_cache_dataframe") - - # Create test DataFrame - df = pd.DataFrame( - { - "name": ["Alice", "Bob", "Charlie"], - "age": [25, 30, 35], - "city": ["New York", "London", "Tokyo"], - }, - ) - - # Test save - assert cache.save("test_df", df) is True - assert cache.exists("test_df") is True - - # Test load - loaded_df = cache.load("test_df") - assert loaded_df is not None - assert isinstance(loaded_df, pd.DataFrame) - assert len(loaded_df) == 3 - assert list(loaded_df.columns) == ["name", "age", "city"] - assert loaded_df.iloc[0]["name"] == "Alice" - logger.info("✓ DataFrame save and load test passed") - - # Test get_info - info = cache.get_info("test_df") - assert info is not None - assert info["data_type"] == "DataFrame" - assert "row_count" in info - assert info["row_count"] == 3 - assert "column_count" in info - assert info["column_count"] == 3 - logger.info("✓ DataFrame get_info test passed") - - # Cleanup - cache.clear_all() - shutil.rmtree("test_cache_dataframe", ignore_errors=True) - - -def test_dict(): - """Test dict save and load.""" - logger.info("=" * 50) - logger.info("Testing dict save and load") - logger.info("=" * 50) - - cache = CacheHandler(cache_dir="test_cache_dict") - - # Create test dict - data_dict = { - "name": "Alice", - "age": 30, - "city": "New York", - "hobbies": ["reading", "coding"], - "metadata": {"key1": "value1", "key2": 42}, - } - - # Test save - assert cache.save("test_dict", data_dict) is True - assert cache.exists("test_dict") is True - - # Test load - loaded_dict = cache.load("test_dict") - assert loaded_dict is not None - assert isinstance(loaded_dict, dict) - assert loaded_dict["name"] == "Alice" - assert loaded_dict["age"] == 30 - assert loaded_dict["hobbies"] == ["reading", "coding"] - assert loaded_dict["metadata"]["key1"] == "value1" - logger.info("✓ Dict save and load test passed") - - # Test get_info - info = cache.get_info("test_dict") - assert info is not None - assert info["data_type"] == "dict" - assert "key_count" in info - assert info["key_count"] == 5 - logger.info("✓ Dict get_info test passed") - - # Cleanup - cache.clear_all() - shutil.rmtree("test_cache_dict", ignore_errors=True) - - -def test_list(): - """Test list save and load.""" - logger.info("=" * 50) - logger.info("Testing list save and load") - logger.info("=" * 50) - - cache = CacheHandler(cache_dir="test_cache_list") - - # Create test list - data_list = ["apple", "banana", "cherry", {"nested": "dict"}, [1, 2, 3]] - - # Test save - assert cache.save("test_list", data_list) is True - assert cache.exists("test_list") is True - - # Test load - loaded_list = cache.load("test_list") - assert loaded_list is not None - assert isinstance(loaded_list, list) - assert len(loaded_list) == 5 - assert loaded_list[0] == "apple" - assert loaded_list[3]["nested"] == "dict" - assert loaded_list[4] == [1, 2, 3] - logger.info("✓ List save and load test passed") - - # Test get_info - info = cache.get_info("test_list") - assert info is not None - assert info["data_type"] == "list" - assert "item_count" in info - assert info["item_count"] == 5 - logger.info("✓ List get_info test passed") - - # Cleanup - cache.clear_all() - shutil.rmtree("test_cache_list", ignore_errors=True) - - -def test_string(): - """Test string save and load.""" - logger.info("=" * 50) - logger.info("Testing string save and load") - logger.info("=" * 50) - - cache = CacheHandler(cache_dir="test_cache_string") - - # Create test string - data_string = "This is a test string\nwith multiple lines\nand special chars: 你好世界 🎉" - - # Test save - assert cache.save("test_string", data_string) is True - assert cache.exists("test_string") is True - - # Test load - loaded_string = cache.load("test_string") - assert loaded_string is not None - assert isinstance(loaded_string, str) - assert loaded_string == data_string - logger.info("✓ String save and load test passed") - - # Test get_info - info = cache.get_info("test_string") - assert info is not None - assert info["data_type"] == "str" - assert "char_count" in info - assert info["char_count"] == len(data_string) - logger.info("✓ String get_info test passed") - - # Cleanup - cache.clear_all() - shutil.rmtree("test_cache_string", ignore_errors=True) - - -def test_expiration(): - """Test expiration functionality.""" - logger.info("=" * 50) - logger.info("Testing expiration functionality") - logger.info("=" * 50) - - cache = CacheHandler(cache_dir="test_cache_expiration") - - # Save data with short expiration time (0.001 hours = 3.6 seconds) - test_data = {"key": "value"} - assert cache.save("expired_data", test_data, expire_hours=0.001) is True - assert cache.exists("expired_data") is True - - # Wait for expiration - sleep(4) - assert cache.exists("expired_data") is False - assert cache.load("expired_data") is None - logger.info("✓ Expiration test passed") - - # Save data without expiration - assert cache.save("permanent_data", test_data) is True - sleep(1) - assert cache.exists("permanent_data") is True - assert cache.load("permanent_data") is not None - logger.info("✓ No expiration test passed") - - # Cleanup - cache.clear_all() - shutil.rmtree("test_cache_expiration", ignore_errors=True) - - -def test_delete(): - """Test delete functionality.""" - logger.info("=" * 50) - logger.info("Testing delete functionality") - logger.info("=" * 50) - - cache = CacheHandler(cache_dir="test_cache_delete") - - # Save multiple items - cache.save("item1", {"key1": "value1"}) - cache.save("item2", {"key2": "value2"}) - cache.save("item3", "test string") - - assert cache.exists("item1") is True - assert cache.exists("item2") is True - assert cache.exists("item3") is True - - # Delete one item - assert cache.delete("item1") is True - assert cache.exists("item1") is False - assert cache.load("item1") is None - assert cache.exists("item2") is True - assert cache.exists("item3") is True - logger.info("✓ Delete test passed") - - # Cleanup - cache.clear_all() - shutil.rmtree("test_cache_delete", ignore_errors=True) - - -def test_clean_expired(): - """Test clean_expired functionality.""" - logger.info("=" * 50) - logger.info("Testing clean_expired functionality") - logger.info("=" * 50) - - cache = CacheHandler(cache_dir="test_cache_clean") - - # Save expired data - cache.save("expired1", {"key": "value"}, expire_hours=0.001) - cache.save("expired2", "test", expire_hours=0.001) - - # Save permanent data - cache.save("permanent1", {"key": "value"}) - cache.save("permanent2", "test") - - # Wait for expiration - sleep(4) - - # Clean expired - cleaned_count = cache.clean_expired() - assert cleaned_count == 2 - assert cache.exists("expired1") is False - assert cache.exists("expired2") is False - assert cache.exists("permanent1") is True - assert cache.exists("permanent2") is True - logger.info("✓ Clean expired test passed") - - # Cleanup - cache.clear_all() - shutil.rmtree("test_cache_clean", ignore_errors=True) - - -def test_list_all(): - """Test list_all functionality.""" - logger.info("=" * 50) - logger.info("Testing list_all functionality") - logger.info("=" * 50) - - cache = CacheHandler(cache_dir="test_cache_list_all") - - # Save multiple items - cache.save("item1", {"key1": "value1"}) - cache.save("item2", {"key2": "value2"}) - cache.save("item3", "test string", expire_hours=0.001) - - # Wait for item3 to expire (0.001 hours = 3.6 seconds) - sleep(4) - - # Test list_all without expired - all_items = cache.list_all(include_expired=False) - assert len(all_items) == 2 - assert "item1" in all_items - assert "item2" in all_items - assert "item3" not in all_items - logger.info("✓ list_all (exclude expired) test passed") - - # Test list_all with expired - all_items_with_expired = cache.list_all(include_expired=True) - assert len(all_items_with_expired) == 3 - assert "item1" in all_items_with_expired - assert "item2" in all_items_with_expired - assert "item3" in all_items_with_expired - logger.info("✓ list_all (include expired) test passed") - - # Cleanup - cache.clear_all() - shutil.rmtree("test_cache_list_all", ignore_errors=True) - - -def test_get_cache_stats(): - """Test get_cache_stats functionality.""" - logger.info("=" * 50) - logger.info("Testing get_cache_stats functionality") - logger.info("=" * 50) - - cache = CacheHandler(cache_dir="test_cache_stats") - - # Save multiple items - cache.save("item1", {"key1": "value1"}) - cache.save("item2", {"key2": "value2"}) - cache.save("item3", "test string", expire_hours=0.001) - - # Wait for item3 to expire (0.001 hours = 3.6 seconds) - sleep(4) - - stats = cache.get_cache_stats() - assert stats["total_count"] == 3 - assert stats["expired_count"] == 1 - assert stats["active_count"] == 2 - assert stats["total_size_bytes"] > 0 - # Check that cache_dir is a valid path (check if it contains the expected directory name) - assert "test_cache_stats" in stats["cache_dir"] - assert Path(stats["cache_dir"]).exists() or Path(stats["cache_dir"]).parent.exists() - logger.info("✓ get_cache_stats test passed") - logger.info(f" Stats: {stats}") - - # Cleanup - cache.clear_all() - shutil.rmtree("test_cache_stats", ignore_errors=True) - - -def test_clear_all(): - """Test clear_all functionality.""" - logger.info("=" * 50) - logger.info("Testing clear_all functionality") - logger.info("=" * 50) - - cache = CacheHandler(cache_dir="test_cache_clear") - - # Save multiple items of different types - cache.save("df1", pd.DataFrame({"a": [1, 2, 3]})) - cache.save("dict1", {"key": "value"}) - cache.save("list1", [1, 2, 3]) - cache.save("str1", "test") - - assert len(cache.list_all()) == 4 - - # Clear all - assert cache.clear_all() is True - assert len(cache.list_all()) == 0 - assert cache.exists("df1") is False - assert cache.exists("dict1") is False - assert cache.exists("list1") is False - assert cache.exists("str1") is False - logger.info("✓ clear_all test passed") - - # Cleanup - shutil.rmtree("test_cache_clear", ignore_errors=True) - - -def test_auto_clean_expired(): - """Test auto_clean_expired in load.""" - logger.info("=" * 50) - logger.info("Testing auto_clean_expired in load") - logger.info("=" * 50) - - cache = CacheHandler(cache_dir="test_cache_auto_clean") - - # Save expired data - cache.save("expired_data", {"key": "value"}, expire_hours=0.001) - - # Wait for expiration - sleep(4) - - # Load with auto_clean_expired=True (default) - result = cache.load("expired_data") - assert result is None - assert cache.exists("expired_data") is False - logger.info("✓ auto_clean_expired test passed") - - # Test with auto_clean_expired=False - cache.save("expired_data2", {"key": "value"}, expire_hours=0.001) - sleep(4) - result = cache.load("expired_data2", auto_clean_expired=False) - assert result is None - # File should still exist (not auto-deleted) - assert cache.exists("expired_data2", check_expired=False) is True - logger.info("✓ auto_clean_expired=False test passed") - - # Cleanup - cache.clear_all() - shutil.rmtree("test_cache_auto_clean", ignore_errors=True) - - -def test_multiple_types(): - """Test saving and loading multiple data types in the same cache.""" - logger.info("=" * 50) - logger.info("Testing multiple data types in same cache") - logger.info("=" * 50) - - cache = CacheHandler(cache_dir="test_cache_multiple") - - # Save different types - df = pd.DataFrame({"a": [1, 2], "b": [3, 4]}) - dict_data = {"key": "value"} - list_data = [1, 2, 3] - string_data = "hello world" - - cache.save("df", df) - cache.save("dict", dict_data) - cache.save("list", list_data) - cache.save("string", string_data) - - # Load and verify - loaded_df = cache.load("df") - assert isinstance(loaded_df, pd.DataFrame) - assert len(loaded_df) == 2 - - loaded_dict = cache.load("dict") - assert isinstance(loaded_dict, dict) - assert loaded_dict["key"] == "value" - - loaded_list = cache.load("list") - assert isinstance(loaded_list, list) - assert loaded_list == [1, 2, 3] - - loaded_string = cache.load("string") - assert isinstance(loaded_string, str) - assert loaded_string == "hello world" - - logger.info("✓ Multiple types test passed") - - # Cleanup - cache.clear_all() - shutil.rmtree("test_cache_multiple", ignore_errors=True) - - -def test_nonexistent_key(): - """Test handling of nonexistent keys.""" - logger.info("=" * 50) - logger.info("Testing nonexistent key handling") - logger.info("=" * 50) - - cache = CacheHandler(cache_dir="test_cache_nonexistent") - - # Test load nonexistent key - assert cache.load("nonexistent") is None - logger.info("✓ Load nonexistent key test passed") - - # Test exists nonexistent key - assert cache.exists("nonexistent") is False - logger.info("✓ Exists nonexistent key test passed") - - # Test get_info nonexistent key - assert cache.get_info("nonexistent") is None - logger.info("✓ Get info nonexistent key test passed") - - # Test delete nonexistent key (should not raise error) - assert cache.delete("nonexistent") is True - logger.info("✓ Delete nonexistent key test passed") - - # Cleanup - shutil.rmtree("test_cache_nonexistent", ignore_errors=True) - - -def test_update_existing_key(): - """Test updating an existing key.""" - logger.info("=" * 50) - logger.info("Testing update existing key") - logger.info("=" * 50) - - cache = CacheHandler(cache_dir="test_cache_update") - - # Save initial data - cache.save("test_key", {"old": "value"}) - assert cache.load("test_key")["old"] == "value" - - # Update with new data - cache.save("test_key", {"new": "value"}) - loaded = cache.load("test_key") - assert "old" not in loaded - assert loaded["new"] == "value" - logger.info("✓ Update existing key test passed") - - # Cleanup - cache.clear_all() - shutil.rmtree("test_cache_update", ignore_errors=True) - - -def main(): - """Run all tests.""" - logger.info("\n" + "=" * 70) - logger.info("Starting CacheHandler Tests") - logger.info("=" * 70 + "\n") - - try: - test_dataframe() - test_dict() - test_list() - test_string() - test_expiration() - test_delete() - test_clean_expired() - test_list_all() - test_get_cache_stats() - test_clear_all() - test_auto_clean_expired() - test_multiple_types() - test_nonexistent_key() - test_update_existing_key() - - logger.info("\n" + "=" * 70) - logger.info("All CacheHandler Tests Passed! ✓") - logger.info("=" * 70 + "\n") - - except Exception as e: - logger.error(f"Test failed: {e}") - raise - - -if __name__ == "__main__": - main() diff --git a/tests/test_ds.py b/tests/test_ds.py deleted file mode 100644 index e8dcbd1..0000000 --- a/tests/test_ds.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Test script for DashScope DeepSeek model with streaming and reasoning.""" - -import os - -from openai import OpenAI - -from flowllm.core.utils import load_env - -load_env() - -# 初始化OpenAI客户端 -client = OpenAI( - # 如果没有配置环境变量，请用阿里云百炼API Key替换：api_key="sk-xxx" - api_key=os.getenv("DASHSCOPE_API_KEY"), - base_url="https://dashscope.aliyuncs.com/compatible-mode/v1", -) - -messages = [{"role": "user", "content": "你是谁"}] -completion = client.chat.completions.create( - model="deepseek-v3.2-exp", - messages=messages, - extra_body={"enable_thinking": True}, - stream=True, - stream_options={ - "include_usage": True, - }, -) - -reasoning_content = "" # 完整思考过程 -answer_content = "" # 完整回复 -is_answering = False # 是否进入回复阶段 -print("\n" + "=" * 20 + "思考过程" + "=" * 20 + "\n") - -for chunk in completion: - if not chunk.choices: - print("\n" + "=" * 20 + "Token 消耗" + "=" * 20 + "\n") - print(chunk.usage) - continue - - delta = chunk.choices[0].delta - - # 只收集思考内容 - if hasattr(delta, "reasoning_content") and delta.reasoning_content is not None: - if not is_answering: - print(delta.reasoning_content, end="", flush=True) - reasoning_content += delta.reasoning_content - - # 收到content，开始进行回复 - if hasattr(delta, "content") and delta.content: - if not is_answering: - print("\n" + "=" * 20 + "完整回复" + "=" * 20 + "\n") - is_answering = True - print(delta.content, end="", flush=True) - answer_content += delta.content diff --git a/tests/test_expression_parser.py b/tests/test_expression_parser.py deleted file mode 100644 index a5583ec..0000000 --- a/tests/test_expression_parser.py +++ /dev/null @@ -1,233 +0,0 @@ -"""Test suite for flow expression parser functionality. - -This module contains tests for parsing various flow expressions including -sequential, parallel, and mixed operations with different syntax patterns. -""" - -from flowllm.core.context import C -from flowllm.core.op import BaseOp -from flowllm.core.utils import parse_flow_expression - - -# Define minimal ops for testing and register them with stable names -@C.register_op() -class ContainerOp(BaseOp): - """A container operation for testing purposes.""" - - def execute(self): - """Execute the container operation.""" - return self.name - - -@C.register_op() -class SearchOp(BaseOp): - """A search operation for testing purposes.""" - - def execute(self): - """Execute the search operation.""" - return self.name - - -@C.register_op() -class FindOp(BaseOp): - """A find operation for testing purposes.""" - - def execute(self): - """Execute the find operation.""" - return self.name - - -@C.register_op() -class Op1Op(BaseOp): - """Operation 1 for testing purposes.""" - - def execute(self): - """Execute operation 1.""" - return self.name - - -@C.register_op() -class Op2Op(BaseOp): - """Operation 2 for testing purposes.""" - - def execute(self): - """Execute operation 2.""" - return self.name - - -@C.register_op() -class Op3Op(BaseOp): - """Operation 3 for testing purposes.""" - - def execute(self): - """Execute operation 3.""" - return self.name - - -def test_expression_parser_single_line_sequential(): - """Test parsing a single-line sequential flow expression.""" - flow = "Op1Op() >> Op2Op() >> Op3Op()" - result = parse_flow_expression(flow) - # Should be a SequentialOp with three ops - from flowllm.core.op import SequentialOp - - assert isinstance(result, SequentialOp) - assert len(result.ops) == 3 - assert isinstance(result.ops[0], Op1Op) - assert isinstance(result.ops[1], Op2Op) - assert isinstance(result.ops[2], Op3Op) - - -def test_expression_parser_multiline_exec_and_eval_independent(): - """Test parsing multiline flow with assignments that don't affect the final expression.""" - flow = """ -op = ContainerOp() -op.ops.search = SearchOp() -op.ops.find = FindOp() -Op1Op() >> Op2Op() >> Op3Op() -""".strip() - - result = parse_flow_expression(flow) - from flowllm.core.op import SequentialOp - - # The assignments should execute without affecting the final expression - assert isinstance(result, SequentialOp) - assert len(result.ops) == 3 - assert isinstance(result.ops[0], Op1Op) - assert isinstance(result.ops[1], Op2Op) - assert isinstance(result.ops[2], Op3Op) - - -def test_expression_parser_multiline_return_assigned_op(): - """Test parsing multiline flow that returns an assigned operation.""" - flow = """ -op = ContainerOp() -op.ops.search = SearchOp() -op.ops.find = FindOp() -op -""".strip() - - result = parse_flow_expression(flow) - assert isinstance(result, ContainerOp) - # Verify assignments were applied to the container's ops - assert hasattr(result.ops, "search") - assert hasattr(result.ops, "find") - assert isinstance(result.ops.search, SearchOp) - assert isinstance(result.ops.find, FindOp) - - -def test_expression_parser_multiline_variable_reassignment_and_return(): - """Test parsing multiline flow with variable reassignment.""" - flow = """ -opx = Op1Op() >> Op2Op() -opx = opx >> Op3Op() -opx -""".strip() - - result = parse_flow_expression(flow) - from flowllm.core.op import SequentialOp - - assert isinstance(result, SequentialOp) - assert len(result.ops) == 3 - assert isinstance(result.ops[0], Op1Op) - assert isinstance(result.ops[1], Op2Op) - assert isinstance(result.ops[2], Op3Op) - - -def test_expression_parser_parallel_basic(): - """Test parsing a basic parallel flow expression.""" - flow = "Op1Op() | Op2Op()" - result = parse_flow_expression(flow) - from flowllm.core.op import ParallelOp - - assert isinstance(result, ParallelOp) - assert len(result.ops) == 2 - assert isinstance(result.ops[0], Op1Op) - assert isinstance(result.ops[1], Op2Op) - - -def test_expression_parser_mixed_with_parentheses(): - """Test parsing a mixed sequential and parallel flow with parentheses.""" - flow = "Op1Op() >> (Op2Op() | Op3Op()) >> Op1Op()" - result = parse_flow_expression(flow) - from flowllm.core.op import SequentialOp - from flowllm.core.op import ParallelOp - - assert isinstance(result, SequentialOp) - assert len(result.ops) == 3 - assert isinstance(result.ops[0], Op1Op) - assert isinstance(result.ops[1], ParallelOp) - assert isinstance(result.ops[2], Op1Op) - # Check inner parallel contents - inner = result.ops[1] - assert isinstance(inner.ops[0], Op2Op) - assert isinstance(inner.ops[1], Op3Op) - - -def test_expression_parser_multiline_multiple_attribute_assignments_mixed_chain(): - """Test parsing multiline flow with multiple attribute assignments in a mixed chain.""" - flow = """ -op1 = Op1Op() -op1.ops.search = SearchOp() -op1.ops.find = FindOp() -(op1 | Op2Op()) >> Op3Op() -""".strip() - - result = parse_flow_expression(flow) - from flowllm.core.op import SequentialOp - from flowllm.core.op import ParallelOp - - assert isinstance(result, SequentialOp) - assert isinstance(result.ops[0], ParallelOp) - assert isinstance(result.ops[1], Op3Op) - # attribute assignments applied - assert hasattr(result.ops[0].ops[0].ops, "search") - assert hasattr(result.ops[0].ops[0].ops, "find") - assert isinstance(result.ops[0].ops[0].ops.search, SearchOp) - assert isinstance(result.ops[0].ops[0].ops.find, FindOp) - - -def test_expression_parser_complex_left_shift_parallel_and_sequential(): - """Test parsing a complex flow with left-shift operator, parallel and sequential operations.""" - flow = """ -op = ContainerOp() -op << {"search": Op1Op(), "find": Op2Op()} -(op | Op2Op()) >> (Op1Op() | Op3Op()) >> op -""".strip() - - result = parse_flow_expression(flow) - from flowllm.core.op import SequentialOp - from flowllm.core.op import ParallelOp - - # Structure: Sequential with three parts: Parallel, Parallel, Container - assert isinstance(result, SequentialOp) - assert len(result.ops) == 3 - - first = result.ops[0] - second = result.ops[1] - third = result.ops[2] - - # First parallel: (op | op2) - assert isinstance(first, ParallelOp) - assert len(first.ops) == 2 - assert isinstance(first.ops[0], ContainerOp) - assert isinstance(first.ops[1], Op2Op) - - # Ensure left-shift attached named children to container - container = first.ops[0] - assert hasattr(container.ops, "search") - assert hasattr(container.ops, "find") - assert isinstance(container.ops.search, Op1Op) - assert isinstance(container.ops.find, Op2Op) - - # Second parallel: (op1 | op3) - assert isinstance(second, ParallelOp) - assert len(second.ops) == 2 - assert isinstance(second.ops[0], Op1Op) - assert isinstance(second.ops[1], Op3Op) - - # Final op is the container op - assert isinstance(third, ContainerOp) - - -# pytest -q tests/test_expression_parser.py diff --git a/tests/test_extract_content.py b/tests/test_extract_content.py deleted file mode 100644 index 3a67504..0000000 --- a/tests/test_extract_content.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Test script for common_utils. - -This script provides test functions for common_utils module. -It can be run directly with: python test_common_utils.py -""" - -import pytest - -from flowllm.core.utils import extract_content - - -def test_extract_content_json_no_space(): - """Test extract_content with JSON code block and no space between ``` and json.""" - text = '```json\n{"key": "value"}\n```' - result = extract_content(text, language_tag="json") - assert result == {"key": "value"} - - -def test_extract_content_json_one_space(): - """Test extract_content with JSON code block and one space between ``` and json.""" - text = '``` json\n{"key": "value"}\n```' - result = extract_content(text, language_tag="json") - assert result == {"key": "value"} - - -def test_extract_content_json_two_spaces(): - """Test extract_content with JSON code block and two spaces between ``` and json.""" - text = '``` json\n{"key": "value"}\n```' - result = extract_content(text, language_tag="json") - assert result == {"key": "value"} - - -def test_extract_content_json_multiple_spaces(): - """Test extract_content with JSON code block and multiple spaces between ``` and json.""" - text = '``` json\n{"key": "value"}\n```' - result = extract_content(text, language_tag="json") - assert result == {"key": "value"} - - -def test_extract_content_python_no_space(): - """Test extract_content with Python code block and no space between ``` and python.""" - text = "```python\nprint('hello')\n```" - result = extract_content(text, language_tag="python") - assert result == "print('hello')" - - -def test_extract_content_python_one_space(): - """Test extract_content with Python code block and one space between ``` and python.""" - text = "``` python\nprint('hello')\n```" - result = extract_content(text, language_tag="python") - assert result == "print('hello')" - - -def test_extract_content_python_two_spaces(): - """Test extract_content with Python code block and two spaces between ``` and python.""" - text = "``` python\nprint('hello')\n```" - result = extract_content(text, language_tag="python") - assert result == "print('hello')" - - -def test_extract_content_json_complex(): - """Test extract_content with complex JSON.""" - text = '```json\n{"name": "test", "value": 123, "nested": {"key": "val"}}\n```' - result = extract_content(text, language_tag="json") - assert result == {"name": "test", "value": 123, "nested": {"key": "val"}} - - -def test_extract_content_json_with_spaces_complex(): - """Test extract_content with complex JSON and spaces.""" - text = '``` json\n{"name": "test", "value": 123}\n```' - result = extract_content(text, language_tag="json") - assert result == {"name": "test", "value": 123} - - -def test_extract_content_no_match(): - """Test extract_content when no matching code block is found.""" - text = "This is plain text without code blocks" - result = extract_content(text, language_tag="json") - # When language_tag is "json" and no code block is found, it tries to parse - # the text as JSON, which fails and returns None - assert result is None - - -def test_extract_content_invalid_json(): - """Test extract_content with invalid JSON.""" - text = "```json\n{invalid json}\n```" - result = extract_content(text, language_tag="json") - assert result is None - - -def main(): - """Run all tests.""" - pytest.main([__file__, "-v"]) - - -if __name__ == "__main__": - main() diff --git a/tests/test_fastmcp_client.py b/tests/test_fastmcp_client.py deleted file mode 100644 index df29012..0000000 --- a/tests/test_fastmcp_client.py +++ /dev/null @@ -1,65 +0,0 @@ -"""Test script for FastMcpClient. - -This script provides test functions for FastMcpClient class. -It can be run directly with: python test_fastmcp_client.py -""" - -import asyncio -import os - -from flowllm.core.utils.common_utils import load_env -from flowllm.core.utils.fastmcp_client import FastMcpClient - -load_env() - - -async def main(): - """Test function to demonstrate FastMcpClient usage.""" - # Example 1: Using SSE HTTP connection - # config = { - # "type": "sse", - # "url": "http://11.160.132.45:8010/sse", - # "headers": {}, - # } - # - # async with FastMcpClient("mcp", config) as client: - # tool_calls = await client.list_tool_calls() - # for tool_call in tool_calls: - # print(tool_call.model_dump_json()) - # - # result = await client.call_tool("ant_search", arguments={"query": "半导体行业PE中位数", "entity": "半导体"}) - # print(result) - - # Example 2: Using stdio connection - # config = { - # "command": "npx", - # "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"], - # } - # async with FastMcpClient("mcp", config) as client: - # tools = await client.list_tools() - # tool_calls = await client.list_tool_calls() - # for tool_call in tool_calls: - # print(tool_call.simple_input_dump()) - - # Example 3: Using SSE HTTP connection - config = { - "type": "sse", - "url": "https://dashscope.aliyuncs.com/api/v1/mcps/WebSearch/sse", - "headers": { - "Authorization": f"Bearer {os.getenv('FLOW_BAILIAN_API_KEY')}", - }, - } - - async with FastMcpClient("mcp", config) as client: - tool_calls = await client.list_tool_calls() - tool_call = None - for tool_call in tool_calls: - print(tool_call.model_dump_json()) - - if tool_call is not None: - result = await client.call_tool(tool_call.name, arguments={"query": "半导体行业PE中位数"}) - print(result) - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/tests/test_lite_llm.py b/tests/test_lite_llm.py deleted file mode 100644 index 6a980df..0000000 --- a/tests/test_lite_llm.py +++ /dev/null @@ -1,683 +0,0 @@ -"""Test suite for LiteLLM. - -This module provides comprehensive tests for LiteLLM including: -- Synchronous and asynchronous chat operations -- Streaming responses -- Tool calling functionality -- Error handling -- Different configurations - -Requires proper environment variables: -- FLOW_LLM_API_KEY: API key for authentication -- FLOW_LLM_BASE_URL: Base URL for the API endpoint -""" - -import asyncio - -from flowllm.core.enumeration import Role -from flowllm.core.llm.lite_llm import LiteLLM -from flowllm.core.schema import Message, ToolCall -from flowllm.core.utils import load_env - -load_env() - -# Test model names -MODEL_NAME_1 = "qwen-max-2025-01-25" -MODEL_NAME_2 = "qwen3-30b-a3b-thinking-2507" - - -def create_llm(): - """Create a default LiteLLM instance for testing.""" - return LiteLLM(model_name=MODEL_NAME_1) - - -def create_llm_with_thinking(): - """Create a LiteLLM instance with thinking enabled.""" - return LiteLLM(model_name=MODEL_NAME_2, enable_thinking=True) - - -def create_weather_tool(): - """Create a weather tool definition for testing.""" - return ToolCall( - type="function", - function={ - "name": "get_current_weather", - "description": "Get the current weather in a given location", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "The city and state, e.g. San Francisco, CA", - }, - "unit": { - "type": "string", - "description": "The temperature unit, either 'celsius' or 'fahrenheit'", - "enum": ["celsius", "fahrenheit"], - }, - }, - "required": ["location"], - }, - }, - ) - - -def create_calculator_tool(): - """Create a calculator tool definition for testing.""" - return ToolCall( - type="function", - function={ - "name": "calculator", - "description": "Perform basic arithmetic operations", - "parameters": { - "type": "object", - "properties": { - "operation": { - "type": "string", - "description": "The operation to perform: add, subtract, multiply, or divide", - "enum": ["add", "subtract", "multiply", "divide"], - }, - "a": { - "type": "number", - "description": "The first number", - }, - "b": { - "type": "number", - "description": "The second number", - }, - }, - "required": ["operation", "a", "b"], - }, - }, - ) - - -def create_search_tool(): - """Create a search tool definition for testing.""" - return ToolCall( - type="function", - function={ - "name": "web_search", - "description": "Search the web for information", - "parameters": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "The search query", - }, - "max_results": { - "type": "number", - "description": "Maximum number of results to return", - }, - }, - "required": ["query"], - }, - }, - ) - - -def test_chat_basic(): - """Test basic synchronous chat operation.""" - print("\n=== Test: chat_basic ===") - llm = create_llm() - message: Message = llm.chat( - [Message(role=Role.USER, content="Hello! How are you?")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Content length: {len(message.content)}") - - -def test_chat_with_thinking(): - """Test synchronous chat with thinking enabled.""" - print("\n=== Test: chat_with_thinking ===") - llm = create_llm_with_thinking() - message: Message = llm.chat( - [Message(role=Role.USER, content="What is 2+2?")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Reasoning content: {message.reasoning_content}") - - -def test_chat_with_tools(): - """Test synchronous chat with tool calling.""" - print("\n=== Test: chat_with_tools ===") - llm = create_llm() - weather_tool = create_weather_tool() - message: Message = llm.chat( - [Message(role=Role.USER, content="What's the weather in Beijing?")], - [weather_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Tool calls: {message.tool_calls}") - print(f"✓ Tool calls count: {len(message.tool_calls)}") - if message.tool_calls: - print(f"✓ First tool call name: {message.tool_calls[0].name}") - print(f"✓ First tool call arguments: {message.tool_calls[0].arguments}") - - -def test_chat_multiple_tools(): - """Test synchronous chat with multiple tools.""" - print("\n=== Test: chat_multiple_tools ===") - llm = create_llm() - weather_tool = create_weather_tool() - calculator_tool = create_calculator_tool() - message: Message = llm.chat( - [Message(role=Role.USER, content="What's 10 + 20?")], - [calculator_tool, weather_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Tool calls: {message.tool_calls}") - print(f"✓ Tool calls count: {len(message.tool_calls)}") - if message.tool_calls: - for i, tc in enumerate(message.tool_calls): - print(f"✓ Tool call {i + 1}: name={tc.name}, arguments={tc.arguments}") - - -def test_chat_streaming(): - """Test synchronous chat with streaming enabled.""" - print("\n=== Test: chat_streaming ===") - llm = create_llm() - message: Message = llm.chat( - [Message(role=Role.USER, content="Tell me a short story.")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Content length: {len(message.content)}") - - -def test_chat_empty_message(): - """Test chat with empty message.""" - print("\n=== Test: chat_empty_message ===") - llm = create_llm() - message: Message = llm.chat( - [Message(role=Role.USER, content="")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - - -def test_chat_long_conversation(): - """Test chat with a multi-turn conversation.""" - print("\n=== Test: chat_long_conversation ===") - llm = create_llm() - messages = [ - Message(role=Role.USER, content="My name is Bob."), - Message(role=Role.ASSISTANT, content="Nice to meet you, Bob!"), - Message(role=Role.USER, content="What's my name?"), - ] - message: Message = llm.chat(messages, [], enable_stream_print=False) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - contains_bob = "Bob" in message.content or "bob" in message.content.lower() - print(f"✓ Contains 'Bob': {contains_bob}") - - -def test_chat_complex_query(): - """Test chat with a complex query.""" - print("\n=== Test: chat_complex_query ===") - llm = create_llm() - message: Message = llm.chat( - [ - Message( - role=Role.USER, - content="Explain the difference between machine learning and deep learning in one sentence.", - ), - ], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Content length: {len(message.content)}") - - -async def test_achat_basic(): - """Test basic asynchronous chat operation.""" - print("\n=== Test: achat_basic ===") - llm = create_llm() - message: Message = await llm.achat( - [Message(role=Role.USER, content="Hello! How are you?")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Content length: {len(message.content)}") - - -async def test_achat_with_thinking(): - """Test asynchronous chat with thinking enabled.""" - print("\n=== Test: achat_with_thinking ===") - llm = create_llm_with_thinking() - message: Message = await llm.achat( - [Message(role=Role.USER, content="What is 2+2?")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Reasoning content: {message.reasoning_content}") - - -async def test_achat_with_tools(): - """Test asynchronous chat with tool calling.""" - print("\n=== Test: achat_with_tools ===") - llm = create_llm() - weather_tool = create_weather_tool() - message: Message = await llm.achat( - [Message(role=Role.USER, content="What's the weather in Beijing?")], - [weather_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Tool calls: {message.tool_calls}") - print(f"✓ Tool calls count: {len(message.tool_calls)}") - if message.tool_calls: - for i, tc in enumerate(message.tool_calls): - print(f"✓ Tool call {i + 1}: name={tc.name}, arguments={tc.arguments}") - - -async def test_achat_multiple_tools(): - """Test asynchronous chat with multiple tools.""" - print("\n=== Test: achat_multiple_tools ===") - llm = create_llm() - weather_tool = create_weather_tool() - calculator_tool = create_calculator_tool() - message: Message = await llm.achat( - [Message(role=Role.USER, content="What's 10 * 20?")], - [calculator_tool, weather_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Tool calls: {message.tool_calls}") - print(f"✓ Tool calls count: {len(message.tool_calls)}") - if message.tool_calls: - for i, tc in enumerate(message.tool_calls): - print(f"✓ Tool call {i + 1}: name={tc.name}, arguments={tc.arguments}") - - -async def test_achat_streaming(): - """Test asynchronous chat with streaming.""" - print("\n=== Test: achat_streaming ===") - llm = create_llm() - message: Message = await llm.achat( - [Message(role=Role.USER, content="Tell me a joke.")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - - -async def test_achat_long_conversation(): - """Test async chat with a multi-turn conversation.""" - print("\n=== Test: achat_long_conversation ===") - llm = create_llm() - messages = [ - Message(role=Role.USER, content="Count from 1 to 3."), - Message(role=Role.ASSISTANT, content="1, 2, 3"), - Message(role=Role.USER, content="What number comes after 3?"), - ] - message: Message = await llm.achat(messages, [], enable_stream_print=False) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - contains_four = "4" in message.content or "four" in message.content.lower() - print(f"✓ Contains '4' or 'four': {contains_four}") - - -async def test_achat_concurrent_requests(): - """Test multiple concurrent async requests.""" - print("\n=== Test: achat_concurrent_requests ===") - llm = create_llm() - tasks = [llm.achat([Message(role=Role.USER, content="Say hello")], [], enable_stream_print=False) for _ in range(3)] - messages = await asyncio.gather(*tasks) - print(f"✓ Messages count: {len(messages)}") - for i, message in enumerate(messages): - print(f"✓ Message {i + 1} role: {message.role}") - print(f"✓ Message {i + 1} content: {message.content}") - - -def test_stream_chat_basic(): - """Test basic streaming chat with stream printing.""" - print("\n=== Test: stream_chat_basic ===") - llm = create_llm() - message: Message = llm.chat( - [Message(role=Role.USER, content="Count to 5.")], - [], - enable_stream_print=True, - ) - print(f"\n✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Content length: {len(message.content)}") - - -async def test_astream_chat_basic(): - """Test basic async streaming chat with stream printing.""" - print("\n=== Test: astream_chat_basic ===") - llm = create_llm() - message: Message = await llm.achat( - [Message(role=Role.USER, content="Count to 3.")], - [], - enable_stream_print=True, - ) - print(f"\n✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Content length: {len(message.content)}") - - -def test_stream_chat_with_tools(): - """Test streaming chat with tools and stream printing.""" - print("\n=== Test: stream_chat_with_tools ===") - llm = create_llm() - calculator_tool = create_calculator_tool() - message: Message = llm.chat( - [Message(role=Role.USER, content="Calculate 5 + 3")], - [calculator_tool], - enable_stream_print=True, - ) - print(f"\n✓ Message role: {message.role}") - print(f"✓ Tool calls: {message.tool_calls}") - print(f"✓ Tool calls count: {len(message.tool_calls)}") - if message.tool_calls: - for i, tc in enumerate(message.tool_calls): - print(f"✓ Tool call {i + 1}: name={tc.name}, arguments={tc.arguments}") - - -async def test_astream_chat_with_tools(): - """Test async streaming chat with tools and stream printing.""" - print("\n=== Test: astream_chat_with_tools ===") - llm = create_llm() - weather_tool = create_weather_tool() - message: Message = await llm.achat( - [Message(role=Role.USER, content="What's the weather in Shanghai?")], - [weather_tool], - enable_stream_print=True, - ) - print(f"\n✓ Message role: {message.role}") - print(f"✓ Tool calls: {message.tool_calls}") - print(f"✓ Tool calls count: {len(message.tool_calls)}") - if message.tool_calls: - for i, tc in enumerate(message.tool_calls): - print(f"✓ Tool call {i + 1}: name={tc.name}, arguments={tc.arguments}") - - -def test_stream_chat_thinking(): - """Test streaming chat with thinking enabled and stream printing.""" - print("\n=== Test: stream_chat_thinking ===") - llm = create_llm_with_thinking() - message: Message = llm.chat( - [Message(role=Role.USER, content="What is 5 * 7?")], - [], - enable_stream_print=True, - ) - print(f"\n✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Reasoning content: {message.reasoning_content}") - print(f"✓ Reasoning content length: {len(message.reasoning_content)}") - - -def test_custom_temperature(): - """Test LLM with custom temperature.""" - print("\n=== Test: custom_temperature ===") - llm = LiteLLM(model_name=MODEL_NAME_1, temperature=0.5) - message: Message = llm.chat( - [Message(role=Role.USER, content="Hello!")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - - -def test_custom_seed(): - """Test LLM with custom seed.""" - print("\n=== Test: custom_seed ===") - llm = LiteLLM(model_name=MODEL_NAME_1, seed=123) - message: Message = llm.chat( - [Message(role=Role.USER, content="Hello!")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - - -def test_custom_top_p(): - """Test LLM with custom top_p.""" - print("\n=== Test: custom_top_p ===") - llm = LiteLLM(model_name=MODEL_NAME_1, top_p=0.9) - message: Message = llm.chat( - [Message(role=Role.USER, content="Hello!")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - - -def test_custom_provider(): - """Test LLM with custom provider.""" - print("\n=== Test: custom_provider ===") - llm = LiteLLM(model_name=MODEL_NAME_1, custom_llm_provider="openai") - message: Message = llm.chat( - [Message(role=Role.USER, content="Hello!")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - - -def test_presence_penalty(): - """Test LLM with presence penalty.""" - print("\n=== Test: presence_penalty ===") - llm = LiteLLM(model_name=MODEL_NAME_1, presence_penalty=0.1) - message: Message = llm.chat( - [Message(role=Role.USER, content="Hello!")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - - -def test_tool_call_weather(): - """Test tool calling with weather tool.""" - print("\n=== Test: tool_call_weather ===") - llm = create_llm() - weather_tool = create_weather_tool() - message: Message = llm.chat( - [Message(role=Role.USER, content="Get the weather for Beijing")], - [weather_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - if message.tool_calls: - print(f"✓ Tool calls count: {len(message.tool_calls)}") - tool_call = message.tool_calls[0] - print(f"✓ Tool name: {tool_call.name}") - has_location = "location" in tool_call.arguments.lower() or "beijing" in tool_call.arguments.lower() - print(f"✓ Has location in arguments: {has_location}") - - -def test_tool_call_calculator(): - """Test tool calling with calculator tool.""" - print("\n=== Test: tool_call_calculator ===") - llm = create_llm() - calculator_tool = create_calculator_tool() - message: Message = llm.chat( - [Message(role=Role.USER, content="Calculate 25 multiplied by 4")], - [calculator_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - if message.tool_calls: - print(f"✓ Tool calls count: {len(message.tool_calls)}") - tool_call = message.tool_calls[0] - print(f"✓ Tool name: {tool_call.name}") - has_multiply = "multiply" in tool_call.arguments.lower() or "25" in tool_call.arguments - print(f"✓ Has multiply/25 in arguments: {has_multiply}") - - -def test_tool_call_search(): - """Test tool calling with search tool.""" - print("\n=== Test: tool_call_search ===") - llm = create_llm() - search_tool = create_search_tool() - message: Message = llm.chat( - [Message(role=Role.USER, content="Search for information about Python")], - [search_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - if message.tool_calls: - print(f"✓ Tool calls count: {len(message.tool_calls)}") - tool_call = message.tool_calls[0] - print(f"✓ Tool name: {tool_call.name}") - - -async def test_tool_call_async(): - """Test async tool calling.""" - print("\n=== Test: tool_call_async ===") - llm = create_llm() - weather_tool = create_weather_tool() - message: Message = await llm.achat( - [Message(role=Role.USER, content="What's the weather in Tokyo?")], - [weather_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - if message.tool_calls: - print(f"✓ Tool calls count: {len(message.tool_calls)}") - - -def test_tool_call_multiple_tools_choice(): - """Test tool calling when multiple tools are available.""" - print("\n=== Test: tool_call_multiple_tools_choice ===") - llm = create_llm() - weather_tool = create_weather_tool() - calculator_tool = create_calculator_tool() - message: Message = llm.chat( - [Message(role=Role.USER, content="What is 15 + 27?")], - [calculator_tool, weather_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - if message.tool_calls: - tool_names = [tc.name for tc in message.tool_calls] - print(f"✓ Tool names: {tool_names}") - for i, tc in enumerate(message.tool_calls): - print(f"✓ Tool call {i + 1}: name={tc.name}, arguments={tc.arguments}") - - -def test_tool_call_parallel(): - """Test parallel tool calling with multiple tools.""" - print("\n=== Test: tool_call_parallel ===") - llm = create_llm() - weather_tool = create_weather_tool() - search_tool = create_search_tool() - message: Message = llm.chat( - [ - Message( - role=Role.USER, - content="Get the weather for Beijing and search for information about weather APIs", - ), - ], - [weather_tool, search_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - if message.tool_calls: - print(f"✓ Tool calls count: {len(message.tool_calls)}") - tool_names = [tc.name for tc in message.tool_calls] - print(f"✓ Tool names: {tool_names}") - for i, tc in enumerate(message.tool_calls): - print(f"✓ Tool call {i + 1}: name={tc.name}, arguments={tc.arguments}") - - -def test_tool_call_with_enum(): - """Test tool calling with enum parameter.""" - print("\n=== Test: tool_call_with_enum ===") - llm = create_llm() - calculator_tool = create_calculator_tool() - message: Message = llm.chat( - [Message(role=Role.USER, content="Subtract 10 from 50")], - [calculator_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - if message.tool_calls: - tool_call = message.tool_calls[0] - print(f"✓ Tool name: {tool_call.name}") - args = tool_call.argument_dict if tool_call.check_argument() else {} - if "operation" in args: - print(f"✓ Operation: {args['operation']}") - valid_ops = ["add", "subtract", "multiply", "divide"] - print(f"✓ Valid operation: {args['operation'] in valid_ops}") - - -def main(): - """Run all tests.""" - print("=" * 60) - print("Running LiteLLM Tests") - print("=" * 60) - - # Sync tests - test_chat_basic() - test_chat_with_thinking() - test_chat_with_tools() - test_chat_multiple_tools() - test_chat_streaming() - test_chat_empty_message() - test_chat_long_conversation() - test_chat_complex_query() - - # Streaming tests - test_stream_chat_basic() - test_stream_chat_with_tools() - test_stream_chat_thinking() - - # Configuration tests - test_custom_temperature() - test_custom_seed() - test_custom_top_p() - test_custom_provider() - test_presence_penalty() - - # Tool calling tests - test_tool_call_weather() - test_tool_call_calculator() - test_tool_call_search() - test_tool_call_multiple_tools_choice() - test_tool_call_parallel() - test_tool_call_with_enum() - - # Async tests - print("\n" + "=" * 60) - print("Running Async Tests") - print("=" * 60) - asyncio.run(test_achat_basic()) - asyncio.run(test_achat_with_thinking()) - asyncio.run(test_achat_with_tools()) - asyncio.run(test_achat_multiple_tools()) - asyncio.run(test_achat_streaming()) - asyncio.run(test_achat_long_conversation()) - asyncio.run(test_achat_concurrent_requests()) - asyncio.run(test_astream_chat_basic()) - asyncio.run(test_astream_chat_with_tools()) - asyncio.run(test_tool_call_async()) - - print("\n" + "=" * 60) - print("All tests completed!") - print("=" * 60) - - -if __name__ == "__main__": - main() diff --git a/tests/test_memory_vector_store.py b/tests/test_memory_vector_store.py deleted file mode 100644 index a385a85..0000000 --- a/tests/test_memory_vector_store.py +++ /dev/null @@ -1,797 +0,0 @@ -"""Test script for all VectorStore implementations. - -This script provides test classes for both synchronous and asynchronous -vector store operations. It supports testing all VectorStore implementations -through a configurable factory pattern. - -Supported vector stores: -- memory: MemoryVectorStore (in-memory storage) -- local: LocalVectorStore (local file-based storage) -- chroma: ChromaVectorStore (ChromaDB backend) -- elasticsearch: EsVectorStore (Elasticsearch backend) -- qdrant: QdrantVectorStore (Qdrant backend) -- pgvector: PgVectorStore (PostgreSQL pgvector backend) - -Requires proper environment variables: -- FLOW_EMBEDDING_API_KEY: API key for authentication -- FLOW_EMBEDDING_BASE_URL: Base URL for the API endpoint - -For specific backends: -- Elasticsearch: FLOW_ES_HOSTS (default: http://localhost:9200) -- Qdrant: FLOW_QDRANT_HOST, FLOW_QDRANT_PORT (default: localhost:6333) -- PgVector: FLOW_PGVECTOR_CONNECTION_STRING (default: postgresql://localhost/postgres) - -Usage: - python test_memory_vector_store.py # Test specific store (sync + async) - python test_memory_vector_store.py --all # Test all vector stores - python test_memory_vector_store.py delete # Delete all test-generated files - -Examples: - python test_memory_vector_store.py memory # Test MemoryVectorStore - python test_memory_vector_store.py local # Test LocalVectorStore - python test_memory_vector_store.py chroma # Test ChromaVectorStore - python test_memory_vector_store.py elasticsearch # Test EsVectorStore - python test_memory_vector_store.py qdrant # Test QdrantVectorStore - python test_memory_vector_store.py pgvector # Test PgVectorStore - python test_memory_vector_store.py --all # Test all vector stores - python test_memory_vector_store.py delete # Delete all test files/directories -""" - -import asyncio -import shutil -import sys -from abc import ABC, abstractmethod -from pathlib import Path -from typing import List, Literal, get_args - -from loguru import logger - -from flowllm.core.embedding_model.openai_compatible_embedding_model import ( - OpenAICompatibleEmbeddingModel, -) -from flowllm.core.schema.vector_node import VectorNode -from flowllm.core.utils import load_env -from flowllm.core.vector_store.base_vector_store import BaseVectorStore -from flowllm.core.vector_store.chroma_vector_store import ChromaVectorStore -from flowllm.core.vector_store.es_vector_store import EsVectorStore -from flowllm.core.vector_store.local_vector_store import LocalVectorStore -from flowllm.core.vector_store.memory_vector_store import MemoryVectorStore -from flowllm.core.vector_store.pgvector_vector_store import PgVectorStore -from flowllm.core.vector_store.qdrant_vector_store import QdrantVectorStore - -load_env() - -# ==================== Vector Store Factory ==================== - - -# All supported store types -STORE_TYPES = Literal["memory", "local", "chroma", "elasticsearch", "qdrant", "pgvector"] - - -class VectorStoreFactory: - """Factory class for creating vector store instances.""" - - @staticmethod - def get_all_store_types() -> List[str]: - """Get all supported store types. - - Returns: - List[str]: All supported store type names. - """ - return list(get_args(STORE_TYPES)) - - @staticmethod - def create( - store_type: STORE_TYPES, - embedding_model: OpenAICompatibleEmbeddingModel, - store_dir: str = "./test_vector_store", - es_hosts: str = "http://11.160.132.46:8200", - qdrant_url: str = "http://11.160.132.46:6333", - pg_connection_string: str = "postgresql://localhost/postgres", - pg_async_connection_string: str = "postgresql://localhost/postgres", - ) -> BaseVectorStore: - """Create a vector store instance based on the specified type. - - Args: - store_type: Type of vector store to create - embedding_model: Embedding model instance to use - store_dir: Directory for file-based storage (LocalVectorStore, ChromaVectorStore) - es_hosts: Elasticsearch host URL (for EsVectorStore) - qdrant_url: Qdrant server URL (for QdrantVectorStore) - pg_connection_string: PostgreSQL connection string (for PgVectorStore) - pg_async_connection_string: Async PostgreSQL connection string (for PgVectorStore) - - Returns: - BaseVectorStore: The created vector store instance - - Raises: - ValueError: If the store type is unknown - """ - if store_type == "memory": - return MemoryVectorStore(embedding_model=embedding_model) - elif store_type == "local": - return LocalVectorStore(embedding_model=embedding_model, store_dir=store_dir) - elif store_type == "chroma": - return ChromaVectorStore(embedding_model=embedding_model, store_dir=store_dir) - elif store_type == "elasticsearch": - return EsVectorStore(hosts=es_hosts, embedding_model=embedding_model) - elif store_type == "qdrant": - return QdrantVectorStore(embedding_model=embedding_model, url=qdrant_url) - elif store_type == "pgvector": - return PgVectorStore( - connection_string=pg_connection_string, - async_connection_string=pg_async_connection_string, - embedding_model=embedding_model, - ) - else: - raise ValueError(f"Unknown store type: {store_type}") - - -# ==================== Sample Data Generator ==================== - - -class SampleDataGenerator: - """Generator for sample test data.""" - - @staticmethod - def create_sample_nodes(workspace_id: str, prefix: str = "") -> List[VectorNode]: - """Create sample VectorNode instances for testing. - - Args: - workspace_id: The workspace ID to assign to nodes - prefix: Optional prefix for unique_id to avoid conflicts - - Returns: - List[VectorNode]: List of sample nodes - """ - id_prefix = f"{prefix}_" if prefix else "" - return [ - VectorNode( - unique_id=f"{id_prefix}node1", - workspace_id=workspace_id, - content="Artificial intelligence is a technology that simulates human intelligence.", - metadata={ - "node_type": "tech", - "category": "AI", - }, - ), - VectorNode( - unique_id=f"{id_prefix}node2", - workspace_id=workspace_id, - content="Machine learning is a subset of artificial intelligence.", - metadata={ - "node_type": "tech", - "category": "ML", - }, - ), - VectorNode( - unique_id=f"{id_prefix}node3", - workspace_id=workspace_id, - content="Machine learning is a subset of artificial intelligence.", - metadata={ - "node_type": "tech_new", - "category": "ML", - }, - ), - VectorNode( - unique_id=f"{id_prefix}node4", - workspace_id=workspace_id, - content="I love eating delicious seafood, especially fresh fish.", - metadata={ - "node_type": "food", - "category": "preference", - }, - ), - VectorNode( - unique_id=f"{id_prefix}node5", - workspace_id=workspace_id, - content="Deep learning uses neural networks with multiple layers.", - metadata={ - "node_type": "tech", - "category": "DL", - }, - ), - ] - - -# ==================== Base Test Class ==================== - - -class BaseVectorStoreTest(ABC): - """Abstract base class for vector store tests.""" - - def __init__( - self, - store_type: STORE_TYPES, - workspace_prefix: str = "test", - ): - """Initialize the test class. - - Args: - store_type: Type of vector store to test - workspace_prefix: Prefix for workspace IDs - """ - self.store_type = store_type - self.workspace_prefix = workspace_prefix - self.embedding_model = OpenAICompatibleEmbeddingModel( - dimensions=64, - model_name="text-embedding-v4", - ) - self.store_dir = f"./{store_type}_test_vector_store" - self.client: BaseVectorStore = None - - def _get_workspace_id(self, suffix: str = "") -> str: - """Generate a workspace ID with the configured prefix.""" - base_id = f"{self.workspace_prefix}_{self.store_type}_workspace" - return f"{base_id}_{suffix}" if suffix else base_id - - @abstractmethod - def run_all_tests(self): - """Run all tests. Must be implemented by subclasses.""" - - -# ==================== Synchronous Test Class ==================== - - -class SyncVectorStoreTest(BaseVectorStoreTest): - """Synchronous test class for vector stores.""" - - def __init__( - self, - store_type: STORE_TYPES, - workspace_prefix: str = "sync", - ): - super().__init__(store_type, workspace_prefix) - self.client = VectorStoreFactory.create( - store_type=store_type, - embedding_model=self.embedding_model, - store_dir=self.store_dir, - ) - - def test_create_workspace(self) -> str: - """Test workspace creation.""" - logger.info("=" * 20 + " CREATE WORKSPACE TEST " + "=" * 20) - workspace_id = self._get_workspace_id() - - # Clean up if exists - if self.client.exist_workspace(workspace_id): - self.client.delete_workspace(workspace_id) - - self.client.create_workspace(workspace_id) - assert self.client.exist_workspace(workspace_id), "Workspace should exist after creation" - logger.info(f"Created workspace: {workspace_id}") - return workspace_id - - def test_insert(self, workspace_id: str) -> List[VectorNode]: - """Test node insertion.""" - logger.info("=" * 20 + " INSERT TEST " + "=" * 20) - sample_nodes = SampleDataGenerator.create_sample_nodes(workspace_id, self.workspace_prefix) - self.client.insert(sample_nodes, workspace_id) - logger.info(f"Inserted {len(sample_nodes)} nodes") - return sample_nodes - - def test_search(self, workspace_id: str): - """Test vector search.""" - logger.info("=" * 20 + " SEARCH TEST " + "=" * 20) - results = self.client.search( - "What is artificial intelligence?", - workspace_id=workspace_id, - top_k=3, - ) - logger.info(f"Search returned {len(results)} results") - for i, r in enumerate(results, 1): - logger.info(f"Result {i}: {r.model_dump(exclude={'vector'})}") - assert len(results) > 0, "Search should return results" - - def test_search_with_filter(self, workspace_id: str): - """Test vector search with filter.""" - logger.info("=" * 20 + " FILTER SEARCH TEST " + "=" * 20) - filter_dict = {"metadata.node_type": ["tech", "tech_new"]} - results = self.client.search( - "What is artificial intelligence?", - workspace_id=workspace_id, - top_k=5, - filter_dict=filter_dict, - ) - logger.info(f"Filtered search returned {len(results)} results (node_type in [tech, tech_new])") - for i, r in enumerate(results, 1): - logger.info(f"Filtered Result {i}: {r.model_dump(exclude={'vector'})}") - assert r.metadata.get("node_type") in [ - "tech", - "tech_new", - ], "All results should have node_type in [tech, tech_new]" - - def test_search_with_id(self, workspace_id: str): - """Test vector search by unique_id with empty query.""" - logger.info("=" * 20 + " SEARCH BY ID TEST " + "=" * 20) - target_unique_id = f"{self.workspace_prefix}_node1" - filter_dict = {"unique_id": target_unique_id} - results = self.client.search( - "", - workspace_id=workspace_id, - top_k=1, - filter_dict=filter_dict, - ) - logger.info(f"Search by ID returned {len(results)} results") - assert len(results) == 1, "Should return exactly one result" - assert results[0].unique_id == target_unique_id, f"Result should have unique_id={target_unique_id}" - logger.info(f"Found node: {results[0].model_dump(exclude={'vector'})}") - - def test_update(self, workspace_id: str): - """Test node update (insert with existing unique_id).""" - logger.info("=" * 20 + " UPDATE TEST " + "=" * 20) - updated_node = VectorNode( - unique_id=f"{self.workspace_prefix}_node2", - workspace_id=workspace_id, - content="Machine learning is a powerful subset of AI that learns from data.", - metadata={ - "node_type": "tech", - "category": "ML", - "updated": True, - }, - ) - self.client.insert(updated_node, workspace_id) - - # Search to verify update - results = self.client.search("machine learning", workspace_id=workspace_id, top_k=2) - logger.info(f"After update, search returned {len(results)} results") - for i, r in enumerate(results, 1): - logger.info(f"Updated Result {i}: {r.model_dump(exclude={'vector'})}") - - def test_delete(self, workspace_id: str): - """Test node deletion.""" - logger.info("=" * 20 + " DELETE TEST " + "=" * 20) - node_to_delete = f"{self.workspace_prefix}_node3" - self.client.delete(node_to_delete, workspace_id=workspace_id) - - # Search for deleted content - results = self.client.search("food fish", workspace_id=workspace_id, top_k=5) - logger.info(f"After deletion, found {len(results)} food-related results") - - def test_list_workspace_nodes(self, workspace_id: str): - """Test listing all nodes in a workspace.""" - logger.info("=" * 20 + " LIST WORKSPACE NODES TEST " + "=" * 20) - nodes = self.client.list_workspace_nodes(workspace_id) - logger.info(f"Workspace has {len(nodes)} nodes") - for i, node in enumerate(nodes, 1): - logger.info(f"Node {i}: unique_id={node.unique_id}, content={node.content[:50]}...") - - def test_dump_workspace(self, workspace_id: str): - """Test dumping workspace to disk.""" - logger.info("=" * 20 + " DUMP WORKSPACE TEST " + "=" * 20) - dump_result = self.client.dump_workspace(workspace_id, path=self.store_dir) - logger.info(f"Dumped {dump_result.get('size', 0)} nodes to disk") - return dump_result - - def test_load_workspace(self, workspace_id: str): - """Test loading workspace from disk.""" - logger.info("=" * 20 + " LOAD WORKSPACE TEST " + "=" * 20) - - # Use a separate dump path to avoid being deleted by delete_workspace - dump_path = f"{self.store_dir}_dump" - - # First dump, then delete from memory, then load - self.client.dump_workspace(workspace_id, path=dump_path) - self.client.delete_workspace(workspace_id) - logger.info(f"Deleted workspace from memory, exists: {self.client.exist_workspace(workspace_id)}") - - load_result = self.client.load_workspace(workspace_id, path=dump_path) - logger.info(f"Loaded {load_result.get('size', 0)} nodes from disk") - - # Verify loaded data - results = self.client.search("AI technology", workspace_id=workspace_id, top_k=3) - logger.info(f"After load, search returned {len(results)} results") - for i, r in enumerate(results, 1): - logger.info(f"Loaded Result {i}: {r.model_dump(exclude={'vector'})}") - - def test_copy_workspace(self, workspace_id: str): - """Test copying workspace.""" - logger.info("=" * 20 + " COPY WORKSPACE TEST " + "=" * 20) - copy_workspace_id = self._get_workspace_id("copy") - - # Clean up if exists - if self.client.exist_workspace(copy_workspace_id): - self.client.delete_workspace(copy_workspace_id) - - copy_result = self.client.copy_workspace(workspace_id, copy_workspace_id) - logger.info(f"Copied {copy_result.get('size', 0)} nodes to new workspace") - - # Search in copied workspace - results = self.client.search("AI technology", workspace_id=copy_workspace_id, top_k=2) - logger.info(f"Copy workspace search returned {len(results)} results") - for i, r in enumerate(results, 1): - logger.info(f"Copy Result {i}: {r.model_dump(exclude={'vector'})}") - - # Clean up copy - self.client.delete_workspace(copy_workspace_id) - - def test_list_workspace(self): - """Test listing all workspaces.""" - logger.info("=" * 20 + " LIST WORKSPACE TEST " + "=" * 20) - workspaces = self.client.list_workspace() - logger.info(f"Found {len(workspaces)} workspaces: {workspaces}") - - def test_exist_workspace(self, workspace_id: str): - """Test workspace existence check.""" - logger.info("=" * 20 + " EXIST WORKSPACE TEST " + "=" * 20) - exists = self.client.exist_workspace(workspace_id) - logger.info(f"Workspace {workspace_id} exists: {exists}") - assert exists, "Workspace should exist" - - def test_delete_workspace(self, workspace_id: str): - """Test workspace deletion.""" - logger.info("=" * 20 + " DELETE WORKSPACE TEST " + "=" * 20) - self.client.delete_workspace(workspace_id) - exists = self.client.exist_workspace(workspace_id) - logger.info(f"After deletion, workspace exists: {exists}") - assert not exists, "Workspace should not exist after deletion" - - def cleanup(self, workspace_id: str): - """Clean up test resources.""" - logger.info("=" * 20 + " CLEANUP " + "=" * 20) - if self.client.exist_workspace(workspace_id): - self.client.delete_workspace(workspace_id) - copy_workspace_id = self._get_workspace_id("copy") - if self.client.exist_workspace(copy_workspace_id): - self.client.delete_workspace(copy_workspace_id) - self.client.close() - logger.info("Cleanup completed") - - def run_all_tests(self): - """Run all synchronous tests.""" - logger.info("=" * 50 + f" SYNC {self.store_type.upper()} TESTS " + "=" * 50) - - workspace_id = self.test_create_workspace() - try: - self.test_insert(workspace_id) - self.test_search(workspace_id) - self.test_search_with_filter(workspace_id) - self.test_search_with_id(workspace_id) - self.test_update(workspace_id) - self.test_delete(workspace_id) - self.test_list_workspace_nodes(workspace_id) - self.test_exist_workspace(workspace_id) - self.test_list_workspace() - self.test_copy_workspace(workspace_id) - self.test_dump_workspace(workspace_id) - self.test_load_workspace(workspace_id) - finally: - self.cleanup(workspace_id) - - -# ==================== Asynchronous Test Class ==================== - - -class AsyncVectorStoreTest(BaseVectorStoreTest): - """Asynchronous test class for vector stores.""" - - def __init__( - self, - store_type: STORE_TYPES, - workspace_prefix: str = "async", - ): - super().__init__(store_type, workspace_prefix) - self.client = VectorStoreFactory.create( - store_type=store_type, - embedding_model=self.embedding_model, - store_dir=self.store_dir, - ) - - async def test_create_workspace(self) -> str: - """Test async workspace creation.""" - logger.info("ASYNC - " + "=" * 20 + " CREATE WORKSPACE TEST " + "=" * 20) - workspace_id = self._get_workspace_id() - - # Clean up if exists - if await self.client.async_exist_workspace(workspace_id): - await self.client.async_delete_workspace(workspace_id) - - await self.client.async_create_workspace(workspace_id) - assert await self.client.async_exist_workspace(workspace_id), "Workspace should exist after creation" - logger.info(f"Created workspace: {workspace_id}") - return workspace_id - - async def test_insert(self, workspace_id: str) -> List[VectorNode]: - """Test async node insertion.""" - logger.info("ASYNC - " + "=" * 20 + " INSERT TEST " + "=" * 20) - sample_nodes = SampleDataGenerator.create_sample_nodes(workspace_id, self.workspace_prefix) - await self.client.async_insert(sample_nodes, workspace_id) - logger.info(f"Inserted {len(sample_nodes)} nodes") - return sample_nodes - - async def test_search(self, workspace_id: str): - """Test async vector search.""" - logger.info("ASYNC - " + "=" * 20 + " SEARCH TEST " + "=" * 20) - results = await self.client.async_search( - "What is artificial intelligence?", - workspace_id=workspace_id, - top_k=3, - ) - logger.info(f"Search returned {len(results)} results") - for i, r in enumerate(results, 1): - logger.info(f"Result {i}: {r.model_dump(exclude={'vector'})}") - assert len(results) > 0, "Search should return results" - - async def test_search_with_filter(self, workspace_id: str): - """Test async vector search with filter.""" - logger.info("ASYNC - " + "=" * 20 + " FILTER SEARCH TEST " + "=" * 20) - filter_dict = {"metadata.node_type": ["tech", "tech_new"]} - results = await self.client.async_search( - "What is artificial intelligence?", - workspace_id=workspace_id, - top_k=5, - filter_dict=filter_dict, - ) - logger.info(f"Filtered search returned {len(results)} results (node_type in [tech, tech_new])") - for i, r in enumerate(results, 1): - logger.info(f"Filtered Result {i}: {r.model_dump(exclude={'vector'})}") - assert r.metadata.get("node_type") in [ - "tech", - "tech_new", - ], "All results should have node_type in [tech, tech_new]" - - async def test_search_with_id(self, workspace_id: str): - """Test async vector search by unique_id with empty query.""" - logger.info("ASYNC - " + "=" * 20 + " SEARCH BY ID TEST " + "=" * 20) - target_unique_id = f"{self.workspace_prefix}_node1" - filter_dict = {"unique_id": target_unique_id} - results = await self.client.async_search( - "", - workspace_id=workspace_id, - top_k=1, - filter_dict=filter_dict, - ) - logger.info(f"Search by ID returned {len(results)} results") - assert len(results) == 1, "Should return exactly one result" - assert results[0].unique_id == target_unique_id, f"Result should have unique_id={target_unique_id}" - logger.info(f"Found node: {results[0].model_dump(exclude={'vector'})}") - - async def test_update(self, workspace_id: str): - """Test async node update (insert with existing unique_id).""" - logger.info("ASYNC - " + "=" * 20 + " UPDATE TEST " + "=" * 20) - updated_node = VectorNode( - unique_id=f"{self.workspace_prefix}_node2", - workspace_id=workspace_id, - content="Machine learning is a powerful subset of AI that learns from data.", - metadata={ - "node_type": "tech", - "category": "ML", - "updated": True, - }, - ) - await self.client.async_insert(updated_node, workspace_id) - - # Search to verify update - results = await self.client.async_search("machine learning", workspace_id=workspace_id, top_k=2) - logger.info(f"After update, search returned {len(results)} results") - for i, r in enumerate(results, 1): - logger.info(f"Updated Result {i}: {r.model_dump(exclude={'vector'})}") - - async def test_delete(self, workspace_id: str): - """Test async node deletion.""" - logger.info("ASYNC - " + "=" * 20 + " DELETE TEST " + "=" * 20) - node_to_delete = f"{self.workspace_prefix}_node3" - await self.client.async_delete(node_to_delete, workspace_id=workspace_id) - - # Search for deleted content - results = await self.client.async_search("food fish", workspace_id=workspace_id, top_k=5) - logger.info(f"After deletion, found {len(results)} food-related results") - - async def test_list_workspace_nodes(self, workspace_id: str): - """Test async listing all nodes in a workspace.""" - logger.info("ASYNC - " + "=" * 20 + " LIST WORKSPACE NODES TEST " + "=" * 20) - nodes = await self.client.async_list_workspace_nodes(workspace_id) - logger.info(f"Workspace has {len(nodes)} nodes") - for i, node in enumerate(nodes, 1): - logger.info(f"Node {i}: unique_id={node.unique_id}, content={node.content[:50]}...") - - async def test_dump_workspace(self, workspace_id: str): - """Test async dumping workspace to disk.""" - logger.info("ASYNC - " + "=" * 20 + " DUMP WORKSPACE TEST " + "=" * 20) - dump_result = await self.client.async_dump_workspace(workspace_id, path=self.store_dir) - logger.info(f"Dumped {dump_result.get('size', 0)} nodes to disk") - return dump_result - - async def test_load_workspace(self, workspace_id: str): - """Test async loading workspace from disk.""" - logger.info("ASYNC - " + "=" * 20 + " LOAD WORKSPACE TEST " + "=" * 20) - - # Use a separate dump path to avoid being deleted by delete_workspace - dump_path = f"{self.store_dir}_dump" - - # First dump, then delete from memory, then load - await self.client.async_dump_workspace(workspace_id, path=dump_path) - await self.client.async_delete_workspace(workspace_id) - logger.info(f"Deleted workspace from memory, exists: {await self.client.async_exist_workspace(workspace_id)}") - - load_result = await self.client.async_load_workspace(workspace_id, path=dump_path) - logger.info(f"Loaded {load_result.get('size', 0)} nodes from disk") - - # Verify loaded data - results = await self.client.async_search("AI technology", workspace_id=workspace_id, top_k=3) - logger.info(f"After load, search returned {len(results)} results") - for i, r in enumerate(results, 1): - logger.info(f"Loaded Result {i}: {r.model_dump(exclude={'vector'})}") - - async def test_copy_workspace(self, workspace_id: str): - """Test async copying workspace.""" - logger.info("ASYNC - " + "=" * 20 + " COPY WORKSPACE TEST " + "=" * 20) - copy_workspace_id = self._get_workspace_id("copy") - - # Clean up if exists - if await self.client.async_exist_workspace(copy_workspace_id): - await self.client.async_delete_workspace(copy_workspace_id) - - copy_result = await self.client.async_copy_workspace(workspace_id, copy_workspace_id) - logger.info(f"Copied {copy_result.get('size', 0)} nodes to new workspace") - - # Search in copied workspace - results = await self.client.async_search("AI technology", workspace_id=copy_workspace_id, top_k=2) - logger.info(f"Copy workspace search returned {len(results)} results") - for i, r in enumerate(results, 1): - logger.info(f"Copy Result {i}: {r.model_dump(exclude={'vector'})}") - - # Clean up copy - await self.client.async_delete_workspace(copy_workspace_id) - - async def test_list_workspace(self): - """Test async listing all workspaces.""" - logger.info("ASYNC - " + "=" * 20 + " LIST WORKSPACE TEST " + "=" * 20) - workspaces = await self.client.async_list_workspace() - logger.info(f"Found {len(workspaces)} workspaces: {workspaces}") - - async def test_exist_workspace(self, workspace_id: str): - """Test async workspace existence check.""" - logger.info("ASYNC - " + "=" * 20 + " EXIST WORKSPACE TEST " + "=" * 20) - exists = await self.client.async_exist_workspace(workspace_id) - logger.info(f"Workspace {workspace_id} exists: {exists}") - assert exists, "Workspace should exist" - - async def test_delete_workspace(self, workspace_id: str): - """Test async workspace deletion.""" - logger.info("ASYNC - " + "=" * 20 + " DELETE WORKSPACE TEST " + "=" * 20) - await self.client.async_delete_workspace(workspace_id) - exists = await self.client.async_exist_workspace(workspace_id) - logger.info(f"After deletion, workspace exists: {exists}") - assert not exists, "Workspace should not exist after deletion" - - async def cleanup(self, workspace_id: str): - """Clean up test resources.""" - logger.info("ASYNC - " + "=" * 20 + " CLEANUP " + "=" * 20) - if await self.client.async_exist_workspace(workspace_id): - await self.client.async_delete_workspace(workspace_id) - copy_workspace_id = self._get_workspace_id("copy") - if await self.client.async_exist_workspace(copy_workspace_id): - await self.client.async_delete_workspace(copy_workspace_id) - await self.client.async_close() - # Close the embedding model's async client to prevent "Event loop is closed" errors - await self.embedding_model.async_close() - logger.info("Async cleanup completed") - - async def _run_all_tests_async(self): - """Run all asynchronous tests.""" - logger.info("=" * 50 + f" ASYNC {self.store_type.upper()} TESTS " + "=" * 50) - - workspace_id = await self.test_create_workspace() - try: - await self.test_insert(workspace_id) - await self.test_search(workspace_id) - await self.test_search_with_filter(workspace_id) - await self.test_search_with_id(workspace_id) - await self.test_update(workspace_id) - await self.test_delete(workspace_id) - await self.test_list_workspace_nodes(workspace_id) - await self.test_exist_workspace(workspace_id) - await self.test_list_workspace() - await self.test_copy_workspace(workspace_id) - await self.test_dump_workspace(workspace_id) - await self.test_load_workspace(workspace_id) - finally: - await self.cleanup(workspace_id) - - def run_all_tests(self): - """Run all asynchronous tests.""" - asyncio.run(self._run_all_tests_async()) - - -# ==================== Test Runner ==================== - - -def run_tests( - store_types: List[str], -): - """Run vector store tests (both sync and async). - - Args: - store_types: List of store types to test (e.g., ["memory", "local"]) - """ - for st in store_types: - logger.info(f"\n{'=' * 60}\nTesting {st.upper()} VectorStore\n{'=' * 60}") - try: - # Run sync tests - sync_test = SyncVectorStoreTest(store_type=st) - sync_test.run_all_tests() - - # Run async tests - async_test = AsyncVectorStoreTest(store_type=st) - async_test.run_all_tests() - except Exception as e: - logger.error(f"Failed to test {st}: {e}") - if len(store_types) == 1: - raise - - -def delete_test_files(): - """Delete all test-generated files and directories. - - This function removes all directories that match test patterns: - - *_test_vector_store - - *_test_vector_store_dump - - chroma_test_db - """ - # Get the project root directory (parent of tests/) - project_root = Path(__file__).parent.parent - - # Patterns for test directories to delete - test_patterns = [ - "*_test_vector_store", - "*_test_vector_store_dump", - "chroma_test_db", - ] - - deleted_count = 0 - - for pattern in test_patterns: - for path in project_root.glob(pattern): - if path.is_dir(): - try: - shutil.rmtree(path) - logger.info(f"Deleted directory: {path}") - deleted_count += 1 - except Exception as e: - logger.error(f"Failed to delete {path}: {e}") - elif path.is_file(): - try: - path.unlink() - logger.info(f"Deleted file: {path}") - deleted_count += 1 - except Exception as e: - logger.error(f"Failed to delete {path}: {e}") - - if deleted_count == 0: - logger.info("No test files/directories found to delete.") - else: - logger.info(f"Successfully deleted {deleted_count} test files/directories.") - - -def print_usage(): - """Print usage information.""" - print(__doc__) - print("\nAvailable store types:") - for store_type in VectorStoreFactory.get_all_store_types(): - print(f" - {store_type}") - print("\nOther commands:") - print(" - delete: Delete all test-generated files and directories") - - -if __name__ == "__main__": - args = sys.argv[1:] - valid_store_types = VectorStoreFactory.get_all_store_types() - - if not args or args[0] in ("-h", "--help"): - print_usage() - sys.exit(0) - - if args[0] == "--all": - # Test all vector stores - run_tests(store_types=valid_store_types) - elif args[0] == "delete": - # Delete all test-generated files - delete_test_files() - elif args[0] in valid_store_types: - # Test specific vector store - run_tests(store_types=[args[0]]) - else: - print(f"Unknown argument: {args[0]}") - print_usage() - sys.exit(1) diff --git a/tests/test_openai_compatible_embedding_model.py b/tests/test_openai_compatible_embedding_model.py deleted file mode 100644 index 385a955..0000000 --- a/tests/test_openai_compatible_embedding_model.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Test script for OpenAICompatibleEmbeddingModel. - -This script provides test functions for both synchronous and asynchronous -embedding operations. It can be run directly with: python test_openai_compatible_embedding_model.py - -Requires proper environment variables: -- FLOW_EMBEDDING_API_KEY: API key for authentication -- FLOW_EMBEDDING_BASE_URL: Base URL for the API endpoint -""" - -import asyncio - -from flowllm.core.embedding_model.openai_compatible_embedding_model import ( - OpenAICompatibleEmbeddingModel, -) -from flowllm.core.utils import load_env - -load_env() - - -def main(): - """Test function for synchronous embedding operations. - - This function demonstrates how to use OpenAICompatibleEmbeddingModel - to generate embeddings for both single text strings and lists of text strings. - """ - model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4") - res1 = model.get_embeddings( - "The clothes are of good quality and look good, definitely worth the wait. I love them.", - ) - res2 = model.get_embeddings(["aa", "bb"]) - print(res1) - print(res2) - - -async def async_main(): - """Test function for asynchronous embedding operations. - - This function demonstrates how to use OpenAICompatibleEmbeddingModel - to generate embeddings asynchronously for both single text strings - and lists of text strings. - """ - model = OpenAICompatibleEmbeddingModel(dimensions=64, model_name="text-embedding-v4") - - # Test async single text embedding - res1 = await model.async_get_embeddings( - "The clothes are of good quality and look good, definitely worth the wait. I love them.", - ) - - # Test async batch text embedding - res2 = await model.async_get_embeddings(["aa", "bb"]) - - print("Async results:") - print(res1) - print(res2) - - -if __name__ == "__main__": - # Run async test by default - asyncio.run(async_main()) - - main() diff --git a/tests/test_openai_compatible_llm.py b/tests/test_openai_compatible_llm.py deleted file mode 100644 index fd91ce3..0000000 --- a/tests/test_openai_compatible_llm.py +++ /dev/null @@ -1,521 +0,0 @@ -"""Test suite for OpenAICompatibleLLM. - -This module provides comprehensive tests for OpenAICompatibleLLM including: -- Synchronous and asynchronous chat operations -- Streaming responses -- Tool calling functionality -- Error handling -- Different configurations - -Requires proper environment variables: -- FLOW_LLM_API_KEY: API key for authentication -- FLOW_LLM_BASE_URL: Base URL for the API endpoint -""" - -import asyncio - -from flowllm.core.enumeration import Role -from flowllm.core.llm.openai_compatible_llm import OpenAICompatibleLLM -from flowllm.core.schema import Message, ToolCall -from flowllm.core.utils import load_env - -load_env() - -# Test model names -MODEL_NAME_1 = "qwen-max-2025-01-25" -MODEL_NAME_2 = "qwen3-30b-a3b-thinking-2507" - - -def create_llm(): - """Create a default OpenAICompatibleLLM instance for testing.""" - return OpenAICompatibleLLM(model_name=MODEL_NAME_1, temperature=0) - - -def create_llm_with_thinking(): - """Create an OpenAICompatibleLLM instance with thinking enabled.""" - return OpenAICompatibleLLM(model_name=MODEL_NAME_2, temperature=0) - - -def create_weather_tool(): - """Create a weather tool definition for testing.""" - return ToolCall( - **{ - "type": "function", - "function": { - "name": "get_current_weather", - "description": "Get the current weather in a given location", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "The city and state, e.g. San Francisco, CA", - }, - "unit": { - "type": "string", - "description": "The temperature unit, either 'celsius' or 'fahrenheit'", - "enum": ["celsius", "fahrenheit"], - }, - }, - "required": ["location"], - }, - }, - }, - ) - - -def create_calculator_tool(): - """Create a calculator tool definition for testing.""" - return ToolCall( - **{ - "type": "function", - "function": { - "name": "calculator", - "description": "Perform basic arithmetic operations", - "parameters": { - "type": "object", - "properties": { - "operation": { - "type": "string", - "description": "The operation to perform: add, subtract, multiply, or divide", - "enum": ["add", "subtract", "multiply", "divide"], - }, - "a": { - "type": "number", - "description": "The first number", - }, - "b": { - "type": "number", - "description": "The second number", - }, - }, - "required": ["operation", "a", "b"], - }, - }, - }, - ) - - -def test_chat_basic(): - """Test basic synchronous chat operation.""" - print("\n=== Test: chat_basic ===") - llm = create_llm() - message: Message = llm.chat( - [Message(role=Role.USER, content="Hello!")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Content length: {len(message.content)}") - - -def test_chat_with_thinking(): - """Test synchronous chat with thinking enabled.""" - print("\n=== Test: chat_with_thinking ===") - llm = create_llm_with_thinking() - message: Message = llm.chat( - [Message(role=Role.USER, content="What is 2+2?")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Reasoning content: {message.reasoning_content}") - - -def test_chat_with_tools(): - """Test synchronous chat with tool calling.""" - print("\n=== Test: chat_with_tools ===") - llm = create_llm() - weather_tool = create_weather_tool() - message: Message = llm.chat( - [Message(role=Role.USER, content="What's the weather in Beijing?")], - [weather_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Tool calls: {message.tool_calls}") - print(f"✓ Tool calls count: {len(message.tool_calls)}") - if message.tool_calls: - print(f"✓ First tool call name: {message.tool_calls[0].name}") - print(f"✓ First tool call arguments: {message.tool_calls[0].arguments}") - - -def test_chat_multiple_tools(): - """Test synchronous chat with multiple tools.""" - print("\n=== Test: chat_multiple_tools ===") - llm = create_llm() - weather_tool = create_weather_tool() - calculator_tool = create_calculator_tool() - message: Message = llm.chat( - [Message(role=Role.USER, content="What's 10 + 20?")], - [calculator_tool, weather_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Tool calls: {message.tool_calls}") - print(f"✓ Tool calls count: {len(message.tool_calls)}") - if message.tool_calls: - for i, tc in enumerate(message.tool_calls): - print(f"✓ Tool call {i + 1}: name={tc.name}, arguments={tc.arguments}") - - -def test_chat_streaming(): - """Test synchronous chat with streaming enabled.""" - print("\n=== Test: chat_streaming ===") - llm = create_llm() - message: Message = llm.chat( - [Message(role=Role.USER, content="Tell me a short story.")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Content length: {len(message.content)}") - - -def test_chat_empty_message(): - """Test chat with empty message.""" - print("\n=== Test: chat_empty_message ===") - llm = create_llm() - message: Message = llm.chat( - [Message(role=Role.USER, content="")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - - -def test_chat_long_conversation(): - """Test chat with a multi-turn conversation.""" - print("\n=== Test: chat_long_conversation ===") - llm = create_llm() - messages = [ - Message(role=Role.USER, content="My name is Alice."), - Message(role=Role.ASSISTANT, content="Nice to meet you, Alice!"), - Message(role=Role.USER, content="What's my name?"), - ] - message: Message = llm.chat(messages, [], enable_stream_print=False) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - contains_alice = "Alice" in message.content or "alice" in message.content.lower() - print(f"✓ Contains 'Alice': {contains_alice}") - - -async def test_achat_basic(): - """Test basic asynchronous chat operation.""" - print("\n=== Test: achat_basic ===") - llm = create_llm() - message: Message = await llm.achat( - [Message(role=Role.USER, content="Hello!")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Content length: {len(message.content)}") - - -async def test_achat_with_thinking(): - """Test asynchronous chat with thinking enabled.""" - print("\n=== Test: achat_with_thinking ===") - llm = create_llm_with_thinking() - message: Message = await llm.achat( - [Message(role=Role.USER, content="What is 2+2?")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Reasoning content: {message.reasoning_content}") - - -async def test_achat_with_tools(): - """Test asynchronous chat with tool calling.""" - print("\n=== Test: achat_with_tools ===") - llm = create_llm() - weather_tool = create_weather_tool() - message: Message = await llm.achat( - [Message(role=Role.USER, content="What's the weather in Beijing?")], - [weather_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Tool calls: {message.tool_calls}") - print(f"✓ Tool calls count: {len(message.tool_calls)}") - if message.tool_calls: - for i, tc in enumerate(message.tool_calls): - print(f"✓ Tool call {i + 1}: name={tc.name}, arguments={tc.arguments}") - - -async def test_achat_multiple_tools(): - """Test asynchronous chat with multiple tools.""" - print("\n=== Test: achat_multiple_tools ===") - llm = create_llm() - weather_tool = create_weather_tool() - calculator_tool = create_calculator_tool() - message: Message = await llm.achat( - [Message(role=Role.USER, content="What's 10 * 20?")], - [calculator_tool, weather_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Tool calls: {message.tool_calls}") - print(f"✓ Tool calls count: {len(message.tool_calls)}") - if message.tool_calls: - for i, tc in enumerate(message.tool_calls): - print(f"✓ Tool call {i + 1}: name={tc.name}, arguments={tc.arguments}") - - -async def test_achat_streaming(): - """Test asynchronous chat with streaming.""" - print("\n=== Test: achat_streaming ===") - llm = create_llm() - message: Message = await llm.achat( - [Message(role=Role.USER, content="Tell me a joke.")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - - -async def test_achat_long_conversation(): - """Test async chat with a multi-turn conversation.""" - print("\n=== Test: achat_long_conversation ===") - llm = create_llm() - messages = [ - Message(role=Role.USER, content="Count from 1 to 3."), - Message(role=Role.ASSISTANT, content="1, 2, 3"), - Message(role=Role.USER, content="What number comes after 3?"), - ] - message: Message = await llm.achat(messages, [], enable_stream_print=False) - print(f"✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - contains_four = "4" in message.content or "four" in message.content.lower() - print(f"✓ Contains '4' or 'four': {contains_four}") - - -def test_stream_chat_basic(): - """Test basic streaming chat with stream printing.""" - print("\n=== Test: stream_chat_basic ===") - llm = create_llm() - message: Message = llm.chat( - [Message(role=Role.USER, content="Count to 5.")], - [], - enable_stream_print=True, - ) - print(f"\n✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Content length: {len(message.content)}") - - -async def test_astream_chat_basic(): - """Test basic async streaming chat with stream printing.""" - print("\n=== Test: astream_chat_basic ===") - llm = create_llm() - message: Message = await llm.achat( - [Message(role=Role.USER, content="Count to 3.")], - [], - enable_stream_print=True, - ) - print(f"\n✓ Message role: {message.role}") - print(f"✓ Content: {message.content}") - print(f"✓ Content length: {len(message.content)}") - - -def test_stream_chat_with_tools(): - """Test streaming chat with tools and stream printing.""" - print("\n=== Test: stream_chat_with_tools ===") - llm = create_llm() - calculator_tool = create_calculator_tool() - message: Message = llm.chat( - [Message(role=Role.USER, content="Calculate 5 + 3")], - [calculator_tool], - enable_stream_print=True, - ) - print(f"\n✓ Message role: {message.role}") - print(f"✓ Tool calls: {message.tool_calls}") - print(f"✓ Tool calls count: {len(message.tool_calls)}") - if message.tool_calls: - for i, tc in enumerate(message.tool_calls): - print(f"✓ Tool call {i + 1}: name={tc.name}, arguments={tc.arguments}") - - -async def test_astream_chat_with_tools(): - """Test async streaming chat with tools and stream printing.""" - print("\n=== Test: astream_chat_with_tools ===") - llm = create_llm() - weather_tool = create_weather_tool() - message: Message = await llm.achat( - [Message(role=Role.USER, content="What's the weather in Shanghai?")], - [weather_tool], - enable_stream_print=True, - ) - print(f"\n✓ Message role: {message.role}") - print(f"✓ Tool calls: {message.tool_calls}") - print(f"✓ Tool calls count: {len(message.tool_calls)}") - if message.tool_calls: - for i, tc in enumerate(message.tool_calls): - print(f"✓ Tool call {i + 1}: name={tc.name}, arguments={tc.arguments}") - - -def test_custom_temperature(): - """Test LLM with custom temperature.""" - print("\n=== Test: custom_temperature ===") - llm = OpenAICompatibleLLM(model_name=MODEL_NAME_1, temperature=0.5) - message: Message = llm.chat( - [Message(role=Role.USER, content="Hello!")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - - -def test_custom_seed(): - """Test LLM with custom seed.""" - print("\n=== Test: custom_seed ===") - llm = OpenAICompatibleLLM(model_name=MODEL_NAME_1, seed=123) - message: Message = llm.chat( - [Message(role=Role.USER, content="Hello!")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - - -def test_custom_top_p(): - """Test LLM with custom top_p.""" - print("\n=== Test: custom_top_p ===") - llm = OpenAICompatibleLLM(model_name=MODEL_NAME_1, top_p=0.9) - message: Message = llm.chat( - [Message(role=Role.USER, content="Hello!")], - [], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - - -def test_tool_call_weather(): - """Test tool calling with weather tool.""" - print("\n=== Test: tool_call_weather ===") - llm = create_llm() - weather_tool = create_weather_tool() - message: Message = llm.chat( - [Message(role=Role.USER, content="Get the weather for Beijing")], - [weather_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - if message.tool_calls: - print(f"✓ Tool calls count: {len(message.tool_calls)}") - tool_call = message.tool_calls[0] - print(f"✓ Tool name: {tool_call.name}") - has_location = "location" in tool_call.arguments.lower() or "beijing" in tool_call.arguments.lower() - print(f"✓ Has location in arguments: {has_location}") - - -def test_tool_call_calculator(): - """Test tool calling with calculator tool.""" - print("\n=== Test: tool_call_calculator ===") - llm = create_llm() - calculator_tool = create_calculator_tool() - message: Message = llm.chat( - [Message(role=Role.USER, content="Calculate 25 multiplied by 4")], - [calculator_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - if message.tool_calls: - print(f"✓ Tool calls count: {len(message.tool_calls)}") - tool_call = message.tool_calls[0] - print(f"✓ Tool name: {tool_call.name}") - has_multiply = "multiply" in tool_call.arguments.lower() or "25" in tool_call.arguments - print(f"✓ Has multiply/25 in arguments: {has_multiply}") - - -async def test_tool_call_async(): - """Test async tool calling.""" - print("\n=== Test: tool_call_async ===") - llm = create_llm() - weather_tool = create_weather_tool() - message: Message = await llm.achat( - [Message(role=Role.USER, content="What's the weather in Tokyo?")], - [weather_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - if message.tool_calls: - print(f"✓ Tool calls count: {len(message.tool_calls)}") - - -def test_tool_call_multiple_tools_choice(): - """Test tool calling when multiple tools are available.""" - print("\n=== Test: tool_call_multiple_tools_choice ===") - llm = create_llm() - weather_tool = create_weather_tool() - calculator_tool = create_calculator_tool() - message: Message = llm.chat( - [Message(role=Role.USER, content="What is 15 + 27?")], - [calculator_tool, weather_tool], - enable_stream_print=False, - ) - print(f"✓ Message role: {message.role}") - if message.tool_calls: - tool_names = [tc.name for tc in message.tool_calls] - print(f"✓ Tool names: {tool_names}") - for i, tc in enumerate(message.tool_calls): - print(f"✓ Tool call {i + 1}: name={tc.name}, arguments={tc.arguments}") - - -def main(): - """Run all tests.""" - print("=" * 60) - print("Running OpenAICompatibleLLM Tests") - print("=" * 60) - - # # Sync tests - # test_chat_basic() - # test_chat_with_thinking() - # test_chat_with_tools() - # test_chat_multiple_tools() - # test_chat_streaming() - # test_chat_empty_message() - # test_chat_long_conversation() - # - # # Streaming tests - # test_stream_chat_basic() - # test_stream_chat_with_tools() - # - # # Configuration tests - # test_custom_temperature() - # test_custom_seed() - # test_custom_top_p() - # - # # Tool calling tests - # test_tool_call_weather() - # test_tool_call_calculator() - # test_tool_call_multiple_tools_choice() - - # Async tests - print("\n" + "=" * 60) - print("Running Async Tests") - print("=" * 60) - asyncio.run(test_achat_basic()) - asyncio.run(test_achat_with_thinking()) - asyncio.run(test_achat_with_tools()) - asyncio.run(test_achat_multiple_tools()) - asyncio.run(test_achat_streaming()) - asyncio.run(test_achat_long_conversation()) - asyncio.run(test_astream_chat_basic()) - asyncio.run(test_astream_chat_with_tools()) - asyncio.run(test_tool_call_async()) - - print("\n" + "=" * 60) - print("All tests completed!") - print("=" * 60) - - -if __name__ == "__main__": - main() diff --git a/tests/test_openai_token_usage.py b/tests/test_openai_token_usage.py deleted file mode 100644 index fd4947f..0000000 --- a/tests/test_openai_token_usage.py +++ /dev/null @@ -1,26 +0,0 @@ -"""Simple script to exercise OpenAIToken.token_count.""" - -from flowllm.core.enumeration import Role -from flowllm.core.schema import Message - - -def main(): - """Create a couple of messages and print their token usage.""" - - # from flowllm.core.token import OpenAIToken - # token_counter = OpenAIToken(model_name="gpt-4o-mini") - - from flowllm.core.token import HuggingFaceToken - - token_counter = HuggingFaceToken(use_mirror=True, model_name="Qwen/Qwen3-Next-80B-A3B-Instruct") - - messages = [ - Message(role=Role.SYSTEM, content="你是一个乐于助人的助手。"), - Message(role=Role.USER, content="请简单回答：token计数好吗？"), - ] - count = token_counter.token_count(messages=messages) - print(f"Token count result: {count}") - - -if __name__ == "__main__": - main() diff --git a/tests/test_parse_message_by_keys.py b/tests/test_parse_message_by_keys.py deleted file mode 100644 index a654c93..0000000 --- a/tests/test_parse_message_by_keys.py +++ /dev/null @@ -1,162 +0,0 @@ -"""Test script for parse_message_by_keys function. - -This script provides test functions for parse_message_by_keys in llm_utils module. -It can be run directly with: python test_parse_message_by_keys.py -""" - -import pytest - -from flowllm.core.utils import parse_message_by_keys - - -def test_parse_message_by_keys_basic(): - """Test parse_message_by_keys with two keys (basic case).""" - content = "prefixkey1middlekey2suffix" - keys = ["key1", "key2"] - result = parse_message_by_keys(content, keys) - assert result == {"key1": "middle", "key2": "suffix"} - - -def test_parse_message_by_keys_single_key(): - """Test parse_message_by_keys with a single key.""" - content = "prefixkey1suffix" - keys = ["key1"] - result = parse_message_by_keys(content, keys) - assert result == {"key1": "suffix"} - - -def test_parse_message_by_keys_three_keys(): - """Test parse_message_by_keys with three keys.""" - content = "startkey1middle1key2middle2key3end" - keys = ["key1", "key2", "key3"] - result = parse_message_by_keys(content, keys) - assert result == {"key1": "middle1", "key2": "middle2", "key3": "end"} - - -def test_parse_message_by_keys_with_spaces(): - """Test parse_message_by_keys with content containing spaces (should be stripped).""" - content = " prefixkey1 middle key2suffix " - keys = ["key1", "key2"] - result = parse_message_by_keys(content, keys) - # strip() is applied to origin_content, so spaces at start/end are removed - # but spaces in the middle are preserved - assert result == {"key1": "middle ", "key2": "suffix"} - - -def test_parse_message_by_keys_key_at_start(): - """Test parse_message_by_keys when first key is at the start of content.""" - content = "key1middlekey2suffix" - keys = ["key1", "key2"] - result = parse_message_by_keys(content, keys) - assert result == {"key1": "middle", "key2": "suffix"} - - -def test_parse_message_by_keys_key_not_found_first(): - """Test parse_message_by_keys when first key is not found.""" - content = "prefixmiddlekey2suffix" - keys = ["key1", "key2"] - result = parse_message_by_keys(content, keys) - # When first key is not found, content is kept for next iteration - # So "prefixmiddlekey2suffix" is split by "key2" - assert result == {"key1": "prefixmiddle", "key2": "suffix"} - - -def test_parse_message_by_keys_key_not_found_middle(): - """Test parse_message_by_keys when a middle key is not found.""" - content = "prefixkey1middlekey3end" - keys = ["key1", "key2", "key3"] - result = parse_message_by_keys(content, keys) - # key2 not found, so remaining content "middlekey3end" is assigned to key1 - # then when processing key3, it splits "middlekey3end" into ["middle", "end"] - # so key2 gets "middle" and key3 gets "end" - assert result == {"key1": "middlekey3end", "key2": "middle", "key3": "end"} - - -def test_parse_message_by_keys_key_not_found_last(): - """Test parse_message_by_keys when last key is not found.""" - content = "prefixkey1middle" - keys = ["key1", "key2"] - result = parse_message_by_keys(content, keys) - # Last key not found, remaining content assigned to previous key - assert result == {"key1": "middle", "key2": ""} - - -def test_parse_message_by_keys_all_keys_not_found(): - """Test parse_message_by_keys when all keys are not found.""" - content = "some content without any keys" - keys = ["key1", "key2"] - result = parse_message_by_keys(content, keys) - # First key not found, content kept. Second key not found, content assigned to first key - assert result == {"key1": "some content without any keys", "key2": ""} - - -def test_parse_message_by_keys_empty_content(): - """Test parse_message_by_keys with empty content.""" - content = "" - keys = ["key1", "key2"] - result = parse_message_by_keys(content, keys) - assert result == {"key1": "", "key2": ""} - - -def test_parse_message_by_keys_empty_keys(): - """Test parse_message_by_keys with empty keys list.""" - content = "some content" - keys = [] - result = parse_message_by_keys(content, keys) - # zip([None] + [], []) creates empty iterator, so result is empty dict - assert not result - - -def test_parse_message_by_keys_multiple_occurrences(): - """Test parse_message_by_keys when keys appear multiple times (only first occurrence used).""" - content = "key1firstkey2middlekey1secondkey2end" - keys = ["key1", "key2"] - result = parse_message_by_keys(content, keys) - # Only first occurrence of each key is used - # After splitting by first key1, content becomes "firstkey2middlekey1secondkey2end" - # After splitting by first key2, content becomes "middlekey1secondkey2end" - assert result == {"key1": "first", "key2": "middlekey1secondkey2end"} - - -def test_parse_message_by_keys_with_newlines(): - """Test parse_message_by_keys with content containing newlines.""" - content = "prefix\nkey1\nmiddle\nkey2\nsuffix" - keys = ["key1", "key2"] - result = parse_message_by_keys(content, keys) - # After splitting by key1, content becomes "\nmiddle\nkey2\nsuffix" - # strip() removes leading newline, so it becomes "middle\nkey2\nsuffix" - # After splitting by key2, key1 gets "middle\n" (before key2), key2 gets "\nsuffix" - assert result == {"key1": "middle\n", "key2": "\nsuffix"} - - -def test_parse_message_by_keys_special_characters(): - """Test parse_message_by_keys with special characters in keys and content.""" - content = "prefix@key1middle#key2suffix" - keys = ["@key1", "#key2"] - result = parse_message_by_keys(content, keys) - assert result == {"@key1": "middle", "#key2": "suffix"} - - -def test_parse_message_by_keys_unicode(): - """Test parse_message_by_keys with unicode characters.""" - content = "前缀key1中间key2后缀" - keys = ["key1", "key2"] - result = parse_message_by_keys(content, keys) - assert result == {"key1": "中间", "key2": "后缀"} - - -def test_parse_message_by_keys_consecutive_keys(): - """Test parse_message_by_keys with consecutive keys (no content between).""" - content = "prefixkey1key2suffix" - keys = ["key1", "key2"] - result = parse_message_by_keys(content, keys) - assert result == {"key1": "", "key2": "suffix"} - - -def main(): - """Run all tests.""" - pytest.main([__file__, "-v"]) - - -if __name__ == "__main__": - main() diff --git a/tests/test_skill_agent_flow.py b/tests/test_skill_agent_flow.py deleted file mode 100644 index 86b4110..0000000 --- a/tests/test_skill_agent_flow.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Test module for skill_agent_flow functionality. - -This module contains a test demonstrating the usage of HttpClient -for executing skill_agent_flow with a PDF filling task. -""" - -import asyncio - -from flowllm.core.utils import HttpClient - - -async def main(): - """Test function for skill_agent_flow. - - This function demonstrates how to use HttpClient to execute - skill_agent_flow with a PDF filling task. - curl -X POST http://localhost:8002/skill_agent_flow \ - -H "Content-Type: application/json" \ - -d '{ - "query": "xxxx", - "skill_dir": "../skills" - }' - """ - async with HttpClient("http://0.0.0.0:8002") as client: - query = ( - "Fill Sample-Fillable-PDF.pdf with: name='Alice Johnson', select first choice from dropdown, " - "check options 1 and 3, dependent name='Bob Johnson', age='12'. Save as filled-sample.pdf" - ) - # query = "把我做一个简单的述职PPT" - # query = "帮我创建一个写小说的skills" - skill_dir = "../skills" - print("=" * 50) - print("Testing skill_agent_flow endpoint...") - response = await client.execute_flow( - "skill_agent", - query=query, - skill_dir=skill_dir, - ) - print(f"Result: {response.answer}") - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/tests/test_tool_call.py b/tests/test_tool_call.py deleted file mode 100644 index dc74547..0000000 --- a/tests/test_tool_call.py +++ /dev/null @@ -1,52 +0,0 @@ -"""Test script for ToolCall. - -This script provides test functions for ToolCall class serialization and deserialization. -It can be run directly with: python test_tool_call.py -""" - -import json - -from flowllm.core.schema.tool_call import ToolCall - - -def main(): - """Test function to demonstrate ToolCall serialization and deserialization.""" - data = { - "id": "call_0fb6077ad56f4647b0b04a", - "function": { - "arguments": '{"symbol": "ZETA"}', - "name": "get_stock_info", - }, - "type": "function", - "index": 0, - } - tool_call = ToolCall(**data) - output_data = tool_call.simple_output_dump() - assert output_data == data - print(json.dumps(output_data, ensure_ascii=False)) - - data = { - "type": "function", - "function": { - "name": "get_current_weather", - "description": "It is very useful when you want to check the weather of a specified city.", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "Cities or counties, such as Beijing, Hangzhou, Yuhang District, etc.", - }, - }, - "required": ["location"], - }, - }, - } - tool_call = ToolCall(**data) - input_data = tool_call.simple_input_dump() - assert input_data == data - print(json.dumps(input_data, ensure_ascii=False)) - - -if __name__ == "__main__": - main() diff --git a/tests/test_tushare_client.py b/tests/test_tushare_client.py deleted file mode 100644 index b8feadb..0000000 --- a/tests/test_tushare_client.py +++ /dev/null @@ -1,18 +0,0 @@ -"""Test module for Tushare client.""" - -from loguru import logger - -from flowllm.core.utils import TushareClient -from flowllm.main import FlowLLMApp - - -def main(): - """Test the Tushare client by requesting daily data.""" - with FlowLLMApp(): - client = TushareClient() - df = client.request(api_name="daily", ts_code="000001.SZ") - logger.info(df) - - -if __name__ == "__main__": - main() diff --git a/tests/test_vision_audio.py b/tests/test_vision_audio.py deleted file mode 100644 index 3decc9a..0000000 --- a/tests/test_vision_audio.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Test script for vision and audio capabilities with OpenAI-compatible LLM.""" - -import base64 - -from flowllm.core.llm import OpenAICompatibleLLM -from flowllm.core.schema import Message - - -def run_sync_openai_image(): - """Test image understanding with Qwen VL model.""" - - def encode_image(image_path): - with open(image_path, "rb") as image_file: - return base64.b64encode(image_file.read()).decode("utf-8") - - # 将xxxx/eagle.png替换为你本地图像的绝对路径 - base64_image = encode_image("/Users/yuli/Documents/20251128144329.jpg") - - llm = OpenAICompatibleLLM(model_name="qwen3-vl-plus") - messages = [ - Message( - **{ - "role": "user", - "content": [ - { - "type": "image_url", - "image_url": { - # "url": "https://img.alicdn.com/imgextra/ - # i1/O1CN01gDEY8M1W114Hi3XcN_!!6000000002727-0-tps-1024-406.jpg" - "url": f"data:[image/jpeg];base64,{base64_image}", - }, - }, - # {"type": "text", "text": "这道题怎么解答？"}, - {"type": "text", "text": "看看用户在干什么？"}, - ], - }, - ), - ] - - print(messages) - - print(llm.chat(messages, enable_stream_print=True)) - - -def run_sync_openai_audio(): - """Test audio understanding with Qwen Omni model.""" - - def encode_audio(audio_path): - with open(audio_path, "rb") as audio_file: - return base64.b64encode(audio_file.read()).decode("utf-8") - - # 请将 ABSOLUTE_PATH/welcome.mp3 替换为本地音频的绝对路径 - audio_file_path = "/Users/yuli/Documents/111.wav" - base64_audio = encode_audio(audio_file_path) - - messages = [ - Message( - **{ - "role": "user", - "content": [ - { - "type": "input_audio", - "input_audio": { - # "data": "https://help-static-aliyun-doc.aliyuncs.com/ - # file-manage-files/zh-CN/20250211/tixcef/cherry.wav", - "data": f"data:;base64,{base64_audio}", - }, - }, - { - "type": "text", - "text": '这段音频在说什么? 语气是什么样的？以json格式输出{"content": "", "tone": ""}', - }, - ], - }, - ), - ] - llm = OpenAICompatibleLLM(model_name="qwen3-omni-flash") - - print(llm.chat(messages, enable_stream_print=True)) - - -if __name__ == "__main__": - run_sync_openai_image() - run_sync_openai_audio() diff --git a/tests/unit/test_base_component.py b/tests/unit/test_base_component.py new file mode 100644 index 0000000..3ada44f --- /dev/null +++ b/tests/unit/test_base_component.py @@ -0,0 +1,372 @@ +"""Unit tests for BaseComponent lifecycle and dependency injection.""" + +# pylint: disable=protected-access,missing-function-docstring,missing-class-docstring,attribute-defined-outside-init + +import asyncio +import os +import tempfile + +import pytest + +from flowllm.components.base_component import BaseComponent, ComponentMixin, Dependency +from flowllm.enumeration import ComponentEnum + + +# Test subclasses + + +class StubComponent(BaseComponent): + component_type = ComponentEnum.FILE_CHUNKER + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.start_count = 0 + self.close_count = 0 + + async def _start(self): + self.start_count += 1 + + async def _close(self): + self.close_count += 1 + + +class DepTarget(BaseComponent): + component_type = ComponentEnum.KEYWORD_INDEX + + +class RequiredDepTarget(BaseComponent): + component_type = ComponentEnum.FILE_GRAPH + + +class FailCloseComponent(StubComponent): + async def _close(self): + await super()._close() + raise RuntimeError("close failed") + + +# Dependency + + +def test_dependency_repr_optional(): + dep = Dependency(ComponentEnum.FILE_CHUNKER, "my_parser", optional=True) + assert "?" in repr(dep) + assert "file_chunker" in repr(dep) + + +def test_dependency_repr_required(): + dep = Dependency(ComponentEnum.FILE_CHUNKER, "my_parser", optional=False) + assert "?" not in repr(dep) + + +def test_dependency_getattr_raises(): + dep = Dependency(ComponentEnum.FILE_CHUNKER, "my_parser") + with pytest.raises(RuntimeError, match="accessed before start"): + _ = dep.some_method + + +# bind + + +def test_bind_returns_none_for_empty_name(): + result = BaseComponent.bind(None, DepTarget) + assert result is None + + result = BaseComponent.bind("", DepTarget) + assert result is None + + +def test_bind_returns_dependency_placeholder(): + result = BaseComponent.bind("my_index", DepTarget) + assert isinstance(result, Dependency) + assert result.ctype == ComponentEnum.KEYWORD_INDEX + assert result.name == "my_index" + + +def test_bind_rejects_base_component_type(): + class BadTarget(BaseComponent): + component_type = ComponentEnum.BASE + + with pytest.raises(TypeError, match="non-BASE"): + BaseComponent.bind("x", BadTarget) + + +def test_bind_rejects_no_component_type(): + class NoType: + pass + + with pytest.raises(TypeError, match="non-BASE"): + BaseComponent.bind("x", NoType) + + +def test_bind_with_default_factory(): + def factory(): + return DepTarget(name="default") + + result = BaseComponent.bind("idx", DepTarget, default_factory=factory) + assert isinstance(result, Dependency) + assert result.default_factory is factory + + +def test_bind_optional_flag(): + dep = BaseComponent.bind("idx", DepTarget, optional=False) + assert dep.optional is False + + +# dependencies property + + +def test_dependencies_lists_unresolved(): + comp = StubComponent() + comp.dep1 = Dependency(ComponentEnum.KEYWORD_INDEX, "a") + comp.dep2 = Dependency(ComponentEnum.FILE_GRAPH, "b") + comp.normal_attr = "not a dep" + deps = comp.dependencies + assert len(deps) == 2 + + +# lifecycle + + +def test_start_close_idempotent(): + async def run(): + comp = StubComponent() + await comp.start() + await comp.start() + assert comp.start_count == 1 + assert comp.is_started is True + + await comp.close() + await comp.close() + assert comp.close_count == 1 + assert comp.is_started is False + + asyncio.run(run()) + + +def test_restart(): + async def run(): + comp = StubComponent() + await comp.start() + await comp.restart() + assert comp.start_count == 2 + assert comp.close_count == 1 + assert comp.is_started is True + await comp.close() + + asyncio.run(run()) + + +def test_async_context_manager(): + async def run(): + comp = StubComponent() + async with comp as c: + assert c is comp + assert comp.is_started is True + assert comp.is_started is False + + asyncio.run(run()) + + +def test_close_closes_owned_when_parent_close_fails(): + async def run(): + owned = StubComponent(name="owned") + parent = FailCloseComponent(name="parent") + parent.dep = BaseComponent.bind( + "sub", + StubComponent, + default_factory=lambda: owned, + ) + await parent.start() + + with pytest.raises(RuntimeError, match="close failed"): + await parent.close() + + assert owned.is_started is False + assert owned.close_count == 1 + assert parent.is_started is False + + asyncio.run(run()) + + +# standalone resolution + + +def test_resolve_standalone_optional_becomes_none(): + async def run(): + comp = StubComponent() + comp.dep = BaseComponent.bind("idx", DepTarget) + await comp.start() + assert comp.dep is None + await comp.close() + + asyncio.run(run()) + + +def test_resolve_standalone_with_default_factory(): + async def run(): + comp = StubComponent() + comp.dep = BaseComponent.bind( + "idx", + DepTarget, + default_factory=lambda: DepTarget(name="auto"), + ) + await comp.start() + assert isinstance(comp.dep, DepTarget) + assert comp.dep.name == "auto" + assert comp.dep in comp._owned + await comp.close() + + asyncio.run(run()) + + +def test_resolve_standalone_required_no_factory_keeps_placeholder(): + async def run(): + comp = StubComponent() + comp.dep = BaseComponent.bind("idx", DepTarget, optional=False) + await comp.start() + assert isinstance(comp.dep, Dependency) + await comp.close() + + asyncio.run(run()) + + +# owned lifecycle cascade + + +def test_owned_components_started_and_closed(): + async def run(): + owned = StubComponent(name="owned") + parent = StubComponent(name="parent") + parent.dep = BaseComponent.bind( + "sub", + StubComponent, + default_factory=lambda: owned, + ) + await parent.start() + assert owned.is_started is True + + await parent.close() + assert owned.is_started is False + + asyncio.run(run()) + + +# context-bound resolution + + +def test_resolve_from_context(): + async def run(): + from flowllm.components.application_context import ApplicationContext + + target = DepTarget(name="real_index") + ctx = ApplicationContext() + ctx.components = {ComponentEnum.KEYWORD_INDEX: {"real_index": target}} + + comp = StubComponent(app_context=ctx) + comp.dep = BaseComponent.bind("real_index", DepTarget) + await comp.start() + assert comp.dep is target + await comp.close() + + asyncio.run(run()) + + +def test_resolve_from_context_optional_missing(): + async def run(): + from flowllm.components.application_context import ApplicationContext + + ctx = ApplicationContext() + ctx.components = {} + comp = StubComponent(app_context=ctx) + comp.dep = BaseComponent.bind("missing", DepTarget, optional=True) + await comp.start() + assert comp.dep is None + await comp.close() + + asyncio.run(run()) + + +def test_resolve_from_context_required_missing_raises(): + async def run(): + from flowllm.components.application_context import ApplicationContext + + ctx = ApplicationContext() + ctx.components = {} + comp = StubComponent(app_context=ctx) + comp.dep = BaseComponent.bind("missing", RequiredDepTarget, optional=False) + with pytest.raises(ValueError, match="not found"): + await comp.start() + + asyncio.run(run()) + + +# ComponentMixin paths + + +def test_workspace_path_no_context(): + mixin = ComponentMixin() + from pathlib import Path + + assert mixin.workspace_path == Path.cwd() + + +def test_to_workspace_relative_inside_workspace(): + with tempfile.TemporaryDirectory() as tmpdir: + old_cwd = os.getcwd() + os.chdir(tmpdir) + try: + mixin = ComponentMixin() + abs_path = mixin.workspace_path / "sub" / "file.md" + rel = mixin.to_workspace_relative(abs_path) + assert rel == str(abs_path.relative_to(mixin.workspace_path)) + finally: + os.chdir(old_cwd) + + +def test_to_workspace_relative_outside_workspace(): + mixin = ComponentMixin() + result = mixin.to_workspace_relative("/some/other/path") + assert result == "/some/other/path" + + +# workspace metadata paths + + +def test_workspace_metadata_path_no_context(): + comp = StubComponent() + assert comp.workspace_metadata_path.name == "metadata" + + +def test_component_metadata_path(): + comp = StubComponent() + assert comp.component_metadata_path.name == ComponentEnum.FILE_CHUNKER.value + + +if __name__ == "__main__": + print("\n=== BaseComponent Tests ===") + test_dependency_repr_optional() + test_dependency_repr_required() + test_dependency_getattr_raises() + test_bind_returns_none_for_empty_name() + test_bind_returns_dependency_placeholder() + test_bind_rejects_base_component_type() + test_bind_rejects_no_component_type() + test_bind_with_default_factory() + test_bind_optional_flag() + test_dependencies_lists_unresolved() + test_start_close_idempotent() + test_restart() + test_async_context_manager() + test_close_closes_owned_when_parent_close_fails() + test_resolve_standalone_optional_becomes_none() + test_resolve_standalone_with_default_factory() + test_resolve_standalone_required_no_factory_keeps_placeholder() + test_owned_components_started_and_closed() + test_resolve_from_context() + test_resolve_from_context_optional_missing() + test_resolve_from_context_required_missing_raises() + test_workspace_path_no_context() + test_to_workspace_relative_outside_workspace() + test_workspace_metadata_path_no_context() + test_component_metadata_path() + print("\n所有测试通过!") diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py new file mode 100644 index 0000000..e23cfa1 --- /dev/null +++ b/tests/unit/test_cli.py @@ -0,0 +1,46 @@ +"""Unit tests for CLI call_server function.""" + +from flowllm import application as flowllm_module + + +def test_call_server_passes_client_kwargs_to_client(monkeypatch, capsys): + """Verify call_server separates client kwargs from action payload.""" + seen = {} + + class FakeClient: + """Mock client for testing call_server.""" + + def __init__(self, **kwargs): + seen["client_kwargs"] = kwargs + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + return None + + async def __call__(self, action: str, **kwargs): + seen["action"] = action + seen["payload"] = kwargs + yield "ok" + + monkeypatch.setattr(flowllm_module.R, "get", lambda component_type, backend: FakeClient) + + async def run(): + await flowllm_module.call_server( + "search", + backend="http", + host="127.0.0.2", + port=2444, + timeout=1.5, + query="hello", + ) + + import asyncio + + asyncio.run(run()) + + assert seen["client_kwargs"] == {"host": "127.0.0.2", "port": 2444, "timeout": 1.5} + assert seen["action"] == "search" + assert seen["payload"] == {"query": "hello"} + assert capsys.readouterr().out == "ok\n" diff --git a/tests/unit/test_common_steps.py b/tests/unit/test_common_steps.py new file mode 100644 index 0000000..01e35c4 --- /dev/null +++ b/tests/unit/test_common_steps.py @@ -0,0 +1,94 @@ +"""Tests for common steps.""" + +# pylint: disable=protected-access + +import asyncio +import warnings + +from flowllm.components.agent_wrapper import BaseAgentWrapper +from flowllm.steps.common.add import AddStep +from flowllm.steps.common.health_check import _file_graph_status +from flowllm.steps.common.llm_demo import LLMDemoStep + +warnings.filterwarnings("ignore", category=DeprecationWarning, module="jieba") +warnings.filterwarnings("ignore", category=DeprecationWarning, module="pkg_resources") + + +def _run(coro): + """Run a coroutine synchronously.""" + asyncio.run(coro) + + +def test_add_step_coerces_numeric_inputs(): + """Numeric strings are coerced, not concatenated.""" + + async def run(): + step = AddStep() + resp = await step(a="1", b="2.5") + assert resp.success is True + assert resp.answer == "3.5" + assert resp.metadata["result"] == 3.5 + print("✓ test_add_step_coerces_numeric_inputs passed") + + _run(run()) + + +def test_add_step_rejects_invalid_inputs(): + """Invalid arguments return a failed response.""" + + async def run(): + step = AddStep() + resp = await step(a="one", b=2) + assert resp.success is False + assert "Invalid add arguments" in resp.answer + print("✓ test_add_step_rejects_invalid_inputs passed") + + _run(run()) + + +class _FakeAgentWrapper(BaseAgentWrapper): + """Captures reply kwargs without a real model.""" + + def __init__(self): + super().__init__() + self.last_kwargs = None + + async def reply(self, inputs, **kwargs) -> dict: + self.last_kwargs = kwargs + return {"result": "ok"} + + +def test_llm_demo_always_registers_add_tool(): + """Add job is always registered as a tool.""" + + async def run(): + wrapper = _FakeAgentWrapper() + step = LLMDemoStep() + resp = await step(query="hello", agent_wrapper=wrapper) + assert resp.success is True + assert wrapper.last_kwargs["job_tools"] == ["add"] + assert "job_tools" not in resp.metadata + print("✓ test_llm_demo_always_registers_add_tool passed") + + _run(run()) + + +def test_file_graph_health_reports_neo4j_cached_counts(): + """Neo4j graph reports cached counts correctly.""" + + class FakeNeo4jGraph: + """Neo4j graph stub with cached counters.""" + + is_started = True + _driver = object() + _uri = "bolt://example" + _database = "neo4j" + _n_nodes = 3 + _n_edges = 4 + _n_virtual = 1 + + status = _file_graph_status(FakeNeo4jGraph()) + assert status["n_nodes"] == 3 + assert status["n_edges"] == 4 + assert status["n_virtual"] == 1 + print("✓ test_file_graph_health_reports_neo4j_cached_counts passed") diff --git a/tests/unit/test_component_registry.py b/tests/unit/test_component_registry.py new file mode 100644 index 0000000..9b2177a --- /dev/null +++ b/tests/unit/test_component_registry.py @@ -0,0 +1,151 @@ +"""Tests for ComponentRegistry.""" + +# pylint: disable=missing-function-docstring,missing-class-docstring,protected-access,unused-argument + +import pytest + +from flowllm.components.base_component import BaseComponent +from flowllm.components.component_registry import ComponentRegistry +from flowllm.enumeration import ComponentEnum + + +class _DummyComponent(BaseComponent): + component_type = ComponentEnum.FILE_CHUNKER + + +class _AnotherComponent(BaseComponent): + component_type = ComponentEnum.KEYWORD_INDEX + + +class _NoComponentType: + pass + + +class _BaseComponentType(BaseComponent): + component_type = ComponentEnum.BASE + + +# register & get + + +def test_register_direct_with_explicit_name(): + reg = ComponentRegistry() + reg.register(_DummyComponent, "my_parser") + assert reg.get(ComponentEnum.FILE_CHUNKER, "my_parser") is _DummyComponent + + +def test_register_direct_defaults_to_class_name(): + reg = ComponentRegistry() + reg.register(_DummyComponent) + assert reg.get(ComponentEnum.FILE_CHUNKER, "_DummyComponent") is _DummyComponent + + +def test_register_decorator(): + reg = ComponentRegistry() + + @reg.register("alias") + class MyParser(BaseComponent): + component_type = ComponentEnum.FILE_CHUNKER + + assert reg.get(ComponentEnum.FILE_CHUNKER, "alias") is MyParser + + +def test_register_overwrite_warns(caplog): + reg = ComponentRegistry() + reg.register(_DummyComponent, "dup") + reg.register(_DummyComponent, "dup") + assert reg.get(ComponentEnum.FILE_CHUNKER, "dup") is _DummyComponent + + +def test_register_rejects_missing_component_type(): + reg = ComponentRegistry() + with pytest.raises(TypeError, match="ComponentEnum"): + reg.register(_NoComponentType, "bad") + + +def test_register_rejects_empty_name(): + reg = ComponentRegistry() + with pytest.raises(ValueError, match="empty"): + reg._do_register(_DummyComponent, "") + + +def test_register_rejects_non_class_non_string(): + reg = ComponentRegistry() + with pytest.raises(TypeError, match="Expected a class or string"): + reg.register(42) + + +# get_all + + +def test_get_all_returns_copy(): + reg = ComponentRegistry() + reg.register(_DummyComponent, "a") + reg.register(_AnotherComponent, "b") + + parsers = reg.get_all(ComponentEnum.FILE_CHUNKER) + assert parsers == {"a": _DummyComponent} + + indexes = reg.get_all(ComponentEnum.KEYWORD_INDEX) + assert indexes == {"b": _AnotherComponent} + + # Mutating the copy doesn't affect the registry. + parsers["hacked"] = _DummyComponent + assert "hacked" not in reg.get_all(ComponentEnum.FILE_CHUNKER) + + +def test_get_all_unknown_type_returns_empty(): + reg = ComponentRegistry() + assert not reg.get_all(ComponentEnum.AS_LLM) + + +# get (miss) + + +def test_get_nonexistent_returns_none(): + reg = ComponentRegistry() + assert reg.get(ComponentEnum.FILE_CHUNKER, "nope") is None + + +# unregister + + +def test_unregister_existing(): + reg = ComponentRegistry() + reg.register(_DummyComponent, "x") + assert reg.unregister(ComponentEnum.FILE_CHUNKER, "x") is True + assert reg.get(ComponentEnum.FILE_CHUNKER, "x") is None + + +def test_unregister_missing_returns_false(): + reg = ComponentRegistry() + assert reg.unregister(ComponentEnum.FILE_CHUNKER, "nope") is False + + +# clear + + +def test_clear(): + reg = ComponentRegistry() + reg.register(_DummyComponent, "a") + reg.register(_AnotherComponent, "b") + reg.clear() + assert not reg.get_all(ComponentEnum.FILE_CHUNKER) + assert not reg.get_all(ComponentEnum.KEYWORD_INDEX) + + +if __name__ == "__main__": + print("\n=== ComponentRegistry Tests ===") + test_register_direct_with_explicit_name() + test_register_direct_defaults_to_class_name() + test_register_decorator() + test_register_rejects_missing_component_type() + test_register_rejects_empty_name() + test_register_rejects_non_class_non_string() + test_get_all_returns_copy() + test_get_all_unknown_type_returns_empty() + test_get_nonexistent_returns_none() + test_unregister_existing() + test_unregister_missing_returns_false() + test_clear() + print("\n所有测试通过!") diff --git a/tests/unit/test_config_parser.py b/tests/unit/test_config_parser.py new file mode 100644 index 0000000..b503b43 --- /dev/null +++ b/tests/unit/test_config_parser.py @@ -0,0 +1,66 @@ +"""Tests for config parsing.""" + +from pathlib import Path + +import pytest + +from flowllm.config.config_parser import ( + _expand_env_vars, + _load_config, + _read_config_file, + parse_args, + parse_dot_notation, +) + + +def test_load_builtin_config_by_filename_with_suffix(): + """Config names may include the YAML suffix.""" + cfg = _load_config("default.yaml") + + assert cfg["service"]["backend"] == "http" + + +def test_parse_args_rejects_non_key_value_extra_argument(): + """Extra arguments must use key=value syntax.""" + with pytest.raises(ValueError, match="expected key=value"): + parse_args("search", "hello") + + +@pytest.mark.parametrize("item", ["=1", ".a=1", "a.=1", "a..b=1"]) +def test_parse_dot_notation_rejects_empty_key_segments(item): + """Empty path segments are rejected.""" + with pytest.raises(ValueError, match="Invalid dot notation key"): + parse_dot_notation([item]) + + +def test_read_config_file_rejects_non_mapping_root(tmp_path: Path): + """Non-mapping root is rejected.""" + config_path = tmp_path / "bad.yaml" + config_path.write_text("- item\n", encoding="utf-8") + + with pytest.raises(ValueError, match="Config root must be a mapping"): + _read_config_file(config_path) + + +def test_expand_env_vars_converts_expanded_scalar_types(monkeypatch): + """Expanded env values keep YAML scalar typing.""" + monkeypatch.setenv("PORT", "18080") + monkeypatch.setenv("ENABLED", "false") + + expanded = _expand_env_vars( + { + "port": "${PORT}", + "enabled": "${ENABLED}", + "zip": "${ZIP:-007}", + "url": "http://${HOST:-localhost}:${PORT}", + "string_bool": '${STRING_BOOL:-"false"}', + }, + ) + + assert expanded == { + "port": 18080, + "enabled": False, + "zip": "007", + "url": "http://localhost:18080", + "string_bool": "false", + } diff --git a/tests/unit/test_runtime_context.py b/tests/unit/test_runtime_context.py new file mode 100644 index 0000000..f9b5139 --- /dev/null +++ b/tests/unit/test_runtime_context.py @@ -0,0 +1,181 @@ +"""Tests for RuntimeContext.""" + +# pylint: disable=protected-access,missing-function-docstring + +import asyncio + +import pytest + +from flowllm.components.runtime_context import RuntimeContext +from flowllm.enumeration import ChunkEnum + + +# dict-like access + + +def test_getitem_setitem(): + ctx = RuntimeContext(foo="bar") + assert ctx["foo"] == "bar" + ctx["baz"] = 42 + assert ctx["baz"] == 42 + + +def test_getitem_missing_raises(): + ctx = RuntimeContext() + with pytest.raises(KeyError): + _ = ctx["nope"] + + +def test_contains(): + ctx = RuntimeContext(a=1) + assert "a" in ctx + assert "b" not in ctx + + +def test_delitem(): + ctx = RuntimeContext(a=1) + del ctx["a"] + assert "a" not in ctx + + +def test_get_with_default(): + ctx = RuntimeContext(a=1) + assert ctx.get("a") == 1 + assert ctx.get("b", "fallback") == "fallback" + assert ctx.get("b") is None + + +def test_update_merges_and_returns_self(): + ctx = RuntimeContext(a=1) + result = ctx.update({"b": 2, "c": 3}) + assert result is ctx + assert ctx["b"] == 2 + assert ctx["c"] == 3 + + +# from_context + + +def test_from_context_creates_new_when_none(): + ctx = RuntimeContext.from_context(None, x=10) + assert ctx["x"] == 10 + + +def test_from_context_reuses_existing(): + original = RuntimeContext(a=1) + reused = RuntimeContext.from_context(original, b=2) + assert reused is original + assert reused["a"] == 1 + assert reused["b"] == 2 + + +# apply_mapping + + +def test_apply_mapping_copies_values(): + ctx = RuntimeContext(src="hello") + result = ctx.apply_mapping({"src": "dst"}) + assert result is ctx + assert ctx["dst"] == "hello" + assert ctx["src"] == "hello" + + +def test_apply_mapping_skips_missing_source(): + ctx = RuntimeContext(a=1) + ctx.apply_mapping({"missing_key": "target"}) + assert "target" not in ctx + + +def test_apply_mapping_empty_is_noop(): + ctx = RuntimeContext(a=1) + result = ctx.apply_mapping({}) + assert result is ctx + + +# streaming + + +def test_stream_property(): + ctx_no_queue = RuntimeContext() + assert ctx_no_queue.stream is False + + ctx_with_queue = RuntimeContext(stream_queue=asyncio.Queue()) + assert ctx_with_queue.stream is True + + +def test_enqueue_raises_without_queue(): + async def run(): + ctx = RuntimeContext() + with pytest.raises(RuntimeError, match="Stream queue not initialized"): + await ctx._enqueue(None) + + asyncio.run(run()) + + +def test_add_stream_string(): + async def run(): + q = asyncio.Queue() + ctx = RuntimeContext(stream_queue=q) + result = await ctx.add_stream_string("hello", ChunkEnum.CONTENT) + assert result is ctx + + chunk = q.get_nowait() + assert chunk.chunk == "hello" + assert chunk.chunk_type == ChunkEnum.CONTENT + assert chunk.done is False + + asyncio.run(run()) + + +def test_add_stream_done(): + async def run(): + q = asyncio.Queue() + ctx = RuntimeContext(stream_queue=q) + result = await ctx.add_stream_done() + assert result is ctx + + chunk = q.get_nowait() + assert chunk.chunk_type == ChunkEnum.DONE + assert chunk.done is True + + asyncio.run(run()) + + +# response + + +def test_default_response(): + ctx = RuntimeContext() + assert ctx.response.success is True + assert ctx.response.answer == "" + + +def test_custom_response(): + from flowllm.schema import Response + + resp = Response(answer="ok", success=False) + ctx = RuntimeContext(response=resp) + assert ctx.response is resp + assert ctx.response.success is False + + +if __name__ == "__main__": + print("\n=== RuntimeContext Tests ===") + test_getitem_setitem() + test_getitem_missing_raises() + test_contains() + test_delitem() + test_get_with_default() + test_update_merges_and_returns_self() + test_from_context_creates_new_when_none() + test_from_context_reuses_existing() + test_apply_mapping_copies_values() + test_apply_mapping_skips_missing_source() + test_apply_mapping_empty_is_noop() + test_stream_property() + test_enqueue_raises_without_queue() + test_add_stream_string() + test_add_stream_done() + test_default_response() + test_custom_response() + print("\n所有测试通过!") diff --git a/tests/unit/test_service.py b/tests/unit/test_service.py new file mode 100644 index 0000000..c197df5 --- /dev/null +++ b/tests/unit/test_service.py @@ -0,0 +1,43 @@ +"""Tests for service job registration.""" + +from types import SimpleNamespace + +from flowllm.components.job import BaseJob, StreamJob +from flowllm.components.service import MCPService + + +def _dummy_app(): + """Minimal MCPService.build_service stub.""" + + async def start(): + return None + + async def close(): + return None + + return SimpleNamespace( + config=SimpleNamespace(app_name="test"), + context=SimpleNamespace(metadata={}), + start=start, + close=close, + ) + + +def test_mcp_service_registers_job_with_empty_parameters(): + """Empty parameters remain a dict for FastMCP validation.""" + service = MCPService() + service.build_service(_dummy_app()) + + job = BaseJob(name="empty_params", parameters={}) + + assert service.add_job(job) is True + + +def test_mcp_service_reports_stream_job_skipped(): + """StreamJob tools are not exposed by MCPService.""" + service = MCPService() + service.build_service(_dummy_app()) + + job = StreamJob(name="stream") + + assert service.add_job(job) is False diff --git a/tests_op/test_dashscope_search_op.py b/tests_op/test_dashscope_search_op.py deleted file mode 100644 index 973f42b..0000000 --- a/tests_op/test_dashscope_search_op.py +++ /dev/null @@ -1,22 +0,0 @@ -"""Test script for DashscopeSearchOp. - -This script provides test cases for DashscopeSearchOp class. -It can be run directly with: python test_dashscope_search_op.py -""" - -import asyncio - -from flowllm.gallery.dashscope_search_op import DashscopeSearchOp -from flowllm.main import FlowLLMApp - - -async def async_main(): - """Test function for DashscopeSearchOp.""" - async with FlowLLMApp(): - op = DashscopeSearchOp(model="qwen3-max", enable_role_prompt=True) - await op.async_call(query="藏格矿业的业务主要有哪几块？营收和利润的角度分析 雪球") - print(op.output) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_edit_op.py b/tests_op/test_edit_op.py deleted file mode 100644 index ffdc2ad..0000000 --- a/tests_op/test_edit_op.py +++ /dev/null @@ -1,119 +0,0 @@ -"""Test script for EditOp. - -This script provides test cases for EditOp class. -It can be run directly with: python test_edit_op.py -""" - -import asyncio -import json -import tempfile -from pathlib import Path - -from flowllm.extensions.file_tool import EditOp -from flowllm.main import FlowLLMApp - - -async def test_simple_input_schema(): - """Test case 1: Test tool_call.simple_input_dump().""" - print("=" * 80) - print("Test Case 1: Testing tool_call.simple_input_dump()") - print("=" * 80) - - async with FlowLLMApp(): - op = EditOp() - tool_call = op.build_tool_call() - tool_call.name = op.short_name - - # Test simple_input_dump - input_schema = tool_call.simple_input_dump() - print("\nInput Schema (simple_input_dump):") - print(json.dumps(input_schema, indent=2, ensure_ascii=False)) - - # Verify structure - assert "type" in input_schema - assert "function" in input_schema - assert "name" in input_schema["function"] - assert "description" in input_schema["function"] - assert "parameters" in input_schema["function"] - assert "properties" in input_schema["function"]["parameters"] - assert "required" in input_schema["function"]["parameters"] - - print("\n✓ Test Case 1 passed: simple_input_dump() works correctly") - print("=" * 80) - - -async def test_normal_async_execute(): - """Test case 2: Test normal async_execute.""" - print("\n" + "=" * 80) - print("Test Case 2: Testing normal async_execute") - print("=" * 80) - - async with FlowLLMApp(): - # Create a temporary file for testing - with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: - test_file = Path(f.name) - f.write("Hello, world!\nThis is a test file.\nLine 3\n") - - try: - op = EditOp() - # Test replacing text - await op.async_call( - file_path=str(test_file), - old_string="Hello, world!", - new_string="Hello, Python!", - ) - - print("\nOutput:") - print(op.output) - - # Verify the file was modified - content = test_file.read_text(encoding="utf-8") - assert "Hello, Python!" in content - assert "Hello, world!" not in content - - print("\n✓ Test Case 2 passed: async_execute() works correctly") - print("=" * 80) - finally: - # Clean up - if test_file.exists(): - test_file.unlink() - - -async def test_exception_async_default_execute(): - """Test case 3: Test exception handling, executing async_default_execute.""" - print("\n" + "=" * 80) - print("Test Case 3: Testing exception handling (async_default_execute)") - print("=" * 80) - - async with FlowLLMApp(): - op = EditOp() - # Test with a non-existent file (should trigger exception) - await op.async_call( - file_path="/nonexistent/path/file.txt", - old_string="test", - new_string="test2", - ) - - print("\nOutput after exception:") - print(op.output) - - # Verify that async_default_execute was called - assert "Failed to edit file" in op.output - - print("\n✓ Test Case 3 passed: async_default_execute() was called on exception") - print("=" * 80) - - -async def async_main(): - """Run all test cases.""" - await test_simple_input_schema() - await test_normal_async_execute() - await test_exception_async_default_execute() - - print("\n" + "=" * 80) - print("All tests passed!") - print("=" * 80) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_execute_code_op.py b/tests_op/test_execute_code_op.py deleted file mode 100644 index 6c0e1b8..0000000 --- a/tests_op/test_execute_code_op.py +++ /dev/null @@ -1,29 +0,0 @@ -"""Test script for ExecuteCodeOp. - -This script moves the inline main tests from gallery/execute_code_op.py -into the tests directory so they can be executed independently. -""" - -import asyncio - -from flowllm.gallery.execute_code_op import ExecuteCodeOp -from flowllm.main import FlowLLMApp - - -async def async_main(): - """Test function for ExecuteCodeOp.""" - async with FlowLLMApp(): - op = ExecuteCodeOp() - print(op.tool_call.model_dump_json(exclude_none=True)) - print(op.tool_call.simple_input_dump()) - print(op.tool_call.simple_output_dump()) - - await op.async_call(code="print('Hello World')") - print(op.output) - - await op.async_call(code="print('Hello World!'") - print(op.output) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_exit_plan_mode_op.py b/tests_op/test_exit_plan_mode_op.py deleted file mode 100644 index 32ce5ca..0000000 --- a/tests_op/test_exit_plan_mode_op.py +++ /dev/null @@ -1,94 +0,0 @@ -"""Test script for ExitPlanModeOp. - -This script provides test cases for ExitPlanModeOp class. -It can be run directly with: python test_exit_plan_mode_op.py -""" - -import asyncio -import json - -from flowllm.extensions.file_tool import ExitPlanModeOp -from flowllm.main import FlowLLMApp - - -async def test_simple_input_schema(): - """Test case 1: Test tool_call.simple_input_dump().""" - print("=" * 80) - print("Test Case 1: Testing tool_call.simple_input_dump()") - print("=" * 80) - - async with FlowLLMApp(): - op = ExitPlanModeOp() - tool_call = op.build_tool_call() - tool_call.name = op.short_name - - # Test simple_input_dump - input_schema = tool_call.simple_input_dump() - print("\nInput Schema (simple_input_dump):") - print(json.dumps(input_schema, indent=2, ensure_ascii=False)) - - # Verify structure - assert "type" in input_schema - assert "function" in input_schema - assert "name" in input_schema["function"] - assert "description" in input_schema["function"] - assert "parameters" in input_schema["function"] - assert "properties" in input_schema["function"]["parameters"] - assert "required" in input_schema["function"]["parameters"] - - print("\n✓ Test Case 1 passed: simple_input_dump() works correctly") - print("=" * 80) - - -async def test_normal_async_execute(): - """Test case 2: Test normal async_execute.""" - print("\n" + "=" * 80) - print("Test Case 2: Testing normal async_execute") - print("=" * 80) - - async with FlowLLMApp(): - op = ExitPlanModeOp() - # Test with a valid plan - test_plan = """ -1. First step: Implement feature A -2. Second step: Add tests for feature A -3. Third step: Update documentation -""" - await op.async_call(plan=test_plan) - - print("\nOutput:") - print(op.output) - print("\n✓ Test Case 2 passed: async_execute() works correctly") - print("=" * 80) - - -async def test_exception_async_default_execute(): - """Test case 3: Test exception handling, executing async_default_execute.""" - print("\n" + "=" * 80) - print("Test Case 3: Testing exception handling (async_default_execute)") - print("=" * 80) - - async with FlowLLMApp(): - op = ExitPlanModeOp() - # Test with empty plan, which should trigger ValueError - await op.async_call(plan="") - - print("\nOutput after exception:") - print(op.output) - print("\n✓ Test Case 3 passed: async_default_execute() was called on exception") - print("=" * 80) - - -async def async_main(): - """Run all test cases.""" - await test_simple_input_schema() - await test_normal_async_execute() - await test_exception_async_default_execute() - - print("\n" + "=" * 80) - print("All tests passed!") - print("=" * 80) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_gen_system_prompt_op.py b/tests_op/test_gen_system_prompt_op.py deleted file mode 100644 index e34c6ad..0000000 --- a/tests_op/test_gen_system_prompt_op.py +++ /dev/null @@ -1,122 +0,0 @@ -"""Test script for GenSystemPromptOp. - -This script provides test cases for GenSystemPromptOp class. -It can be run directly with: python test_gen_system_prompt_op.py -""" - -import asyncio - -from flowllm.core.enumeration import Role -from flowllm.core.schema import Message -from flowllm.core.utils import format_messages -from flowllm.gallery.gen_system_prompt_op import GenSystemPromptOp -from flowllm.main import FlowLLMApp - - -async def test_query_parameter(): - """Test with query parameter.""" - print("=" * 80) - print("Test 1: Testing with query parameter") - print("=" * 80) - - op = GenSystemPromptOp() - await op.async_call(query="What is Python?") - - print(f"Query: {op.context.query}") - print(f"System Prompt: {op.context.system_prompt}") - print(f"System Prompt type: {type(op.context.system_prompt)}") - assert op.context.system_prompt is not None, "System prompt should not be None" - assert isinstance(op.context.system_prompt, str), "System prompt should be a string" - assert len(op.context.system_prompt) > 0, "System prompt should not be empty" - print("✓ Test 1 passed\n") - - -async def test_messages_parameter(): - """Test with messages parameter.""" - print("=" * 80) - print("Test 2: Testing with messages parameter") - print("=" * 80) - - op2 = GenSystemPromptOp() - messages = [ - Message(role=Role.USER, content="I'm feeling frustrated with my code"), - Message(role=Role.ASSISTANT, content="I understand. Let's work through this together."), - Message(role=Role.USER, content="It keeps giving me errors"), - ] - await op2.async_call(messages=[msg.model_dump() for msg in messages]) - - print(f"Messages count: {len(op2.context.messages)}") - print(f"System Prompt: {op2.context.system_prompt}") - assert op2.context.system_prompt is not None, "System prompt should not be None" - assert isinstance(op2.context.system_prompt, str), "System prompt should be a string" - assert len(op2.context.system_prompt) > 0, "System prompt should not be empty" - print("✓ Test 2 passed\n") - - -async def test_format_messages(): - """Test format_messages utility function.""" - print("=" * 80) - print("Test 3: Testing format_messages utility function") - print("=" * 80) - - test_messages = [ - Message(role=Role.USER, content="Hello"), - Message(role=Role.ASSISTANT, content="Hi there!"), - ] - formatted = format_messages(test_messages) - print(f"Formatted messages:\n{formatted}") - assert ( - "user: Hello".lower() in formatted.lower() - ), f"Formatted messages should contain user message. formatted={formatted}" - assert ( - "assistant: Hi there!".lower() in formatted.lower() - ), f"Formatted messages should contain assistant message. formatted={formatted}" - print("✓ Test 3 passed\n") - - -async def test_empty_input(): - """Test with empty query and empty messages (should raise assertion error).""" - print("=" * 80) - print("Test 4: Testing with empty query and empty messages (should raise assertion)") - print("=" * 80) - - op4 = GenSystemPromptOp() - try: - await op4.async_call() - print("✗ Test 4 failed: Should have raised an assertion error") - except AssertionError as e: - print(f"✓ Test 4 passed: Correctly raised AssertionError: {e}\n") - - -async def test_technical_query(): - """Test with technical query.""" - print("=" * 80) - print("Test 5: Testing with technical query") - print("=" * 80) - - op5 = GenSystemPromptOp() - await op5.async_call(query="How do I implement a binary search tree in Python?") - - print(f"Query: {op5.context.query}") - print(f"System Prompt: {op5.context.system_prompt}") - assert op5.context.system_prompt is not None, "System prompt should not be None" - assert len(op5.context.system_prompt) > 0, "System prompt should not be empty" - print("✓ Test 5 passed\n") - - -async def async_main(): - """Test function for GenSystemPromptOp.""" - async with FlowLLMApp(): - await test_query_parameter() - await test_messages_parameter() - await test_format_messages() - await test_empty_input() - await test_technical_query() - - print("=" * 80) - print("All tests passed!") - print("=" * 80) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_glob_op.py b/tests_op/test_glob_op.py deleted file mode 100644 index 63f2aab..0000000 --- a/tests_op/test_glob_op.py +++ /dev/null @@ -1,88 +0,0 @@ -"""Test script for GlobOp. - -This script provides test cases for GlobOp class. -It can be run directly with: python test_glob_op.py -""" - -import asyncio -import json - -from flowllm.extensions.file_tool import GlobOp -from flowllm.main import FlowLLMApp - - -async def test_simple_input_schema(): - """Test case 1: Test tool_call.simple_input_dump().""" - print("=" * 80) - print("Test Case 1: Testing tool_call.simple_input_dump()") - print("=" * 80) - - async with FlowLLMApp(): - op = GlobOp() - tool_call = op.build_tool_call() - tool_call.name = op.short_name - - # Test simple_input_dump - input_schema = tool_call.simple_input_dump() - print("\nInput Schema (simple_input_dump):") - print(json.dumps(input_schema, indent=2, ensure_ascii=False)) - - # Verify structure - assert "type" in input_schema - assert "function" in input_schema - assert "name" in input_schema["function"] - assert "description" in input_schema["function"] - assert "parameters" in input_schema["function"] - assert "properties" in input_schema["function"]["parameters"] - assert "required" in input_schema["function"]["parameters"] - - print("\n✓ Test Case 1 passed: simple_input_dump() works correctly") - print("=" * 80) - - -async def test_normal_async_execute(): - """Test case 2: Test normal async_execute.""" - print("\n" + "=" * 80) - print("Test Case 2: Testing normal async_execute") - print("=" * 80) - - async with FlowLLMApp(): - op = GlobOp() - # Test with a pattern that should find files - await op.async_call(pattern="*.py", dir_path=".") - - print("\nOutput:") - print(op.output) - print("\n✓ Test Case 2 passed: async_execute() works correctly") - print("=" * 80) - - -async def test_exception_async_default_execute(): - """Test case 3: Test exception handling, executing async_default_execute.""" - print("\n" + "=" * 80) - print("Test Case 3: Testing exception handling (async_default_execute)") - print("=" * 80) - - async with FlowLLMApp(): - op = GlobOp() - await op.async_call(pattern="*xxxx.py") - - print("\nOutput after exception:") - print(op.output) - print("\n✓ Test Case 3 passed: async_default_execute() was called on exception") - print("=" * 80) - - -async def async_main(): - """Run all test cases.""" - await test_simple_input_schema() - await test_normal_async_execute() - await test_exception_async_default_execute() - - print("\n" + "=" * 80) - print("All tests passed!") - print("=" * 80) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_grep_op.py b/tests_op/test_grep_op.py deleted file mode 100644 index 0df84cc..0000000 --- a/tests_op/test_grep_op.py +++ /dev/null @@ -1,89 +0,0 @@ -"""Test script for GrepOp. - -This script provides test cases for GrepOp class. -It can be run directly with: python test_grep_op.py -""" - -import asyncio -import json - -from flowllm.extensions.file_tool import GrepOp -from flowllm.main import FlowLLMApp - - -async def test_simple_input_schema(): - """Test case 1: Test tool_call.simple_input_dump().""" - print("=" * 80) - print("Test Case 1: Testing tool_call.simple_input_dump()") - print("=" * 80) - - async with FlowLLMApp(): - op = GrepOp() - tool_call = op.build_tool_call() - tool_call.name = op.short_name - - # Test simple_input_dump - input_schema = tool_call.simple_input_dump() - print("\nInput Schema (simple_input_dump):") - print(json.dumps(input_schema, indent=2, ensure_ascii=False)) - - # Verify structure - assert "type" in input_schema - assert "function" in input_schema - assert "name" in input_schema["function"] - assert "description" in input_schema["function"] - assert "parameters" in input_schema["function"] - assert "properties" in input_schema["function"]["parameters"] - assert "required" in input_schema["function"]["parameters"] - - print("\n✓ Test Case 1 passed: simple_input_dump() works correctly") - print("=" * 80) - - -async def test_normal_async_execute(): - """Test case 2: Test normal async_execute.""" - print("\n" + "=" * 80) - print("Test Case 2: Testing normal async_execute") - print("=" * 80) - - async with FlowLLMApp(): - op = GrepOp() - # Test with a pattern that should find matches (search for "def" in Python files) - await op.async_call(pattern="def ", glob="*.py") - - print("\nOutput:") - print(op.output) - print("\n✓ Test Case 2 passed: async_execute() works correctly") - print("=" * 80) - - -async def test_exception_async_default_execute(): - """Test case 3: Test exception handling, executing async_default_execute.""" - print("\n" + "=" * 80) - print("Test Case 3: Testing exception handling (async_default_execute)") - print("=" * 80) - - async with FlowLLMApp(): - op = GrepOp() - # Test with an invalid path to trigger exception - await op.async_call(pattern="test", path="/nonexistent/path/that/does/not/exist") - - print("\nOutput after exception:") - print(op.output) - print("\n✓ Test Case 3 passed: async_default_execute() was called on exception") - print("=" * 80) - - -async def async_main(): - """Run all test cases.""" - await test_simple_input_schema() - await test_normal_async_execute() - await test_exception_async_default_execute() - - print("\n" + "=" * 80) - print("All tests passed!") - print("=" * 80) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_ls_op.py b/tests_op/test_ls_op.py deleted file mode 100644 index 679d1c2..0000000 --- a/tests_op/test_ls_op.py +++ /dev/null @@ -1,89 +0,0 @@ -"""Test script for LSOp. - -This script provides test cases for LSOp class. -It can be run directly with: python test_ls_op.py -""" - -import asyncio -import json - -from flowllm.extensions.file_tool import LSOp -from flowllm.main import FlowLLMApp - - -async def test_simple_input_schema(): - """Test case 1: Test tool_call.simple_input_dump().""" - print("=" * 80) - print("Test Case 1: Testing tool_call.simple_input_dump()") - print("=" * 80) - - async with FlowLLMApp(): - op = LSOp() - tool_call = op.build_tool_call() - tool_call.name = op.short_name - - # Test simple_input_dump - input_schema = tool_call.simple_input_dump() - print("\nInput Schema (simple_input_dump):") - print(json.dumps(input_schema, indent=2, ensure_ascii=False)) - - # Verify structure - assert "type" in input_schema - assert "function" in input_schema - assert "name" in input_schema["function"] - assert "description" in input_schema["function"] - assert "parameters" in input_schema["function"] - assert "properties" in input_schema["function"]["parameters"] - assert "required" in input_schema["function"]["parameters"] - - print("\n✓ Test Case 1 passed: simple_input_dump() works correctly") - print("=" * 80) - - -async def test_normal_async_execute(): - """Test case 2: Test normal async_execute.""" - print("\n" + "=" * 80) - print("Test Case 2: Testing normal async_execute") - print("=" * 80) - - async with FlowLLMApp(): - op = LSOp() - # Test with a path that should exist - await op.async_call(path=".") - - print("\nOutput:") - print(op.output) - print("\n✓ Test Case 2 passed: async_execute() works correctly") - print("=" * 80) - - -async def test_exception_async_default_execute(): - """Test case 3: Test exception handling, executing async_default_execute.""" - print("\n" + "=" * 80) - print("Test Case 3: Testing exception handling (async_default_execute)") - print("=" * 80) - - async with FlowLLMApp(): - op = LSOp() - # Test with a non-existent path to trigger exception - await op.async_call(path="/nonexistent/path/that/does/not/exist") - - print("\nOutput after exception:") - print(op.output) - print("\n✓ Test Case 3 passed: async_default_execute() was called on exception") - print("=" * 80) - - -async def async_main(): - """Run all test cases.""" - await test_simple_input_schema() - await test_normal_async_execute() - await test_exception_async_default_execute() - - print("\n" + "=" * 80) - print("All tests passed!") - print("=" * 80) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_mock_search_op.py b/tests_op/test_mock_search_op.py deleted file mode 100644 index 3630cec..0000000 --- a/tests_op/test_mock_search_op.py +++ /dev/null @@ -1,47 +0,0 @@ -"""Test script for MockSearchOp. - -This script provides a simple test case for MockSearchOp class. -It can be run directly with: python test_mock_search_op.py -""" - -import asyncio -import json - -from flowllm.gallery.mock_search_op import MockSearchOp -from flowllm.main import FlowLLMApp - - -async def async_main(): - """Test function for MockSearchOp.""" - async with FlowLLMApp(): - # Single test: run mock search with a query and validate output - print("=" * 80) - print("Test: Running MockSearchOp with a query and validating JSON output") - print("=" * 80) - - op = MockSearchOp() - await op.async_call(query="latest AI research news") - - print("Output (raw):") - print(op.context.response.answer) - - # Basic assertions - assert op.context.response.answer is not None, "Output should not be None" - assert isinstance(op.context.response.answer, str), "Output should be a string" - assert len(op.context.response.answer) > 0, "Output should not be empty" - - # Validate it is valid JSON - try: - parsed = json.loads(op.context.response.answer) - assert isinstance(parsed, (list, dict)), "Parsed JSON should be list or dict" - except json.JSONDecodeError as e: - raise AssertionError(f"Output should be valid JSON, got decode error: {e}") from e - - print("\n✓ Test passed\n") - print("=" * 80) - print("All tests passed!") - print("=" * 80) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_react_search_op.py b/tests_op/test_react_search_op.py deleted file mode 100644 index 6504c2d..0000000 --- a/tests_op/test_react_search_op.py +++ /dev/null @@ -1,22 +0,0 @@ -"""Test script for ReactSearchOp. - -This script provides test cases for ReactSearchOp class. -It can be run directly with: python test_react_search_op.py -""" - -import asyncio - -from flowllm.gallery.agent import ReactSearchOp -from flowllm.main import FlowLLMApp - - -async def async_main(): - """Test function for ReactSearchOp.""" - async with FlowLLMApp(): - op = ReactSearchOp(add_think_tool=True, language="zh") - await op.async_call(query="小米股价为什么一直跌？现在还值得买吗？") - print(op.output) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_read_file_op.py b/tests_op/test_read_file_op.py deleted file mode 100644 index db6ea25..0000000 --- a/tests_op/test_read_file_op.py +++ /dev/null @@ -1,91 +0,0 @@ -"""Test script for ReadFileOp. - -This script provides test cases for ReadFileOp class. -It can be run directly with: python test_read_file_op.py -""" - -import asyncio -import json -from pathlib import Path - -from flowllm.extensions.file_tool import ReadFileOp -from flowllm.main import FlowLLMApp - - -async def test_simple_input_schema(): - """Test case 1: Test tool_call.simple_input_dump().""" - print("=" * 80) - print("Test Case 1: Testing tool_call.simple_input_dump()") - print("=" * 80) - - async with FlowLLMApp(): - op = ReadFileOp() - tool_call = op.build_tool_call() - tool_call.name = op.short_name - - # Test simple_input_dump - input_schema = tool_call.simple_input_dump() - print("\nInput Schema (simple_input_dump):") - print(json.dumps(input_schema, indent=2, ensure_ascii=False)) - - # Verify structure - assert "type" in input_schema - assert "function" in input_schema - assert "name" in input_schema["function"] - assert "description" in input_schema["function"] - assert "parameters" in input_schema["function"] - assert "properties" in input_schema["function"]["parameters"] - assert "required" in input_schema["function"]["parameters"] - - print("\n✓ Test Case 1 passed: simple_input_dump() works correctly") - print("=" * 80) - - -async def test_normal_async_execute(): - """Test case 2: Test normal async_execute.""" - print("\n" + "=" * 80) - print("Test Case 2: Testing normal async_execute") - print("=" * 80) - - async with FlowLLMApp(): - op = ReadFileOp() - # Test with a file that should exist (read this test file itself) - test_file_path = Path(__file__).resolve() - await op.async_call(absolute_path=str(test_file_path)) - - print("\nOutput:") - print(op.output[:200] + "..." if len(op.output) > 200 else op.output) - print("\n✓ Test Case 2 passed: async_execute() works correctly") - print("=" * 80) - - -async def test_exception_async_default_execute(): - """Test case 3: Test exception handling, executing async_default_execute.""" - print("\n" + "=" * 80) - print("Test Case 3: Testing exception handling (async_default_execute)") - print("=" * 80) - - async with FlowLLMApp(): - op = ReadFileOp() - # Test with a non-existent file to trigger exception - await op.async_call(absolute_path="/nonexistent/file/that/does/not/exist.txt") - - print("\nOutput after exception:") - print(op.output) - print("\n✓ Test Case 3 passed: async_default_execute() was called on exception") - print("=" * 80) - - -async def async_main(): - """Run all test cases.""" - await test_simple_input_schema() - await test_normal_async_execute() - await test_exception_async_default_execute() - - print("\n" + "=" * 80) - print("All tests passed!") - print("=" * 80) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_rip_grep_op.py b/tests_op/test_rip_grep_op.py deleted file mode 100644 index 4c791e1..0000000 --- a/tests_op/test_rip_grep_op.py +++ /dev/null @@ -1,104 +0,0 @@ -"""Test script for RipGrepOp. - -This script provides test cases for RipGrepOp class. -It can be run directly with: python test_rip_grep_op.py -""" - -import asyncio -import json -import shutil - -from flowllm.extensions.file_tool import RipGrepOp -from flowllm.main import FlowLLMApp - - -def is_ripgrep_available(): - """Check if ripgrep (rg) is available in the system.""" - return shutil.which("rg") is not None - - -async def test_simple_input_schema(): - """Test case 1: Test tool_call.simple_input_dump().""" - print("=" * 80) - print("Test Case 1: Testing tool_call.simple_input_dump()") - print("=" * 80) - - async with FlowLLMApp(): - op = RipGrepOp() - tool_call = op.build_tool_call() - tool_call.name = op.short_name - - # Test simple_input_dump - input_schema = tool_call.simple_input_dump() - print("\nInput Schema (simple_input_dump):") - print(json.dumps(input_schema, indent=2, ensure_ascii=False)) - - # Verify structure - assert "type" in input_schema - assert "function" in input_schema - assert "name" in input_schema["function"] - assert "description" in input_schema["function"] - assert "parameters" in input_schema["function"] - assert "properties" in input_schema["function"]["parameters"] - assert "required" in input_schema["function"]["parameters"] - - print("\n✓ Test Case 1 passed: simple_input_dump() works correctly") - print("=" * 80) - - -async def test_normal_async_execute(): - """Test case 2: Test normal async_execute.""" - print("\n" + "=" * 80) - print("Test Case 2: Testing normal async_execute") - print("=" * 80) - - if not is_ripgrep_available(): - print("\n⚠ Warning: ripgrep (rg) is not available. Skipping this test.") - print("Please install ripgrep to run this test case.") - print("=" * 80) - return - - async with FlowLLMApp(): - op = RipGrepOp() - # Test with a pattern that should find matches - await op.async_call(pattern="def test", path=".") - - print("\nOutput:") - print(op.output) - print("\n✓ Test Case 2 passed: async_execute() works correctly") - print("=" * 80) - - -async def test_exception_async_default_execute(): - """Test case 3: Test exception handling, executing async_default_execute.""" - print("\n" + "=" * 80) - print("Test Case 3: Testing exception handling (async_default_execute)") - print("=" * 80) - - async with FlowLLMApp(): - op = RipGrepOp() - # Test with a non-existent path to trigger exception - # This will raise ValueError in async_execute, which will trigger async_default_execute - await op.async_call(pattern="test", path="/nonexistent/path/xxxx") - - print("\nOutput after exception:") - print(op.output) - # Verify that async_default_execute was called (output should contain error message) - assert "Failed to search" in op.output or "does not exist" in op.output - print("\n✓ Test Case 3 passed: async_default_execute() was called on exception") - print("=" * 80) - - -async def async_main(): - """Run all test cases.""" - await test_simple_input_schema() - await test_normal_async_execute() - await test_exception_async_default_execute() - - print("\n" + "=" * 80) - print("All tests passed!") - print("=" * 80) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_search_op.py b/tests_op/test_search_op.py deleted file mode 100644 index 6001d24..0000000 --- a/tests_op/test_search_op.py +++ /dev/null @@ -1,41 +0,0 @@ -"""Test module for search operations. - -This module provides test functions for various search operations including -Tongyi MCP search, BochaAI MCP search, and Tavily search. -""" - -import asyncio - -from flowllm.main import FlowLLMApp - - -async def main(): - """Test function for various search operations. - - This function tests different search operations by executing sample queries - and printing the results for Tongyi MCP search, BochaAI MCP search, and Tavily search. - """ - async with FlowLLMApp("config=default,search"): - query = "what is ai?" - - from flowllm.gallery.search import TongyiMcpSearchOp - - op = TongyiMcpSearchOp() - await op.async_call(query=query) - print("tongyi:", op.output) - - from flowllm.gallery.search import BochaMcpSearchOp - - op = BochaMcpSearchOp() - await op.async_call(query=query) - print("bocha:", op.output) - - from flowllm.gallery.search import TavilySearchOp - - op = TavilySearchOp() - await op.async_call(query=query) - print("tavily:", op.output) - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/tests_op/test_shell_op.py b/tests_op/test_shell_op.py deleted file mode 100644 index 010c4ac..0000000 --- a/tests_op/test_shell_op.py +++ /dev/null @@ -1,121 +0,0 @@ -"""Test script for ShellOp. - -This script provides test cases for ShellOp class. -It can be run directly with: python test_shell_op.py -""" - -import asyncio -import json - -from flowllm.extensions.file_tool import ShellOp -from flowllm.main import FlowLLMApp - - -async def test_simple_input_schema(): - """Test case 1: Test tool_call.simple_input_dump().""" - print("=" * 80) - print("Test Case 1: Testing tool_call.simple_input_dump()") - print("=" * 80) - - async with FlowLLMApp(): - op = ShellOp() - tool_call = op.build_tool_call() - tool_call.name = op.short_name - - # Test simple_input_dump - input_schema = tool_call.simple_input_dump() - print("\nInput Schema (simple_input_dump):") - print(json.dumps(input_schema, indent=2, ensure_ascii=False)) - - # Verify structure - assert "type" in input_schema - assert "function" in input_schema - assert "name" in input_schema["function"] - assert "description" in input_schema["function"] - assert "parameters" in input_schema["function"] - assert "properties" in input_schema["function"]["parameters"] - assert "required" in input_schema["function"]["parameters"] - - print("\n✓ Test Case 1 passed: simple_input_dump() works correctly") - print("=" * 80) - - -async def test_normal_async_execute(): - """Test case 2: Test normal async_execute (foreground).""" - print("\n" + "=" * 80) - print("Test Case 2: Testing normal async_execute (foreground)") - print("=" * 80) - - async with FlowLLMApp(): - op = ShellOp() - # Test with a simple command that should succeed - await op.async_call(command="echo 'Hello, World!'", is_background=False) - - print("\nOutput:") - print(op.output) - # Verify output contains expected content - assert "Hello, World!" in op.output or "Command:" in op.output - print("\n✓ Test Case 2 passed: async_execute() works correctly") - print("=" * 80) - - -async def test_background_execute(): - """Test case 3: Test background execution.""" - print("\n" + "=" * 80) - print("Test Case 3: Testing background execution") - print("=" * 80) - - async with FlowLLMApp(): - op = ShellOp() - # Test with a background command - await op.async_call( - command="echo 'Background test'", - is_background=True, - description="Test background execution", - ) - - print("\nOutput:") - print(op.output) - # Verify output contains background execution info - assert "background" in op.output.lower() or "PID:" in op.output - print("\n✓ Test Case 3 passed: background execution works correctly") - print("=" * 80) - - -async def test_exception_async_default_execute(): - """Test case 4: Test exception handling, executing async_default_execute.""" - print("\n" + "=" * 80) - print("Test Case 4: Testing exception handling (async_default_execute)") - print("=" * 80) - - async with FlowLLMApp(): - op = ShellOp() - # Test with a non-existent directory to trigger exception - await op.async_call( - command="echo test", - is_background=False, - directory="/nonexistent/path/xxxx", - ) - - print("\nOutput after exception:") - print(op.output) - # Verify that async_default_execute was called (output should contain error message) - assert "Failed to execute" in op.output or "does not exist" in op.output - print("\n✓ Test Case 4 passed: async_default_execute() was called on exception") - print("=" * 80) - - -async def async_main(): - """Run all test cases.""" - await test_simple_input_schema() - await test_normal_async_execute() - await test_background_execute() - await test_exception_async_default_execute() - - print("\n" + "=" * 80) - print("All tests passed!") - print("=" * 80) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_skill_agent_op.py b/tests_op/test_skill_agent_op.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests_op/test_smart_edit_op.py b/tests_op/test_smart_edit_op.py deleted file mode 100644 index 30535dd..0000000 --- a/tests_op/test_smart_edit_op.py +++ /dev/null @@ -1,177 +0,0 @@ -"""Test script for SmartEditOp. - -This script provides test cases for SmartEditOp class. -It can be run directly with: python test_smart_edit_op.py -""" - -import asyncio -import json -import os -import tempfile - -from flowllm.extensions.file_tool import SmartEditOp -from flowllm.main import FlowLLMApp - - -async def test_simple_input_schema(): - """Test case 1: Test tool_call.simple_input_dump().""" - print("=" * 80) - print("Test Case 1: Testing tool_call.simple_input_dump()") - print("=" * 80) - - async with FlowLLMApp(): - op = SmartEditOp() - tool_call = op.build_tool_call() - tool_call.name = op.short_name - - # Test simple_input_dump - input_schema = tool_call.simple_input_dump() - print("\nInput Schema (simple_input_dump):") - print(json.dumps(input_schema, indent=2, ensure_ascii=False)) - - # Verify structure - assert "type" in input_schema - assert "function" in input_schema - assert "name" in input_schema["function"] - assert "description" in input_schema["function"] - assert "parameters" in input_schema["function"] - assert "properties" in input_schema["function"]["parameters"] - assert "required" in input_schema["function"]["parameters"] - - print("\n✓ Test Case 1 passed: simple_input_dump() works correctly") - print("=" * 80) - - -async def test_normal_async_execute(): - """Test case 2: Test normal async_execute.""" - print("\n" + "=" * 80) - print("Test Case 2: Testing normal async_execute") - print("=" * 80) - - async with FlowLLMApp(): - op = SmartEditOp() - - # Create a temporary file for testing - with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as tmp_file: - tmp_file_path = tmp_file.name - tmp_file.write( - """def hello(): - print("Hello, World!") - return True - -def goodbye(): - print("Goodbye!") -""", - ) - - try: - # Test editing an existing file - await op.async_call( - file_path=tmp_file_path, - old_string='def hello():\n print("Hello, World!")\n return True', - new_string='def hello():\n print("Hello, Updated World!")\n return True', - instruction="Update the hello function to print a different message", - ) - - print("\nOutput:") - print(op.output) - - # Verify the file was modified - with open(tmp_file_path, "r", encoding="utf-8") as f: - content = f.read() - assert "Hello, Updated World!" in content - assert "Hello, World!" not in content - - print("\n✓ Test Case 2 passed: async_execute() works correctly") - finally: - # Clean up - if os.path.exists(tmp_file_path): - os.unlink(tmp_file_path) - - print("=" * 80) - - -async def test_create_new_file(): - """Test case 2b: Test creating a new file.""" - print("\n" + "=" * 80) - print("Test Case 2b: Testing creating a new file") - print("=" * 80) - - async with FlowLLMApp(): - op = SmartEditOp() - - # Create a temporary file path that doesn't exist - with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as tmp_file: - tmp_file_path = tmp_file.name - - # Remove the file so it doesn't exist - if os.path.exists(tmp_file_path): - os.unlink(tmp_file_path) - - try: - # Test creating a new file - await op.async_call( - file_path=tmp_file_path, - old_string="", - new_string="""def new_function(): - print("This is a new file!") - return True -""", - instruction="Create a new file with a simple function", - ) - - print("\nOutput:") - print(op.output) - - # Verify the file was created - assert os.path.exists(tmp_file_path) - with open(tmp_file_path, "r", encoding="utf-8") as f: - content = f.read() - assert "new_function" in content - assert "This is a new file!" in content - - print("\n✓ Test Case 2b passed: Creating new file works correctly") - finally: - # Clean up - if os.path.exists(tmp_file_path): - os.unlink(tmp_file_path) - - print("=" * 80) - - -async def test_exception_async_default_execute(): - """Test case 3: Test exception handling, executing async_default_execute.""" - print("\n" + "=" * 80) - print("Test Case 3: Testing exception handling (async_default_execute)") - print("=" * 80) - - async with FlowLLMApp(): - op = SmartEditOp() - # Test with a non-existent file and non-empty old_string to trigger exception - await op.async_call( - file_path="/nonexistent/file/that/does/not/exist.txt", - old_string="some text", - new_string="new text", - instruction="This should fail because file doesn't exist", - ) - - print("\nOutput after exception:") - print(op.output) - print("\n✓ Test Case 3 passed: async_default_execute() was called on exception") - print("=" * 80) - - -async def async_main(): - """Run all test cases.""" - await test_simple_input_schema() - await test_normal_async_execute() - await test_create_new_file() - await test_exception_async_default_execute() - - print("\n" + "=" * 80) - print("All tests passed!") - print("=" * 80) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_stream_chat_op.py b/tests_op/test_stream_chat_op.py deleted file mode 100644 index 7c58ca1..0000000 --- a/tests_op/test_stream_chat_op.py +++ /dev/null @@ -1,87 +0,0 @@ -"""Test script for StreamChatOp. - -This script provides test cases for StreamChatOp class. -It can be run directly with: python test_stream_chat_op.py -""" - -import asyncio - -from flowllm.core.enumeration import Role, ChunkEnum -from flowllm.core.schema import Message -from flowllm.gallery.stream_chat_op import StreamChatOp -from flowllm.main import FlowLLMApp - - -async def async_task(op: StreamChatOp, messages, system_prompt: str, stream_queue: asyncio.Queue): - """Wrapper function to call async_call and add done chunk.""" - await op.async_call( - messages=messages, - system_prompt=system_prompt, - stream_queue=stream_queue, - ) - await op.context.add_stream_done() - - -async def async_main(): - """Test function for StreamChatOp.""" - async with FlowLLMApp(): - # Test 1: Test basic streaming with ANSWER chunks - print("=" * 80) - print("Test 1: Testing basic streaming with ANSWER chunks") - print("=" * 80) - - op = StreamChatOp() - messages = [ - Message(role=Role.USER, content="Hello, how are you?"), - ] - - stream_queue = asyncio.Queue() - task = asyncio.create_task( - async_task( - op=op, - messages=messages, - system_prompt="You are a helpful assistant.", - stream_queue=stream_queue, - ), - ) - - # Collect stream chunks - stream_chunks = [] - while True: - stream_chunk = await stream_queue.get() - if stream_chunk.done: - print("\nend") - break - stream_chunks.append(stream_chunk) - print(stream_chunk.chunk, end="", flush=True) - - await task - - assert len(stream_chunks) > 0, "Should receive at least one stream chunk" - assert any(chunk.chunk_type == ChunkEnum.ANSWER for chunk in stream_chunks), "Should receive ANSWER chunks" - print("\n✓ Test 1 passed\n") - - # Test 2: Test assertion error with invalid messages - print("=" * 80) - print("Test 2: Testing assertion error with invalid messages") - print("=" * 80) - - op2 = StreamChatOp() - try: - await op2.async_call( - messages=["not a Message object"], - system_prompt="Test prompt", - stream_queue=asyncio.Queue(), - ) - print("✗ Test 2 failed: Should have raised an assertion error") - assert False, "Should have raised AssertionError" - except AssertionError as e: - print(f"✓ Test 2 passed: Correctly raised AssertionError: {e}\n") - - print("=" * 80) - print("All tests passed!") - print("=" * 80) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_token_count_op.py b/tests_op/test_token_count_op.py deleted file mode 100644 index 45440ac..0000000 --- a/tests_op/test_token_count_op.py +++ /dev/null @@ -1,42 +0,0 @@ -"""Test script for TokenCountOp. - -This script provides a simple test case for TokenCountOp class. -It can be run directly with: python test_token_count_op.py -""" - -import asyncio - -from flowllm.gallery import TokenCountOp -from flowllm.main import FlowLLMApp - - -async def async_main(): - """Test function for TokenCountOp.""" - async with FlowLLMApp(): - print("=" * 80) - print("Test: Running TokenCountOp with sample messages") - print("=" * 80) - - op = TokenCountOp() - - # Prepare messages - messages = [ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "Hello, how are you?"}, - {"role": "assistant", "content": "I'm doing well, thank you for asking!"}, - ] - - print("Messages to count:") - for msg in messages: - print(f" - {msg['role']}: {msg['content']}") - print() - - # Execute the operation - await op.async_call(messages=messages) - - print("\n✓ Test completed successfully") - print("=" * 80) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_write_file_op.py b/tests_op/test_write_file_op.py deleted file mode 100644 index 00a8115..0000000 --- a/tests_op/test_write_file_op.py +++ /dev/null @@ -1,209 +0,0 @@ -"""Test script for WriteFileOp. - -This script provides test cases for WriteFileOp class. -It can be run directly with: python test_write_file_op.py -""" - -import asyncio -import json -import os -import tempfile - -from flowllm.extensions.file_tool import WriteFileOp -from flowllm.main import FlowLLMApp - - -async def test_simple_input_schema(): - """Test case 1: Test tool_call.simple_input_dump().""" - print("=" * 80) - print("Test Case 1: Testing tool_call.simple_input_dump()") - print("=" * 80) - - async with FlowLLMApp(): - op = WriteFileOp() - tool_call = op.build_tool_call() - tool_call.name = op.short_name - - # Test simple_input_dump - input_schema = tool_call.simple_input_dump() - print("\nInput Schema (simple_input_dump):") - print(json.dumps(input_schema, indent=2, ensure_ascii=False)) - - # Verify structure - assert "type" in input_schema - assert "function" in input_schema - assert "name" in input_schema["function"] - assert "description" in input_schema["function"] - assert "parameters" in input_schema["function"] - assert "properties" in input_schema["function"]["parameters"] - assert "required" in input_schema["function"]["parameters"] - - print("\n✓ Test Case 1 passed: simple_input_dump() works correctly") - print("=" * 80) - - -async def test_create_new_file(): - """Test case 2: Test creating a new file.""" - print("\n" + "=" * 80) - print("Test Case 2: Testing creating a new file") - print("=" * 80) - - async with FlowLLMApp(): - op = WriteFileOp() - - # Create a temporary file path that doesn't exist - with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as tmp_file: - tmp_file_path = tmp_file.name - - # Remove the file so it doesn't exist - if os.path.exists(tmp_file_path): - os.unlink(tmp_file_path) - - try: - # Test creating a new file - test_content = """def hello(): - print("Hello, World!") - return True -""" - await op.async_call( - file_path=tmp_file_path, - content=test_content, - ) - - print("\nOutput:") - print(op.output) - - # Verify the file was created - assert os.path.exists(tmp_file_path) - with open(tmp_file_path, "r", encoding="utf-8") as f: - content = f.read() - assert content == test_content - assert "def hello()" in content - - print("\n✓ Test Case 2 passed: Creating new file works correctly") - finally: - # Clean up - if os.path.exists(tmp_file_path): - os.unlink(tmp_file_path) - - print("=" * 80) - - -async def test_overwrite_existing_file(): - """Test case 2b: Test overwriting an existing file.""" - print("\n" + "=" * 80) - print("Test Case 2b: Testing overwriting an existing file") - print("=" * 80) - - async with FlowLLMApp(): - op = WriteFileOp() - - # Create a temporary file with initial content - with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as tmp_file: - tmp_file_path = tmp_file.name - tmp_file.write("Initial content\n") - - try: - # Test overwriting the file - new_content = """def updated(): - print("This is updated content!") - return True -""" - await op.async_call( - file_path=tmp_file_path, - content=new_content, - ) - - print("\nOutput:") - print(op.output) - - # Verify the file was overwritten - with open(tmp_file_path, "r", encoding="utf-8") as f: - content = f.read() - assert content == new_content - assert "Initial content" not in content - assert "def updated()" in content - - print("\n✓ Test Case 2b passed: Overwriting existing file works correctly") - finally: - # Clean up - if os.path.exists(tmp_file_path): - os.unlink(tmp_file_path) - - print("=" * 80) - - -async def test_create_file_with_parent_dirs(): - """Test case 2c: Test creating a file with parent directories.""" - print("\n" + "=" * 80) - print("Test Case 2c: Testing creating file with parent directories") - print("=" * 80) - - async with FlowLLMApp(): - op = WriteFileOp() - - # Create a temporary directory and a nested path - with tempfile.TemporaryDirectory() as tmp_dir: - nested_path = os.path.join(tmp_dir, "nested", "subdir", "test_file.txt") - - try: - # Test creating a file with nested directories - test_content = "Content in nested directory\n" - await op.async_call( - file_path=nested_path, - content=test_content, - ) - - print("\nOutput:") - print(op.output) - - # Verify the file was created with parent directories - assert os.path.exists(nested_path) - with open(nested_path, "r", encoding="utf-8") as f: - content = f.read() - assert content == test_content - - print("\n✓ Test Case 2c passed: Creating file with parent directories works correctly") - finally: - # Clean up is handled by TemporaryDirectory context manager - pass - - print("=" * 80) - - -async def test_exception_async_default_execute(): - """Test case 3: Test exception handling, executing async_default_execute.""" - print("\n" + "=" * 80) - print("Test Case 3: Testing exception handling (async_default_execute)") - print("=" * 80) - - async with FlowLLMApp(): - op = WriteFileOp() - - # Test with an empty file_path to trigger exception - await op.async_call( - file_path="", - content="some content", - ) - - print("\nOutput after exception:") - print(op.output) - print("\n✓ Test Case 3 passed: async_default_execute() was called on exception") - print("=" * 80) - - -async def async_main(): - """Run all test cases.""" - await test_simple_input_schema() - await test_create_new_file() - await test_overwrite_existing_file() - await test_create_file_with_parent_dirs() - await test_exception_async_default_execute() - - print("\n" + "=" * 80) - print("All tests passed!") - print("=" * 80) - - -if __name__ == "__main__": - asyncio.run(async_main()) diff --git a/tests_op/test_write_todos_op.py b/tests_op/test_write_todos_op.py deleted file mode 100644 index 4e13542..0000000 --- a/tests_op/test_write_todos_op.py +++ /dev/null @@ -1,121 +0,0 @@ -"""Test script for WriteTodosOp. - -This script provides test cases for WriteTodosOp class. -It can be run directly with: python test_write_todos_op.py -""" - -import asyncio -import json - -from flowllm.extensions.file_tool import WriteTodosOp -from flowllm.main import FlowLLMApp - - -async def test_simple_input_schema(): - """Test case 1: Test tool_call.simple_input_dump().""" - print("=" * 80) - print("Test Case 1: Testing tool_call.simple_input_dump()") - print("=" * 80) - - async with FlowLLMApp(): - op = WriteTodosOp() - tool_call = op.build_tool_call() - tool_call.name = op.short_name - - # Test simple_input_dump - input_schema = tool_call.simple_input_dump() - print("\nInput Schema (simple_input_dump):") - print(json.dumps(input_schema, indent=2, ensure_ascii=False)) - - # Verify structure - assert "type" in input_schema - assert "function" in input_schema - assert "name" in input_schema["function"] - assert "description" in input_schema["function"] - assert "parameters" in input_schema["function"] - assert "properties" in input_schema["function"]["parameters"] - assert "required" in input_schema["function"]["parameters"] - - # Verify todos parameter structure - todos_param = input_schema["function"]["parameters"]["properties"]["todos"] - assert todos_param["type"] == "array" - assert "items" in todos_param - assert todos_param["items"]["type"] == "object" - assert "properties" in todos_param["items"] - assert "description" in todos_param["items"]["properties"] - assert "status" in todos_param["items"]["properties"] - assert "enum" in todos_param["items"]["properties"]["status"] - - print("\n✓ Test Case 1 passed: simple_input_dump() works correctly") - print("=" * 80) - - -async def test_normal_async_execute(): - """Test case 2: Test normal async_execute.""" - print("\n" + "=" * 80) - print("Test Case 2: Testing normal async_execute") - print("=" * 80) - - async with FlowLLMApp(): - op = WriteTodosOp() - - # Test 1: Empty todo list - print("\nTest 2.1: Empty todo list") - await op.async_call(todos=[]) - print("\nOutput:") - print(op.output) - assert "cleared the todo list" in op.output.lower() - - # Test 2: Todo list with multiple items - print("\nTest 2.2: Todo list with multiple items") - todos = [ - {"description": "Task 1", "status": "pending"}, - {"description": "Task 2", "status": "in_progress"}, - {"description": "Task 3", "status": "completed"}, - {"description": "Task 4", "status": "cancelled"}, - ] - await op.async_call(todos=todos) - print("\nOutput:") - print(op.output) - assert "Successfully updated the todo list" in op.output - assert "[pending] Task 1" in op.output - assert "[in_progress] Task 2" in op.output - assert "[completed] Task 3" in op.output - assert "[cancelled] Task 4" in op.output - - print("\n✓ Test Case 2 passed: async_execute() works correctly") - print("=" * 80) - - -async def test_exception_async_default_execute(): - """Test case 3: Test exception handling, executing async_default_execute.""" - print("\n" + "=" * 80) - print("Test Case 3: Testing exception handling (async_default_execute)") - print("=" * 80) - - async with FlowLLMApp(): - op = WriteTodosOp() - - # Test: todos is not a list - await op.async_call(todos="not a list") - print("\nOutput after exception:") - print(op.output) - assert "Failed to update the todo list" in op.output - - print("\n✓ Test Case 3 passed: async_default_execute() was called on exception") - print("=" * 80) - - -async def async_main(): - """Run all test cases.""" - await test_simple_input_schema() - await test_normal_async_execute() - await test_exception_async_default_execute() - - print("\n" + "=" * 80) - print("All tests passed!") - print("=" * 80) - - -if __name__ == "__main__": - asyncio.run(async_main())