agent-brain-rag 1.2.0__tar.gz → 2.0.0__tar.gz
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 → agent_brain_rag-2.0.0}/PKG-INFO +54 -16
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/README.md +42 -12
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/__init__.py +1 -1
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/api/main.py +30 -2
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/api/routers/health.py +1 -0
- agent_brain_rag-2.0.0/agent_brain_server/config/provider_config.py +308 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/config/settings.py +12 -1
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/indexing/__init__.py +21 -0
- agent_brain_rag-2.0.0/agent_brain_server/indexing/embedding.py +225 -0
- agent_brain_rag-2.0.0/agent_brain_server/indexing/graph_extractors.py +582 -0
- agent_brain_rag-2.0.0/agent_brain_server/indexing/graph_index.py +536 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/models/__init__.py +9 -0
- agent_brain_rag-2.0.0/agent_brain_server/models/graph.py +253 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/models/health.py +15 -3
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/models/query.py +14 -1
- agent_brain_rag-2.0.0/agent_brain_server/providers/__init__.py +64 -0
- agent_brain_rag-2.0.0/agent_brain_server/providers/base.py +251 -0
- agent_brain_rag-2.0.0/agent_brain_server/providers/embedding/__init__.py +23 -0
- agent_brain_rag-2.0.0/agent_brain_server/providers/embedding/cohere.py +163 -0
- agent_brain_rag-2.0.0/agent_brain_server/providers/embedding/ollama.py +150 -0
- agent_brain_rag-2.0.0/agent_brain_server/providers/embedding/openai.py +118 -0
- agent_brain_rag-2.0.0/agent_brain_server/providers/exceptions.py +95 -0
- agent_brain_rag-2.0.0/agent_brain_server/providers/factory.py +157 -0
- agent_brain_rag-2.0.0/agent_brain_server/providers/summarization/__init__.py +41 -0
- agent_brain_rag-2.0.0/agent_brain_server/providers/summarization/anthropic.py +87 -0
- agent_brain_rag-2.0.0/agent_brain_server/providers/summarization/gemini.py +96 -0
- agent_brain_rag-2.0.0/agent_brain_server/providers/summarization/grok.py +95 -0
- agent_brain_rag-2.0.0/agent_brain_server/providers/summarization/ollama.py +114 -0
- agent_brain_rag-2.0.0/agent_brain_server/providers/summarization/openai.py +87 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/services/indexing_service.py +39 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/services/query_service.py +203 -0
- agent_brain_rag-2.0.0/agent_brain_server/storage/__init__.py +21 -0
- agent_brain_rag-2.0.0/agent_brain_server/storage/graph_store.py +519 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/storage/vector_store.py +35 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/storage_paths.py +2 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/pyproject.toml +19 -4
- agent_brain_rag-1.2.0/agent_brain_server/indexing/embedding.py +0 -274
- agent_brain_rag-1.2.0/agent_brain_server/storage/__init__.py +0 -5
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/api/__init__.py +0 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/api/routers/__init__.py +0 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/api/routers/index.py +0 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/api/routers/query.py +0 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/config/__init__.py +0 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/indexing/bm25_index.py +0 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/indexing/chunking.py +0 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/indexing/document_loader.py +0 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/locking.py +0 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/models/index.py +0 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/project_root.py +0 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/runtime.py +0 -0
- {agent_brain_rag-1.2.0 → agent_brain_rag-2.0.0}/agent_brain_server/services/__init__.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.3
|
|
2
2
|
Name: agent-brain-rag
|
|
3
|
-
Version:
|
|
3
|
+
Version: 2.0.0
|
|
4
4
|
Summary: Agent Brain RAG - Intelligent document indexing and semantic search server that gives AI agents long-term memory
|
|
5
5
|
License: MIT
|
|
6
6
|
Keywords: agent-brain,rag,semantic-search,ai-memory,llm-memory,documentation,indexing,llama-index,chromadb,ai-agent,claude-code,agent-memory
|
|
@@ -16,12 +16,19 @@ Classifier: Programming Language :: Python :: 3.12
|
|
|
16
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,25 @@ 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: Homepage, https://github.com/SpillwaveSolutions/
|
|
38
|
-
Project-URL: Repository, https://github.com/SpillwaveSolutions/
|
|
44
|
+
Project-URL: Documentation, https://github.com/SpillwaveSolutions/agent-brain/wiki
|
|
45
|
+
Project-URL: Homepage, https://github.com/SpillwaveSolutions/agent-brain
|
|
46
|
+
Project-URL: Repository, https://github.com/SpillwaveSolutions/agent-brain
|
|
39
47
|
Description-Content-Type: text/markdown
|
|
40
48
|
|
|
41
49
|
# Agent Brain RAG Server
|
|
42
50
|
|
|
43
|
-
|
|
51
|
+
> **Agent Brain** (formerly doc-serve) is an intelligent document indexing and semantic search system designed to give AI agents long-term memory.
|
|
52
|
+
|
|
53
|
+
AI agents need persistent memory to be truly useful. Agent Brain provides the retrieval infrastructure that enables context-aware, knowledge-grounded AI interactions.
|
|
54
|
+
|
|
55
|
+
[](https://pypi.org/project/agent-brain-rag/)
|
|
56
|
+
[](https://www.python.org/downloads/)
|
|
57
|
+
[](https://opensource.org/licenses/MIT)
|
|
44
58
|
|
|
45
59
|
## Installation
|
|
46
60
|
|
|
@@ -65,28 +79,37 @@ The server will start at `http://127.0.0.1:8000`.
|
|
|
65
79
|
|
|
66
80
|
> **Note**: The legacy command `doc-serve` is still available but deprecated. Please use `agent-brain-serve` for new installations.
|
|
67
81
|
|
|
82
|
+
## Search Capabilities
|
|
83
|
+
|
|
84
|
+
Agent Brain provides multiple search strategies to match your retrieval needs:
|
|
85
|
+
|
|
86
|
+
| Search Type | Description | Best For |
|
|
87
|
+
|-------------|-------------|----------|
|
|
88
|
+
| **Semantic Search** | Natural language queries using OpenAI embeddings (`text-embedding-3-large`) | Conceptual questions, finding related content |
|
|
89
|
+
| **Keyword Search (BM25)** | Traditional keyword matching with TF-IDF ranking | Exact matches, technical terms, code identifiers |
|
|
90
|
+
| **Hybrid Search** | Combines vector + BM25 for best of both approaches | General-purpose queries, balanced recall/precision |
|
|
91
|
+
| **GraphRAG** | Knowledge graph-based retrieval for relationship-aware queries | Understanding connections, multi-hop reasoning |
|
|
92
|
+
|
|
68
93
|
## Features
|
|
69
94
|
|
|
70
95
|
- **Document Indexing**: Load and index documents from folders (PDF, Markdown, TXT, DOCX, HTML)
|
|
71
96
|
- **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
|
|
97
|
+
- **Multi-Strategy Retrieval**: Semantic, keyword, hybrid, and graph-based search
|
|
74
98
|
- **OpenAI Embeddings**: Uses `text-embedding-3-large` for high-quality embeddings
|
|
99
|
+
- **Claude Summarization**: AI-powered code summaries for better context
|
|
75
100
|
- **Chroma Vector Store**: Persistent, thread-safe vector database
|
|
76
101
|
- **FastAPI**: Modern, high-performance REST API with OpenAPI documentation
|
|
77
102
|
|
|
78
|
-
##
|
|
79
|
-
|
|
80
|
-
### Prerequisites
|
|
103
|
+
## Prerequisites
|
|
81
104
|
|
|
82
105
|
- Python 3.10+
|
|
83
|
-
-
|
|
84
|
-
-
|
|
106
|
+
- OpenAI API key (for embeddings)
|
|
107
|
+
- Anthropic API key (for summarization)
|
|
85
108
|
|
|
86
|
-
|
|
109
|
+
## Development Installation
|
|
87
110
|
|
|
88
111
|
```bash
|
|
89
|
-
cd
|
|
112
|
+
cd agent-brain-server
|
|
90
113
|
poetry install
|
|
91
114
|
```
|
|
92
115
|
|
|
@@ -101,6 +124,7 @@ cp ../.env.example .env
|
|
|
101
124
|
|
|
102
125
|
Required environment variables:
|
|
103
126
|
- `OPENAI_API_KEY`: Your OpenAI API key for embeddings
|
|
127
|
+
- `ANTHROPIC_API_KEY`: Your Anthropic API key for summarization
|
|
104
128
|
|
|
105
129
|
### Running the Server
|
|
106
130
|
|
|
@@ -112,8 +136,6 @@ poetry run uvicorn agent_brain_server.api.main:app --reload
|
|
|
112
136
|
poetry run agent-brain-serve
|
|
113
137
|
```
|
|
114
138
|
|
|
115
|
-
The server will start at `http://127.0.0.1:8000`.
|
|
116
|
-
|
|
117
139
|
### API Documentation
|
|
118
140
|
|
|
119
141
|
Once running, visit:
|
|
@@ -199,6 +221,22 @@ poetry run ruff check agent_brain_server/
|
|
|
199
221
|
poetry run mypy agent_brain_server/
|
|
200
222
|
```
|
|
201
223
|
|
|
224
|
+
## Documentation
|
|
225
|
+
|
|
226
|
+
- [User Guide](https://github.com/SpillwaveSolutions/agent-brain/wiki/User-Guide) - Getting started and usage
|
|
227
|
+
- [Developer Guide](https://github.com/SpillwaveSolutions/agent-brain/wiki/Developer-Guide) - Contributing and development
|
|
228
|
+
- [API Reference](https://github.com/SpillwaveSolutions/agent-brain/wiki/API-Reference) - Full API documentation
|
|
229
|
+
|
|
230
|
+
## Release Information
|
|
231
|
+
|
|
232
|
+
- **Current Version**: See [pyproject.toml](./pyproject.toml)
|
|
233
|
+
- **Release Notes**: [GitHub Releases](https://github.com/SpillwaveSolutions/agent-brain/releases)
|
|
234
|
+
- **Changelog**: [Latest Release](https://github.com/SpillwaveSolutions/agent-brain/releases/latest)
|
|
235
|
+
|
|
236
|
+
## Related Packages
|
|
237
|
+
|
|
238
|
+
- [agent-brain-cli](https://pypi.org/project/agent-brain-cli/) - Command-line interface for Agent Brain
|
|
239
|
+
|
|
202
240
|
## License
|
|
203
241
|
|
|
204
242
|
MIT
|
|
@@ -1,6 +1,12 @@
|
|
|
1
1
|
# Agent Brain RAG Server
|
|
2
2
|
|
|
3
|
-
|
|
3
|
+
> **Agent Brain** (formerly doc-serve) is an intelligent document indexing and semantic search system designed to give AI agents long-term memory.
|
|
4
|
+
|
|
5
|
+
AI agents need persistent memory to be truly useful. Agent Brain provides the retrieval infrastructure that enables context-aware, knowledge-grounded AI interactions.
|
|
6
|
+
|
|
7
|
+
[](https://pypi.org/project/agent-brain-rag/)
|
|
8
|
+
[](https://www.python.org/downloads/)
|
|
9
|
+
[](https://opensource.org/licenses/MIT)
|
|
4
10
|
|
|
5
11
|
## Installation
|
|
6
12
|
|
|
@@ -25,28 +31,37 @@ The server will start at `http://127.0.0.1:8000`.
|
|
|
25
31
|
|
|
26
32
|
> **Note**: The legacy command `doc-serve` is still available but deprecated. Please use `agent-brain-serve` for new installations.
|
|
27
33
|
|
|
34
|
+
## Search Capabilities
|
|
35
|
+
|
|
36
|
+
Agent Brain provides multiple search strategies to match your retrieval needs:
|
|
37
|
+
|
|
38
|
+
| Search Type | Description | Best For |
|
|
39
|
+
|-------------|-------------|----------|
|
|
40
|
+
| **Semantic Search** | Natural language queries using OpenAI embeddings (`text-embedding-3-large`) | Conceptual questions, finding related content |
|
|
41
|
+
| **Keyword Search (BM25)** | Traditional keyword matching with TF-IDF ranking | Exact matches, technical terms, code identifiers |
|
|
42
|
+
| **Hybrid Search** | Combines vector + BM25 for best of both approaches | General-purpose queries, balanced recall/precision |
|
|
43
|
+
| **GraphRAG** | Knowledge graph-based retrieval for relationship-aware queries | Understanding connections, multi-hop reasoning |
|
|
44
|
+
|
|
28
45
|
## Features
|
|
29
46
|
|
|
30
47
|
- **Document Indexing**: Load and index documents from folders (PDF, Markdown, TXT, DOCX, HTML)
|
|
31
48
|
- **AST-Aware Code Ingestion**: Smart parsing for Python, TypeScript, JavaScript, Java, Go, Rust, C, C++
|
|
32
|
-
- **
|
|
33
|
-
- **Semantic Search**: Query indexed documents using natural language
|
|
49
|
+
- **Multi-Strategy Retrieval**: Semantic, keyword, hybrid, and graph-based search
|
|
34
50
|
- **OpenAI Embeddings**: Uses `text-embedding-3-large` for high-quality embeddings
|
|
51
|
+
- **Claude Summarization**: AI-powered code summaries for better context
|
|
35
52
|
- **Chroma Vector Store**: Persistent, thread-safe vector database
|
|
36
53
|
- **FastAPI**: Modern, high-performance REST API with OpenAPI documentation
|
|
37
54
|
|
|
38
|
-
##
|
|
39
|
-
|
|
40
|
-
### Prerequisites
|
|
55
|
+
## Prerequisites
|
|
41
56
|
|
|
42
57
|
- Python 3.10+
|
|
43
|
-
-
|
|
44
|
-
-
|
|
58
|
+
- OpenAI API key (for embeddings)
|
|
59
|
+
- Anthropic API key (for summarization)
|
|
45
60
|
|
|
46
|
-
|
|
61
|
+
## Development Installation
|
|
47
62
|
|
|
48
63
|
```bash
|
|
49
|
-
cd
|
|
64
|
+
cd agent-brain-server
|
|
50
65
|
poetry install
|
|
51
66
|
```
|
|
52
67
|
|
|
@@ -61,6 +76,7 @@ cp ../.env.example .env
|
|
|
61
76
|
|
|
62
77
|
Required environment variables:
|
|
63
78
|
- `OPENAI_API_KEY`: Your OpenAI API key for embeddings
|
|
79
|
+
- `ANTHROPIC_API_KEY`: Your Anthropic API key for summarization
|
|
64
80
|
|
|
65
81
|
### Running the Server
|
|
66
82
|
|
|
@@ -72,8 +88,6 @@ poetry run uvicorn agent_brain_server.api.main:app --reload
|
|
|
72
88
|
poetry run agent-brain-serve
|
|
73
89
|
```
|
|
74
90
|
|
|
75
|
-
The server will start at `http://127.0.0.1:8000`.
|
|
76
|
-
|
|
77
91
|
### API Documentation
|
|
78
92
|
|
|
79
93
|
Once running, visit:
|
|
@@ -159,6 +173,22 @@ poetry run ruff check agent_brain_server/
|
|
|
159
173
|
poetry run mypy agent_brain_server/
|
|
160
174
|
```
|
|
161
175
|
|
|
176
|
+
## Documentation
|
|
177
|
+
|
|
178
|
+
- [User Guide](https://github.com/SpillwaveSolutions/agent-brain/wiki/User-Guide) - Getting started and usage
|
|
179
|
+
- [Developer Guide](https://github.com/SpillwaveSolutions/agent-brain/wiki/Developer-Guide) - Contributing and development
|
|
180
|
+
- [API Reference](https://github.com/SpillwaveSolutions/agent-brain/wiki/API-Reference) - Full API documentation
|
|
181
|
+
|
|
182
|
+
## Release Information
|
|
183
|
+
|
|
184
|
+
- **Current Version**: See [pyproject.toml](./pyproject.toml)
|
|
185
|
+
- **Release Notes**: [GitHub Releases](https://github.com/SpillwaveSolutions/agent-brain/releases)
|
|
186
|
+
- **Changelog**: [Latest Release](https://github.com/SpillwaveSolutions/agent-brain/releases/latest)
|
|
187
|
+
|
|
188
|
+
## Related Packages
|
|
189
|
+
|
|
190
|
+
- [agent-brain-cli](https://pypi.org/project/agent-brain-cli/) - Command-line interface for Agent Brain
|
|
191
|
+
|
|
162
192
|
## License
|
|
163
193
|
|
|
164
194
|
MIT
|
|
@@ -22,6 +22,10 @@ from fastapi.middleware.cors import CORSMiddleware
|
|
|
22
22
|
|
|
23
23
|
from agent_brain_server import __version__
|
|
24
24
|
from agent_brain_server.config import settings
|
|
25
|
+
from agent_brain_server.config.provider_config import (
|
|
26
|
+
load_provider_settings,
|
|
27
|
+
validate_provider_config,
|
|
28
|
+
)
|
|
25
29
|
from agent_brain_server.indexing.bm25_index import BM25IndexManager
|
|
26
30
|
from agent_brain_server.locking import (
|
|
27
31
|
acquire_lock,
|
|
@@ -66,6 +70,30 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
|
66
70
|
|
|
67
71
|
logger.info("Starting Agent Brain RAG server...")
|
|
68
72
|
|
|
73
|
+
# Load and validate provider configuration
|
|
74
|
+
try:
|
|
75
|
+
provider_settings = load_provider_settings()
|
|
76
|
+
validation_errors = validate_provider_config(provider_settings)
|
|
77
|
+
|
|
78
|
+
if validation_errors:
|
|
79
|
+
for error in validation_errors:
|
|
80
|
+
logger.warning(f"Provider config warning: {error}")
|
|
81
|
+
# Log but don't fail - providers may work if keys are set later
|
|
82
|
+
# or if using Ollama which doesn't need keys
|
|
83
|
+
|
|
84
|
+
# Log active provider configuration
|
|
85
|
+
logger.info(
|
|
86
|
+
f"Embedding provider: {provider_settings.embedding.provider} "
|
|
87
|
+
f"(model: {provider_settings.embedding.model})"
|
|
88
|
+
)
|
|
89
|
+
logger.info(
|
|
90
|
+
f"Summarization provider: {provider_settings.summarization.provider} "
|
|
91
|
+
f"(model: {provider_settings.summarization.model})"
|
|
92
|
+
)
|
|
93
|
+
except Exception as e:
|
|
94
|
+
logger.error(f"Failed to load provider configuration: {e}")
|
|
95
|
+
# Continue with defaults - EmbeddingGenerator will handle provider creation
|
|
96
|
+
|
|
69
97
|
if settings.OPENAI_API_KEY:
|
|
70
98
|
os.environ["OPENAI_API_KEY"] = settings.OPENAI_API_KEY
|
|
71
99
|
|
|
@@ -166,7 +194,7 @@ app = FastAPI(
|
|
|
166
194
|
"RAG-based document indexing and semantic search API. "
|
|
167
195
|
"Index documents from folders and query them using natural language."
|
|
168
196
|
),
|
|
169
|
-
version="
|
|
197
|
+
version="2.0.0",
|
|
170
198
|
lifespan=lifespan,
|
|
171
199
|
docs_url="/docs",
|
|
172
200
|
redoc_url="/redoc",
|
|
@@ -193,7 +221,7 @@ async def root() -> dict[str, str]:
|
|
|
193
221
|
"""Root endpoint redirects to docs."""
|
|
194
222
|
return {
|
|
195
223
|
"name": "Agent Brain RAG API",
|
|
196
|
-
"version": "
|
|
224
|
+
"version": "2.0.0",
|
|
197
225
|
"docs": "/docs",
|
|
198
226
|
"health": "/health",
|
|
199
227
|
}
|
|
@@ -0,0 +1,308 @@
|
|
|
1
|
+
"""Provider configuration models and YAML loader.
|
|
2
|
+
|
|
3
|
+
This module provides Pydantic models for embedding and summarization
|
|
4
|
+
provider configuration, and functions to load configuration from YAML files.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
import logging
|
|
8
|
+
import os
|
|
9
|
+
from functools import lru_cache
|
|
10
|
+
from pathlib import Path
|
|
11
|
+
from typing import Any, Optional
|
|
12
|
+
|
|
13
|
+
import yaml
|
|
14
|
+
from pydantic import BaseModel, Field, field_validator
|
|
15
|
+
|
|
16
|
+
from agent_brain_server.providers.base import (
|
|
17
|
+
EmbeddingProviderType,
|
|
18
|
+
SummarizationProviderType,
|
|
19
|
+
)
|
|
20
|
+
|
|
21
|
+
logger = logging.getLogger(__name__)
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class EmbeddingConfig(BaseModel):
|
|
25
|
+
"""Configuration for embedding provider."""
|
|
26
|
+
|
|
27
|
+
provider: EmbeddingProviderType = Field(
|
|
28
|
+
default=EmbeddingProviderType.OPENAI,
|
|
29
|
+
description="Embedding provider to use",
|
|
30
|
+
)
|
|
31
|
+
model: str = Field(
|
|
32
|
+
default="text-embedding-3-large",
|
|
33
|
+
description="Model name for embeddings",
|
|
34
|
+
)
|
|
35
|
+
api_key_env: Optional[str] = Field(
|
|
36
|
+
default="OPENAI_API_KEY",
|
|
37
|
+
description="Environment variable name containing API key",
|
|
38
|
+
)
|
|
39
|
+
base_url: Optional[str] = Field(
|
|
40
|
+
default=None,
|
|
41
|
+
description="Custom base URL (for Ollama or compatible APIs)",
|
|
42
|
+
)
|
|
43
|
+
params: dict[str, Any] = Field(
|
|
44
|
+
default_factory=dict,
|
|
45
|
+
description="Provider-specific parameters",
|
|
46
|
+
)
|
|
47
|
+
|
|
48
|
+
model_config = {"use_enum_values": True}
|
|
49
|
+
|
|
50
|
+
@field_validator("provider", mode="before")
|
|
51
|
+
@classmethod
|
|
52
|
+
def validate_provider(cls, v: Any) -> EmbeddingProviderType:
|
|
53
|
+
"""Convert string to enum if needed."""
|
|
54
|
+
if isinstance(v, str):
|
|
55
|
+
return EmbeddingProviderType(v.lower())
|
|
56
|
+
if isinstance(v, EmbeddingProviderType):
|
|
57
|
+
return v
|
|
58
|
+
return EmbeddingProviderType(v)
|
|
59
|
+
|
|
60
|
+
def get_api_key(self) -> Optional[str]:
|
|
61
|
+
"""Resolve API key from environment variable.
|
|
62
|
+
|
|
63
|
+
Returns:
|
|
64
|
+
API key value or None if not found/not needed
|
|
65
|
+
"""
|
|
66
|
+
if self.provider == EmbeddingProviderType.OLLAMA:
|
|
67
|
+
return None # Ollama doesn't need API key
|
|
68
|
+
if self.api_key_env:
|
|
69
|
+
return os.getenv(self.api_key_env)
|
|
70
|
+
return None
|
|
71
|
+
|
|
72
|
+
def get_base_url(self) -> Optional[str]:
|
|
73
|
+
"""Get base URL with defaults for specific providers.
|
|
74
|
+
|
|
75
|
+
Returns:
|
|
76
|
+
Base URL for the provider
|
|
77
|
+
"""
|
|
78
|
+
if self.base_url:
|
|
79
|
+
return self.base_url
|
|
80
|
+
if self.provider == EmbeddingProviderType.OLLAMA:
|
|
81
|
+
return "http://localhost:11434/v1"
|
|
82
|
+
return None
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
class SummarizationConfig(BaseModel):
|
|
86
|
+
"""Configuration for summarization provider."""
|
|
87
|
+
|
|
88
|
+
provider: SummarizationProviderType = Field(
|
|
89
|
+
default=SummarizationProviderType.ANTHROPIC,
|
|
90
|
+
description="Summarization provider to use",
|
|
91
|
+
)
|
|
92
|
+
model: str = Field(
|
|
93
|
+
default="claude-haiku-4-5-20251001",
|
|
94
|
+
description="Model name for summarization",
|
|
95
|
+
)
|
|
96
|
+
api_key_env: Optional[str] = Field(
|
|
97
|
+
default="ANTHROPIC_API_KEY",
|
|
98
|
+
description="Environment variable name containing API key",
|
|
99
|
+
)
|
|
100
|
+
base_url: Optional[str] = Field(
|
|
101
|
+
default=None,
|
|
102
|
+
description="Custom base URL (for Grok or Ollama)",
|
|
103
|
+
)
|
|
104
|
+
params: dict[str, Any] = Field(
|
|
105
|
+
default_factory=dict,
|
|
106
|
+
description="Provider-specific parameters (max_tokens, temperature)",
|
|
107
|
+
)
|
|
108
|
+
|
|
109
|
+
model_config = {"use_enum_values": True}
|
|
110
|
+
|
|
111
|
+
@field_validator("provider", mode="before")
|
|
112
|
+
@classmethod
|
|
113
|
+
def validate_provider(cls, v: Any) -> SummarizationProviderType:
|
|
114
|
+
"""Convert string to enum if needed."""
|
|
115
|
+
if isinstance(v, str):
|
|
116
|
+
return SummarizationProviderType(v.lower())
|
|
117
|
+
if isinstance(v, SummarizationProviderType):
|
|
118
|
+
return v
|
|
119
|
+
return SummarizationProviderType(v)
|
|
120
|
+
|
|
121
|
+
def get_api_key(self) -> Optional[str]:
|
|
122
|
+
"""Resolve API key from environment variable.
|
|
123
|
+
|
|
124
|
+
Returns:
|
|
125
|
+
API key value or None if not found/not needed
|
|
126
|
+
"""
|
|
127
|
+
if self.provider == SummarizationProviderType.OLLAMA:
|
|
128
|
+
return None # Ollama doesn't need API key
|
|
129
|
+
if self.api_key_env:
|
|
130
|
+
return os.getenv(self.api_key_env)
|
|
131
|
+
return None
|
|
132
|
+
|
|
133
|
+
def get_base_url(self) -> Optional[str]:
|
|
134
|
+
"""Get base URL with defaults for specific providers.
|
|
135
|
+
|
|
136
|
+
Returns:
|
|
137
|
+
Base URL for the provider
|
|
138
|
+
"""
|
|
139
|
+
if self.base_url:
|
|
140
|
+
return self.base_url
|
|
141
|
+
if self.provider == SummarizationProviderType.OLLAMA:
|
|
142
|
+
return "http://localhost:11434/v1"
|
|
143
|
+
if self.provider == SummarizationProviderType.GROK:
|
|
144
|
+
return "https://api.x.ai/v1"
|
|
145
|
+
return None
|
|
146
|
+
|
|
147
|
+
|
|
148
|
+
class ProviderSettings(BaseModel):
|
|
149
|
+
"""Top-level provider configuration."""
|
|
150
|
+
|
|
151
|
+
embedding: EmbeddingConfig = Field(
|
|
152
|
+
default_factory=EmbeddingConfig,
|
|
153
|
+
description="Embedding provider configuration",
|
|
154
|
+
)
|
|
155
|
+
summarization: SummarizationConfig = Field(
|
|
156
|
+
default_factory=SummarizationConfig,
|
|
157
|
+
description="Summarization provider configuration",
|
|
158
|
+
)
|
|
159
|
+
|
|
160
|
+
|
|
161
|
+
def _find_config_file() -> Optional[Path]:
|
|
162
|
+
"""Find the configuration file in standard locations.
|
|
163
|
+
|
|
164
|
+
Search order:
|
|
165
|
+
1. DOC_SERVE_CONFIG environment variable
|
|
166
|
+
2. Current directory config.yaml
|
|
167
|
+
3. State directory config.yaml (if DOC_SERVE_STATE_DIR set)
|
|
168
|
+
4. Project root config.yaml
|
|
169
|
+
|
|
170
|
+
Returns:
|
|
171
|
+
Path to config file or None if not found
|
|
172
|
+
"""
|
|
173
|
+
# 1. Environment variable override
|
|
174
|
+
env_config = os.getenv("DOC_SERVE_CONFIG")
|
|
175
|
+
if env_config:
|
|
176
|
+
path = Path(env_config)
|
|
177
|
+
if path.exists():
|
|
178
|
+
return path
|
|
179
|
+
logger.warning(f"DOC_SERVE_CONFIG points to non-existent file: {env_config}")
|
|
180
|
+
|
|
181
|
+
# 2. Current directory
|
|
182
|
+
cwd_config = Path.cwd() / "config.yaml"
|
|
183
|
+
if cwd_config.exists():
|
|
184
|
+
return cwd_config
|
|
185
|
+
|
|
186
|
+
# 3. State directory
|
|
187
|
+
state_dir = os.getenv("DOC_SERVE_STATE_DIR")
|
|
188
|
+
if state_dir:
|
|
189
|
+
state_config = Path(state_dir) / "config.yaml"
|
|
190
|
+
if state_config.exists():
|
|
191
|
+
return state_config
|
|
192
|
+
|
|
193
|
+
# 4. .claude/doc-serve directory (project root pattern)
|
|
194
|
+
claude_dir = Path.cwd() / ".claude" / "doc-serve"
|
|
195
|
+
if claude_dir.exists():
|
|
196
|
+
claude_config = claude_dir / "config.yaml"
|
|
197
|
+
if claude_config.exists():
|
|
198
|
+
return claude_config
|
|
199
|
+
|
|
200
|
+
return None
|
|
201
|
+
|
|
202
|
+
|
|
203
|
+
def _load_yaml_config(path: Path) -> dict[str, Any]:
|
|
204
|
+
"""Load YAML configuration from file.
|
|
205
|
+
|
|
206
|
+
Args:
|
|
207
|
+
path: Path to YAML config file
|
|
208
|
+
|
|
209
|
+
Returns:
|
|
210
|
+
Configuration dictionary
|
|
211
|
+
|
|
212
|
+
Raises:
|
|
213
|
+
ConfigurationError: If YAML parsing fails
|
|
214
|
+
"""
|
|
215
|
+
from agent_brain_server.providers.exceptions import ConfigurationError
|
|
216
|
+
|
|
217
|
+
try:
|
|
218
|
+
with open(path) as f:
|
|
219
|
+
config = yaml.safe_load(f)
|
|
220
|
+
return config if config else {}
|
|
221
|
+
except yaml.YAMLError as e:
|
|
222
|
+
raise ConfigurationError(
|
|
223
|
+
f"Failed to parse config file {path}: {e}",
|
|
224
|
+
"config",
|
|
225
|
+
) from e
|
|
226
|
+
except OSError as e:
|
|
227
|
+
raise ConfigurationError(
|
|
228
|
+
f"Failed to read config file {path}: {e}",
|
|
229
|
+
"config",
|
|
230
|
+
) from e
|
|
231
|
+
|
|
232
|
+
|
|
233
|
+
@lru_cache
|
|
234
|
+
def load_provider_settings() -> ProviderSettings:
|
|
235
|
+
"""Load provider settings from YAML config or defaults.
|
|
236
|
+
|
|
237
|
+
This function:
|
|
238
|
+
1. Searches for config.yaml in standard locations
|
|
239
|
+
2. Parses YAML and validates against Pydantic models
|
|
240
|
+
3. Falls back to defaults (OpenAI embeddings + Anthropic summarization)
|
|
241
|
+
|
|
242
|
+
Returns:
|
|
243
|
+
Validated ProviderSettings instance
|
|
244
|
+
"""
|
|
245
|
+
config_path = _find_config_file()
|
|
246
|
+
|
|
247
|
+
if config_path:
|
|
248
|
+
logger.info(f"Loading provider config from {config_path}")
|
|
249
|
+
raw_config = _load_yaml_config(config_path)
|
|
250
|
+
settings = ProviderSettings(**raw_config)
|
|
251
|
+
else:
|
|
252
|
+
logger.info("No config file found, using default providers")
|
|
253
|
+
settings = ProviderSettings()
|
|
254
|
+
|
|
255
|
+
# Log active configuration
|
|
256
|
+
logger.info(
|
|
257
|
+
f"Active embedding provider: {settings.embedding.provider} "
|
|
258
|
+
f"(model: {settings.embedding.model})"
|
|
259
|
+
)
|
|
260
|
+
logger.info(
|
|
261
|
+
f"Active summarization provider: {settings.summarization.provider} "
|
|
262
|
+
f"(model: {settings.summarization.model})"
|
|
263
|
+
)
|
|
264
|
+
|
|
265
|
+
return settings
|
|
266
|
+
|
|
267
|
+
|
|
268
|
+
def clear_settings_cache() -> None:
|
|
269
|
+
"""Clear the cached provider settings (for testing)."""
|
|
270
|
+
load_provider_settings.cache_clear()
|
|
271
|
+
|
|
272
|
+
|
|
273
|
+
def validate_provider_config(settings: ProviderSettings) -> list[str]:
|
|
274
|
+
"""Validate provider configuration and return list of errors.
|
|
275
|
+
|
|
276
|
+
Checks:
|
|
277
|
+
- API keys are available for providers that need them
|
|
278
|
+
- Models are known for the selected provider
|
|
279
|
+
|
|
280
|
+
Args:
|
|
281
|
+
settings: Provider settings to validate
|
|
282
|
+
|
|
283
|
+
Returns:
|
|
284
|
+
List of validation error messages (empty if valid)
|
|
285
|
+
"""
|
|
286
|
+
errors: list[str] = []
|
|
287
|
+
|
|
288
|
+
# Validate embedding provider
|
|
289
|
+
if settings.embedding.provider != EmbeddingProviderType.OLLAMA:
|
|
290
|
+
api_key = settings.embedding.get_api_key()
|
|
291
|
+
if not api_key:
|
|
292
|
+
env_var = settings.embedding.api_key_env or "OPENAI_API_KEY"
|
|
293
|
+
errors.append(
|
|
294
|
+
f"Missing API key for {settings.embedding.provider} embeddings. "
|
|
295
|
+
f"Set {env_var} environment variable."
|
|
296
|
+
)
|
|
297
|
+
|
|
298
|
+
# Validate summarization provider
|
|
299
|
+
if settings.summarization.provider != SummarizationProviderType.OLLAMA:
|
|
300
|
+
api_key = settings.summarization.get_api_key()
|
|
301
|
+
if not api_key:
|
|
302
|
+
env_var = settings.summarization.api_key_env or "ANTHROPIC_API_KEY"
|
|
303
|
+
errors.append(
|
|
304
|
+
f"Missing API key for {settings.summarization.provider} summarization. "
|
|
305
|
+
f"Set {env_var} environment variable."
|
|
306
|
+
)
|
|
307
|
+
|
|
308
|
+
return errors
|
|
@@ -26,7 +26,7 @@ class Settings(BaseSettings):
|
|
|
26
26
|
|
|
27
27
|
# Anthropic Configuration
|
|
28
28
|
ANTHROPIC_API_KEY: str = ""
|
|
29
|
-
CLAUDE_MODEL: str = "claude-
|
|
29
|
+
CLAUDE_MODEL: str = "claude-haiku-4-5-20251001" # Claude 4.5 Haiku (latest)
|
|
30
30
|
|
|
31
31
|
# Chroma Configuration
|
|
32
32
|
CHROMA_PERSIST_DIR: str = "./chroma_db"
|
|
@@ -51,6 +51,17 @@ class Settings(BaseSettings):
|
|
|
51
51
|
DOC_SERVE_STATE_DIR: Optional[str] = None # Override state directory
|
|
52
52
|
DOC_SERVE_MODE: str = "project" # "project" or "shared"
|
|
53
53
|
|
|
54
|
+
# GraphRAG Configuration (Feature 113)
|
|
55
|
+
ENABLE_GRAPH_INDEX: bool = False # Master switch for graph indexing
|
|
56
|
+
GRAPH_STORE_TYPE: str = "simple" # "simple" (in-memory) or "kuzu" (persistent)
|
|
57
|
+
GRAPH_INDEX_PATH: str = "./graph_index" # Path for graph persistence
|
|
58
|
+
GRAPH_EXTRACTION_MODEL: str = "claude-haiku-4-5" # Model for entity extraction
|
|
59
|
+
GRAPH_MAX_TRIPLETS_PER_CHUNK: int = 10 # Max triplets per document chunk
|
|
60
|
+
GRAPH_USE_CODE_METADATA: bool = True # Use AST metadata for code entities
|
|
61
|
+
GRAPH_USE_LLM_EXTRACTION: bool = True # Use LLM for additional extraction
|
|
62
|
+
GRAPH_TRAVERSAL_DEPTH: int = 2 # Depth for graph traversal in queries
|
|
63
|
+
GRAPH_RRF_K: int = 60 # Reciprocal Rank Fusion constant for multi-retrieval
|
|
64
|
+
|
|
54
65
|
model_config = SettingsConfigDict(
|
|
55
66
|
env_file=[
|
|
56
67
|
".env", # Current directory
|
|
@@ -7,6 +7,18 @@ from agent_brain_server.indexing.embedding import (
|
|
|
7
7
|
EmbeddingGenerator,
|
|
8
8
|
get_embedding_generator,
|
|
9
9
|
)
|
|
10
|
+
from agent_brain_server.indexing.graph_extractors import (
|
|
11
|
+
CodeMetadataExtractor,
|
|
12
|
+
LLMEntityExtractor,
|
|
13
|
+
get_code_extractor,
|
|
14
|
+
get_llm_extractor,
|
|
15
|
+
reset_extractors,
|
|
16
|
+
)
|
|
17
|
+
from agent_brain_server.indexing.graph_index import (
|
|
18
|
+
GraphIndexManager,
|
|
19
|
+
get_graph_index_manager,
|
|
20
|
+
reset_graph_index_manager,
|
|
21
|
+
)
|
|
10
22
|
|
|
11
23
|
__all__ = [
|
|
12
24
|
"DocumentLoader",
|
|
@@ -16,4 +28,13 @@ __all__ = [
|
|
|
16
28
|
"get_embedding_generator",
|
|
17
29
|
"BM25IndexManager",
|
|
18
30
|
"get_bm25_manager",
|
|
31
|
+
# Graph indexing (Feature 113)
|
|
32
|
+
"LLMEntityExtractor",
|
|
33
|
+
"CodeMetadataExtractor",
|
|
34
|
+
"get_llm_extractor",
|
|
35
|
+
"get_code_extractor",
|
|
36
|
+
"reset_extractors",
|
|
37
|
+
"GraphIndexManager",
|
|
38
|
+
"get_graph_index_manager",
|
|
39
|
+
"reset_graph_index_manager",
|
|
19
40
|
]
|