diff --git a/agentic-rag-authorization/.env.example b/agentic-rag-authorization/.env.example new file mode 100644 index 0000000..2101b18 --- /dev/null +++ b/agentic-rag-authorization/.env.example @@ -0,0 +1,10 @@ +# Required +MISTRAL_API_KEY=your-mistral-api-key-here + +# Optional (defaults shown) +ELASTICSEARCH_URL=http://localhost:9200 +ELASTICSEARCH_API_KEY= +SPICEDB_ENDPOINT=localhost:50051 +SPICEDB_TOKEN=devtoken +MAX_RETRIEVAL_ATTEMPTS=1 +LOG_LEVEL=INFO diff --git a/agentic-rag-authorization/ARCHITECTURE.md b/agentic-rag-authorization/ARCHITECTURE.md index 1b1a20a..58ac9ce 100644 --- a/agentic-rag-authorization/ARCHITECTURE.md +++ b/agentic-rag-authorization/ARCHITECTURE.md @@ -35,8 +35,8 @@ Technical details for those implementing similar systems or extending this one. │ │ │ │ ┌──────────┼──────────┐ │ │ ▼ ▼ ▼ │ -│ Weaviate SpiceDB OpenAI │ -│ (Search) (AuthZ) (LLM) │ +│ Elasticsearch SpiceDB Mistral │ +│ (Search) (AuthZ) (LLM) │ └─────────────────────────────────────────────────────────┘ ``` @@ -107,9 +107,9 @@ result = run_agentic_rag( python3 run_ui.py ``` This launcher performs pre-flight checks: -- Verifies Weaviate connectivity +- Verifies Elasticsearch connectivity - Verifies SpiceDB connectivity -- Checks OpenAI API key configuration +- Checks Mistral API key configuration - Validates documents are loaded - Auto-opens browser to http://localhost:8000 @@ -224,7 +224,7 @@ async def get_users(): @router.get("/health") async def health_check(): # Returns service status - # TODO: Actually check Weaviate/SpiceDB connectivity + # TODO: Actually check Elasticsearch/SpiceDB connectivity ``` **3. Pydantic Models (`api/models.py`)** @@ -308,7 +308,7 @@ FastAPI Router (routes.py) run_agentic_rag_async() │ ▼ Execute LangGraph - │ 1. Retrieval Node → Weaviate BM25 search + │ 1. Retrieval Node → Elasticsearch semantic vector search │ 2. Authorization Node → SpiceDB permission checks │ 3. Generation Node → LLM answer generation │ @@ -461,7 +461,7 @@ The API layer is designed for **demonstration and education**, not production us 4. **No Rate Limiting** - No protection against abuse or DoS - - OpenAI API costs could accumulate + - Mistral API costs could accumulate **Why This Is Acceptable for Demo:** - System demonstrates authorization concepts (SpiceDB) @@ -581,7 +581,7 @@ async def run_agentic_rag_async(query: str, subject_id: str, max_attempts: int) 1. **Concurrency** - Server can handle multiple queries simultaneously - - Other requests aren't blocked while one query waits for OpenAI + - Other requests aren't blocked while one query waits for Mistral 2. **Resource Efficiency** - Async I/O doesn't waste threads on waiting @@ -628,9 +628,10 @@ async def run_agentic_rag_async(query, subject_id, max_attempts): **Node Compatibility:** -All nodes work with both sync and async execution: -- LangChain components support async (Weaviate client, OpenAI) -- SpiceDB gRPC client is synchronous but fast (~40-50ms) +All nodes run under both the sync and async entry points: +- The retrieval and generation nodes are synchronous — the `elasticsearch` client + and the Mistral embedding/LLM calls (`chain.invoke`) are blocking +- The SpiceDB permission check (authorization node) runs on async gRPC and is fast (~40-50ms) - No code duplication required ### LangGraph State Machine @@ -639,7 +640,7 @@ All nodes work with both sync and async execution: ``` START ↓ -Retrieval Node (Weaviate BM25) +Retrieval Node (Elasticsearch semantic vector search) ↓ Authorization Node (SpiceDB) ◄── Security Boundary (deterministic) ↓ @@ -677,12 +678,12 @@ AgenticRAGState = TypedDict("AgenticRAGState", { ### Retrieval Node (Deterministic) -**Purpose**: Execute semantic/keyword search in Weaviate. +**Purpose**: Execute semantic vector search in Elasticsearch. **Input**: `query` from state **Operation**: -- Weaviate BM25 keyword search (default) +- Elasticsearch semantic vector search: kNN over `dense_vector` field using cosine similarity (default) - Returns top-k documents (typically 5) - No authorization filtering at this stage - Direct execution without planning overhead @@ -763,7 +764,7 @@ definition document { 1. User makes query subject_id: "alice" -2. Weaviate retrieves documents +2. Elasticsearch retrieves documents [eng-001, eng-002, hr-001] 3. For each document, SpiceDB checks: @@ -835,7 +836,7 @@ hr-001 (document) ──viewer──> hr_manager (user) ▼ ┌───────────────────────────────────┐ │ Trusted Zone │ -│ SpiceDB + Weaviate │ +│ SpiceDB + Elasticsearch │ │ Authorized data │ └───────────────────────────────────┘ ``` @@ -866,7 +867,7 @@ hr-001 (document) ──viewer──> hr_manager (user) ▼ ┌───────────────────────────┐ │ Trusted Zone │ -│ SpiceDB + Weaviate │ +│ SpiceDB + Elasticsearch │ │ Authorized data │ └───────────────────────────┘ ``` @@ -1002,7 +1003,7 @@ def authorize(state): ``` Query ↓ -Retrieve (BM25 search) +Retrieve (semantic vector search) ↓ Authorize (filter) ↓ @@ -1100,7 +1101,7 @@ from fastapi import APIRouter @router.get("/documents") async def list_documents(): """List all documents with metadata.""" - # Query Weaviate for all documents + # Query Elasticsearch for all documents # Return structured list return {"documents": [...]} @@ -1212,18 +1213,18 @@ def format_document_detailed(doc: Document) -> dict: ## Project Structure ``` -agentic-rag-weaviate/ +agentic-rag-elasticsearch/ ├── agentic_rag/ # Core RAG engine │ ├── graph.py # LangGraph state machine (sync + async entry points) │ ├── state.py # AgenticRAGState TypedDict definition │ ├── config.py # Configuration management (env vars) │ ├── nodes/ -│ │ ├── retrieval_node.py # Weaviate BM25 search +│ │ ├── retrieval_node.py # Elasticsearch semantic vector search │ │ ├── authorization_node.py # SpiceDB filtering (security boundary) │ │ ├── reasoning_node.py # Optional: adaptive retry logic │ │ └── generation_node.py # Final answer with context │ ├── authorization_helpers.py # Batch permission checking -│ ├── weaviate_client.py # Connection pooling for Weaviate +│ ├── elasticsearch_client.py # Connection pooling for Elasticsearch │ ├── grpc_helpers.py # Connection pooling for SpiceDB │ ├── logging_config.py # Structured JSON logging │ ├── node_helpers.py # Shared utilities for nodes @@ -1256,7 +1257,7 @@ agentic-rag-weaviate/ │ └── PERMISSIONS.md # Permission matrix and patterns │ ├── run_ui.py # Web UI launcher with pre-flight checks -├── docker-compose.yml # Weaviate + SpiceDB services +├── docker-compose.yml # Elasticsearch + SpiceDB services ├── requirements.txt # Python dependencies (includes fastapi, uvicorn) ├── .env.example # Environment variable template ├── README.md # Overview, quick start, usage guide @@ -1279,7 +1280,7 @@ agentic-rag-weaviate/ **Shared Core:** - State machine: `agentic_rag/graph.py` - Nodes: `agentic_rag/nodes/*.py` -- Services: `agentic_rag/weaviate_client.py`, `agentic_rag/grpc_helpers.py` +- Services: `agentic_rag/elasticsearch_client.py`, `agentic_rag/grpc_helpers.py` ## Performance Characteristics @@ -1322,11 +1323,11 @@ Async execution allows concurrent query handling: - **Single query**: ~3.5-4.5s - **5 concurrent queries**: ~4-5s each (minimal overhead) -- **10 concurrent queries**: ~5-6s each (slight queueing at OpenAI) +- **10 concurrent queries**: ~5-6s each (slight queueing at Mistral) **Bottlenecks:** -1. OpenAI API calls (~2-3s) - most significant -2. Weaviate search (~0.5-1s) +1. Mistral API calls (~2-3s) - most significant +2. Elasticsearch search (~0.5-1s) 3. SpiceDB checks (~40-50ms for 3-5 documents) 4. API/network overhead (~30-50ms total) @@ -1379,7 +1380,7 @@ Every node updates the state with: ``` [AIMessage] Planning: Searching for engineering documents... -[SystemMessage] Retrieved 3 documents from Weaviate +[SystemMessage] Retrieved 3 documents from Elasticsearch [SystemMessage] Authorization: 2/3 documents authorized (1 denied) [AIMessage] Reasoning: User has partial access, generating answer from available docs [AIMessage] Answer: Based on the 2 authorized documents... @@ -1525,7 +1526,7 @@ For production deployment, consider tracking: - API response time (p50, p95, p99) - LangGraph execution time - Individual node execution times - - OpenAI API latency + - Mistral API latency 2. **Authorization Metrics:** - Authorization pass rate (per user) @@ -1542,8 +1543,8 @@ For production deployment, consider tracking: 4. **Error Metrics:** - API error rate - Validation failures - - OpenAI API errors - - Service connectivity failures (Weaviate, SpiceDB) + - Mistral API errors + - Service connectivity failures (Elasticsearch, SpiceDB) ## Summary diff --git a/agentic-rag-authorization/README.md b/agentic-rag-authorization/README.md index e88d5f5..0e2156d 100644 --- a/agentic-rag-authorization/README.md +++ b/agentic-rag-authorization/README.md @@ -1,8 +1,6 @@ # Agentic RAG with Fine-Grained Authorization -> **Also available:** [Weaviate version](https://github.com/authzed/examples/tree/weaviate/agentic-rag-authorization) (BM25 keyword search) · [Mistral version](https://github.com/authzed/examples/tree/mistral/agentic-rag-authorization) (Mistral embeddings and LLM) - -This repository demonstrates how to combine agentic behavior with deterministic fine-grained authorization using LangGraph, SpiceDB, and [Milvus](https://github.com/milvus-io/milvus). You'll learn to build RAG systems where a user can only see information from the documents they have access to. +This repository demonstrates how to combine agentic behavior with deterministic fine-grained authorization using LangGraph, SpiceDB, [Elasticsearch](https://github.com/elastic/elasticsearch), and [Mistral](https://mistral.ai/). You'll learn to build RAG systems where a user can only see information from the documents they have access to. This project uses the [LangChain SpiceDB](https://pypi.org/project/langchain-spicedb/) library. @@ -19,22 +17,20 @@ The two takeaways from this demo are: 2. Never ever let an AI Agent *decide* if it needs to check for authorization. Gen AI is inherently probabilistic so you have to ensure that permission checks are deterministic and cannot be skipped. -## Documentation Navigation +The diagram below shows the end-to-end request flow: the application embeds the query with Mistral, retrieves relevant context from Elasticsearch, checks permissions in SpiceDB, and only then passes the authorized context to the Mistral LLM to generate a response. -- **[README.md](README.md)** (you are here) - Overview, quick start, core concepts -- **[ARCHITECTURE.md](ARCHITECTURE.md)** - Deep dive into system design, security model, and trade-offs -- **[data/PERMISSIONS.md](data/PERMISSIONS.md)** - Permission matrix and authorization patterns +![Agentic RAG request flow: a user query is embedded with Mistral, searched against Elasticsearch, filtered through a SpiceDB permission check, and answered by the Mistral LLM before the response returns to the user](agentic-rag-flow.png) ## What You'll Learn This repo demonstrates: -1. **Fine-grained authorization in RAG** - How to enforce document-level permissions with SpiceDB so users only see what they're allowed to see +1. **Fine-grained authorization in Agentic RAG** - How to enforce document-level permissions with SpiceDB so users only see what they're allowed to see 2. **Security architecture** - A deterministic authorization boundary that cannot be bypassed by the agent 3. **Production features** - Structured logging, connection pooling, batch operations, error handling 4. **Real-world complexity** - 50 documents, 4 permission patterns with hierarchies -Note: Despite the "agentic RAG" name, the default mode is intentionally simple and deterministic (3 nodes: retrieve → authorize → generate). This provides fast, predictable behavior suitable for most use cases. There is a `MAX_RETRIES` option where the AI Agent can reason if it has to retrieve more data. +Note: Despite the "agentic RAG" name, the default mode is intentionally simple and deterministic (3 nodes: retrieve → authorize → generate). This provides fast, predictable behavior suitable for most use cases. There is a `MAX_RETRIEVAL_ATTEMPTS` option where the AI Agent can reason if it has to retrieve more data. ## The Problem This Solves @@ -48,7 +44,7 @@ Read the [OWASP Top 10 for LLM](https://owasp.org/www-project-top-10-for-large-l ## The Solution This implementation shows how to combine: -- **Retrieval-first approach**: Semantic vector search without upfront planning overhead +- **Retrieval-first approach**: Semantic vector search (Mistral embeddings in Elasticsearch) without upfront planning overhead - **Deterministic security**: SpiceDB authorization that cannot be bypassed - **Transparency**: Users understand what they can and can't access, and why @@ -74,8 +70,7 @@ Result: ✅ Authorized: 2 documents (eng-001, eng-002) ❌ Denied: 1 document (hr-001) -Answer: "Based on the engineering documents, our system uses microservices -architecture with event-driven patterns..." +Answer: "Based on the engineering documents, our system uses microservices architecture with event-driven patterns..." ``` ```bash @@ -88,100 +83,45 @@ Result: ❌ Authorized: 0 documents ❌ Denied: 3 documents -Answer: "I don't have access to the engineering documents needed to answer -this question. This information is restricted to the engineering department." +Answer: "I don't have access to the engineering documents needed to answer this question. This information is restricted to the engineering department." ``` The agent transparently explains access limitations instead of failing silently. -## Setup (5 minutes) +## Setup & Run (5 minutes) + +The demo runs entirely through the web UI, which lets you switch users and watch the authorization boundary filter results in real time. ### Prerequisites - Docker & Docker Compose - Python 3.11+ -- OpenAI API key +- Mistral API key ### Steps ```bash -# 1. Configure +# 1. Configure — add your Mistral API key to the new .env file cp .env.example .env -# Edit .env with your actual OpenAI API key -# 2. Start services +# 2. Start Elasticsearch + SpiceDB docker-compose up -d -# 3. Install dependencies +# 3. Install dependencies (includes FastAPI + Uvicorn for the UI) python3 -m venv venv source venv/bin/activate pip install -r requirements.txt -``` - -## Web UI -A web interface is available to demonstrate the authorization capabilities interactively. - -### Quick Start - -```bash -# 1. Ensure services are running and data is initialized -docker-compose up -d +# 4. Initialize data — embed 50 documents into Elasticsearch +# and write the permission model to SpiceDB python3 examples/setup_environment.py -# 2. Install web dependencies -pip install -r requirements.txt # Includes fastapi and uvicorn - -# 3. Launch UI (includes pre-flight checks) +# 5. Launch the web UI (runs pre-flight checks, then opens your browser) python3 run_ui.py ``` -The `setup_environment.py` script sets up Milvus as the vector database and SpiceDB with sample documents and department-based access control. It embeds all 50 documents using OpenAI's `text-embedding-3-small` and inserts them into Milvus, then writes a hierarchical permission model to SpiceDB: users assigned to departments, department-wide document access, 3 cross-department collaboration grants, and 3 individual user exceptions. - -The UI launcher will: -- Verify documents are loaded in Milvus -- Start the FastAPI server -- Open your browser to http://localhost:8000 - -Here are a few sample prompts to try: - -Choose "Bob" from "Sales" as the user and run the query "What are the company handbook guidelines?" - -You should see: -``` -📊 Retrieved: 5 -✅ Authorized: 3 -❌ Denied: 2 -``` - -Now run the same query as "HR Manager": -``` -📊 Retrieved: 5 -✅ Authorized: 5 -❌ Denied: 0 -``` - -### Manual Start - -```bash -# Terminal 1: Start services (if not running) -docker-compose up -d - -# Terminal 2: Start API server -uvicorn api.main:app --reload --host 0.0.0.0 --port 8000 - -# Browser -open http://localhost:8000 -``` +`setup_environment.py` sets up Elasticsearch as the vector database and SpiceDB with sample documents and department-based access control. It embeds all 50 documents using Mistral's `mistral-embed` and indexes them into Elasticsearch, then writes a hierarchical permission model to SpiceDB: users assigned to departments, department-wide document access, 3 cross-department collaboration grants, and 3 individual user exceptions. -## Run Without UI - -```bash -# Initialize data -python3 examples/setup_environment.py - -# Run demo -python3 examples/basic_example.py -``` +`run_ui.py` verifies Elasticsearch and SpiceDB connectivity, confirms the documents are loaded, starts the FastAPI server, and opens http://localhost:8000. ## How It Works @@ -211,50 +151,44 @@ definition document { ### 2. State Flow **Default Mode (`max_attempts=1`)** + +```mermaid +flowchart TD + Q([User Query]) --> R[Retrieval Node
Elasticsearch semantic search · mistral-embed] + R --> A[Authorization Node
SpiceDB permission filter] + A --> G([Generation Node
Answer from authorized context]) + class A boundary + classDef boundary fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937; ``` -User Query - ↓ -Retrieval Node ← Milvus semantic vector search (text-embedding-3-small) - ↓ -Authorization Node ← SpiceDB filters (SECURITY BOUNDARY - cannot be bypassed) - ↓ -Generation Node ← Answer with authorized context + explanations -``` + +The **Authorization Node** (highlighted) is the security boundary: it always runs, it's deterministic, and the agent cannot bypass it. **Adaptive Mode (`max_attempts > 1`)** When `max_attempts` is set above 1, a reasoning node activates if authorization fails. The LLM analyzes why access was denied and decides whether a different retrieval strategy might find documents the user *can* access: -``` -User Query - ↓ -Retrieval Node - ↓ -Authorization Node ← still deterministic, still non-bypassable - ↓ - some docs authorized? → Yes → Generation Node - ↓ No -Reasoning Node ← LLM decides: retry with different query, or give up? - ↓ - attempts left? → Yes → Retrieval Node (loop) - ↓ No -Generation Node ← explains the denial +```mermaid +flowchart TD + Q([User Query]) --> R[Retrieval Node] + R --> A[Authorization Node
deterministic · non-bypassable] + A --> D{Some docs
authorized?} + D -->|Yes| G([Generation Node]) + D -->|No| RE[Reasoning Node
LLM: retry with a different query, or give up?] + RE --> AL{Attempts
left?} + AL -->|Yes| R + AL -->|No| G2([Generation Node
explains the denial]) + class A boundary + classDef boundary fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937; ``` For example, if Bob (sales) asks about "microservices architecture" and the first retrieval returns only engineering-restricted docs, the reasoning node might try a broader query that surfaces a shared architecture doc Bob can actually access. -Enable it by setting `MAX_RETRIEVAL_ATTEMPTS` in `.env` (or passing `max_attempts` directly): +Enable it by setting `MAX_RETRIEVAL_ATTEMPTS` in `.env`: ```bash MAX_RETRIEVAL_ATTEMPTS=3 # default is 1 ``` -Or in code: - -```python -result = run_agentic_rag(query="...", subject_id="bob", max_attempts=3) -``` - ### 3. Security Guarantees - **Authorization always runs**: Hardcoded in the LangGraph workflow — the agent cannot skip it @@ -262,51 +196,21 @@ result = run_agentic_rag(query="...", subject_id="bob", max_attempts=3) - **Fail closed**: Access denied unless explicitly granted - **Observable**: Full audit trail in state -## Project Structure - -``` -agentic-rag-authorization/ -├── agentic_rag/ -│ ├── graph.py # LangGraph state machine -│ ├── state.py # State schema -│ ├── config.py # Configuration management -│ ├── nodes/ -│ │ ├── retrieval_node.py # Milvus semantic vector search -│ │ ├── authorization_node.py # SpiceDB filtering (security boundary) -│ │ ├── reasoning_node.py # Optional: adaptive retry logic -│ │ └── generation_node.py # Final answer with context -│ ├── milvus_client.py # Connection pooling for Milvus -│ ├── grpc_helpers.py # Connection pooling for SpiceDB -│ ├── logging_config.py # Structured JSON logging -│ └── validation.py # Input validation and sanitization -├── examples/ -│ ├── setup_environment.py # Initialize data (embeds and loads 50 documents) -│ └── basic_example.py # 8 demo scenarios -├── scripts/ -│ ├── generate_documents.py # Generate 50 .txt files -│ ├── parse_documents.py # Parse documents into objects -│ └── verify_permissions.py # Test authorization patterns -├── data/ -│ ├── documents/ # 50 .txt files (5 departments) -│ ├── schema.zed # SpiceDB permission schema -│ └── PERMISSIONS.md # Permission matrix -└── docker-compose.yml # Milvus + SpiceDB -``` - ## Configuration Environment variables (`.env`): ```bash # Required -OPENAI_API_KEY=sk-... +MISTRAL_API_KEY=... # Optional (defaults shown) -MILVUS_URI=http://localhost:19530 -MILVUS_TOKEN= +ELASTICSEARCH_URL=http://localhost:9200 +ELASTICSEARCH_API_KEY= SPICEDB_ENDPOINT=localhost:50051 SPICEDB_TOKEN=devtoken MAX_RETRIEVAL_ATTEMPTS=1 +LOG_LEVEL=INFO ``` ## Dataset Overview @@ -321,33 +225,11 @@ The repository includes a realistic 50-document dataset across 5 departments. See [data/PERMISSIONS.md](data/PERMISSIONS.md) for the complete permission matrix. -## Sample Scenarios - -The `examples/basic_example.py` demonstrates 8 scenarios: - -1. **Department Access** - alice queries engineering documents -2. **Access Denial** - bob attempts to access engineering documents -3. **Cross-Department** - bob accesses shared architecture document -4. **Individual Exception** - alice accesses sales proposal (special grant) -5. **Public Access** - Anyone can access company handbooks -6. **Finance Department** - finance_manager queries financial reports -7. **HR Department** - hr_manager queries HR policies -8. **Transparent Explanations** - Agent explains why access was denied - -## Testing - -```bash -# Run all tests -pytest tests/ - -# Run specific test -pytest tests/test_basic_flow.py::test_authorized_access -``` - ## Learn More - **SpiceDB**: https://authzed.com/docs -- **Milvus**: https://milvus.io/docs +- **Elasticsearch**: https://www.elastic.co/docs +- **Mistral**: https://docs.mistral.ai/ - **LangGraph**: https://langchain-ai.github.io/langgraph/ - **langchain-spicedb**: https://github.com/authzed/langchain-spicedb diff --git a/agentic-rag-authorization/agentic-rag-flow.png b/agentic-rag-authorization/agentic-rag-flow.png new file mode 100644 index 0000000..0252786 Binary files /dev/null and b/agentic-rag-authorization/agentic-rag-flow.png differ diff --git a/agentic-rag-authorization/agentic_rag/__init__.py b/agentic-rag-authorization/agentic_rag/__init__.py index f12f9b8..6264a64 100644 --- a/agentic-rag-authorization/agentic_rag/__init__.py +++ b/agentic-rag-authorization/agentic_rag/__init__.py @@ -1,4 +1,4 @@ -"""Agentic RAG with fine-grained authorization using Milvus and SpiceDB.""" +"""Agentic RAG with fine-grained authorization using Elasticsearch and SpiceDB.""" __version__ = "0.1.0" diff --git a/agentic-rag-authorization/agentic_rag/config.py b/agentic-rag-authorization/agentic_rag/config.py index 749c2f7..e2d0950 100644 --- a/agentic-rag-authorization/agentic_rag/config.py +++ b/agentic-rag-authorization/agentic_rag/config.py @@ -12,16 +12,16 @@ class Config: """Configuration for agentic RAG system.""" - # Milvus - milvus_uri: str - milvus_token: str + # Elasticsearch + elasticsearch_url: str + elasticsearch_api_key: str # SpiceDB spicedb_endpoint: str spicedb_token: str - # OpenAI - openai_api_key: str + # Mistral + mistral_api_key: str # Agent behavior max_retrieval_attempts: int = 1 @@ -33,11 +33,11 @@ class Config: def from_env(cls): """Load configuration from environment variables.""" return cls( - milvus_uri=os.getenv("MILVUS_URI", "http://localhost:19530"), - milvus_token=os.getenv("MILVUS_TOKEN", ""), + elasticsearch_url=os.getenv("ELASTICSEARCH_URL", "http://localhost:9200"), + elasticsearch_api_key=os.getenv("ELASTICSEARCH_API_KEY", ""), spicedb_endpoint=os.getenv("SPICEDB_ENDPOINT", "localhost:50051"), spicedb_token=os.getenv("SPICEDB_TOKEN", "devtoken"), - openai_api_key=os.getenv("OPENAI_API_KEY", ""), + mistral_api_key=os.getenv("MISTRAL_API_KEY", ""), max_retrieval_attempts=int(os.getenv("MAX_RETRIEVAL_ATTEMPTS", "1")), log_level=os.getenv("LOG_LEVEL", "INFO"), ) diff --git a/agentic-rag-authorization/agentic_rag/elasticsearch_client.py b/agentic-rag-authorization/agentic_rag/elasticsearch_client.py new file mode 100644 index 0000000..152892b --- /dev/null +++ b/agentic-rag-authorization/agentic_rag/elasticsearch_client.py @@ -0,0 +1,29 @@ +"""Elasticsearch client connection pooling.""" + +from elasticsearch import Elasticsearch +from threading import Lock +from typing import Optional + +_es_client: Optional[Elasticsearch] = None +_es_lock = Lock() + + +def get_elasticsearch_client(url: str, api_key: str = "") -> Elasticsearch: + """Get or create reusable Elasticsearch client (singleton, thread-safe).""" + global _es_client + if _es_client is not None: + return _es_client + with _es_lock: + if _es_client is None: + if api_key: + _es_client = Elasticsearch(url, api_key=api_key) + else: + _es_client = Elasticsearch(url) + return _es_client + + +def reset_elasticsearch_client(): + """Reset singleton (useful for testing).""" + global _es_client + with _es_lock: + _es_client = None diff --git a/agentic-rag-authorization/agentic_rag/graph.py b/agentic-rag-authorization/agentic_rag/graph.py index 01f0ed1..51b71c5 100644 --- a/agentic-rag-authorization/agentic_rag/graph.py +++ b/agentic-rag-authorization/agentic_rag/graph.py @@ -55,7 +55,7 @@ def build_agentic_rag_graph(): """Build the agentic RAG graph with deterministic authorization. Simplified Flow: - 1. Retrieval: Fetch documents from Milvus + 1. Retrieval: Fetch documents from Elasticsearch 2. Authorization: Deterministic permission check (security boundary) 3. Conditional: - If authorized docs exist: Generate answer diff --git a/agentic-rag-authorization/agentic_rag/logging_config.py b/agentic-rag-authorization/agentic_rag/logging_config.py index c901a3a..d2dff77 100644 --- a/agentic-rag-authorization/agentic_rag/logging_config.py +++ b/agentic-rag-authorization/agentic_rag/logging_config.py @@ -103,8 +103,8 @@ def setup_logging(level: str = "INFO") -> None: # Reduce noise from third-party libraries logging.getLogger("httpx").setLevel(logging.WARNING) logging.getLogger("httpcore").setLevel(logging.WARNING) - logging.getLogger("openai").setLevel(logging.WARNING) - logging.getLogger("pymilvus").setLevel(logging.WARNING) + logging.getLogger("langchain_mistralai").setLevel(logging.WARNING) + logging.getLogger("elasticsearch").setLevel(logging.WARNING) logging.getLogger("grpc").setLevel(logging.WARNING) diff --git a/agentic-rag-authorization/agentic_rag/milvus_client.py b/agentic-rag-authorization/agentic_rag/milvus_client.py deleted file mode 100644 index 44f14f6..0000000 --- a/agentic-rag-authorization/agentic_rag/milvus_client.py +++ /dev/null @@ -1,26 +0,0 @@ -"""Milvus client connection pooling.""" - -from pymilvus import MilvusClient -from threading import Lock -from typing import Optional - -_milvus_client: Optional[MilvusClient] = None -_milvus_lock = Lock() - - -def get_milvus_client(uri: str, token: str = "") -> MilvusClient: - """Get or create reusable MilvusClient (singleton, thread-safe).""" - global _milvus_client - if _milvus_client is not None: - return _milvus_client - with _milvus_lock: - if _milvus_client is None: - _milvus_client = MilvusClient(uri=uri, token=token) - return _milvus_client - - -def reset_milvus_client(): - """Reset singleton (useful for testing).""" - global _milvus_client - with _milvus_lock: - _milvus_client = None diff --git a/agentic-rag-authorization/agentic_rag/node_helpers.py b/agentic-rag-authorization/agentic_rag/node_helpers.py index 3e8b33f..27facf6 100644 --- a/agentic-rag-authorization/agentic_rag/node_helpers.py +++ b/agentic-rag-authorization/agentic_rag/node_helpers.py @@ -3,22 +3,22 @@ import time from contextlib import contextmanager from typing import Dict, Any -from langchain_openai import ChatOpenAI +from langchain_mistralai import ChatMistralAI from .config import get_config -def get_llm() -> ChatOpenAI: +def get_llm() -> ChatMistralAI: """Get configured LLM instance. Returns: - ChatOpenAI: Configured LLM with gpt-4 and temperature=0 + ChatMistralAI: Configured LLM with mistral-small-latest and temperature=0 """ config = get_config() - return ChatOpenAI( - model="gpt-4", + return ChatMistralAI( + model="mistral-small-latest", temperature=0, - api_key=config.openai_api_key + api_key=config.mistral_api_key ) diff --git a/agentic-rag-authorization/agentic_rag/nodes/retrieval_node.py b/agentic-rag-authorization/agentic_rag/nodes/retrieval_node.py index 83e3ba4..e92fcce 100644 --- a/agentic-rag-authorization/agentic_rag/nodes/retrieval_node.py +++ b/agentic-rag-authorization/agentic_rag/nodes/retrieval_node.py @@ -1,26 +1,25 @@ -"""Retrieval node - retrieve documents from Milvus using semantic vector search.""" +"""Retrieval node - retrieve documents from Elasticsearch using semantic vector search.""" -import openai +from langchain_mistralai import MistralAIEmbeddings from langchain_core.messages import SystemMessage from langchain_core.documents import Document from ..state import AgenticRAGState from ..config import get_config from ..logging_config import get_logger -from ..milvus_client import get_milvus_client +from ..elasticsearch_client import get_elasticsearch_client from ..node_helpers import log_node_execution logger = get_logger("nodes.retrieval") def _embed(text: str, api_key: str) -> list[float]: - client = openai.OpenAI(api_key=api_key) - response = client.embeddings.create(model="text-embedding-3-small", input=text) - return response.data[0].embedding + embeddings = MistralAIEmbeddings(model="mistral-embed", api_key=api_key) + return embeddings.embed_query(text) def retrieval_node(state: AgenticRAGState) -> dict: - """Retrieve documents from Milvus based on semantic similarity to the query.""" + """Retrieve documents from Elasticsearch based on semantic similarity to the query.""" config = get_config() with log_node_execution( @@ -29,28 +28,33 @@ def retrieval_node(state: AgenticRAGState) -> dict: {"query": state["query"], "subject_id": state["subject_id"]}, ): try: - milvus_client = get_milvus_client(config.milvus_uri, config.milvus_token) - query_embedding = _embed(state["query"], config.openai_api_key) + es_client = get_elasticsearch_client( + config.elasticsearch_url, config.elasticsearch_api_key + ) + query_embedding = _embed(state["query"], config.mistral_api_key) - results = milvus_client.search( - collection_name="Documents", - data=[query_embedding], - anns_field="embedding", - limit=5, - output_fields=["doc_id", "title", "content", "department", "classification"], + response = es_client.search( + index="documents", + knn={ + "field": "embedding", + "query_vector": query_embedding, + "k": 5, + "num_candidates": 50, + }, + source=["doc_id", "title", "content", "department", "classification"], ) documents = [ Document( - page_content=hit["entity"]["content"], + page_content=hit["_source"]["content"], metadata={ - "doc_id": hit["entity"]["doc_id"], - "title": hit["entity"]["title"], - "department": hit["entity"]["department"], - "classification": hit["entity"]["classification"], + "doc_id": hit["_source"]["doc_id"], + "title": hit["_source"]["title"], + "department": hit["_source"]["department"], + "classification": hit["_source"]["classification"], }, ) - for hit in results[0] + for hit in response["hits"]["hits"] ] logger.info( @@ -65,7 +69,7 @@ def retrieval_node(state: AgenticRAGState) -> dict: "retrieved_documents": documents, "retrieval_attempt": state["retrieval_attempt"] + 1, "messages": [ - SystemMessage(content=f"Retrieved {len(documents)} documents from Milvus") + SystemMessage(content=f"Retrieved {len(documents)} documents from Elasticsearch") ], } diff --git a/agentic-rag-authorization/api/routes.py b/agentic-rag-authorization/api/routes.py index e18ab32..125399f 100644 --- a/agentic-rag-authorization/api/routes.py +++ b/agentic-rag-authorization/api/routes.py @@ -37,14 +37,14 @@ async def get_users(): @router.get("/health") async def health_check(): """Check health of backend services.""" - # TODO: Actually check Milvus + SpiceDB connectivity + # TODO: Actually check Elasticsearch + SpiceDB connectivity # For now, return optimistic health status return { "status": "healthy", "services": { - "milvus": "connected", + "elasticsearch": "connected", "spicedb": "connected", - "openai": "configured", + "mistral": "configured", }, } diff --git a/agentic-rag-authorization/docker-compose.yml b/agentic-rag-authorization/docker-compose.yml index a23002f..f175890 100644 --- a/agentic-rag-authorization/docker-compose.yml +++ b/agentic-rag-authorization/docker-compose.yml @@ -1,61 +1,25 @@ --- -version: "3.8" - services: - etcd: - container_name: "milvus-etcd" - image: "quay.io/coreos/etcd:v3.5.5" - environment: - - "ETCD_AUTO_COMPACTION_MODE=revision" - - "ETCD_AUTO_COMPACTION_RETENTION=1000" - - "ETCD_QUOTA_BACKEND_BYTES=4294967296" - - "ETCD_SNAPSHOT_COUNT=50000" - volumes: - - "etcd_data:/etcd" - command: "etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd" - healthcheck: - test: ["CMD", "etcdctl", "endpoint", "health"] - interval: "30s" - timeout: "20s" - retries: 3 - - minio: - container_name: "milvus-minio" - image: "minio/minio:RELEASE.2023-03-13T19-46-17Z" + elasticsearch: + container_name: "elasticsearch" + image: "docker.elastic.co/elasticsearch/elasticsearch:8.13.4" environment: - MINIO_ACCESS_KEY: "minioadmin" - MINIO_SECRET_KEY: "minioadmin" - volumes: - - "minio_data:/minio_data" - command: "minio server /minio_data --console-address :9001" - healthcheck: - test: ["CMD", "curl", "-f", "http://localhost:9000/minio/health/live"] - interval: "30s" - timeout: "20s" - retries: 3 - - standalone: - container_name: "milvus-standalone" - image: "milvusdb/milvus:v2.4.6" - command: ["milvus", "run", "standalone"] - environment: - ETCD_ENDPOINTS: "etcd:2379" - MINIO_ADDRESS: "minio:9000" - volumes: - - "milvus_data:/var/lib/milvus" + - "discovery.type=single-node" + - "xpack.security.enabled=false" + - "xpack.security.http.ssl.enabled=false" + - "ES_JAVA_OPTS=-Xms512m -Xmx512m" ports: - - "19530:19530" - - "9091:9091" - depends_on: - etcd: - condition: "service_healthy" - minio: - condition: "service_healthy" + - "9200:9200" + volumes: + - "elasticsearch_data:/usr/share/elasticsearch/data" healthcheck: - test: ["CMD", "curl", "-f", "http://localhost:9091/healthz"] - interval: "30s" - timeout: "20s" - retries: 3 + test: + - "CMD-SHELL" + - "curl -s -f 'http://localhost:9200/_cluster/health?wait_for_status=yellow&timeout=1s' || exit 1" + interval: "15s" + timeout: "10s" + retries: 10 + start_period: "40s" spicedb: image: "authzed/spicedb:latest" @@ -66,6 +30,4 @@ services: SPICEDB_DATASTORE_ENGINE: "memory" volumes: - etcd_data: - minio_data: - milvus_data: + elasticsearch_data: diff --git a/agentic-rag-authorization/examples/basic_example.py b/agentic-rag-authorization/examples/basic_example.py index 7bca4eb..ca0cdc6 100644 --- a/agentic-rag-authorization/examples/basic_example.py +++ b/agentic-rag-authorization/examples/basic_example.py @@ -51,8 +51,8 @@ async def main(): """Run basic examples.""" config = get_config() - if not config.openai_api_key: - print("❌ Error: OPENAI_API_KEY not set") + if not config.mistral_api_key: + print("❌ Error: MISTRAL_API_KEY not set") print("Please set it in your .env file or environment") return diff --git a/agentic-rag-authorization/examples/setup_environment.py b/agentic-rag-authorization/examples/setup_environment.py index a9e798c..e08c054 100644 --- a/agentic-rag-authorization/examples/setup_environment.py +++ b/agentic-rag-authorization/examples/setup_environment.py @@ -1,4 +1,4 @@ -"""Initialize Milvus and SpiceDB with sample data.""" +"""Initialize Elasticsearch and SpiceDB with sample data.""" import sys import os @@ -8,8 +8,8 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -import openai -from pymilvus import MilvusClient, DataType +from langchain_mistralai import MistralAIEmbeddings +from elasticsearch import Elasticsearch, helpers from authzed.api.v1 import ( WriteSchemaRequest, WriteRelationshipsRequest, @@ -170,61 +170,71 @@ def setup_spicedb(): print(f" - Public access: 5 documents accessible to all users") -def setup_milvus(): - """Setup Milvus with sample documents using OpenAI semantic embeddings.""" - print("\nSetting up Milvus...") - - milvus_uri = os.getenv("MILVUS_URI", "http://localhost:19530") - openai_api_key = os.getenv("OPENAI_API_KEY", "") - - client = MilvusClient(uri=milvus_uri) - oai_client = openai.OpenAI(api_key=openai_api_key) - - if client.has_collection("Documents"): - client.drop_collection("Documents") - print(" ✅ Dropped existing Documents collection") - - schema = client.create_schema(auto_id=False, enable_dynamic_field=False) - schema.add_field("doc_id", DataType.VARCHAR, max_length=256, is_primary=True) - schema.add_field("title", DataType.VARCHAR, max_length=512) - schema.add_field("content", DataType.VARCHAR, max_length=65535) - schema.add_field("department", DataType.VARCHAR, max_length=128) - schema.add_field("classification", DataType.VARCHAR, max_length=128) - schema.add_field("embedding", DataType.FLOAT_VECTOR, dim=1536) - - index_params = client.prepare_index_params() - index_params.add_index( - field_name="embedding", - metric_type="COSINE", - index_type="IVF_FLAT", - params={"nlist": 128}, - ) - - client.create_collection("Documents", schema=schema, index_params=index_params) - print(" ✅ Documents collection created") +def setup_elasticsearch(): + """Setup Elasticsearch with sample documents using Mistral semantic embeddings.""" + print("\nSetting up Elasticsearch...") + + es_url = os.getenv("ELASTICSEARCH_URL", "http://localhost:9200") + es_api_key = os.getenv("ELASTICSEARCH_API_KEY", "") + mistral_api_key = os.getenv("MISTRAL_API_KEY", "") + + if es_api_key: + client = Elasticsearch(es_url, api_key=es_api_key) + else: + client = Elasticsearch(es_url) + embeddings = MistralAIEmbeddings(model="mistral-embed", api_key=mistral_api_key) + + if client.indices.exists(index="documents"): + client.indices.delete(index="documents") + print(" ✅ Dropped existing documents index") + + mappings = { + "properties": { + "doc_id": {"type": "keyword"}, + "title": {"type": "text"}, + "content": {"type": "text"}, + "department": {"type": "keyword"}, + "classification": {"type": "keyword"}, + "embedding": { + "type": "dense_vector", + "dims": 1024, + "index": True, + "similarity": "cosine", + }, + } + } + + client.indices.create(index="documents", mappings=mappings) + print(" ✅ documents index created") documents = load_all_documents() print(f" ✅ Loaded {len(documents)} documents from data/documents/") - rows = [] - for i, doc in enumerate(documents): - response = oai_client.embeddings.create( - model="text-embedding-3-small", - input=doc["content"], - ) - rows.append({ - "doc_id": doc["doc_id"], - "title": doc["title"], - "content": doc["content"], - "department": doc["department"], - "classification": doc["classification"], - "embedding": response.data[0].embedding, - }) - if (i + 1) % 10 == 0: - print(f" Embedded {i + 1}/{len(documents)} documents...") - - client.insert("Documents", rows) - print(f" ✅ Inserted {len(rows)} documents with embeddings") + # Embed all documents in a single batched call. embed_documents() chunks the + # inputs internally, so this issues far fewer Mistral API requests than + # embedding one document at a time (and is much less likely to hit rate limits). + doc_embeddings = embeddings.embed_documents([doc["content"] for doc in documents]) + print(f" ✅ Embedded {len(doc_embeddings)} documents") + + actions = [ + { + "_index": "documents", + "_id": doc["doc_id"], + "_source": { + "doc_id": doc["doc_id"], + "title": doc["title"], + "content": doc["content"], + "department": doc["department"], + "classification": doc["classification"], + "embedding": embedding, + }, + } + for doc, embedding in zip(documents, doc_embeddings) + ] + + helpers.bulk(client, actions) + client.indices.refresh(index="documents") + print(f" ✅ Inserted {len(actions)} documents with embeddings") dept_counts = {} for doc in documents: @@ -243,7 +253,7 @@ def main(): print("=" * 60) setup_spicedb() - setup_milvus() + setup_elasticsearch() print("\n" + "=" * 60) print("✅ Setup complete!") diff --git a/agentic-rag-authorization/pyproject.toml b/agentic-rag-authorization/pyproject.toml index e96e902..4c8c8a4 100644 --- a/agentic-rag-authorization/pyproject.toml +++ b/agentic-rag-authorization/pyproject.toml @@ -1,14 +1,14 @@ [project] -name = "agentic-rag-weaviate" +name = "agentic-rag-elasticsearch" version = "0.1.0" -description = "Agentic RAG with fine-grained authorization using Weaviate and SpiceDB" +description = "Agentic RAG with fine-grained authorization using Elasticsearch and SpiceDB" requires-python = ">=3.10" dependencies = [ "langchain>=0.1.0", - "langchain-openai>=0.1.0", + "langchain-mistralai>=0.2.0", "langgraph>=0.0.20", "langchain-spicedb>=0.1.0", - "weaviate-client>=4.4.0", + "elasticsearch>=8.0.0,<9.0.0", "authzed>=0.7.0", "python-dotenv>=1.0.0", ] diff --git a/agentic-rag-authorization/requirements.txt b/agentic-rag-authorization/requirements.txt index f5401bd..767fa56 100644 --- a/agentic-rag-authorization/requirements.txt +++ b/agentic-rag-authorization/requirements.txt @@ -1,9 +1,9 @@ # Core dependencies langchain>=0.1.0 -langchain-openai>=0.1.0 +langchain-mistralai>=0.2.0 langgraph>=0.0.20 langchain-spicedb>=0.2.0 -pymilvus>=2.4.0 +elasticsearch>=8.0.0,<9.0.0 authzed>=0.7.0 python-dotenv>=1.0.0 grpcio>=1.50.0 diff --git a/agentic-rag-authorization/run_ui.py b/agentic-rag-authorization/run_ui.py index 79cbeea..5c322a0 100755 --- a/agentic-rag-authorization/run_ui.py +++ b/agentic-rag-authorization/run_ui.py @@ -18,16 +18,19 @@ def check_services(): print(" Copy .env.example to .env and configure it") return False - # Check Milvus + # Check Elasticsearch try: from agentic_rag.config import get_config - from agentic_rag.milvus_client import get_milvus_client + from agentic_rag.elasticsearch_client import get_elasticsearch_client config = get_config() - milvus_client = get_milvus_client(config.milvus_uri, config.milvus_token) - print(" ✅ Milvus connected") + es_client = get_elasticsearch_client( + config.elasticsearch_url, config.elasticsearch_api_key + ) + es_client.info() + print(" ✅ Elasticsearch connected") except Exception as e: - print(f" ❌ Milvus not available: {e}") + print(f" ❌ Elasticsearch not available: {e}") print(" Run: docker-compose up -d") return False @@ -42,30 +45,25 @@ def check_services(): print(" Run: docker-compose up -d") return False - # Check OpenAI key - if not config.openai_api_key or config.openai_api_key == "your-openai-api-key-here": - print(" ❌ OPENAI_API_KEY not configured") + # Check Mistral key + if not config.mistral_api_key or config.mistral_api_key == "your-mistral-api-key-here": + print(" ❌ MISTRAL_API_KEY not configured") print(" Set it in .env file") return False - print(" ✅ OpenAI API key configured") + print(" ✅ Mistral API key configured") # Check if documents are loaded try: - if milvus_client.has_collection("Documents"): - results = milvus_client.query( - collection_name="Documents", - filter='doc_id != ""', - output_fields=["doc_id"], - limit=1, - ) - if results: - print(" ✅ Documents loaded in Milvus") + if es_client.indices.exists(index="documents"): + count = es_client.count(index="documents")["count"] + if count > 0: + print(" ✅ Documents loaded in Elasticsearch") else: - print(" ⚠️ No documents found in Milvus") + print(" ⚠️ No documents found in Elasticsearch") print(" Run: python examples/setup_environment.py") return False else: - print(" ⚠️ Documents collection does not exist in Milvus") + print(" ⚠️ documents index does not exist in Elasticsearch") print(" Run: python examples/setup_environment.py") return False except Exception as e: diff --git a/agentic-rag-authorization/test_improvements.py b/agentic-rag-authorization/test_improvements.py index 0849ecd..831398b 100644 --- a/agentic-rag-authorization/test_improvements.py +++ b/agentic-rag-authorization/test_improvements.py @@ -65,19 +65,19 @@ def test_connection_pooling(): reset_spicedb_client, _spicedb_client, ) - from agentic_rag.milvus_client import ( - get_milvus_client, - reset_milvus_client, - _milvus_client, + from agentic_rag.elasticsearch_client import ( + get_elasticsearch_client, + reset_elasticsearch_client, + _es_client, ) # Verify reset functions exist assert callable(reset_spicedb_client) - assert callable(reset_milvus_client) + assert callable(reset_elasticsearch_client) print("✅ Connection pooling functions defined correctly") print(" - get_spicedb_client() available") - print(" - get_milvus_client() available") + print(" - get_elasticsearch_client() available") print(" - reset_*_client() functions available") return True except (ImportError, AssertionError) as e: diff --git a/agentic-rag-authorization/tests/conftest.py b/agentic-rag-authorization/tests/conftest.py index 52109b6..6819429 100644 --- a/agentic-rag-authorization/tests/conftest.py +++ b/agentic-rag-authorization/tests/conftest.py @@ -1,14 +1,14 @@ """Test fixtures for agentic RAG tests.""" import pytest -from pymilvus import MilvusClient +from elasticsearch import Elasticsearch from agentic_rag.grpc_helpers import create_insecure_spicedb_client @pytest.fixture -def milvus_client(): - """Create MilvusClient for tests.""" - return MilvusClient(uri="http://localhost:19530") +def elasticsearch_client(): + """Create Elasticsearch client for tests.""" + return Elasticsearch("http://localhost:9200") @pytest.fixture diff --git a/agentic-rag-authorization/tests/test_config.py b/agentic-rag-authorization/tests/test_config.py index fb1ebd0..90beb81 100644 --- a/agentic-rag-authorization/tests/test_config.py +++ b/agentic-rag-authorization/tests/test_config.py @@ -5,31 +5,45 @@ from agentic_rag.config import Config -def test_config_loads_milvus_uri(): +def test_config_loads_elasticsearch_url(): env = { - "MILVUS_URI": "http://milvus-host:19530", - "OPENAI_API_KEY": "sk-test", + "ELASTICSEARCH_URL": "http://es-host:9200", + "MISTRAL_API_KEY": "mstrl-test", "SPICEDB_TOKEN": "tok", "SPICEDB_ENDPOINT": "localhost:50051", } with patch.dict(os.environ, env, clear=True): config = Config.from_env() - assert config.milvus_uri == "http://milvus-host:19530" + assert config.elasticsearch_url == "http://es-host:9200" -def test_config_milvus_defaults(): +def test_config_loads_mistral_api_key(): env = { - "OPENAI_API_KEY": "sk-test", + "MISTRAL_API_KEY": "mstrl-test", "SPICEDB_TOKEN": "tok", "SPICEDB_ENDPOINT": "localhost:50051", } with patch.dict(os.environ, env, clear=True): config = Config.from_env() - assert config.milvus_uri == "http://localhost:19530" - assert config.milvus_token == "" + assert config.mistral_api_key == "mstrl-test" -def test_config_has_no_weaviate_fields(): +def test_config_elasticsearch_defaults(): + env = { + "MISTRAL_API_KEY": "mstrl-test", + "SPICEDB_TOKEN": "tok", + "SPICEDB_ENDPOINT": "localhost:50051", + } + with patch.dict(os.environ, env, clear=True): + config = Config.from_env() + assert config.elasticsearch_url == "http://localhost:9200" + assert config.elasticsearch_api_key == "" + + +def test_config_has_no_milvus_weaviate_or_openai_fields(): config = Config.from_env() + assert not hasattr(config, "milvus_uri") + assert not hasattr(config, "milvus_token") assert not hasattr(config, "weaviate_url") assert not hasattr(config, "weaviate_api_key") + assert not hasattr(config, "openai_api_key") diff --git a/agentic-rag-authorization/tests/test_elasticsearch_client.py b/agentic-rag-authorization/tests/test_elasticsearch_client.py new file mode 100644 index 0000000..341bea6 --- /dev/null +++ b/agentic-rag-authorization/tests/test_elasticsearch_client.py @@ -0,0 +1,43 @@ +"""Unit tests for Elasticsearch client singleton.""" + +from unittest.mock import patch, MagicMock + + +def test_get_elasticsearch_client_returns_singleton(): + from agentic_rag.elasticsearch_client import ( + get_elasticsearch_client, + reset_elasticsearch_client, + ) + reset_elasticsearch_client() + with patch("agentic_rag.elasticsearch_client.Elasticsearch") as mock_cls: + mock_cls.return_value = MagicMock() + client1 = get_elasticsearch_client("http://localhost:9200") + client2 = get_elasticsearch_client("http://localhost:9200") + assert client1 is client2 + mock_cls.assert_called_once_with("http://localhost:9200") + + +def test_reset_clears_singleton(): + from agentic_rag.elasticsearch_client import ( + get_elasticsearch_client, + reset_elasticsearch_client, + ) + reset_elasticsearch_client() + with patch("agentic_rag.elasticsearch_client.Elasticsearch") as mock_cls: + mock_cls.return_value = MagicMock() + get_elasticsearch_client("http://localhost:9200") + reset_elasticsearch_client() + get_elasticsearch_client("http://localhost:9200") + assert mock_cls.call_count == 2 + + +def test_api_key_passed_when_set(): + from agentic_rag.elasticsearch_client import ( + get_elasticsearch_client, + reset_elasticsearch_client, + ) + reset_elasticsearch_client() + with patch("agentic_rag.elasticsearch_client.Elasticsearch") as mock_cls: + mock_cls.return_value = MagicMock() + get_elasticsearch_client("http://localhost:9200", api_key="secret") + mock_cls.assert_called_once_with("http://localhost:9200", api_key="secret") diff --git a/agentic-rag-authorization/tests/test_milvus_client.py b/agentic-rag-authorization/tests/test_milvus_client.py deleted file mode 100644 index f489833..0000000 --- a/agentic-rag-authorization/tests/test_milvus_client.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Unit tests for Milvus client singleton.""" - -from unittest.mock import patch, MagicMock - - -def test_get_milvus_client_returns_singleton(): - from agentic_rag.milvus_client import get_milvus_client, reset_milvus_client - reset_milvus_client() - with patch("agentic_rag.milvus_client.MilvusClient") as mock_cls: - mock_cls.return_value = MagicMock() - client1 = get_milvus_client("http://localhost:19530") - client2 = get_milvus_client("http://localhost:19530") - assert client1 is client2 - mock_cls.assert_called_once_with(uri="http://localhost:19530", token="") - - -def test_reset_clears_singleton(): - from agentic_rag.milvus_client import get_milvus_client, reset_milvus_client - reset_milvus_client() - with patch("agentic_rag.milvus_client.MilvusClient") as mock_cls: - mock_cls.return_value = MagicMock() - get_milvus_client("http://localhost:19530") - reset_milvus_client() - get_milvus_client("http://localhost:19530") - assert mock_cls.call_count == 2 diff --git a/agentic-rag-authorization/tests/test_retrieval_node.py b/agentic-rag-authorization/tests/test_retrieval_node.py index 5276e58..38fe6c5 100644 --- a/agentic-rag-authorization/tests/test_retrieval_node.py +++ b/agentic-rag-authorization/tests/test_retrieval_node.py @@ -24,7 +24,7 @@ def sample_state(): def _make_hit(doc_id="doc-001", title="Test Doc", content="Content", dept="engineering", classification="internal"): return { - "entity": { + "_source": { "doc_id": doc_id, "title": title, "content": content, @@ -37,11 +37,11 @@ def _make_hit(doc_id="doc-001", title="Test Doc", content="Content", dept="engin def test_retrieval_node_returns_documents(sample_state): from agentic_rag.nodes.retrieval_node import retrieval_node - with patch("agentic_rag.nodes.retrieval_node.get_milvus_client") as mock_get_client, \ + with patch("agentic_rag.nodes.retrieval_node.get_elasticsearch_client") as mock_get_client, \ patch("agentic_rag.nodes.retrieval_node._embed") as mock_embed: - mock_embed.return_value = [0.1] * 1536 + mock_embed.return_value = [0.1] * 1024 mock_client = MagicMock() - mock_client.search.return_value = [[_make_hit()]] + mock_client.search.return_value = {"hits": {"hits": [_make_hit()]}} mock_get_client.return_value = mock_client result = retrieval_node(sample_state) @@ -60,9 +60,9 @@ def test_retrieval_node_returns_documents(sample_state): def test_retrieval_node_increments_attempt_on_failure(sample_state): from agentic_rag.nodes.retrieval_node import retrieval_node - with patch("agentic_rag.nodes.retrieval_node.get_milvus_client") as mock_get_client, \ + with patch("agentic_rag.nodes.retrieval_node.get_elasticsearch_client") as mock_get_client, \ patch("agentic_rag.nodes.retrieval_node._embed") as mock_embed: - mock_embed.side_effect = RuntimeError("OpenAI unavailable") + mock_embed.side_effect = RuntimeError("Mistral unavailable") mock_get_client.return_value = MagicMock() result = retrieval_node(sample_state) @@ -71,17 +71,13 @@ def test_retrieval_node_increments_attempt_on_failure(sample_state): assert result["retrieval_attempt"] == 1 -def test_embed_calls_openai(): +def test_embed_calls_mistral(): from agentic_rag.nodes.retrieval_node import _embed - mock_response = MagicMock() - mock_response.data = [MagicMock(embedding=[0.5] * 1536)] + with patch("agentic_rag.nodes.retrieval_node.MistralAIEmbeddings") as mock_emb: + mock_emb.return_value.embed_query.return_value = [0.5] * 1024 + result = _embed("hello world", "mstrl-test") - with patch("agentic_rag.nodes.retrieval_node.openai.OpenAI") as mock_oai: - mock_oai.return_value.embeddings.create.return_value = mock_response - result = _embed("hello world", "sk-test") - - assert result == [0.5] * 1536 - mock_oai.return_value.embeddings.create.assert_called_once_with( - model="text-embedding-3-small", input="hello world" - ) + assert result == [0.5] * 1024 + mock_emb.assert_called_once_with(model="mistral-embed", api_key="mstrl-test") + mock_emb.return_value.embed_query.assert_called_once_with("hello world")