Skip to content

Commit b886b16

Browse files
committed
feat: support time series forecasting tool in agent
1 parent efbb0dc commit b886b16

42 files changed

Lines changed: 474 additions & 130 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

README.md

Lines changed: 23 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -14,18 +14,19 @@ it provides chat agents and workflows that support natural language to SQL conve
1414
3. **Data Visualization**: Generate intuitive data visualizations (via plotly)
1515
4. **Data Catalog Management**: Automatically discovers and indexes database table structures, supports flexible catalog
1616
storage backends, and easily maintains business explanations for tables and columns as well as optimizes Prompts.
17-
5. **Knowledge Base Integration**: Answer complex questions by combining catalog based knowledge retrival and external
18-
knowledge base retrival (via MCP tools)
17+
5. **Time Series Forecasting**: Forecasting models deployed in-house that can be called as tools
1918
6. **Code Execution**: Execute Python code for data analysis and visualization
2019
7. **Interactive Problem-Solving**: Proactively ask users for more context when information is incomplete
2120
8. **Persistent Memory**: Conversation management and user characteristic memory based on LangGraph checkpointing
2221
9. **MCP Support**: Integration with MCP tools by configuration
23-
10. **Web UI Interface**: Provide 2 sample UI: simple and streaming web interfaces using Gradio and Streamlit, easy to
22+
10. **Knowledge Base Integration**: Answer complex questions by combining catalog based knowledge retrival and external
23+
knowledge base retrival (via MCP tools)
24+
11. **Web UI Interface**: Provide 2 sample UI: simple and streaming web interfaces using Gradio and Streamlit, easy to
2425
integrate with other web applications
2526

2627
## Roadmap
2728

28-
1. **Time Series Forecasting**: Forecasting models deployed in-house
29+
1. **Anomaly Detection Algorithm**: Time series anomaly detection
2930
2. **Root Cause Analysis Algorithm**: Multi-dimensional drill-down capabilities for anomaly investigation
3031

3132
# Getting started
@@ -219,6 +220,7 @@ OpenChatBI supports sophisticated customization through prompt engineering and c
219220
- **Prompt Engineering Configuration**: Customize system prompts, business glossaries, and data warehouse introductions
220221
- **Data Catalog Management**: Configure table metadata, column descriptions, and SQL generation rules
221222
- **Business Rules**: Define table selection criteria and domain-specific SQL constraints
223+
- **Forecasting Service**: Configure the forecasting service url and prompt based on your own deployment
222224

223225
For detailed configuration options and examples, see the [Advanced Features](#advanced-features) section.
224226

@@ -264,16 +266,21 @@ openchatbi/
264266
│ ├── constants.py # Application constants
265267
│ ├── agent_graph.py # Main LangGraph workflow
266268
│ ├── graph_state.py # State definition for workflows
269+
│ ├── context_config.py # Context management configuration
270+
│ ├── context_manager.py # Context window and token management
271+
│ ├── text_segmenter.py # Text segmentation utilities
267272
│ ├── utils.py # Utility functions
268273
│ ├── catalog/ # Data catalog management
269274
│ │ ├── __init__.py # Package initialization
270275
│ │ ├── catalog_loader.py # Catalog loading logic
271276
│ │ ├── catalog_store.py # Catalog storage interface
272-
│ │ ├── entry.py # Catalog entry points
273277
│ │ ├── factory.py # Catalog factory patterns
274278
│ │ ├── helper.py # Catalog helper functions
279+
│ │ ├── retrival_helper.py # Retrieval helper utilities
275280
│ │ ├── schema_retrival.py # Schema retrieval logic
276-
│ │ └── token_service.py # Token service integration
281+
│ │ ├── token_service.py # Token service integration
282+
│ │ └── store/ # Catalog storage implementations
283+
│ │ └── file_system.py # File system-based catalog storage
277284
│ ├── code/ # Code execution framework
278285
│ │ ├── __init__.py # Package initialization
279286
│ │ ├── executor_base.py # Base executor interface
@@ -290,7 +297,8 @@ openchatbi/
290297
│ │ ├── system_prompt.py # System prompt management
291298
│ │ ├── summary_prompt.md # Summary conversation prompts
292299
│ │ ├── table_selection_prompt.md # Table selection prompts
293-
│ │ └── text2sql_prompt.md # Text-to-SQL prompts
300+
│ │ ├── text2sql_prompt.md # Text-to-SQL prompts
301+
│ │ └── sql_dialect/ # SQL dialect-specific prompts
294302
│ ├── text2sql/ # Text-to-SQL conversion pipeline
295303
│ │ ├── __init__.py # Package initialization
296304
│ │ ├── data.py # Data and retriever for Text-to-SQL
@@ -306,7 +314,8 @@ openchatbi/
306314
│ ├── mcp_tools.py # MCP (Model Context Protocol) integration
307315
│ ├── run_python_code.py # Configurable Python code execution
308316
│ ├── save_report.py # Report saving functionality
309-
│ └── search_knowledge.py # Knowledge base search
317+
│ ├── search_knowledge.py # Knowledge base search
318+
│ └── timeseries_forecast.py # Time series forecasting tool
310319
├── sample_api/ # API implementations
311320
│ └── async_api.py # Asynchronous FastAPI example
312321
├── sample_ui/ # Web interface implementations
@@ -325,6 +334,9 @@ openchatbi/
325334
│ ├── sql_example.yaml # SQL examples for retrieval
326335
│ ├── table_selection_example.csv # Table selection examples
327336
│ └── tracking_orders.sqlite # Sample SQLite database
337+
├── timeseries_forecasting/ # Time series forecasting service
338+
│ ├── README.md # Forecasting service documentation
339+
│ └── ... # Forecasting service implementation
328340
├── tests/ # Test suite
329341
│ ├── __init__.py # Package initialization
330342
│ ├── conftest.py # Test configuration
@@ -489,6 +501,9 @@ The file system catalog store organizes metadata across multiple files for maint
489501
ORDER BY date;
490502
```
491503

504+
### Time Series Forecasting Service Setup
505+
1. Setup forecasting docker container, see [README](timeseries_forecasting/README.md)
506+
2. Config the `timeseries_forecast` tool prompt
492507

493508
### Python Code Execution Configuration
494509

@@ -535,7 +550,6 @@ For production deployments or when running untrusted code, the Docker executor p
535550
- Automatic container cleanup
536551
- Resource isolation from host system
537552

538-
539553
## Development & Testing
540554

541555
### Code Quality Tools

docs/source/conf.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
project = "OpenChatBI"
1414
copyright = "2025, Yu Zhong"
1515
author = "Yu Zhong"
16-
release = "0.1.0"
16+
release = "0.2.0"
1717

1818
# -- General configuration ---------------------------------------------------
1919
# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

docs/source/index.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ OpenChatBI Documentation
2424
Code Execution <code>
2525
LLM Integration <llm>
2626
Tools and Utilities <tools>
27+
Time Series Forecasting Service <timeseries>
2728

2829
Indices and tables
2930
==================

docs/source/timeseries.rst

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
Time Series Forecasting Service
2+
========================
3+
4+
`GitHub Repository <https://github.com/zhongyu09/openchatbi/timeseries_forecasting>`_
5+
6+
.. include:: ../../timeseries_forecasting/README.md
7+
:parser: myst_parser.sphinx_

example/bi.yaml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
1-
extra_tool_use_rule: |
1+
extra_tool_use_rule: |
22
- Try your best to give appropriate parameters when calling tools.
3+
- timeseries_forecast tool requires at least 96 time points in input data. If no enough input data, set input_len to 96 to pad with zeros.
34
45
table_selection_extra_rule: |
56
- When users ask about orders, consider if they need customer information (join with Customers table)

openchatbi/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ def get_default_graph():
2626
return None
2727

2828
from langgraph.checkpoint.memory import MemorySaver
29+
2930
from openchatbi.agent_graph import build_agent_graph_sync
3031
from openchatbi.tool.memory import get_sync_memory_store
3132

openchatbi/agent_graph.py

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from typing import Any
88

99
from langchain_core.language_models import BaseChatModel
10-
from langchain_core.messages import AIMessage, SystemMessage, HumanMessage
10+
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
1111
from langchain_core.tools import StructuredTool
1212
from langchain_openai.chat_models.base import BaseChatOpenAI
1313
from langgraph.constants import START
@@ -34,8 +34,8 @@
3434
from openchatbi.tool.run_python_code import run_python_code
3535
from openchatbi.tool.save_report import save_report
3636
from openchatbi.tool.search_knowledge import search_knowledge, show_schema
37-
from openchatbi.utils import log
38-
from openchatbi.utils import recover_incomplete_tool_calls
37+
from openchatbi.tool.timeseries_forecast import check_forecast_service_health, timeseries_forecast
38+
from openchatbi.utils import log, recover_incomplete_tool_calls
3939

4040
logger = logging.getLogger(__name__)
4141

@@ -313,7 +313,12 @@ def _build_graph_core(
313313
manage_memory_tool,
314314
search_memory_tool,
315315
save_report,
316-
] + mcp_tools
316+
]
317+
if check_forecast_service_health:
318+
normal_tools.append(timeseries_forecast)
319+
else:
320+
logger.warning("Time series forecasting service is not healthy. Skipping timeseries_forecast tool.")
321+
normal_tools.extend(mcp_tools)
317322

318323
# Initialize context manager if enabled
319324
context_manager = None

openchatbi/catalog/retrival_helper.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@
33
from rank_bm25 import BM25Okapi
44

55
from openchatbi.llm.llm import get_embedding_model
6-
from openchatbi.utils import log, create_vector_db
76
from openchatbi.text_segmenter import _segmenter
7+
from openchatbi.utils import create_vector_db, log
88

99

1010
def get_columns_metadata(catalog):

openchatbi/catalog/schema_retrival.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@
66
import Levenshtein
77

88
from openchatbi import config
9-
from openchatbi.text_segmenter import _segmenter
109
from openchatbi.catalog.retrival_helper import build_column_tables_mapping, build_columns_retriever
10+
from openchatbi.text_segmenter import _segmenter
1111
from openchatbi.utils import log
1212

1313
# Skip build during documentation build

openchatbi/code/docker_executor.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,16 @@
1-
import docker
2-
import tempfile
31
import os
4-
import subprocess
52
import shutil
3+
import subprocess
4+
import tempfile
65
from pathlib import Path
7-
from typing import Tuple
6+
7+
import docker
88
from docker.errors import ContainerError
99

1010
from openchatbi.code.executor_base import ExecutorBase
1111

1212

13-
def check_docker_status() -> Tuple[bool, str]:
13+
def check_docker_status() -> tuple[bool, str]:
1414
"""
1515
Check Docker installation and status without initializing DockerExecutor.
1616
@@ -118,7 +118,7 @@ def _ensure_image_exists(self):
118118
)
119119
print(f"Docker image '{self.image_name}' built successfully.")
120120

121-
def run_code(self, code: str) -> Tuple[bool, str]:
121+
def run_code(self, code: str) -> tuple[bool, str]:
122122
"""Execute Python code in a Docker container."""
123123
try:
124124
# Create a temporary file with the code

0 commit comments

Comments
 (0)