agent-brain-rag 1.1.0__tar.gz → 2.0.0__tar.gz

{agent_brain_rag-1.1.0 → agent_brain_rag-2.0.0}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.3
 Name: agent-brain-rag
-Version: 1.1.0
-Summary: RAG-based document indexing and semantic search server for AI agents
+Version: 2.0.0
+Summary: Agent Brain RAG - Intelligent document indexing and semantic search server that gives AI agents long-term memory
 License: MIT
-Keywords: rag,semantic-search,documentation,indexing,llama-index,chromadb,ai-agent,brain
+Keywords: agent-brain,rag,semantic-search,ai-memory,llm-memory,documentation,indexing,llama-index,chromadb,ai-agent,claude-code,agent-memory
 Author: Spillwave Solutions
 Requires-Python: >=3.10,<4.0
 Classifier: Development Status :: 4 - Beta
@@ -16,12 +16,19 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Software Development :: Documentation
 Classifier: Topic :: Text Processing :: Indexing
+Provides-Extra: graphrag
+Provides-Extra: graphrag-all
+Provides-Extra: graphrag-kuzu
 Requires-Dist: anthropic (>=0.40.0,<0.41.0)
 Requires-Dist: chromadb (>=0.5.0,<0.6.0)
 Requires-Dist: click (>=8.1.0,<9.0.0)
+Requires-Dist: cohere (>=5.0.0,<6.0.0)
 Requires-Dist: fastapi (>=0.115.0,<0.116.0)
+Requires-Dist: google-generativeai (>=0.8.0,<0.9.0)
+Requires-Dist: langextract (>=1.0.0,<2.0.0) ; extra == "graphrag" or extra == "graphrag-all"
 Requires-Dist: llama-index-core (>=0.14.0,<0.15.0)
 Requires-Dist: llama-index-embeddings-openai (>=0.5.0,<0.6.0)
+Requires-Dist: llama-index-graph-stores-kuzu (>=0.9.0,<0.10.0) ; extra == "graphrag-kuzu" or extra == "graphrag-all"
 Requires-Dist: llama-index-llms-openai (>=0.6.12,<0.7.0)
 Requires-Dist: llama-index-readers-file (>=0.5.0,<0.6.0)
 Requires-Dist: llama-index-retrievers-bm25 (>=0.6.0,<0.7.0)
@@ -29,23 +36,30 @@ Requires-Dist: openai (>=1.57.0,<2.0.0)
 Requires-Dist: pydantic (>=2.10.0,<3.0.0)
 Requires-Dist: pydantic-settings (>=2.6.0,<3.0.0)
 Requires-Dist: python-dotenv (>=1.0.0,<2.0.0)
+Requires-Dist: pyyaml (>=6.0.0,<7.0.0)
 Requires-Dist: rank-bm25 (>=0.2.2,<0.3.0)
 Requires-Dist: tiktoken (>=0.8.0,<0.9.0)
 Requires-Dist: tree-sitter-language-pack (>=0.7.3,<0.8.0)
 Requires-Dist: uvicorn[standard] (>=0.32.0,<0.33.0)
-Project-URL: Documentation, https://github.com/SpillwaveSolutions/doc-serve#readme
-Project-URL: Homepage, https://github.com/SpillwaveSolutions/doc-serve
-Project-URL: Repository, https://github.com/SpillwaveSolutions/doc-serve
+Project-URL: Documentation, https://github.com/SpillwaveSolutions/agent-brain/wiki
+Project-URL: Homepage, https://github.com/SpillwaveSolutions/agent-brain
+Project-URL: Repository, https://github.com/SpillwaveSolutions/agent-brain
 Description-Content-Type: text/markdown
 
-# Doc-Serve Server
+# Agent Brain RAG Server
 
-RAG-based document indexing and semantic search REST API service.
+> **Agent Brain** (formerly doc-serve) is an intelligent document indexing and semantic search system designed to give AI agents long-term memory.
+
+AI agents need persistent memory to be truly useful. Agent Brain provides the retrieval infrastructure that enables context-aware, knowledge-grounded AI interactions.
+
+[![PyPI version](https://badge.fury.io/py/agent-brain-rag.svg)](https://pypi.org/project/agent-brain-rag/)
+[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 
 ## Installation
 
 ```bash
-pip install doc-serve
+pip install agent-brain-rag
 ```
 
 ## Quick Start
@@ -58,32 +72,44 @@ pip install doc-serve
 
 2. Start the server:
    ```bash
-   doc-serve
+   agent-brain-serve
    ```
 
 The server will start at `http://127.0.0.1:8000`.
 
+> **Note**: The legacy command `doc-serve` is still available but deprecated. Please use `agent-brain-serve` for new installations.
+
+## Search Capabilities
+
+Agent Brain provides multiple search strategies to match your retrieval needs:
+
+| Search Type | Description | Best For |
+|-------------|-------------|----------|
+| **Semantic Search** | Natural language queries using OpenAI embeddings (`text-embedding-3-large`) | Conceptual questions, finding related content |
+| **Keyword Search (BM25)** | Traditional keyword matching with TF-IDF ranking | Exact matches, technical terms, code identifiers |
+| **Hybrid Search** | Combines vector + BM25 for best of both approaches | General-purpose queries, balanced recall/precision |
+| **GraphRAG** | Knowledge graph-based retrieval for relationship-aware queries | Understanding connections, multi-hop reasoning |
+
 ## Features
 
 - **Document Indexing**: Load and index documents from folders (PDF, Markdown, TXT, DOCX, HTML)
-- **Context-Aware Chunking**: Smart text splitting with configurable chunk sizes and overlap
-- **Semantic Search**: Query indexed documents using natural language
+- **AST-Aware Code Ingestion**: Smart parsing for Python, TypeScript, JavaScript, Java, Go, Rust, C, C++
+- **Multi-Strategy Retrieval**: Semantic, keyword, hybrid, and graph-based search
 - **OpenAI Embeddings**: Uses `text-embedding-3-large` for high-quality embeddings
+- **Claude Summarization**: AI-powered code summaries for better context
 - **Chroma Vector Store**: Persistent, thread-safe vector database
 - **FastAPI**: Modern, high-performance REST API with OpenAPI documentation
 
-## Quick Start
-
-### Prerequisites
+## Prerequisites
 
 - Python 3.10+
-- Poetry
-- OpenAI API key
+- OpenAI API key (for embeddings)
+- Anthropic API key (for summarization)
 
-### Installation
+## Development Installation
 
 ```bash
-cd doc-serve-server
+cd agent-brain-server
 poetry install
 ```
 
@@ -98,19 +124,18 @@ cp ../.env.example .env
 
 Required environment variables:
 - `OPENAI_API_KEY`: Your OpenAI API key for embeddings
+- `ANTHROPIC_API_KEY`: Your Anthropic API key for summarization
 
 ### Running the Server
 
 ```bash
 # Development mode
-poetry run uvicorn doc_serve_server.api.main:app --reload
+poetry run uvicorn agent_brain_server.api.main:app --reload
 
 # Or use the entry point
-poetry run doc-serve
+poetry run agent-brain-serve
 ```
 
-The server will start at `http://127.0.0.1:8000`.
-
 ### API Documentation
 
 Once running, visit:
@@ -157,7 +182,7 @@ curl -X POST http://localhost:8000/query \
 ## Architecture
 
 ```
-doc_serve_server/
+agent_brain_server/
 ├── api/
 │   ├── main.py           # FastAPI application
 │   └── routers/          # Endpoint handlers
@@ -186,16 +211,32 @@ poetry run pytest
 ### Code Formatting
 
 ```bash
-poetry run black doc_serve_server/
-poetry run ruff check doc_serve_server/
+poetry run black agent_brain_server/
+poetry run ruff check agent_brain_server/
 ```
 
 ### Type Checking
 
 ```bash
-poetry run mypy doc_serve_server/
+poetry run mypy agent_brain_server/
 ```
 
+## Documentation
+
+- [User Guide](https://github.com/SpillwaveSolutions/agent-brain/wiki/User-Guide) - Getting started and usage
+- [Developer Guide](https://github.com/SpillwaveSolutions/agent-brain/wiki/Developer-Guide) - Contributing and development
+- [API Reference](https://github.com/SpillwaveSolutions/agent-brain/wiki/API-Reference) - Full API documentation
+
+## Release Information
+
+- **Current Version**: See [pyproject.toml](./pyproject.toml)
+- **Release Notes**: [GitHub Releases](https://github.com/SpillwaveSolutions/agent-brain/releases)
+- **Changelog**: [Latest Release](https://github.com/SpillwaveSolutions/agent-brain/releases/latest)
+
+## Related Packages
+
+- [agent-brain-cli](https://pypi.org/project/agent-brain-cli/) - Command-line interface for Agent Brain
+
 ## License
 
 MIT

agent_brain_rag-2.0.0/README.md

@@ -0,0 +1,194 @@
+# Agent Brain RAG Server
+
+> **Agent Brain** (formerly doc-serve) is an intelligent document indexing and semantic search system designed to give AI agents long-term memory.
+
+AI agents need persistent memory to be truly useful. Agent Brain provides the retrieval infrastructure that enables context-aware, knowledge-grounded AI interactions.
+
+[![PyPI version](https://badge.fury.io/py/agent-brain-rag.svg)](https://pypi.org/project/agent-brain-rag/)
+[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+
+## Installation
+
+```bash
+pip install agent-brain-rag
+```
+
+## Quick Start
+
+1. Set environment variables:
+   ```bash
+   export OPENAI_API_KEY=your-key
+   export ANTHROPIC_API_KEY=your-key
+   ```
+
+2. Start the server:
+   ```bash
+   agent-brain-serve
+   ```
+
+The server will start at `http://127.0.0.1:8000`.
+
+> **Note**: The legacy command `doc-serve` is still available but deprecated. Please use `agent-brain-serve` for new installations.
+
+## Search Capabilities
+
+Agent Brain provides multiple search strategies to match your retrieval needs:
+
+| Search Type | Description | Best For |
+|-------------|-------------|----------|
+| **Semantic Search** | Natural language queries using OpenAI embeddings (`text-embedding-3-large`) | Conceptual questions, finding related content |
+| **Keyword Search (BM25)** | Traditional keyword matching with TF-IDF ranking | Exact matches, technical terms, code identifiers |
+| **Hybrid Search** | Combines vector + BM25 for best of both approaches | General-purpose queries, balanced recall/precision |
+| **GraphRAG** | Knowledge graph-based retrieval for relationship-aware queries | Understanding connections, multi-hop reasoning |
+
+## Features
+
+- **Document Indexing**: Load and index documents from folders (PDF, Markdown, TXT, DOCX, HTML)
+- **AST-Aware Code Ingestion**: Smart parsing for Python, TypeScript, JavaScript, Java, Go, Rust, C, C++
+- **Multi-Strategy Retrieval**: Semantic, keyword, hybrid, and graph-based search
+- **OpenAI Embeddings**: Uses `text-embedding-3-large` for high-quality embeddings
+- **Claude Summarization**: AI-powered code summaries for better context
+- **Chroma Vector Store**: Persistent, thread-safe vector database
+- **FastAPI**: Modern, high-performance REST API with OpenAPI documentation
+
+## Prerequisites
+
+- Python 3.10+
+- OpenAI API key (for embeddings)
+- Anthropic API key (for summarization)
+
+## Development Installation
+
+```bash
+cd agent-brain-server
+poetry install
+```
+
+### Configuration
+
+Copy the environment template and configure:
+
+```bash
+cp ../.env.example .env
+# Edit .env with your API keys
+```
+
+Required environment variables:
+- `OPENAI_API_KEY`: Your OpenAI API key for embeddings
+- `ANTHROPIC_API_KEY`: Your Anthropic API key for summarization
+
+### Running the Server
+
+```bash
+# Development mode
+poetry run uvicorn agent_brain_server.api.main:app --reload
+
+# Or use the entry point
+poetry run agent-brain-serve
+```
+
+### API Documentation
+
+Once running, visit:
+- Swagger UI: http://127.0.0.1:8000/docs
+- ReDoc: http://127.0.0.1:8000/redoc
+- OpenAPI JSON: http://127.0.0.1:8000/openapi.json
+
+## API Endpoints
+
+### Health
+
+- `GET /health` - Server health status
+- `GET /health/status` - Detailed indexing status
+
+### Indexing
+
+- `POST /index` - Start indexing documents from a folder
+- `POST /index/add` - Add documents to existing index
+- `DELETE /index` - Reset the index
+
+### Querying
+
+- `POST /query` - Semantic search query
+- `GET /query/count` - Get indexed document count
+
+## Example Usage
+
+### Index Documents
+
+```bash
+curl -X POST http://localhost:8000/index \
+  -H "Content-Type: application/json" \
+  -d '{"folder_path": "/path/to/docs"}'
+```
+
+### Query Documents
+
+```bash
+curl -X POST http://localhost:8000/query \
+  -H "Content-Type: application/json" \
+  -d '{"query": "How do I configure authentication?", "top_k": 5}'
+```
+
+## Architecture
+
+```
+agent_brain_server/
+├── api/
+│   ├── main.py           # FastAPI application
+│   └── routers/          # Endpoint handlers
+├── config/
+│   └── settings.py       # Configuration management
+├── models/               # Pydantic request/response models
+├── indexing/
+│   ├── document_loader.py  # Document loading
+│   ├── chunking.py         # Text chunking
+│   └── embedding.py        # Embedding generation
+├── services/
+│   ├── indexing_service.py # Indexing orchestration
+│   └── query_service.py    # Query execution
+└── storage/
+    └── vector_store.py     # Chroma vector store
+```
+
+## Development
+
+### Running Tests
+
+```bash
+poetry run pytest
+```
+
+### Code Formatting
+
+```bash
+poetry run black agent_brain_server/
+poetry run ruff check agent_brain_server/
+```
+
+### Type Checking
+
+```bash
+poetry run mypy agent_brain_server/
+```
+
+## Documentation
+
+- [User Guide](https://github.com/SpillwaveSolutions/agent-brain/wiki/User-Guide) - Getting started and usage
+- [Developer Guide](https://github.com/SpillwaveSolutions/agent-brain/wiki/Developer-Guide) - Contributing and development
+- [API Reference](https://github.com/SpillwaveSolutions/agent-brain/wiki/API-Reference) - Full API documentation
+
+## Release Information
+
+- **Current Version**: See [pyproject.toml](./pyproject.toml)
+- **Release Notes**: [GitHub Releases](https://github.com/SpillwaveSolutions/agent-brain/releases)
+- **Changelog**: [Latest Release](https://github.com/SpillwaveSolutions/agent-brain/releases/latest)
+
+## Related Packages
+
+- [agent-brain-cli](https://pypi.org/project/agent-brain-cli/) - Command-line interface for Agent Brain
+
+## License
+
+MIT

{agent_brain_rag-1.1.0/doc_serve_server → agent_brain_rag-2.0.0/agent_brain_server}/__init__.py

@@ -1,3 +1,3 @@
 """Doc-Serve Server - RAG-based document indexing and query service."""
 
-__version__ = "1.1.0"
+__version__ = "2.0.0"

{agent_brain_rag-1.1.0/doc_serve_server → agent_brain_rag-2.0.0/agent_brain_server}/api/main.py

@@ -1,8 +1,15 @@
-"""FastAPI application entry point."""
+"""FastAPI application entry point.
+
+This module provides the Agent Brain RAG server, a FastAPI application
+for document indexing and semantic search. The primary entry point is
+`agent-brain-serve`, with `doc-serve` provided for backward compatibility.
+"""
 
 import logging
 import os
 import socket
+import sys
+import warnings
 from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from pathlib import Path
@@ -13,15 +20,24 @@ import uvicorn
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 
-from doc_serve_server import __version__
-from doc_serve_server.config import settings
-from doc_serve_server.indexing.bm25_index import BM25IndexManager
-from doc_serve_server.locking import acquire_lock, cleanup_stale, is_stale, release_lock
-from doc_serve_server.project_root import resolve_project_root
-from doc_serve_server.runtime import RuntimeState, delete_runtime, write_runtime
-from doc_serve_server.services import IndexingService, QueryService
-from doc_serve_server.storage import VectorStoreManager
-from doc_serve_server.storage_paths import resolve_state_dir, resolve_storage_paths
+from agent_brain_server import __version__
+from agent_brain_server.config import settings
+from agent_brain_server.config.provider_config import (
+    load_provider_settings,
+    validate_provider_config,
+)
+from agent_brain_server.indexing.bm25_index import BM25IndexManager
+from agent_brain_server.locking import (
+    acquire_lock,
+    cleanup_stale,
+    is_stale,
+    release_lock,
+)
+from agent_brain_server.project_root import resolve_project_root
+from agent_brain_server.runtime import RuntimeState, delete_runtime, write_runtime
+from agent_brain_server.services import IndexingService, QueryService
+from agent_brain_server.storage import VectorStoreManager
+from agent_brain_server.storage_paths import resolve_state_dir, resolve_storage_paths
 
 from .routers import health_router, index_router, query_router
 
@@ -52,7 +68,31 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
     """
     global _runtime_state, _state_dir
 
-    logger.info("Starting Doc-Serve server...")
+    logger.info("Starting Agent Brain RAG server...")
+
+    # Load and validate provider configuration
+    try:
+        provider_settings = load_provider_settings()
+        validation_errors = validate_provider_config(provider_settings)
+
+        if validation_errors:
+            for error in validation_errors:
+                logger.warning(f"Provider config warning: {error}")
+            # Log but don't fail - providers may work if keys are set later
+            # or if using Ollama which doesn't need keys
+
+        # Log active provider configuration
+        logger.info(
+            f"Embedding provider: {provider_settings.embedding.provider} "
+            f"(model: {provider_settings.embedding.model})"
+        )
+        logger.info(
+            f"Summarization provider: {provider_settings.summarization.provider} "
+            f"(model: {provider_settings.summarization.model})"
+        )
+    except Exception as e:
+        logger.error(f"Failed to load provider configuration: {e}")
+        # Continue with defaults - EmbeddingGenerator will handle provider creation
 
     if settings.OPENAI_API_KEY:
         os.environ["OPENAI_API_KEY"] = settings.OPENAI_API_KEY
@@ -138,7 +178,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
 
     yield
 
-    logger.info("Shutting down Doc-Serve server...")
+    logger.info("Shutting down Agent Brain RAG server...")
 
     # Cleanup for per-project mode
     if state_dir is not None:
@@ -149,12 +189,12 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]:
 
 # Create FastAPI application
 app = FastAPI(
-    title="Doc-Serve API",
+    title="Agent Brain RAG API",
     description=(
         "RAG-based document indexing and semantic search API. "
         "Index documents from folders and query them using natural language."
     ),
-    version="1.1.0",
+    version="2.0.0",
     lifespan=lifespan,
     docs_url="/docs",
     redoc_url="/redoc",
@@ -180,8 +220,8 @@ app.include_router(query_router, prefix="/query", tags=["Querying"])
 async def root() -> dict[str, str]:
     """Root endpoint redirects to docs."""
     return {
-        "name": "Doc-Serve API",
-        "version": "1.1.0",
+        "name": "Agent Brain RAG API",
+        "version": "2.0.0",
         "docs": "/docs",
         "health": "/health",
     }
@@ -246,7 +286,7 @@ def run(
         logger.info(f"Per-project mode enabled: {_state_dir}")
 
     uvicorn.run(
-        "doc_serve_server.api.main:app",
+        "agent_brain_server.api.main:app",
         host=resolved_host,
         port=resolved_port,
         reload=reload if reload is not None else settings.DEBUG,
@@ -254,7 +294,7 @@ def run(
 
 
 @click.command()
-@click.version_option(version=__version__, prog_name="doc-serve")
+@click.version_option(version=__version__, prog_name="agent-brain-serve")
 @click.option(
     "--host",
     "-h",
@@ -292,19 +332,19 @@ def cli(
     state_dir: Optional[str],
     project_dir: Optional[str],
 ) -> None:
-    """Doc-Serve Server - RAG-based document indexing and semantic search API.
+    """Agent Brain RAG Server - Document indexing and semantic search API.
 
     Start the FastAPI server for document indexing and querying.
 
     \b
     Examples:
-      doc-serve                           # Start with default settings
-      doc-serve --port 8080               # Start on port 8080
-      doc-serve --port 0                  # Auto-assign an available port
-      doc-serve --host 0.0.0.0            # Bind to all interfaces
-      doc-serve --reload                  # Enable auto-reload
-      doc-serve --project-dir /my/project # Per-project mode (auto state-dir)
-      doc-serve --state-dir /path/.claude/doc-serve  # Explicit state directory
+      agent-brain-serve                           # Start with default settings
+      agent-brain-serve --port 8080               # Start on port 8080
+      agent-brain-serve --port 0                  # Auto-assign an available port
+      agent-brain-serve --host 0.0.0.0            # Bind to all interfaces
+      agent-brain-serve --reload                  # Enable auto-reload
+      agent-brain-serve --project-dir /my/project # Per-project mode
+      agent-brain-serve --state-dir /path/.claude/doc-serve  # Explicit state dir
 
     \b
     Environment Variables:
@@ -328,5 +368,29 @@ def cli(
     run(host=host, port=port, reload=reload, state_dir=resolved_state_dir)
 
 
+def cli_deprecated() -> None:
+    """Deprecated entry point for doc-serve command.
+
+    Shows a deprecation warning and then runs the main CLI.
+    """
+    warnings.warn(
+        "\n"
+        "WARNING: 'doc-serve' is deprecated and will be removed in v2.0.\n"
+        "Please use 'agent-brain-serve' instead.\n"
+        "\n"
+        "Migration guide: docs/MIGRATION.md\n"
+        "Online: https://github.com/SpillwaveSolutions/agent-brain/blob/main/docs/MIGRATION.md\n",
+        DeprecationWarning,
+        stacklevel=1,
+    )
+    # Print to stderr for visibility since warnings may be filtered
+    print(
+        "\033[93mWARNING: 'doc-serve' is deprecated. "
+        "Use 'agent-brain-serve' instead. See docs/MIGRATION.md\033[0m",
+        file=sys.stderr,
+    )
+    cli()
+
+
 if __name__ == "__main__":
     cli()

{agent_brain_rag-1.1.0/doc_serve_server → agent_brain_rag-2.0.0/agent_brain_server}/api/routers/health.py

@@ -5,7 +5,8 @@ from typing import Literal
 
 from fastapi import APIRouter, Request
 
-from doc_serve_server.models import HealthStatus, IndexingStatus
+from agent_brain_server import __version__
+from agent_brain_server.models import HealthStatus, IndexingStatus
 
 router = APIRouter()
 
@@ -54,7 +55,7 @@ async def health_check(request: Request) -> HealthStatus:
         status=status,
         message=message,
         timestamp=datetime.now(timezone.utc),
-        version="1.1.0",
+        version=__version__,
         mode=mode,
         instance_id=instance_id,
         project_id=project_id,
@@ -97,4 +98,5 @@ async def indexing_status(request: Request) -> IndexingStatus:
         ),
         indexed_folders=status["indexed_folders"],
         supported_languages=status.get("supported_languages", []),
+        graph_index=status.get("graph_index"),
     )

{agent_brain_rag-1.1.0/doc_serve_server → agent_brain_rag-2.0.0/agent_brain_server}/api/routers/index.py

@@ -5,7 +5,7 @@ from pathlib import Path
 
 from fastapi import APIRouter, HTTPException, Request, status
 
-from doc_serve_server.models import IndexRequest, IndexResponse
+from agent_brain_server.models import IndexRequest, IndexResponse
 
 router = APIRouter()
 

{agent_brain_rag-1.1.0/doc_serve_server → agent_brain_rag-2.0.0/agent_brain_server}/api/routers/query.py

@@ -4,7 +4,7 @@ import logging
 
 from fastapi import APIRouter, HTTPException, Request, status
 
-from doc_serve_server.models import QueryRequest, QueryResponse
+from agent_brain_server.models import QueryRequest, QueryResponse
 
 logger = logging.getLogger(__name__)
 
@@ -33,8 +33,8 @@ async def query_documents(
         400: Invalid query (empty or too long)
         503: Index not ready (indexing in progress or not initialized)
     """
-    from doc_serve_server.services import QueryService
-    from doc_serve_server.services.indexing_service import IndexingService
+    from agent_brain_server.services import QueryService
+    from agent_brain_server.services.indexing_service import IndexingService
 
     query_service: QueryService = request.app.state.query_service
     indexing_service: IndexingService = request.app.state.indexing_service