agent-brain-rag 2.0.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-2.0.0.dist-info → agent_brain_rag-3.0.0.dist-info}/METADATA +3 -4
- {agent_brain_rag-2.0.0.dist-info → agent_brain_rag-3.0.0.dist-info}/RECORD +26 -20
- {agent_brain_rag-2.0.0.dist-info → agent_brain_rag-3.0.0.dist-info}/WHEEL +1 -1
- {agent_brain_rag-2.0.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 +118 -45
- agent_brain_server/api/routers/__init__.py +2 -0
- agent_brain_server/api/routers/health.py +85 -22
- 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 +63 -19
- agent_brain_server/config/settings.py +10 -4
- agent_brain_server/indexing/bm25_index.py +15 -2
- agent_brain_server/indexing/document_loader.py +45 -4
- 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 +19 -0
- agent_brain_server/models/health.py +15 -0
- agent_brain_server/models/job.py +289 -0
- agent_brain_server/models/query.py +2 -2
- agent_brain_server/project_root.py +1 -1
- agent_brain_server/runtime.py +2 -2
- agent_brain_server/storage_paths.py +3 -3
|
@@ -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,7 +14,6 @@ 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
19
|
Provides-Extra: graphrag
|
|
@@ -42,7 +42,6 @@ Requires-Dist: tiktoken (>=0.8.0,<0.9.0)
|
|
|
42
42
|
Requires-Dist: tree-sitter-language-pack (>=0.7.3,<0.8.0)
|
|
43
43
|
Requires-Dist: uvicorn[standard] (>=0.32.0,<0.33.0)
|
|
44
44
|
Project-URL: Documentation, https://github.com/SpillwaveSolutions/agent-brain/wiki
|
|
45
|
-
Project-URL: Homepage, https://github.com/SpillwaveSolutions/agent-brain
|
|
46
45
|
Project-URL: Repository, https://github.com/SpillwaveSolutions/agent-brain
|
|
47
46
|
Description-Content-Type: text/markdown
|
|
48
47
|
|
|
@@ -1,27 +1,33 @@
|
|
|
1
|
-
agent_brain_server/__init__.py,sha256=
|
|
1
|
+
agent_brain_server/__init__.py,sha256=Wjd2w3VnDYURTECaKxFQ5x324-7tZ_vK2Qv3yKfxRGg,95
|
|
2
2
|
agent_brain_server/api/__init__.py,sha256=nvTvO_ahHAAsRDlV3dL_JlNruSdan4kav_P5sT_1PFk,93
|
|
3
|
-
agent_brain_server/api/main.py,sha256=
|
|
4
|
-
agent_brain_server/api/routers/__init__.py,sha256=
|
|
5
|
-
agent_brain_server/api/routers/health.py,sha256=
|
|
6
|
-
agent_brain_server/api/routers/index.py,sha256=
|
|
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
|
|
7
8
|
agent_brain_server/api/routers/query.py,sha256=TAbpuCybjacUo_-7Zm3Jt8EW8EMDTYJDLv6aGtjaUzQ,2793
|
|
8
9
|
agent_brain_server/config/__init__.py,sha256=zzDErZGUBwUm5Fk43OHJN7eWpeIw_1kWdnhsN6QQqSc,84
|
|
9
|
-
agent_brain_server/config/provider_config.py,sha256
|
|
10
|
-
agent_brain_server/config/settings.py,sha256=
|
|
10
|
+
agent_brain_server/config/provider_config.py,sha256=-uHj28i8OvAhQxlZ0JudYPAyibdJWmEJfvOYuRAdSHw,11535
|
|
11
|
+
agent_brain_server/config/settings.py,sha256=XfkAi-5aV_vqbk4ltkG8mkJnPARPwl5g6JxvE0ngDjo,3726
|
|
11
12
|
agent_brain_server/indexing/__init__.py,sha256=tQiFW4945h-ENd0pmMoGhCbm1Fz3AaREKBgajHeW-Us,1169
|
|
12
|
-
agent_brain_server/indexing/bm25_index.py,sha256=
|
|
13
|
+
agent_brain_server/indexing/bm25_index.py,sha256=_CZdTF1zoLyILTJ1nTLSmovIqN2NVQuYUiMEWh4o3-o,5740
|
|
13
14
|
agent_brain_server/indexing/chunking.py,sha256=kW-ZOdwOS-PfXuPAYb8qL4-7yAY-rpY9a1el2YCN9CA,29848
|
|
14
|
-
agent_brain_server/indexing/document_loader.py,sha256=
|
|
15
|
+
agent_brain_server/indexing/document_loader.py,sha256=l43Z5xHEykr6hbQqkMWXZGlT_CDz_tdhkI3PsxTMsgQ,17543
|
|
15
16
|
agent_brain_server/indexing/embedding.py,sha256=i3bN3SivvVB8iNiYDjrDeB8nF3WkFoOy5wHrDaNosGs,7410
|
|
16
17
|
agent_brain_server/indexing/graph_extractors.py,sha256=XEWBqlqDffxbs5ZLaijGICXeaptGtkUGNEhJblP3jBU,19339
|
|
17
18
|
agent_brain_server/indexing/graph_index.py,sha256=ZKwICIqR2stcPEqNQ_SSuSk-O0Hj1lUmiUMxm5udhjg,18584
|
|
18
|
-
agent_brain_server/
|
|
19
|
-
agent_brain_server/
|
|
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
|
|
20
25
|
agent_brain_server/models/graph.py,sha256=YzdjJj_X7hueJt4v8Gsoly4C-qquJLZi1d9epIh_Mos,7973
|
|
21
|
-
agent_brain_server/models/health.py,sha256=
|
|
26
|
+
agent_brain_server/models/health.py,sha256=ppvDF_yUaddzYPe0BPc9emVDrR17TVlQLW6ggsOOzPo,4592
|
|
22
27
|
agent_brain_server/models/index.py,sha256=pjDv7phLS6dpiHLlEtcAuXQN_aHIfA_4lMkAZ-NkXZQ,5400
|
|
23
|
-
agent_brain_server/models/
|
|
24
|
-
agent_brain_server/
|
|
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
|
|
25
31
|
agent_brain_server/providers/__init__.py,sha256=04wJ_QhGKRM55R4yrbw3J3kCtKX44_FHzh6xZq1xFdc,1670
|
|
26
32
|
agent_brain_server/providers/base.py,sha256=yLIpz9W0xlilBXxuFJRNP_gbGDsZ9HeRSqEnIsVGQjU,6933
|
|
27
33
|
agent_brain_server/providers/embedding/__init__.py,sha256=WRUAkOLLgQB9AR9I-_m3QWa55hSeC5tDdp3gZQFfMRM,982
|
|
@@ -36,15 +42,15 @@ agent_brain_server/providers/summarization/gemini.py,sha256=t2pOC4R2ae7WQAq7WCy8
|
|
|
36
42
|
agent_brain_server/providers/summarization/grok.py,sha256=nK3i46NbUcVxBqWNbJ-IifahB450Blnqm0rBhSQnQOw,2875
|
|
37
43
|
agent_brain_server/providers/summarization/ollama.py,sha256=mw4zvx6o2nC4xoz1W6F47x3fk7gidl7aeh1uAnSTHO8,3673
|
|
38
44
|
agent_brain_server/providers/summarization/openai.py,sha256=HjZJIjwc0usHejRAFW3JT6a7cmM_Ib2FdC05y_iEwA4,2632
|
|
39
|
-
agent_brain_server/runtime.py,sha256=
|
|
45
|
+
agent_brain_server/runtime.py,sha256=SVQa-8rgNFBeHyJe23bvzZ49Gf-4PqkuSJ-kJ1FPkSM,3081
|
|
40
46
|
agent_brain_server/services/__init__.py,sha256=E4VPN9Rqa2mxGQQEQn-5IYj63LSPTrA8aIx8ENO5xcc,296
|
|
41
47
|
agent_brain_server/services/indexing_service.py,sha256=d_oj3n7gNgua7nmiLNTM-hDp4wvcrkos9kCRuOTzlpc,20338
|
|
42
48
|
agent_brain_server/services/query_service.py,sha256=oc_BNqTNjAtz17y_F7r-egP7Z0GozmP_mlfSs1jKbMM,22910
|
|
43
49
|
agent_brain_server/storage/__init__.py,sha256=xYk4MfrEj-IpGZ0n4myfy5dwyjKEvprS7h68bet0ooU,557
|
|
44
50
|
agent_brain_server/storage/graph_store.py,sha256=boW6_L0suvHuOdfZdPMdSNSb8U-am7KLelsa3K2bKl4,18168
|
|
45
51
|
agent_brain_server/storage/vector_store.py,sha256=rcI--g6g7mqUbZzewOsaMdb7uqw173bAZau1p0oYZ68,12135
|
|
46
|
-
agent_brain_server/storage_paths.py,sha256=
|
|
47
|
-
agent_brain_rag-
|
|
48
|
-
agent_brain_rag-
|
|
49
|
-
agent_brain_rag-
|
|
50
|
-
agent_brain_rag-
|
|
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
|
|
@@ -23,10 +24,12 @@ from fastapi.middleware.cors import CORSMiddleware
|
|
|
23
24
|
from agent_brain_server import __version__
|
|
24
25
|
from agent_brain_server.config import settings
|
|
25
26
|
from agent_brain_server.config.provider_config import (
|
|
27
|
+
clear_settings_cache,
|
|
26
28
|
load_provider_settings,
|
|
27
29
|
validate_provider_config,
|
|
28
30
|
)
|
|
29
31
|
from agent_brain_server.indexing.bm25_index import BM25IndexManager
|
|
32
|
+
from agent_brain_server.job_queue import JobQueueService, JobQueueStore, JobWorker
|
|
30
33
|
from agent_brain_server.locking import (
|
|
31
34
|
acquire_lock,
|
|
32
35
|
cleanup_stale,
|
|
@@ -39,7 +42,7 @@ from agent_brain_server.services import IndexingService, QueryService
|
|
|
39
42
|
from agent_brain_server.storage import VectorStoreManager
|
|
40
43
|
from agent_brain_server.storage_paths import resolve_state_dir, resolve_storage_paths
|
|
41
44
|
|
|
42
|
-
from .routers import health_router, index_router, query_router
|
|
45
|
+
from .routers import health_router, index_router, jobs_router, query_router
|
|
43
46
|
|
|
44
47
|
# Configure logging
|
|
45
48
|
logging.basicConfig(
|
|
@@ -52,6 +55,9 @@ logger = logging.getLogger(__name__)
|
|
|
52
55
|
_runtime_state: Optional[RuntimeState] = None
|
|
53
56
|
_state_dir: Optional[Path] = None
|
|
54
57
|
|
|
58
|
+
# Module-level reference to job worker for cleanup
|
|
59
|
+
_job_worker: Optional[JobWorker] = None
|
|
60
|
+
|
|
55
61
|
|
|
56
62
|
@asynccontextmanager
|
|
57
63
|
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
@@ -64,13 +70,16 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
|
64
70
|
- Resolves project root and state directory
|
|
65
71
|
- Acquires lock (with stale detection)
|
|
66
72
|
- Writes runtime.json with server info
|
|
73
|
+
- Initializes job queue system
|
|
67
74
|
- Cleans up on shutdown
|
|
68
75
|
"""
|
|
69
|
-
global _runtime_state, _state_dir
|
|
76
|
+
global _runtime_state, _state_dir, _job_worker
|
|
70
77
|
|
|
71
78
|
logger.info("Starting Agent Brain RAG server...")
|
|
72
79
|
|
|
73
80
|
# Load and validate provider configuration
|
|
81
|
+
# Clear cache first to ensure we pick up env vars set by CLI
|
|
82
|
+
clear_settings_cache()
|
|
74
83
|
try:
|
|
75
84
|
provider_settings = load_provider_settings()
|
|
76
85
|
validation_errors = validate_provider_config(provider_settings)
|
|
@@ -98,8 +107,14 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
|
98
107
|
os.environ["OPENAI_API_KEY"] = settings.OPENAI_API_KEY
|
|
99
108
|
|
|
100
109
|
# Determine mode and resolve paths
|
|
101
|
-
mode = settings.
|
|
102
|
-
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
|
+
|
|
103
118
|
storage_paths: Optional[dict[str, Path]] = None
|
|
104
119
|
|
|
105
120
|
if state_dir is not None:
|
|
@@ -114,13 +129,19 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
|
114
129
|
# Acquire exclusive lock
|
|
115
130
|
if not acquire_lock(state_dir):
|
|
116
131
|
raise RuntimeError(
|
|
117
|
-
f"Another
|
|
132
|
+
f"Another Agent Brain instance is already running for {state_dir}"
|
|
118
133
|
)
|
|
119
134
|
|
|
120
135
|
# Resolve storage paths (creates directories)
|
|
121
136
|
storage_paths = resolve_storage_paths(state_dir)
|
|
122
137
|
logger.info(f"State directory: {state_dir}")
|
|
123
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
|
+
|
|
124
145
|
try:
|
|
125
146
|
# Determine persistence directories
|
|
126
147
|
chroma_dir = (
|
|
@@ -149,10 +170,28 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
|
149
170
|
app.state.bm25_manager = bm25_manager
|
|
150
171
|
logger.info("BM25 index manager initialized")
|
|
151
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
|
+
|
|
152
190
|
# Create indexing service with injected deps
|
|
153
191
|
indexing_service = IndexingService(
|
|
154
192
|
vector_store=vector_store,
|
|
155
193
|
bm25_manager=bm25_manager,
|
|
194
|
+
document_loader=document_loader,
|
|
156
195
|
)
|
|
157
196
|
app.state.indexing_service = indexing_service
|
|
158
197
|
|
|
@@ -163,6 +202,57 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
|
163
202
|
)
|
|
164
203
|
app.state.query_service = query_service
|
|
165
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
|
+
|
|
166
256
|
# Set multi-instance metadata on app.state for health endpoint
|
|
167
257
|
app.state.mode = mode
|
|
168
258
|
app.state.instance_id = _runtime_state.instance_id if _runtime_state else None
|
|
@@ -180,6 +270,12 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
|
180
270
|
|
|
181
271
|
logger.info("Shutting down Agent Brain RAG server...")
|
|
182
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
|
+
|
|
183
279
|
# Cleanup for per-project mode
|
|
184
280
|
if state_dir is not None:
|
|
185
281
|
delete_runtime(state_dir)
|
|
@@ -194,7 +290,7 @@ app = FastAPI(
|
|
|
194
290
|
"RAG-based document indexing and semantic search API. "
|
|
195
291
|
"Index documents from folders and query them using natural language."
|
|
196
292
|
),
|
|
197
|
-
version=
|
|
293
|
+
version=__version__,
|
|
198
294
|
lifespan=lifespan,
|
|
199
295
|
docs_url="/docs",
|
|
200
296
|
redoc_url="/redoc",
|
|
@@ -213,6 +309,7 @@ app.add_middleware(
|
|
|
213
309
|
# Include routers
|
|
214
310
|
app.include_router(health_router, prefix="/health", tags=["Health"])
|
|
215
311
|
app.include_router(index_router, prefix="/index", tags=["Indexing"])
|
|
312
|
+
app.include_router(jobs_router, prefix="/index/jobs", tags=["Jobs"])
|
|
216
313
|
app.include_router(query_router, prefix="/query", tags=["Querying"])
|
|
217
314
|
|
|
218
315
|
|
|
@@ -221,7 +318,7 @@ async def root() -> dict[str, str]:
|
|
|
221
318
|
"""Root endpoint redirects to docs."""
|
|
222
319
|
return {
|
|
223
320
|
"name": "Agent Brain RAG API",
|
|
224
|
-
"version":
|
|
321
|
+
"version": __version__,
|
|
225
322
|
"docs": "/docs",
|
|
226
323
|
"health": "/health",
|
|
227
324
|
}
|
|
@@ -271,7 +368,7 @@ def run(
|
|
|
271
368
|
# Create runtime state
|
|
272
369
|
_runtime_state = RuntimeState(
|
|
273
370
|
mode="project",
|
|
274
|
-
project_root=str(_state_dir.parent.parent.parent), # .claude/
|
|
371
|
+
project_root=str(_state_dir.parent.parent.parent), # .claude/agent-brain
|
|
275
372
|
bind_host=resolved_host,
|
|
276
373
|
port=resolved_port,
|
|
277
374
|
pid=os.getpid(),
|
|
@@ -323,7 +420,7 @@ def run(
|
|
|
323
420
|
"--project-dir",
|
|
324
421
|
"-d",
|
|
325
422
|
default=None,
|
|
326
|
-
help="Project directory (auto-resolves state-dir to .claude/
|
|
423
|
+
help="Project directory (auto-resolves state-dir to .claude/agent-brain)",
|
|
327
424
|
)
|
|
328
425
|
def cli(
|
|
329
426
|
host: Optional[str],
|
|
@@ -344,15 +441,15 @@ def cli(
|
|
|
344
441
|
agent-brain-serve --host 0.0.0.0 # Bind to all interfaces
|
|
345
442
|
agent-brain-serve --reload # Enable auto-reload
|
|
346
443
|
agent-brain-serve --project-dir /my/project # Per-project mode
|
|
347
|
-
agent-brain-serve --state-dir /path/.claude/
|
|
444
|
+
agent-brain-serve --state-dir /path/.claude/agent-brain # Explicit state dir
|
|
348
445
|
|
|
349
446
|
\b
|
|
350
447
|
Environment Variables:
|
|
351
|
-
API_HOST
|
|
352
|
-
API_PORT
|
|
353
|
-
DEBUG
|
|
354
|
-
|
|
355
|
-
|
|
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'
|
|
356
453
|
"""
|
|
357
454
|
# Resolve state directory from options
|
|
358
455
|
resolved_state_dir = state_dir
|
|
@@ -361,36 +458,12 @@ def cli(
|
|
|
361
458
|
# Auto-resolve state-dir from project directory
|
|
362
459
|
project_root = resolve_project_root(Path(project_dir))
|
|
363
460
|
resolved_state_dir = str(resolve_state_dir(project_root))
|
|
364
|
-
elif settings.
|
|
461
|
+
elif settings.AGENT_BRAIN_STATE_DIR and not state_dir:
|
|
365
462
|
# Use environment variable if set
|
|
366
|
-
resolved_state_dir = settings.
|
|
463
|
+
resolved_state_dir = settings.AGENT_BRAIN_STATE_DIR
|
|
367
464
|
|
|
368
465
|
run(host=host, port=port, reload=reload, state_dir=resolved_state_dir)
|
|
369
466
|
|
|
370
467
|
|
|
371
|
-
def cli_deprecated() -> None:
|
|
372
|
-
"""Deprecated entry point for doc-serve command.
|
|
373
|
-
|
|
374
|
-
Shows a deprecation warning and then runs the main CLI.
|
|
375
|
-
"""
|
|
376
|
-
warnings.warn(
|
|
377
|
-
"\n"
|
|
378
|
-
"WARNING: 'doc-serve' is deprecated and will be removed in v2.0.\n"
|
|
379
|
-
"Please use 'agent-brain-serve' instead.\n"
|
|
380
|
-
"\n"
|
|
381
|
-
"Migration guide: docs/MIGRATION.md\n"
|
|
382
|
-
"Online: https://github.com/SpillwaveSolutions/agent-brain/blob/main/docs/MIGRATION.md\n",
|
|
383
|
-
DeprecationWarning,
|
|
384
|
-
stacklevel=1,
|
|
385
|
-
)
|
|
386
|
-
# Print to stderr for visibility since warnings may be filtered
|
|
387
|
-
print(
|
|
388
|
-
"\033[93mWARNING: 'doc-serve' is deprecated. "
|
|
389
|
-
"Use 'agent-brain-serve' instead. See docs/MIGRATION.md\033[0m",
|
|
390
|
-
file=sys.stderr,
|
|
391
|
-
)
|
|
392
|
-
cli()
|
|
393
|
-
|
|
394
|
-
|
|
395
468
|
if __name__ == "__main__":
|
|
396
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
|
]
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
"""Health check endpoints."""
|
|
1
|
+
"""Health check endpoints with non-blocking queue status."""
|
|
2
2
|
|
|
3
3
|
from datetime import datetime, timezone
|
|
4
4
|
from typing import Literal
|
|
@@ -20,6 +20,8 @@ router = APIRouter()
|
|
|
20
20
|
async def health_check(request: Request) -> HealthStatus:
|
|
21
21
|
"""Check server health status.
|
|
22
22
|
|
|
23
|
+
This endpoint never blocks and always returns quickly.
|
|
24
|
+
|
|
23
25
|
Returns:
|
|
24
26
|
HealthStatus with current status:
|
|
25
27
|
- healthy: Server is running and ready for queries
|
|
@@ -27,20 +29,35 @@ async def health_check(request: Request) -> HealthStatus:
|
|
|
27
29
|
- degraded: Server is up but some services are unavailable
|
|
28
30
|
- unhealthy: Server is not operational
|
|
29
31
|
"""
|
|
30
|
-
indexing_service = request.app.state.indexing_service
|
|
31
32
|
vector_store = request.app.state.vector_store
|
|
33
|
+
job_service = getattr(request.app.state, "job_service", None)
|
|
32
34
|
|
|
33
|
-
# Determine status
|
|
35
|
+
# Determine status using queue service (non-blocking)
|
|
34
36
|
status: Literal["healthy", "indexing", "degraded", "unhealthy"]
|
|
35
|
-
|
|
37
|
+
message: str
|
|
38
|
+
|
|
39
|
+
# Check queue status (non-blocking)
|
|
40
|
+
is_indexing = False
|
|
41
|
+
current_folder = None
|
|
42
|
+
if job_service:
|
|
43
|
+
try:
|
|
44
|
+
queue_stats = await job_service.get_queue_stats()
|
|
45
|
+
is_indexing = queue_stats.running > 0
|
|
46
|
+
if is_indexing and queue_stats.current_job_id:
|
|
47
|
+
# Get current job details for message
|
|
48
|
+
current_job = await job_service.get_job(queue_stats.current_job_id)
|
|
49
|
+
if current_job:
|
|
50
|
+
current_folder = current_job.folder_path
|
|
51
|
+
except Exception:
|
|
52
|
+
# Non-blocking: don't fail health check if queue service errors
|
|
53
|
+
pass
|
|
54
|
+
|
|
55
|
+
if is_indexing:
|
|
36
56
|
status = "indexing"
|
|
37
|
-
message = f"Indexing in progress: {
|
|
57
|
+
message = f"Indexing in progress: {current_folder or 'unknown'}"
|
|
38
58
|
elif not vector_store.is_initialized:
|
|
39
59
|
status = "degraded"
|
|
40
60
|
message = "Vector store not initialized"
|
|
41
|
-
elif indexing_service.state.error:
|
|
42
|
-
status = "degraded"
|
|
43
|
-
message = f"Last indexing failed: {indexing_service.state.error}"
|
|
44
61
|
else:
|
|
45
62
|
status = "healthy"
|
|
46
63
|
message = "Server is running and ready for queries"
|
|
@@ -67,36 +84,82 @@ async def health_check(request: Request) -> HealthStatus:
|
|
|
67
84
|
"/status",
|
|
68
85
|
response_model=IndexingStatus,
|
|
69
86
|
summary="Indexing Status",
|
|
70
|
-
description="Returns detailed indexing status information.",
|
|
87
|
+
description="Returns detailed indexing status information. Never blocks.",
|
|
71
88
|
)
|
|
72
89
|
async def indexing_status(request: Request) -> IndexingStatus:
|
|
73
90
|
"""Get detailed indexing status.
|
|
74
91
|
|
|
92
|
+
This endpoint never blocks and always returns quickly, even during indexing.
|
|
93
|
+
|
|
75
94
|
Returns:
|
|
76
95
|
IndexingStatus with:
|
|
77
96
|
- total_documents: Number of documents indexed
|
|
78
97
|
- total_chunks: Number of chunks in vector store
|
|
79
98
|
- indexing_in_progress: Boolean indicating active indexing
|
|
99
|
+
- queue_pending: Number of pending jobs
|
|
100
|
+
- queue_running: Number of running jobs (0 or 1)
|
|
101
|
+
- current_job_running_time_ms: How long current job has been running
|
|
80
102
|
- last_indexed_at: Timestamp of last indexing operation
|
|
81
103
|
- indexed_folders: List of folders that have been indexed
|
|
82
104
|
"""
|
|
83
105
|
indexing_service = request.app.state.indexing_service
|
|
84
|
-
|
|
106
|
+
vector_store = request.app.state.vector_store
|
|
107
|
+
job_service = getattr(request.app.state, "job_service", None)
|
|
108
|
+
|
|
109
|
+
# Get vector store count (non-blocking read)
|
|
110
|
+
try:
|
|
111
|
+
total_chunks = (
|
|
112
|
+
await vector_store.get_count() if vector_store.is_initialized else 0
|
|
113
|
+
)
|
|
114
|
+
except Exception:
|
|
115
|
+
total_chunks = 0
|
|
116
|
+
|
|
117
|
+
# Get queue status (non-blocking)
|
|
118
|
+
queue_pending = 0
|
|
119
|
+
queue_running = 0
|
|
120
|
+
current_job_id = None
|
|
121
|
+
current_job_running_time_ms = None
|
|
122
|
+
progress_percent = 0.0
|
|
123
|
+
|
|
124
|
+
if job_service:
|
|
125
|
+
try:
|
|
126
|
+
queue_stats = await job_service.get_queue_stats()
|
|
127
|
+
queue_pending = queue_stats.pending
|
|
128
|
+
queue_running = queue_stats.running
|
|
129
|
+
current_job_id = queue_stats.current_job_id
|
|
130
|
+
current_job_running_time_ms = queue_stats.current_job_running_time_ms
|
|
131
|
+
|
|
132
|
+
# Get progress from current job
|
|
133
|
+
if current_job_id:
|
|
134
|
+
current_job = await job_service.get_job(current_job_id)
|
|
135
|
+
if current_job and current_job.progress:
|
|
136
|
+
progress_percent = current_job.progress.percent_complete
|
|
137
|
+
except Exception:
|
|
138
|
+
# Non-blocking: don't fail status if queue service errors
|
|
139
|
+
pass
|
|
140
|
+
|
|
141
|
+
# Get indexing service status for historical data
|
|
142
|
+
# This is read-only and non-blocking
|
|
143
|
+
service_status = await indexing_service.get_status()
|
|
85
144
|
|
|
86
145
|
return IndexingStatus(
|
|
87
|
-
total_documents=
|
|
88
|
-
total_chunks=
|
|
89
|
-
total_doc_chunks=
|
|
90
|
-
total_code_chunks=
|
|
91
|
-
indexing_in_progress=
|
|
92
|
-
current_job_id=
|
|
93
|
-
progress_percent=
|
|
146
|
+
total_documents=service_status.get("total_documents", 0),
|
|
147
|
+
total_chunks=total_chunks,
|
|
148
|
+
total_doc_chunks=service_status.get("total_doc_chunks", 0),
|
|
149
|
+
total_code_chunks=service_status.get("total_code_chunks", 0),
|
|
150
|
+
indexing_in_progress=queue_running > 0,
|
|
151
|
+
current_job_id=current_job_id,
|
|
152
|
+
progress_percent=progress_percent,
|
|
94
153
|
last_indexed_at=(
|
|
95
|
-
datetime.fromisoformat(
|
|
96
|
-
if
|
|
154
|
+
datetime.fromisoformat(service_status["completed_at"])
|
|
155
|
+
if service_status.get("completed_at")
|
|
97
156
|
else None
|
|
98
157
|
),
|
|
99
|
-
indexed_folders=
|
|
100
|
-
supported_languages=
|
|
101
|
-
graph_index=
|
|
158
|
+
indexed_folders=service_status.get("indexed_folders", []),
|
|
159
|
+
supported_languages=service_status.get("supported_languages", []),
|
|
160
|
+
graph_index=service_status.get("graph_index"),
|
|
161
|
+
# Queue status (Feature 115)
|
|
162
|
+
queue_pending=queue_pending,
|
|
163
|
+
queue_running=queue_running,
|
|
164
|
+
current_job_running_time_ms=current_job_running_time_ms,
|
|
102
165
|
)
|