Compare commits

..

5 Commits

Author SHA1 Message Date
laundry ba6198f3ec test: fix unit test error
Change-Id: I3dd7a6179132e5497a30ada443d88de0c47af3d4
2025-05-20 11:59:00 +08:00
laundry 28a01dfe0e test: fix unit test error
Change-Id: If4c4cd10673e76a30945674c7cda198aeabf28d0
2025-05-20 11:58:39 +08:00
laundry 3b1db26507 test: fix test error
Change-Id: I3997dc53a2cfaa35501a1fbda5902ee15528124e
2025-05-19 15:47:08 +08:00
laundry e927b556d6 test: add background node unit test
Change-Id: I9aabcf02ff04fda40c56f3ea22abe6b8f93bf9b6
2025-05-19 15:36:24 +08:00
laundry c2b8dd8e6a test: add background node unit test
Change-Id: Ia99f5a1687464387dcb01bbee04deaa371c6e490
2025-05-19 15:33:42 +08:00
169 changed files with 1885 additions and 22520 deletions
-17
View File
@@ -7,29 +7,12 @@ NEXT_PUBLIC_API_URL="http://localhost:8000/api"
AGENT_RECURSION_LIMIT=30 AGENT_RECURSION_LIMIT=30
# CORS settings
# Comma-separated list of allowed origins for CORS requests
# Example: ALLOWED_ORIGINS=http://localhost:3000,http://example.com
ALLOWED_ORIGINS=http://localhost:3000
# Search Engine, Supported values: tavily (recommended), duckduckgo, brave_search, arxiv # Search Engine, Supported values: tavily (recommended), duckduckgo, brave_search, arxiv
SEARCH_API=tavily SEARCH_API=tavily
TAVILY_API_KEY=tvly-xxx TAVILY_API_KEY=tvly-xxx
# BRAVE_SEARCH_API_KEY=xxx # Required only if SEARCH_API is brave_search # BRAVE_SEARCH_API_KEY=xxx # Required only if SEARCH_API is brave_search
# JINA_API_KEY=jina_xxx # Optional, default is None # JINA_API_KEY=jina_xxx # Optional, default is None
# Optional, RAG provider
# RAG_PROVIDER=vikingdb_knowledge_base
# VIKINGDB_KNOWLEDGE_BASE_API_URL="api-knowledgebase.mlp.cn-beijing.volces.com"
# VIKINGDB_KNOWLEDGE_BASE_API_AK="AKxxx"
# VIKINGDB_KNOWLEDGE_BASE_API_SK=""
# VIKINGDB_KNOWLEDGE_BASE_RETRIEVAL_SIZE=15
# RAG_PROVIDER=ragflow
# RAGFLOW_API_URL="http://localhost:9388"
# RAGFLOW_API_KEY="ragflow-xxx"
# RAGFLOW_RETRIEVAL_SIZE=10
# Optional, volcengine TTS for generating podcast # Optional, volcengine TTS for generating podcast
VOLCENGINE_TTS_APPID=xxx VOLCENGINE_TTS_APPID=xxx
VOLCENGINE_TTS_ACCESS_TOKEN=xxx VOLCENGINE_TTS_ACCESS_TOKEN=xxx
-95
View File
@@ -1,95 +0,0 @@
name: Publish Containers
on:
push:
branches:
- main
release:
types: [published]
workflow_dispatch:
jobs:
backend-container:
runs-on: ubuntu-latest
permissions:
contents: read
packages: write
attestations: write
id-token: write
env:
REGISTRY: ghcr.io
IMAGE_NAME: ${{ github.repository }}
steps:
- name: Checkout repository
uses: actions/checkout@v4
- name: Log in to the Container registry
uses: docker/login-action@v3
with:
registry: ${{ env.REGISTRY }}
username: ${{ github.actor }}
password: ${{ secrets.GITHUB_TOKEN }}
- name: Extract metadata (tags, labels) for Docker
id: meta
uses: docker/metadata-action@v5
with:
images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
- name: Build and push Docker image
id: push
uses: docker/build-push-action@v6
with:
context: .
file: Dockerfile
push: true
tags: ${{ steps.meta.outputs.tags }}
labels: ${{ steps.meta.outputs.labels }}
- name: Generate artifact attestation
uses: actions/attest-build-provenance@v2
with:
subject-name: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME}}
subject-digest: ${{ steps.push.outputs.digest }}
push-to-registry: true
frontend-container:
runs-on: ubuntu-latest
permissions:
contents: read
packages: write
attestations: write
id-token: write
env:
REGISTRY: ghcr.io
IMAGE_NAME: ${{ github.repository }}-web
steps:
- name: Checkout repository
uses: actions/checkout@v4
- name: Log in to the Container registry
uses: docker/login-action@74a5d142397b4f367a81961eba4e8cd7edddf772 #v3.4.0
with:
registry: ${{ env.REGISTRY }}
username: ${{ github.actor }}
password: ${{ secrets.GITHUB_TOKEN }}
- name: Extract metadata (tags, labels) for Docker
id: meta
uses: docker/metadata-action@902fa8ec7d6ecbf8d84d538b9b233a880e428804 #v5.7.0
with:
images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
- name: Build and push Docker image
id: push
uses: docker/build-push-action@263435318d21b8e681c14492fe198d362a7d2c83 #v6.18.0
with:
context: web
file: web/Dockerfile
push: true
tags: ${{ steps.meta.outputs.tags }}
labels: ${{ steps.meta.outputs.labels }}
- name: Generate artifact attestation
uses: actions/attest-build-provenance@v2
with:
subject-name: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME}}
subject-digest: ${{ steps.push.outputs.digest }}
push-to-registry: true
+1 -4
View File
@@ -6,9 +6,6 @@ on:
pull_request: pull_request:
branches: [ '*' ] branches: [ '*' ]
permissions:
contents: read
jobs: jobs:
lint: lint:
runs-on: ubuntu-latest runs-on: ubuntu-latest
@@ -16,7 +13,7 @@ jobs:
- uses: actions/checkout@v3 - uses: actions/checkout@v3
- name: Install the latest version of uv - name: Install the latest version of uv
uses: astral-sh/setup-uv@v6.3.1 uses: astral-sh/setup-uv@v5
with: with:
version: "latest" version: "latest"
+1 -4
View File
@@ -6,9 +6,6 @@ on:
pull_request: pull_request:
branches: [ '*' ] branches: [ '*' ]
permissions:
contents: read
jobs: jobs:
test: test:
runs-on: ubuntu-latest runs-on: ubuntu-latest
@@ -16,7 +13,7 @@ jobs:
- uses: actions/checkout@v3 - uses: actions/checkout@v3
- name: Install the latest version of uv - name: Install the latest version of uv
uses: astral-sh/setup-uv@v6.3.1 uses: astral-sh/setup-uv@v5
with: with:
version: "latest" version: "latest"
-2
View File
@@ -6,13 +6,11 @@ dist/
wheels/ wheels/
*.egg-info *.egg-info
.coverage .coverage
.coverage.*
agent_history.gif agent_history.gif
static/browser_history/*.gif static/browser_history/*.gif
# Virtual environments # Virtual environments
.venv .venv
venv/
# Environment variables # Environment variables
.env .env
-30
View File
@@ -1,36 +1,6 @@
{ {
"version": "0.2.0", "version": "0.2.0",
"configurations": [ "configurations": [
{
"name": "Debug Tests",
"type": "debugpy",
"request": "launch",
"module": "pytest",
"args": [
"${workspaceFolder}/tests",
"-v",
"-s"
],
"console": "integratedTerminal",
"justMyCode": false,
"env": {
"PYTHONPATH": "${workspaceFolder}"
}
},
{
"name": "Debug Current Test File",
"type": "debugpy",
"request": "launch",
"module": "pytest",
"args": [
"${file}",
"-v",
"-s"
],
"console": "integratedTerminal",
"justMyCode": false
},
{ {
"name": "Python: 当前文件", "name": "Python: 当前文件",
"type": "debugpy", "type": "debugpy",
-7
View File
@@ -1,7 +0,0 @@
{
"python.testing.pytestArgs": [
"tests"
],
"python.testing.unittestEnabled": false,
"python.testing.pytestEnabled": true
}
-11
View File
@@ -17,14 +17,11 @@ There are many ways you can contribute to DeerFlow:
1. Fork the repository 1. Fork the repository
2. Clone your fork: 2. Clone your fork:
```bash ```bash
git clone https://github.com/bytedance/deer-flow.git git clone https://github.com/bytedance/deer-flow.git
cd deer-flow cd deer-flow
``` ```
3. Set up your development environment: 3. Set up your development environment:
```bash ```bash
# Install dependencies, uv will take care of the python interpreter and venv creation # Install dependencies, uv will take care of the python interpreter and venv creation
uv sync uv sync
@@ -33,9 +30,7 @@ There are many ways you can contribute to DeerFlow:
uv pip install -e ".[dev]" uv pip install -e ".[dev]"
uv pip install -e ".[test]" uv pip install -e ".[test]"
``` ```
4. Configure pre-commit hooks: 4. Configure pre-commit hooks:
```bash ```bash
chmod +x pre-commit chmod +x pre-commit
ln -s ../../pre-commit .git/hooks/pre-commit ln -s ../../pre-commit .git/hooks/pre-commit
@@ -44,7 +39,6 @@ There are many ways you can contribute to DeerFlow:
## Development Process ## Development Process
1. Create a new branch: 1. Create a new branch:
```bash ```bash
git checkout -b feature/amazing-feature git checkout -b feature/amazing-feature
``` ```
@@ -56,7 +50,6 @@ There are many ways you can contribute to DeerFlow:
- Update documentation as needed - Update documentation as needed
3. Run tests and checks: 3. Run tests and checks:
```bash ```bash
make test # Run tests make test # Run tests
make lint # Run linting make lint # Run linting
@@ -65,13 +58,11 @@ There are many ways you can contribute to DeerFlow:
``` ```
4. Commit your changes: 4. Commit your changes:
```bash ```bash
git commit -m 'Add some amazing feature' git commit -m 'Add some amazing feature'
``` ```
5. Push to your fork: 5. Push to your fork:
```bash ```bash
git push origin feature/amazing-feature git push origin feature/amazing-feature
``` ```
@@ -99,7 +90,6 @@ There are many ways you can contribute to DeerFlow:
## Testing ## Testing
Run the test suite: Run the test suite:
```bash ```bash
# Run all tests # Run all tests
make test make test
@@ -132,7 +122,6 @@ make format
## Need Help? ## Need Help?
If you need help with anything: If you need help with anything:
- Check existing issues and discussions - Check existing issues and discussions
- Join our community channels - Join our community channels
- Ask questions in discussions - Ask questions in discussions
+1 -1
View File
@@ -1,4 +1,4 @@
FROM ghcr.io/astral-sh/uv:python3.12-bookworm FROM ghcr.io/astral-sh/uv:python3.12-bookworm-slim
# Install uv. # Install uv.
COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv
-1
View File
@@ -8,7 +8,6 @@ format:
lint: lint:
uv run black --check . uv run black --check .
uv run ruff check .
serve: serve:
uv run server.py --reload uv run server.py --reload
+6 -28
View File
@@ -12,16 +12,13 @@
**DeerFlow** (**D**eep **E**xploration and **E**fficient **R**esearch **Flow**) is a community-driven Deep Research framework that builds upon the incredible work of the open source community. Our goal is to combine language models with specialized tools for tasks like web search, crawling, and Python code execution, while giving back to the community that made this possible. **DeerFlow** (**D**eep **E**xploration and **E**fficient **R**esearch **Flow**) is a community-driven Deep Research framework that builds upon the incredible work of the open source community. Our goal is to combine language models with specialized tools for tasks like web search, crawling, and Python code execution, while giving back to the community that made this possible.
Currently, DeerFlow has officially entered the [FaaS Application Center of Volcengine](https://console.volcengine.com/vefaas/region:vefaas+cn-beijing/market). Users can experience it online through the [experience link](https://console.volcengine.com/vefaas/region:vefaas+cn-beijing/market/deerflow/?channel=github&source=deerflow) to intuitively feel its powerful functions and convenient operations. At the same time, to meet the deployment needs of different users, DeerFlow supports one-click deployment based on Volcengine. Click the [deployment link](https://console.volcengine.com/vefaas/region:vefaas+cn-beijing/application/create?templateId=683adf9e372daa0008aaed5c&channel=github&source=deerflow) to quickly complete the deployment process and start an efficient research journey.
Please visit [our official website](https://deerflow.tech/) for more details. Please visit [our official website](https://deerflow.tech/) for more details.
## Demo ## Demo
### Video ### Video
<https://github.com/user-attachments/assets/f3786598-1f2a-4d07-919e-8b99dfa1de3e> https://github.com/user-attachments/assets/f3786598-1f2a-4d07-919e-8b99dfa1de3e
In this demo, we showcase how to use DeerFlow to: In this demo, we showcase how to use DeerFlow to:
@@ -147,18 +144,19 @@ Explore more details in the [`web`](./web/) directory.
## Supported Search Engines ## Supported Search Engines
### Web Search
DeerFlow supports multiple search engines that can be configured in your `.env` file using the `SEARCH_API` variable: DeerFlow supports multiple search engines that can be configured in your `.env` file using the `SEARCH_API` variable:
- **Tavily** (default): A specialized search API for AI applications - **Tavily** (default): A specialized search API for AI applications
- Requires `TAVILY_API_KEY` in your `.env` file - Requires `TAVILY_API_KEY` in your `.env` file
- Sign up at: https://app.tavily.com/home - Sign up at: https://app.tavily.com/home
- **DuckDuckGo**: Privacy-focused search engine - **DuckDuckGo**: Privacy-focused search engine
- No API key required - No API key required
- **Brave Search**: Privacy-focused search engine with advanced features - **Brave Search**: Privacy-focused search engine with advanced features
- Requires `BRAVE_SEARCH_API_KEY` in your `.env` file - Requires `BRAVE_SEARCH_API_KEY` in your `.env` file
- Sign up at: https://brave.com/search/api/ - Sign up at: https://brave.com/search/api/
@@ -173,19 +171,6 @@ To configure your preferred search engine, set the `SEARCH_API` variable in your
SEARCH_API=tavily SEARCH_API=tavily
``` ```
### Private Knowledgebase
DeerFlow support private knowledgebase such as ragflow and vikingdb, so that you can use your private documents to answer questions.
- **[RAGFlow](https://ragflow.io/docs/dev/)**open source RAG engine
```
# examples in .env.example
RAG_PROVIDER=ragflow
RAGFLOW_API_URL="http://localhost:9388"
RAGFLOW_API_KEY="ragflow-xxx"
RAGFLOW_RETRIEVAL_SIZE=10
```
## Features ## Features
### Core Capabilities ### Core Capabilities
@@ -199,14 +184,10 @@ DeerFlow support private knowledgebase such as ragflow and vikingdb, so that you
### Tools and MCP Integrations ### Tools and MCP Integrations
- 🔍 **Search and Retrieval** - 🔍 **Search and Retrieval**
- Web search via Tavily, Brave Search and more - Web search via Tavily, Brave Search and more
- Crawling with Jina - Crawling with Jina
- Advanced content extraction - Advanced content extraction
- Support for private knowledgebase
- 📃 **RAG Integration**
- Supports mentioning files from [RAGFlow](https://github.com/infiniflow/ragflow) within the input box. [Start up RAGFlow server](https://ragflow.io/docs/dev/).
- 🔗 **MCP Seamless Integration** - 🔗 **MCP Seamless Integration**
- Expand capabilities for private domain access, knowledge graph, web browsing and more - Expand capabilities for private domain access, knowledge graph, web browsing and more
@@ -215,6 +196,7 @@ DeerFlow support private knowledgebase such as ragflow and vikingdb, so that you
### Human Collaboration ### Human Collaboration
- 🧠 **Human-in-the-loop** - 🧠 **Human-in-the-loop**
- Supports interactive modification of research plans using natural language - Supports interactive modification of research plans using natural language
- Supports auto-acceptance of research plans - Supports auto-acceptance of research plans
@@ -370,7 +352,6 @@ When you submit a research topic in the Studio UI, you'll be able to see the ent
DeerFlow supports LangSmith tracing to help you debug and monitor your workflows. To enable LangSmith tracing: DeerFlow supports LangSmith tracing to help you debug and monitor your workflows. To enable LangSmith tracing:
1. Make sure your `.env` file has the following configurations (see `.env.example`): 1. Make sure your `.env` file has the following configurations (see `.env.example`):
```bash ```bash
LANGSMITH_TRACING=true LANGSMITH_TRACING=true
LANGSMITH_ENDPOINT="https://api.smith.langchain.com" LANGSMITH_ENDPOINT="https://api.smith.langchain.com"
@@ -522,7 +503,6 @@ DeerFlow includes a human in the loop mechanism that allows you to review, edit,
- Via API: Set `auto_accepted_plan: true` in your request - Via API: Set `auto_accepted_plan: true` in your request
4. **API Integration**: When using the API, you can provide feedback through the `feedback` parameter: 4. **API Integration**: When using the API, you can provide feedback through the `feedback` parameter:
```json ```json
{ {
"messages": [{ "role": "user", "content": "What is quantum computing?" }], "messages": [{ "role": "user", "content": "What is quantum computing?" }],
@@ -558,8 +538,6 @@ We would like to extend our sincere appreciation to the following projects for t
- **[LangChain](https://github.com/langchain-ai/langchain)**: Their exceptional framework powers our LLM interactions and chains, enabling seamless integration and functionality. - **[LangChain](https://github.com/langchain-ai/langchain)**: Their exceptional framework powers our LLM interactions and chains, enabling seamless integration and functionality.
- **[LangGraph](https://github.com/langchain-ai/langgraph)**: Their innovative approach to multi-agent orchestration has been instrumental in enabling DeerFlow's sophisticated workflows. - **[LangGraph](https://github.com/langchain-ai/langgraph)**: Their innovative approach to multi-agent orchestration has been instrumental in enabling DeerFlow's sophisticated workflows.
- **[Novel](https://github.com/steven-tey/novel)**: Their Notion-style WYSIWYG editor supports our report editing and AI-assisted rewriting.
- **[RAGFlow](https://github.com/infiniflow/ragflow)**: We have achieved support for research on users' private knowledge bases through integration with RAGFlow.
These projects exemplify the transformative power of open-source collaboration, and we are proud to build upon their foundations. These projects exemplify the transformative power of open-source collaboration, and we are proud to build upon their foundations.
+11 -20
View File
@@ -11,18 +11,15 @@
**DeerFlow** (**D**eep **E**xploration and **E**fficient **R**esearch **Flow**) ist ein Community-getriebenes Framework für tiefgehende Recherche, das auf der großartigen Arbeit der Open-Source-Community aufbaut. Unser Ziel ist es, Sprachmodelle mit spezialisierten Werkzeugen für Aufgaben wie Websuche, Crawling und Python-Code-Ausführung zu kombinieren und gleichzeitig der Community, die dies möglich gemacht hat, etwas zurückzugeben. **DeerFlow** (**D**eep **E**xploration and **E**fficient **R**esearch **Flow**) ist ein Community-getriebenes Framework für tiefgehende Recherche, das auf der großartigen Arbeit der Open-Source-Community aufbaut. Unser Ziel ist es, Sprachmodelle mit spezialisierten Werkzeugen für Aufgaben wie Websuche, Crawling und Python-Code-Ausführung zu kombinieren und gleichzeitig der Community, die dies möglich gemacht hat, etwas zurückzugeben.
Derzeit ist DeerFlow offiziell in das FaaS-Anwendungszentrum von Volcengine eingezogen. Benutzer können es über den Erfahrungslink online erleben, um seine leistungsstarken Funktionen und bequemen Operationen intuitiv zu spüren. Gleichzeitig unterstützt DeerFlow zur Erfüllung der Bereitstellungsanforderungen verschiedener Benutzer die Ein-Klick-Bereitstellung basierend auf Volcengine. Klicken Sie auf den Bereitstellungslink, um den Bereitstellungsprozess schnell abzuschließen und eine effiziente Forschungsreise zu beginnen.
Besuchen Sie [unsere offizielle Website](https://deerflow.tech/) für weitere Details. Besuchen Sie [unsere offizielle Website](https://deerflow.tech/) für weitere Details.
## Demo ## Demo
### Video ### Video
<https://github.com/user-attachments/assets/f3786598-1f2a-4d07-919e-8b99dfa1de3e> https://github.com/user-attachments/assets/f3786598-1f2a-4d07-919e-8b99dfa1de3e
In dieser Demo zeigen wir, wie man DeerFlow nutzt, um: In dieser Demo zeigen wir, wie man DeerFlow nutzt, um:
- Nahtlos mit MCP-Diensten zu integrieren - Nahtlos mit MCP-Diensten zu integrieren
- Den Prozess der tiefgehenden Recherche durchzuführen und einen umfassenden Bericht mit Bildern zu erstellen - Den Prozess der tiefgehenden Recherche durchzuführen und einen umfassenden Bericht mit Bildern zu erstellen
- Podcast-Audio basierend auf dem generierten Bericht zu erstellen - Podcast-Audio basierend auf dem generierten Bericht zu erstellen
@@ -37,6 +34,7 @@ In dieser Demo zeigen wir, wie man DeerFlow nutzt, um:
--- ---
## 📑 Inhaltsverzeichnis ## 📑 Inhaltsverzeichnis
- [🚀 Schnellstart](#schnellstart) - [🚀 Schnellstart](#schnellstart)
@@ -50,12 +48,12 @@ In dieser Demo zeigen wir, wie man DeerFlow nutzt, um:
- [💖 Danksagungen](#danksagungen) - [💖 Danksagungen](#danksagungen)
- [⭐ Star-Verlauf](#star-verlauf) - [⭐ Star-Verlauf](#star-verlauf)
## Schnellstart ## Schnellstart
DeerFlow ist in Python entwickelt und kommt mit einer in Node.js geschriebenen Web-UI. Um einen reibungslosen Einrichtungsprozess zu gewährleisten, empfehlen wir die Verwendung der folgenden Tools: DeerFlow ist in Python entwickelt und kommt mit einer in Node.js geschriebenen Web-UI. Um einen reibungslosen Einrichtungsprozess zu gewährleisten, empfehlen wir die Verwendung der folgenden Tools:
### Empfohlene Tools ### Empfohlene Tools
- **[`uv`](https://docs.astral.sh/uv/getting-started/installation/):** - **[`uv`](https://docs.astral.sh/uv/getting-started/installation/):**
Vereinfacht die Verwaltung von Python-Umgebungen und Abhängigkeiten. `uv` erstellt automatisch eine virtuelle Umgebung im Stammverzeichnis und installiert alle erforderlichen Pakete für Sie—keine manuelle Installation von Python-Umgebungen notwendig. Vereinfacht die Verwaltung von Python-Umgebungen und Abhängigkeiten. `uv` erstellt automatisch eine virtuelle Umgebung im Stammverzeichnis und installiert alle erforderlichen Pakete für Sie—keine manuelle Installation von Python-Umgebungen notwendig.
@@ -66,14 +64,11 @@ DeerFlow ist in Python entwickelt und kommt mit einer in Node.js geschriebenen W
Installieren und verwalten Sie Abhängigkeiten des Node.js-Projekts. Installieren und verwalten Sie Abhängigkeiten des Node.js-Projekts.
### Umgebungsanforderungen ### Umgebungsanforderungen
Stellen Sie sicher, dass Ihr System die folgenden Mindestanforderungen erfüllt: Stellen Sie sicher, dass Ihr System die folgenden Mindestanforderungen erfüllt:
- **[Python](https://www.python.org/downloads/):** Version `3.12+` - **[Python](https://www.python.org/downloads/):** Version `3.12+`
- **[Node.js](https://nodejs.org/en/download/):** Version `22+` - **[Node.js](https://nodejs.org/en/download/):** Version `22+`
### Installation ### Installation
```bash ```bash
# Repository klonen # Repository klonen
git clone https://github.com/bytedance/deer-flow.git git clone https://github.com/bytedance/deer-flow.git
@@ -141,20 +136,21 @@ bootstrap.bat -d
Weitere Details finden Sie im Verzeichnis [`web`](./web/). Weitere Details finden Sie im Verzeichnis [`web`](./web/).
## Unterstützte Suchmaschinen ## Unterstützte Suchmaschinen
DeerFlow unterstützt mehrere Suchmaschinen, die in Ihrer `.env`-Datei über die Variable `SEARCH_API` konfiguriert werden können: DeerFlow unterstützt mehrere Suchmaschinen, die in Ihrer `.env`-Datei über die Variable `SEARCH_API` konfiguriert werden können:
- **Tavily** (Standard): Eine spezialisierte Such-API für KI-Anwendungen - **Tavily** (Standard): Eine spezialisierte Such-API für KI-Anwendungen
- Erfordert `TAVILY_API_KEY` in Ihrer `.env`-Datei - Erfordert `TAVILY_API_KEY` in Ihrer `.env`-Datei
- Registrieren Sie sich unter: <https://app.tavily.com/home> - Registrieren Sie sich unter: https://app.tavily.com/home
- **DuckDuckGo**: Datenschutzorientierte Suchmaschine - **DuckDuckGo**: Datenschutzorientierte Suchmaschine
- Kein API-Schlüssel erforderlich - Kein API-Schlüssel erforderlich
- **Brave Search**: Datenschutzorientierte Suchmaschine mit erweiterten Funktionen - **Brave Search**: Datenschutzorientierte Suchmaschine mit erweiterten Funktionen
- Erfordert `BRAVE_SEARCH_API_KEY` in Ihrer `.env`-Datei - Erfordert `BRAVE_SEARCH_API_KEY` in Ihrer `.env`-Datei
- Registrieren Sie sich unter: <https://brave.com/search/api/> - Registrieren Sie sich unter: https://brave.com/search/api/
- **Arxiv**: Wissenschaftliche Papiersuche für akademische Forschung - **Arxiv**: Wissenschaftliche Papiersuche für akademische Forschung
- Kein API-Schlüssel erforderlich - Kein API-Schlüssel erforderlich
@@ -206,6 +202,7 @@ SEARCH_API=tavily
- Automatisierte Erstellung einfacher PowerPoint-Präsentationen - Automatisierte Erstellung einfacher PowerPoint-Präsentationen
- Anpassbare Vorlagen für maßgeschneiderte Inhalte - Anpassbare Vorlagen für maßgeschneiderte Inhalte
## Architektur ## Architektur
DeerFlow implementiert eine modulare Multi-Agenten-Systemarchitektur, die für automatisierte Forschung und Codeanalyse konzipiert ist. Das System basiert auf LangGraph und ermöglicht einen flexiblen zustandsbasierten Workflow, bei dem Komponenten über ein klar definiertes Nachrichtenübermittlungssystem kommunizieren. DeerFlow implementiert eine modulare Multi-Agenten-Systemarchitektur, die für automatisierte Forschung und Codeanalyse konzipiert ist. Das System basiert auf LangGraph und ermöglicht einen flexiblen zustandsbasierten Workflow, bei dem Komponenten über ein klar definiertes Nachrichtenübermittlungssystem kommunizieren.
@@ -256,6 +253,7 @@ curl --location 'http://localhost:8000/api/tts' \
--output speech.mp3 --output speech.mp3
``` ```
## Entwicklung ## Entwicklung
### Testen ### Testen
@@ -313,10 +311,9 @@ langgraph dev
``` ```
Nach dem Start des LangGraph-Servers sehen Sie mehrere URLs im Terminal: Nach dem Start des LangGraph-Servers sehen Sie mehrere URLs im Terminal:
- API: http://127.0.0.1:2024
- API: <http://127.0.0.1:2024> - Studio UI: https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024
- Studio UI: <https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024> - API-Dokumentation: http://127.0.0.1:2024/docs
- API-Dokumentation: <http://127.0.0.1:2024/docs>
Öffnen Sie den Studio UI-Link in Ihrem Browser, um auf die Debugging-Schnittstelle zuzugreifen. Öffnen Sie den Studio UI-Link in Ihrem Browser, um auf die Debugging-Schnittstelle zuzugreifen.
@@ -331,7 +328,6 @@ In der Studio UI können Sie:
5. Feedback während der Planungsphase geben, um Forschungspläne zu verfeinern 5. Feedback während der Planungsphase geben, um Forschungspläne zu verfeinern
Wenn Sie ein Forschungsthema in der Studio UI einreichen, können Sie die gesamte Workflow-Ausführung sehen, einschließlich: Wenn Sie ein Forschungsthema in der Studio UI einreichen, können Sie die gesamte Workflow-Ausführung sehen, einschließlich:
- Die Planungsphase, in der der Forschungsplan erstellt wird - Die Planungsphase, in der der Forschungsplan erstellt wird
- Die Feedback-Schleife, in der Sie den Plan ändern können - Die Feedback-Schleife, in der Sie den Plan ändern können
- Die Forschungs- und Schreibphasen für jeden Abschnitt - Die Forschungs- und Schreibphasen für jeden Abschnitt
@@ -342,7 +338,6 @@ Wenn Sie ein Forschungsthema in der Studio UI einreichen, können Sie die gesamt
DeerFlow unterstützt LangSmith-Tracing, um Ihnen beim Debuggen und Überwachen Ihrer Workflows zu helfen. Um LangSmith-Tracing zu aktivieren: DeerFlow unterstützt LangSmith-Tracing, um Ihnen beim Debuggen und Überwachen Ihrer Workflows zu helfen. Um LangSmith-Tracing zu aktivieren:
1. Stellen Sie sicher, dass Ihre `.env`-Datei die folgenden Konfigurationen enthält (siehe `.env.example`): 1. Stellen Sie sicher, dass Ihre `.env`-Datei die folgenden Konfigurationen enthält (siehe `.env.example`):
```bash ```bash
LANGSMITH_TRACING=true LANGSMITH_TRACING=true
LANGSMITH_ENDPOINT="https://api.smith.langchain.com" LANGSMITH_ENDPOINT="https://api.smith.langchain.com"
@@ -351,7 +346,6 @@ DeerFlow unterstützt LangSmith-Tracing, um Ihnen beim Debuggen und Überwachen
``` ```
2. Starten Sie das Tracing mit LangSmith lokal, indem Sie folgenden Befehl ausführen: 2. Starten Sie das Tracing mit LangSmith lokal, indem Sie folgenden Befehl ausführen:
```bash ```bash
langgraph dev langgraph dev
``` ```
@@ -425,7 +419,6 @@ uv run main.py --help
Die Anwendung unterstützt jetzt einen interaktiven Modus mit eingebauten Fragen in Englisch und Chinesisch: Die Anwendung unterstützt jetzt einen interaktiven Modus mit eingebauten Fragen in Englisch und Chinesisch:
1. Starten Sie den interaktiven Modus: 1. Starten Sie den interaktiven Modus:
```bash ```bash
uv run main.py --interactive uv run main.py --interactive
``` ```
@@ -451,7 +444,6 @@ DeerFlow enthält einen Mensch-in-der-Schleife-Mechanismus, der es Ihnen ermögl
- Über API: Setzen Sie `auto_accepted_plan: true` in Ihrer Anfrage - Über API: Setzen Sie `auto_accepted_plan: true` in Ihrer Anfrage
4. **API-Integration**: Bei Verwendung der API können Sie Feedback über den Parameter `feedback` geben: 4. **API-Integration**: Bei Verwendung der API können Sie Feedback über den Parameter `feedback` geben:
```json ```json
{ {
"messages": [{"role": "user", "content": "Was ist Quantencomputing?"}], "messages": [{"role": "user", "content": "Was ist Quantencomputing?"}],
@@ -491,7 +483,6 @@ Wir möchten unsere aufrichtige Wertschätzung den folgenden Projekten für ihre
Diese Projekte veranschaulichen die transformative Kraft der Open-Source-Zusammenarbeit, und wir sind stolz darauf, auf ihren Grundlagen aufzubauen. Diese Projekte veranschaulichen die transformative Kraft der Open-Source-Zusammenarbeit, und wir sind stolz darauf, auf ihren Grundlagen aufzubauen.
### Hauptmitwirkende ### Hauptmitwirkende
Ein herzliches Dankeschön geht an die Hauptautoren von `DeerFlow`, deren Vision, Leidenschaft und Engagement dieses Projekt zum Leben erweckt haben: Ein herzliches Dankeschön geht an die Hauptautoren von `DeerFlow`, deren Vision, Leidenschaft und Engagement dieses Projekt zum Leben erweckt haben:
- **[Daniel Walnut](https://github.com/hetaoBackend/)** - **[Daniel Walnut](https://github.com/hetaoBackend/)**
+6 -11
View File
@@ -11,15 +11,13 @@
**DeerFlow** (**D**eep **E**xploration and **E**fficient **R**esearch **Flow**) es un marco de Investigación Profunda impulsado por la comunidad que se basa en el increíble trabajo de la comunidad de código abierto. Nuestro objetivo es combinar modelos de lenguaje con herramientas especializadas para tareas como búsqueda web, rastreo y ejecución de código Python, mientras devolvemos a la comunidad que hizo esto posible. **DeerFlow** (**D**eep **E**xploration and **E**fficient **R**esearch **Flow**) es un marco de Investigación Profunda impulsado por la comunidad que se basa en el increíble trabajo de la comunidad de código abierto. Nuestro objetivo es combinar modelos de lenguaje con herramientas especializadas para tareas como búsqueda web, rastreo y ejecución de código Python, mientras devolvemos a la comunidad que hizo esto posible.
Actualmente, DeerFlow ha ingresado oficialmente al Centro de Aplicaciones FaaS de Volcengine. Los usuarios pueden experimentarlo en línea a través del enlace de experiencia para sentir intuitivamente sus potentes funciones y operaciones convenientes. Al mismo tiempo, para satisfacer las necesidades de implementación de diferentes usuarios, DeerFlow admite la implementación con un clic basada en Volcengine. Haga clic en el enlace de implementación para completar rápidamente el proceso de implementación y comenzar un viaje de investigación eficiente.
Por favor, visita [nuestra página web oficial](https://deerflow.tech/) para más detalles. Por favor, visita [nuestra página web oficial](https://deerflow.tech/) para más detalles.
## Demostración ## Demostración
### Video ### Video
<https://github.com/user-attachments/assets/f3786598-1f2a-4d07-919e-8b99dfa1de3e> https://github.com/user-attachments/assets/f3786598-1f2a-4d07-919e-8b99dfa1de3e
En esta demostración, mostramos cómo usar DeerFlow para: En esta demostración, mostramos cómo usar DeerFlow para:
@@ -150,7 +148,7 @@ DeerFlow soporta múltiples motores de búsqueda que pueden configurarse en tu a
- **Tavily** (predeterminado): Una API de búsqueda especializada para aplicaciones de IA - **Tavily** (predeterminado): Una API de búsqueda especializada para aplicaciones de IA
- Requiere `TAVILY_API_KEY` en tu archivo `.env` - Requiere `TAVILY_API_KEY` en tu archivo `.env`
- Regístrate en: <https://app.tavily.com/home> - Regístrate en: https://app.tavily.com/home
- **DuckDuckGo**: Motor de búsqueda centrado en la privacidad - **DuckDuckGo**: Motor de búsqueda centrado en la privacidad
@@ -159,7 +157,7 @@ DeerFlow soporta múltiples motores de búsqueda que pueden configurarse en tu a
- **Brave Search**: Motor de búsqueda centrado en la privacidad con características avanzadas - **Brave Search**: Motor de búsqueda centrado en la privacidad con características avanzadas
- Requiere `BRAVE_SEARCH_API_KEY` en tu archivo `.env` - Requiere `BRAVE_SEARCH_API_KEY` en tu archivo `.env`
- Regístrate en: <https://brave.com/search/api/> - Regístrate en: https://brave.com/search/api/
- **Arxiv**: Búsqueda de artículos científicos para investigación académica - **Arxiv**: Búsqueda de artículos científicos para investigación académica
- No requiere clave API - No requiere clave API
@@ -325,9 +323,9 @@ langgraph dev
Después de iniciar el servidor LangGraph, verás varias URLs en la terminal: Después de iniciar el servidor LangGraph, verás varias URLs en la terminal:
- API: <http://127.0.0.1:2024> - API: http://127.0.0.1:2024
- UI de Studio: <https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024> - UI de Studio: https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024
- Docs de API: <http://127.0.0.1:2024/docs> - Docs de API: http://127.0.0.1:2024/docs
Abre el enlace de UI de Studio en tu navegador para acceder a la interfaz de depuración. Abre el enlace de UI de Studio en tu navegador para acceder a la interfaz de depuración.
@@ -353,7 +351,6 @@ Cuando envías un tema de investigación en la UI de Studio, podrás ver toda la
DeerFlow soporta el rastreo de LangSmith para ayudarte a depurar y monitorear tus flujos de trabajo. Para habilitar el rastreo de LangSmith: DeerFlow soporta el rastreo de LangSmith para ayudarte a depurar y monitorear tus flujos de trabajo. Para habilitar el rastreo de LangSmith:
1. Asegúrate de que tu archivo `.env` tenga las siguientes configuraciones (ver `.env.example`): 1. Asegúrate de que tu archivo `.env` tenga las siguientes configuraciones (ver `.env.example`):
```bash ```bash
LANGSMITH_TRACING=true LANGSMITH_TRACING=true
LANGSMITH_ENDPOINT="https://api.smith.langchain.com" LANGSMITH_ENDPOINT="https://api.smith.langchain.com"
@@ -362,7 +359,6 @@ DeerFlow soporta el rastreo de LangSmith para ayudarte a depurar y monitorear tu
``` ```
2. Inicia el rastreo y visualiza el grafo localmente con LangSmith ejecutando: 2. Inicia el rastreo y visualiza el grafo localmente con LangSmith ejecutando:
```bash ```bash
langgraph dev langgraph dev
``` ```
@@ -506,7 +502,6 @@ DeerFlow incluye un mecanismo de humano en el bucle que te permite revisar, edit
- Vía API: Establece `auto_accepted_plan: true` en tu solicitud - Vía API: Establece `auto_accepted_plan: true` en tu solicitud
4. **Integración API**: Cuando uses la API, puedes proporcionar retroalimentación a través del parámetro `feedback`: 4. **Integración API**: Cuando uses la API, puedes proporcionar retroalimentación a través del parámetro `feedback`:
```json ```json
{ {
"messages": [{ "role": "user", "content": "¿Qué es la computación cuántica?" }], "messages": [{ "role": "user", "content": "¿Qué es la computación cuántica?" }],
+31 -54
View File
@@ -9,19 +9,17 @@
**DeerFlow****D**eep **E**xploration and **E**fficient **R**esearch **Flow**)は、オープンソースコミュニティの素晴らしい成果の上に構築されたコミュニティ主導の深層研究フレームワークです。私たちの目標は、言語モデルとウェブ検索、クローリング、Python コード実行などの専門ツールを組み合わせながら、これを可能にしたコミュニティに貢献することです。 **DeerFlow****D**eep **E**xploration and **E**fficient **R**esearch **Flow**)は、オープンソースコミュニティの素晴らしい成果の上に構築されたコミュニティ主導の深層研究フレームワークです。私たちの目標は、言語モデルとウェブ検索、クローリング、Python コード実行などの専門ツールを組み合わせながら、これを可能にしたコミュニティに貢献することです。
現在、DeerFlow は火山引擎の FaaS アプリケーションセンターに正式に入居しています。ユーザーは体験リンクを通じてオンラインで体験し、その強力な機能と便利な操作を直感的に感じることができます。同時に、さまざまなユーザーの展開ニーズを満たすため、DeerFlow は火山引擎に基づくワンクリック展開をサポートしています。展開リンクをクリックして展開プロセスを迅速に完了し、効率的な研究の旅を始めましょう。
詳細については[DeerFlow の公式ウェブサイト](https://deerflow.tech/)をご覧ください。 詳細については[DeerFlow の公式ウェブサイト](https://deerflow.tech/)をご覧ください。
## デモ ## デモ
### ビデオ ### ビデオ
<https://github.com/user-attachments/assets/f3786598-1f2a-4d07-919e-8b99dfa1de3e> https://github.com/user-attachments/assets/f3786598-1f2a-4d07-919e-8b99dfa1de3e
このデモでは、DeerFlowの使用方法を紹介しています: このデモでは、DeerFlow の使用方法を紹介しています:
- MCPサービスとのシームレスな統合 - MCP サービスとのシームレスな統合
- 深層研究プロセスの実施と画像を含む包括的なレポートの作成 - 深層研究プロセスの実施と画像を含む包括的なレポートの作成
- 生成されたレポートに基づくポッドキャストオーディオの作成 - 生成されたレポートに基づくポッドキャストオーディオの作成
@@ -145,18 +143,21 @@ bootstrap.bat -d
DeerFlow は複数の検索エンジンをサポートしており、`.env`ファイルの`SEARCH_API`変数で設定できます: DeerFlow は複数の検索エンジンをサポートしており、`.env`ファイルの`SEARCH_API`変数で設定できます:
- **Tavily**(デフォルト):AI アプリケーション向けの専門検索 API - **Tavily**(デフォルト):AI アプリケーション向けの専門検索 API
- `.env`ファイルに`TAVILY_API_KEY`が必要 - `.env`ファイルに`TAVILY_API_KEY`が必要
- 登録先:<https://app.tavily.com/home> - 登録先:https://app.tavily.com/home
- **DuckDuckGo**:プライバシー重視の検索エンジン - **DuckDuckGo**:プライバシー重視の検索エンジン
- APIキー不要
- API キー不要
- **Brave Search**:高度な機能を備えたプライバシー重視の検索エンジン - **Brave Search**:高度な機能を備えたプライバシー重視の検索エンジン
- `.env`ファイルに`BRAVE_SEARCH_API_KEY`が必要 - `.env`ファイルに`BRAVE_SEARCH_API_KEY`が必要
- 登録先:<https://brave.com/search/api/> - 登録先:https://brave.com/search/api/
- **Arxiv**:学術研究用の科学論文検索 - **Arxiv**:学術研究用の科学論文検索
- APIキー不要 - API キー不要
- 科学・学術論文専用 - 科学・学術論文専用
お好みの検索エンジンを設定するには、`.env`ファイルで`SEARCH_API`変数を設定します: お好みの検索エンジンを設定するには、`.env`ファイルで`SEARCH_API`変数を設定します:
@@ -170,39 +171,41 @@ SEARCH_API=tavily
### コア機能 ### コア機能
- 🤖 **LLM統合** - 🤖 **LLM 統合**
- [litellm](https://docs.litellm.ai/docs/providers)を通じてほとんどのモデルの統合をサポート - [litellm](https://docs.litellm.ai/docs/providers)を通じてほとんどのモデルの統合をサポート
- Qwenなどのオープンソースモデルをサポート - Qwen などのオープンソースモデルをサポート
- OpenAI互換のAPIインターフェース - OpenAI 互換の API インターフェース
- 異なるタスクの複雑さに対応するマルチティアLLMシステム - 異なるタスクの複雑さに対応するマルチティア LLM システム
### ツールと MCP 統合 ### ツールと MCP 統合
- 🔍 **検索と取得** - 🔍 **検索と取得**
- Tavily、Brave Searchなどを通じたWeb検索
- Jinaを使用したクローリング - Tavily、Brave Search などを通じた Web 検索
- Jina を使用したクローリング
- 高度なコンテンツ抽出 - 高度なコンテンツ抽出
- 🔗 **MCPシームレス統合** - 🔗 **MCP シームレス統合**
- プライベートドメインアクセス、ナレッジグラフ、Webブラウジングなどの機能を拡張 - プライベートドメインアクセス、ナレッジグラフ、Web ブラウジングなどの機能を拡張
- 多様な研究ツールと方法論の統合を促進 - 多様な研究ツールと方法論の統合を促進
### 人間との協力 ### 人間との協力
- 🧠 **人間参加型ループ** - 🧠 **人間参加型ループ**
- 自然言語を使用した研究計画の対話的修正をサポート - 自然言語を使用した研究計画の対話的修正をサポート
- 研究計画の自動承認をサポート - 研究計画の自動承認をサポート
- 📝 **レポート後編集** - 📝 **レポート後編集**
- Notionライクなブロック編集をサポート - Notion ライクなブロック編集をサポート
- AI支援による洗練、文の短縮、拡張などのAI改良を可能に - AI 支援による洗練、文の短縮、拡張などの AI 改良を可能に
- [tiptap](https://tiptap.dev/)を活用 - [tiptap](https://tiptap.dev/)を活用
### コンテンツ作成 ### コンテンツ作成
- 🎙️ **ポッドキャストとプレゼンテーション生成** - 🎙️ **ポッドキャストとプレゼンテーション生成**
- AI駆動のポッドキャストスクリプト生成と音声合成 - AI 駆動のポッドキャストスクリプト生成と音声合成
- シンプルなPowerPointプレゼンテーションの自動作成 - シンプルな PowerPoint プレゼンテーションの自動作成
- カスタマイズ可能なテンプレートで個別のコンテンツに対応 - カスタマイズ可能なテンプレートで個別のコンテンツに対応
## アーキテクチャ ## アーキテクチャ
@@ -238,27 +241,6 @@ DeerFlow は、自動研究とコード分析のためのモジュラーなマ
- 収集した情報を処理および構造化 - 収集した情報を処理および構造化
- 包括的な研究レポートを生成 - 包括的な研究レポートを生成
## テキスト読み上げ統合
DeerFlowには現在、研究レポートを音声に変換できるテキスト読み上げ(TTS)機能が含まれています。この機能は火山引擎TTS APIを使用して高品質なテキストオーディオを生成します。速度、音量、ピッチなどの特性もカスタマイズ可能です。
### TTS APIの使用
`/api/tts`エンドポイントからTTS機能にアクセスできます:
```bash
# curlを使用したAPI呼び出し例
curl --location 'http://localhost:8000/api/tts' \
--header 'Content-Type: application/json' \
--data '{
"text": "これはテキスト読み上げ機能のテストです。",
"speed_ratio": 1.0,
"volume_ratio": 1.0,
"pitch_ratio": 1.0
}' \
--output speech.mp3
```
## 開発 ## 開発
### テスト ### テスト
@@ -315,15 +297,11 @@ pip install -U "langgraph-cli[inmem]"
langgraph dev langgraph dev
``` ```
LangGraphサーバーを開始すると、端末にいくつかのURLが表示されます: LangGraph サーバーを開始すると、端末にいくつかの URL が表示されます:
- API: <http://127.0.0.1:2024> - API: http://127.0.0.1:2024
- Studio UI: <https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024> - Studio UI: https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024
- APIドキュメント: <http://127.0.0.1:2024/docs> - API ドキュメント: http://127.0.0.1:2024/docs
- API: <http://127.0.0.1:2024>
- Studio UI: <https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024>
- APIドキュメント: <http://127.0.0.1:2024/docs>
ブラウザで Studio UI リンクを開いてデバッグインターフェースにアクセスします。 ブラウザで Studio UI リンクを開いてデバッグインターフェースにアクセスします。
@@ -337,7 +315,7 @@ Studio UI では、次のことができます:
4. 各コンポーネントの入力と出力を検査して問題をデバッグ 4. 各コンポーネントの入力と出力を検査して問題をデバッグ
5. 計画段階でフィードバックを提供して研究計画を洗練 5. 計画段階でフィードバックを提供して研究計画を洗練
Studio UIで研究トピックを送信すると、次を含む全ワークフロー実行プロセスを見ることができます: Studio UI で研究トピックを送信すると、次を含む全ワークフロー実行プロセスを見ることができます:
- 研究計画を作成する計画段階 - 研究計画を作成する計画段階
- 計画を修正できるフィードバックループ - 計画を修正できるフィードバックループ
@@ -349,7 +327,6 @@ Studio UIで研究トピックを送信すると、次を含む全ワークフ
DeerFlow は LangSmith トレース機能をサポートしており、ワークフローのデバッグとモニタリングに役立ちます。LangSmith トレースを有効にするには: DeerFlow は LangSmith トレース機能をサポートしており、ワークフローのデバッグとモニタリングに役立ちます。LangSmith トレースを有効にするには:
1. `.env` ファイルに次の設定があることを確認してください(`.env.example` を参照): 1. `.env` ファイルに次の設定があることを確認してください(`.env.example` を参照):
```bash ```bash
LANGSMITH_TRACING=true LANGSMITH_TRACING=true
LANGSMITH_ENDPOINT="https://api.smith.langchain.com" LANGSMITH_ENDPOINT="https://api.smith.langchain.com"
@@ -358,7 +335,6 @@ DeerFlow は LangSmith トレース機能をサポートしており、ワーク
``` ```
2. 次のコマンドを実行して LangSmith トレースを開始します: 2. 次のコマンドを実行して LangSmith トレースを開始します:
```bash ```bash
langgraph dev langgraph dev
``` ```
@@ -520,8 +496,9 @@ DeerFlow には人間参加型ループメカニズムが含まれており、
3. **自動承認**:レビュープロセスをスキップするために自動承認を有効にできます: 3. **自動承認**:レビュープロセスをスキップするために自動承認を有効にできます:
4. **API統合**APIを使用する場合、`feedback`パラメータでフィードバックを提供できます: - API 経由:リクエストで`auto_accepted_plan: true`を設定
4. **API 統合**API を使用する場合、`feedback`パラメータでフィードバックを提供できます:
```json ```json
{ {
"messages": [ "messages": [
+10 -10
View File
@@ -12,15 +12,13 @@
**DeerFlow** (**D**eep **E**xploration and **E**fficient **R**esearch **Flow**) é um framework de Pesquisa Profunda orientado-a-comunidade que baseia-se em um íncrivel trabalho da comunidade open source. Nosso objetivo é combinar modelos de linguagem com ferramentas especializadas para tarefas como busca na web, crawling, e execução de código Python, enquanto retribui com a comunidade que o tornou possível. **DeerFlow** (**D**eep **E**xploration and **E**fficient **R**esearch **Flow**) é um framework de Pesquisa Profunda orientado-a-comunidade que baseia-se em um íncrivel trabalho da comunidade open source. Nosso objetivo é combinar modelos de linguagem com ferramentas especializadas para tarefas como busca na web, crawling, e execução de código Python, enquanto retribui com a comunidade que o tornou possível.
Atualmente, o DeerFlow entrou oficialmente no Centro de Aplicações FaaS da Volcengine. Os usuários podem experimentá-lo online através do link de experiência para sentir intuitivamente suas funções poderosas e operações convenientes. Ao mesmo tempo, para atender às necessidades de implantação de diferentes usuários, o DeerFlow suporta implantação com um clique baseada na Volcengine. Clique no link de implantação para completar rapidamente o processo de implantação e iniciar uma jornada de pesquisa eficiente.
Por favor, visite [Nosso Site Oficial](https://deerflow.tech/) para maiores detalhes. Por favor, visite [Nosso Site Oficial](https://deerflow.tech/) para maiores detalhes.
## Demo ## Demo
### Video ### Video
<https://github.com/user-attachments/assets/f3786598-1f2a-4d07-919e-8b99dfa1de3e> https://github.com/user-attachments/assets/f3786598-1f2a-4d07-919e-8b99dfa1de3e
Nesse demo, nós demonstramos como usar o DeerFlow para: Nesse demo, nós demonstramos como usar o DeerFlow para:
In this demo, we showcase how to use DeerFlow to: In this demo, we showcase how to use DeerFlow to:
@@ -148,12 +146,13 @@ Explore mais detalhes no diretório [`web`](./web/) .
## Mecanismos de Busca Suportados ## Mecanismos de Busca Suportados
DeerFlow suporta múltiplos mecanismos de busca que podem ser configurados no seu arquivo `.env` usando a variável `SEARCH_API`: DeerFlow suporta múltiplos mecanismos de busca que podem ser configurados no seu arquivo `.env` usando a variável `SEARCH_API`:
- **Tavily** (padrão): Uma API de busca especializada para aplicações de IA - **Tavily** (padrão): Uma API de busca especializada para aplicações de IA
- Requer `TAVILY_API_KEY` no seu arquivo `.env` - Requer `TAVILY_API_KEY` no seu arquivo `.env`
- Inscreva-se em: <https://app.tavily.com/home> - Inscreva-se em: https://app.tavily.com/home
- **DuckDuckGo**: Mecanismo de busca focado em privacidade - **DuckDuckGo**: Mecanismo de busca focado em privacidade
@@ -162,7 +161,7 @@ DeerFlow suporta múltiplos mecanismos de busca que podem ser configurados no se
- **Brave Search**: Mecanismo de busca focado em privacidade com funcionalidades avançadas - **Brave Search**: Mecanismo de busca focado em privacidade com funcionalidades avançadas
- Requer `BRAVE_SEARCH_API_KEY` no seu arquivo `.env` - Requer `BRAVE_SEARCH_API_KEY` no seu arquivo `.env`
- Inscreva-se em: <https://brave.com/search/api/> - Inscreva-se em: https://brave.com/search/api/
- **Arxiv**: Busca de artigos científicos para pesquisa acadêmica - **Arxiv**: Busca de artigos científicos para pesquisa acadêmica
- Não requer chave API - Não requer chave API
@@ -203,6 +202,7 @@ SEARCH_API=tavily
- 🧠 **Humano-no-processo** - 🧠 **Humano-no-processo**
- Suporta modificação interativa de planos de pesquisa usando linguagem natural - Suporta modificação interativa de planos de pesquisa usando linguagem natural
- Suporta auto-aceite de planos de pesquisa - Suporta auto-aceite de planos de pesquisa
@@ -331,9 +331,9 @@ langgraph dev
Após iniciar o servidor LangGraph, você verá diversas URLs no seu terminal: Após iniciar o servidor LangGraph, você verá diversas URLs no seu terminal:
- API: <http://127.0.0.1:2024> - API: http://127.0.0.1:2024
- Studio UI: <https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024> - Studio UI: https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024
- API Docs: <http://127.0.0.1:2024/docs> - API Docs: http://127.0.0.1:2024/docs
Abra o link do Studio UI no seu navegador para acessar a interface de depuração. Abra o link do Studio UI no seu navegador para acessar a interface de depuração.
@@ -341,6 +341,7 @@ Abra o link do Studio UI no seu navegador para acessar a interface de depuraçã
No Studio UI, você pode: No Studio UI, você pode:
1. Visualizar o grafo do fluxo de trabalho e como seus componentes se conectam 1. Visualizar o grafo do fluxo de trabalho e como seus componentes se conectam
2. Rastrear a execução em tempo-real e ver como os dados fluem através do sistema 2. Rastrear a execução em tempo-real e ver como os dados fluem através do sistema
3. Inspecionar o estado de cada passo do fluxo de trabalho 3. Inspecionar o estado de cada passo do fluxo de trabalho
@@ -388,7 +389,7 @@ docker compose build
docker compose up docker compose up
``` ```
## Exemplos ## Exemplos:
Os seguintes exemplos demonstram as capacidades do DeerFlow: Os seguintes exemplos demonstram as capacidades do DeerFlow:
@@ -492,7 +493,6 @@ DeerFlow inclue um mecanismo de humano no processo que permite a você revisar,
- Via API: Defina `auto_accepted_plan: true` na sua requisição - Via API: Defina `auto_accepted_plan: true` na sua requisição
4. **Integração de API**: Quanto usar a API, você pode fornecer um feedback através do parâmetro `feedback`: 4. **Integração de API**: Quanto usar a API, você pode fornecer um feedback através do parâmetro `feedback`:
```json ```json
{ {
"messages": [{ "role": "user", "content": "O que é computação quântica?" }], "messages": [{ "role": "user", "content": "O que é computação quântica?" }],
+6 -11
View File
@@ -11,15 +11,13 @@
**DeerFlow** (**D**eep **E**xploration and **E**fficient **R**esearch **Flow**) - это фреймворк для глубокого исследования, разработанный сообществом и основанный на впечатляющей работе сообщества открытого кода. Наша цель - объединить языковые модели со специализированными инструментами для таких задач, как веб-поиск, сканирование и выполнение кода Python, одновременно возвращая пользу сообществу, которое сделало это возможным. **DeerFlow** (**D**eep **E**xploration and **E**fficient **R**esearch **Flow**) - это фреймворк для глубокого исследования, разработанный сообществом и основанный на впечатляющей работе сообщества открытого кода. Наша цель - объединить языковые модели со специализированными инструментами для таких задач, как веб-поиск, сканирование и выполнение кода Python, одновременно возвращая пользу сообществу, которое сделало это возможным.
В настоящее время DeerFlow официально вошел в Центр приложений FaaS Volcengine. Пользователи могут испытать его онлайн через ссылку для опыта, чтобы интуитивно почувствовать его мощные функции и удобные операции. В то же время, для удовлетворения потребностей развертывания различных пользователей, DeerFlow поддерживает развертывание одним кликом на основе Volcengine. Нажмите на ссылку развертывания, чтобы быстро завершить процесс развертывания и начать эффективное исследовательское путешествие.
Пожалуйста, посетите [наш официальный сайт](https://deerflow.tech/) для получения дополнительной информации. Пожалуйста, посетите [наш официальный сайт](https://deerflow.tech/) для получения дополнительной информации.
## Демонстрация ## Демонстрация
### Видео ### Видео
<https://github.com/user-attachments/assets/f3786598-1f2a-4d07-919e-8b99dfa1de3e> https://github.com/user-attachments/assets/f3786598-1f2a-4d07-919e-8b99dfa1de3e
В этой демонстрации мы показываем, как использовать DeerFlow для: В этой демонстрации мы показываем, как использовать DeerFlow для:
@@ -150,7 +148,7 @@ DeerFlow поддерживает несколько поисковых сист
- **Tavily** (по умолчанию): Специализированный поисковый API для приложений ИИ - **Tavily** (по умолчанию): Специализированный поисковый API для приложений ИИ
- Требуется `TAVILY_API_KEY` в вашем файле `.env` - Требуется `TAVILY_API_KEY` в вашем файле `.env`
- Зарегистрируйтесь на: <https://app.tavily.com/home> - Зарегистрируйтесь на: https://app.tavily.com/home
- **DuckDuckGo**: Поисковая система, ориентированная на конфиденциальность - **DuckDuckGo**: Поисковая система, ориентированная на конфиденциальность
@@ -159,7 +157,7 @@ DeerFlow поддерживает несколько поисковых сист
- **Brave Search**: Поисковая система, ориентированная на конфиденциальность, с расширенными функциями - **Brave Search**: Поисковая система, ориентированная на конфиденциальность, с расширенными функциями
- Требуется `BRAVE_SEARCH_API_KEY` в вашем файле `.env` - Требуется `BRAVE_SEARCH_API_KEY` в вашем файле `.env`
- Зарегистрируйтесь на: <https://brave.com/search/api/> - Зарегистрируйтесь на: https://brave.com/search/api/
- **Arxiv**: Поиск научных статей для академических исследований - **Arxiv**: Поиск научных статей для академических исследований
- Не требуется API-ключ - Не требуется API-ключ
@@ -325,9 +323,9 @@ langgraph dev
После запуска сервера LangGraph вы увидите несколько URL в терминале: После запуска сервера LangGraph вы увидите несколько URL в терминале:
- API: <http://127.0.0.1:2024> - API: http://127.0.0.1:2024
- Studio UI: <https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024> - Studio UI: https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024
- API Docs: <http://127.0.0.1:2024/docs> - API Docs: http://127.0.0.1:2024/docs
Откройте ссылку Studio UI в вашем браузере для доступа к интерфейсу отладки. Откройте ссылку Studio UI в вашем браузере для доступа к интерфейсу отладки.
@@ -353,7 +351,6 @@ langgraph dev
DeerFlow поддерживает трассировку LangSmith, чтобы помочь вам отладить и контролировать ваши рабочие процессы. Чтобы включить трассировку LangSmith: DeerFlow поддерживает трассировку LangSmith, чтобы помочь вам отладить и контролировать ваши рабочие процессы. Чтобы включить трассировку LangSmith:
1. Убедитесь, что в вашем файле `.env` есть следующие конфигурации (см. `.env.example`): 1. Убедитесь, что в вашем файле `.env` есть следующие конфигурации (см. `.env.example`):
```bash ```bash
LANGSMITH_TRACING=true LANGSMITH_TRACING=true
LANGSMITH_ENDPOINT="https://api.smith.langchain.com" LANGSMITH_ENDPOINT="https://api.smith.langchain.com"
@@ -362,7 +359,6 @@ DeerFlow поддерживает трассировку LangSmith, чтобы
``` ```
2. Запустите трассировку и визуализируйте граф локально с LangSmith, выполнив: 2. Запустите трассировку и визуализируйте граф локально с LangSmith, выполнив:
```bash ```bash
langgraph dev langgraph dev
``` ```
@@ -506,7 +502,6 @@ DeerFlow включает механизм "человек в контуре",
- Через API: Установите `auto_accepted_plan: true` в вашем запросе - Через API: Установите `auto_accepted_plan: true` в вашем запросе
4. **Интеграция API**: При использовании API вы можете предоставить обратную связь через параметр `feedback`: 4. **Интеграция API**: При использовании API вы можете предоставить обратную связь через параметр `feedback`:
```json ```json
{ {
"messages": [{ "role": "user", "content": "Что такое квантовые вычисления?" }], "messages": [{ "role": "user", "content": "Что такое квантовые вычисления?" }],
+26 -69
View File
@@ -9,15 +9,13 @@
**DeerFlow****D**eep **E**xploration and **E**fficient **R**esearch **Flow**)是一个社区驱动的深度研究框架,它建立在开源社区的杰出工作基础之上。我们的目标是将语言模型与专业工具(如网络搜索、爬虫和 Python 代码执行)相结合,同时回馈使这一切成为可能的社区。 **DeerFlow****D**eep **E**xploration and **E**fficient **R**esearch **Flow**)是一个社区驱动的深度研究框架,它建立在开源社区的杰出工作基础之上。我们的目标是将语言模型与专业工具(如网络搜索、爬虫和 Python 代码执行)相结合,同时回馈使这一切成为可能的社区。
目前,DeerFlow 已正式入驻[火山引擎的 FaaS 应用中心](https://console.volcengine.com/vefaas/region:vefaas+cn-beijing/market),用户可通过[体验链接](https://console.volcengine.com/vefaas/region:vefaas+cn-beijing/market/deerflow/?channel=github&source=deerflow)进行在线体验,直观感受其强大功能与便捷操作;同时,为满足不同用户的部署需求,DeerFlow 支持基于火山引擎一键部署,点击[部署链接](https://console.volcengine.com/vefaas/region:vefaas+cn-beijing/application/create?templateId=683adf9e372daa0008aaed5c&channel=github&source=deerflow)即可快速完成部署流程,开启高效研究之旅。
请访问[DeerFlow 的官方网站](https://deerflow.tech/)了解更多详情。 请访问[DeerFlow 的官方网站](https://deerflow.tech/)了解更多详情。
## 演示 ## 演示
### 视频 ### 视频
<https://github.com/user-attachments/assets/f3786598-1f2a-4d07-919e-8b99dfa1de3e> https://github.com/user-attachments/assets/f3786598-1f2a-4d07-919e-8b99dfa1de3e
在此演示中,我们展示了如何使用 DeerFlow: 在此演示中,我们展示了如何使用 DeerFlow:
@@ -46,7 +44,7 @@
- [❓ 常见问题](#常见问题) - [❓ 常见问题](#常见问题)
- [📜 许可证](#许可证) - [📜 许可证](#许可证)
- [💖 致谢](#致谢) - [💖 致谢](#致谢)
- [⭐ Star History](#star-history) - [⭐ Star History](#star-History)
## 快速开始 ## 快速开始
@@ -108,7 +106,7 @@ pnpm install
请参阅[配置指南](docs/configuration_guide.md)获取更多详情。 请参阅[配置指南](docs/configuration_guide.md)获取更多详情。
> [! 注意] > [!注意]
> 在启动项目之前,请仔细阅读指南,并更新配置以匹配您的特定设置和要求。 > 在启动项目之前,请仔细阅读指南,并更新配置以匹配您的特定设置和要求。
### 控制台 UI ### 控制台 UI
@@ -123,7 +121,8 @@ uv run main.py
### Web UI ### Web UI
本项目还包括一个 Web UI,提供更加动态和引人入胜的交互体验。 本项目还包括一个 Web UI,提供更加动态和引人入胜的交互体验。
> [! 注意]
> [!注意]
> 您需要先安装 Web UI 的依赖。 > 您需要先安装 Web UI 的依赖。
```bash ```bash
@@ -141,20 +140,21 @@ bootstrap.bat -d
## 支持的搜索引擎 ## 支持的搜索引擎
### 公域搜索引擎
DeerFlow 支持多种搜索引擎,可以在`.env`文件中通过`SEARCH_API`变量进行配置: DeerFlow 支持多种搜索引擎,可以在`.env`文件中通过`SEARCH_API`变量进行配置:
- **Tavily**(默认):专为 AI 应用设计的专业搜索 API - **Tavily**(默认):专为 AI 应用设计的专业搜索 API
- 需要在`.env`文件中设置`TAVILY_API_KEY` - 需要在`.env`文件中设置`TAVILY_API_KEY`
- 注册地址:<https://app.tavily.com/home> - 注册地址:https://app.tavily.com/home
- **DuckDuckGo**:注重隐私的搜索引擎 - **DuckDuckGo**:注重隐私的搜索引擎
- 无需 API 密钥 - 无需 API 密钥
- **Brave Search**:具有高级功能的注重隐私的搜索引擎 - **Brave Search**:具有高级功能的注重隐私的搜索引擎
- 需要在`.env`文件中设置`BRAVE_SEARCH_API_KEY` - 需要在`.env`文件中设置`BRAVE_SEARCH_API_KEY`
- 注册地址:<https://brave.com/search/api/> - 注册地址:https://brave.com/search/api/
- **Arxiv**:用于学术研究的科学论文搜索 - **Arxiv**:用于学术研究的科学论文搜索
- 无需 API 密钥 - 无需 API 密钥
@@ -167,30 +167,6 @@ DeerFlow 支持多种搜索引擎,可以在`.env`文件中通过`SEARCH_API`
SEARCH_API=tavily SEARCH_API=tavily
``` ```
### 私域知识库引擎
DeerFlow 支持基于私有域知识的检索,您可以将文档上传到多种私有知识库中,以便在研究过程中使用,当前支持的私域知识库有:
- **[RAGFlow](https://ragflow.io/docs/dev/)**:开源的基于检索增强生成的知识库引擎
```
# 参照示例进行配置 .env.example
RAG_PROVIDER=ragflow
RAGFLOW_API_URL="http://localhost:9388"
RAGFLOW_API_KEY="ragflow-xxx"
RAGFLOW_RETRIEVAL_SIZE=10
```
- **[VikingDB 知识库](https://www.volcengine.com/docs/84313/1254457)**:火山引擎提供的公有云知识库引擎
> 注意先从 [火山引擎](https://www.volcengine.com/docs/84313/1254485) 获取账号 AK/SK
```
# 参照示例进行配置 .env.example
RAG_PROVIDER=vikingdb_knowledge_base
VIKINGDB_KNOWLEDGE_BASE_API_URL="api-knowledgebase.mlp.cn-beijing.volces.com"
VIKINGDB_KNOWLEDGE_BASE_API_AK="volcengine-ak-xxx"
VIKINGDB_KNOWLEDGE_BASE_API_SK="volcengine-sk-xxx"
VIKINGDB_KNOWLEDGE_BASE_RETRIEVAL_SIZE=15
```
## 特性 ## 特性
### 核心能力 ### 核心能力
@@ -204,14 +180,10 @@ DeerFlow 支持基于私有域知识的检索,您可以将文档上传到多
### 工具和 MCP 集成 ### 工具和 MCP 集成
- 🔍 **搜索和检索** - 🔍 **搜索和检索**
- 通过 Tavily、Brave Search 等进行网络搜索 - 通过 Tavily、Brave Search 等进行网络搜索
- 使用 Jina 进行爬取 - 使用 Jina 进行爬取
- 高级内容提取 - 高级内容提取
- 支持检索指定私有知识库
- 📃 **RAG 集成**
- 支持 [RAGFlow](https://github.com/infiniflow/ragflow) 知识库
- 支持 [VikingDB](https://www.volcengine.com/docs/84313/1254457) 火山知识库
- 🔗 **MCP 无缝集成** - 🔗 **MCP 无缝集成**
- 扩展私有域访问、知识图谱、网页浏览等能力 - 扩展私有域访问、知识图谱、网页浏览等能力
@@ -220,6 +192,7 @@ DeerFlow 支持基于私有域知识的检索,您可以将文档上传到多
### 人机协作 ### 人机协作
- 🧠 **人在环中** - 🧠 **人在环中**
- 支持使用自然语言交互式修改研究计划 - 支持使用自然语言交互式修改研究计划
- 支持自动接受研究计划 - 支持自动接受研究计划
@@ -258,6 +231,7 @@ DeerFlow 实现了一个模块化的多智能体系统架构,专为自动化
- 管理研究流程并决定何时生成最终报告 - 管理研究流程并决定何时生成最终报告
3. **研究团队**:执行计划的专业智能体集合: 3. **研究团队**:执行计划的专业智能体集合:
- **研究员**:使用网络搜索引擎、爬虫甚至 MCP 服务等工具进行网络搜索和信息收集。 - **研究员**:使用网络搜索引擎、爬虫甚至 MCP 服务等工具进行网络搜索和信息收集。
- **编码员**:使用 Python REPL 工具处理代码分析、执行和技术任务。 - **编码员**:使用 Python REPL 工具处理代码分析、执行和技术任务。
每个智能体都可以访问针对其角色优化的特定工具,并在 LangGraph 框架内运行 每个智能体都可以访问针对其角色优化的特定工具,并在 LangGraph 框架内运行
@@ -267,27 +241,6 @@ DeerFlow 实现了一个模块化的多智能体系统架构,专为自动化
- 处理和组织收集的信息 - 处理和组织收集的信息
- 生成全面的研究报告 - 生成全面的研究报告
## 文本转语音集成
DeerFlow 现在包含一个文本转语音 (TTS) 功能,允许您将研究报告转换为语音。此功能使用火山引擎 TTS API 生成高质量的文本音频。速度、音量和音调等特性也可以自定义。
### 使用 TTS API
您可以通过`/api/tts`端点访问 TTS 功能:
```bash
# 使用curl的API调用示例
curl --location 'http://localhost:8000/api/tts' \
--header 'Content-Type: application/json' \
--data '{
"text": "这是文本转语音功能的测试。",
"speed_ratio": 1.0,
"volume_ratio": 1.0,
"pitch_ratio": 1.0
}' \
--output speech.mp3
```
## 开发 ## 开发
### 测试 ### 测试
@@ -346,9 +299,9 @@ langgraph dev
启动 LangGraph 服务器后,您将在终端中看到几个 URL: 启动 LangGraph 服务器后,您将在终端中看到几个 URL:
- API: <http://127.0.0.1:2024> - API: http://127.0.0.1:2024
- Studio UI: <https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024> - Studio UI: https://smith.langchain.com/studio/?baseUrl=http://127.0.0.1:2024
- API 文档<http://127.0.0.1:2024/docs> - API 文档: http://127.0.0.1:2024/docs
在浏览器中打开 Studio UI 链接以访问调试界面。 在浏览器中打开 Studio UI 链接以访问调试界面。
@@ -374,7 +327,6 @@ langgraph dev
DeerFlow 支持 LangSmith 追踪功能,帮助您调试和监控工作流。要启用 LangSmith 追踪: DeerFlow 支持 LangSmith 追踪功能,帮助您调试和监控工作流。要启用 LangSmith 追踪:
1. 确保您的 `.env` 文件中有以下配置(参见 `.env.example`): 1. 确保您的 `.env` 文件中有以下配置(参见 `.env.example`):
```bash ```bash
LANGSMITH_TRACING=true LANGSMITH_TRACING=true
LANGSMITH_ENDPOINT="https://api.smith.langchain.com" LANGSMITH_ENDPOINT="https://api.smith.langchain.com"
@@ -383,7 +335,6 @@ DeerFlow 支持 LangSmith 追踪功能,帮助您调试和监控工作流。要
``` ```
2. 通过运行以下命令本地启动 LangSmith 追踪: 2. 通过运行以下命令本地启动 LangSmith 追踪:
```bash ```bash
langgraph dev langgraph dev
``` ```
@@ -426,7 +377,7 @@ docker compose up
## 文本转语音集成 ## 文本转语音集成
DeerFlow 现在包含一个文本转语音 (TTS) 功能,允许您将研究报告转换为语音。此功能使用火山引擎 TTS API 生成高质量的文本音频。速度、音量和音调等特性也可以自定义。 DeerFlow 现在包含一个文本转语音(TTS)功能,允许您将研究报告转换为语音。此功能使用火山引擎 TTS API 生成高质量的文本音频。速度、音量和音调等特性也可以自定义。
### 使用 TTS API ### 使用 TTS API
@@ -452,14 +403,17 @@ curl --location 'http://localhost:8000/api/tts' \
### 研究报告 ### 研究报告
1. **OpenAI Sora 报告** - OpenAI 的 Sora AI 工具分析 1. **OpenAI Sora 报告** - OpenAI 的 Sora AI 工具分析
- 讨论功能、访问方式、提示工程、限制和伦理考虑 - 讨论功能、访问方式、提示工程、限制和伦理考虑
- [查看完整报告](examples/openai_sora_report.md) - [查看完整报告](examples/openai_sora_report.md)
2. **Google 的 Agent to Agent 协议报告** - Google 的 Agent to Agent (A2A) 协议概述 2. **Google 的 Agent to Agent 协议报告** - Google 的 Agent to Agent (A2A)协议概述
- 讨论其在 AI 智能体通信中的作用及其与 Anthropic 的 Model Context Protocol (MCP) 的关系
- 讨论其在 AI 智能体通信中的作用及其与 Anthropic 的 Model Context Protocol (MCP)的关系
- [查看完整报告](examples/what_is_agent_to_agent_protocol.md) - [查看完整报告](examples/what_is_agent_to_agent_protocol.md)
3. **什么是 MCP** - 对"MCP"一词在多个上下文中的全面分析 3. **什么是 MCP** - 对"MCP"一词在多个上下文中的全面分析
- 探讨 AI 中的 Model Context Protocol、化学中的 Monocalcium Phosphate 和电子学中的 Micro-channel Plate - 探讨 AI 中的 Model Context Protocol、化学中的 Monocalcium Phosphate 和电子学中的 Micro-channel Plate
- [查看完整报告](examples/what_is_mcp.md) - [查看完整报告](examples/what_is_mcp.md)
@@ -470,14 +424,17 @@ curl --location 'http://localhost:8000/api/tts' \
- [查看完整报告](examples/bitcoin_price_fluctuation.md) - [查看完整报告](examples/bitcoin_price_fluctuation.md)
5. **什么是 LLM?** - 对大型语言模型的深入探索 5. **什么是 LLM?** - 对大型语言模型的深入探索
- 讨论架构、训练、应用和伦理考虑 - 讨论架构、训练、应用和伦理考虑
- [查看完整报告](examples/what_is_llm.md) - [查看完整报告](examples/what_is_llm.md)
6. **如何使用 Claude 进行深度研究?** - 在深度研究中使用 Claude 的最佳实践和工作流程 6. **如何使用 Claude 进行深度研究?** - 在深度研究中使用 Claude 的最佳实践和工作流程
- 涵盖提示工程、数据分析和与其他工具的集成 - 涵盖提示工程、数据分析和与其他工具的集成
- [查看完整报告](examples/how_to_use_claude_deep_research.md) - [查看完整报告](examples/how_to_use_claude_deep_research.md)
7. **医疗保健中的 AI 采用:影响因素** - 影响医疗保健中 AI 采用的因素分析 7. **医疗保健中的 AI 采用:影响因素** - 影响医疗保健中 AI 采用的因素分析
- 讨论 AI 技术、数据质量、伦理考虑、经济评估、组织准备度和数字基础设施 - 讨论 AI 技术、数据质量、伦理考虑、经济评估、组织准备度和数字基础设施
- [查看完整报告](examples/AI_adoption_in_healthcare.md) - [查看完整报告](examples/AI_adoption_in_healthcare.md)
@@ -538,10 +495,10 @@ DeerFlow 包含一个人在环中机制,允许您在执行研究计划前审
- 系统将整合您的反馈并生成修订后的计划 - 系统将整合您的反馈并生成修订后的计划
3. **自动接受**:您可以启用自动接受以跳过审查过程: 3. **自动接受**:您可以启用自动接受以跳过审查过程:
- 通过 API:在请求中设置`auto_accepted_plan: true` - 通过 API:在请求中设置`auto_accepted_plan: true`
4. **API 集成**:使用 API 时,您可以通过`feedback`参数提供反馈: 4. **API 集成**:使用 API 时,您可以通过`feedback`参数提供反馈:
```json ```json
{ {
"messages": [{ "role": "user", "content": "什么是量子计算?" }], "messages": [{ "role": "user", "content": "什么是量子计算?" }],
+3 -34
View File
@@ -1,40 +1,9 @@
# [!NOTE] # [!NOTE]
# Read the `docs/configuration_guide.md` carefully, and update the # Read the `docs/configuration_guide.md` carefully, and update the configurations to match your specific settings and requirements.
# configurations to match your specific settings and requirements. # - Replace `api_key` with your own credentials
# - Replace `api_key` with your own credentials. # - Replace `base_url` and `model` name if you want to use a custom model
# - Replace `base_url` and `model` name if you want to use a custom model.
# - Set `verify_ssl` to `false` if your LLM server uses self-signed certificates
# - A restart is required every time you change the `config.yaml` file.
BASIC_MODEL: BASIC_MODEL:
base_url: https://ark.cn-beijing.volces.com/api/v3 base_url: https://ark.cn-beijing.volces.com/api/v3
model: "doubao-1-5-pro-32k-250115" model: "doubao-1-5-pro-32k-250115"
api_key: xxxx api_key: xxxx
# max_retries: 3 # Maximum number of retries for LLM calls
# verify_ssl: false # Uncomment this line to disable SSL certificate verification for self-signed certificates
# Reasoning model is optional.
# Uncomment the following settings if you want to use reasoning model
# for planning.
# REASONING_MODEL:
# base_url: https://ark.cn-beijing.volces.com/api/v3
# model: "doubao-1-5-thinking-pro-m-250428"
# api_key: xxxx
# max_retries: 3 # Maximum number of retries for LLM calls
# OTHER SETTINGS:
# Search engine configuration (Only supports Tavily currently)
# SEARCH_ENGINE:
# engine: tavily
# # Only include results from these domains
# include_domains:
# - example.com
# - trusted-news.com
# - reliable-source.org
# - gov.cn
# - edu.cn
# # Exclude results from these domains
# exclude_domains:
# - example.com
+1 -1
View File
@@ -9,7 +9,7 @@ services:
env_file: env_file:
- .env - .env
volumes: volumes:
- ./conf.yaml:/app/conf.yaml:ro - ./conf.yaml:/app/conf.yaml
restart: unless-stopped restart: unless-stopped
networks: networks:
- deer-flow-network - deer-flow-network
+10 -51
View File
@@ -11,7 +11,7 @@ cp conf.yaml.example conf.yaml
## Which models does DeerFlow support? ## Which models does DeerFlow support?
In DeerFlow, we currently only support non-reasoning models. This means models like OpenAI's o1/o3 or DeepSeek's R1 are not supported yet, but we plan to add support for them in the future. Additionally, all Gemma-3 models are currently unsupported due to the lack of tool usage capabilities. In DeerFlow, currently we only support non-reasoning models, which means models like OpenAI's o1/o3 or DeepSeek's R1 are not supported yet, but we will add support for them in the future.
### Supported Models ### Supported Models
@@ -49,7 +49,7 @@ BASIC_MODEL:
BASIC_MODEL: BASIC_MODEL:
base_url: "https://api.deepseek.com" base_url: "https://api.deepseek.com"
model: "deepseek-chat" model: "deepseek-chat"
api_key: YOUR_API_KEY api_key: YOU_API_KEY
# An example of Google Gemini models using OpenAI-Compatible interface # An example of Google Gemini models using OpenAI-Compatible interface
BASIC_MODEL: BASIC_MODEL:
@@ -58,31 +58,15 @@ BASIC_MODEL:
api_key: YOUR_API_KEY api_key: YOUR_API_KEY
``` ```
### How to use models with self-signed SSL certificates?
If your LLM server uses self-signed SSL certificates, you can disable SSL certificate verification by adding the `verify_ssl: false` parameter to your model configuration:
```yaml
BASIC_MODEL:
base_url: "https://your-llm-server.com/api/v1"
model: "your-model-name"
api_key: YOUR_API_KEY
verify_ssl: false # Disable SSL certificate verification for self-signed certificates
```
> [!WARNING]
> Disabling SSL certificate verification reduces security and should only be used in development environments or when you trust the LLM server. In production environments, it's recommended to use properly signed SSL certificates.
### How to use Ollama models? ### How to use Ollama models?
DeerFlow supports the integration of Ollama models. You can refer to [litellm Ollama](https://docs.litellm.ai/docs/providers/ollama). <br> DeerFlow supports the integration of Ollama models. You can refer to [litellm Ollama](https://docs.litellm.ai/docs/providers/ollama). <br>
The following is a configuration example of `conf.yaml` for using Ollama models(you might need to run the 'ollama serve' first): The following is a configuration example of `conf.yaml` for using Ollama models:
```yaml ```yaml
BASIC_MODEL: BASIC_MODEL:
model: "model-name" # Model name, which supports the completions API(important), such as: qwen3:8b, mistral-small3.1:24b, qwen2.5:3b model: "ollama/ollama-model-name"
base_url: "http://localhost:11434/v1" # Local service address of Ollama, which can be started/viewed via ollama serve base_url: "http://localhost:11434" # Local service address of Ollama, which can be started/viewed via ollama serve
api_key: "whatever" # Mandatory, fake api_key with a random string you like :-)
``` ```
### How to use OpenRouter models? ### How to use OpenRouter models?
@@ -105,38 +89,13 @@ BASIC_MODEL:
Note: The available models and their exact names may change over time. Please verify the currently available models and their correct identifiers in [OpenRouter's official documentation](https://openrouter.ai/docs). Note: The available models and their exact names may change over time. Please verify the currently available models and their correct identifiers in [OpenRouter's official documentation](https://openrouter.ai/docs).
### How to use Azure models?
### How to use Azure OpenAI chat models? DeerFlow supports the integration of Azure models. You can refer to [litellm Azure](https://docs.litellm.ai/docs/providers/azure). Configuration example of `conf.yaml`:
DeerFlow supports the integration of Azure OpenAI chat models. You can refer to [AzureChatOpenAI](https://python.langchain.com/api_reference/openai/chat_models/langchain_openai.chat_models.azure.AzureChatOpenAI.html). Configuration example of `conf.yaml`:
```yaml ```yaml
BASIC_MODEL: BASIC_MODEL:
model: "azure/gpt-4o-2024-08-06" model: "azure/gpt-4o-2024-08-06"
azure_endpoint: $AZURE_OPENAI_ENDPOINT api_base: $AZURE_API_BASE
api_version: $OPENAI_API_VERSION api_version: $AZURE_API_VERSION
api_key: $AZURE_OPENAI_API_KEY api_key: $AZURE_API_KEY
``` ```
## About Search Engine
### How to control search domains for Tavily?
DeerFlow allows you to control which domains are included or excluded in Tavily search results through the configuration file. This helps improve search result quality and reduce hallucinations by focusing on trusted sources.
`Tips`: it only supports Tavily currently.
You can configure domain filtering in your `conf.yaml` file as follows:
```yaml
SEARCH_ENGINE:
engine: tavily
# Only include results from these domains (whitelist)
include_domains:
- trusted-news.com
- gov.org
- reliable-source.edu
# Exclude results from these domains (blacklist)
exclude_domains:
- unreliable-site.com
- spam-domain.net
-4
View File
@@ -140,11 +140,7 @@ if __name__ == "__main__":
if args.query: if args.query:
user_query = " ".join(args.query) user_query = " ".join(args.query)
else: else:
# Loop until user provides non-empty input
while True:
user_query = input("Enter your query: ") user_query = input("Enter your query: ")
if user_query is not None and user_query != "":
break
# Run the agent workflow with the provided parameters # Run the agent workflow with the provided parameters
ask( ask(
-7
View File
@@ -32,25 +32,18 @@ dependencies = [
"arxiv>=2.2.0", "arxiv>=2.2.0",
"mcp>=1.6.0", "mcp>=1.6.0",
"langchain-mcp-adapters>=0.0.9", "langchain-mcp-adapters>=0.0.9",
"langchain-deepseek>=0.1.3",
"volcengine>=1.0.191",
] ]
[project.optional-dependencies] [project.optional-dependencies]
dev = [ dev = [
"ruff",
"black>=24.2.0", "black>=24.2.0",
"langgraph-cli[inmem]>=0.2.10", "langgraph-cli[inmem]>=0.2.10",
] ]
test = [ test = [
"pytest>=7.4.0", "pytest>=7.4.0",
"pytest-cov>=4.1.0", "pytest-cov>=4.1.0",
"pytest-asyncio>=1.0.0",
] ]
[tool.uv]
required-version = ">=0.6.15"
[tool.pytest.ini_options] [tool.pytest.ini_options]
testpaths = ["tests"] testpaths = ["tests"]
python_files = ["test_*.py"] python_files = ["test_*.py"]
+4 -18
View File
@@ -7,8 +7,7 @@ Server script for running the DeerFlow API.
import argparse import argparse
import logging import logging
import signal
import sys
import uvicorn import uvicorn
# Configure logging # Configure logging
@@ -19,17 +18,6 @@ logging.basicConfig(
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
def handle_shutdown(signum, frame):
"""Handle graceful shutdown on SIGTERM/SIGINT"""
logger.info("Received shutdown signal. Starting graceful shutdown...")
sys.exit(0)
# Register signal handlers
signal.signal(signal.SIGTERM, handle_shutdown)
signal.signal(signal.SIGINT, handle_shutdown)
if __name__ == "__main__": if __name__ == "__main__":
# Parse command line arguments # Parse command line arguments
parser = argparse.ArgumentParser(description="Run the DeerFlow API server") parser = argparse.ArgumentParser(description="Run the DeerFlow API server")
@@ -62,11 +50,12 @@ if __name__ == "__main__":
# Determine reload setting # Determine reload setting
reload = False reload = False
# Command line arguments override defaults
if args.reload: if args.reload:
reload = True reload = True
try: logger.info("Starting DeerFlow API server")
logger.info(f"Starting DeerFlow API server on {args.host}:{args.port}")
uvicorn.run( uvicorn.run(
"src.server:app", "src.server:app",
host=args.host, host=args.host,
@@ -74,6 +63,3 @@ if __name__ == "__main__":
reload=reload, reload=reload,
log_level=args.log_level, log_level=args.log_level,
) )
except Exception as e:
logger.error(f"Failed to start server: {str(e)}")
sys.exit(1)
+4 -5
View File
@@ -1,8 +1,8 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT # SPDX-License-Identifier: MIT
from .loader import load_yaml_config
from .tools import SELECTED_SEARCH_ENGINE, SearchEngine from .tools import SELECTED_SEARCH_ENGINE, SearchEngine
from .loader import load_yaml_config
from .questions import BUILT_IN_QUESTIONS, BUILT_IN_QUESTIONS_ZH_CN from .questions import BUILT_IN_QUESTIONS, BUILT_IN_QUESTIONS_ZH_CN
from dotenv import load_dotenv from dotenv import load_dotenv
@@ -11,7 +11,7 @@ from dotenv import load_dotenv
load_dotenv() load_dotenv()
# Team configuration # Team configuration
TEAM_MEMBER_CONFIGURATIONS = { TEAM_MEMBER_CONFIGRATIONS = {
"researcher": { "researcher": {
"name": "researcher", "name": "researcher",
"desc": ( "desc": (
@@ -36,15 +36,14 @@ TEAM_MEMBER_CONFIGURATIONS = {
}, },
} }
TEAM_MEMBERS = list(TEAM_MEMBER_CONFIGURATIONS.keys()) TEAM_MEMBERS = list(TEAM_MEMBER_CONFIGRATIONS.keys())
__all__ = [ __all__ = [
# Other configurations # Other configurations
"TEAM_MEMBERS", "TEAM_MEMBERS",
"TEAM_MEMBER_CONFIGURATIONS", "TEAM_MEMBER_CONFIGRATIONS",
"SELECTED_SEARCH_ENGINE", "SELECTED_SEARCH_ENGINE",
"SearchEngine", "SearchEngine",
"BUILT_IN_QUESTIONS", "BUILT_IN_QUESTIONS",
"BUILT_IN_QUESTIONS_ZH_CN", "BUILT_IN_QUESTIONS_ZH_CN",
load_yaml_config,
] ]
-1
View File
@@ -16,5 +16,4 @@ AGENT_LLM_MAP: dict[str, LLMType] = {
"podcast_script_writer": "basic", "podcast_script_writer": "basic",
"ppt_composer": "basic", "ppt_composer": "basic",
"prose_writer": "basic", "prose_writer": "basic",
"prompt_enhancer": "basic",
} }
+1 -9
View File
@@ -2,28 +2,20 @@
# SPDX-License-Identifier: MIT # SPDX-License-Identifier: MIT
import os import os
from dataclasses import dataclass, field, fields from dataclasses import dataclass, fields
from typing import Any, Optional from typing import Any, Optional
from langchain_core.runnables import RunnableConfig from langchain_core.runnables import RunnableConfig
from src.rag.retriever import Resource
from src.config.report_style import ReportStyle
@dataclass(kw_only=True) @dataclass(kw_only=True)
class Configuration: class Configuration:
"""The configurable fields.""" """The configurable fields."""
resources: list[Resource] = field(
default_factory=list
) # Resources to be used for the research
max_plan_iterations: int = 1 # Maximum number of plan iterations max_plan_iterations: int = 1 # Maximum number of plan iterations
max_step_num: int = 3 # Maximum number of steps in a plan max_step_num: int = 3 # Maximum number of steps in a plan
max_search_results: int = 3 # Maximum number of search results max_search_results: int = 3 # Maximum number of search results
mcp_settings: dict = None # MCP settings, including dynamic loaded tools mcp_settings: dict = None # MCP settings, including dynamic loaded tools
report_style: str = ReportStyle.ACADEMIC.value # Report style
enable_deep_thinking: bool = False # Whether to enable deep thinking
@classmethod @classmethod
def from_runnable_config( def from_runnable_config(
+1 -3
View File
@@ -12,14 +12,12 @@ def replace_env_vars(value: str) -> str:
return value return value
if value.startswith("$"): if value.startswith("$"):
env_var = value[1:] env_var = value[1:]
return os.getenv(env_var, env_var) return os.getenv(env_var, value)
return value return value
def process_dict(config: Dict[str, Any]) -> Dict[str, Any]: def process_dict(config: Dict[str, Any]) -> Dict[str, Any]:
"""Recursively process dictionary to replace environment variables.""" """Recursively process dictionary to replace environment variables."""
if not config:
return {}
result = {} result = {}
for key, value in config.items(): for key, value in config.items():
if isinstance(value, dict): if isinstance(value, dict):
-8
View File
@@ -1,8 +0,0 @@
import enum
class ReportStyle(enum.Enum):
ACADEMIC = "academic"
POPULAR_SCIENCE = "popular_science"
NEWS = "news"
SOCIAL_MEDIA = "social_media"
-8
View File
@@ -17,11 +17,3 @@ class SearchEngine(enum.Enum):
# Tool configuration # Tool configuration
SELECTED_SEARCH_ENGINE = os.getenv("SEARCH_API", SearchEngine.TAVILY.value) SELECTED_SEARCH_ENGINE = os.getenv("SEARCH_API", SearchEngine.TAVILY.value)
class RAGProvider(enum.Enum):
RAGFLOW = "ragflow"
VIKINGDB_KNOWLEDGE_BASE = "vikingdb_knowledge_base"
SELECTED_RAG_PROVIDER = os.getenv("RAG_PROVIDER")
+4 -3
View File
@@ -3,7 +3,8 @@
from .article import Article from .article import Article
from .crawler import Crawler from .crawler import Crawler
from .jina_client import JinaClient
from .readability_extractor import ReadabilityExtractor
__all__ = ["Article", "Crawler", "JinaClient", "ReadabilityExtractor"] __all__ = [
"Article",
"Crawler",
]
+11
View File
@@ -1,6 +1,7 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT # SPDX-License-Identifier: MIT
import sys
from .article import Article from .article import Article
from .jina_client import JinaClient from .jina_client import JinaClient
@@ -25,3 +26,13 @@ class Crawler:
article = extractor.extract_article(html) article = extractor.extract_article(html)
article.url = url article.url = url
return article return article
if __name__ == "__main__":
if len(sys.argv) == 2:
url = sys.argv[1]
else:
url = "https://fintel.io/zh-hant/s/br/nvdc34"
crawler = Crawler()
article = crawler.crawl(url)
print(article.to_markdown())
-32
View File
@@ -3,7 +3,6 @@
from langgraph.graph import StateGraph, START, END from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver from langgraph.checkpoint.memory import MemorySaver
from src.prompts.planner_model import StepType
from .types import State from .types import State
from .nodes import ( from .nodes import (
@@ -18,31 +17,6 @@ from .nodes import (
) )
def continue_to_running_research_team(state: State):
current_plan = state.get("current_plan")
if not current_plan or not current_plan.steps:
return "planner"
if all(step.execution_res for step in current_plan.steps):
return "planner"
# Find first incomplete step
incomplete_step = None
for step in current_plan.steps:
if not step.execution_res:
incomplete_step = step
break
if not incomplete_step:
return "planner"
if incomplete_step.step_type == StepType.RESEARCH:
return "researcher"
if incomplete_step.step_type == StepType.PROCESSING:
return "coder"
return "planner"
def _build_base_graph(): def _build_base_graph():
"""Build and return the base state graph with all nodes and edges.""" """Build and return the base state graph with all nodes and edges."""
builder = StateGraph(State) builder = StateGraph(State)
@@ -55,12 +29,6 @@ def _build_base_graph():
builder.add_node("researcher", researcher_node) builder.add_node("researcher", researcher_node)
builder.add_node("coder", coder_node) builder.add_node("coder", coder_node)
builder.add_node("human_feedback", human_feedback_node) builder.add_node("human_feedback", human_feedback_node)
builder.add_edge("background_investigator", "planner")
builder.add_conditional_edges(
"research_team",
continue_to_running_research_team,
["planner", "researcher", "coder"],
)
builder.add_edge("reporter", END) builder.add_edge("reporter", END)
return builder return builder
+53 -69
View File
@@ -17,14 +17,13 @@ from src.tools.search import LoggedTavilySearch
from src.tools import ( from src.tools import (
crawl_tool, crawl_tool,
get_web_search_tool, get_web_search_tool,
get_retriever_tool,
python_repl_tool, python_repl_tool,
) )
from src.config.agents import AGENT_LLM_MAP from src.config.agents import AGENT_LLM_MAP
from src.config.configuration import Configuration from src.config.configuration import Configuration
from src.llms.llm import get_llm_by_type from src.llms.llm import get_llm_by_type
from src.prompts.planner_model import Plan from src.prompts.planner_model import Plan, StepType
from src.prompts.template import apply_prompt_template from src.prompts.template import apply_prompt_template
from src.utils.json_utils import repair_json_output from src.utils.json_utils import repair_json_output
@@ -36,7 +35,7 @@ logger = logging.getLogger(__name__)
@tool @tool
def handoff_to_planner( def handoff_to_planner(
research_topic: Annotated[str, "The topic of the research task to be handed off."], task_title: Annotated[str, "The title of the task to be handed off."],
locale: Annotated[str, "The user's detected language locale (e.g., en-US, zh-CN)."], locale: Annotated[str, "The user's detected language locale (e.g., en-US, zh-CN)."],
): ):
"""Handoff to planner agent to do plan.""" """Handoff to planner agent to do plan."""
@@ -45,24 +44,22 @@ def handoff_to_planner(
return return
def background_investigation_node(state: State, config: RunnableConfig): def background_investigation_node(
state: State, config: RunnableConfig
) -> Command[Literal["planner"]]:
logger.info("background investigation node is running.") logger.info("background investigation node is running.")
configurable = Configuration.from_runnable_config(config) configurable = Configuration.from_runnable_config(config)
query = state.get("research_topic") query = state["messages"][-1].content
background_investigation_results = None if SELECTED_SEARCH_ENGINE == SearchEngine.TAVILY:
if SELECTED_SEARCH_ENGINE == SearchEngine.TAVILY.value:
searched_content = LoggedTavilySearch( searched_content = LoggedTavilySearch(
max_results=configurable.max_search_results max_results=configurable.max_search_results
).invoke(query) ).invoke({"query": query})
background_investigation_results = None
if isinstance(searched_content, list): if isinstance(searched_content, list):
background_investigation_results = [ background_investigation_results = [
f"## {elem['title']}\n\n{elem['content']}" for elem in searched_content {"title": elem["title"], "content": elem["content"]}
for elem in searched_content
] ]
return {
"background_investigation_results": "\n\n".join(
background_investigation_results
)
}
else: else:
logger.error( logger.error(
f"Tavily search returned malformed response: {searched_content}" f"Tavily search returned malformed response: {searched_content}"
@@ -71,11 +68,14 @@ def background_investigation_node(state: State, config: RunnableConfig):
background_investigation_results = get_web_search_tool( background_investigation_results = get_web_search_tool(
configurable.max_search_results configurable.max_search_results
).invoke(query) ).invoke(query)
return { return Command(
update={
"background_investigation_results": json.dumps( "background_investigation_results": json.dumps(
background_investigation_results, ensure_ascii=False background_investigation_results, ensure_ascii=False
) )
} },
goto="planner",
)
def planner_node( def planner_node(
@@ -87,8 +87,10 @@ def planner_node(
plan_iterations = state["plan_iterations"] if state.get("plan_iterations", 0) else 0 plan_iterations = state["plan_iterations"] if state.get("plan_iterations", 0) else 0
messages = apply_prompt_template("planner", state, configurable) messages = apply_prompt_template("planner", state, configurable)
if state.get("enable_background_investigation") and state.get( if (
"background_investigation_results" plan_iterations == 0
and state.get("enable_background_investigation")
and state.get("background_investigation_results")
): ):
messages += [ messages += [
{ {
@@ -101,10 +103,8 @@ def planner_node(
} }
] ]
if configurable.enable_deep_thinking: if AGENT_LLM_MAP["planner"] == "basic":
llm = get_llm_by_type("reasoning") llm = get_llm_by_type(AGENT_LLM_MAP["planner"]).with_structured_output(
elif AGENT_LLM_MAP["planner"] == "basic":
llm = get_llm_by_type("basic").with_structured_output(
Plan, Plan,
method="json_mode", method="json_mode",
) )
@@ -116,7 +116,7 @@ def planner_node(
return Command(goto="reporter") return Command(goto="reporter")
full_response = "" full_response = ""
if AGENT_LLM_MAP["planner"] == "basic" and not configurable.enable_deep_thinking: if AGENT_LLM_MAP["planner"] == "basic":
response = llm.invoke(messages) response = llm.invoke(messages)
full_response = response.model_dump_json(indent=4, exclude_none=True) full_response = response.model_dump_json(indent=4, exclude_none=True)
else: else:
@@ -134,7 +134,7 @@ def planner_node(
return Command(goto="reporter") return Command(goto="reporter")
else: else:
return Command(goto="__end__") return Command(goto="__end__")
if isinstance(curr_plan, dict) and curr_plan.get("has_enough_context"): if curr_plan.get("has_enough_context"):
logger.info("Planner response has enough context.") logger.info("Planner response has enough context.")
new_plan = Plan.model_validate(curr_plan) new_plan = Plan.model_validate(curr_plan)
return Command( return Command(
@@ -186,9 +186,11 @@ def human_feedback_node(
plan_iterations += 1 plan_iterations += 1
# parse the plan # parse the plan
new_plan = json.loads(current_plan) new_plan = json.loads(current_plan)
if new_plan["has_enough_context"]:
goto = "reporter"
except json.JSONDecodeError: except json.JSONDecodeError:
logger.warning("Planner response is not a valid JSON") logger.warning("Planner response is not a valid JSON")
if plan_iterations > 1: # the plan_iterations is increased before this check if plan_iterations > 0:
return Command(goto="reporter") return Command(goto="reporter")
else: else:
return Command(goto="__end__") return Command(goto="__end__")
@@ -204,11 +206,10 @@ def human_feedback_node(
def coordinator_node( def coordinator_node(
state: State, config: RunnableConfig state: State,
) -> Command[Literal["planner", "background_investigator", "__end__"]]: ) -> Command[Literal["planner", "background_investigator", "__end__"]]:
"""Coordinator node that communicate with customers.""" """Coordinator node that communicate with customers."""
logger.info("Coordinator talking.") logger.info("Coordinator talking.")
configurable = Configuration.from_runnable_config(config)
messages = apply_prompt_template("coordinator", state) messages = apply_prompt_template("coordinator", state)
response = ( response = (
get_llm_by_type(AGENT_LLM_MAP["coordinator"]) get_llm_by_type(AGENT_LLM_MAP["coordinator"])
@@ -219,7 +220,6 @@ def coordinator_node(
goto = "__end__" goto = "__end__"
locale = state.get("locale", "en-US") # Default locale if not specified locale = state.get("locale", "en-US") # Default locale if not specified
research_topic = state.get("research_topic", "")
if len(response.tool_calls) > 0: if len(response.tool_calls) > 0:
goto = "planner" goto = "planner"
@@ -230,11 +230,8 @@ def coordinator_node(
for tool_call in response.tool_calls: for tool_call in response.tool_calls:
if tool_call.get("name", "") != "handoff_to_planner": if tool_call.get("name", "") != "handoff_to_planner":
continue continue
if tool_call.get("args", {}).get("locale") and tool_call.get( if tool_locale := tool_call.get("args", {}).get("locale"):
"args", {} locale = tool_locale
).get("research_topic"):
locale = tool_call.get("args", {}).get("locale")
research_topic = tool_call.get("args", {}).get("research_topic")
break break
except Exception as e: except Exception as e:
logger.error(f"Error processing tool calls: {e}") logger.error(f"Error processing tool calls: {e}")
@@ -243,24 +240,16 @@ def coordinator_node(
"Coordinator response contains no tool calls. Terminating workflow execution." "Coordinator response contains no tool calls. Terminating workflow execution."
) )
logger.debug(f"Coordinator response: {response}") logger.debug(f"Coordinator response: {response}")
messages = state.get("messages", [])
if response.content:
messages.append(HumanMessage(content=response.content, name="coordinator"))
return Command( return Command(
update={ update={"locale": locale},
"messages": messages,
"locale": locale,
"research_topic": research_topic,
"resources": configurable.resources,
},
goto=goto, goto=goto,
) )
def reporter_node(state: State, config: RunnableConfig): def reporter_node(state: State):
"""Reporter node that write a final report.""" """Reporter node that write a final report."""
logger.info("Reporter write final report") logger.info("Reporter write final report")
configurable = Configuration.from_runnable_config(config)
current_plan = state.get("current_plan") current_plan = state.get("current_plan")
input_ = { input_ = {
"messages": [ "messages": [
@@ -270,7 +259,7 @@ def reporter_node(state: State, config: RunnableConfig):
], ],
"locale": state.get("locale", "en-US"), "locale": state.get("locale", "en-US"),
} }
invoke_messages = apply_prompt_template("reporter", input_, configurable) invoke_messages = apply_prompt_template("reporter", input_)
observations = state.get("observations", []) observations = state.get("observations", [])
# Add a reminder about the new report format, citation style, and table usage # Add a reminder about the new report format, citation style, and table usage
@@ -296,10 +285,24 @@ def reporter_node(state: State, config: RunnableConfig):
return {"final_report": response_content} return {"final_report": response_content}
def research_team_node(state: State): def research_team_node(
state: State,
) -> Command[Literal["planner", "researcher", "coder"]]:
"""Research team node that collaborates on tasks.""" """Research team node that collaborates on tasks."""
logger.info("Research team is collaborating on tasks.") logger.info("Research team is collaborating on tasks.")
pass current_plan = state.get("current_plan")
if not current_plan or not current_plan.steps:
return Command(goto="planner")
if all(step.execution_res for step in current_plan.steps):
return Command(goto="planner")
for step in current_plan.steps:
if not step.execution_res:
break
if step.step_type and step.step_type == StepType.RESEARCH:
return Command(goto="researcher")
if step.step_type and step.step_type == StepType.PROCESSING:
return Command(goto="coder")
return Command(goto="planner")
async def _execute_agent_step( async def _execute_agent_step(
@@ -323,14 +326,14 @@ async def _execute_agent_step(
logger.warning("No unexecuted step found") logger.warning("No unexecuted step found")
return Command(goto="research_team") return Command(goto="research_team")
logger.info(f"Executing step: {current_step.title}, agent: {agent_name}") logger.info(f"Executing step: {current_step.title}")
# Format completed steps information # Format completed steps information
completed_steps_info = "" completed_steps_info = ""
if completed_steps: if completed_steps:
completed_steps_info = "# Existing Research Findings\n\n" completed_steps_info = "# Existing Research Findings\n\n"
for i, step in enumerate(completed_steps): for i, step in enumerate(completed_steps):
completed_steps_info += f"## Existing Finding {i + 1}: {step.title}\n\n" completed_steps_info += f"## Existing Finding {i+1}: {step.title}\n\n"
completed_steps_info += f"<finding>\n{step.execution_res}\n</finding>\n\n" completed_steps_info += f"<finding>\n{step.execution_res}\n</finding>\n\n"
# Prepare the input for the agent with completed steps info # Prepare the input for the agent with completed steps info
@@ -344,19 +347,6 @@ async def _execute_agent_step(
# Add citation reminder for researcher agent # Add citation reminder for researcher agent
if agent_name == "researcher": if agent_name == "researcher":
if state.get("resources"):
resources_info = "**The user mentioned the following resource files:**\n\n"
for resource in state.get("resources"):
resources_info += f"- {resource.title} ({resource.description})\n"
agent_input["messages"].append(
HumanMessage(
content=resources_info
+ "\n\n"
+ "You MUST use the **local_search_tool** to retrieve the information from the resource files.",
)
)
agent_input["messages"].append( agent_input["messages"].append(
HumanMessage( HumanMessage(
content="IMPORTANT: DO NOT include inline citations in the text. Instead, track all sources and include a References section at the end using link reference format. Include an empty line between each citation for better readability. Use this format for each reference:\n- [Source Title](URL)\n\n- [Another Source](URL)", content="IMPORTANT: DO NOT include inline citations in the text. Instead, track all sources and include a References section at the end using link reference format. Include an empty line between each citation for better readability. Use this format for each reference:\n- [Source Title](URL)\n\n- [Another Source](URL)",
@@ -387,7 +377,6 @@ async def _execute_agent_step(
) )
recursion_limit = default_recursion_limit recursion_limit = default_recursion_limit
logger.info(f"Agent input: {agent_input}")
result = await agent.ainvoke( result = await agent.ainvoke(
input=agent_input, config={"recursion_limit": recursion_limit} input=agent_input, config={"recursion_limit": recursion_limit}
) )
@@ -479,16 +468,11 @@ async def researcher_node(
"""Researcher node that do research""" """Researcher node that do research"""
logger.info("Researcher node is researching.") logger.info("Researcher node is researching.")
configurable = Configuration.from_runnable_config(config) configurable = Configuration.from_runnable_config(config)
tools = [get_web_search_tool(configurable.max_search_results), crawl_tool]
retriever_tool = get_retriever_tool(state.get("resources", []))
if retriever_tool:
tools.insert(0, retriever_tool)
logger.info(f"Researcher tools: {tools}")
return await _setup_and_execute_agent_step( return await _setup_and_execute_agent_step(
state, state,
config, config,
"researcher", "researcher",
tools, [get_web_search_tool(configurable.max_search_results), crawl_tool],
) )
+2 -3
View File
@@ -1,11 +1,12 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT # SPDX-License-Identifier: MIT
import operator
from typing import Annotated
from langgraph.graph import MessagesState from langgraph.graph import MessagesState
from src.prompts.planner_model import Plan from src.prompts.planner_model import Plan
from src.rag import Resource
class State(MessagesState): class State(MessagesState):
@@ -13,9 +14,7 @@ class State(MessagesState):
# Runtime Variables # Runtime Variables
locale: str = "en-US" locale: str = "en-US"
research_topic: str = ""
observations: list[str] = [] observations: list[str] = []
resources: list[Resource] = []
plan_iterations: int = 0 plan_iterations: int = 0
current_plan: Plan | str = None current_plan: Plan | str = None
final_report: str = "" final_report: str = ""
+21 -116
View File
@@ -3,148 +3,53 @@
from pathlib import Path from pathlib import Path
from typing import Any, Dict from typing import Any, Dict
import os
import httpx
from langchain_core.language_models import BaseChatModel from langchain_openai import ChatOpenAI
from langchain_openai import ChatOpenAI, AzureChatOpenAI
from langchain_deepseek import ChatDeepSeek
from typing import get_args
from src.config import load_yaml_config from src.config import load_yaml_config
from src.config.agents import LLMType from src.config.agents import LLMType
# Cache for LLM instances # Cache for LLM instances
_llm_cache: dict[LLMType, BaseChatModel] = {} _llm_cache: dict[LLMType, ChatOpenAI] = {}
def _get_config_file_path() -> str: def _create_llm_use_conf(llm_type: LLMType, conf: Dict[str, Any]) -> ChatOpenAI:
"""Get the path to the configuration file.""" llm_type_map = {
return str((Path(__file__).parent.parent.parent / "conf.yaml").resolve()) "reasoning": conf.get("REASONING_MODEL"),
"basic": conf.get("BASIC_MODEL"),
"vision": conf.get("VISION_MODEL"),
def _get_llm_type_config_keys() -> dict[str, str]:
"""Get mapping of LLM types to their configuration keys."""
return {
"reasoning": "REASONING_MODEL",
"basic": "BASIC_MODEL",
"vision": "VISION_MODEL",
} }
llm_conf = llm_type_map.get(llm_type)
if not llm_conf:
def _get_env_llm_conf(llm_type: str) -> Dict[str, Any]:
"""
Get LLM configuration from environment variables.
Environment variables should follow the format: {LLM_TYPE}__{KEY}
e.g., BASIC_MODEL__api_key, BASIC_MODEL__base_url
"""
prefix = f"{llm_type.upper()}_MODEL__"
conf = {}
for key, value in os.environ.items():
if key.startswith(prefix):
conf_key = key[len(prefix) :].lower()
conf[conf_key] = value
return conf
def _create_llm_use_conf(llm_type: LLMType, conf: Dict[str, Any]) -> BaseChatModel:
"""Create LLM instance using configuration."""
llm_type_config_keys = _get_llm_type_config_keys()
config_key = llm_type_config_keys.get(llm_type)
if not config_key:
raise ValueError(f"Unknown LLM type: {llm_type}") raise ValueError(f"Unknown LLM type: {llm_type}")
llm_conf = conf.get(config_key, {})
if not isinstance(llm_conf, dict): if not isinstance(llm_conf, dict):
raise ValueError(f"Invalid LLM configuration for {llm_type}: {llm_conf}") raise ValueError(f"Invalid LLM Conf: {llm_type}")
return ChatOpenAI(**llm_conf)
# Get configuration from environment variables
env_conf = _get_env_llm_conf(llm_type)
# Merge configurations, with environment variables taking precedence
merged_conf = {**llm_conf, **env_conf}
if not merged_conf:
raise ValueError(f"No configuration found for LLM type: {llm_type}")
# Add max_retries to handle rate limit errors
if "max_retries" not in merged_conf:
merged_conf["max_retries"] = 3
if llm_type == "reasoning":
merged_conf["api_base"] = merged_conf.pop("base_url", None)
# Handle SSL verification settings
verify_ssl = merged_conf.pop("verify_ssl", True)
# Create custom HTTP client if SSL verification is disabled
if not verify_ssl:
http_client = httpx.Client(verify=False)
http_async_client = httpx.AsyncClient(verify=False)
merged_conf["http_client"] = http_client
merged_conf["http_async_client"] = http_async_client
if "azure_endpoint" in merged_conf or os.getenv("AZURE_OPENAI_ENDPOINT"):
return AzureChatOpenAI(**merged_conf)
if llm_type == "reasoning":
return ChatDeepSeek(**merged_conf)
else:
return ChatOpenAI(**merged_conf)
def get_llm_by_type( def get_llm_by_type(
llm_type: LLMType, llm_type: LLMType,
) -> BaseChatModel: ) -> ChatOpenAI:
""" """
Get LLM instance by type. Returns cached instance if available. Get LLM instance by type. Returns cached instance if available.
""" """
if llm_type in _llm_cache: if llm_type in _llm_cache:
return _llm_cache[llm_type] return _llm_cache[llm_type]
conf = load_yaml_config(_get_config_file_path()) conf = load_yaml_config(
str((Path(__file__).parent.parent.parent / "conf.yaml").resolve())
)
llm = _create_llm_use_conf(llm_type, conf) llm = _create_llm_use_conf(llm_type, conf)
_llm_cache[llm_type] = llm _llm_cache[llm_type] = llm
return llm return llm
def get_configured_llm_models() -> dict[str, list[str]]:
"""
Get all configured LLM models grouped by type.
Returns:
Dictionary mapping LLM type to list of configured model names.
"""
try:
conf = load_yaml_config(_get_config_file_path())
llm_type_config_keys = _get_llm_type_config_keys()
configured_models: dict[str, list[str]] = {}
for llm_type in get_args(LLMType):
# Get configuration from YAML file
config_key = llm_type_config_keys.get(llm_type, "")
yaml_conf = conf.get(config_key, {}) if config_key else {}
# Get configuration from environment variables
env_conf = _get_env_llm_conf(llm_type)
# Merge configurations, with environment variables taking precedence
merged_conf = {**yaml_conf, **env_conf}
# Check if model is configured
model_name = merged_conf.get("model")
if model_name:
configured_models.setdefault(llm_type, []).append(model_name)
return configured_models
except Exception as e:
# Log error and return empty dict to avoid breaking the application
print(f"Warning: Failed to load LLM configuration: {e}")
return {}
# In the future, we will use reasoning_llm and vl_llm for different purposes # In the future, we will use reasoning_llm and vl_llm for different purposes
# reasoning_llm = get_llm_by_type("reasoning") # reasoning_llm = get_llm_by_type("reasoning")
# vl_llm = get_llm_by_type("vision") # vl_llm = get_llm_by_type("vision")
if __name__ == "__main__":
# Initialize LLMs for different purposes - now these will be cached
basic_llm = get_llm_by_type("basic")
print(basic_llm.invoke("Hello"))
+1
View File
@@ -1,6 +1,7 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT # SPDX-License-Identifier: MIT
from typing import Optional
from langgraph.graph import MessagesState from langgraph.graph import MessagesState
-4
View File
@@ -1,4 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
"""Prompt enhancer module for improving user prompts."""
-25
View File
@@ -1,25 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
from langgraph.graph import StateGraph
from src.prompt_enhancer.graph.enhancer_node import prompt_enhancer_node
from src.prompt_enhancer.graph.state import PromptEnhancerState
def build_graph():
"""Build and return the prompt enhancer workflow graph."""
# Build state graph
builder = StateGraph(PromptEnhancerState)
# Add the enhancer node
builder.add_node("enhancer", prompt_enhancer_node)
# Set entry point
builder.set_entry_point("enhancer")
# Set finish point
builder.set_finish_point("enhancer")
# Compile and return the graph
return builder.compile()
@@ -1,83 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import logging
import re
from langchain.schema import HumanMessage
from src.config.agents import AGENT_LLM_MAP
from src.llms.llm import get_llm_by_type
from src.prompts.template import apply_prompt_template
from src.prompt_enhancer.graph.state import PromptEnhancerState
logger = logging.getLogger(__name__)
def prompt_enhancer_node(state: PromptEnhancerState):
"""Node that enhances user prompts using AI analysis."""
logger.info("Enhancing user prompt...")
model = get_llm_by_type(AGENT_LLM_MAP["prompt_enhancer"])
try:
# Create messages with context if provided
context_info = ""
if state.get("context"):
context_info = f"\n\nAdditional context: {state['context']}"
original_prompt_message = HumanMessage(
content=f"Please enhance this prompt:{context_info}\n\nOriginal prompt: {state['prompt']}"
)
messages = apply_prompt_template(
"prompt_enhancer/prompt_enhancer",
{
"messages": [original_prompt_message],
"report_style": state.get("report_style"),
},
)
# Get the response from the model
response = model.invoke(messages)
# Extract content from response
response_content = response.content.strip()
logger.debug(f"Response content: {response_content}")
# Try to extract content from XML tags first
xml_match = re.search(
r"<enhanced_prompt>(.*?)</enhanced_prompt>", response_content, re.DOTALL
)
if xml_match:
# Extract content from XML tags and clean it up
enhanced_prompt = xml_match.group(1).strip()
logger.debug("Successfully extracted enhanced prompt from XML tags")
else:
# Fallback to original logic if no XML tags found
enhanced_prompt = response_content
logger.warning("No XML tags found in response, using fallback parsing")
# Remove common prefixes that might be added by the model
prefixes_to_remove = [
"Enhanced Prompt:",
"Enhanced prompt:",
"Here's the enhanced prompt:",
"Here is the enhanced prompt:",
"**Enhanced Prompt**:",
"**Enhanced prompt**:",
]
for prefix in prefixes_to_remove:
if enhanced_prompt.startswith(prefix):
enhanced_prompt = enhanced_prompt[len(prefix) :].strip()
break
logger.info("Prompt enhancement completed successfully")
logger.debug(f"Enhanced prompt: {enhanced_prompt}")
return {"output": enhanced_prompt}
except Exception as e:
logger.error(f"Error in prompt enhancement: {str(e)}")
return {"output": state["prompt"]}
-14
View File
@@ -1,14 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
from typing import TypedDict, Optional
from src.config.report_style import ReportStyle
class PromptEnhancerState(TypedDict):
"""State for the prompt enhancer workflow."""
prompt: str # Original prompt to enhance
context: Optional[str] # Additional context
report_style: Optional[ReportStyle] # Report style preference
output: Optional[str] # Enhanced prompt result
+9 -10
View File
@@ -57,15 +57,14 @@ Before creating a detailed plan, assess if there is sufficient context to answer
Different types of steps have different web search requirements: Different types of steps have different web search requirements:
1. **Research Steps** (`need_search: true`): 1. **Research Steps** (`need_web_search: true`):
- Retrieve information from the file with the URL with `rag://` or `http://` prefix specified by the user
- Gathering market data or industry trends - Gathering market data or industry trends
- Finding historical information - Finding historical information
- Collecting competitor analysis - Collecting competitor analysis
- Researching current events or news - Researching current events or news
- Finding statistical data or reports - Finding statistical data or reports
2. **Data Processing Steps** (`need_search: false`): 2. **Data Processing Steps** (`need_web_search: false`):
- API calls and data extraction - API calls and data extraction
- Database queries - Database queries
- Raw data collection from existing sources - Raw data collection from existing sources
@@ -144,8 +143,8 @@ When planning information gathering, consider these key aspects and ensure COMPR
- Ensure each step is substantial and covers related information categories - Ensure each step is substantial and covers related information categories
- Prioritize breadth and depth within the {{ max_step_num }}-step constraint - Prioritize breadth and depth within the {{ max_step_num }}-step constraint
- For each step, carefully assess if web search is needed: - For each step, carefully assess if web search is needed:
- Research and external data gathering: Set `need_search: true` - Research and external data gathering: Set `need_web_search: true`
- Internal data processing: Set `need_search: false` - Internal data processing: Set `need_web_search: false`
- Specify the exact data to be collected in step's `description`. Include a `note` if necessary. - Specify the exact data to be collected in step's `description`. Include a `note` if necessary.
- Prioritize depth and volume of relevant information - limited information is not acceptable. - Prioritize depth and volume of relevant information - limited information is not acceptable.
- Use the same language as the user to generate the plan. - Use the same language as the user to generate the plan.
@@ -157,9 +156,9 @@ Directly output the raw JSON format of `Plan` without "```json". The `Plan` inte
```ts ```ts
interface Step { interface Step {
need_search: boolean; // Must be explicitly set for each step need_web_search: boolean; // Must be explicitly set for each step
title: string; title: string;
description: string; // Specify exactly what data to collect. If the user input contains a link, please retain the full Markdown format when necessary. description: string; // Specify exactly what data to collect
step_type: "research" | "processing"; // Indicates the nature of the step step_type: "research" | "processing"; // Indicates the nature of the step
} }
@@ -180,8 +179,8 @@ interface Plan {
- Prioritize BOTH breadth (covering essential aspects) AND depth (detailed information on each aspect) - Prioritize BOTH breadth (covering essential aspects) AND depth (detailed information on each aspect)
- Never settle for minimal information - the goal is a comprehensive, detailed final report - Never settle for minimal information - the goal is a comprehensive, detailed final report
- Limited or insufficient information will lead to an inadequate final report - Limited or insufficient information will lead to an inadequate final report
- Carefully assess each step's web search or retrieve from URL requirement based on its nature: - Carefully assess each step's web search requirement based on its nature:
- Research steps (`need_search: true`) for gathering information - Research steps (`need_web_search: true`) for gathering information
- Processing steps (`need_search: false`) for calculations and data processing - Processing steps (`need_web_search: false`) for calculations and data processing
- Default to gathering more information unless the strictest sufficient context criteria are met - Default to gathering more information unless the strictest sufficient context criteria are met
- Always use the language specified by the locale = **{{ locale }}**. - Always use the language specified by the locale = **{{ locale }}**.
+4 -2
View File
@@ -13,7 +13,9 @@ class StepType(str, Enum):
class Step(BaseModel): class Step(BaseModel):
need_search: bool = Field(..., description="Must be explicitly set for each step") need_web_search: bool = Field(
..., description="Must be explicitly set for each step"
)
title: str title: str
description: str = Field(..., description="Specify exactly what data to collect") description: str = Field(..., description="Specify exactly what data to collect")
step_type: StepType = Field(..., description="Indicates the nature of the step") step_type: StepType = Field(..., description="Indicates the nature of the step")
@@ -45,7 +47,7 @@ class Plan(BaseModel):
"title": "AI Market Research Plan", "title": "AI Market Research Plan",
"steps": [ "steps": [
{ {
"need_search": True, "need_web_search": True,
"title": "Current AI Market Analysis", "title": "Current AI Market Analysis",
"description": ( "description": (
"Collect data on market size, growth rates, major players, and investment trends in AI sector." "Collect data on market size, growth rates, major players, and investment trends in AI sector."
@@ -1,135 +0,0 @@
---
CURRENT_TIME: {{ CURRENT_TIME }}
---
You are an expert prompt engineer. Your task is to enhance user prompts to make them more effective, specific, and likely to produce high-quality results from AI systems.
# Your Role
- Analyze the original prompt for clarity, specificity, and completeness
- Enhance the prompt by adding relevant details, context, and structure
- Make the prompt more actionable and results-oriented
- Preserve the user's original intent while improving effectiveness
{% if report_style == "academic" %}
# Enhancement Guidelines for Academic Style
1. **Add methodological rigor**: Include research methodology, scope, and analytical framework
2. **Specify academic structure**: Organize with clear thesis, literature review, analysis, and conclusions
3. **Clarify scholarly expectations**: Specify citation requirements, evidence standards, and academic tone
4. **Add theoretical context**: Include relevant theoretical frameworks and disciplinary perspectives
5. **Ensure precision**: Use precise terminology and avoid ambiguous language
6. **Include limitations**: Acknowledge scope limitations and potential biases
{% elif report_style == "popular_science" %}
# Enhancement Guidelines for Popular Science Style
1. **Add accessibility**: Transform technical concepts into relatable analogies and examples
2. **Improve narrative structure**: Organize as an engaging story with clear beginning, middle, and end
3. **Clarify audience expectations**: Specify general audience level and engagement goals
4. **Add human context**: Include real-world applications and human interest elements
5. **Make it compelling**: Ensure the prompt guides toward fascinating and wonder-inspiring content
6. **Include visual elements**: Suggest use of metaphors and descriptive language for complex concepts
{% elif report_style == "news" %}
# Enhancement Guidelines for News Style
1. **Add journalistic rigor**: Include fact-checking requirements, source verification, and objectivity standards
2. **Improve news structure**: Organize with inverted pyramid structure (most important information first)
3. **Clarify reporting expectations**: Specify timeliness, accuracy, and balanced perspective requirements
4. **Add contextual background**: Include relevant background information and broader implications
5. **Make it newsworthy**: Ensure the prompt focuses on current relevance and public interest
6. **Include attribution**: Specify source requirements and quote standards
{% elif report_style == "social_media" %}
# Enhancement Guidelines for Social Media Style
1. **Add engagement focus**: Include attention-grabbing elements, hooks, and shareability factors
2. **Improve platform structure**: Organize for specific platform requirements (character limits, hashtags, etc.)
3. **Clarify audience expectations**: Specify target demographic and engagement goals
4. **Add viral elements**: Include trending topics, relatable content, and interactive elements
5. **Make it shareable**: Ensure the prompt guides toward content that encourages sharing and discussion
6. **Include visual considerations**: Suggest emoji usage, formatting, and visual appeal elements
{% else %}
# General Enhancement Guidelines
1. **Add specificity**: Include relevant details, scope, and constraints
2. **Improve structure**: Organize the request logically with clear sections if needed
3. **Clarify expectations**: Specify desired output format, length, or style
4. **Add context**: Include background information that would help generate better results
5. **Make it actionable**: Ensure the prompt guides toward concrete, useful outputs
{% endif %}
# Output Requirements
- You may include thoughts or reasoning before your final answer
- Wrap the final enhanced prompt in XML tags: <enhanced_prompt></enhanced_prompt>
- Do NOT include any explanations, comments, or meta-text within the XML tags
- Do NOT use phrases like "Enhanced Prompt:" or "Here's the enhanced version:" within the XML tags
- The content within the XML tags should be ready to use directly as a prompt
{% if report_style == "academic" %}
# Academic Style Examples
**Original**: "Write about AI"
**Enhanced**:
<enhanced_prompt>
Conduct a comprehensive academic analysis of artificial intelligence applications across three key sectors: healthcare, education, and business. Employ a systematic literature review methodology to examine peer-reviewed sources from the past five years. Structure your analysis with: (1) theoretical framework defining AI and its taxonomies, (2) sector-specific case studies with quantitative performance metrics, (3) critical evaluation of implementation challenges and ethical considerations, (4) comparative analysis across sectors, and (5) evidence-based recommendations for future research directions. Maintain academic rigor with proper citations, acknowledge methodological limitations, and present findings with appropriate hedging language. Target length: 3000-4000 words with APA formatting.
</enhanced_prompt>
**Original**: "Explain climate change"
**Enhanced**:
<enhanced_prompt>
Provide a rigorous academic examination of anthropogenic climate change, synthesizing current scientific consensus and recent research developments. Structure your analysis as follows: (1) theoretical foundations of greenhouse effect and radiative forcing mechanisms, (2) systematic review of empirical evidence from paleoclimatic, observational, and modeling studies, (3) critical analysis of attribution studies linking human activities to observed warming, (4) evaluation of climate sensitivity estimates and uncertainty ranges, (5) assessment of projected impacts under different emission scenarios, and (6) discussion of research gaps and methodological limitations. Include quantitative data, statistical significance levels, and confidence intervals where appropriate. Cite peer-reviewed sources extensively and maintain objective, third-person academic voice throughout.
</enhanced_prompt>
{% elif report_style == "popular_science" %}
# Popular Science Style Examples
**Original**: "Write about AI"
**Enhanced**:
<enhanced_prompt>
Tell the fascinating story of how artificial intelligence is quietly revolutionizing our daily lives in ways most people never realize. Take readers on an engaging journey through three surprising realms: the hospital where AI helps doctors spot diseases faster than ever before, the classroom where intelligent tutors adapt to each student's learning style, and the boardroom where algorithms are making million-dollar decisions. Use vivid analogies (like comparing neural networks to how our brains work) and real-world examples that readers can relate to. Include 'wow factor' moments that showcase AI's incredible capabilities, but also honest discussions about current limitations. Write with infectious enthusiasm while maintaining scientific accuracy, and conclude with exciting possibilities that await us in the near future. Aim for 1500-2000 words that feel like a captivating conversation with a brilliant friend.
</enhanced_prompt>
**Original**: "Explain climate change"
**Enhanced**:
<enhanced_prompt>
Craft a compelling narrative that transforms the complex science of climate change into an accessible and engaging story for curious readers. Begin with a relatable scenario (like why your hometown weather feels different than when you were a kid) and use this as a gateway to explore the fascinating science behind our changing planet. Employ vivid analogies - compare Earth's atmosphere to a blanket, greenhouse gases to invisible heat-trapping molecules, and climate feedback loops to a snowball rolling downhill. Include surprising facts and 'aha moments' that will make readers think differently about the world around them. Weave in human stories of scientists making discoveries, communities adapting to change, and innovative solutions being developed. Balance the serious implications with hope and actionable insights, concluding with empowering steps readers can take. Write with wonder and curiosity, making complex concepts feel approachable and personally relevant.
</enhanced_prompt>
{% elif report_style == "news" %}
# News Style Examples
**Original**: "Write about AI"
**Enhanced**:
<enhanced_prompt>
Report on the current state and immediate impact of artificial intelligence across three critical sectors: healthcare, education, and business. Lead with the most newsworthy developments and recent breakthroughs that are affecting people today. Structure using inverted pyramid format: start with key findings and immediate implications, then provide essential background context, followed by detailed analysis and expert perspectives. Include specific, verifiable data points, recent statistics, and quotes from credible sources including industry leaders, researchers, and affected stakeholders. Address both benefits and concerns with balanced reporting, fact-check all claims, and provide proper attribution for all information. Focus on timeliness and relevance to current events, highlighting what's happening now and what readers need to know. Maintain journalistic objectivity while making the significance clear to a general news audience. Target 800-1200 words following AP style guidelines.
</enhanced_prompt>
**Original**: "Explain climate change"
**Enhanced**:
<enhanced_prompt>
Provide comprehensive news coverage of climate change that explains the current scientific understanding and immediate implications for readers. Lead with the most recent and significant developments in climate science, policy, or impacts that are making headlines today. Structure the report with: breaking developments first, essential background for understanding the issue, current scientific consensus with specific data and timeframes, real-world impacts already being observed, policy responses and debates, and what experts say comes next. Include quotes from credible climate scientists, policy makers, and affected communities. Present information objectively while clearly communicating the scientific consensus, fact-check all claims, and provide proper source attribution. Address common misconceptions with factual corrections. Focus on what's happening now, why it matters to readers, and what they can expect in the near future. Follow journalistic standards for accuracy, balance, and timeliness.
</enhanced_prompt>
{% elif report_style == "social_media" %}
# Social Media Style Examples
**Original**: "Write about AI"
**Enhanced**:
<enhanced_prompt>
Create engaging social media content about AI that will stop the scroll and spark conversations! Start with an attention-grabbing hook like 'You won't believe what AI just did in hospitals this week 🤯' and structure as a compelling thread or post series. Include surprising facts, relatable examples (like AI helping doctors spot diseases or personalizing your Netflix recommendations), and interactive elements that encourage sharing and comments. Use strategic hashtags (#AI #Technology #Future), incorporate relevant emojis for visual appeal, and include questions that prompt audience engagement ('Have you noticed AI in your daily life? Drop examples below! 👇'). Make complex concepts digestible with bite-sized explanations, trending analogies, and shareable quotes. Include a clear call-to-action and optimize for the specific platform (Twitter threads, Instagram carousel, LinkedIn professional insights, or TikTok-style quick facts). Aim for high shareability with content that feels both informative and entertaining.
</enhanced_prompt>
**Original**: "Explain climate change"
**Enhanced**:
<enhanced_prompt>
Develop viral-worthy social media content that makes climate change accessible and shareable without being preachy. Open with a scroll-stopping hook like 'The weather app on your phone is telling a bigger story than you think 📱🌡️' and break down complex science into digestible, engaging chunks. Use relatable comparisons (Earth's fever, atmosphere as a blanket), trending formats (before/after visuals, myth-busting series, quick facts), and interactive elements (polls, questions, challenges). Include strategic hashtags (#ClimateChange #Science #Environment), eye-catching emojis, and shareable graphics or infographics. Address common questions and misconceptions with clear, factual responses. Create content that encourages positive action rather than climate anxiety, ending with empowering steps followers can take. Optimize for platform-specific features (Instagram Stories, TikTok trends, Twitter threads) and include calls-to-action that drive engagement and sharing.
</enhanced_prompt>
{% else %}
# General Examples
**Original**: "Write about AI"
**Enhanced**:
<enhanced_prompt>
Write a comprehensive 1000-word analysis of artificial intelligence's current applications in healthcare, education, and business. Include specific examples of AI tools being used in each sector, discuss both benefits and challenges, and provide insights into future trends. Structure the response with clear sections for each industry and conclude with key takeaways.
</enhanced_prompt>
**Original**: "Explain climate change"
**Enhanced**:
<enhanced_prompt>
Provide a detailed explanation of climate change suitable for a general audience. Cover the scientific mechanisms behind global warming, major causes including greenhouse gas emissions, observable effects we're seeing today, and projected future impacts. Include specific data and examples, and explain the difference between weather and climate. Organize the response with clear headings and conclude with actionable steps individuals can take.
</enhanced_prompt>
{% endif %}
+2 -159
View File
@@ -2,21 +2,7 @@
CURRENT_TIME: {{ CURRENT_TIME }} CURRENT_TIME: {{ CURRENT_TIME }}
--- ---
{% if report_style == "academic" %} You are a professional reporter responsible for writing clear, comprehensive reports based ONLY on provided information and verifiable facts.
You are a distinguished academic researcher and scholarly writer. Your report must embody the highest standards of academic rigor and intellectual discourse. Write with the precision of a peer-reviewed journal article, employing sophisticated analytical frameworks, comprehensive literature synthesis, and methodological transparency. Your language should be formal, technical, and authoritative, utilizing discipline-specific terminology with exactitude. Structure arguments logically with clear thesis statements, supporting evidence, and nuanced conclusions. Maintain complete objectivity, acknowledge limitations, and present balanced perspectives on controversial topics. The report should demonstrate deep scholarly engagement and contribute meaningfully to academic knowledge.
{% elif report_style == "popular_science" %}
You are an award-winning science communicator and storyteller. Your mission is to transform complex scientific concepts into captivating narratives that spark curiosity and wonder in everyday readers. Write with the enthusiasm of a passionate educator, using vivid analogies, relatable examples, and compelling storytelling techniques. Your tone should be warm, approachable, and infectious in its excitement about discovery. Break down technical jargon into accessible language without sacrificing accuracy. Use metaphors, real-world comparisons, and human interest angles to make abstract concepts tangible. Think like a National Geographic writer or a TED Talk presenter - engaging, enlightening, and inspiring.
{% elif report_style == "news" %}
You are an NBC News correspondent and investigative journalist with decades of experience in breaking news and in-depth reporting. Your report must exemplify the gold standard of American broadcast journalism: authoritative, meticulously researched, and delivered with the gravitas and credibility that NBC News is known for. Write with the precision of a network news anchor, employing the classic inverted pyramid structure while weaving compelling human narratives. Your language should be clear, authoritative, and accessible to prime-time television audiences. Maintain NBC's tradition of balanced reporting, thorough fact-checking, and ethical journalism. Think like Lester Holt or Andrea Mitchell - delivering complex stories with clarity, context, and unwavering integrity.
{% elif report_style == "social_media" %}
{% if locale == "zh-CN" %}
You are a popular 小红书 (Xiaohongshu) content creator specializing in lifestyle and knowledge sharing. Your report should embody the authentic, personal, and engaging style that resonates with 小红书 users. Write with genuine enthusiasm and a "姐妹们" (sisters) tone, as if sharing exciting discoveries with close friends. Use abundant emojis, create "种草" (grass-planting/recommendation) moments, and structure content for easy mobile consumption. Your writing should feel like a personal diary entry mixed with expert insights - warm, relatable, and irresistibly shareable. Think like a top 小红书 blogger who effortlessly combines personal experience with valuable information, making readers feel like they've discovered a hidden gem.
{% else %}
You are a viral Twitter content creator and digital influencer specializing in breaking down complex topics into engaging, shareable threads. Your report should be optimized for maximum engagement and viral potential across social media platforms. Write with energy, authenticity, and a conversational tone that resonates with global online communities. Use strategic hashtags, create quotable moments, and structure content for easy consumption and sharing. Think like a successful Twitter thought leader who can make any topic accessible, engaging, and discussion-worthy while maintaining credibility and accuracy.
{% endif %}
{% else %}
You are a professional reporter responsible for writing clear, comprehensive reports based ONLY on provided information and verifiable facts. Your report should adopt a professional tone.
{% endif %}
# Role # Role
@@ -57,40 +43,10 @@ Structure your report in the following format:
- **Including images from the previous steps in the report is very helpful.** - **Including images from the previous steps in the report is very helpful.**
5. **Survey Note** (for more comprehensive reports) 5. **Survey Note** (for more comprehensive reports)
{% if report_style == "academic" %}
- **Literature Review & Theoretical Framework**: Comprehensive analysis of existing research and theoretical foundations
- **Methodology & Data Analysis**: Detailed examination of research methods and analytical approaches
- **Critical Discussion**: In-depth evaluation of findings with consideration of limitations and implications
- **Future Research Directions**: Identification of gaps and recommendations for further investigation
{% elif report_style == "popular_science" %}
- **The Bigger Picture**: How this research fits into the broader scientific landscape
- **Real-World Applications**: Practical implications and potential future developments
- **Behind the Scenes**: Interesting details about the research process and challenges faced
- **What's Next**: Exciting possibilities and upcoming developments in the field
{% elif report_style == "news" %}
- **NBC News Analysis**: In-depth examination of the story's broader implications and significance
- **Impact Assessment**: How these developments affect different communities, industries, and stakeholders
- **Expert Perspectives**: Insights from credible sources, analysts, and subject matter experts
- **Timeline & Context**: Chronological background and historical context essential for understanding
- **What's Next**: Expected developments, upcoming milestones, and stories to watch
{% elif report_style == "social_media" %}
{% if locale == "zh-CN" %}
- **【种草时刻】**: 最值得关注的亮点和必须了解的核心信息
- **【数据震撼】**: 用小红书风格展示重要统计数据和发现
- **【姐妹们的看法】**: 社区热议话题和大家的真实反馈
- **【行动指南】**: 实用建议和读者可以立即行动的清单
{% else %}
- **Thread Highlights**: Key takeaways formatted for maximum shareability
- **Data That Matters**: Important statistics and findings presented for viral potential
- **Community Pulse**: Trending discussions and reactions from the online community
- **Action Steps**: Practical advice and immediate next steps for readers
{% endif %}
{% else %}
- A more detailed, academic-style analysis. - A more detailed, academic-style analysis.
- Include comprehensive sections covering all aspects of the topic. - Include comprehensive sections covering all aspects of the topic.
- Can include comparative analysis, tables, and detailed feature breakdowns. - Can include comparative analysis, tables, and detailed feature breakdowns.
- This section is optional for shorter reports. - This section is optional for shorter reports.
{% endif %}
6. **Key Citations** 6. **Key Citations**
- List all references at the end in link reference format. - List all references at the end in link reference format.
@@ -100,64 +56,7 @@ Structure your report in the following format:
# Writing Guidelines # Writing Guidelines
1. Writing style: 1. Writing style:
{% if report_style == "academic" %} - Use professional tone.
**Academic Excellence Standards:**
- Employ sophisticated, formal academic discourse with discipline-specific terminology
- Construct complex, nuanced arguments with clear thesis statements and logical progression
- Use third-person perspective and passive voice where appropriate for objectivity
- Include methodological considerations and acknowledge research limitations
- Reference theoretical frameworks and cite relevant scholarly work patterns
- Maintain intellectual rigor with precise, unambiguous language
- Avoid contractions, colloquialisms, and informal expressions entirely
- Use hedging language appropriately ("suggests," "indicates," "appears to")
{% elif report_style == "popular_science" %}
**Science Communication Excellence:**
- Write with infectious enthusiasm and genuine curiosity about discoveries
- Transform technical jargon into vivid, relatable analogies and metaphors
- Use active voice and engaging narrative techniques to tell scientific stories
- Include "wow factor" moments and surprising revelations to maintain interest
- Employ conversational tone while maintaining scientific accuracy
- Use rhetorical questions to engage readers and guide their thinking
- Include human elements: researcher personalities, discovery stories, real-world impacts
- Balance accessibility with intellectual respect for your audience
{% elif report_style == "news" %}
**NBC News Editorial Standards:**
- Open with a compelling lede that captures the essence of the story in 25-35 words
- Use the classic inverted pyramid: most newsworthy information first, supporting details follow
- Write in clear, conversational broadcast style that sounds natural when read aloud
- Employ active voice and strong, precise verbs that convey action and urgency
- Attribute every claim to specific, credible sources using NBC's attribution standards
- Use present tense for ongoing situations, past tense for completed events
- Maintain NBC's commitment to balanced reporting with multiple perspectives
- Include essential context and background without overwhelming the main story
- Verify information through at least two independent sources when possible
- Clearly label speculation, analysis, and ongoing investigations
- Use transitional phrases that guide readers smoothly through the narrative
{% elif report_style == "social_media" %}
{% if locale == "zh-CN" %}
**小红书风格写作标准:**
- 用"姐妹们!"、"宝子们!"等亲切称呼开头,营造闺蜜聊天氛围
- 大量使用emoji表情符号增强表达力和视觉吸引力 ✨
- 采用"种草"语言:"真的绝了!"、"必须安利给大家!"、"不看后悔系列!"
- 使用小红书特色标题格式:"【干货分享】"、"【亲测有效】"、"【避雷指南】"
- 穿插个人感受和体验:"我当时看到这个数据真的震惊了!"
- 用数字和符号增强视觉效果:①②③、✅❌、🔥💡⭐
- 创造"金句"和可截图分享的内容段落
- 结尾用互动性语言:"你们觉得呢?"、"评论区聊聊!"、"记得点赞收藏哦!"
{% else %}
**Twitter/X Engagement Standards:**
- Open with attention-grabbing hooks that stop the scroll
- Use thread-style formatting with numbered points (1/n, 2/n, etc.)
- Incorporate strategic hashtags for discoverability and trending topics
- Write quotable, tweetable snippets that beg to be shared
- Use conversational, authentic voice with personality and wit
- Include relevant emojis to enhance meaning and visual appeal 🧵📊💡
- Create "thread-worthy" content with clear progression and payoff
- End with engagement prompts: "What do you think?", "Retweet if you agree"
{% endif %}
{% else %}
- Use a professional tone.
{% endif %}
- Be concise and precise. - Be concise and precise.
- Avoid speculation. - Avoid speculation.
- Support claims with evidence. - Support claims with evidence.
@@ -178,62 +77,6 @@ Structure your report in the following format:
- Use horizontal rules (---) to separate major sections. - Use horizontal rules (---) to separate major sections.
- Track the sources of information but keep the main text clean and readable. - Track the sources of information but keep the main text clean and readable.
{% if report_style == "academic" %}
**Academic Formatting Specifications:**
- Use formal section headings with clear hierarchical structure (## Introduction, ### Methodology, #### Subsection)
- Employ numbered lists for methodological steps and logical sequences
- Use block quotes for important definitions or key theoretical concepts
- Include detailed tables with comprehensive headers and statistical data
- Use footnote-style formatting for additional context or clarifications
- Maintain consistent academic citation patterns throughout
- Use `code blocks` for technical specifications, formulas, or data samples
{% elif report_style == "popular_science" %}
**Science Communication Formatting:**
- Use engaging, descriptive headings that spark curiosity ("The Surprising Discovery That Changed Everything")
- Employ creative formatting like callout boxes for "Did You Know?" facts
- Use bullet points for easy-to-digest key findings
- Include visual breaks with strategic use of bold text for emphasis
- Format analogies and metaphors prominently to aid understanding
- Use numbered lists for step-by-step explanations of complex processes
- Highlight surprising statistics or findings with special formatting
{% elif report_style == "news" %}
**NBC News Formatting Standards:**
- Craft headlines that are informative yet compelling, following NBC's style guide
- Use NBC-style datelines and bylines for professional credibility
- Structure paragraphs for broadcast readability (1-2 sentences for digital, 2-3 for print)
- Employ strategic subheadings that advance the story narrative
- Format direct quotes with proper attribution and context
- Use bullet points sparingly, primarily for breaking news updates or key facts
- Include "BREAKING" or "DEVELOPING" labels for ongoing stories
- Format source attribution clearly: "according to NBC News," "sources tell NBC News"
- Use italics for emphasis on key terms or breaking developments
- Structure the story with clear sections: Lede, Context, Analysis, Looking Ahead
{% elif report_style == "social_media" %}
{% if locale == "zh-CN" %}
**小红书格式优化标准:**
- 使用吸睛标题配合emoji:"🔥【重磅】这个发现太震撼了!"
- 关键数据用醒目格式突出:「 重点数据 」或 ⭐ 核心发现 ⭐
- 适度使用大写强调:真的YYDS!、绝绝子!
- 用emoji作为分点符号:✨、🌟、、💯
- 创建话题标签区域:#科技前沿 #必看干货 #涨知识了
- 设置"划重点"总结区域,方便快速阅读
- 利用换行和空白营造手机阅读友好的版式
- 制作"金句卡片"格式,便于截图分享
- 使用分割线和特殊符号:「」『』【】━━━━━━
{% else %}
**Twitter/X Formatting Standards:**
- Use compelling headlines with strategic emoji placement 🧵⚡️🔥
- Format key insights as standalone, quotable tweet blocks
- Employ thread numbering for multi-part content (1/12, 2/12, etc.)
- Use bullet points with emoji bullets for visual appeal
- Include strategic hashtags at the end: #TechNews #Innovation #MustRead
- Create "TL;DR" summaries for quick consumption
- Use line breaks and white space for mobile readability
- Format "quotable moments" with clear visual separation
- Include call-to-action elements: "🔄 RT to share" "💬 What's your take?"
{% endif %}
{% endif %}
# Data Integrity # Data Integrity
- Only use information explicitly provided in the input. - Only use information explicitly provided in the input.
+1 -4
View File
@@ -11,9 +11,6 @@ You are dedicated to conducting thorough investigations using search tools and p
You have access to two types of tools: You have access to two types of tools:
1. **Built-in Tools**: These are always available: 1. **Built-in Tools**: These are always available:
{% if resources %}
- **local_search_tool**: For retrieving information from the local knowledge base when user mentioned in the messages.
{% endif %}
- **web_search_tool**: For performing web searches - **web_search_tool**: For performing web searches
- **crawl_tool**: For reading content from URLs - **crawl_tool**: For reading content from URLs
@@ -37,7 +34,7 @@ You have access to two types of tools:
3. **Plan the Solution**: Determine the best approach to solve the problem using the available tools. 3. **Plan the Solution**: Determine the best approach to solve the problem using the available tools.
4. **Execute the Solution**: 4. **Execute the Solution**:
- Forget your previous knowledge, so you **should leverage the tools** to retrieve the information. - Forget your previous knowledge, so you **should leverage the tools** to retrieve the information.
- Use the {% if resources %}**local_search_tool** or{% endif %}**web_search_tool** or other suitable search tool to perform a search with the provided keywords. - Use the **web_search_tool** or other suitable search tool to perform a search with the provided keywords.
- When the task includes time range requirements: - When the task includes time range requirements:
- Incorporate appropriate time-based search parameters in your queries (e.g., "after:2020", "before:2023", or specific date ranges) - Incorporate appropriate time-based search parameters in your queries (e.g., "after:2020", "before:2023", or specific date ranges)
- Ensure search results respect the specified time constraints. - Ensure search results respect the specified time constraints.
-17
View File
@@ -1,17 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
from .retriever import Retriever, Document, Resource, Chunk
from .ragflow import RAGFlowProvider
from .vikingdb_knowledge_base import VikingDBKnowledgeBaseProvider
from .builder import build_retriever
__all__ = [
Retriever,
Document,
Resource,
RAGFlowProvider,
VikingDBKnowledgeBaseProvider,
Chunk,
build_retriever,
]
-17
View File
@@ -1,17 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
from src.config.tools import SELECTED_RAG_PROVIDER, RAGProvider
from src.rag.ragflow import RAGFlowProvider
from src.rag.vikingdb_knowledge_base import VikingDBKnowledgeBaseProvider
from src.rag.retriever import Retriever
def build_retriever() -> Retriever | None:
if SELECTED_RAG_PROVIDER == RAGProvider.RAGFLOW.value:
return RAGFlowProvider()
elif SELECTED_RAG_PROVIDER == RAGProvider.VIKINGDB_KNOWLEDGE_BASE.value:
return VikingDBKnowledgeBaseProvider()
elif SELECTED_RAG_PROVIDER:
raise ValueError(f"Unsupported RAG provider: {SELECTED_RAG_PROVIDER}")
return None
-124
View File
@@ -1,124 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import os
import requests
from src.rag.retriever import Chunk, Document, Resource, Retriever
from urllib.parse import urlparse
class RAGFlowProvider(Retriever):
"""
RAGFlowProvider is a provider that uses RAGFlow to retrieve documents.
"""
api_url: str
api_key: str
page_size: int = 10
def __init__(self):
api_url = os.getenv("RAGFLOW_API_URL")
if not api_url:
raise ValueError("RAGFLOW_API_URL is not set")
self.api_url = api_url
api_key = os.getenv("RAGFLOW_API_KEY")
if not api_key:
raise ValueError("RAGFLOW_API_KEY is not set")
self.api_key = api_key
page_size = os.getenv("RAGFLOW_PAGE_SIZE")
if page_size:
self.page_size = int(page_size)
def query_relevant_documents(
self, query: str, resources: list[Resource] = []
) -> list[Document]:
headers = {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
}
dataset_ids: list[str] = []
document_ids: list[str] = []
for resource in resources:
dataset_id, document_id = parse_uri(resource.uri)
dataset_ids.append(dataset_id)
if document_id:
document_ids.append(document_id)
payload = {
"question": query,
"dataset_ids": dataset_ids,
"document_ids": document_ids,
"page_size": self.page_size,
}
response = requests.post(
f"{self.api_url}/api/v1/retrieval", headers=headers, json=payload
)
if response.status_code != 200:
raise Exception(f"Failed to query documents: {response.text}")
result = response.json()
data = result.get("data", {})
doc_aggs = data.get("doc_aggs", [])
docs: dict[str, Document] = {
doc.get("doc_id"): Document(
id=doc.get("doc_id"),
title=doc.get("doc_name"),
chunks=[],
)
for doc in doc_aggs
}
for chunk in data.get("chunks", []):
doc = docs.get(chunk.get("document_id"))
if doc:
doc.chunks.append(
Chunk(
content=chunk.get("content"),
similarity=chunk.get("similarity"),
)
)
return list(docs.values())
def list_resources(self, query: str | None = None) -> list[Resource]:
headers = {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
}
params = {}
if query:
params["name"] = query
response = requests.get(
f"{self.api_url}/api/v1/datasets", headers=headers, params=params
)
if response.status_code != 200:
raise Exception(f"Failed to list resources: {response.text}")
result = response.json()
resources = []
for item in result.get("data", []):
item = Resource(
uri=f"rag://dataset/{item.get('id')}",
title=item.get("name", ""),
description=item.get("description", ""),
)
resources.append(item)
return resources
def parse_uri(uri: str) -> tuple[str, str]:
parsed = urlparse(uri)
if parsed.scheme != "rag":
raise ValueError(f"Invalid URI: {uri}")
return parsed.path.split("/")[1], parsed.fragment
-80
View File
@@ -1,80 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import abc
from pydantic import BaseModel, Field
class Chunk:
content: str
similarity: float
def __init__(self, content: str, similarity: float):
self.content = content
self.similarity = similarity
class Document:
"""
Document is a class that represents a document.
"""
id: str
url: str | None = None
title: str | None = None
chunks: list[Chunk] = []
def __init__(
self,
id: str,
url: str | None = None,
title: str | None = None,
chunks: list[Chunk] = [],
):
self.id = id
self.url = url
self.title = title
self.chunks = chunks
def to_dict(self) -> dict:
d = {
"id": self.id,
"content": "\n\n".join([chunk.content for chunk in self.chunks]),
}
if self.url:
d["url"] = self.url
if self.title:
d["title"] = self.title
return d
class Resource(BaseModel):
"""
Resource is a class that represents a resource.
"""
uri: str = Field(..., description="The URI of the resource")
title: str = Field(..., description="The title of the resource")
description: str | None = Field("", description="The description of the resource")
class Retriever(abc.ABC):
"""
Define a RAG provider, which can be used to query documents and resources.
"""
@abc.abstractmethod
def list_resources(self, query: str | None = None) -> list[Resource]:
"""
List resources from the rag provider.
"""
pass
@abc.abstractmethod
def query_relevant_documents(
self, query: str, resources: list[Resource] = []
) -> list[Document]:
"""
Query relevant documents from the resources.
"""
pass
-208
View File
@@ -1,208 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import os
import requests
import json
from src.rag.retriever import Chunk, Document, Resource, Retriever
from urllib.parse import urlparse
from volcengine.auth.SignerV4 import SignerV4
from volcengine.base.Request import Request
from volcengine.Credentials import Credentials
class VikingDBKnowledgeBaseProvider(Retriever):
"""
VikingDBKnowledgeBaseProvider is a provider that uses VikingDB Knowledge base API to retrieve documents.
"""
api_url: str
api_ak: str
api_sk: str
retrieval_size: int = 10
def __init__(self):
api_url = os.getenv("VIKINGDB_KNOWLEDGE_BASE_API_URL")
if not api_url:
raise ValueError("VIKINGDB_KNOWLEDGE_BASE_API_URL is not set")
self.api_url = api_url
api_ak = os.getenv("VIKINGDB_KNOWLEDGE_BASE_API_AK")
if not api_ak:
raise ValueError("VIKINGDB_KNOWLEDGE_BASE_API_AK is not set")
self.api_ak = api_ak
api_sk = os.getenv("VIKINGDB_KNOWLEDGE_BASE_API_SK")
if not api_sk:
raise ValueError("VIKINGDB_KNOWLEDGE_BASE_API_SK is not set")
self.api_sk = api_sk
retrieval_size = os.getenv("VIKINGDB_KNOWLEDGE_BASE_RETRIEVAL_SIZE")
if retrieval_size:
self.retrieval_size = int(retrieval_size)
def prepare_request(self, method, path, params=None, data=None, doseq=0):
"""
Prepare signed request using volcengine auth
"""
if params:
for key in params:
if (
type(params[key]) is int
or type(params[key]) is float
or type(params[key]) is bool
):
params[key] = str(params[key])
elif type(params[key]) is list:
if not doseq:
params[key] = ",".join(params[key])
r = Request()
r.set_shema("https")
r.set_method(method)
r.set_connection_timeout(10)
r.set_socket_timeout(10)
mheaders = {
"Accept": "application/json",
"Content-Type": "application/json",
}
r.set_headers(mheaders)
if params:
r.set_query(params)
r.set_path(path)
if data is not None:
r.set_body(json.dumps(data))
credentials = Credentials(self.api_ak, self.api_sk, "air", "cn-north-1")
SignerV4.sign(r, credentials)
return r
def query_relevant_documents(
self, query: str, resources: list[Resource] = []
) -> list[Document]:
"""
Query relevant documents from the knowledge base
"""
if not resources:
return []
all_documents = {}
for resource in resources:
resource_id, document_id = parse_uri(resource.uri)
request_params = {
"resource_id": resource_id,
"query": query,
"limit": self.retrieval_size,
"dense_weight": 0.5,
"pre_processing": {
"need_instruction": True,
"rewrite": False,
"return_token_usage": True,
},
"post_processing": {
"rerank_switch": True,
"chunk_diffusion_count": 0,
"chunk_group": True,
"get_attachment_link": True,
},
}
if document_id:
doc_filter = {"op": "must", "field": "doc_id", "conds": [document_id]}
query_param = {"doc_filter": doc_filter}
request_params["query_param"] = query_param
method = "POST"
path = "/api/knowledge/collection/search_knowledge"
info_req = self.prepare_request(
method=method, path=path, data=request_params
)
rsp = requests.request(
method=info_req.method,
url="http://{}{}".format(self.api_url, info_req.path),
headers=info_req.headers,
data=info_req.body,
)
try:
response = json.loads(rsp.text)
except json.JSONDecodeError as e:
raise ValueError(f"Failed to parse JSON response: {e}")
if response["code"] != 0:
raise ValueError(
f"Failed to query documents from resource: {response['message']}"
)
rsp_data = response.get("data", {})
if "result_list" not in rsp_data:
continue
result_list = rsp_data["result_list"]
for item in result_list:
doc_info = item.get("doc_info", {})
doc_id = doc_info.get("doc_id")
if not doc_id:
continue
if doc_id not in all_documents:
all_documents[doc_id] = Document(
id=doc_id, title=doc_info.get("doc_name"), chunks=[]
)
chunk = Chunk(
content=item.get("content", ""), similarity=item.get("score", 0.0)
)
all_documents[doc_id].chunks.append(chunk)
return list(all_documents.values())
def list_resources(self, query: str | None = None) -> list[Resource]:
"""
List resources (knowledge bases) from the knowledge base service
"""
method = "POST"
path = "/api/knowledge/collection/list"
info_req = self.prepare_request(method=method, path=path)
rsp = requests.request(
method=info_req.method,
url="http://{}{}".format(self.api_url, info_req.path),
headers=info_req.headers,
data=info_req.body,
)
try:
response = json.loads(rsp.text)
except json.JSONDecodeError as e:
raise ValueError(f"Failed to parse JSON response: {e}")
if response["code"] != 0:
raise Exception(f"Failed to list resources: {response["message"]}")
resources = []
rsp_data = response.get("data", {})
collection_list = rsp_data.get("collection_list", [])
for item in collection_list:
collection_name = item.get("collection_name", "")
description = item.get("description", "")
if query and query.lower() not in collection_name.lower():
continue
resource_id = item.get("resource_id", "")
resource = Resource(
uri=f"rag://dataset/{resource_id}",
title=collection_name,
description=description,
)
resources.append(resource)
return resources
def parse_uri(uri: str) -> tuple[str, str]:
parsed = urlparse(uri)
if parsed.scheme != "rag":
raise ValueError(f"Invalid URI: {uri}")
return parsed.path.split("/")[1], parsed.fragment
+22 -122
View File
@@ -5,47 +5,33 @@ import base64
import json import json
import logging import logging
import os import os
from typing import Annotated, List, cast from typing import List, cast
from uuid import uuid4 from uuid import uuid4
from fastapi import FastAPI, HTTPException, Query from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response, StreamingResponse from fastapi.responses import Response, StreamingResponse
from langchain_core.messages import AIMessageChunk, BaseMessage, ToolMessage from langchain_core.messages import AIMessageChunk, ToolMessage, BaseMessage
from langgraph.types import Command from langgraph.types import Command
from src.config.report_style import ReportStyle
from src.config.tools import SELECTED_RAG_PROVIDER
from src.graph.builder import build_graph_with_memory from src.graph.builder import build_graph_with_memory
from src.llms.llm import get_configured_llm_models
from src.podcast.graph.builder import build_graph as build_podcast_graph from src.podcast.graph.builder import build_graph as build_podcast_graph
from src.ppt.graph.builder import build_graph as build_ppt_graph from src.ppt.graph.builder import build_graph as build_ppt_graph
from src.prompt_enhancer.graph.builder import build_graph as build_prompt_enhancer_graph
from src.prose.graph.builder import build_graph as build_prose_graph from src.prose.graph.builder import build_graph as build_prose_graph
from src.rag.builder import build_retriever
from src.rag.retriever import Resource
from src.server.chat_request import ( from src.server.chat_request import (
ChatMessage,
ChatRequest, ChatRequest,
EnhancePromptRequest,
GeneratePodcastRequest, GeneratePodcastRequest,
GeneratePPTRequest, GeneratePPTRequest,
GenerateProseRequest, GenerateProseRequest,
TTSRequest, TTSRequest,
) )
from src.server.config_request import ConfigResponse
from src.server.mcp_request import MCPServerMetadataRequest, MCPServerMetadataResponse from src.server.mcp_request import MCPServerMetadataRequest, MCPServerMetadataResponse
from src.server.mcp_utils import load_mcp_tools from src.server.mcp_utils import load_mcp_tools
from src.server.rag_request import (
RAGConfigResponse,
RAGResourceRequest,
RAGResourcesResponse,
)
from src.tools import VolcengineTTS from src.tools import VolcengineTTS
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
INTERNAL_SERVER_ERROR_DETAIL = "Internal Server Error"
app = FastAPI( app = FastAPI(
title="DeerFlow API", title="DeerFlow API",
description="API for Deer", description="API for Deer",
@@ -53,17 +39,12 @@ app = FastAPI(
) )
# Add CORS middleware # Add CORS middleware
# It's recommended to load the allowed origins from an environment variable
# for better security and flexibility across different environments.
allowed_origins_str = os.getenv("ALLOWED_ORIGINS", "http://localhost:3000")
allowed_origins = [origin.strip() for origin in allowed_origins_str.split(",")]
app.add_middleware( app.add_middleware(
CORSMiddleware, CORSMiddleware,
allow_origins=allowed_origins, # Restrict to specific origins allow_origins=["*"], # Allows all origins
allow_credentials=True, allow_credentials=True,
allow_methods=["GET", "POST"], # Be specific about allowed methods allow_methods=["*"], # Allows all methods
allow_headers=["Content-Type", "Authorization", "X-Requested-With"], # Be specific allow_headers=["*"], # Allows all headers
) )
graph = build_graph_with_memory() graph = build_graph_with_memory()
@@ -78,7 +59,6 @@ async def chat_stream(request: ChatRequest):
_astream_workflow_generator( _astream_workflow_generator(
request.model_dump()["messages"], request.model_dump()["messages"],
thread_id, thread_id,
request.resources,
request.max_plan_iterations, request.max_plan_iterations,
request.max_step_num, request.max_step_num,
request.max_search_results, request.max_search_results,
@@ -86,26 +66,21 @@ async def chat_stream(request: ChatRequest):
request.interrupt_feedback, request.interrupt_feedback,
request.mcp_settings, request.mcp_settings,
request.enable_background_investigation, request.enable_background_investigation,
request.report_style,
request.enable_deep_thinking,
), ),
media_type="text/event-stream", media_type="text/event-stream",
) )
async def _astream_workflow_generator( async def _astream_workflow_generator(
messages: List[dict], messages: List[ChatMessage],
thread_id: str, thread_id: str,
resources: List[Resource],
max_plan_iterations: int, max_plan_iterations: int,
max_step_num: int, max_step_num: int,
max_search_results: int, max_search_results: int,
auto_accepted_plan: bool, auto_accepted_plan: bool,
interrupt_feedback: str, interrupt_feedback: str,
mcp_settings: dict, mcp_settings: dict,
enable_background_investigation: bool, enable_background_investigation,
report_style: ReportStyle,
enable_deep_thinking: bool,
): ):
input_ = { input_ = {
"messages": messages, "messages": messages,
@@ -115,7 +90,6 @@ async def _astream_workflow_generator(
"observations": [], "observations": [],
"auto_accepted_plan": auto_accepted_plan, "auto_accepted_plan": auto_accepted_plan,
"enable_background_investigation": enable_background_investigation, "enable_background_investigation": enable_background_investigation,
"research_topic": messages[-1]["content"] if messages else "",
} }
if not auto_accepted_plan and interrupt_feedback: if not auto_accepted_plan and interrupt_feedback:
resume_msg = f"[{interrupt_feedback}]" resume_msg = f"[{interrupt_feedback}]"
@@ -127,13 +101,10 @@ async def _astream_workflow_generator(
input_, input_,
config={ config={
"thread_id": thread_id, "thread_id": thread_id,
"resources": resources,
"max_plan_iterations": max_plan_iterations, "max_plan_iterations": max_plan_iterations,
"max_step_num": max_step_num, "max_step_num": max_step_num,
"max_search_results": max_search_results, "max_search_results": max_search_results,
"mcp_settings": mcp_settings, "mcp_settings": mcp_settings,
"report_style": report_style.value,
"enable_deep_thinking": enable_deep_thinking,
}, },
stream_mode=["messages", "updates"], stream_mode=["messages", "updates"],
subgraphs=True, subgraphs=True,
@@ -158,21 +129,13 @@ async def _astream_workflow_generator(
message_chunk, message_metadata = cast( message_chunk, message_metadata = cast(
tuple[BaseMessage, dict[str, any]], event_data tuple[BaseMessage, dict[str, any]], event_data
) )
# Handle empty agent tuple gracefully
agent_name = "unknown"
if agent and len(agent) > 0:
agent_name = agent[0].split(":")[0] if ":" in agent[0] else agent[0]
event_stream_message: dict[str, any] = { event_stream_message: dict[str, any] = {
"thread_id": thread_id, "thread_id": thread_id,
"agent": agent_name, "agent": agent[0].split(":")[0],
"id": message_chunk.id, "id": message_chunk.id,
"role": "assistant", "role": "assistant",
"content": message_chunk.content, "content": message_chunk.content,
} }
if message_chunk.additional_kwargs.get("reasoning_content"):
event_stream_message["reasoning_content"] = message_chunk.additional_kwargs[
"reasoning_content"
]
if message_chunk.response_metadata.get("finish_reason"): if message_chunk.response_metadata.get("finish_reason"):
event_stream_message["finish_reason"] = message_chunk.response_metadata.get( event_stream_message["finish_reason"] = message_chunk.response_metadata.get(
"finish_reason" "finish_reason"
@@ -210,16 +173,17 @@ def _make_event(event_type: str, data: dict[str, any]):
@app.post("/api/tts") @app.post("/api/tts")
async def text_to_speech(request: TTSRequest): async def text_to_speech(request: TTSRequest):
"""Convert text to speech using volcengine TTS API.""" """Convert text to speech using volcengine TTS API."""
try:
app_id = os.getenv("VOLCENGINE_TTS_APPID", "") app_id = os.getenv("VOLCENGINE_TTS_APPID", "")
if not app_id: if not app_id:
raise HTTPException(status_code=400, detail="VOLCENGINE_TTS_APPID is not set") raise HTTPException(
status_code=400, detail="VOLCENGINE_TTS_APPID is not set"
)
access_token = os.getenv("VOLCENGINE_TTS_ACCESS_TOKEN", "") access_token = os.getenv("VOLCENGINE_TTS_ACCESS_TOKEN", "")
if not access_token: if not access_token:
raise HTTPException( raise HTTPException(
status_code=400, detail="VOLCENGINE_TTS_ACCESS_TOKEN is not set" status_code=400, detail="VOLCENGINE_TTS_ACCESS_TOKEN is not set"
) )
try:
cluster = os.getenv("VOLCENGINE_TTS_CLUSTER", "volcano_tts") cluster = os.getenv("VOLCENGINE_TTS_CLUSTER", "volcano_tts")
voice_type = os.getenv("VOLCENGINE_TTS_VOICE_TYPE", "BV700_V2_streaming") voice_type = os.getenv("VOLCENGINE_TTS_VOICE_TYPE", "BV700_V2_streaming")
@@ -257,10 +221,9 @@ async def text_to_speech(request: TTSRequest):
) )
}, },
) )
except Exception as e: except Exception as e:
logger.exception(f"Error in TTS endpoint: {str(e)}") logger.exception(f"Error in TTS endpoint: {str(e)}")
raise HTTPException(status_code=500, detail=INTERNAL_SERVER_ERROR_DETAIL) raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/podcast/generate") @app.post("/api/podcast/generate")
@@ -274,7 +237,7 @@ async def generate_podcast(request: GeneratePodcastRequest):
return Response(content=audio_bytes, media_type="audio/mp3") return Response(content=audio_bytes, media_type="audio/mp3")
except Exception as e: except Exception as e:
logger.exception(f"Error occurred during podcast generation: {str(e)}") logger.exception(f"Error occurred during podcast generation: {str(e)}")
raise HTTPException(status_code=500, detail=INTERNAL_SERVER_ERROR_DETAIL) raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/ppt/generate") @app.post("/api/ppt/generate")
@@ -293,14 +256,13 @@ async def generate_ppt(request: GeneratePPTRequest):
) )
except Exception as e: except Exception as e:
logger.exception(f"Error occurred during ppt generation: {str(e)}") logger.exception(f"Error occurred during ppt generation: {str(e)}")
raise HTTPException(status_code=500, detail=INTERNAL_SERVER_ERROR_DETAIL) raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/prose/generate") @app.post("/api/prose/generate")
async def generate_prose(request: GenerateProseRequest): async def generate_prose(request: GenerateProseRequest):
try: try:
sanitized_prompt = request.prompt.replace("\r\n", "").replace("\n", "") logger.info(f"Generating prose for prompt: {request.prompt}")
logger.info(f"Generating prose for prompt: {sanitized_prompt}")
workflow = build_prose_graph() workflow = build_prose_graph()
events = workflow.astream( events = workflow.astream(
{ {
@@ -317,47 +279,7 @@ async def generate_prose(request: GenerateProseRequest):
) )
except Exception as e: except Exception as e:
logger.exception(f"Error occurred during prose generation: {str(e)}") logger.exception(f"Error occurred during prose generation: {str(e)}")
raise HTTPException(status_code=500, detail=INTERNAL_SERVER_ERROR_DETAIL) raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/prompt/enhance")
async def enhance_prompt(request: EnhancePromptRequest):
try:
sanitized_prompt = request.prompt.replace("\r\n", "").replace("\n", "")
logger.info(f"Enhancing prompt: {sanitized_prompt}")
# Convert string report_style to ReportStyle enum
report_style = None
if request.report_style:
try:
# Handle both uppercase and lowercase input
style_mapping = {
"ACADEMIC": ReportStyle.ACADEMIC,
"POPULAR_SCIENCE": ReportStyle.POPULAR_SCIENCE,
"NEWS": ReportStyle.NEWS,
"SOCIAL_MEDIA": ReportStyle.SOCIAL_MEDIA,
}
report_style = style_mapping.get(
request.report_style.upper(), ReportStyle.ACADEMIC
)
except Exception:
# If invalid style, default to ACADEMIC
report_style = ReportStyle.ACADEMIC
else:
report_style = ReportStyle.ACADEMIC
workflow = build_prompt_enhancer_graph()
final_state = workflow.invoke(
{
"prompt": request.prompt,
"context": request.context,
"report_style": report_style,
}
)
return {"result": final_state["output"]}
except Exception as e:
logger.exception(f"Error occurred during prompt enhancement: {str(e)}")
raise HTTPException(status_code=500, detail=INTERNAL_SERVER_ERROR_DETAIL)
@app.post("/api/mcp/server/metadata", response_model=MCPServerMetadataResponse) @app.post("/api/mcp/server/metadata", response_model=MCPServerMetadataResponse)
@@ -393,29 +315,7 @@ async def mcp_server_metadata(request: MCPServerMetadataRequest):
return response return response
except Exception as e: except Exception as e:
if not isinstance(e, HTTPException):
logger.exception(f"Error in MCP server metadata endpoint: {str(e)}") logger.exception(f"Error in MCP server metadata endpoint: {str(e)}")
raise HTTPException(status_code=500, detail=INTERNAL_SERVER_ERROR_DETAIL) raise HTTPException(status_code=500, detail=str(e))
raise
@app.get("/api/rag/config", response_model=RAGConfigResponse)
async def rag_config():
"""Get the config of the RAG."""
return RAGConfigResponse(provider=SELECTED_RAG_PROVIDER)
@app.get("/api/rag/resources", response_model=RAGResourcesResponse)
async def rag_resources(request: Annotated[RAGResourceRequest, Query()]):
"""Get the resources of the RAG."""
retriever = build_retriever()
if retriever:
return RAGResourcesResponse(resources=retriever.list_resources(request.query))
return RAGResourcesResponse(resources=[])
@app.get("/api/config", response_model=ConfigResponse)
async def config():
"""Get the config of the server."""
return ConfigResponse(
rag=RAGConfigResponse(provider=SELECTED_RAG_PROVIDER),
models=get_configured_llm_models(),
)
-22
View File
@@ -5,9 +5,6 @@ from typing import List, Optional, Union
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from src.rag.retriever import Resource
from src.config.report_style import ReportStyle
class ContentItem(BaseModel): class ContentItem(BaseModel):
type: str = Field(..., description="The type of content (text, image, etc.)") type: str = Field(..., description="The type of content (text, image, etc.)")
@@ -31,9 +28,6 @@ class ChatRequest(BaseModel):
messages: Optional[List[ChatMessage]] = Field( messages: Optional[List[ChatMessage]] = Field(
[], description="History of messages between the user and the assistant" [], description="History of messages between the user and the assistant"
) )
resources: Optional[List[Resource]] = Field(
[], description="Resources to be used for the research"
)
debug: Optional[bool] = Field(False, description="Whether to enable debug logging") debug: Optional[bool] = Field(False, description="Whether to enable debug logging")
thread_id: Optional[str] = Field( thread_id: Optional[str] = Field(
"__default__", description="A specific conversation identifier" "__default__", description="A specific conversation identifier"
@@ -59,12 +53,6 @@ class ChatRequest(BaseModel):
enable_background_investigation: Optional[bool] = Field( enable_background_investigation: Optional[bool] = Field(
True, description="Whether to get background investigation before plan" True, description="Whether to get background investigation before plan"
) )
report_style: Optional[ReportStyle] = Field(
ReportStyle.ACADEMIC, description="The style of the report"
)
enable_deep_thinking: Optional[bool] = Field(
False, description="Whether to enable deep thinking"
)
class TTSRequest(BaseModel): class TTSRequest(BaseModel):
@@ -97,13 +85,3 @@ class GenerateProseRequest(BaseModel):
command: Optional[str] = Field( command: Optional[str] = Field(
"", description="The user custom command of the prose writer" "", description="The user custom command of the prose writer"
) )
class EnhancePromptRequest(BaseModel):
prompt: str = Field(..., description="The original prompt to enhance")
context: Optional[str] = Field(
"", description="Additional context about the intended use"
)
report_style: Optional[str] = Field(
"academic", description="The style of the report"
)
-13
View File
@@ -1,13 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
from pydantic import BaseModel, Field
from src.server.rag_request import RAGConfigResponse
class ConfigResponse(BaseModel):
"""Response model for server config."""
rag: RAGConfigResponse = Field(..., description="The config of the RAG")
models: dict[str, list[str]] = Field(..., description="The configured models")
+1 -1
View File
@@ -3,7 +3,7 @@
import logging import logging
from datetime import timedelta from datetime import timedelta
from typing import Any, Dict, List, Optional from typing import Any, Dict, List, Optional, Tuple
from fastapi import HTTPException from fastapi import HTTPException
from mcp import ClientSession, StdioServerParameters from mcp import ClientSession, StdioServerParameters
-28
View File
@@ -1,28 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
from pydantic import BaseModel, Field
from src.rag.retriever import Resource
class RAGConfigResponse(BaseModel):
"""Response model for RAG config."""
provider: str | None = Field(
None, description="The provider of the RAG, default is ragflow"
)
class RAGResourceRequest(BaseModel):
"""Request model for RAG resource."""
query: str | None = Field(
None, description="The query of the resource need to be searched"
)
class RAGResourcesResponse(BaseModel):
"""Response model for RAG resources."""
resources: list[Resource] = Field(..., description="The resources of the RAG")
+2 -2
View File
@@ -1,9 +1,10 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT # SPDX-License-Identifier: MIT
import os
from .crawl import crawl_tool from .crawl import crawl_tool
from .python_repl import python_repl_tool from .python_repl import python_repl_tool
from .retriever import get_retriever_tool
from .search import get_web_search_tool from .search import get_web_search_tool
from .tts import VolcengineTTS from .tts import VolcengineTTS
@@ -11,6 +12,5 @@ __all__ = [
"crawl_tool", "crawl_tool",
"python_repl_tool", "python_repl_tool",
"get_web_search_tool", "get_web_search_tool",
"get_retriever_tool",
"VolcengineTTS", "VolcengineTTS",
] ]
-62
View File
@@ -1,62 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import logging
from typing import List, Optional, Type
from langchain_core.tools import BaseTool
from langchain_core.callbacks import (
AsyncCallbackManagerForToolRun,
CallbackManagerForToolRun,
)
from pydantic import BaseModel, Field
from src.config.tools import SELECTED_RAG_PROVIDER
from src.rag import Document, Retriever, Resource, build_retriever
logger = logging.getLogger(__name__)
class RetrieverInput(BaseModel):
keywords: str = Field(description="search keywords to look up")
class RetrieverTool(BaseTool):
name: str = "local_search_tool"
description: str = (
"Useful for retrieving information from the file with `rag://` uri prefix, it should be higher priority than the web search or writing code. Input should be a search keywords."
)
args_schema: Type[BaseModel] = RetrieverInput
retriever: Retriever = Field(default_factory=Retriever)
resources: list[Resource] = Field(default_factory=list)
def _run(
self,
keywords: str,
run_manager: Optional[CallbackManagerForToolRun] = None,
) -> list[Document]:
logger.info(
f"Retriever tool query: {keywords}", extra={"resources": self.resources}
)
documents = self.retriever.query_relevant_documents(keywords, self.resources)
if not documents:
return "No results found from the local knowledge base."
return [doc.to_dict() for doc in documents]
async def _arun(
self,
keywords: str,
run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
) -> list[Document]:
return self._run(keywords, run_manager.get_sync())
def get_retriever_tool(resources: List[Resource]) -> RetrieverTool | None:
if not resources:
return None
logger.info(f"create retriever tool: {SELECTED_RAG_PROVIDER}")
retriever = build_retriever()
if not retriever:
return None
return RetrieverTool(retriever=retriever, resources=resources)
+9 -24
View File
@@ -1,16 +1,15 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT # SPDX-License-Identifier: MIT
import json
import logging import logging
import os import os
from typing import List, Optional
from langchain_community.tools import BraveSearch, DuckDuckGoSearchResults from langchain_community.tools import BraveSearch, DuckDuckGoSearchResults
from langchain_community.tools.arxiv import ArxivQueryRun from langchain_community.tools.arxiv import ArxivQueryRun
from langchain_community.utilities import ArxivAPIWrapper, BraveSearchWrapper from langchain_community.utilities import ArxivAPIWrapper, BraveSearchWrapper
from src.config import SearchEngine, SELECTED_SEARCH_ENGINE from src.config import SearchEngine, SELECTED_SEARCH_ENGINE
from src.config import load_yaml_config
from src.tools.tavily_search.tavily_search_results_with_images import ( from src.tools.tavily_search.tavily_search_results_with_images import (
TavilySearchResultsWithImages, TavilySearchResultsWithImages,
) )
@@ -26,39 +25,18 @@ LoggedBraveSearch = create_logged_tool(BraveSearch)
LoggedArxivSearch = create_logged_tool(ArxivQueryRun) LoggedArxivSearch = create_logged_tool(ArxivQueryRun)
def get_search_config():
config = load_yaml_config("conf.yaml")
search_config = config.get("SEARCH_ENGINE", {})
return search_config
# Get the selected search tool # Get the selected search tool
def get_web_search_tool(max_search_results: int): def get_web_search_tool(max_search_results: int):
search_config = get_search_config()
if SELECTED_SEARCH_ENGINE == SearchEngine.TAVILY.value: if SELECTED_SEARCH_ENGINE == SearchEngine.TAVILY.value:
# Only get and apply include/exclude domains for Tavily
include_domains: Optional[List[str]] = search_config.get("include_domains", [])
exclude_domains: Optional[List[str]] = search_config.get("exclude_domains", [])
logger.info(
f"Tavily search configuration loaded: include_domains={include_domains}, exclude_domains={exclude_domains}"
)
return LoggedTavilySearch( return LoggedTavilySearch(
name="web_search", name="web_search",
max_results=max_search_results, max_results=max_search_results,
include_raw_content=True, include_raw_content=True,
include_images=True, include_images=True,
include_image_descriptions=True, include_image_descriptions=True,
include_domains=include_domains,
exclude_domains=exclude_domains,
) )
elif SELECTED_SEARCH_ENGINE == SearchEngine.DUCKDUCKGO.value: elif SELECTED_SEARCH_ENGINE == SearchEngine.DUCKDUCKGO.value:
return LoggedDuckDuckGoSearch( return LoggedDuckDuckGoSearch(name="web_search", max_results=max_search_results)
name="web_search",
num_results=max_search_results,
)
elif SELECTED_SEARCH_ENGINE == SearchEngine.BRAVE_SEARCH.value: elif SELECTED_SEARCH_ENGINE == SearchEngine.BRAVE_SEARCH.value:
return LoggedBraveSearch( return LoggedBraveSearch(
name="web_search", name="web_search",
@@ -78,3 +56,10 @@ def get_web_search_tool(max_search_results: int):
) )
else: else:
raise ValueError(f"Unsupported search engine: {SELECTED_SEARCH_ENGINE}") raise ValueError(f"Unsupported search engine: {SELECTED_SEARCH_ENGINE}")
if __name__ == "__main__":
results = LoggedDuckDuckGoSearch(
name="web_search", max_results=3, output_format="list"
).invoke("cute panda")
print(json.dumps(results, indent=2, ensure_ascii=False))
@@ -1,7 +1,3 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import json import json
from typing import Dict, List, Optional from typing import Dict, List, Optional
@@ -74,7 +70,7 @@ class EnhancedTavilySearchAPIWrapper(OriginalTavilySearchAPIWrapper):
"include_images": include_images, "include_images": include_images,
"include_image_descriptions": include_image_descriptions, "include_image_descriptions": include_image_descriptions,
} }
async with aiohttp.ClientSession(trust_env=True) as session: async with aiohttp.ClientSession() as session:
async with session.post(f"{TAVILY_API_URL}/search", json=params) as res: async with session.post(f"{TAVILY_API_URL}/search", json=params) as res:
if res.status == 200: if res.status == 200:
data = await res.text() data = await res.text()
@@ -111,3 +107,9 @@ class EnhancedTavilySearchAPIWrapper(OriginalTavilySearchAPIWrapper):
} }
clean_results.append(clean_result) clean_results.append(clean_result)
return clean_results return clean_results
if __name__ == "__main__":
wrapper = EnhancedTavilySearchAPIWrapper()
results = wrapper.raw_results("cute panda", include_images=True)
print(json.dumps(results, indent=2, ensure_ascii=False))
@@ -1,6 +1,3 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import json import json
from typing import Dict, List, Optional, Tuple, Union from typing import Dict, List, Optional, Tuple, Union
+2 -3
View File
@@ -102,8 +102,7 @@ class VolcengineTTS:
} }
try: try:
sanitized_text = text.replace("\r\n", "").replace("\n", "") logger.debug(f"Sending TTS request for text: {text[:50]}...")
logger.debug(f"Sending TTS request for text: {sanitized_text[:50]}...")
response = requests.post( response = requests.post(
self.api_url, json.dumps(request_json), headers=self.header self.api_url, json.dumps(request_json), headers=self.header
) )
@@ -129,4 +128,4 @@ class VolcengineTTS:
except Exception as e: except Exception as e:
logger.exception(f"Error in TTS API call: {str(e)}") logger.exception(f"Error in TTS API call: {str(e)}")
return {"success": False, "error": "TTS API call error", "audio_data": None} return {"success": False, "error": str(e), "audio_data": None}
+12 -8
View File
@@ -19,17 +19,21 @@ def repair_json_output(content: str) -> str:
str: Repaired JSON string, or original content if not JSON str: Repaired JSON string, or original content if not JSON
""" """
content = content.strip() content = content.strip()
if content.startswith(("{", "[")) or "```json" in content or "```ts" in content:
try: try:
# If content is wrapped in ```json code block, extract the JSON part
if content.startswith("```json"):
content = content.removeprefix("```json")
if content.startswith("```ts"):
content = content.removeprefix("```ts")
if content.endswith("```"):
content = content.removesuffix("```")
# Try to repair and parse JSON # Try to repair and parse JSON
repaired_content = json_repair.loads(content) repaired_content = json_repair.loads(content)
if not isinstance(repaired_content, dict) and not isinstance( return json.dumps(repaired_content, ensure_ascii=False)
repaired_content, list
):
logger.warning("Repaired content is not a valid JSON object or array.")
return content
content = json.dumps(repaired_content, ensure_ascii=False)
except Exception as e: except Exception as e:
logger.warning(f"JSON repair failed: {e}") logger.warning(f"JSON repair failed: {e}")
return content return content
+1
View File
@@ -1,6 +1,7 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT # SPDX-License-Identifier: MIT
import asyncio
import logging import logging
from src.graph import build_graph from src.graph import build_graph
+1
View File
@@ -1,6 +1,7 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT # SPDX-License-Identifier: MIT
import pytest
from src.crawler import Crawler from src.crawler import Crawler
File diff suppressed because it is too large Load Diff
@@ -1,6 +1,7 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT # SPDX-License-Identifier: MIT
import pytest
from src.tools.python_repl import python_repl_tool from src.tools.python_repl import python_repl_tool
+2 -37
View File
@@ -101,43 +101,8 @@ def test_current_time_format():
messages = apply_prompt_template("coder", test_state) messages = apply_prompt_template("coder", test_state)
system_content = messages[0]["content"] system_content = messages[0]["content"]
# Time format should be like: Mon Jan 01 2024 12:34:56 +0000
time_format = r"\w{3} \w{3} \d{2} \d{4} \d{2}:\d{2}:\d{2}"
assert any( assert any(
line.strip().startswith("CURRENT_TIME:") for line in system_content.split("\n") line.strip().startswith("CURRENT_TIME:") for line in system_content.split("\n")
) )
def test_apply_prompt_template_reporter():
"""Test reporter template rendering with different styles and locale"""
test_state_news = {
"messages": [],
"task": "test reporter task",
"workspace_context": "test reporter context",
"report_style": "news",
"locale": "en-US",
}
messages_news = apply_prompt_template("reporter", test_state_news)
system_content_news = messages_news[0]["content"]
assert "NBC News" in system_content_news
test_state_social_media_en = {
"messages": [],
"task": "test reporter task",
"workspace_context": "test reporter context",
"report_style": "social_media",
"locale": "en-US",
}
messages_default = apply_prompt_template("reporter", test_state_social_media_en)
system_content_default = messages_default[0]["content"]
assert "Twitter/X" in system_content_default
test_state_social_media_cn = {
"messages": [],
"task": "test reporter task",
"workspace_context": "test reporter context",
"report_style": "social_media",
"locale": "zh-CN",
}
messages_cn = apply_prompt_template("reporter", test_state_social_media_cn)
system_content_cn = messages_cn[0]["content"]
assert "小红书" in system_content_cn
+2 -18
View File
@@ -2,7 +2,9 @@
# SPDX-License-Identifier: MIT # SPDX-License-Identifier: MIT
import json import json
import pytest
from unittest.mock import patch, MagicMock from unittest.mock import patch, MagicMock
import uuid
import base64 import base64
from src.tools.tts import VolcengineTTS from src.tools.tts import VolcengineTTS
@@ -227,21 +229,3 @@ class TestVolcengineTTS:
args, kwargs = mock_post.call_args args, kwargs = mock_post.call_args
request_json = json.loads(args[1]) request_json = json.loads(args[1])
assert request_json["user"]["uid"] == str(mock_uuid_value) assert request_json["user"]["uid"] == str(mock_uuid_value)
@patch("src.tools.tts.requests.post")
def test_text_to_speech_request_exception(self, mock_post):
"""Test error handling when requests.post raises an exception."""
# Mock requests.post to raise an exception
mock_post.side_effect = Exception("Network error")
# Create TTS client
tts = VolcengineTTS(
appid="test_appid",
access_token="test_token",
)
# Call the method
result = tts.text_to_speech("Hello, world!")
# Verify the result
assert result["success"] is False
# The TTS error is caught and returned as a string
assert result["error"] == "TTS API call error"
assert result["audio_data"] is None
+5 -7
View File
@@ -1,9 +1,7 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates import pytest
# SPDX-License-Identifier: MIT
import sys import sys
import os import os
from typing import Annotated from typing import Annotated, List, Optional
# Import MessagesState directly from langgraph rather than through our application # Import MessagesState directly from langgraph rather than through our application
from langgraph.graph import MessagesState from langgraph.graph import MessagesState
@@ -16,8 +14,8 @@ class StepType:
class Step: class Step:
def __init__(self, need_search, title, description, step_type): def __init__(self, need_web_search, title, description, step_type):
self.need_search = need_search self.need_web_search = need_web_search
self.title = title self.title = title
self.description = description self.description = description
self.step_type = step_type self.step_type = step_type
@@ -92,7 +90,7 @@ def test_state_initialization():
def test_state_with_custom_values(): def test_state_with_custom_values():
"""Test that State can be initialized with custom values.""" """Test that State can be initialized with custom values."""
test_step = Step( test_step = Step(
need_search=True, need_web_search=True,
title="Test Step", title="Test Step",
description="Step description", description="Step description",
step_type=StepType.RESEARCH, step_type=StepType.RESEARCH,
-90
View File
@@ -1,90 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import sys
import types
from src.config.configuration import Configuration
# Patch sys.path so relative import works
# Patch Resource for import
mock_resource = type("Resource", (), {})
# Patch src.rag.retriever.Resource for import
module_name = "src.rag.retriever"
if module_name not in sys.modules:
retriever_mod = types.ModuleType(module_name)
retriever_mod.Resource = mock_resource
sys.modules[module_name] = retriever_mod
# Relative import of Configuration
def test_default_configuration():
config = Configuration()
assert config.resources == []
assert config.max_plan_iterations == 1
assert config.max_step_num == 3
assert config.max_search_results == 3
assert config.mcp_settings is None
def test_from_runnable_config_with_config_dict(monkeypatch):
config_dict = {
"configurable": {
"max_plan_iterations": 5,
"max_step_num": 7,
"max_search_results": 10,
"mcp_settings": {"foo": "bar"},
}
}
config = Configuration.from_runnable_config(config_dict)
assert config.max_plan_iterations == 5
assert config.max_step_num == 7
assert config.max_search_results == 10
assert config.mcp_settings == {"foo": "bar"}
def test_from_runnable_config_with_env_override(monkeypatch):
monkeypatch.setenv("MAX_PLAN_ITERATIONS", "9")
monkeypatch.setenv("MAX_STEP_NUM", "11")
config_dict = {
"configurable": {
"max_plan_iterations": 2,
"max_step_num": 3,
"max_search_results": 4,
}
}
config = Configuration.from_runnable_config(config_dict)
# Environment variables take precedence and are strings
assert config.max_plan_iterations == "9"
assert config.max_step_num == "11"
assert config.max_search_results == 4 # not overridden
# Clean up
monkeypatch.delenv("MAX_PLAN_ITERATIONS")
monkeypatch.delenv("MAX_STEP_NUM")
def test_from_runnable_config_with_none_and_falsy(monkeypatch):
config_dict = {
"configurable": {
"max_plan_iterations": None,
"max_step_num": 0, # falsy, should be skipped
"max_search_results": "",
}
}
config = Configuration.from_runnable_config(config_dict)
# Should fall back to defaults for skipped/falsy values
assert config.max_plan_iterations == 1
assert config.max_step_num == 3
assert config.max_search_results == 3
def test_from_runnable_config_with_no_config():
config = Configuration.from_runnable_config()
assert config.max_plan_iterations == 1
assert config.max_step_num == 3
assert config.max_search_results == 3
assert config.resources == []
assert config.mcp_settings is None
-81
View File
@@ -1,81 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import os
import tempfile
from src.config.loader import load_yaml_config, process_dict, replace_env_vars
def test_replace_env_vars_with_env(monkeypatch):
monkeypatch.setenv("TEST_ENV", "env_value")
assert replace_env_vars("$TEST_ENV") == "env_value"
def test_replace_env_vars_without_env(monkeypatch):
monkeypatch.delenv("NOT_SET_ENV", raising=False)
assert replace_env_vars("$NOT_SET_ENV") == "NOT_SET_ENV"
def test_replace_env_vars_non_string():
assert replace_env_vars(123) == 123
def test_replace_env_vars_regular_string():
assert replace_env_vars("no_env") == "no_env"
def test_process_dict_nested(monkeypatch):
monkeypatch.setenv("FOO", "bar")
config = {"a": "$FOO", "b": {"c": "$FOO", "d": 42, "e": "$NOT_SET_ENV"}}
processed = process_dict(config)
assert processed["a"] == "bar"
assert processed["b"]["c"] == "bar"
assert processed["b"]["d"] == 42
assert processed["b"]["e"] == "NOT_SET_ENV"
def test_process_dict_empty():
assert process_dict({}) == {}
def test_load_yaml_config_file_not_exist():
assert load_yaml_config("non_existent_file.yaml") == {}
def test_load_yaml_config(monkeypatch):
monkeypatch.setenv("MY_ENV", "my_value")
yaml_content = """
key1: value1
key2: $MY_ENV
nested:
key3: $MY_ENV
key4: 123
"""
with tempfile.NamedTemporaryFile("w+", delete=False) as tmp:
tmp.write(yaml_content)
tmp_path = tmp.name
try:
config = load_yaml_config(tmp_path)
assert config["key1"] == "value1"
assert config["key2"] == "my_value"
assert config["nested"]["key3"] == "my_value"
assert config["nested"]["key4"] == 123
finally:
os.remove(tmp_path)
def test_load_yaml_config_cache(monkeypatch):
monkeypatch.setenv("CACHE_ENV", "cache_value")
yaml_content = "foo: $CACHE_ENV"
with tempfile.NamedTemporaryFile("w+", delete=False) as tmp:
tmp.write(yaml_content)
tmp_path = tmp.name
try:
config1 = load_yaml_config(tmp_path)
config2 = load_yaml_config(tmp_path)
assert config1 is config2 # Should be cached (same object)
assert config1["foo"] == "cache_value"
finally:
os.remove(tmp_path)
-73
View File
@@ -1,73 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
from src.crawler.article import Article
class DummyMarkdownify:
"""A dummy markdownify replacement for patching if needed."""
@staticmethod
def markdownify(html):
return html
def test_to_markdown_includes_title(monkeypatch):
article = Article("Test Title", "<p>Hello <b>world</b>!</p>")
result = article.to_markdown(including_title=True)
assert result.startswith("# Test Title")
assert "Hello" in result
def test_to_markdown_excludes_title():
article = Article("Test Title", "<p>Hello <b>world</b>!</p>")
result = article.to_markdown(including_title=False)
assert not result.startswith("# Test Title")
assert "Hello" in result
def test_to_message_with_text_only():
article = Article("Test Title", "<p>Hello world!</p>")
article.url = "https://example.com/"
result = article.to_message()
assert isinstance(result, list)
assert any(item["type"] == "text" for item in result)
assert all("type" in item for item in result)
def test_to_message_with_image(monkeypatch):
html = '<p>Intro</p><img src="img/pic.png"/>'
article = Article("Title", html)
article.url = "https://host.com/path/"
# The markdownify library will convert <img> to markdown image syntax
result = article.to_message()
# Should have both text and image_url types
types = [item["type"] for item in result]
assert "image_url" in types
assert "text" in types
# Check that the image_url is correctly joined
image_items = [item for item in result if item["type"] == "image_url"]
assert image_items
assert image_items[0]["image_url"]["url"] == "https://host.com/path/img/pic.png"
def test_to_message_multiple_images():
html = '<p>Start</p><img src="a.png"/><p>Mid</p><img src="b.jpg"/>End'
article = Article("Title", html)
article.url = "http://x/"
result = article.to_message()
image_urls = [
item["image_url"]["url"] for item in result if item["type"] == "image_url"
]
assert "http://x/a.png" in image_urls
assert "http://x/b.jpg" in image_urls
text_items = [item for item in result if item["type"] == "text"]
assert any("Start" in item["text"] for item in text_items)
assert any("Mid" in item["text"] for item in text_items)
def test_to_message_handles_empty_html():
article = Article("Empty", "")
article.url = "http://test/"
result = article.to_message()
assert isinstance(result, list)
assert result[0]["type"] == "text"
-70
View File
@@ -1,70 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import src.crawler as crawler_module
def test_crawler_sets_article_url(monkeypatch):
"""Test that the crawler sets the article.url field correctly."""
class DummyArticle:
def __init__(self):
self.url = None
def to_markdown(self):
return "# Dummy"
class DummyJinaClient:
def crawl(self, url, return_format=None):
return "<html>dummy</html>"
class DummyReadabilityExtractor:
def extract_article(self, html):
return DummyArticle()
monkeypatch.setattr("src.crawler.crawler.JinaClient", DummyJinaClient)
monkeypatch.setattr(
"src.crawler.crawler.ReadabilityExtractor", DummyReadabilityExtractor
)
crawler = crawler_module.Crawler()
url = "http://example.com"
article = crawler.crawl(url)
assert article.url == url
assert article.to_markdown() == "# Dummy"
def test_crawler_calls_dependencies(monkeypatch):
"""Test that Crawler calls JinaClient.crawl and ReadabilityExtractor.extract_article."""
calls = {}
class DummyJinaClient:
def crawl(self, url, return_format=None):
calls["jina"] = (url, return_format)
return "<html>dummy</html>"
class DummyReadabilityExtractor:
def extract_article(self, html):
calls["extractor"] = html
class DummyArticle:
url = None
def to_markdown(self):
return "# Dummy"
return DummyArticle()
monkeypatch.setattr("src.crawler.crawler.JinaClient", DummyJinaClient)
monkeypatch.setattr(
"src.crawler.crawler.ReadabilityExtractor", DummyReadabilityExtractor
)
crawler = crawler_module.Crawler()
url = "http://example.com"
crawler.crawl(url)
assert "jina" in calls
assert calls["jina"][0] == url
assert calls["jina"][1] == "html"
assert "extractor" in calls
assert calls["extractor"] == "<html>dummy</html>"
-132
View File
@@ -1,132 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import pytest
from unittest.mock import MagicMock, patch
import importlib
import sys
import src.graph.builder as builder_mod
@pytest.fixture
def mock_state():
class Step:
def __init__(self, execution_res=None, step_type=None):
self.execution_res = execution_res
self.step_type = step_type
class Plan:
def __init__(self, steps):
self.steps = steps
return {
"Step": Step,
"Plan": Plan,
}
def test_continue_to_running_research_team_no_plan(mock_state):
state = {"current_plan": None}
assert builder_mod.continue_to_running_research_team(state) == "planner"
def test_continue_to_running_research_team_no_steps(mock_state):
state = {"current_plan": mock_state["Plan"](steps=[])}
assert builder_mod.continue_to_running_research_team(state) == "planner"
def test_continue_to_running_research_team_all_executed(mock_state):
Step = mock_state["Step"]
Plan = mock_state["Plan"]
steps = [Step(execution_res=True), Step(execution_res=True)]
state = {"current_plan": Plan(steps=steps)}
assert builder_mod.continue_to_running_research_team(state) == "planner"
def test_continue_to_running_research_team_next_researcher(mock_state):
Step = mock_state["Step"]
Plan = mock_state["Plan"]
steps = [
Step(execution_res=True),
Step(execution_res=None, step_type=builder_mod.StepType.RESEARCH),
]
state = {"current_plan": Plan(steps=steps)}
assert builder_mod.continue_to_running_research_team(state) == "researcher"
def test_continue_to_running_research_team_next_coder(mock_state):
Step = mock_state["Step"]
Plan = mock_state["Plan"]
steps = [
Step(execution_res=True),
Step(execution_res=None, step_type=builder_mod.StepType.PROCESSING),
]
state = {"current_plan": Plan(steps=steps)}
assert builder_mod.continue_to_running_research_team(state) == "coder"
def test_continue_to_running_research_team_next_coder_withresult(mock_state):
Step = mock_state["Step"]
Plan = mock_state["Plan"]
steps = [
Step(execution_res=True),
Step(execution_res=True, step_type=builder_mod.StepType.PROCESSING),
]
state = {"current_plan": Plan(steps=steps)}
assert builder_mod.continue_to_running_research_team(state) == "planner"
def test_continue_to_running_research_team_default_planner(mock_state):
Step = mock_state["Step"]
Plan = mock_state["Plan"]
steps = [Step(execution_res=True), Step(execution_res=None, step_type=None)]
state = {"current_plan": Plan(steps=steps)}
assert builder_mod.continue_to_running_research_team(state) == "planner"
@patch("src.graph.builder.StateGraph")
def test_build_base_graph_adds_nodes_and_edges(MockStateGraph):
mock_builder = MagicMock()
MockStateGraph.return_value = mock_builder
builder_mod._build_base_graph()
# Check that all nodes and edges are added
assert mock_builder.add_edge.call_count >= 2
assert mock_builder.add_node.call_count >= 8
mock_builder.add_conditional_edges.assert_called_once()
@patch("src.graph.builder._build_base_graph")
@patch("src.graph.builder.MemorySaver")
def test_build_graph_with_memory_uses_memory(MockMemorySaver, mock_build_base_graph):
mock_builder = MagicMock()
mock_build_base_graph.return_value = mock_builder
mock_memory = MagicMock()
MockMemorySaver.return_value = mock_memory
builder_mod.build_graph_with_memory()
mock_builder.compile.assert_called_once_with(checkpointer=mock_memory)
@patch("src.graph.builder._build_base_graph")
def test_build_graph_without_memory(mock_build_base_graph):
mock_builder = MagicMock()
mock_build_base_graph.return_value = mock_builder
builder_mod.build_graph()
mock_builder.compile.assert_called_once_with()
def test_graph_is_compiled():
# The graph object should be the result of build_graph()
with patch("src.graph.builder._build_base_graph") as mock_base:
mock_builder = MagicMock()
mock_base.return_value = mock_builder
mock_builder.compile.return_value = "compiled_graph"
# reload the module to re-run the graph assignment
importlib.reload(sys.modules["src.graph.builder"])
assert builder_mod.graph is not None
-86
View File
@@ -1,86 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import pytest
from src.llms import llm
class DummyChatOpenAI:
def __init__(self, **kwargs):
self.kwargs = kwargs
def invoke(self, msg):
return f"Echo: {msg}"
@pytest.fixture(autouse=True)
def patch_chat_openai(monkeypatch):
monkeypatch.setattr(llm, "ChatOpenAI", DummyChatOpenAI)
@pytest.fixture
def dummy_conf():
return {
"BASIC_MODEL": {"api_key": "test_key", "base_url": "http://test"},
"REASONING_MODEL": {"api_key": "reason_key"},
"VISION_MODEL": {"api_key": "vision_key"},
}
def test_get_env_llm_conf(monkeypatch):
# Clear any existing environment variables that might interfere
monkeypatch.delenv("BASIC_MODEL__API_KEY", raising=False)
monkeypatch.delenv("BASIC_MODEL__BASE_URL", raising=False)
monkeypatch.delenv("BASIC_MODEL__MODEL", raising=False)
monkeypatch.setenv("BASIC_MODEL__API_KEY", "env_key")
monkeypatch.setenv("BASIC_MODEL__BASE_URL", "http://env")
conf = llm._get_env_llm_conf("basic")
assert conf["api_key"] == "env_key"
assert conf["base_url"] == "http://env"
def test_create_llm_use_conf_merges_env(monkeypatch, dummy_conf):
# Clear any existing environment variables that might interfere
monkeypatch.delenv("BASIC_MODEL__BASE_URL", raising=False)
monkeypatch.delenv("BASIC_MODEL__MODEL", raising=False)
monkeypatch.setenv("BASIC_MODEL__API_KEY", "env_key")
result = llm._create_llm_use_conf("basic", dummy_conf)
assert isinstance(result, DummyChatOpenAI)
assert result.kwargs["api_key"] == "env_key"
assert result.kwargs["base_url"] == "http://test"
def test_create_llm_use_conf_invalid_type(monkeypatch, dummy_conf):
# Clear any existing environment variables that might interfere
monkeypatch.delenv("BASIC_MODEL__API_KEY", raising=False)
monkeypatch.delenv("BASIC_MODEL__BASE_URL", raising=False)
monkeypatch.delenv("BASIC_MODEL__MODEL", raising=False)
with pytest.raises(ValueError):
llm._create_llm_use_conf("unknown", dummy_conf)
def test_create_llm_use_conf_empty_conf(monkeypatch):
# Clear any existing environment variables that might interfere
monkeypatch.delenv("BASIC_MODEL__API_KEY", raising=False)
monkeypatch.delenv("BASIC_MODEL__BASE_URL", raising=False)
monkeypatch.delenv("BASIC_MODEL__MODEL", raising=False)
with pytest.raises(ValueError):
llm._create_llm_use_conf("basic", {})
def test_get_llm_by_type_caches(monkeypatch, dummy_conf):
called = {}
def fake_load_yaml_config(path):
called["called"] = True
return dummy_conf
monkeypatch.setattr(llm, "load_yaml_config", fake_load_yaml_config)
llm._llm_cache.clear()
inst1 = llm.get_llm_by_type("basic")
inst2 = llm.get_llm_by_type("basic")
assert inst1 is inst2
assert called["called"]
-2
View File
@@ -1,2 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
@@ -1,2 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
@@ -1,155 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import pytest
from unittest.mock import patch, MagicMock
from src.prompt_enhancer.graph.builder import build_graph
from src.prompt_enhancer.graph.state import PromptEnhancerState
class TestBuildGraph:
"""Test cases for build_graph function."""
@patch("src.prompt_enhancer.graph.builder.StateGraph")
def test_build_graph_structure(self, mock_state_graph):
"""Test that build_graph creates the correct graph structure."""
mock_builder = MagicMock()
mock_compiled_graph = MagicMock()
mock_state_graph.return_value = mock_builder
mock_builder.compile.return_value = mock_compiled_graph
result = build_graph()
# Verify StateGraph was created with correct state type
mock_state_graph.assert_called_once_with(PromptEnhancerState)
# Verify entry point was set
mock_builder.set_entry_point.assert_called_once_with("enhancer")
# Verify finish point was set
mock_builder.set_finish_point.assert_called_once_with("enhancer")
# Verify graph was compiled
mock_builder.compile.assert_called_once()
# Verify return value
assert result == mock_compiled_graph
@patch("src.prompt_enhancer.graph.builder.StateGraph")
@patch("src.prompt_enhancer.graph.builder.prompt_enhancer_node")
def test_build_graph_node_function(self, mock_enhancer_node, mock_state_graph):
"""Test that the correct node function is added to the graph."""
mock_builder = MagicMock()
mock_compiled_graph = MagicMock()
mock_state_graph.return_value = mock_builder
mock_builder.compile.return_value = mock_compiled_graph
build_graph()
# Verify the correct node function was added
mock_builder.add_node.assert_called_once_with("enhancer", mock_enhancer_node)
def test_build_graph_returns_compiled_graph(self):
"""Test that build_graph returns a compiled graph object."""
with patch("src.prompt_enhancer.graph.builder.StateGraph") as mock_state_graph:
mock_builder = MagicMock()
mock_compiled_graph = MagicMock()
mock_state_graph.return_value = mock_builder
mock_builder.compile.return_value = mock_compiled_graph
result = build_graph()
assert result is mock_compiled_graph
@patch("src.prompt_enhancer.graph.builder.StateGraph")
def test_build_graph_call_sequence(self, mock_state_graph):
"""Test that build_graph calls methods in the correct sequence."""
mock_builder = MagicMock()
mock_compiled_graph = MagicMock()
mock_state_graph.return_value = mock_builder
mock_builder.compile.return_value = mock_compiled_graph
# Track call order
call_order = []
def track_add_node(*args, **kwargs):
call_order.append("add_node")
def track_set_entry_point(*args, **kwargs):
call_order.append("set_entry_point")
def track_set_finish_point(*args, **kwargs):
call_order.append("set_finish_point")
def track_compile(*args, **kwargs):
call_order.append("compile")
return mock_compiled_graph
mock_builder.add_node.side_effect = track_add_node
mock_builder.set_entry_point.side_effect = track_set_entry_point
mock_builder.set_finish_point.side_effect = track_set_finish_point
mock_builder.compile.side_effect = track_compile
build_graph()
# Verify the correct call sequence
expected_order = ["add_node", "set_entry_point", "set_finish_point", "compile"]
assert call_order == expected_order
def test_build_graph_integration(self):
"""Integration test to verify the graph can be built without mocking."""
# This test verifies that all imports and dependencies are correct
try:
graph = build_graph()
assert graph is not None
# The graph should be a compiled LangGraph object
assert hasattr(graph, "invoke") or hasattr(graph, "stream")
except ImportError as e:
pytest.skip(f"Skipping integration test due to missing dependencies: {e}")
except Exception as e:
# If there are configuration issues (like missing LLM config),
# we still consider the test successful if the graph structure is built
if "LLM" in str(e) or "configuration" in str(e).lower():
pytest.skip(
f"Skipping integration test due to configuration issues: {e}"
)
else:
raise
@patch("src.prompt_enhancer.graph.builder.StateGraph")
def test_build_graph_single_node_workflow(self, mock_state_graph):
"""Test that the graph is configured as a single-node workflow."""
mock_builder = MagicMock()
mock_compiled_graph = MagicMock()
mock_state_graph.return_value = mock_builder
mock_builder.compile.return_value = mock_compiled_graph
build_graph()
# Verify only one node is added
assert mock_builder.add_node.call_count == 1
# Verify entry and finish points are the same node
mock_builder.set_entry_point.assert_called_once_with("enhancer")
mock_builder.set_finish_point.assert_called_once_with("enhancer")
@patch("src.prompt_enhancer.graph.builder.StateGraph")
def test_build_graph_state_type(self, mock_state_graph):
"""Test that the graph is initialized with the correct state type."""
mock_builder = MagicMock()
mock_compiled_graph = MagicMock()
mock_state_graph.return_value = mock_builder
mock_builder.compile.return_value = mock_compiled_graph
build_graph()
# Verify StateGraph was initialized with PromptEnhancerState
args, kwargs = mock_state_graph.call_args
assert args[0] == PromptEnhancerState
@@ -1,525 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import pytest
from unittest.mock import patch, MagicMock
from langchain.schema import HumanMessage, SystemMessage
from src.prompt_enhancer.graph.enhancer_node import prompt_enhancer_node
from src.prompt_enhancer.graph.state import PromptEnhancerState
from src.config.report_style import ReportStyle
@pytest.fixture
def mock_llm():
"""Mock LLM that returns a test response."""
llm = MagicMock()
llm.invoke.return_value = MagicMock(
content="""Thoughts: LLM thinks a lot
<enhanced_prompt>
Enhanced test prompt
</enhanced_prompt>
"""
)
return llm
@pytest.fixture
def mock_llm_xml_with_whitespace():
"""Mock LLM that returns XML response with extra whitespace."""
llm = MagicMock()
llm.invoke.return_value = MagicMock(
content="""
Some thoughts here...
<enhanced_prompt>
Enhanced prompt with whitespace
</enhanced_prompt>
Additional content after XML
"""
)
return llm
@pytest.fixture
def mock_llm_xml_multiline():
"""Mock LLM that returns XML response with multiline content."""
llm = MagicMock()
llm.invoke.return_value = MagicMock(
content="""
<enhanced_prompt>
This is a multiline enhanced prompt
that spans multiple lines
and includes various formatting.
It should preserve the structure.
</enhanced_prompt>
"""
)
return llm
@pytest.fixture
def mock_llm_no_xml():
"""Mock LLM that returns response without XML tags."""
llm = MagicMock()
llm.invoke.return_value = MagicMock(
content="Enhanced Prompt: This is an enhanced prompt without XML tags"
)
return llm
@pytest.fixture
def mock_llm_malformed_xml():
"""Mock LLM that returns response with malformed XML."""
llm = MagicMock()
llm.invoke.return_value = MagicMock(
content="""
<enhanced_prompt>
This XML tag is not properly closed
<enhanced_prompt>
"""
)
return llm
@pytest.fixture
def mock_messages():
"""Mock messages returned by apply_prompt_template."""
return [
SystemMessage(content="System prompt template"),
HumanMessage(content="Test human message"),
]
class TestPromptEnhancerNode:
"""Test cases for prompt_enhancer_node function."""
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_basic_prompt_enhancement(
self, mock_get_llm, mock_apply_template, mock_llm, mock_messages
):
"""Test basic prompt enhancement without context or report style."""
mock_get_llm.return_value = mock_llm
mock_apply_template.return_value = mock_messages
state = PromptEnhancerState(prompt="Write about AI")
result = prompt_enhancer_node(state)
# Verify LLM was called
mock_get_llm.assert_called_once_with("basic")
mock_llm.invoke.assert_called_once_with(mock_messages)
# Verify apply_prompt_template was called correctly
mock_apply_template.assert_called_once()
call_args = mock_apply_template.call_args
assert call_args[0][0] == "prompt_enhancer/prompt_enhancer"
assert "messages" in call_args[0][1]
assert "report_style" in call_args[0][1]
# Verify result
assert result == {"output": "Enhanced test prompt"}
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_prompt_enhancement_with_report_style(
self, mock_get_llm, mock_apply_template, mock_llm, mock_messages
):
"""Test prompt enhancement with report style."""
mock_get_llm.return_value = mock_llm
mock_apply_template.return_value = mock_messages
state = PromptEnhancerState(
prompt="Write about AI", report_style=ReportStyle.ACADEMIC
)
result = prompt_enhancer_node(state)
# Verify apply_prompt_template was called with report_style
mock_apply_template.assert_called_once()
call_args = mock_apply_template.call_args
assert call_args[0][0] == "prompt_enhancer/prompt_enhancer"
assert call_args[0][1]["report_style"] == ReportStyle.ACADEMIC
# Verify result
assert result == {"output": "Enhanced test prompt"}
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_prompt_enhancement_with_context(
self, mock_get_llm, mock_apply_template, mock_llm, mock_messages
):
"""Test prompt enhancement with additional context."""
mock_get_llm.return_value = mock_llm
mock_apply_template.return_value = mock_messages
state = PromptEnhancerState(
prompt="Write about AI", context="Focus on machine learning applications"
)
result = prompt_enhancer_node(state)
# Verify apply_prompt_template was called
mock_apply_template.assert_called_once()
call_args = mock_apply_template.call_args
# Check that the context was included in the human message
messages_arg = call_args[0][1]["messages"]
assert len(messages_arg) == 1
human_message = messages_arg[0]
assert isinstance(human_message, HumanMessage)
assert "Focus on machine learning applications" in human_message.content
assert result == {"output": "Enhanced test prompt"}
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_error_handling(
self, mock_get_llm, mock_apply_template, mock_llm, mock_messages
):
"""Test error handling when LLM call fails."""
mock_get_llm.return_value = mock_llm
mock_apply_template.return_value = mock_messages
# Mock LLM to raise an exception
mock_llm.invoke.side_effect = Exception("LLM error")
state = PromptEnhancerState(prompt="Test prompt")
result = prompt_enhancer_node(state)
# Should return original prompt on error
assert result == {"output": "Test prompt"}
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_template_error_handling(
self, mock_get_llm, mock_apply_template, mock_llm, mock_messages
):
"""Test error handling when template application fails."""
mock_get_llm.return_value = mock_llm
# Mock apply_prompt_template to raise an exception
mock_apply_template.side_effect = Exception("Template error")
state = PromptEnhancerState(prompt="Test prompt")
result = prompt_enhancer_node(state)
# Should return original prompt on error
assert result == {"output": "Test prompt"}
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_prefix_removal(
self, mock_get_llm, mock_apply_template, mock_llm, mock_messages
):
"""Test that common prefixes are removed from LLM response."""
mock_get_llm.return_value = mock_llm
mock_apply_template.return_value = mock_messages
# Test different prefixes that should be removed
test_cases = [
"Enhanced Prompt: This is the enhanced prompt",
"Enhanced prompt: This is the enhanced prompt",
"Here's the enhanced prompt: This is the enhanced prompt",
"Here is the enhanced prompt: This is the enhanced prompt",
"**Enhanced Prompt**: This is the enhanced prompt",
"**Enhanced prompt**: This is the enhanced prompt",
]
for response_with_prefix in test_cases:
mock_llm.invoke.return_value = MagicMock(content=response_with_prefix)
state = PromptEnhancerState(prompt="Test prompt")
result = prompt_enhancer_node(state)
assert result == {"output": "This is the enhanced prompt"}
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_whitespace_handling(
self, mock_get_llm, mock_apply_template, mock_llm, mock_messages
):
"""Test that whitespace is properly stripped from LLM response."""
mock_get_llm.return_value = mock_llm
mock_apply_template.return_value = mock_messages
# Mock LLM response with extra whitespace
mock_llm.invoke.return_value = MagicMock(
content=" \n\n Enhanced prompt \n\n "
)
state = PromptEnhancerState(prompt="Test prompt")
result = prompt_enhancer_node(state)
assert result == {"output": "Enhanced prompt"}
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_xml_with_whitespace_handling(
self,
mock_get_llm,
mock_apply_template,
mock_llm_xml_with_whitespace,
mock_messages,
):
"""Test XML extraction with extra whitespace inside tags."""
mock_get_llm.return_value = mock_llm_xml_with_whitespace
mock_apply_template.return_value = mock_messages
state = PromptEnhancerState(prompt="Test prompt")
result = prompt_enhancer_node(state)
assert result == {"output": "Enhanced prompt with whitespace"}
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_xml_multiline_content(
self, mock_get_llm, mock_apply_template, mock_llm_xml_multiline, mock_messages
):
"""Test XML extraction with multiline content."""
mock_get_llm.return_value = mock_llm_xml_multiline
mock_apply_template.return_value = mock_messages
state = PromptEnhancerState(prompt="Test prompt")
result = prompt_enhancer_node(state)
expected_output = """This is a multiline enhanced prompt
that spans multiple lines
and includes various formatting.
It should preserve the structure."""
assert result == {"output": expected_output}
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_fallback_to_prefix_removal(
self, mock_get_llm, mock_apply_template, mock_llm_no_xml, mock_messages
):
"""Test fallback to prefix removal when no XML tags are found."""
mock_get_llm.return_value = mock_llm_no_xml
mock_apply_template.return_value = mock_messages
state = PromptEnhancerState(prompt="Test prompt")
result = prompt_enhancer_node(state)
assert result == {"output": "This is an enhanced prompt without XML tags"}
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_malformed_xml_fallback(
self, mock_get_llm, mock_apply_template, mock_llm_malformed_xml, mock_messages
):
"""Test handling of malformed XML tags."""
mock_get_llm.return_value = mock_llm_malformed_xml
mock_apply_template.return_value = mock_messages
state = PromptEnhancerState(prompt="Test prompt")
result = prompt_enhancer_node(state)
# Should fall back to using the entire content since XML is malformed
expected_content = """<enhanced_prompt>
This XML tag is not properly closed
<enhanced_prompt>"""
assert result == {"output": expected_content}
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_case_sensitive_prefix_removal(
self, mock_get_llm, mock_apply_template, mock_llm, mock_messages
):
"""Test that prefix removal is case-sensitive."""
mock_get_llm.return_value = mock_llm
mock_apply_template.return_value = mock_messages
# Test case variations that should NOT be removed
test_cases = [
"ENHANCED PROMPT: This should not be removed",
"enhanced prompt: This should not be removed",
"Enhanced Prompt This should not be removed", # Missing colon
"Enhanced Prompt :: This should not be removed", # Double colon
]
for response_content in test_cases:
mock_llm.invoke.return_value = MagicMock(content=response_content)
state = PromptEnhancerState(prompt="Test prompt")
result = prompt_enhancer_node(state)
# Should return the full content since prefix doesn't match exactly
assert result == {"output": response_content}
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_prefix_with_extra_whitespace(
self, mock_get_llm, mock_apply_template, mock_llm, mock_messages
):
"""Test prefix removal with extra whitespace after colon."""
mock_get_llm.return_value = mock_llm
mock_apply_template.return_value = mock_messages
test_cases = [
("Enhanced Prompt: This has extra spaces", "This has extra spaces"),
("Enhanced prompt:\t\tThis has tabs", "This has tabs"),
("Here's the enhanced prompt:\n\nThis has newlines", "This has newlines"),
]
for response_content, expected_output in test_cases:
mock_llm.invoke.return_value = MagicMock(content=response_content)
state = PromptEnhancerState(prompt="Test prompt")
result = prompt_enhancer_node(state)
assert result == {"output": expected_output}
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_xml_with_special_characters(
self, mock_get_llm, mock_apply_template, mock_llm, mock_messages
):
"""Test XML extraction with special characters and symbols."""
mock_get_llm.return_value = mock_llm
mock_apply_template.return_value = mock_messages
special_content = """<enhanced_prompt>
Enhanced prompt with special chars: @#$%^&*()
Unicode: 🚀 ✨ 💡
Quotes: "double" and 'single'
Backslashes: \\n \\t \\r
</enhanced_prompt>"""
mock_llm.invoke.return_value = MagicMock(content=special_content)
state = PromptEnhancerState(prompt="Test prompt")
result = prompt_enhancer_node(state)
expected_output = """Enhanced prompt with special chars: @#$%^&*()
Unicode: 🚀 ✨ 💡
Quotes: "double" and 'single'
Backslashes: \\n \\t \\r"""
assert result == {"output": expected_output}
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_very_long_response(
self, mock_get_llm, mock_apply_template, mock_llm, mock_messages
):
"""Test handling of very long LLM responses."""
mock_get_llm.return_value = mock_llm
mock_apply_template.return_value = mock_messages
# Create a very long response
long_content = "This is a very long enhanced prompt. " * 100
xml_response = f"<enhanced_prompt>\n{long_content}\n</enhanced_prompt>"
mock_llm.invoke.return_value = MagicMock(content=xml_response)
state = PromptEnhancerState(prompt="Test prompt")
result = prompt_enhancer_node(state)
assert result == {"output": long_content.strip()}
assert len(result["output"]) > 1000 # Verify it's actually long
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_empty_response_content(
self, mock_get_llm, mock_apply_template, mock_llm, mock_messages
):
"""Test handling of empty response content."""
mock_get_llm.return_value = mock_llm
mock_apply_template.return_value = mock_messages
mock_llm.invoke.return_value = MagicMock(content="")
state = PromptEnhancerState(prompt="Test prompt")
result = prompt_enhancer_node(state)
assert result == {"output": ""}
@patch("src.prompt_enhancer.graph.enhancer_node.apply_prompt_template")
@patch("src.prompt_enhancer.graph.enhancer_node.get_llm_by_type")
@patch(
"src.prompt_enhancer.graph.enhancer_node.AGENT_LLM_MAP",
{"prompt_enhancer": "basic"},
)
def test_only_whitespace_response(
self, mock_get_llm, mock_apply_template, mock_llm, mock_messages
):
"""Test handling of response with only whitespace."""
mock_get_llm.return_value = mock_llm
mock_apply_template.return_value = mock_messages
mock_llm.invoke.return_value = MagicMock(content=" \n\n\t\t ")
state = PromptEnhancerState(prompt="Test prompt")
result = prompt_enhancer_node(state)
assert result == {"output": ""}
@@ -1,107 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
from src.prompt_enhancer.graph.state import PromptEnhancerState
from src.config.report_style import ReportStyle
def test_prompt_enhancer_state_creation():
"""Test that PromptEnhancerState can be created with required fields."""
state = PromptEnhancerState(
prompt="Test prompt", context=None, report_style=None, output=None
)
assert state["prompt"] == "Test prompt"
assert state["context"] is None
assert state["report_style"] is None
assert state["output"] is None
def test_prompt_enhancer_state_with_all_fields():
"""Test PromptEnhancerState with all fields populated."""
state = PromptEnhancerState(
prompt="Write about AI",
context="Additional context about AI research",
report_style=ReportStyle.ACADEMIC,
output="Enhanced prompt about AI research",
)
assert state["prompt"] == "Write about AI"
assert state["context"] == "Additional context about AI research"
assert state["report_style"] == ReportStyle.ACADEMIC
assert state["output"] == "Enhanced prompt about AI research"
def test_prompt_enhancer_state_minimal():
"""Test PromptEnhancerState with only required prompt field."""
state = PromptEnhancerState(prompt="Minimal prompt")
assert state["prompt"] == "Minimal prompt"
# Optional fields should not be present if not specified
assert "context" not in state
assert "report_style" not in state
assert "output" not in state
def test_prompt_enhancer_state_with_different_report_styles():
"""Test PromptEnhancerState with different ReportStyle values."""
styles = [
ReportStyle.ACADEMIC,
ReportStyle.POPULAR_SCIENCE,
ReportStyle.NEWS,
ReportStyle.SOCIAL_MEDIA,
]
for style in styles:
state = PromptEnhancerState(prompt="Test prompt", report_style=style)
assert state["report_style"] == style
def test_prompt_enhancer_state_update():
"""Test updating PromptEnhancerState fields."""
state = PromptEnhancerState(prompt="Original prompt")
# Update with new fields
state.update(
{
"context": "New context",
"report_style": ReportStyle.NEWS,
"output": "Enhanced output",
}
)
assert state["prompt"] == "Original prompt"
assert state["context"] == "New context"
assert state["report_style"] == ReportStyle.NEWS
assert state["output"] == "Enhanced output"
def test_prompt_enhancer_state_get_method():
"""Test using get() method on PromptEnhancerState."""
state = PromptEnhancerState(prompt="Test prompt", report_style=ReportStyle.ACADEMIC)
# Test get with existing keys
assert state.get("prompt") == "Test prompt"
assert state.get("report_style") == ReportStyle.ACADEMIC
# Test get with non-existing keys
assert state.get("context") is None
assert state.get("output") is None
assert state.get("nonexistent", "default") == "default"
def test_prompt_enhancer_state_type_annotations():
"""Test that the state accepts correct types."""
# This test ensures the TypedDict structure is working correctly
state = PromptEnhancerState(
prompt="Test prompt",
context="Test context",
report_style=ReportStyle.POPULAR_SCIENCE,
output="Test output",
)
# Verify types
assert isinstance(state["prompt"], str)
assert isinstance(state["context"], str)
assert isinstance(state["report_style"], ReportStyle)
assert isinstance(state["output"], str)
-155
View File
@@ -1,155 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import pytest
from unittest.mock import patch, MagicMock
from src.rag.ragflow import RAGFlowProvider, parse_uri
# Dummy classes to mock dependencies
class DummyResource:
def __init__(self, uri, title="", description=""):
self.uri = uri
self.title = title
self.description = description
class DummyChunk:
def __init__(self, content, similarity):
self.content = content
self.similarity = similarity
class DummyDocument:
def __init__(self, id, title, chunks=None):
self.id = id
self.title = title
self.chunks = chunks or []
# Patch imports in ragflow.py to use dummy classes
@pytest.fixture(autouse=True)
def patch_imports(monkeypatch):
import src.rag.ragflow as ragflow
ragflow.Resource = DummyResource
ragflow.Chunk = DummyChunk
ragflow.Document = DummyDocument
yield
def test_parse_uri_valid():
uri = "rag://dataset/123#abc"
dataset_id, document_id = parse_uri(uri)
assert dataset_id == "123"
assert document_id == "abc"
def test_parse_uri_invalid():
with pytest.raises(ValueError):
parse_uri("http://dataset/123#abc")
def test_init_env_vars(monkeypatch):
monkeypatch.setenv("RAGFLOW_API_URL", "http://api")
monkeypatch.setenv("RAGFLOW_API_KEY", "key")
monkeypatch.delenv("RAGFLOW_PAGE_SIZE", raising=False)
provider = RAGFlowProvider()
assert provider.api_url == "http://api"
assert provider.api_key == "key"
assert provider.page_size == 10
def test_init_page_size(monkeypatch):
monkeypatch.setenv("RAGFLOW_API_URL", "http://api")
monkeypatch.setenv("RAGFLOW_API_KEY", "key")
monkeypatch.setenv("RAGFLOW_PAGE_SIZE", "5")
provider = RAGFlowProvider()
assert provider.page_size == 5
def test_init_missing_env(monkeypatch):
monkeypatch.delenv("RAGFLOW_API_URL", raising=False)
monkeypatch.setenv("RAGFLOW_API_KEY", "key")
with pytest.raises(ValueError):
RAGFlowProvider()
monkeypatch.setenv("RAGFLOW_API_URL", "http://api")
monkeypatch.delenv("RAGFLOW_API_KEY", raising=False)
with pytest.raises(ValueError):
RAGFlowProvider()
@patch("src.rag.ragflow.requests.post")
def test_query_relevant_documents_success(mock_post, monkeypatch):
monkeypatch.setenv("RAGFLOW_API_URL", "http://api")
monkeypatch.setenv("RAGFLOW_API_KEY", "key")
provider = RAGFlowProvider()
resource = DummyResource("rag://dataset/123#doc456")
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"data": {
"doc_aggs": [{"doc_id": "doc456", "doc_name": "Doc Title"}],
"chunks": [
{"document_id": "doc456", "content": "chunk text", "similarity": 0.9}
],
}
}
mock_post.return_value = mock_response
docs = provider.query_relevant_documents("query", [resource])
assert len(docs) == 1
assert docs[0].id == "doc456"
assert docs[0].title == "Doc Title"
assert len(docs[0].chunks) == 1
assert docs[0].chunks[0].content == "chunk text"
assert docs[0].chunks[0].similarity == 0.9
@patch("src.rag.ragflow.requests.post")
def test_query_relevant_documents_error(mock_post, monkeypatch):
monkeypatch.setenv("RAGFLOW_API_URL", "http://api")
monkeypatch.setenv("RAGFLOW_API_KEY", "key")
provider = RAGFlowProvider()
mock_response = MagicMock()
mock_response.status_code = 400
mock_response.text = "error"
mock_post.return_value = mock_response
with pytest.raises(Exception):
provider.query_relevant_documents("query", [])
@patch("src.rag.ragflow.requests.get")
def test_list_resources_success(mock_get, monkeypatch):
monkeypatch.setenv("RAGFLOW_API_URL", "http://api")
monkeypatch.setenv("RAGFLOW_API_KEY", "key")
provider = RAGFlowProvider()
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {
"data": [
{"id": "123", "name": "Dataset1", "description": "desc1"},
{"id": "456", "name": "Dataset2", "description": "desc2"},
]
}
mock_get.return_value = mock_response
resources = provider.list_resources()
assert len(resources) == 2
assert resources[0].uri == "rag://dataset/123"
assert resources[0].title == "Dataset1"
assert resources[0].description == "desc1"
assert resources[1].uri == "rag://dataset/456"
assert resources[1].title == "Dataset2"
assert resources[1].description == "desc2"
@patch("src.rag.ragflow.requests.get")
def test_list_resources_error(mock_get, monkeypatch):
monkeypatch.setenv("RAGFLOW_API_URL", "http://api")
monkeypatch.setenv("RAGFLOW_API_KEY", "key")
provider = RAGFlowProvider()
mock_response = MagicMock()
mock_response.status_code = 500
mock_response.text = "fail"
mock_get.return_value = mock_response
with pytest.raises(Exception):
provider.list_resources()
-72
View File
@@ -1,72 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import pytest
from src.rag.retriever import Chunk, Document, Resource, Retriever
def test_chunk_init():
chunk = Chunk(content="test content", similarity=0.9)
assert chunk.content == "test content"
assert chunk.similarity == 0.9
def test_document_init_and_to_dict():
chunk1 = Chunk(content="chunk1", similarity=0.8)
chunk2 = Chunk(content="chunk2", similarity=0.7)
doc = Document(
id="doc1", url="http://example.com", title="Title", chunks=[chunk1, chunk2]
)
assert doc.id == "doc1"
assert doc.url == "http://example.com"
assert doc.title == "Title"
assert doc.chunks == [chunk1, chunk2]
d = doc.to_dict()
assert d["id"] == "doc1"
assert d["content"] == "chunk1\n\nchunk2"
assert d["url"] == "http://example.com"
assert d["title"] == "Title"
def test_document_to_dict_optional_fields():
chunk = Chunk(content="only chunk", similarity=1.0)
doc = Document(id="doc2", chunks=[chunk])
d = doc.to_dict()
assert d["id"] == "doc2"
assert d["content"] == "only chunk"
assert "url" not in d
assert "title" not in d
def test_resource_model():
resource = Resource(uri="uri1", title="Resource Title")
assert resource.uri == "uri1"
assert resource.title == "Resource Title"
assert resource.description == ""
def test_resource_model_with_description():
resource = Resource(uri="uri2", title="Resource2", description="desc")
assert resource.description == "desc"
def test_retriever_abstract_methods():
class DummyRetriever(Retriever):
def list_resources(self, query=None):
return [Resource(uri="uri", title="title")]
def query_relevant_documents(self, query, resources=[]):
return [Document(id="id", chunks=[])]
retriever = DummyRetriever()
resources = retriever.list_resources()
assert isinstance(resources, list)
assert isinstance(resources[0], Resource)
docs = retriever.query_relevant_documents("query", resources)
assert isinstance(docs, list)
assert isinstance(docs[0], Document)
def test_retriever_cannot_instantiate():
with pytest.raises(TypeError):
Retriever()
@@ -1,503 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import os
import pytest
import json
from unittest.mock import patch, MagicMock
from src.rag.vikingdb_knowledge_base import VikingDBKnowledgeBaseProvider, parse_uri
# Dummy classes to mock dependencies
class MockResource:
def __init__(self, uri, title="", description=""):
self.uri = uri
self.title = title
self.description = description
class MockChunk:
def __init__(self, content, similarity):
self.content = content
self.similarity = similarity
class MockDocument:
def __init__(self, id, title, chunks=None):
self.id = id
self.title = title
self.chunks = chunks or []
# Patch the imports to use mock classes
@pytest.fixture(autouse=True)
def patch_imports():
with (
patch("src.rag.vikingdb_knowledge_base.Resource", MockResource),
patch("src.rag.vikingdb_knowledge_base.Chunk", MockChunk),
patch("src.rag.vikingdb_knowledge_base.Document", MockDocument),
):
yield
@pytest.fixture
def env_vars():
"""Fixture to set up environment variables"""
with patch.dict(
os.environ,
{
"VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com",
"VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak",
"VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk",
"VIKINGDB_KNOWLEDGE_BASE_RETRIEVAL_SIZE": "10",
},
):
yield
class TestParseUri:
def test_parse_uri_valid_with_fragment(self):
"""Test parsing valid URI with fragment"""
uri = "rag://dataset/123#doc456"
resource_id, document_id = parse_uri(uri)
assert resource_id == "123"
assert document_id == "doc456"
def test_parse_uri_valid_without_fragment(self):
"""Test parsing valid URI without fragment"""
uri = "rag://dataset/123"
resource_id, document_id = parse_uri(uri)
assert resource_id == "123"
assert document_id == ""
def test_parse_uri_invalid_scheme(self):
"""Test parsing URI with invalid scheme"""
with pytest.raises(ValueError, match="Invalid URI"):
parse_uri("http://dataset/123#abc")
def test_parse_uri_malformed(self):
"""Test parsing malformed URI"""
with pytest.raises(ValueError, match="Invalid URI"):
parse_uri("invalid_uri")
class TestVikingDBKnowledgeBaseProviderInit:
def test_init_success_with_all_env_vars(self, env_vars):
"""Test successful initialization with all environment variables"""
provider = VikingDBKnowledgeBaseProvider()
assert provider.api_url == "api-test.example.com"
assert provider.api_ak == "test_ak"
assert provider.api_sk == "test_sk"
assert provider.retrieval_size == 10
def test_init_success_without_retrieval_size(self):
"""Test initialization without VIKINGDB_KNOWLEDGE_BASE_RETRIEVAL_SIZE (should use default)"""
with patch.dict(
os.environ,
{
"VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com",
"VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak",
"VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk",
},
clear=True,
):
provider = VikingDBKnowledgeBaseProvider()
assert provider.retrieval_size == 10
def test_init_custom_retrieval_size(self):
"""Test initialization with custom retrieval size"""
with patch.dict(
os.environ,
{
"VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com",
"VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak",
"VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk",
"VIKINGDB_KNOWLEDGE_BASE_RETRIEVAL_SIZE": "5",
},
):
provider = VikingDBKnowledgeBaseProvider()
assert provider.retrieval_size == 5
def test_init_missing_api_url(self):
"""Test initialization fails when API URL is missing"""
with patch.dict(
os.environ,
{
"VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak",
"VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk",
},
clear=True,
):
with pytest.raises(
ValueError, match="VIKINGDB_KNOWLEDGE_BASE_API_URL is not set"
):
VikingDBKnowledgeBaseProvider()
def test_init_missing_api_ak(self):
"""Test initialization fails when API AK is missing"""
with patch.dict(
os.environ,
{
"VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com",
"VIKINGDB_KNOWLEDGE_BASE_API_SK": "test_sk",
},
clear=True,
):
with pytest.raises(
ValueError, match="VIKINGDB_KNOWLEDGE_BASE_API_AK is not set"
):
VikingDBKnowledgeBaseProvider()
def test_init_missing_api_sk(self):
"""Test initialization fails when API SK is missing"""
with patch.dict(
os.environ,
{
"VIKINGDB_KNOWLEDGE_BASE_API_URL": "api-test.example.com",
"VIKINGDB_KNOWLEDGE_BASE_API_AK": "test_ak",
},
clear=True,
):
with pytest.raises(
ValueError, match="VIKINGDB_KNOWLEDGE_BASE_API_SK is not set"
):
VikingDBKnowledgeBaseProvider()
class TestVikingDBKnowledgeBaseProviderPrepareRequest:
@pytest.fixture
def provider(self, env_vars):
return VikingDBKnowledgeBaseProvider()
def test_prepare_request_basic(self, provider):
"""Test basic request preparation"""
with (
patch("src.rag.vikingdb_knowledge_base.Request") as mock_request,
patch("src.rag.vikingdb_knowledge_base.Credentials") as _mock_credentials,
patch("src.rag.vikingdb_knowledge_base.SignerV4.sign") as _mock_sign,
):
mock_req_instance = MagicMock()
mock_request.return_value = mock_req_instance
result = provider.prepare_request("POST", "/test/path")
assert result == mock_req_instance
mock_req_instance.set_shema.assert_called_once_with("https")
mock_req_instance.set_method.assert_called_once_with("POST")
mock_req_instance.set_path.assert_called_once_with("/test/path")
def test_prepare_request_with_params(self, provider):
"""Test request preparation with parameters"""
with (
patch("src.rag.vikingdb_knowledge_base.Request") as mock_request,
patch("src.rag.vikingdb_knowledge_base.Credentials"),
patch("src.rag.vikingdb_knowledge_base.SignerV4.sign"),
):
mock_req_instance = MagicMock()
mock_request.return_value = mock_req_instance
params = {"key": "value", "number": 123, "boolean": True}
provider.prepare_request("GET", "/test", params=params)
expected_params = {"key": "value", "number": "123", "boolean": "True"}
mock_req_instance.set_query.assert_called_once_with(expected_params)
def test_prepare_request_with_data(self, provider):
"""Test request preparation with data"""
with (
patch("src.rag.vikingdb_knowledge_base.Request") as mock_request,
patch("src.rag.vikingdb_knowledge_base.Credentials"),
patch("src.rag.vikingdb_knowledge_base.SignerV4.sign"),
):
mock_req_instance = MagicMock()
mock_request.return_value = mock_req_instance
data = {"test": "data"}
provider.prepare_request("POST", "/test", data=data)
mock_req_instance.set_body.assert_called_once_with(json.dumps(data))
class TestVikingDBKnowledgeBaseProviderQueryRelevantDocuments:
@pytest.fixture
def provider(self, env_vars):
return VikingDBKnowledgeBaseProvider()
def test_query_relevant_documents_empty_resources(self, provider):
"""Test querying with empty resources list"""
result = provider.query_relevant_documents("test query", [])
assert result == []
@patch("src.rag.vikingdb_knowledge_base.requests.request")
def test_query_relevant_documents_success(self, mock_request, provider):
"""Test successful document query"""
# Mock response
mock_response = MagicMock()
mock_response.text = json.dumps(
{
"code": 0,
"data": {
"result_list": [
{
"doc_info": {
"doc_id": "doc123",
"doc_name": "Test Document",
},
"content": "Test content",
"score": 0.95,
}
]
},
}
)
mock_request.return_value = mock_response
# Mock prepare_request
with patch.object(provider, "prepare_request") as mock_prepare:
mock_req = MagicMock()
mock_req.method = "POST"
mock_req.path = "/api/knowledge/collection/search_knowledge"
mock_req.headers = {}
mock_req.body = "{}"
mock_prepare.return_value = mock_req
resources = [MockResource("rag://dataset/123")]
result = provider.query_relevant_documents("test query", resources)
assert len(result) == 1
assert result[0].id == "doc123"
assert result[0].title == "Test Document"
assert len(result[0].chunks) == 1
assert result[0].chunks[0].content == "Test content"
assert result[0].chunks[0].similarity == 0.95
@patch("src.rag.vikingdb_knowledge_base.requests.request")
def test_query_relevant_documents_with_document_filter(
self, mock_request, provider
):
"""Test document query with document ID filter"""
mock_response = MagicMock()
mock_response.text = json.dumps({"code": 0, "data": {"result_list": []}})
mock_request.return_value = mock_response
with patch.object(provider, "prepare_request") as mock_prepare:
mock_req = MagicMock()
mock_prepare.return_value = mock_req
resources = [MockResource("rag://dataset/123#doc456")]
provider.query_relevant_documents("test query", resources)
# Verify that query_param with doc_filter was included in the request
call_args = mock_prepare.call_args
request_data = call_args[1]["data"]
assert "query_param" in request_data
assert "doc_filter" in request_data["query_param"]
doc_filter = request_data["query_param"]["doc_filter"]
assert doc_filter["op"] == "must"
assert doc_filter["field"] == "doc_id"
assert doc_filter["conds"] == ["doc456"]
@patch("src.rag.vikingdb_knowledge_base.requests.request")
def test_query_relevant_documents_api_error(self, mock_request, provider):
"""Test handling of API error response"""
mock_response = MagicMock()
mock_response.text = json.dumps({"code": 1, "message": "API Error"})
mock_request.return_value = mock_response
with patch.object(provider, "prepare_request"):
resources = [MockResource("rag://dataset/123")]
with pytest.raises(
ValueError, match="Failed to query documents from resource: API Error"
):
provider.query_relevant_documents("test query", resources)
@patch("src.rag.vikingdb_knowledge_base.requests.request")
def test_query_relevant_documents_json_decode_error(self, mock_request, provider):
"""Test handling of JSON decode error"""
mock_response = MagicMock()
mock_response.text = "invalid json"
mock_request.return_value = mock_response
with patch.object(provider, "prepare_request"):
resources = [MockResource("rag://dataset/123")]
with pytest.raises(ValueError, match="Failed to parse JSON response"):
provider.query_relevant_documents("test query", resources)
@patch("src.rag.vikingdb_knowledge_base.requests.request")
def test_query_relevant_documents_multiple_resources(self, mock_request, provider):
"""Test querying multiple resources and merging results"""
# Mock responses for different resources
responses = [
json.dumps(
{
"code": 0,
"data": {
"result_list": [
{
"doc_info": {
"doc_id": "doc1",
"doc_name": "Document 1",
},
"content": "Content 1",
"score": 0.9,
}
]
},
}
),
json.dumps(
{
"code": 0,
"data": {
"result_list": [
{
"doc_info": {
"doc_id": "doc1",
"doc_name": "Document 1",
},
"content": "Content 2",
"score": 0.8,
},
{
"doc_info": {
"doc_id": "doc2",
"doc_name": "Document 2",
},
"content": "Content 3",
"score": 0.7,
},
]
},
}
),
]
mock_request.side_effect = [MagicMock(text=resp) for resp in responses]
with patch.object(provider, "prepare_request"):
resources = [
MockResource("rag://dataset/123"),
MockResource("rag://dataset/456"),
]
result = provider.query_relevant_documents("test query", resources)
# Should have 2 documents: doc1 (with 2 chunks) and doc2 (with 1 chunk)
assert len(result) == 2
doc1 = next(doc for doc in result if doc.id == "doc1")
doc2 = next(doc for doc in result if doc.id == "doc2")
assert len(doc1.chunks) == 2
assert len(doc2.chunks) == 1
class TestVikingDBKnowledgeBaseProviderListResources:
@pytest.fixture
def provider(self, env_vars):
return VikingDBKnowledgeBaseProvider()
@patch("src.rag.vikingdb_knowledge_base.requests.request")
def test_list_resources_success(self, mock_request, provider):
"""Test successful resource listing"""
mock_response = MagicMock()
mock_response.text = json.dumps(
{
"code": 0,
"data": {
"collection_list": [
{
"resource_id": "123",
"collection_name": "Dataset 1",
"description": "Description 1",
},
{
"resource_id": "456",
"collection_name": "Dataset 2",
"description": "Description 2",
},
]
},
}
)
mock_request.return_value = mock_response
with patch.object(provider, "prepare_request") as mock_prepare:
mock_req = MagicMock()
mock_prepare.return_value = mock_req
result = provider.list_resources()
assert len(result) == 2
assert result[0].uri == "rag://dataset/123"
assert result[0].title == "Dataset 1"
assert result[0].description == "Description 1"
assert result[1].uri == "rag://dataset/456"
assert result[1].title == "Dataset 2"
assert result[1].description == "Description 2"
@patch("src.rag.vikingdb_knowledge_base.requests.request")
def test_list_resources_with_query_filter(self, mock_request, provider):
"""Test resource listing with query filter"""
mock_response = MagicMock()
mock_response.text = json.dumps(
{
"code": 0,
"data": {
"collection_list": [
{
"resource_id": "123",
"collection_name": "Test Dataset",
"description": "Description",
},
{
"resource_id": "456",
"collection_name": "Other Dataset",
"description": "Description",
},
]
},
}
)
mock_request.return_value = mock_response
with patch.object(provider, "prepare_request"):
result = provider.list_resources("test")
# Should only return the dataset with "test" in the name
assert len(result) == 1
assert result[0].title == "Test Dataset"
@patch("src.rag.vikingdb_knowledge_base.requests.request")
def test_list_resources_api_error(self, mock_request, provider):
"""Test handling of API error in list_resources"""
mock_response = MagicMock()
mock_response.text = json.dumps({"code": 1, "message": "API Error"})
mock_request.return_value = mock_response
with patch.object(provider, "prepare_request"):
with pytest.raises(Exception, match="Failed to list resources: API Error"):
provider.list_resources()
@patch("src.rag.vikingdb_knowledge_base.requests.request")
def test_list_resources_json_decode_error(self, mock_request, provider):
"""Test handling of JSON decode error in list_resources"""
mock_response = MagicMock()
mock_response.text = "invalid json"
mock_request.return_value = mock_response
with patch.object(provider, "prepare_request"):
with pytest.raises(ValueError, match="Failed to parse JSON response"):
provider.list_resources()
@patch("src.rag.vikingdb_knowledge_base.requests.request")
def test_list_resources_empty_response(self, mock_request, provider):
"""Test handling of empty response"""
mock_response = MagicMock()
mock_response.text = json.dumps({"code": 0, "data": {"collection_list": []}})
mock_request.return_value = mock_response
with patch.object(provider, "prepare_request"):
result = provider.list_resources()
assert result == []
-722
View File
@@ -1,722 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import base64
import os
from unittest.mock import MagicMock, patch, mock_open
import pytest
from fastapi.testclient import TestClient
from fastapi import HTTPException
from src.server.app import app, _make_event, _astream_workflow_generator
from src.config.report_style import ReportStyle
from langgraph.types import Command
from langchain_core.messages import ToolMessage
from langchain_core.messages import AIMessageChunk
@pytest.fixture
def client():
return TestClient(app)
class TestMakeEvent:
def test_make_event_with_content(self):
event_type = "message_chunk"
data = {"content": "Hello", "role": "assistant"}
result = _make_event(event_type, data)
expected = (
'event: message_chunk\ndata: {"content": "Hello", "role": "assistant"}\n\n'
)
assert result == expected
def test_make_event_with_empty_content(self):
event_type = "message_chunk"
data = {"content": "", "role": "assistant"}
result = _make_event(event_type, data)
expected = 'event: message_chunk\ndata: {"role": "assistant"}\n\n'
assert result == expected
def test_make_event_without_content(self):
event_type = "tool_calls"
data = {"role": "assistant", "tool_calls": []}
result = _make_event(event_type, data)
expected = (
'event: tool_calls\ndata: {"role": "assistant", "tool_calls": []}\n\n'
)
assert result == expected
class TestTTSEndpoint:
@patch.dict(
os.environ,
{
"VOLCENGINE_TTS_APPID": "test_app_id",
"VOLCENGINE_TTS_ACCESS_TOKEN": "test_token",
"VOLCENGINE_TTS_CLUSTER": "test_cluster",
"VOLCENGINE_TTS_VOICE_TYPE": "test_voice",
},
)
@patch("src.server.app.VolcengineTTS")
def test_tts_success(self, mock_tts_class, client):
mock_tts_instance = MagicMock()
mock_tts_class.return_value = mock_tts_instance
# Mock successful TTS response
audio_data_b64 = base64.b64encode(b"fake_audio_data").decode()
mock_tts_instance.text_to_speech.return_value = {
"success": True,
"audio_data": audio_data_b64,
}
request_data = {
"text": "Hello world",
"encoding": "mp3",
"speed_ratio": 1.0,
"volume_ratio": 1.0,
"pitch_ratio": 1.0,
"text_type": "plain",
"with_frontend": True,
"frontend_type": "unitTson",
}
response = client.post("/api/tts", json=request_data)
assert response.status_code == 200
assert response.headers["content-type"] == "audio/mp3"
assert b"fake_audio_data" in response.content
@patch.dict(os.environ, {}, clear=True)
def test_tts_missing_app_id(self, client):
request_data = {"text": "Hello world", "encoding": "mp3"}
response = client.post("/api/tts", json=request_data)
assert response.status_code == 400
assert "VOLCENGINE_TTS_APPID is not set" in response.json()["detail"]
@patch.dict(
os.environ,
{"VOLCENGINE_TTS_APPID": "test_app_id", "VOLCENGINE_TTS_ACCESS_TOKEN": ""},
)
def test_tts_missing_access_token(self, client):
request_data = {"text": "Hello world", "encoding": "mp3"}
response = client.post("/api/tts", json=request_data)
assert response.status_code == 400
assert "VOLCENGINE_TTS_ACCESS_TOKEN is not set" in response.json()["detail"]
@patch.dict(
os.environ,
{
"VOLCENGINE_TTS_APPID": "test_app_id",
"VOLCENGINE_TTS_ACCESS_TOKEN": "test_token",
},
)
@patch("src.server.app.VolcengineTTS")
def test_tts_api_error(self, mock_tts_class, client):
mock_tts_instance = MagicMock()
mock_tts_class.return_value = mock_tts_instance
# Mock TTS error response
mock_tts_instance.text_to_speech.return_value = {
"success": False,
"error": "TTS API error",
}
request_data = {"text": "Hello world", "encoding": "mp3"}
response = client.post("/api/tts", json=request_data)
assert response.status_code == 500
assert "Internal Server Error" in response.json()["detail"]
@pytest.mark.skip(reason="TTS server exception is catched")
@patch("src.server.app.VolcengineTTS")
def test_tts_api_exception(self, mock_tts_class, client):
mock_tts_instance = MagicMock()
mock_tts_class.return_value = mock_tts_instance
# Mock TTS error response
mock_tts_instance.side_effect = Exception("TTS API error")
request_data = {"text": "Hello world", "encoding": "mp3"}
response = client.post("/api/tts", json=request_data)
assert response.status_code == 500
assert "Internal Server Error" in response.json()["detail"]
class TestPodcastEndpoint:
@patch("src.server.app.build_podcast_graph")
def test_generate_podcast_success(self, mock_build_graph, client):
mock_workflow = MagicMock()
mock_build_graph.return_value = mock_workflow
mock_workflow.invoke.return_value = {"output": b"fake_audio_data"}
request_data = {"content": "Test content for podcast"}
response = client.post("/api/podcast/generate", json=request_data)
assert response.status_code == 200
assert response.headers["content-type"] == "audio/mp3"
assert response.content == b"fake_audio_data"
@patch("src.server.app.build_podcast_graph")
def test_generate_podcast_error(self, mock_build_graph, client):
mock_build_graph.side_effect = Exception("Podcast generation failed")
request_data = {"content": "Test content"}
response = client.post("/api/podcast/generate", json=request_data)
assert response.status_code == 500
assert response.json()["detail"] == "Internal Server Error"
class TestPPTEndpoint:
@patch("src.server.app.build_ppt_graph")
@patch("builtins.open", new_callable=mock_open, read_data=b"fake_ppt_data")
def test_generate_ppt_success(self, mock_file, mock_build_graph, client):
mock_workflow = MagicMock()
mock_build_graph.return_value = mock_workflow
mock_workflow.invoke.return_value = {
"generated_file_path": "/fake/path/test.pptx"
}
request_data = {"content": "Test content for PPT"}
response = client.post("/api/ppt/generate", json=request_data)
assert response.status_code == 200
assert (
"application/vnd.openxmlformats-officedocument.presentationml.presentation"
in response.headers["content-type"]
)
assert response.content == b"fake_ppt_data"
@patch("src.server.app.build_ppt_graph")
def test_generate_ppt_error(self, mock_build_graph, client):
mock_build_graph.side_effect = Exception("PPT generation failed")
request_data = {"content": "Test content"}
response = client.post("/api/ppt/generate", json=request_data)
assert response.status_code == 500
assert response.json()["detail"] == "Internal Server Error"
class TestEnhancePromptEndpoint:
@patch("src.server.app.build_prompt_enhancer_graph")
def test_enhance_prompt_success(self, mock_build_graph, client):
mock_workflow = MagicMock()
mock_build_graph.return_value = mock_workflow
mock_workflow.invoke.return_value = {"output": "Enhanced prompt"}
request_data = {
"prompt": "Original prompt",
"context": "Some context",
"report_style": "academic",
}
response = client.post("/api/prompt/enhance", json=request_data)
assert response.status_code == 200
assert response.json()["result"] == "Enhanced prompt"
@patch("src.server.app.build_prompt_enhancer_graph")
def test_enhance_prompt_with_different_styles(self, mock_build_graph, client):
mock_workflow = MagicMock()
mock_build_graph.return_value = mock_workflow
mock_workflow.invoke.return_value = {"output": "Enhanced prompt"}
styles = [
"ACADEMIC",
"popular_science",
"NEWS",
"social_media",
"invalid_style",
]
for style in styles:
request_data = {"prompt": "Test prompt", "report_style": style}
response = client.post("/api/prompt/enhance", json=request_data)
assert response.status_code == 200
@patch("src.server.app.build_prompt_enhancer_graph")
def test_enhance_prompt_error(self, mock_build_graph, client):
mock_build_graph.side_effect = Exception("Enhancement failed")
request_data = {"prompt": "Test prompt"}
response = client.post("/api/prompt/enhance", json=request_data)
assert response.status_code == 500
assert response.json()["detail"] == "Internal Server Error"
class TestMCPEndpoint:
@patch("src.server.app.load_mcp_tools")
def test_mcp_server_metadata_success(self, mock_load_tools, client):
mock_load_tools.return_value = [
{"name": "test_tool", "description": "Test tool"}
]
request_data = {
"transport": "stdio",
"command": "test_command",
"args": ["arg1", "arg2"],
"env": {"ENV_VAR": "value"},
}
response = client.post("/api/mcp/server/metadata", json=request_data)
assert response.status_code == 200
response_data = response.json()
assert response_data["transport"] == "stdio"
assert response_data["command"] == "test_command"
assert len(response_data["tools"]) == 1
@patch("src.server.app.load_mcp_tools")
def test_mcp_server_metadata_with_custom_timeout(self, mock_load_tools, client):
mock_load_tools.return_value = []
request_data = {
"transport": "stdio",
"command": "test_command",
"timeout_seconds": 600,
}
response = client.post("/api/mcp/server/metadata", json=request_data)
assert response.status_code == 200
mock_load_tools.assert_called_once()
@patch("src.server.app.load_mcp_tools")
def test_mcp_server_metadata_with_exception(self, mock_load_tools, client):
mock_load_tools.side_effect = HTTPException(
status_code=400, detail="MCP Server Error"
)
request_data = {
"transport": "stdio",
"command": "test_command",
"args": ["arg1", "arg2"],
"env": {"ENV_VAR": "value"},
}
response = client.post("/api/mcp/server/metadata", json=request_data)
assert response.status_code == 500
assert response.json()["detail"] == "Internal Server Error"
class TestRAGEndpoints:
@patch("src.server.app.SELECTED_RAG_PROVIDER", "test_provider")
def test_rag_config(self, client):
response = client.get("/api/rag/config")
assert response.status_code == 200
assert response.json()["provider"] == "test_provider"
@patch("src.server.app.build_retriever")
def test_rag_resources_with_retriever(self, mock_build_retriever, client):
mock_retriever = MagicMock()
mock_retriever.list_resources.return_value = [
{
"uri": "test_uri",
"title": "Test Resource",
"description": "Test Description",
}
]
mock_build_retriever.return_value = mock_retriever
response = client.get("/api/rag/resources?query=test")
assert response.status_code == 200
assert len(response.json()["resources"]) == 1
@patch("src.server.app.build_retriever")
def test_rag_resources_without_retriever(self, mock_build_retriever, client):
mock_build_retriever.return_value = None
response = client.get("/api/rag/resources")
assert response.status_code == 200
assert response.json()["resources"] == []
class TestChatStreamEndpoint:
@patch("src.server.app.graph")
def test_chat_stream_with_default_thread_id(self, mock_graph, client):
# Mock the async stream
async def mock_astream(*args, **kwargs):
yield ("agent1", "step1", {"test": "data"})
mock_graph.astream = mock_astream
request_data = {
"thread_id": "__default__",
"messages": [{"role": "user", "content": "Hello"}],
"resources": [],
"max_plan_iterations": 3,
"max_step_num": 10,
"max_search_results": 5,
"auto_accepted_plan": True,
"interrupt_feedback": "",
"mcp_settings": {},
"enable_background_investigation": False,
"report_style": "academic",
}
response = client.post("/api/chat/stream", json=request_data)
assert response.status_code == 200
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
class TestAstreamWorkflowGenerator:
@pytest.mark.asyncio
@patch("src.server.app.graph")
async def test_astream_workflow_generator_basic_flow(self, mock_graph):
# Mock AI message chunk
mock_message = AIMessageChunk(content="Hello world")
mock_message.id = "msg_123"
mock_message.response_metadata = {}
mock_message.tool_calls = []
mock_message.tool_call_chunks = []
# Mock the async stream - yield messages in the correct format
async def mock_astream(*args, **kwargs):
# Yield a tuple (message, metadata) instead of just [message]
yield ("agent1:subagent", "messages", (mock_message, {}))
mock_graph.astream = mock_astream
messages = [{"role": "user", "content": "Hello"}]
thread_id = "test_thread"
resources = []
generator = _astream_workflow_generator(
messages=messages,
thread_id=thread_id,
resources=resources,
max_plan_iterations=3,
max_step_num=10,
max_search_results=5,
auto_accepted_plan=True,
interrupt_feedback="",
mcp_settings={},
enable_background_investigation=False,
report_style=ReportStyle.ACADEMIC,
enable_deep_thinking=False,
)
events = []
async for event in generator:
events.append(event)
assert len(events) == 1
assert "event: message_chunk" in events[0]
assert "Hello world" in events[0]
# Check for the actual agent name that appears in the output
assert '"agent": "a"' in events[0]
@pytest.mark.asyncio
@patch("src.server.app.graph")
async def test_astream_workflow_generator_with_interrupt_feedback(self, mock_graph):
# Mock the async stream
async def mock_astream(*args, **kwargs):
# Verify that Command is passed as input when interrupt_feedback is provided
assert isinstance(args[0], Command)
assert "[edit_plan] Hello" in args[0].resume
yield ("agent1", "step1", {"test": "data"})
mock_graph.astream = mock_astream
messages = [{"role": "user", "content": "Hello"}]
generator = _astream_workflow_generator(
messages=messages,
thread_id="test_thread",
resources=[],
max_plan_iterations=3,
max_step_num=10,
max_search_results=5,
auto_accepted_plan=False,
interrupt_feedback="edit_plan",
mcp_settings={},
enable_background_investigation=False,
report_style=ReportStyle.ACADEMIC,
enable_deep_thinking=False,
)
events = []
async for event in generator:
events.append(event)
@pytest.mark.asyncio
@patch("src.server.app.graph")
async def test_astream_workflow_generator_interrupt_event(self, mock_graph):
# Mock interrupt data
mock_interrupt = MagicMock()
mock_interrupt.ns = ["interrupt_id"]
mock_interrupt.value = "Plan requires approval"
interrupt_data = {"__interrupt__": [mock_interrupt]}
async def mock_astream(*args, **kwargs):
yield ("agent1", "step1", interrupt_data)
mock_graph.astream = mock_astream
generator = _astream_workflow_generator(
messages=[],
thread_id="test_thread",
resources=[],
max_plan_iterations=3,
max_step_num=10,
max_search_results=5,
auto_accepted_plan=True,
interrupt_feedback="",
mcp_settings={},
enable_background_investigation=False,
report_style=ReportStyle.ACADEMIC,
enable_deep_thinking=False,
)
events = []
async for event in generator:
events.append(event)
assert len(events) == 1
assert "event: interrupt" in events[0]
assert "Plan requires approval" in events[0]
assert "interrupt_id" in events[0]
@pytest.mark.asyncio
@patch("src.server.app.graph")
async def test_astream_workflow_generator_tool_message(self, mock_graph):
# Mock tool message
mock_tool_message = ToolMessage(content="Tool result", tool_call_id="tool_123")
mock_tool_message.id = "msg_456"
async def mock_astream(*args, **kwargs):
yield ("agent1:subagent", "step1", (mock_tool_message, {}))
mock_graph.astream = mock_astream
generator = _astream_workflow_generator(
messages=[],
thread_id="test_thread",
resources=[],
max_plan_iterations=3,
max_step_num=10,
max_search_results=5,
auto_accepted_plan=True,
interrupt_feedback="",
mcp_settings={},
enable_background_investigation=False,
report_style=ReportStyle.ACADEMIC,
enable_deep_thinking=False,
)
events = []
async for event in generator:
events.append(event)
assert len(events) == 1
assert "event: tool_call_result" in events[0]
assert "Tool result" in events[0]
assert "tool_123" in events[0]
@pytest.mark.asyncio
@patch("src.server.app.graph")
async def test_astream_workflow_generator_ai_message_with_tool_calls(
self, mock_graph
):
# Mock AI message with tool calls
mock_ai_message = AIMessageChunk(content="Making tool call")
mock_ai_message.id = "msg_789"
mock_ai_message.response_metadata = {"finish_reason": "tool_calls"}
mock_ai_message.tool_calls = [{"name": "search", "args": {"query": "test"}}]
mock_ai_message.tool_call_chunks = [{"name": "search"}]
async def mock_astream(*args, **kwargs):
yield ("agent1:subagent", "step1", (mock_ai_message, {}))
mock_graph.astream = mock_astream
generator = _astream_workflow_generator(
messages=[],
thread_id="test_thread",
resources=[],
max_plan_iterations=3,
max_step_num=10,
max_search_results=5,
auto_accepted_plan=True,
interrupt_feedback="",
mcp_settings={},
enable_background_investigation=False,
report_style=ReportStyle.ACADEMIC,
enable_deep_thinking=False,
)
events = []
async for event in generator:
events.append(event)
assert len(events) == 1
assert "event: tool_calls" in events[0]
assert "Making tool call" in events[0]
assert "tool_calls" in events[0]
@pytest.mark.asyncio
@patch("src.server.app.graph")
async def test_astream_workflow_generator_ai_message_with_tool_call_chunks(
self, mock_graph
):
# Mock AI message with only tool call chunks
mock_ai_message = AIMessageChunk(content="Streaming tool call")
mock_ai_message.id = "msg_101"
mock_ai_message.response_metadata = {}
mock_ai_message.tool_calls = []
mock_ai_message.tool_call_chunks = [{"name": "search", "index": 0}]
async def mock_astream(*args, **kwargs):
yield ("agent1:subagent", "step1", (mock_ai_message, {}))
mock_graph.astream = mock_astream
generator = _astream_workflow_generator(
messages=[],
thread_id="test_thread",
resources=[],
max_plan_iterations=3,
max_step_num=10,
max_search_results=5,
auto_accepted_plan=True,
interrupt_feedback="",
mcp_settings={},
enable_background_investigation=False,
report_style=ReportStyle.ACADEMIC,
enable_deep_thinking=False,
)
events = []
async for event in generator:
events.append(event)
assert len(events) == 1
assert "event: tool_call_chunks" in events[0]
assert "Streaming tool call" in events[0]
@pytest.mark.asyncio
@patch("src.server.app.graph")
async def test_astream_workflow_generator_with_finish_reason(self, mock_graph):
# Mock AI message with finish reason
mock_ai_message = AIMessageChunk(content="Complete response")
mock_ai_message.id = "msg_finish"
mock_ai_message.response_metadata = {"finish_reason": "stop"}
mock_ai_message.tool_calls = []
mock_ai_message.tool_call_chunks = []
async def mock_astream(*args, **kwargs):
yield ("agent1:subagent", "step1", (mock_ai_message, {}))
mock_graph.astream = mock_astream
generator = _astream_workflow_generator(
messages=[],
thread_id="test_thread",
resources=[],
max_plan_iterations=3,
max_step_num=10,
max_search_results=5,
auto_accepted_plan=True,
interrupt_feedback="",
mcp_settings={},
enable_background_investigation=False,
report_style=ReportStyle.ACADEMIC,
enable_deep_thinking=False,
)
events = []
async for event in generator:
events.append(event)
assert len(events) == 1
assert "event: message_chunk" in events[0]
assert "finish_reason" in events[0]
assert "stop" in events[0]
@pytest.mark.asyncio
@patch("src.server.app.graph")
async def test_astream_workflow_generator_config_passed_correctly(self, mock_graph):
mock_ai_message = AIMessageChunk(content="Test")
mock_ai_message.id = "test_id"
mock_ai_message.response_metadata = {}
mock_ai_message.tool_calls = []
mock_ai_message.tool_call_chunks = []
async def verify_config(*args, **kwargs):
config = kwargs.get("config", {})
assert config["thread_id"] == "test_thread"
assert config["max_plan_iterations"] == 5
assert config["max_step_num"] == 20
assert config["max_search_results"] == 10
assert config["report_style"] == ReportStyle.NEWS.value
yield ("agent1", "messages", [mock_ai_message])
class TestGenerateProseEndpoint:
@patch("src.server.app.build_prose_graph")
def test_generate_prose_success(self, mock_build_graph, client):
# Mock the workflow and its astream method
mock_workflow = MagicMock()
mock_build_graph.return_value = mock_workflow
class MockEvent:
def __init__(self, content):
self.content = content
async def mock_astream(*args, **kwargs):
yield (None, [MockEvent("Generated prose 1")])
yield (None, [MockEvent("Generated prose 2")])
mock_workflow.astream.return_value = mock_astream()
request_data = {
"prompt": "Write a story.",
"option": "default",
"command": "generate",
}
response = client.post("/api/prose/generate", json=request_data)
assert response.status_code == 200
assert response.headers["content-type"].startswith("text/event-stream")
# Read the streaming response content
content = b"".join(response.iter_bytes())
assert b"Generated prose 1" in content or b"Generated prose 2" in content
@patch("src.server.app.build_prose_graph")
def test_generate_prose_error(self, mock_build_graph, client):
mock_build_graph.side_effect = Exception("Prose generation failed")
request_data = {
"prompt": "Write a story.",
"option": "default",
"command": "generate",
}
response = client.post("/api/prose/generate", json=request_data)
assert response.status_code == 500
assert response.json()["detail"] == "Internal Server Error"
-167
View File
@@ -1,167 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import pytest
from pydantic import ValidationError
from src.config.report_style import ReportStyle
from src.rag.retriever import Resource
from unittest.mock import AsyncMock, patch, MagicMock
from fastapi import HTTPException
from src.server.chat_request import (
ContentItem,
ChatMessage,
ChatRequest,
TTSRequest,
GeneratePodcastRequest,
GeneratePPTRequest,
GenerateProseRequest,
EnhancePromptRequest,
)
import src.server.mcp_utils as mcp_utils # Assuming mcp_utils is the module to test
def test_content_item_text_and_image():
item_text = ContentItem(type="text", text="hello")
assert item_text.type == "text"
assert item_text.text == "hello"
assert item_text.image_url is None
item_image = ContentItem(type="image", image_url="http://img.com/1.png")
assert item_image.type == "image"
assert item_image.text is None
assert item_image.image_url == "http://img.com/1.png"
def test_chat_message_with_string_content():
msg = ChatMessage(role="user", content="Hello!")
assert msg.role == "user"
assert msg.content == "Hello!"
def test_chat_message_with_content_items():
items = [ContentItem(type="text", text="hi")]
msg = ChatMessage(role="assistant", content=items)
assert msg.role == "assistant"
assert isinstance(msg.content, list)
assert msg.content[0].type == "text"
def test_chat_request_defaults():
req = ChatRequest()
assert req.messages == []
assert req.resources == []
assert req.debug is False
assert req.thread_id == "__default__"
assert req.max_plan_iterations == 1
assert req.max_step_num == 3
assert req.max_search_results == 3
assert req.auto_accepted_plan is False
assert req.interrupt_feedback is None
assert req.mcp_settings is None
assert req.enable_background_investigation is True
assert req.report_style == ReportStyle.ACADEMIC
def test_chat_request_with_values():
resource = Resource(
name="test", type="doc", uri="some-uri-value", title="some-title-value"
)
msg = ChatMessage(role="user", content="hi")
req = ChatRequest(
messages=[msg],
resources=[resource],
debug=True,
thread_id="tid",
max_plan_iterations=2,
max_step_num=5,
max_search_results=10,
auto_accepted_plan=True,
interrupt_feedback="stop",
mcp_settings={"foo": "bar"},
enable_background_investigation=False,
report_style="academic",
)
assert req.messages[0].role == "user"
assert req.debug is True
assert req.thread_id == "tid"
assert req.max_plan_iterations == 2
assert req.max_step_num == 5
assert req.max_search_results == 10
assert req.auto_accepted_plan is True
assert req.interrupt_feedback == "stop"
assert req.mcp_settings == {"foo": "bar"}
assert req.enable_background_investigation is False
assert req.report_style == ReportStyle.ACADEMIC
def test_tts_request_defaults():
req = TTSRequest(text="hello")
assert req.text == "hello"
assert req.voice_type == "BV700_V2_streaming"
assert req.encoding == "mp3"
assert req.speed_ratio == 1.0
assert req.volume_ratio == 1.0
assert req.pitch_ratio == 1.0
assert req.text_type == "plain"
assert req.with_frontend == 1
assert req.frontend_type == "unitTson"
def test_generate_podcast_request():
req = GeneratePodcastRequest(content="Podcast content")
assert req.content == "Podcast content"
def test_generate_ppt_request():
req = GeneratePPTRequest(content="PPT content")
assert req.content == "PPT content"
def test_generate_prose_request():
req = GenerateProseRequest(prompt="Write a poem", option="poet", command="rhyme")
assert req.prompt == "Write a poem"
assert req.option == "poet"
assert req.command == "rhyme"
req2 = GenerateProseRequest(prompt="Write", option="short")
assert req2.command == ""
def test_enhance_prompt_request_defaults():
req = EnhancePromptRequest(prompt="Improve this")
assert req.prompt == "Improve this"
assert req.context == ""
assert req.report_style == "academic"
def test_content_item_validation_error():
with pytest.raises(ValidationError):
ContentItem() # missing required 'type'
def test_chat_message_validation_error():
with pytest.raises(ValidationError):
ChatMessage(role="user") # missing content
def test_tts_request_validation_error():
with pytest.raises(ValidationError):
TTSRequest() # missing required 'text'
@pytest.mark.asyncio
@patch("src.server.mcp_utils._get_tools_from_client_session", new_callable=AsyncMock)
@patch("src.server.mcp_utils.StdioServerParameters")
@patch("src.server.mcp_utils.stdio_client")
async def test_load_mcp_tools_exception_handling(
mock_stdio_client, mock_StdioServerParameters, mock_get_tools
): # Changed to async def
mock_get_tools.side_effect = Exception("unexpected error")
mock_StdioServerParameters.return_value = MagicMock()
mock_stdio_client.return_value = MagicMock()
with pytest.raises(HTTPException) as exc:
await mcp_utils.load_mcp_tools(server_type="stdio", command="foo") # Use await
assert exc.value.status_code == 500
assert "unexpected error" in exc.value.detail
-73
View File
@@ -1,73 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import pytest
from pydantic import ValidationError
from src.server.mcp_request import MCPServerMetadataRequest, MCPServerMetadataResponse
def test_mcp_server_metadata_request_required_fields():
# 'transport' is required
req = MCPServerMetadataRequest(transport="stdio")
assert req.transport == "stdio"
assert req.command is None
assert req.args is None
assert req.url is None
assert req.env is None
assert req.timeout_seconds is None
def test_mcp_server_metadata_request_optional_fields():
req = MCPServerMetadataRequest(
transport="sse",
command="run",
args=["--foo", "bar"],
url="http://localhost:8080",
env={"FOO": "BAR"},
timeout_seconds=30,
)
assert req.transport == "sse"
assert req.command == "run"
assert req.args == ["--foo", "bar"]
assert req.url == "http://localhost:8080"
assert req.env == {"FOO": "BAR"}
assert req.timeout_seconds == 30
def test_mcp_server_metadata_request_missing_transport():
with pytest.raises(ValidationError):
MCPServerMetadataRequest()
def test_mcp_server_metadata_response_required_fields():
resp = MCPServerMetadataResponse(transport="stdio")
assert resp.transport == "stdio"
assert resp.command is None
assert resp.args is None
assert resp.url is None
assert resp.env is None
assert resp.tools == []
def test_mcp_server_metadata_response_optional_fields():
resp = MCPServerMetadataResponse(
transport="sse",
command="run",
args=["--foo", "bar"],
url="http://localhost:8080",
env={"FOO": "BAR"},
tools=["tool1", "tool2"],
)
assert resp.transport == "sse"
assert resp.command == "run"
assert resp.args == ["--foo", "bar"]
assert resp.url == "http://localhost:8080"
assert resp.env == {"FOO": "BAR"}
assert resp.tools == ["tool1", "tool2"]
def test_mcp_server_metadata_response_tools_default_factory():
resp1 = MCPServerMetadataResponse(transport="stdio")
resp2 = MCPServerMetadataResponse(transport="stdio")
resp1.tools.append("toolA")
assert resp2.tools == [] # Should not share list between instances
-121
View File
@@ -1,121 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import pytest
from unittest.mock import AsyncMock, patch, MagicMock
from fastapi import HTTPException
import src.server.mcp_utils as mcp_utils
@pytest.mark.asyncio
@patch("src.server.mcp_utils.ClientSession")
async def test__get_tools_from_client_session_success(mock_ClientSession):
mock_read = AsyncMock()
mock_write = AsyncMock()
mock_context_manager = AsyncMock()
mock_context_manager.__aenter__.return_value = (mock_read, mock_write)
mock_context_manager.__aexit__.return_value = None
mock_session = AsyncMock()
mock_session.__aenter__.return_value = mock_session
mock_session.__aexit__.return_value = None
mock_session.initialize = AsyncMock()
mock_tools_obj = MagicMock()
mock_tools_obj.tools = ["tool1", "tool2"]
mock_session.list_tools = AsyncMock(return_value=mock_tools_obj)
mock_ClientSession.return_value = mock_session
result = await mcp_utils._get_tools_from_client_session(
mock_context_manager, timeout_seconds=5
)
assert result == ["tool1", "tool2"]
mock_session.initialize.assert_awaited_once()
mock_session.list_tools.assert_awaited_once()
@pytest.mark.asyncio
@patch("src.server.mcp_utils._get_tools_from_client_session", new_callable=AsyncMock)
@patch("src.server.mcp_utils.StdioServerParameters")
@patch("src.server.mcp_utils.stdio_client")
async def test_load_mcp_tools_stdio_success(
mock_stdio_client, mock_StdioServerParameters, mock_get_tools
):
mock_get_tools.return_value = ["toolA"]
params = MagicMock()
mock_StdioServerParameters.return_value = params
mock_client = MagicMock()
mock_stdio_client.return_value = mock_client
result = await mcp_utils.load_mcp_tools(
server_type="stdio",
command="echo",
args=["foo"],
env={"FOO": "BAR"},
timeout_seconds=3,
)
assert result == ["toolA"]
mock_StdioServerParameters.assert_called_once_with(
command="echo", args=["foo"], env={"FOO": "BAR"}
)
mock_stdio_client.assert_called_once_with(params)
mock_get_tools.assert_awaited_once_with(mock_client, 3)
@pytest.mark.asyncio
async def test_load_mcp_tools_stdio_missing_command():
with pytest.raises(HTTPException) as exc:
await mcp_utils.load_mcp_tools(server_type="stdio")
assert exc.value.status_code == 400
assert "Command is required" in exc.value.detail
@pytest.mark.asyncio
@patch("src.server.mcp_utils._get_tools_from_client_session", new_callable=AsyncMock)
@patch("src.server.mcp_utils.sse_client")
async def test_load_mcp_tools_sse_success(mock_sse_client, mock_get_tools):
mock_get_tools.return_value = ["toolB"]
mock_client = MagicMock()
mock_sse_client.return_value = mock_client
result = await mcp_utils.load_mcp_tools(
server_type="sse",
url="http://localhost:1234",
timeout_seconds=7,
)
assert result == ["toolB"]
mock_sse_client.assert_called_once_with(url="http://localhost:1234")
mock_get_tools.assert_awaited_once_with(mock_client, 7)
@pytest.mark.asyncio
async def test_load_mcp_tools_sse_missing_url():
with pytest.raises(HTTPException) as exc:
await mcp_utils.load_mcp_tools(server_type="sse")
assert exc.value.status_code == 400
assert "URL is required" in exc.value.detail
@pytest.mark.asyncio
async def test_load_mcp_tools_unsupported_type():
with pytest.raises(HTTPException) as exc:
await mcp_utils.load_mcp_tools(server_type="unknown")
assert exc.value.status_code == 400
assert "Unsupported server type" in exc.value.detail
@pytest.mark.asyncio
@patch("src.server.mcp_utils._get_tools_from_client_session", new_callable=AsyncMock)
@patch("src.server.mcp_utils.StdioServerParameters")
@patch("src.server.mcp_utils.stdio_client")
async def test_load_mcp_tools_exception_handling(
mock_stdio_client, mock_StdioServerParameters, mock_get_tools
):
mock_get_tools.side_effect = Exception("unexpected error")
mock_StdioServerParameters.return_value = MagicMock()
mock_stdio_client.return_value = MagicMock()
with pytest.raises(HTTPException) as exc:
await mcp_utils.load_mcp_tools(server_type="stdio", command="foo")
assert exc.value.status_code == 500
assert "unexpected error" in exc.value.detail
-109
View File
@@ -1,109 +0,0 @@
from unittest.mock import Mock, patch
from src.tools.crawl import crawl_tool
class TestCrawlTool:
@patch("src.tools.crawl.Crawler")
def test_crawl_tool_success(self, mock_crawler_class):
# Arrange
mock_crawler = Mock()
mock_article = Mock()
mock_article.to_markdown.return_value = (
"# Test Article\nThis is test content." * 100
)
mock_crawler.crawl.return_value = mock_article
mock_crawler_class.return_value = mock_crawler
url = "https://example.com"
# Act
result = crawl_tool(url)
# Assert
assert isinstance(result, dict)
assert result["url"] == url
assert "crawled_content" in result
assert len(result["crawled_content"]) <= 1000
mock_crawler_class.assert_called_once()
mock_crawler.crawl.assert_called_once_with(url)
mock_article.to_markdown.assert_called_once()
@patch("src.tools.crawl.Crawler")
def test_crawl_tool_short_content(self, mock_crawler_class):
# Arrange
mock_crawler = Mock()
mock_article = Mock()
short_content = "Short content"
mock_article.to_markdown.return_value = short_content
mock_crawler.crawl.return_value = mock_article
mock_crawler_class.return_value = mock_crawler
url = "https://example.com"
# Act
result = crawl_tool(url)
# Assert
assert result["crawled_content"] == short_content
@patch("src.tools.crawl.Crawler")
@patch("src.tools.crawl.logger")
def test_crawl_tool_crawler_exception(self, mock_logger, mock_crawler_class):
# Arrange
mock_crawler = Mock()
mock_crawler.crawl.side_effect = Exception("Network error")
mock_crawler_class.return_value = mock_crawler
url = "https://example.com"
# Act
result = crawl_tool(url)
# Assert
assert isinstance(result, str)
assert "Failed to crawl" in result
assert "Network error" in result
mock_logger.error.assert_called_once()
@patch("src.tools.crawl.Crawler")
@patch("src.tools.crawl.logger")
def test_crawl_tool_crawler_instantiation_exception(
self, mock_logger, mock_crawler_class
):
# Arrange
mock_crawler_class.side_effect = Exception("Crawler init error")
url = "https://example.com"
# Act
result = crawl_tool(url)
# Assert
assert isinstance(result, str)
assert "Failed to crawl" in result
assert "Crawler init error" in result
mock_logger.error.assert_called_once()
@patch("src.tools.crawl.Crawler")
@patch("src.tools.crawl.logger")
def test_crawl_tool_markdown_conversion_exception(
self, mock_logger, mock_crawler_class
):
# Arrange
mock_crawler = Mock()
mock_article = Mock()
mock_article.to_markdown.side_effect = Exception("Markdown conversion error")
mock_crawler.crawl.return_value = mock_article
mock_crawler_class.return_value = mock_crawler
url = "https://example.com"
# Act
result = crawl_tool(url)
# Assert
assert isinstance(result, str)
assert "Failed to crawl" in result
assert "Markdown conversion error" in result
mock_logger.error.assert_called_once()
-119
View File
@@ -1,119 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
from unittest.mock import Mock, call, patch
from src.tools.decorators import create_logged_tool
class MockBaseTool:
"""Mock base tool class for testing."""
def _run(self, *args, **kwargs):
return "base_result"
class TestLoggedToolMixin:
def test_run_calls_log_operation(self):
"""Test that _run calls _log_operation with correct parameters."""
# Create a logged tool instance
LoggedTool = create_logged_tool(MockBaseTool)
tool = LoggedTool()
# Mock the _log_operation method
tool._log_operation = Mock()
# Call _run with test parameters
args = ("arg1", "arg2")
kwargs = {"key1": "value1", "key2": "value2"}
tool._run(*args, **kwargs)
# Verify _log_operation was called with correct parameters
tool._log_operation.assert_called_once_with("_run", *args, **kwargs)
def test_run_calls_super_run(self):
"""Test that _run calls the parent class _run method."""
# Create a logged tool instance
LoggedTool = create_logged_tool(MockBaseTool)
tool = LoggedTool()
# Mock the parent _run method
with patch.object(
MockBaseTool, "_run", return_value="mocked_result"
) as mock_super_run:
args = ("arg1", "arg2")
kwargs = {"key1": "value1"}
result = tool._run(*args, **kwargs)
# Verify super()._run was called with correct parameters
mock_super_run.assert_called_once_with(*args, **kwargs)
# Verify the result is returned
assert result == "mocked_result"
def test_run_logs_result(self):
"""Test that _run logs the result with debug level."""
LoggedTool = create_logged_tool(MockBaseTool)
tool = LoggedTool()
with patch("src.tools.decorators.logger.debug") as mock_debug:
tool._run("test_arg")
# Verify debug log was called with correct message
mock_debug.assert_has_calls(
[
call("Tool MockBaseTool._run called with parameters: test_arg"),
call("Tool MockBaseTool returned: base_result"),
]
)
def test_run_returns_super_result(self):
"""Test that _run returns the result from parent class."""
LoggedTool = create_logged_tool(MockBaseTool)
tool = LoggedTool()
result = tool._run()
assert result == "base_result"
def test_run_with_no_args(self):
"""Test _run method with no arguments."""
LoggedTool = create_logged_tool(MockBaseTool)
tool = LoggedTool()
with patch("src.tools.decorators.logger.debug") as mock_debug:
tool._log_operation = Mock()
result = tool._run()
# Verify _log_operation called with no args
tool._log_operation.assert_called_once_with("_run")
# Verify result logging
mock_debug.assert_called_once()
assert result == "base_result"
def test_run_with_mixed_args_kwargs(self):
"""Test _run method with both positional and keyword arguments."""
LoggedTool = create_logged_tool(MockBaseTool)
tool = LoggedTool()
tool._log_operation = Mock()
args = ("pos1", "pos2")
kwargs = {"kw1": "val1", "kw2": "val2"}
result = tool._run(*args, **kwargs)
# Verify all arguments passed correctly
tool._log_operation.assert_called_once_with("_run", *args, **kwargs)
assert result == "base_result"
def test_run_class_name_replacement(self):
"""Test that class name 'Logged' prefix is correctly removed in logging."""
LoggedTool = create_logged_tool(MockBaseTool)
tool = LoggedTool()
with patch("src.tools.decorators.logger.debug") as mock_debug:
tool._run()
# Verify the logged class name has 'Logged' prefix removed
call_args = mock_debug.call_args[0][0]
assert "Tool MockBaseTool returned:" in call_args
assert "LoggedMockBaseTool" not in call_args
-147
View File
@@ -1,147 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import pytest
from unittest.mock import patch
from src.tools.python_repl import python_repl_tool
class TestPythonReplTool:
@patch("src.tools.python_repl.repl")
@patch("src.tools.python_repl.logger")
def test_successful_code_execution(self, mock_logger, mock_repl):
# Arrange
code = "print('Hello, World!')"
expected_output = "Hello, World!\n"
mock_repl.run.return_value = expected_output
# Act
result = python_repl_tool(code)
# Assert
mock_repl.run.assert_called_once_with(code)
mock_logger.info.assert_called_with("Code execution successful")
assert "Successfully executed:" in result
assert code in result
assert expected_output in result
@patch("src.tools.python_repl.repl")
@patch("src.tools.python_repl.logger")
def test_invalid_input_type(self, mock_logger, mock_repl):
# Arrange
invalid_code = 123
# Act & Assert - expect ValidationError from LangChain
with pytest.raises(Exception) as exc_info:
python_repl_tool(invalid_code)
# Verify that it's a validation error
assert "ValidationError" in str(
type(exc_info.value)
) or "validation error" in str(exc_info.value)
# The REPL should not be called since validation fails first
mock_repl.run.assert_not_called()
@patch("src.tools.python_repl.repl")
@patch("src.tools.python_repl.logger")
def test_code_execution_with_error_in_result(self, mock_logger, mock_repl):
# Arrange
code = "invalid_function()"
error_result = "NameError: name 'invalid_function' is not defined"
mock_repl.run.return_value = error_result
# Act
result = python_repl_tool(code)
# Assert
mock_repl.run.assert_called_once_with(code)
mock_logger.error.assert_called_with(error_result)
assert "Error executing code:" in result
assert code in result
assert error_result in result
@patch("src.tools.python_repl.repl")
@patch("src.tools.python_repl.logger")
def test_code_execution_with_exception_in_result(self, mock_logger, mock_repl):
# Arrange
code = "1/0"
exception_result = "ZeroDivisionError: division by zero"
mock_repl.run.return_value = exception_result
# Act
result = python_repl_tool(code)
# Assert
mock_repl.run.assert_called_once_with(code)
mock_logger.error.assert_called_with(exception_result)
assert "Error executing code:" in result
assert code in result
assert exception_result in result
@patch("src.tools.python_repl.repl")
@patch("src.tools.python_repl.logger")
def test_code_execution_raises_exception(self, mock_logger, mock_repl):
# Arrange
code = "print('test')"
exception = RuntimeError("REPL failed")
mock_repl.run.side_effect = exception
# Act
result = python_repl_tool(code)
# Assert
mock_repl.run.assert_called_once_with(code)
mock_logger.error.assert_called_with(repr(exception))
assert "Error executing code:" in result
assert code in result
assert repr(exception) in result
@patch("src.tools.python_repl.repl")
@patch("src.tools.python_repl.logger")
def test_successful_execution_with_calculation(self, mock_logger, mock_repl):
# Arrange
code = "result = 2 + 3\nprint(result)"
expected_output = "5\n"
mock_repl.run.return_value = expected_output
# Act
result = python_repl_tool(code)
# Assert
mock_repl.run.assert_called_once_with(code)
mock_logger.info.assert_any_call("Executing Python code")
mock_logger.info.assert_any_call("Code execution successful")
assert "Successfully executed:" in result
assert code in result
assert expected_output in result
@patch("src.tools.python_repl.repl")
@patch("src.tools.python_repl.logger")
def test_empty_string_code(self, mock_logger, mock_repl):
# Arrange
code = ""
mock_repl.run.return_value = ""
# Act
result = python_repl_tool(code)
# Assert
mock_repl.run.assert_called_once_with(code)
mock_logger.info.assert_called_with("Code execution successful")
assert "Successfully executed:" in result
@patch("src.tools.python_repl.repl")
@patch("src.tools.python_repl.logger")
def test_logging_calls(self, mock_logger, mock_repl):
# Arrange
code = "x = 1"
mock_repl.run.return_value = ""
# Act
python_repl_tool(code)
# Assert
mock_logger.info.assert_any_call("Executing Python code")
mock_logger.info.assert_any_call("Code execution successful")
-54
View File
@@ -1,54 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import os
import pytest
from unittest.mock import patch
from src.tools.search import get_web_search_tool
from src.config import SearchEngine
class TestGetWebSearchTool:
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.TAVILY.value)
def test_get_web_search_tool_tavily(self):
tool = get_web_search_tool(max_search_results=5)
assert tool.name == "web_search"
assert tool.max_results == 5
assert tool.include_raw_content is True
assert tool.include_images is True
assert tool.include_image_descriptions is True
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.DUCKDUCKGO.value)
def test_get_web_search_tool_duckduckgo(self):
tool = get_web_search_tool(max_search_results=3)
assert tool.name == "web_search"
assert tool.max_results == 3
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.BRAVE_SEARCH.value)
@patch.dict(os.environ, {"BRAVE_SEARCH_API_KEY": "test_api_key"})
def test_get_web_search_tool_brave(self):
tool = get_web_search_tool(max_search_results=4)
assert tool.name == "web_search"
assert tool.search_wrapper.api_key == "test_api_key"
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.ARXIV.value)
def test_get_web_search_tool_arxiv(self):
tool = get_web_search_tool(max_search_results=2)
assert tool.name == "web_search"
assert tool.api_wrapper.top_k_results == 2
assert tool.api_wrapper.load_max_docs == 2
assert tool.api_wrapper.load_all_available_meta is True
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", "unsupported_engine")
def test_get_web_search_tool_unsupported_engine(self):
with pytest.raises(
ValueError, match="Unsupported search engine: unsupported_engine"
):
get_web_search_tool(max_search_results=1)
@patch("src.tools.search.SELECTED_SEARCH_ENGINE", SearchEngine.BRAVE_SEARCH.value)
@patch.dict(os.environ, {}, clear=True)
def test_get_web_search_tool_brave_no_api_key(self):
tool = get_web_search_tool(max_search_results=1)
assert tool.search_wrapper.api_key == ""
@@ -1,206 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import json
import pytest
from unittest.mock import Mock, patch, AsyncMock, MagicMock
import requests
from src.tools.tavily_search.tavily_search_api_wrapper import (
EnhancedTavilySearchAPIWrapper,
)
class TestEnhancedTavilySearchAPIWrapper:
@pytest.fixture
def wrapper(self):
with patch(
"src.tools.tavily_search.tavily_search_api_wrapper.OriginalTavilySearchAPIWrapper"
):
wrapper = EnhancedTavilySearchAPIWrapper(tavily_api_key="dummy-key")
# The parent class is mocked, so initialization won't fail
return wrapper
@pytest.fixture
def mock_response_data(self):
return {
"results": [
{
"title": "Test Title",
"url": "https://example.com",
"content": "Test content",
"score": 0.9,
"raw_content": "Raw test content",
}
],
"images": [
{
"url": "https://example.com/image.jpg",
"description": "Test image description",
}
],
}
@patch("src.tools.tavily_search.tavily_search_api_wrapper.requests.post")
def test_raw_results_success(self, mock_post, wrapper, mock_response_data):
mock_response = Mock()
mock_response.json.return_value = mock_response_data
mock_response.raise_for_status.return_value = None
mock_post.return_value = mock_response
result = wrapper.raw_results("test query", max_results=10)
assert result == mock_response_data
mock_post.assert_called_once()
call_args = mock_post.call_args
assert "json" in call_args.kwargs
assert call_args.kwargs["json"]["query"] == "test query"
assert call_args.kwargs["json"]["max_results"] == 10
@patch("src.tools.tavily_search.tavily_search_api_wrapper.requests.post")
def test_raw_results_with_all_parameters(
self, mock_post, wrapper, mock_response_data
):
mock_response = Mock()
mock_response.json.return_value = mock_response_data
mock_response.raise_for_status.return_value = None
mock_post.return_value = mock_response
result = wrapper.raw_results(
"test query",
max_results=3,
search_depth="basic",
include_domains=["example.com"],
exclude_domains=["spam.com"],
include_answer=True,
include_raw_content=True,
include_images=True,
include_image_descriptions=True,
)
assert result == mock_response_data
call_args = mock_post.call_args
params = call_args.kwargs["json"]
assert params["include_domains"] == ["example.com"]
assert params["exclude_domains"] == ["spam.com"]
assert params["include_answer"] is True
assert params["include_raw_content"] is True
@patch("src.tools.tavily_search.tavily_search_api_wrapper.requests.post")
def test_raw_results_http_error(self, mock_post, wrapper):
mock_response = Mock()
mock_response.raise_for_status.side_effect = requests.HTTPError("API Error")
mock_post.return_value = mock_response
with pytest.raises(requests.HTTPError):
wrapper.raw_results("test query")
@pytest.mark.asyncio
async def test_raw_results_async_success(self, wrapper, mock_response_data):
# Create a mock that acts as both the response and its context manager
mock_response_cm = AsyncMock()
mock_response_cm.__aenter__ = AsyncMock(return_value=mock_response_cm)
mock_response_cm.__aexit__ = AsyncMock(return_value=None)
mock_response_cm.status = 200
mock_response_cm.text = AsyncMock(return_value=json.dumps(mock_response_data))
# Create mock session that returns the context manager
mock_session = AsyncMock()
mock_session.post = MagicMock(
return_value=mock_response_cm
) # Use MagicMock, not AsyncMock
# Create mock session class
mock_session_cm = AsyncMock()
mock_session_cm.__aenter__ = AsyncMock(return_value=mock_session)
mock_session_cm.__aexit__ = AsyncMock(return_value=None)
with patch(
"src.tools.tavily_search.tavily_search_api_wrapper.aiohttp.ClientSession",
return_value=mock_session_cm,
):
result = await wrapper.raw_results_async("test query")
assert result == mock_response_data
@pytest.mark.asyncio
async def test_raw_results_async_error(self, wrapper):
# Create a mock that acts as both the response and its context manager
mock_response_cm = AsyncMock()
mock_response_cm.__aenter__ = AsyncMock(return_value=mock_response_cm)
mock_response_cm.__aexit__ = AsyncMock(return_value=None)
mock_response_cm.status = 400
mock_response_cm.reason = "Bad Request"
# Create mock session that returns the context manager
mock_session = AsyncMock()
mock_session.post = MagicMock(
return_value=mock_response_cm
) # Use MagicMock, not AsyncMock
# Create mock session class
mock_session_cm = AsyncMock()
mock_session_cm.__aenter__ = AsyncMock(return_value=mock_session)
mock_session_cm.__aexit__ = AsyncMock(return_value=None)
with patch(
"src.tools.tavily_search.tavily_search_api_wrapper.aiohttp.ClientSession",
return_value=mock_session_cm,
):
with pytest.raises(Exception, match="Error 400: Bad Request"):
await wrapper.raw_results_async("test query")
def test_clean_results_with_images(self, wrapper, mock_response_data):
result = wrapper.clean_results_with_images(mock_response_data)
assert len(result) == 2
# Test page result
page_result = result[0]
assert page_result["type"] == "page"
assert page_result["title"] == "Test Title"
assert page_result["url"] == "https://example.com"
assert page_result["content"] == "Test content"
assert page_result["score"] == 0.9
assert page_result["raw_content"] == "Raw test content"
# Test image result
image_result = result[1]
assert image_result["type"] == "image"
assert image_result["image_url"] == "https://example.com/image.jpg"
assert image_result["image_description"] == "Test image description"
def test_clean_results_without_raw_content(self, wrapper):
data = {
"results": [
{
"title": "Test Title",
"url": "https://example.com",
"content": "Test content",
"score": 0.9,
}
],
"images": [],
}
result = wrapper.clean_results_with_images(data)
assert len(result) == 1
assert "raw_content" not in result[0]
def test_clean_results_empty_images(self, wrapper):
data = {
"results": [
{
"title": "Test Title",
"url": "https://example.com",
"content": "Test content",
"score": 0.9,
}
],
"images": [],
}
result = wrapper.clean_results_with_images(data)
assert len(result) == 1
assert result[0]["type"] == "page"
@@ -1,264 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import json
import pytest
from unittest.mock import Mock, patch, AsyncMock
from src.tools.tavily_search.tavily_search_results_with_images import (
TavilySearchResultsWithImages,
)
from src.tools.tavily_search.tavily_search_api_wrapper import (
EnhancedTavilySearchAPIWrapper,
)
class TestTavilySearchResultsWithImages:
@pytest.fixture
def mock_api_wrapper(self):
"""Create a mock API wrapper."""
wrapper = Mock(spec=EnhancedTavilySearchAPIWrapper)
return wrapper
@pytest.fixture
def search_tool(self, mock_api_wrapper):
"""Create a TavilySearchResultsWithImages instance with mocked dependencies."""
tool = TavilySearchResultsWithImages(
max_results=5,
include_answer=True,
include_raw_content=True,
include_images=True,
include_image_descriptions=True,
)
tool.api_wrapper = mock_api_wrapper
return tool
@pytest.fixture
def sample_raw_results(self):
"""Sample raw results from Tavily API."""
return {
"query": "test query",
"answer": "Test answer",
"images": ["https://example.com/image1.jpg"],
"results": [
{
"title": "Test Title",
"url": "https://example.com",
"content": "Test content",
"score": 0.95,
"raw_content": "Raw test content",
}
],
"response_time": 1.5,
}
@pytest.fixture
def sample_cleaned_results(self):
"""Sample cleaned results."""
return [
{
"title": "Test Title",
"url": "https://example.com",
"content": "Test content",
}
]
def test_init_default_values(self):
"""Test initialization with default values."""
tool = TavilySearchResultsWithImages()
assert tool.include_image_descriptions is False
assert isinstance(tool.api_wrapper, EnhancedTavilySearchAPIWrapper)
def test_init_custom_values(self):
"""Test initialization with custom values."""
tool = TavilySearchResultsWithImages(
max_results=10, include_image_descriptions=True
)
assert tool.max_results == 10
assert tool.include_image_descriptions is True
@patch("builtins.print")
def test_run_success(
self,
mock_print,
search_tool,
mock_api_wrapper,
sample_raw_results,
sample_cleaned_results,
):
"""Test successful synchronous run."""
mock_api_wrapper.raw_results.return_value = sample_raw_results
mock_api_wrapper.clean_results_with_images.return_value = sample_cleaned_results
result, raw = search_tool._run("test query")
assert result == sample_cleaned_results
assert raw == sample_raw_results
mock_api_wrapper.raw_results.assert_called_once_with(
"test query",
search_tool.max_results,
search_tool.search_depth,
search_tool.include_domains,
search_tool.exclude_domains,
search_tool.include_answer,
search_tool.include_raw_content,
search_tool.include_images,
search_tool.include_image_descriptions,
)
mock_api_wrapper.clean_results_with_images.assert_called_once_with(
sample_raw_results
)
mock_print.assert_called_once()
@patch("builtins.print")
def test_run_exception(self, mock_print, search_tool, mock_api_wrapper):
"""Test synchronous run with exception."""
mock_api_wrapper.raw_results.side_effect = Exception("API Error")
result, raw = search_tool._run("test query")
assert "API Error" in result
assert raw == {}
mock_api_wrapper.clean_results_with_images.assert_not_called()
@pytest.mark.asyncio
@patch("builtins.print")
async def test_arun_success(
self,
mock_print,
search_tool,
mock_api_wrapper,
sample_raw_results,
sample_cleaned_results,
):
"""Test successful asynchronous run."""
mock_api_wrapper.raw_results_async = AsyncMock(return_value=sample_raw_results)
mock_api_wrapper.clean_results_with_images.return_value = sample_cleaned_results
result, raw = await search_tool._arun("test query")
assert result == sample_cleaned_results
assert raw == sample_raw_results
mock_api_wrapper.raw_results_async.assert_called_once_with(
"test query",
search_tool.max_results,
search_tool.search_depth,
search_tool.include_domains,
search_tool.exclude_domains,
search_tool.include_answer,
search_tool.include_raw_content,
search_tool.include_images,
search_tool.include_image_descriptions,
)
mock_api_wrapper.clean_results_with_images.assert_called_once_with(
sample_raw_results
)
mock_print.assert_called_once()
@pytest.mark.asyncio
@patch("builtins.print")
async def test_arun_exception(self, mock_print, search_tool, mock_api_wrapper):
"""Test asynchronous run with exception."""
mock_api_wrapper.raw_results_async = AsyncMock(
side_effect=Exception("Async API Error")
)
result, raw = await search_tool._arun("test query")
assert "Async API Error" in result
assert raw == {}
mock_api_wrapper.clean_results_with_images.assert_not_called()
@patch("builtins.print")
def test_run_with_run_manager(
self,
mock_print,
search_tool,
mock_api_wrapper,
sample_raw_results,
sample_cleaned_results,
):
"""Test run with callback manager."""
mock_run_manager = Mock()
mock_api_wrapper.raw_results.return_value = sample_raw_results
mock_api_wrapper.clean_results_with_images.return_value = sample_cleaned_results
result, raw = search_tool._run("test query", run_manager=mock_run_manager)
assert result == sample_cleaned_results
assert raw == sample_raw_results
@pytest.mark.asyncio
@patch("builtins.print")
async def test_arun_with_run_manager(
self,
mock_print,
search_tool,
mock_api_wrapper,
sample_raw_results,
sample_cleaned_results,
):
"""Test async run with callback manager."""
mock_run_manager = Mock()
mock_api_wrapper.raw_results_async = AsyncMock(return_value=sample_raw_results)
mock_api_wrapper.clean_results_with_images.return_value = sample_cleaned_results
result, raw = await search_tool._arun(
"test query", run_manager=mock_run_manager
)
assert result == sample_cleaned_results
assert raw == sample_raw_results
@patch("builtins.print")
def test_print_output_format(
self,
mock_print,
search_tool,
mock_api_wrapper,
sample_raw_results,
sample_cleaned_results,
):
"""Test that print outputs correctly formatted JSON."""
mock_api_wrapper.raw_results.return_value = sample_raw_results
mock_api_wrapper.clean_results_with_images.return_value = sample_cleaned_results
search_tool._run("test query")
# Verify print was called with expected format
call_args = mock_print.call_args[0]
assert call_args[0] == "sync"
assert isinstance(call_args[1], str) # Should be JSON string
# Verify it's valid JSON
json_data = json.loads(call_args[1])
assert json_data == sample_cleaned_results
@pytest.mark.asyncio
@patch("builtins.print")
async def test_async_print_output_format(
self,
mock_print,
search_tool,
mock_api_wrapper,
sample_raw_results,
sample_cleaned_results,
):
"""Test that async print outputs correctly formatted JSON."""
mock_api_wrapper.raw_results_async = AsyncMock(return_value=sample_raw_results)
mock_api_wrapper.clean_results_with_images.return_value = sample_cleaned_results
await search_tool._arun("test query")
# Verify print was called with expected format
call_args = mock_print.call_args[0]
assert call_args[0] == "async"
assert isinstance(call_args[1], str) # Should be JSON string
# Verify it's valid JSON
json_data = json.loads(call_args[1])
assert json_data == sample_cleaned_results
-122
View File
@@ -1,122 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
from unittest.mock import Mock, patch
from langchain_core.callbacks import (
CallbackManagerForToolRun,
AsyncCallbackManagerForToolRun,
)
import pytest
from src.tools.retriever import RetrieverInput, RetrieverTool, get_retriever_tool
from src.rag import Document, Retriever, Resource, Chunk
def test_retriever_input_model():
input_data = RetrieverInput(keywords="test keywords")
assert input_data.keywords == "test keywords"
def test_retriever_tool_init():
mock_retriever = Mock(spec=Retriever)
resources = [Resource(uri="test://uri", title="Test")]
tool = RetrieverTool(retriever=mock_retriever, resources=resources)
assert tool.name == "local_search_tool"
assert "retrieving information" in tool.description
assert tool.args_schema == RetrieverInput
assert tool.retriever == mock_retriever
assert tool.resources == resources
def test_retriever_tool_run_with_results():
mock_retriever = Mock(spec=Retriever)
chunk = Chunk(content="test content", similarity=0.9)
doc = Document(id="doc1", chunks=[chunk])
mock_retriever.query_relevant_documents.return_value = [doc]
resources = [Resource(uri="test://uri", title="Test")]
tool = RetrieverTool(retriever=mock_retriever, resources=resources)
result = tool._run("test keywords")
mock_retriever.query_relevant_documents.assert_called_once_with(
"test keywords", resources
)
assert isinstance(result, list)
assert len(result) == 1
assert result[0] == doc.to_dict()
def test_retriever_tool_run_no_results():
mock_retriever = Mock(spec=Retriever)
mock_retriever.query_relevant_documents.return_value = []
resources = [Resource(uri="test://uri", title="Test")]
tool = RetrieverTool(retriever=mock_retriever, resources=resources)
result = tool._run("test keywords")
assert result == "No results found from the local knowledge base."
@pytest.mark.asyncio
async def test_retriever_tool_arun():
mock_retriever = Mock(spec=Retriever)
chunk = Chunk(content="async content", similarity=0.8)
doc = Document(id="doc2", chunks=[chunk])
mock_retriever.query_relevant_documents.return_value = [doc]
resources = [Resource(uri="test://uri", title="Test")]
tool = RetrieverTool(retriever=mock_retriever, resources=resources)
mock_run_manager = Mock(spec=AsyncCallbackManagerForToolRun)
mock_sync_manager = Mock(spec=CallbackManagerForToolRun)
mock_run_manager.get_sync.return_value = mock_sync_manager
result = await tool._arun("async keywords", mock_run_manager)
mock_run_manager.get_sync.assert_called_once()
assert isinstance(result, list)
assert len(result) == 1
assert result[0] == doc.to_dict()
@patch("src.tools.retriever.build_retriever")
def test_get_retriever_tool_success(mock_build_retriever):
mock_retriever = Mock(spec=Retriever)
mock_build_retriever.return_value = mock_retriever
resources = [Resource(uri="test://uri", title="Test")]
tool = get_retriever_tool(resources)
assert isinstance(tool, RetrieverTool)
assert tool.retriever == mock_retriever
assert tool.resources == resources
def test_get_retriever_tool_empty_resources():
result = get_retriever_tool([])
assert result is None
@patch("src.tools.retriever.build_retriever")
def test_get_retriever_tool_no_retriever(mock_build_retriever):
mock_build_retriever.return_value = None
resources = [Resource(uri="test://uri", title="Test")]
result = get_retriever_tool(resources)
assert result is None
def test_retriever_tool_run_with_callback_manager():
mock_retriever = Mock(spec=Retriever)
mock_retriever.query_relevant_documents.return_value = []
resources = [Resource(uri="test://uri", title="Test")]
tool = RetrieverTool(retriever=mock_retriever, resources=resources)
mock_callback_manager = Mock(spec=CallbackManagerForToolRun)
result = tool._run("test keywords", mock_callback_manager)
assert result == "No results found from the local knowledge base."
-108
View File
@@ -1,108 +0,0 @@
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# SPDX-License-Identifier: MIT
import json
from src.utils.json_utils import repair_json_output
class TestRepairJsonOutput:
def test_valid_json_object(self):
"""Test with valid JSON object"""
content = '{"key": "value", "number": 123}'
result = repair_json_output(content)
expected = json.dumps({"key": "value", "number": 123}, ensure_ascii=False)
assert result == expected
def test_valid_json_array(self):
"""Test with valid JSON array"""
content = '[1, 2, 3, "test"]'
result = repair_json_output(content)
expected = json.dumps([1, 2, 3, "test"], ensure_ascii=False)
assert result == expected
def test_json_with_code_block_json(self):
"""Test JSON wrapped in ```json code block"""
content = '```json\n{"key": "value"}\n```'
result = repair_json_output(content)
expected = json.dumps({"key": "value"}, ensure_ascii=False)
assert result == expected
def test_json_with_code_block_ts(self):
"""Test JSON wrapped in ```ts code block"""
content = '```ts\n{"key": "value"}\n```'
result = repair_json_output(content)
expected = json.dumps({"key": "value"}, ensure_ascii=False)
assert result == expected
def test_malformed_json_repair(self):
"""Test with malformed JSON that can be repaired"""
content = '{"key": "value", "incomplete":'
result = repair_json_output(content)
# Should return repaired JSON
assert result.startswith('{"key": "value"')
def test_non_json_content(self):
"""Test with non-JSON content"""
content = "This is just plain text"
result = repair_json_output(content)
assert result == content
def test_empty_string(self):
"""Test with empty string"""
content = ""
result = repair_json_output(content)
assert result == ""
def test_whitespace_only(self):
"""Test with whitespace only"""
content = " \n\t "
result = repair_json_output(content)
assert result == ""
def test_json_with_unicode(self):
"""Test JSON with unicode characters"""
content = '{"name": "测试", "emoji": "🎯"}'
result = repair_json_output(content)
expected = json.dumps({"name": "测试", "emoji": "🎯"}, ensure_ascii=False)
assert result == expected
def test_json_code_block_without_closing(self):
"""Test JSON code block without closing```"""
content = '```json\n{"key": "value"}'
result = repair_json_output(content)
expected = json.dumps({"key": "value"}, ensure_ascii=False)
assert result == expected
def test_json_repair_broken_json(self):
"""Test exception handling when JSON repair fails"""
content = '{"this": "is", "completely": broken and unparseable'
expect = '{"this": "is", "completely": "broken and unparseable"}'
result = repair_json_output(content)
assert result == expect
def test_nested_json_object(self):
"""Test with nested JSON object"""
content = '{"outer": {"inner": {"deep": "value"}}}'
result = repair_json_output(content)
expected = json.dumps(
{"outer": {"inner": {"deep": "value"}}}, ensure_ascii=False
)
assert result == expected
def test_json_array_with_objects(self):
"""Test JSON array containing objects"""
content = '[{"id": 1, "name": "test1"}, {"id": 2, "name": "test2"}]'
result = repair_json_output(content)
expected = json.dumps(
[{"id": 1, "name": "test1"}, {"id": 2, "name": "test2"}], ensure_ascii=False
)
assert result == expected
def test_content_with_json_in_middle(self):
"""Test content that contains ```json in the middle"""
content = 'Some text before ```json {"key": "value"} and after'
result = repair_json_output(content)
# Should attempt to process as JSON since it contains ```json
assert isinstance(result, str)
assert result == '{"key": "value"}'
Generated
+996 -1132
View File
File diff suppressed because it is too large Load Diff
-130
View File
@@ -1,130 +0,0 @@
# 深度思考块功能实现总结
## 🎯 实现的功能
### 核心特性
1. **智能展示逻辑**: 深度思考过程初始展开,计划内容开始时自动折叠
2. **分阶段显示**: 思考阶段只显示思考块,思考结束后才显示计划卡片
3. **动态主题**: 思考阶段使用蓝色主题,完成后切换为默认主题
4. **流式支持**: 实时展示推理内容的流式传输
5. **优雅交互**: 平滑的动画效果和状态切换
### 交互流程
```
用户发送问题 (启用深度思考)
开始接收 reasoning_content
思考块自动展开 + primary 主题 + 加载动画
推理内容流式更新
开始接收 content (计划内容)
思考块自动折叠 + 主题切换
计划卡片优雅出现 (动画效果)
计划内容保持流式更新 (标题→思路→步骤)
完成 (用户可手动展开思考块)
```
## 🔧 技术实现
### 数据结构扩展
- `Message` 接口添加 `reasoningContent``reasoningContentChunks` 字段
- `MessageChunkEvent` 接口添加 `reasoning_content` 字段
- 消息合并逻辑支持推理内容的流式处理
### 组件架构
- `ThoughtBlock`: 可折叠的思考块组件
- `PlanCard`: 更新后的计划卡片,集成思考块
- 智能状态管理和条件渲染
### 状态管理
```typescript
// 关键状态逻辑
const hasMainContent = message.content && message.content.trim() !== "";
const isThinking = reasoningContent && !hasMainContent;
const shouldShowPlan = hasMainContent; // 有内容就显示,保持流式效果
```
### 自动折叠逻辑
```typescript
React.useEffect(() => {
if (hasMainContent && !hasAutoCollapsed) {
setIsOpen(false);
setHasAutoCollapsed(true);
}
}, [hasMainContent, hasAutoCollapsed]);
```
## 🎨 视觉设计
### 统一设计语言
- **字体系统**: 使用 `font-semibold` 与 CardTitle 保持一致
- **圆角规范**: 采用 `rounded-xl` 与其他卡片组件统一
- **间距标准**: 使用 `px-6 py-4` 内边距,`mb-6` 外边距
- **图标尺寸**: 18px 大脑图标,与文字比例协调
### 思考阶段样式
- Primary 主题色边框和背景
- Primary 色图标和文字
- 标准边框样式
- 加载动画
### 完成阶段样式
- 默认 border 和 card 背景
- muted-foreground 图标
- 80% 透明度文字
- 静态图标
### 动画效果
- 展开/折叠动画
- 主题切换过渡
- 颜色变化动画
## 📁 文件更改
### 核心文件
1. `web/src/core/messages/types.ts` - 消息类型扩展
2. `web/src/core/api/types.ts` - API 事件类型扩展
3. `web/src/core/messages/merge-message.ts` - 消息合并逻辑
4. `web/src/core/store/store.ts` - 状态管理更新
5. `web/src/app/chat/components/message-list-view.tsx` - 主要组件实现
### 测试和文档
1. `web/public/mock/reasoning-example.txt` - 测试数据
2. `web/docs/thought-block-feature.md` - 功能文档
3. `web/docs/testing-thought-block.md` - 测试指南
4. `web/docs/interaction-flow-test.md` - 交互流程测试
## 🧪 测试方法
### 快速测试
```
访问: http://localhost:3000?mock=reasoning-example
发送任意消息,观察交互流程
```
### 完整测试
1. 启用深度思考模式
2. 配置 reasoning 模型
3. 发送复杂问题
4. 验证完整交互流程
## 🔄 兼容性
- ✅ 向后兼容:无推理内容时正常显示
- ✅ 渐进增强:功能仅在有推理内容时激活
- ✅ 优雅降级:推理内容为空时不显示思考块
## 🚀 使用建议
1. **启用深度思考**: 点击"Deep Thinking"按钮
2. **观察流程**: 注意思考块的自动展开和折叠
3. **手动控制**: 可随时点击思考块标题栏控制展开/折叠
4. **查看推理**: 展开思考块查看完整的推理过程
这个实现完全满足了用户的需求,提供了直观、流畅的深度思考过程展示体验。
-112
View File
@@ -1,112 +0,0 @@
# 思考块交互流程测试
## 测试场景
### 场景 1: 完整的深度思考流程
**步骤**:
1. 启用深度思考模式
2. 发送问题:"什么是 vibe coding"
3. 观察交互流程
**预期行为**:
#### 阶段 1: 深度思考开始
- ✅ 思考块立即出现并展开
- ✅ 使用蓝色主题(边框、背景、图标、文字)
- ✅ 显示加载动画
- ✅ 不显示计划卡片
- ✅ 推理内容实时流式更新
#### 阶段 2: 思考过程中
- ✅ 思考块保持展开状态
- ✅ 蓝色主题持续显示
- ✅ 推理内容持续增加
- ✅ 加载动画持续显示
- ✅ 计划卡片仍然不显示
#### 阶段 3: 开始接收计划内容
- ✅ 思考块自动折叠
- ✅ 主题从 primary 切换为默认
- ✅ 加载动画消失
- ✅ 计划卡片以优雅动画出现(opacity: 0→1, y: 20→0
- ✅ 计划内容保持流式更新效果
#### 阶段 4: 计划流式输出
- ✅ 标题逐步显示
- ✅ 思路内容流式更新
- ✅ 步骤列表逐项显示
- ✅ 每个步骤的标题和描述分别流式渲染
#### 阶段 5: 计划完成
- ✅ 思考块保持折叠状态
- ✅ 计划卡片完全显示
- ✅ 用户可手动展开思考块查看推理过程
### 场景 2: 手动交互测试
**步骤**:
1. 在思考完成后,手动点击思考块
2. 验证展开/折叠功能
**预期行为**:
- ✅ 点击可正常展开/折叠
- ✅ 动画效果流畅
- ✅ 内容完整显示
- ✅ 不影响计划卡片显示
### 场景 3: 边界情况测试
#### 3.1 只有推理内容,没有计划内容
**预期**: 思考块保持展开,不显示计划卡片
#### 3.2 没有推理内容,只有计划内容
**预期**: 不显示思考块,直接显示计划卡片
#### 3.3 推理内容为空
**预期**: 不显示思考块,直接显示计划卡片
## 验证要点
### 视觉效果
- [ ] Primary 主题色在思考阶段正确显示
- [ ] 主题切换动画流畅
- [ ] 字体权重与 CardTitle 保持一致 (`font-semibold`)
- [ ] 圆角设计与其他卡片统一 (`rounded-xl`)
- [ ] 图标尺寸和颜色正确变化 (18px, primary/muted-foreground)
- [ ] 内边距与设计系统一致 (`px-6 py-4`)
- [ ] 整体视觉层次与页面协调
### 交互逻辑
- [ ] 自动展开/折叠时机正确
- [ ] 手动展开/折叠功能正常
- [ ] 计划卡片显示时机正确
- [ ] 加载动画显示时机正确
### 内容渲染
- [ ] 推理内容正确流式更新
- [ ] Markdown 格式正确渲染
- [ ] 中文内容正确显示
- [ ] 内容不丢失或重复
### 性能表现
- [ ] 动画流畅,无卡顿
- [ ] 内存使用正常
- [ ] 组件重新渲染次数合理
## 故障排除
### 思考块不自动折叠
1. 检查 `hasMainContent` 逻辑
2. 验证 `useEffect` 依赖项
3. 确认 `hasAutoCollapsed` 状态管理
### 计划卡片显示时机错误
1. 检查 `shouldShowPlan` 计算逻辑
2. 验证 `isThinking` 状态判断
3. 确认消息内容解析正确
### 主题切换异常
1. 检查 `isStreaming` 状态
2. 验证 CSS 类名应用
3. 确认条件渲染逻辑
-125
View File
@@ -1,125 +0,0 @@
# 流式输出优化改进
## 🎯 改进目标
确保在深度思考结束后,plan block 保持流式输出效果,提供更流畅丝滑的用户体验。
## 🔧 技术改进
### 状态逻辑优化
**之前的逻辑**:
```typescript
const isThinking = reasoningContent && (!hasMainContent || message.isStreaming);
const shouldShowPlan = hasMainContent && !isThinking;
```
**优化后的逻辑**:
```typescript
const isThinking = reasoningContent && !hasMainContent;
const shouldShowPlan = hasMainContent; // 简化逻辑,有内容就显示
```
### 关键改进点
1. **简化显示逻辑**: 只要有主要内容就显示 plan,不再依赖思考状态
2. **保持流式状态**: plan 组件的 `animated` 属性直接使用 `message.isStreaming`
3. **优雅入场动画**: 添加 motion.div 包装,提供平滑的出现效果
## 🎨 用户体验提升
### 流式输出效果
#### 思考阶段
- ✅ 推理内容实时流式更新
- ✅ 思考块保持展开状态
- ✅ Primary 主题色高亮显示
#### 计划阶段
- ✅ 计划卡片优雅出现(300ms 动画)
- ✅ 标题内容流式渲染
- ✅ 思路内容流式更新
- ✅ 步骤列表逐项显示
- ✅ 每个步骤的标题和描述分别流式渲染
### 动画效果
#### 计划卡片入场动画
```typescript
<motion.div
initial={{ opacity: 0, y: 20 }}
animate={{ opacity: 1, y: 0 }}
transition={{ duration: 0.3, ease: "easeOut" }}
>
```
#### 流式文本动画
- 所有 Markdown 组件都使用 `animated={message.isStreaming}`
- 确保文本逐字符或逐词显示效果
## 📊 性能优化
### 渲染优化
- **减少重新渲染**: 简化状态逻辑,减少不必要的组件重新挂载
- **保持组件实例**: plan 组件一旦出现就保持存在,避免重新创建
- **流式状态传递**: 直接使用消息的流式状态,避免额外的状态计算
### 内存优化
- **组件复用**: 避免频繁的组件销毁和重建
- **状态管理**: 简化状态依赖,减少内存占用
## 🧪 测试验证
### 流式效果验证
1. **思考阶段**: 推理内容应该逐步显示
2. **过渡阶段**: 计划卡片应该平滑出现
3. **计划阶段**: 所有计划内容应该保持流式效果
### 动画效果验证
1. **入场动画**: 计划卡片应该从下方滑入并淡入
2. **文本动画**: 所有文本内容应该有打字机效果
3. **状态切换**: 思考块折叠应该平滑自然
### 性能验证
1. **渲染次数**: 检查组件重新渲染频率
2. **内存使用**: 监控内存占用情况
3. **动画流畅度**: 确保 60fps 的动画效果
## 📝 使用示例
### 完整交互流程
```
1. 用户发送问题 (启用深度思考)
2. 思考块展开,推理内容流式显示
3. 开始接收计划内容
4. 思考块自动折叠
5. 计划卡片优雅出现 (动画效果)
6. 计划内容流式渲染:
- 标题逐步显示
- 思路内容流式更新
- 步骤列表逐项显示
7. 完成,用户可查看完整内容
```
## 🔄 兼容性
-**向后兼容**: 不影响现有的非深度思考模式
-**渐进增强**: 功能仅在有推理内容时激活
-**优雅降级**: 在不支持的环境中正常显示
## 🚀 效果总结
这次优化显著提升了用户体验:
1. **更流畅的过渡**: 从思考到计划的切换更加自然
2. **保持流式效果**: 计划内容保持了原有的流式输出特性
3. **视觉连贯性**: 整个过程的视觉效果更加连贯统一
4. **性能提升**: 减少了不必要的组件重新渲染
用户现在可以享受到完整的流式体验,从深度思考到计划展示都保持了一致的流畅感。
-78
View File
@@ -1,78 +0,0 @@
# 测试思考块功能
## 快速测试
### 方法 1: 使用模拟数据
1. 在浏览器中访问应用并添加 `?mock=reasoning-example` 参数
2. 发送任意消息
3. 观察计划卡片上方是否出现思考块
### 方法 2: 启用深度思考模式
1. 确保配置了 reasoning 模型(如 DeepSeek R1
2. 在聊天界面点击"Deep Thinking"按钮
3. 发送一个需要规划的问题
4. 观察是否出现思考块
## 预期行为
### 思考块外观
- 深度思考开始时自动展开显示
- 思考阶段使用 primary 主题色(边框、背景、文字、图标)
- 带有 18px 大脑图标和"深度思考过程"标题
- 使用 `font-semibold` 字体权重,与 CardTitle 保持一致
- `rounded-xl` 圆角设计,与其他卡片组件统一
- 标准的 `px-6 py-4` 内边距
### 交互行为
- 思考阶段:自动展开,蓝色高亮,显示加载动画
- 计划阶段:自动折叠,切换为默认主题
- 用户可随时手动展开/折叠
- 平滑的展开/折叠动画和主题切换
### 分阶段显示
- 思考阶段:只显示思考块,不显示计划卡片
- 计划阶段:思考块折叠,显示完整计划卡片
### 内容渲染
- 支持 Markdown 格式
- 中文内容正确显示
- 保持原有的换行和格式
## 故障排除
### 思考块不显示
1. 检查消息是否包含 `reasoningContent` 字段
2. 确认 `reasoning_content` 事件是否正确处理
3. 验证消息合并逻辑是否正常工作
### 内容显示异常
1. 检查 Markdown 渲染是否正常
2. 确认 CSS 样式是否正确加载
3. 验证动画效果是否启用
### 流式传输问题
1. 检查 WebSocket 连接状态
2. 确认事件流格式是否正确
3. 验证消息更新逻辑
## 开发调试
### 控制台检查
```javascript
// 检查消息对象
const messages = useStore.getState().messages;
const lastMessage = Array.from(messages.values()).pop();
console.log('Reasoning content:', lastMessage?.reasoningContent);
```
### 网络面板
- 查看 SSE 事件流
- 确认 `reasoning_content` 字段存在
- 检查事件格式是否正确
### React DevTools
- 检查 ThoughtBlock 组件状态
- 验证 props 传递是否正确
- 观察组件重新渲染情况

Some files were not shown because too many files have changed in this diff Show More