agent-brain-rag 1.2.0__py3-none-any.whl → 2.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.
Files changed (38) hide show
  1. {agent_brain_rag-1.2.0.dist-info → agent_brain_rag-2.0.0.dist-info}/METADATA +54 -16
  2. agent_brain_rag-2.0.0.dist-info/RECORD +50 -0
  3. agent_brain_server/__init__.py +1 -1
  4. agent_brain_server/api/main.py +30 -2
  5. agent_brain_server/api/routers/health.py +1 -0
  6. agent_brain_server/config/provider_config.py +308 -0
  7. agent_brain_server/config/settings.py +12 -1
  8. agent_brain_server/indexing/__init__.py +21 -0
  9. agent_brain_server/indexing/embedding.py +86 -135
  10. agent_brain_server/indexing/graph_extractors.py +582 -0
  11. agent_brain_server/indexing/graph_index.py +536 -0
  12. agent_brain_server/models/__init__.py +9 -0
  13. agent_brain_server/models/graph.py +253 -0
  14. agent_brain_server/models/health.py +15 -3
  15. agent_brain_server/models/query.py +14 -1
  16. agent_brain_server/providers/__init__.py +64 -0
  17. agent_brain_server/providers/base.py +251 -0
  18. agent_brain_server/providers/embedding/__init__.py +23 -0
  19. agent_brain_server/providers/embedding/cohere.py +163 -0
  20. agent_brain_server/providers/embedding/ollama.py +150 -0
  21. agent_brain_server/providers/embedding/openai.py +118 -0
  22. agent_brain_server/providers/exceptions.py +95 -0
  23. agent_brain_server/providers/factory.py +157 -0
  24. agent_brain_server/providers/summarization/__init__.py +41 -0
  25. agent_brain_server/providers/summarization/anthropic.py +87 -0
  26. agent_brain_server/providers/summarization/gemini.py +96 -0
  27. agent_brain_server/providers/summarization/grok.py +95 -0
  28. agent_brain_server/providers/summarization/ollama.py +114 -0
  29. agent_brain_server/providers/summarization/openai.py +87 -0
  30. agent_brain_server/services/indexing_service.py +39 -0
  31. agent_brain_server/services/query_service.py +203 -0
  32. agent_brain_server/storage/__init__.py +18 -2
  33. agent_brain_server/storage/graph_store.py +519 -0
  34. agent_brain_server/storage/vector_store.py +35 -0
  35. agent_brain_server/storage_paths.py +2 -0
  36. agent_brain_rag-1.2.0.dist-info/RECORD +0 -31
  37. {agent_brain_rag-1.2.0.dist-info → agent_brain_rag-2.0.0.dist-info}/WHEEL +0 -0
  38. {agent_brain_rag-1.2.0.dist-info → agent_brain_rag-2.0.0.dist-info}/entry_points.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: agent-brain-rag
3
- Version: 1.2.0
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/doc-serve#readme
37
- Project-URL: Homepage, https://github.com/SpillwaveSolutions/doc-serve
38
- Project-URL: Repository, https://github.com/SpillwaveSolutions/doc-serve
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
- Intelligent document indexing and semantic search REST API service that gives AI agents long-term memory.
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
+ [![PyPI version](https://badge.fury.io/py/agent-brain-rag.svg)](https://pypi.org/project/agent-brain-rag/)
56
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
57
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](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
- - **Hybrid Search**: Combines BM25 keyword matching with semantic similarity
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
- ## Quick Start
79
-
80
- ### Prerequisites
103
+ ## Prerequisites
81
104
 
82
105
  - Python 3.10+
83
- - Poetry
84
- - OpenAI API key
106
+ - OpenAI API key (for embeddings)
107
+ - Anthropic API key (for summarization)
85
108
 
86
- ### Installation
109
+ ## Development Installation
87
110
 
88
111
  ```bash
89
- cd doc-serve-server
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
@@ -0,0 +1,50 @@
1
+ agent_brain_server/__init__.py,sha256=W5Jm_aOWbsUD605_og-pejqYr3kW_gh5Mf7e4kLOviU,95
2
+ agent_brain_server/api/__init__.py,sha256=nvTvO_ahHAAsRDlV3dL_JlNruSdan4kav_P5sT_1PFk,93
3
+ agent_brain_server/api/main.py,sha256=KMPrhH9EsvuFa6aAZlui3og8wm6WoxLwZkj4LTKRiHM,12894
4
+ agent_brain_server/api/routers/__init__.py,sha256=Z3PUKDwxeI8T88xMQeTcQ8nq1bQnu0nYjRnqNIsDUyY,254
5
+ agent_brain_server/api/routers/health.py,sha256=xzgjFlSndRx9jT_G4kOl79MiWv06ZUcXTAE7z2jjANs,3538
6
+ agent_brain_server/api/routers/index.py,sha256=jTfeY8EIbM1tksJY6jNDw6qkmKBMaAi1ufDIdrg80WQ,6756
7
+ agent_brain_server/api/routers/query.py,sha256=TAbpuCybjacUo_-7Zm3Jt8EW8EMDTYJDLv6aGtjaUzQ,2793
8
+ agent_brain_server/config/__init__.py,sha256=zzDErZGUBwUm5Fk43OHJN7eWpeIw_1kWdnhsN6QQqSc,84
9
+ agent_brain_server/config/provider_config.py,sha256=6c1CaE-_Az8PGglSxJVJjR0ns2ljEYyTSJyUxAVgZO4,9552
10
+ agent_brain_server/config/settings.py,sha256=cXVsHeCQjGf9I_N1FfkMFBBHEeukqQOdJ7zMIMAWOXg,3380
11
+ agent_brain_server/indexing/__init__.py,sha256=tQiFW4945h-ENd0pmMoGhCbm1Fz3AaREKBgajHeW-Us,1169
12
+ agent_brain_server/indexing/bm25_index.py,sha256=CIeGZIa8PwvOm_Fnz3EMgETvM2Mu6GzANWeUims0P0Q,5268
13
+ agent_brain_server/indexing/chunking.py,sha256=kW-ZOdwOS-PfXuPAYb8qL4-7yAY-rpY9a1el2YCN9CA,29848
14
+ agent_brain_server/indexing/document_loader.py,sha256=D6U3mDl_2jFfT_JxD-5yudtqlJW8Fh0zj5lhz0V0tPs,16090
15
+ agent_brain_server/indexing/embedding.py,sha256=i3bN3SivvVB8iNiYDjrDeB8nF3WkFoOy5wHrDaNosGs,7410
16
+ agent_brain_server/indexing/graph_extractors.py,sha256=XEWBqlqDffxbs5ZLaijGICXeaptGtkUGNEhJblP3jBU,19339
17
+ agent_brain_server/indexing/graph_index.py,sha256=ZKwICIqR2stcPEqNQ_SSuSk-O0Hj1lUmiUMxm5udhjg,18584
18
+ agent_brain_server/locking.py,sha256=yRpswpBem4964gTh1VH4VUPT-vBIRnNEVgWdorJA4Hg,3365
19
+ agent_brain_server/models/__init__.py,sha256=5KAkdcB8rR5B0Z9IY2ZxtrUGSglFtxNeu6ZwGHTVDAU,737
20
+ agent_brain_server/models/graph.py,sha256=YzdjJj_X7hueJt4v8Gsoly4C-qquJLZi1d9epIh_Mos,7973
21
+ agent_brain_server/models/health.py,sha256=kwSuR64ldDnrim7PO_4tZLhejdFRKbqkdu4mDf6Z7gQ,4160
22
+ agent_brain_server/models/index.py,sha256=pjDv7phLS6dpiHLlEtcAuXQN_aHIfA_4lMkAZ-NkXZQ,5400
23
+ agent_brain_server/models/query.py,sha256=mrb8Kjvg6DQZJs8aGEqic54FpzSr62qrSPDFXZeUfpM,6613
24
+ agent_brain_server/project_root.py,sha256=HIY5NMRDYWYIT7K6B_UMOGo1zXn9zwAdGI6ApViKI_8,2194
25
+ agent_brain_server/providers/__init__.py,sha256=04wJ_QhGKRM55R4yrbw3J3kCtKX44_FHzh6xZq1xFdc,1670
26
+ agent_brain_server/providers/base.py,sha256=yLIpz9W0xlilBXxuFJRNP_gbGDsZ9HeRSqEnIsVGQjU,6933
27
+ agent_brain_server/providers/embedding/__init__.py,sha256=WRUAkOLLgQB9AR9I-_m3QWa55hSeC5tDdp3gZQFfMRM,982
28
+ agent_brain_server/providers/embedding/cohere.py,sha256=9ktr0_ZHmocJE8sFqo1Ah68B9J0wgs9-pLS0aKmLtQE,5295
29
+ agent_brain_server/providers/embedding/ollama.py,sha256=oXKbbEivY74ASVjayLdr3xfgXTMvFS_jAjZs5HqCHzc,4725
30
+ agent_brain_server/providers/embedding/openai.py,sha256=B5F6OJ_PSpF7OLUlhhmqGWOvho3Kox29Lq52bvxrWWY,3576
31
+ agent_brain_server/providers/exceptions.py,sha256=kYggL4lG1NmJrEyU040Al5rm3_nWmKROJemkjcdohCM,2836
32
+ agent_brain_server/providers/factory.py,sha256=3sREN_VnTDam9MDQJzf7eYlX5k-tbzgWXHb7DQ6MhRw,5488
33
+ agent_brain_server/providers/summarization/__init__.py,sha256=WfMXywPvY1uTbmy_Z5prK9FNNqv17TeEvdg2MqFyKAc,1545
34
+ agent_brain_server/providers/summarization/anthropic.py,sha256=fZOUfxJCPRUxMdeakBAdw1FY6NaYhvlWBlheTJlyAtc,2704
35
+ agent_brain_server/providers/summarization/gemini.py,sha256=t2pOC4R2ae7WQAq7WCy8ZD6Ku_RPjyqsktBZRuRwGi8,2960
36
+ agent_brain_server/providers/summarization/grok.py,sha256=nK3i46NbUcVxBqWNbJ-IifahB450Blnqm0rBhSQnQOw,2875
37
+ agent_brain_server/providers/summarization/ollama.py,sha256=mw4zvx6o2nC4xoz1W6F47x3fk7gidl7aeh1uAnSTHO8,3673
38
+ agent_brain_server/providers/summarization/openai.py,sha256=HjZJIjwc0usHejRAFW3JT6a7cmM_Ib2FdC05y_iEwA4,2632
39
+ agent_brain_server/runtime.py,sha256=bcchfDBt_tn2_r-lFhxqKp0RncKrz152D-oM5RcKGrU,3076
40
+ agent_brain_server/services/__init__.py,sha256=E4VPN9Rqa2mxGQQEQn-5IYj63LSPTrA8aIx8ENO5xcc,296
41
+ agent_brain_server/services/indexing_service.py,sha256=d_oj3n7gNgua7nmiLNTM-hDp4wvcrkos9kCRuOTzlpc,20338
42
+ agent_brain_server/services/query_service.py,sha256=oc_BNqTNjAtz17y_F7r-egP7Z0GozmP_mlfSs1jKbMM,22910
43
+ agent_brain_server/storage/__init__.py,sha256=xYk4MfrEj-IpGZ0n4myfy5dwyjKEvprS7h68bet0ooU,557
44
+ agent_brain_server/storage/graph_store.py,sha256=boW6_L0suvHuOdfZdPMdSNSb8U-am7KLelsa3K2bKl4,18168
45
+ agent_brain_server/storage/vector_store.py,sha256=rcI--g6g7mqUbZzewOsaMdb7uqw173bAZau1p0oYZ68,12135
46
+ agent_brain_server/storage_paths.py,sha256=OvCXjSDu_L_udlxWo4OoOd_J-mcBcfoCSNwHRCnFHHk,1844
47
+ agent_brain_rag-2.0.0.dist-info/METADATA,sha256=KLRRJQO2nCyqwvZuuovxSh-IrBmb7LUtgOnx-0EqcQs,8155
48
+ agent_brain_rag-2.0.0.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
49
+ agent_brain_rag-2.0.0.dist-info/entry_points.txt,sha256=r6kfBauVV9OtMXQ7trM9ofbW5N_kdzqTouoJ9zc31-I,122
50
+ agent_brain_rag-2.0.0.dist-info/RECORD,,
@@ -1,3 +1,3 @@
1
1
  """Doc-Serve Server - RAG-based document indexing and query service."""
2
2
 
3
- __version__ = "1.2.0"
3
+ __version__ = "2.0.0"
@@ -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="1.2.0",
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": "1.2.0",
224
+ "version": "2.0.0",
197
225
  "docs": "/docs",
198
226
  "health": "/health",
199
227
  }
@@ -98,4 +98,5 @@ async def indexing_status(request: Request) -> IndexingStatus:
98
98
  ),
99
99
  indexed_folders=status["indexed_folders"],
100
100
  supported_languages=status.get("supported_languages", []),
101
+ graph_index=status.get("graph_index"),
101
102
  )
@@ -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-3-5-haiku-20241022" # Claude 3.5 Haiku (latest)
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
  ]