agent-brain-rag 1.2.0__py3-none-any.whl → 3.0.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {agent_brain_rag-1.2.0.dist-info → agent_brain_rag-3.0.0.dist-info}/METADATA +55 -18
- agent_brain_rag-3.0.0.dist-info/RECORD +56 -0
- {agent_brain_rag-1.2.0.dist-info → agent_brain_rag-3.0.0.dist-info}/WHEEL +1 -1
- {agent_brain_rag-1.2.0.dist-info → agent_brain_rag-3.0.0.dist-info}/entry_points.txt +0 -1
- agent_brain_server/__init__.py +1 -1
- agent_brain_server/api/main.py +146 -45
- agent_brain_server/api/routers/__init__.py +2 -0
- agent_brain_server/api/routers/health.py +85 -21
- agent_brain_server/api/routers/index.py +108 -36
- agent_brain_server/api/routers/jobs.py +111 -0
- agent_brain_server/config/provider_config.py +352 -0
- agent_brain_server/config/settings.py +22 -5
- agent_brain_server/indexing/__init__.py +21 -0
- agent_brain_server/indexing/bm25_index.py +15 -2
- agent_brain_server/indexing/document_loader.py +45 -4
- agent_brain_server/indexing/embedding.py +86 -135
- agent_brain_server/indexing/graph_extractors.py +582 -0
- agent_brain_server/indexing/graph_index.py +536 -0
- agent_brain_server/job_queue/__init__.py +11 -0
- agent_brain_server/job_queue/job_service.py +317 -0
- agent_brain_server/job_queue/job_store.py +427 -0
- agent_brain_server/job_queue/job_worker.py +434 -0
- agent_brain_server/locking.py +101 -8
- agent_brain_server/models/__init__.py +28 -0
- agent_brain_server/models/graph.py +253 -0
- agent_brain_server/models/health.py +30 -3
- agent_brain_server/models/job.py +289 -0
- agent_brain_server/models/query.py +16 -3
- agent_brain_server/project_root.py +1 -1
- agent_brain_server/providers/__init__.py +64 -0
- agent_brain_server/providers/base.py +251 -0
- agent_brain_server/providers/embedding/__init__.py +23 -0
- agent_brain_server/providers/embedding/cohere.py +163 -0
- agent_brain_server/providers/embedding/ollama.py +150 -0
- agent_brain_server/providers/embedding/openai.py +118 -0
- agent_brain_server/providers/exceptions.py +95 -0
- agent_brain_server/providers/factory.py +157 -0
- agent_brain_server/providers/summarization/__init__.py +41 -0
- agent_brain_server/providers/summarization/anthropic.py +87 -0
- agent_brain_server/providers/summarization/gemini.py +96 -0
- agent_brain_server/providers/summarization/grok.py +95 -0
- agent_brain_server/providers/summarization/ollama.py +114 -0
- agent_brain_server/providers/summarization/openai.py +87 -0
- agent_brain_server/runtime.py +2 -2
- agent_brain_server/services/indexing_service.py +39 -0
- agent_brain_server/services/query_service.py +203 -0
- agent_brain_server/storage/__init__.py +18 -2
- agent_brain_server/storage/graph_store.py +519 -0
- agent_brain_server/storage/vector_store.py +35 -0
- agent_brain_server/storage_paths.py +5 -3
- agent_brain_rag-1.2.0.dist-info/RECORD +0 -31
|
@@ -1,7 +1,8 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
2
|
Name: agent-brain-rag
|
|
3
|
-
Version:
|
|
3
|
+
Version: 3.0.0
|
|
4
4
|
Summary: Agent Brain RAG - Intelligent document indexing and semantic search server that gives AI agents long-term memory
|
|
5
|
+
Home-page: https://github.com/SpillwaveSolutions/agent-brain
|
|
5
6
|
License: MIT
|
|
6
7
|
Keywords: agent-brain,rag,semantic-search,ai-memory,llm-memory,documentation,indexing,llama-index,chromadb,ai-agent,claude-code,agent-memory
|
|
7
8
|
Author: Spillwave Solutions
|
|
@@ -13,15 +14,21 @@ Classifier: Programming Language :: Python :: 3
|
|
|
13
14
|
Classifier: Programming Language :: Python :: 3.10
|
|
14
15
|
Classifier: Programming Language :: Python :: 3.11
|
|
15
16
|
Classifier: Programming Language :: Python :: 3.12
|
|
16
|
-
Classifier: Programming Language :: Python :: 3.13
|
|
17
17
|
Classifier: Topic :: Software Development :: Documentation
|
|
18
18
|
Classifier: Topic :: Text Processing :: Indexing
|
|
19
|
+
Provides-Extra: graphrag
|
|
20
|
+
Provides-Extra: graphrag-all
|
|
21
|
+
Provides-Extra: graphrag-kuzu
|
|
19
22
|
Requires-Dist: anthropic (>=0.40.0,<0.41.0)
|
|
20
23
|
Requires-Dist: chromadb (>=0.5.0,<0.6.0)
|
|
21
24
|
Requires-Dist: click (>=8.1.0,<9.0.0)
|
|
25
|
+
Requires-Dist: cohere (>=5.0.0,<6.0.0)
|
|
22
26
|
Requires-Dist: fastapi (>=0.115.0,<0.116.0)
|
|
27
|
+
Requires-Dist: google-generativeai (>=0.8.0,<0.9.0)
|
|
28
|
+
Requires-Dist: langextract (>=1.0.0,<2.0.0) ; extra == "graphrag" or extra == "graphrag-all"
|
|
23
29
|
Requires-Dist: llama-index-core (>=0.14.0,<0.15.0)
|
|
24
30
|
Requires-Dist: llama-index-embeddings-openai (>=0.5.0,<0.6.0)
|
|
31
|
+
Requires-Dist: llama-index-graph-stores-kuzu (>=0.9.0,<0.10.0) ; extra == "graphrag-kuzu" or extra == "graphrag-all"
|
|
25
32
|
Requires-Dist: llama-index-llms-openai (>=0.6.12,<0.7.0)
|
|
26
33
|
Requires-Dist: llama-index-readers-file (>=0.5.0,<0.6.0)
|
|
27
34
|
Requires-Dist: llama-index-retrievers-bm25 (>=0.6.0,<0.7.0)
|
|
@@ -29,18 +36,24 @@ Requires-Dist: openai (>=1.57.0,<2.0.0)
|
|
|
29
36
|
Requires-Dist: pydantic (>=2.10.0,<3.0.0)
|
|
30
37
|
Requires-Dist: pydantic-settings (>=2.6.0,<3.0.0)
|
|
31
38
|
Requires-Dist: python-dotenv (>=1.0.0,<2.0.0)
|
|
39
|
+
Requires-Dist: pyyaml (>=6.0.0,<7.0.0)
|
|
32
40
|
Requires-Dist: rank-bm25 (>=0.2.2,<0.3.0)
|
|
33
41
|
Requires-Dist: tiktoken (>=0.8.0,<0.9.0)
|
|
34
42
|
Requires-Dist: tree-sitter-language-pack (>=0.7.3,<0.8.0)
|
|
35
43
|
Requires-Dist: uvicorn[standard] (>=0.32.0,<0.33.0)
|
|
36
|
-
Project-URL: Documentation, https://github.com/SpillwaveSolutions/
|
|
37
|
-
Project-URL:
|
|
38
|
-
Project-URL: Repository, https://github.com/SpillwaveSolutions/doc-serve
|
|
44
|
+
Project-URL: Documentation, https://github.com/SpillwaveSolutions/agent-brain/wiki
|
|
45
|
+
Project-URL: Repository, https://github.com/SpillwaveSolutions/agent-brain
|
|
39
46
|
Description-Content-Type: text/markdown
|
|
40
47
|
|
|
41
48
|
# Agent Brain RAG Server
|
|
42
49
|
|
|
43
|
-
|
|
50
|
+
> **Agent Brain** (formerly doc-serve) is an intelligent document indexing and semantic search system designed to give AI agents long-term memory.
|
|
51
|
+
|
|
52
|
+
AI agents need persistent memory to be truly useful. Agent Brain provides the retrieval infrastructure that enables context-aware, knowledge-grounded AI interactions.
|
|
53
|
+
|
|
54
|
+
[](https://pypi.org/project/agent-brain-rag/)
|
|
55
|
+
[](https://www.python.org/downloads/)
|
|
56
|
+
[](https://opensource.org/licenses/MIT)
|
|
44
57
|
|
|
45
58
|
## Installation
|
|
46
59
|
|
|
@@ -65,28 +78,37 @@ The server will start at `http://127.0.0.1:8000`.
|
|
|
65
78
|
|
|
66
79
|
> **Note**: The legacy command `doc-serve` is still available but deprecated. Please use `agent-brain-serve` for new installations.
|
|
67
80
|
|
|
81
|
+
## Search Capabilities
|
|
82
|
+
|
|
83
|
+
Agent Brain provides multiple search strategies to match your retrieval needs:
|
|
84
|
+
|
|
85
|
+
| Search Type | Description | Best For |
|
|
86
|
+
|-------------|-------------|----------|
|
|
87
|
+
| **Semantic Search** | Natural language queries using OpenAI embeddings (`text-embedding-3-large`) | Conceptual questions, finding related content |
|
|
88
|
+
| **Keyword Search (BM25)** | Traditional keyword matching with TF-IDF ranking | Exact matches, technical terms, code identifiers |
|
|
89
|
+
| **Hybrid Search** | Combines vector + BM25 for best of both approaches | General-purpose queries, balanced recall/precision |
|
|
90
|
+
| **GraphRAG** | Knowledge graph-based retrieval for relationship-aware queries | Understanding connections, multi-hop reasoning |
|
|
91
|
+
|
|
68
92
|
## Features
|
|
69
93
|
|
|
70
94
|
- **Document Indexing**: Load and index documents from folders (PDF, Markdown, TXT, DOCX, HTML)
|
|
71
95
|
- **AST-Aware Code Ingestion**: Smart parsing for Python, TypeScript, JavaScript, Java, Go, Rust, C, C++
|
|
72
|
-
- **
|
|
73
|
-
- **Semantic Search**: Query indexed documents using natural language
|
|
96
|
+
- **Multi-Strategy Retrieval**: Semantic, keyword, hybrid, and graph-based search
|
|
74
97
|
- **OpenAI Embeddings**: Uses `text-embedding-3-large` for high-quality embeddings
|
|
98
|
+
- **Claude Summarization**: AI-powered code summaries for better context
|
|
75
99
|
- **Chroma Vector Store**: Persistent, thread-safe vector database
|
|
76
100
|
- **FastAPI**: Modern, high-performance REST API with OpenAPI documentation
|
|
77
101
|
|
|
78
|
-
##
|
|
79
|
-
|
|
80
|
-
### Prerequisites
|
|
102
|
+
## Prerequisites
|
|
81
103
|
|
|
82
104
|
- Python 3.10+
|
|
83
|
-
-
|
|
84
|
-
-
|
|
105
|
+
- OpenAI API key (for embeddings)
|
|
106
|
+
- Anthropic API key (for summarization)
|
|
85
107
|
|
|
86
|
-
|
|
108
|
+
## Development Installation
|
|
87
109
|
|
|
88
110
|
```bash
|
|
89
|
-
cd
|
|
111
|
+
cd agent-brain-server
|
|
90
112
|
poetry install
|
|
91
113
|
```
|
|
92
114
|
|
|
@@ -101,6 +123,7 @@ cp ../.env.example .env
|
|
|
101
123
|
|
|
102
124
|
Required environment variables:
|
|
103
125
|
- `OPENAI_API_KEY`: Your OpenAI API key for embeddings
|
|
126
|
+
- `ANTHROPIC_API_KEY`: Your Anthropic API key for summarization
|
|
104
127
|
|
|
105
128
|
### Running the Server
|
|
106
129
|
|
|
@@ -112,8 +135,6 @@ poetry run uvicorn agent_brain_server.api.main:app --reload
|
|
|
112
135
|
poetry run agent-brain-serve
|
|
113
136
|
```
|
|
114
137
|
|
|
115
|
-
The server will start at `http://127.0.0.1:8000`.
|
|
116
|
-
|
|
117
138
|
### API Documentation
|
|
118
139
|
|
|
119
140
|
Once running, visit:
|
|
@@ -199,6 +220,22 @@ poetry run ruff check agent_brain_server/
|
|
|
199
220
|
poetry run mypy agent_brain_server/
|
|
200
221
|
```
|
|
201
222
|
|
|
223
|
+
## Documentation
|
|
224
|
+
|
|
225
|
+
- [User Guide](https://github.com/SpillwaveSolutions/agent-brain/wiki/User-Guide) - Getting started and usage
|
|
226
|
+
- [Developer Guide](https://github.com/SpillwaveSolutions/agent-brain/wiki/Developer-Guide) - Contributing and development
|
|
227
|
+
- [API Reference](https://github.com/SpillwaveSolutions/agent-brain/wiki/API-Reference) - Full API documentation
|
|
228
|
+
|
|
229
|
+
## Release Information
|
|
230
|
+
|
|
231
|
+
- **Current Version**: See [pyproject.toml](./pyproject.toml)
|
|
232
|
+
- **Release Notes**: [GitHub Releases](https://github.com/SpillwaveSolutions/agent-brain/releases)
|
|
233
|
+
- **Changelog**: [Latest Release](https://github.com/SpillwaveSolutions/agent-brain/releases/latest)
|
|
234
|
+
|
|
235
|
+
## Related Packages
|
|
236
|
+
|
|
237
|
+
- [agent-brain-cli](https://pypi.org/project/agent-brain-cli/) - Command-line interface for Agent Brain
|
|
238
|
+
|
|
202
239
|
## License
|
|
203
240
|
|
|
204
241
|
MIT
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
agent_brain_server/__init__.py,sha256=Wjd2w3VnDYURTECaKxFQ5x324-7tZ_vK2Qv3yKfxRGg,95
|
|
2
|
+
agent_brain_server/api/__init__.py,sha256=nvTvO_ahHAAsRDlV3dL_JlNruSdan4kav_P5sT_1PFk,93
|
|
3
|
+
agent_brain_server/api/main.py,sha256=TEb4lMSFm8eqNYOhXvoF6Xwi_r20RMuO8HL0zQxKfhY,16086
|
|
4
|
+
agent_brain_server/api/routers/__init__.py,sha256=EQ5xXSTEIqRaZu-CKiOxeR05BMzUVgATMir_PV6zJ-Q,313
|
|
5
|
+
agent_brain_server/api/routers/health.py,sha256=lTk0VXvTIvAV8YyOKAcPvZXhBeZ8jnENCDkZG3xk9Bg,6015
|
|
6
|
+
agent_brain_server/api/routers/index.py,sha256=89_SrTByqEhQuQxsP7SXWiocXgQh2sgJ9n0Y_kNEWjA,9397
|
|
7
|
+
agent_brain_server/api/routers/jobs.py,sha256=GYoZtTUc2BE2_u1HoUSGqUVDg7MAjhs6dG9-gjKoztw,3329
|
|
8
|
+
agent_brain_server/api/routers/query.py,sha256=TAbpuCybjacUo_-7Zm3Jt8EW8EMDTYJDLv6aGtjaUzQ,2793
|
|
9
|
+
agent_brain_server/config/__init__.py,sha256=zzDErZGUBwUm5Fk43OHJN7eWpeIw_1kWdnhsN6QQqSc,84
|
|
10
|
+
agent_brain_server/config/provider_config.py,sha256=-uHj28i8OvAhQxlZ0JudYPAyibdJWmEJfvOYuRAdSHw,11535
|
|
11
|
+
agent_brain_server/config/settings.py,sha256=XfkAi-5aV_vqbk4ltkG8mkJnPARPwl5g6JxvE0ngDjo,3726
|
|
12
|
+
agent_brain_server/indexing/__init__.py,sha256=tQiFW4945h-ENd0pmMoGhCbm1Fz3AaREKBgajHeW-Us,1169
|
|
13
|
+
agent_brain_server/indexing/bm25_index.py,sha256=_CZdTF1zoLyILTJ1nTLSmovIqN2NVQuYUiMEWh4o3-o,5740
|
|
14
|
+
agent_brain_server/indexing/chunking.py,sha256=kW-ZOdwOS-PfXuPAYb8qL4-7yAY-rpY9a1el2YCN9CA,29848
|
|
15
|
+
agent_brain_server/indexing/document_loader.py,sha256=l43Z5xHEykr6hbQqkMWXZGlT_CDz_tdhkI3PsxTMsgQ,17543
|
|
16
|
+
agent_brain_server/indexing/embedding.py,sha256=i3bN3SivvVB8iNiYDjrDeB8nF3WkFoOy5wHrDaNosGs,7410
|
|
17
|
+
agent_brain_server/indexing/graph_extractors.py,sha256=XEWBqlqDffxbs5ZLaijGICXeaptGtkUGNEhJblP3jBU,19339
|
|
18
|
+
agent_brain_server/indexing/graph_index.py,sha256=ZKwICIqR2stcPEqNQ_SSuSk-O0Hj1lUmiUMxm5udhjg,18584
|
|
19
|
+
agent_brain_server/job_queue/__init__.py,sha256=8aUjeXygbB_u76AqPTrlrwFDVVdqfiaWM9ROUAslKM4,240
|
|
20
|
+
agent_brain_server/job_queue/job_service.py,sha256=BAzG_NMCnbHNHluKtKI1lMJxl1MF03G7Rv9mKEcgkWc,10660
|
|
21
|
+
agent_brain_server/job_queue/job_store.py,sha256=xHWZh9rITUeCIAZwWmoEYDtSGCEfQuveduB9ET3L4zU,13699
|
|
22
|
+
agent_brain_server/job_queue/job_worker.py,sha256=bh7pV9XKiCrK_bc1jeESh35DEcgfsiiAsrXo_ihwcLw,17566
|
|
23
|
+
agent_brain_server/locking.py,sha256=ANgzY8aczB6W9OeWEv8zEZ8iOhXlBCOYlZH60xZsiHU,6249
|
|
24
|
+
agent_brain_server/models/__init__.py,sha256=8tCr8Gmx6ieVx-CgX6W_a7s8doqAzWtDs8Jk5lrYmRA,1105
|
|
25
|
+
agent_brain_server/models/graph.py,sha256=YzdjJj_X7hueJt4v8Gsoly4C-qquJLZi1d9epIh_Mos,7973
|
|
26
|
+
agent_brain_server/models/health.py,sha256=ppvDF_yUaddzYPe0BPc9emVDrR17TVlQLW6ggsOOzPo,4592
|
|
27
|
+
agent_brain_server/models/index.py,sha256=pjDv7phLS6dpiHLlEtcAuXQN_aHIfA_4lMkAZ-NkXZQ,5400
|
|
28
|
+
agent_brain_server/models/job.py,sha256=OkNO8tHFFyV82kLJAocs4jKKI7Hde93y7KH8q_mfVyI,11194
|
|
29
|
+
agent_brain_server/models/query.py,sha256=hCdcXIOPhW4wXrSQGunEXkIgiugD3lBXnS4F6fOCe_Q,6613
|
|
30
|
+
agent_brain_server/project_root.py,sha256=js9ju-AhQ2SyC6wlfFhbMHhkNMJ3VNDdGZ3OJONUFjs,2196
|
|
31
|
+
agent_brain_server/providers/__init__.py,sha256=04wJ_QhGKRM55R4yrbw3J3kCtKX44_FHzh6xZq1xFdc,1670
|
|
32
|
+
agent_brain_server/providers/base.py,sha256=yLIpz9W0xlilBXxuFJRNP_gbGDsZ9HeRSqEnIsVGQjU,6933
|
|
33
|
+
agent_brain_server/providers/embedding/__init__.py,sha256=WRUAkOLLgQB9AR9I-_m3QWa55hSeC5tDdp3gZQFfMRM,982
|
|
34
|
+
agent_brain_server/providers/embedding/cohere.py,sha256=9ktr0_ZHmocJE8sFqo1Ah68B9J0wgs9-pLS0aKmLtQE,5295
|
|
35
|
+
agent_brain_server/providers/embedding/ollama.py,sha256=oXKbbEivY74ASVjayLdr3xfgXTMvFS_jAjZs5HqCHzc,4725
|
|
36
|
+
agent_brain_server/providers/embedding/openai.py,sha256=B5F6OJ_PSpF7OLUlhhmqGWOvho3Kox29Lq52bvxrWWY,3576
|
|
37
|
+
agent_brain_server/providers/exceptions.py,sha256=kYggL4lG1NmJrEyU040Al5rm3_nWmKROJemkjcdohCM,2836
|
|
38
|
+
agent_brain_server/providers/factory.py,sha256=3sREN_VnTDam9MDQJzf7eYlX5k-tbzgWXHb7DQ6MhRw,5488
|
|
39
|
+
agent_brain_server/providers/summarization/__init__.py,sha256=WfMXywPvY1uTbmy_Z5prK9FNNqv17TeEvdg2MqFyKAc,1545
|
|
40
|
+
agent_brain_server/providers/summarization/anthropic.py,sha256=fZOUfxJCPRUxMdeakBAdw1FY6NaYhvlWBlheTJlyAtc,2704
|
|
41
|
+
agent_brain_server/providers/summarization/gemini.py,sha256=t2pOC4R2ae7WQAq7WCy8ZD6Ku_RPjyqsktBZRuRwGi8,2960
|
|
42
|
+
agent_brain_server/providers/summarization/grok.py,sha256=nK3i46NbUcVxBqWNbJ-IifahB450Blnqm0rBhSQnQOw,2875
|
|
43
|
+
agent_brain_server/providers/summarization/ollama.py,sha256=mw4zvx6o2nC4xoz1W6F47x3fk7gidl7aeh1uAnSTHO8,3673
|
|
44
|
+
agent_brain_server/providers/summarization/openai.py,sha256=HjZJIjwc0usHejRAFW3JT6a7cmM_Ib2FdC05y_iEwA4,2632
|
|
45
|
+
agent_brain_server/runtime.py,sha256=SVQa-8rgNFBeHyJe23bvzZ49Gf-4PqkuSJ-kJ1FPkSM,3081
|
|
46
|
+
agent_brain_server/services/__init__.py,sha256=E4VPN9Rqa2mxGQQEQn-5IYj63LSPTrA8aIx8ENO5xcc,296
|
|
47
|
+
agent_brain_server/services/indexing_service.py,sha256=d_oj3n7gNgua7nmiLNTM-hDp4wvcrkos9kCRuOTzlpc,20338
|
|
48
|
+
agent_brain_server/services/query_service.py,sha256=oc_BNqTNjAtz17y_F7r-egP7Z0GozmP_mlfSs1jKbMM,22910
|
|
49
|
+
agent_brain_server/storage/__init__.py,sha256=xYk4MfrEj-IpGZ0n4myfy5dwyjKEvprS7h68bet0ooU,557
|
|
50
|
+
agent_brain_server/storage/graph_store.py,sha256=boW6_L0suvHuOdfZdPMdSNSb8U-am7KLelsa3K2bKl4,18168
|
|
51
|
+
agent_brain_server/storage/vector_store.py,sha256=rcI--g6g7mqUbZzewOsaMdb7uqw173bAZau1p0oYZ68,12135
|
|
52
|
+
agent_brain_server/storage_paths.py,sha256=FOpcFu21mlxPY-TCjWg6bMmREYEYDkVru7NumztGZWw,1850
|
|
53
|
+
agent_brain_rag-3.0.0.dist-info/METADATA,sha256=C5ad5HVow94PzzSYAYcHx2gR4z9NelnZk3D7SIqwC7g,8092
|
|
54
|
+
agent_brain_rag-3.0.0.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
|
|
55
|
+
agent_brain_rag-3.0.0.dist-info/entry_points.txt,sha256=Ig6JVYNKNNNnE96d8LefHRUdg11HUIGWfVQ1sAYI47c,69
|
|
56
|
+
agent_brain_rag-3.0.0.dist-info/RECORD,,
|
agent_brain_server/__init__.py
CHANGED
agent_brain_server/api/main.py
CHANGED
|
@@ -1,15 +1,16 @@
|
|
|
1
1
|
"""FastAPI application entry point.
|
|
2
2
|
|
|
3
3
|
This module provides the Agent Brain RAG server, a FastAPI application
|
|
4
|
-
for document indexing and semantic search.
|
|
5
|
-
|
|
4
|
+
for document indexing and semantic search.
|
|
5
|
+
|
|
6
|
+
Note: This server assumes a single uvicorn worker process. If running
|
|
7
|
+
multiple workers, ensure only one worker handles indexing jobs by using
|
|
8
|
+
the single-worker model or a separate job processor service.
|
|
6
9
|
"""
|
|
7
10
|
|
|
8
11
|
import logging
|
|
9
12
|
import os
|
|
10
13
|
import socket
|
|
11
|
-
import sys
|
|
12
|
-
import warnings
|
|
13
14
|
from collections.abc import AsyncIterator
|
|
14
15
|
from contextlib import asynccontextmanager
|
|
15
16
|
from pathlib import Path
|
|
@@ -22,7 +23,13 @@ from fastapi.middleware.cors import CORSMiddleware
|
|
|
22
23
|
|
|
23
24
|
from agent_brain_server import __version__
|
|
24
25
|
from agent_brain_server.config import settings
|
|
26
|
+
from agent_brain_server.config.provider_config import (
|
|
27
|
+
clear_settings_cache,
|
|
28
|
+
load_provider_settings,
|
|
29
|
+
validate_provider_config,
|
|
30
|
+
)
|
|
25
31
|
from agent_brain_server.indexing.bm25_index import BM25IndexManager
|
|
32
|
+
from agent_brain_server.job_queue import JobQueueService, JobQueueStore, JobWorker
|
|
26
33
|
from agent_brain_server.locking import (
|
|
27
34
|
acquire_lock,
|
|
28
35
|
cleanup_stale,
|
|
@@ -35,7 +42,7 @@ from agent_brain_server.services import IndexingService, QueryService
|
|
|
35
42
|
from agent_brain_server.storage import VectorStoreManager
|
|
36
43
|
from agent_brain_server.storage_paths import resolve_state_dir, resolve_storage_paths
|
|
37
44
|
|
|
38
|
-
from .routers import health_router, index_router, query_router
|
|
45
|
+
from .routers import health_router, index_router, jobs_router, query_router
|
|
39
46
|
|
|
40
47
|
# Configure logging
|
|
41
48
|
logging.basicConfig(
|
|
@@ -48,6 +55,9 @@ logger = logging.getLogger(__name__)
|
|
|
48
55
|
_runtime_state: Optional[RuntimeState] = None
|
|
49
56
|
_state_dir: Optional[Path] = None
|
|
50
57
|
|
|
58
|
+
# Module-level reference to job worker for cleanup
|
|
59
|
+
_job_worker: Optional[JobWorker] = None
|
|
60
|
+
|
|
51
61
|
|
|
52
62
|
@asynccontextmanager
|
|
53
63
|
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
@@ -60,18 +70,51 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
|
60
70
|
- Resolves project root and state directory
|
|
61
71
|
- Acquires lock (with stale detection)
|
|
62
72
|
- Writes runtime.json with server info
|
|
73
|
+
- Initializes job queue system
|
|
63
74
|
- Cleans up on shutdown
|
|
64
75
|
"""
|
|
65
|
-
global _runtime_state, _state_dir
|
|
76
|
+
global _runtime_state, _state_dir, _job_worker
|
|
66
77
|
|
|
67
78
|
logger.info("Starting Agent Brain RAG server...")
|
|
68
79
|
|
|
80
|
+
# Load and validate provider configuration
|
|
81
|
+
# Clear cache first to ensure we pick up env vars set by CLI
|
|
82
|
+
clear_settings_cache()
|
|
83
|
+
try:
|
|
84
|
+
provider_settings = load_provider_settings()
|
|
85
|
+
validation_errors = validate_provider_config(provider_settings)
|
|
86
|
+
|
|
87
|
+
if validation_errors:
|
|
88
|
+
for error in validation_errors:
|
|
89
|
+
logger.warning(f"Provider config warning: {error}")
|
|
90
|
+
# Log but don't fail - providers may work if keys are set later
|
|
91
|
+
# or if using Ollama which doesn't need keys
|
|
92
|
+
|
|
93
|
+
# Log active provider configuration
|
|
94
|
+
logger.info(
|
|
95
|
+
f"Embedding provider: {provider_settings.embedding.provider} "
|
|
96
|
+
f"(model: {provider_settings.embedding.model})"
|
|
97
|
+
)
|
|
98
|
+
logger.info(
|
|
99
|
+
f"Summarization provider: {provider_settings.summarization.provider} "
|
|
100
|
+
f"(model: {provider_settings.summarization.model})"
|
|
101
|
+
)
|
|
102
|
+
except Exception as e:
|
|
103
|
+
logger.error(f"Failed to load provider configuration: {e}")
|
|
104
|
+
# Continue with defaults - EmbeddingGenerator will handle provider creation
|
|
105
|
+
|
|
69
106
|
if settings.OPENAI_API_KEY:
|
|
70
107
|
os.environ["OPENAI_API_KEY"] = settings.OPENAI_API_KEY
|
|
71
108
|
|
|
72
109
|
# Determine mode and resolve paths
|
|
73
|
-
mode = settings.
|
|
74
|
-
state_dir = _state_dir # May be set by
|
|
110
|
+
mode = settings.AGENT_BRAIN_MODE
|
|
111
|
+
state_dir = _state_dir # May be set by run() function
|
|
112
|
+
|
|
113
|
+
# If not set via run(), check environment variable (set by CLI subprocess)
|
|
114
|
+
if state_dir is None and settings.AGENT_BRAIN_STATE_DIR:
|
|
115
|
+
state_dir = Path(settings.AGENT_BRAIN_STATE_DIR).resolve()
|
|
116
|
+
logger.info(f"Using state directory from environment: {state_dir}")
|
|
117
|
+
|
|
75
118
|
storage_paths: Optional[dict[str, Path]] = None
|
|
76
119
|
|
|
77
120
|
if state_dir is not None:
|
|
@@ -86,13 +129,19 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
|
86
129
|
# Acquire exclusive lock
|
|
87
130
|
if not acquire_lock(state_dir):
|
|
88
131
|
raise RuntimeError(
|
|
89
|
-
f"Another
|
|
132
|
+
f"Another Agent Brain instance is already running for {state_dir}"
|
|
90
133
|
)
|
|
91
134
|
|
|
92
135
|
# Resolve storage paths (creates directories)
|
|
93
136
|
storage_paths = resolve_storage_paths(state_dir)
|
|
94
137
|
logger.info(f"State directory: {state_dir}")
|
|
95
138
|
|
|
139
|
+
# Determine project root for path validation
|
|
140
|
+
project_root: Optional[Path] = None
|
|
141
|
+
if state_dir is not None:
|
|
142
|
+
# Project root is 3 levels up from .claude/agent-brain
|
|
143
|
+
project_root = state_dir.parent.parent.parent
|
|
144
|
+
|
|
96
145
|
try:
|
|
97
146
|
# Determine persistence directories
|
|
98
147
|
chroma_dir = (
|
|
@@ -121,10 +170,28 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
|
121
170
|
app.state.bm25_manager = bm25_manager
|
|
122
171
|
logger.info("BM25 index manager initialized")
|
|
123
172
|
|
|
173
|
+
# Load project config for exclude patterns
|
|
174
|
+
exclude_patterns = None
|
|
175
|
+
if state_dir:
|
|
176
|
+
from agent_brain_server.config.settings import load_project_config
|
|
177
|
+
|
|
178
|
+
project_config = load_project_config(state_dir)
|
|
179
|
+
exclude_patterns = project_config.get("exclude_patterns")
|
|
180
|
+
if exclude_patterns:
|
|
181
|
+
logger.info(
|
|
182
|
+
f"Using exclude patterns from config: {exclude_patterns[:3]}..."
|
|
183
|
+
)
|
|
184
|
+
|
|
185
|
+
# Create document loader with exclude patterns
|
|
186
|
+
from agent_brain_server.indexing import DocumentLoader
|
|
187
|
+
|
|
188
|
+
document_loader = DocumentLoader(exclude_patterns=exclude_patterns)
|
|
189
|
+
|
|
124
190
|
# Create indexing service with injected deps
|
|
125
191
|
indexing_service = IndexingService(
|
|
126
192
|
vector_store=vector_store,
|
|
127
193
|
bm25_manager=bm25_manager,
|
|
194
|
+
document_loader=document_loader,
|
|
128
195
|
)
|
|
129
196
|
app.state.indexing_service = indexing_service
|
|
130
197
|
|
|
@@ -135,6 +202,57 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
|
135
202
|
)
|
|
136
203
|
app.state.query_service = query_service
|
|
137
204
|
|
|
205
|
+
# Initialize job queue system (Feature 115)
|
|
206
|
+
if state_dir is not None:
|
|
207
|
+
# Initialize job queue store
|
|
208
|
+
job_store = JobQueueStore(state_dir)
|
|
209
|
+
await job_store.initialize()
|
|
210
|
+
logger.info("Job queue store initialized")
|
|
211
|
+
|
|
212
|
+
# Initialize job queue service
|
|
213
|
+
job_service = JobQueueService(
|
|
214
|
+
store=job_store,
|
|
215
|
+
project_root=project_root,
|
|
216
|
+
)
|
|
217
|
+
app.state.job_service = job_service
|
|
218
|
+
logger.info("Job queue service initialized")
|
|
219
|
+
|
|
220
|
+
# Initialize and start job worker
|
|
221
|
+
_job_worker = JobWorker(
|
|
222
|
+
job_store=job_store,
|
|
223
|
+
indexing_service=indexing_service,
|
|
224
|
+
max_runtime_seconds=settings.AGENT_BRAIN_JOB_TIMEOUT,
|
|
225
|
+
progress_checkpoint_interval=settings.AGENT_BRAIN_CHECKPOINT_INTERVAL,
|
|
226
|
+
)
|
|
227
|
+
await _job_worker.start()
|
|
228
|
+
logger.info("Job worker started")
|
|
229
|
+
else:
|
|
230
|
+
# No state directory - create minimal job service for backward compat
|
|
231
|
+
# Jobs will not be persisted in this mode
|
|
232
|
+
logger.warning(
|
|
233
|
+
"No state directory configured - job queue persistence disabled"
|
|
234
|
+
)
|
|
235
|
+
# Create in-memory store with temp directory
|
|
236
|
+
import tempfile
|
|
237
|
+
|
|
238
|
+
temp_dir = Path(tempfile.mkdtemp(prefix="agent-brain-"))
|
|
239
|
+
job_store = JobQueueStore(temp_dir)
|
|
240
|
+
await job_store.initialize()
|
|
241
|
+
|
|
242
|
+
job_service = JobQueueService(
|
|
243
|
+
store=job_store,
|
|
244
|
+
project_root=project_root,
|
|
245
|
+
)
|
|
246
|
+
app.state.job_service = job_service
|
|
247
|
+
|
|
248
|
+
_job_worker = JobWorker(
|
|
249
|
+
job_store=job_store,
|
|
250
|
+
indexing_service=indexing_service,
|
|
251
|
+
max_runtime_seconds=settings.AGENT_BRAIN_JOB_TIMEOUT,
|
|
252
|
+
progress_checkpoint_interval=settings.AGENT_BRAIN_CHECKPOINT_INTERVAL,
|
|
253
|
+
)
|
|
254
|
+
await _job_worker.start()
|
|
255
|
+
|
|
138
256
|
# Set multi-instance metadata on app.state for health endpoint
|
|
139
257
|
app.state.mode = mode
|
|
140
258
|
app.state.instance_id = _runtime_state.instance_id if _runtime_state else None
|
|
@@ -152,6 +270,12 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
|
152
270
|
|
|
153
271
|
logger.info("Shutting down Agent Brain RAG server...")
|
|
154
272
|
|
|
273
|
+
# Stop job worker gracefully
|
|
274
|
+
if _job_worker is not None:
|
|
275
|
+
await _job_worker.stop()
|
|
276
|
+
logger.info("Job worker stopped")
|
|
277
|
+
_job_worker = None
|
|
278
|
+
|
|
155
279
|
# Cleanup for per-project mode
|
|
156
280
|
if state_dir is not None:
|
|
157
281
|
delete_runtime(state_dir)
|
|
@@ -166,7 +290,7 @@ app = FastAPI(
|
|
|
166
290
|
"RAG-based document indexing and semantic search API. "
|
|
167
291
|
"Index documents from folders and query them using natural language."
|
|
168
292
|
),
|
|
169
|
-
version=
|
|
293
|
+
version=__version__,
|
|
170
294
|
lifespan=lifespan,
|
|
171
295
|
docs_url="/docs",
|
|
172
296
|
redoc_url="/redoc",
|
|
@@ -185,6 +309,7 @@ app.add_middleware(
|
|
|
185
309
|
# Include routers
|
|
186
310
|
app.include_router(health_router, prefix="/health", tags=["Health"])
|
|
187
311
|
app.include_router(index_router, prefix="/index", tags=["Indexing"])
|
|
312
|
+
app.include_router(jobs_router, prefix="/index/jobs", tags=["Jobs"])
|
|
188
313
|
app.include_router(query_router, prefix="/query", tags=["Querying"])
|
|
189
314
|
|
|
190
315
|
|
|
@@ -193,7 +318,7 @@ async def root() -> dict[str, str]:
|
|
|
193
318
|
"""Root endpoint redirects to docs."""
|
|
194
319
|
return {
|
|
195
320
|
"name": "Agent Brain RAG API",
|
|
196
|
-
"version":
|
|
321
|
+
"version": __version__,
|
|
197
322
|
"docs": "/docs",
|
|
198
323
|
"health": "/health",
|
|
199
324
|
}
|
|
@@ -243,7 +368,7 @@ def run(
|
|
|
243
368
|
# Create runtime state
|
|
244
369
|
_runtime_state = RuntimeState(
|
|
245
370
|
mode="project",
|
|
246
|
-
project_root=str(_state_dir.parent.parent.parent), # .claude/
|
|
371
|
+
project_root=str(_state_dir.parent.parent.parent), # .claude/agent-brain
|
|
247
372
|
bind_host=resolved_host,
|
|
248
373
|
port=resolved_port,
|
|
249
374
|
pid=os.getpid(),
|
|
@@ -295,7 +420,7 @@ def run(
|
|
|
295
420
|
"--project-dir",
|
|
296
421
|
"-d",
|
|
297
422
|
default=None,
|
|
298
|
-
help="Project directory (auto-resolves state-dir to .claude/
|
|
423
|
+
help="Project directory (auto-resolves state-dir to .claude/agent-brain)",
|
|
299
424
|
)
|
|
300
425
|
def cli(
|
|
301
426
|
host: Optional[str],
|
|
@@ -316,15 +441,15 @@ def cli(
|
|
|
316
441
|
agent-brain-serve --host 0.0.0.0 # Bind to all interfaces
|
|
317
442
|
agent-brain-serve --reload # Enable auto-reload
|
|
318
443
|
agent-brain-serve --project-dir /my/project # Per-project mode
|
|
319
|
-
agent-brain-serve --state-dir /path/.claude/
|
|
444
|
+
agent-brain-serve --state-dir /path/.claude/agent-brain # Explicit state dir
|
|
320
445
|
|
|
321
446
|
\b
|
|
322
447
|
Environment Variables:
|
|
323
|
-
API_HOST
|
|
324
|
-
API_PORT
|
|
325
|
-
DEBUG
|
|
326
|
-
|
|
327
|
-
|
|
448
|
+
API_HOST Server host (default: 127.0.0.1)
|
|
449
|
+
API_PORT Server port (default: 8000)
|
|
450
|
+
DEBUG Enable debug mode (default: false)
|
|
451
|
+
AGENT_BRAIN_STATE_DIR Override state directory
|
|
452
|
+
AGENT_BRAIN_MODE Instance mode: 'project' or 'shared'
|
|
328
453
|
"""
|
|
329
454
|
# Resolve state directory from options
|
|
330
455
|
resolved_state_dir = state_dir
|
|
@@ -333,36 +458,12 @@ def cli(
|
|
|
333
458
|
# Auto-resolve state-dir from project directory
|
|
334
459
|
project_root = resolve_project_root(Path(project_dir))
|
|
335
460
|
resolved_state_dir = str(resolve_state_dir(project_root))
|
|
336
|
-
elif settings.
|
|
461
|
+
elif settings.AGENT_BRAIN_STATE_DIR and not state_dir:
|
|
337
462
|
# Use environment variable if set
|
|
338
|
-
resolved_state_dir = settings.
|
|
463
|
+
resolved_state_dir = settings.AGENT_BRAIN_STATE_DIR
|
|
339
464
|
|
|
340
465
|
run(host=host, port=port, reload=reload, state_dir=resolved_state_dir)
|
|
341
466
|
|
|
342
467
|
|
|
343
|
-
def cli_deprecated() -> None:
|
|
344
|
-
"""Deprecated entry point for doc-serve command.
|
|
345
|
-
|
|
346
|
-
Shows a deprecation warning and then runs the main CLI.
|
|
347
|
-
"""
|
|
348
|
-
warnings.warn(
|
|
349
|
-
"\n"
|
|
350
|
-
"WARNING: 'doc-serve' is deprecated and will be removed in v2.0.\n"
|
|
351
|
-
"Please use 'agent-brain-serve' instead.\n"
|
|
352
|
-
"\n"
|
|
353
|
-
"Migration guide: docs/MIGRATION.md\n"
|
|
354
|
-
"Online: https://github.com/SpillwaveSolutions/agent-brain/blob/main/docs/MIGRATION.md\n",
|
|
355
|
-
DeprecationWarning,
|
|
356
|
-
stacklevel=1,
|
|
357
|
-
)
|
|
358
|
-
# Print to stderr for visibility since warnings may be filtered
|
|
359
|
-
print(
|
|
360
|
-
"\033[93mWARNING: 'doc-serve' is deprecated. "
|
|
361
|
-
"Use 'agent-brain-serve' instead. See docs/MIGRATION.md\033[0m",
|
|
362
|
-
file=sys.stderr,
|
|
363
|
-
)
|
|
364
|
-
cli()
|
|
365
|
-
|
|
366
|
-
|
|
367
468
|
if __name__ == "__main__":
|
|
368
469
|
cli()
|
|
@@ -2,10 +2,12 @@
|
|
|
2
2
|
|
|
3
3
|
from .health import router as health_router
|
|
4
4
|
from .index import router as index_router
|
|
5
|
+
from .jobs import router as jobs_router
|
|
5
6
|
from .query import router as query_router
|
|
6
7
|
|
|
7
8
|
__all__ = [
|
|
8
9
|
"health_router",
|
|
9
10
|
"index_router",
|
|
11
|
+
"jobs_router",
|
|
10
12
|
"query_router",
|
|
11
13
|
]
|