diff --git a/.env.example b/.env.example index f8d9d40..17a4a16 100644 --- a/.env.example +++ b/.env.example @@ -18,12 +18,17 @@ NEO4J_PASSWORD=your_neo4j_password_here # NEO4J_DATABASE=neo4j # ============================================================================== -# REQUIRED: IPFS Storage (Lighthouse) +# REQUIRED: Akave S3-Compatible Storage # ============================================================================== -# Required for distributed storage of component outputs via IPFS -# Get API key at: https://www.lighthouse.storage/ +# Required for distributed storage of component outputs via Akave +# Akave provides decentralized S3-compatible storage with encryption +# Get credentials at: https://akave.ai/ -LIGHTHOUSE_API_KEY=your_lighthouse_api_key_here +AKAVE_ACCESS_KEY=your_akave_access_key_here +AKAVE_SECRET_KEY=your_akave_secret_key_here +AKAVE_BUCKET=agentic-rag +# Optional: Override default endpoint (defaults to o3-rc3.akave.xyz) +# AKAVE_ENDPOINT=https://o3-rc3.akave.xyz # ============================================================================== # OPTIONAL: LLM API Keys (only needed if using specific components) @@ -66,7 +71,7 @@ OPENAI_API_KEY=your_openai_api_key_here # ----------------------- # Required for all pipelines: # - Neo4j: NEO4J_URI, NEO4J_USERNAME, NEO4J_PASSWORD -# - Lighthouse IPFS: LIGHTHOUSE_API_KEY +# - Akave Storage: AKAVE_ACCESS_KEY, AKAVE_SECRET_KEY # # Optional (only needed if using specific components): # - GENERATOR.OPENROUTER requires: OPENROUTER_API_KEY diff --git a/agentic_rag/__init__.py b/agentic_rag/__init__.py index b1cfc58..debc7a0 100644 --- a/agentic_rag/__init__.py +++ b/agentic_rag/__init__.py @@ -2,6 +2,7 @@ from .components import ComponentRegistry, get_default_registry from .config import Config, get_config, get_global_config, set_global_config +from .export import JSONLDExporter from .pipeline import PipelineFactory, PipelineRunner from .types import ( DOCUMENT_STORE, @@ -34,4 +35,5 @@ "ComponentType", "DOCUMENT_STORE", "list_available_components", + "JSONLDExporter", ] diff --git a/agentic_rag/components/gates/gated_component.py b/agentic_rag/components/gates/gated_component.py index 1852759..9b5b631 100644 --- a/agentic_rag/components/gates/gated_component.py +++ b/agentic_rag/components/gates/gated_component.py @@ -1,10 +1,17 @@ -"""Generic wrapper that adds caching gates to any Haystack component.""" +"""Generic wrapper that adds caching gates to any Haystack component. + +FAIR Compliance: +- Supports content_type from ComponentSpec registry +- Supports run_id for provenance tracking +- Passes FAIR fields through to OutGate for storage +""" import asyncio import contextvars import time from typing import Any, Dict, List, Optional +from ...types.node_types import ContentType from ...utils.logger import get_logger from ...utils.metrics import MetricsCollector from ..neo4j_manager import GraphStore @@ -12,6 +19,31 @@ from .outgate import OutGate +def get_content_type_from_registry(component_name: str) -> Optional[ContentType]: + """ + Get the FAIR content_type from the component registry. + + Args: + component_name: Name of the component (e.g., "markdown_aware_chunker") + + Returns: + ContentType enum value or None if not found/specified + """ + from ..registry import get_default_registry + + registry = get_default_registry() + spec = registry.get_component_spec(component_name) + + if spec and spec.produces_content_type: + try: + return ContentType(spec.produces_content_type) + except ValueError: + pass + + # Fallback to None if not in registry or not specified + return None + + class GatedComponent: """ Generic wrapper that adds InGate/OutGate caching to any component. @@ -52,7 +84,9 @@ def __init__( graph_store: GraphStore, username: str, cache_key: Optional[str] = None, - retrieve_from_ipfs: bool = True, + retrieve_from_storage: bool = True, + run_id: Optional[str] = None, + content_type: Optional[ContentType] = None, ): """ Initialize gated component wrapper. @@ -64,7 +98,9 @@ def __init__( graph_store: Neo4j graph store username: User who owns this pipeline/data cache_key: Pipeline-agnostic cache key for cache lookups - retrieve_from_ipfs: If True, retrieves cached data from IPFS (default: False) + retrieve_from_storage: If True, retrieves cached data from Akave (default: False) + run_id: FAIR provenance - Run ID to link all outputs to + content_type: FAIR semantic type (auto-inferred from component_name if not provided) """ self.component = component self.component_id = component_id @@ -77,13 +113,20 @@ def __init__( self.logger = get_logger(f"{__name__}.{component_name}", username=username) self.metrics = MetricsCollector(username=username) - # Create gates with IPFS retrieval enabled (use cache_key for lookups) + # FAIR compliance fields + self.run_id = run_id + # Get content_type from registry if not provided + self.content_type = content_type or get_content_type_from_registry( + component_name + ) + + # Create gates with storage retrieval enabled (use cache_key for lookups) self.ingate = InGate( graph_store=graph_store, component_id=self.cache_key, # Use cache_key for lookups component_name=component_name, username=username, - retrieve_from_ipfs=True, # Always retrieve from IPFS for cached results + retrieve_from_storage=True, # Always retrieve from Akave for cached results ) self.outgate = OutGate( @@ -93,6 +136,17 @@ def __init__( username=username, ) + def set_run_id(self, run_id: str) -> None: + """ + Set the run_id for FAIR provenance tracking. + + This can be called after construction to link outputs to a specific Run. + + Args: + run_id: The Run ID to link all outputs to + """ + self.run_id = run_id + def run(self, **kwargs: Any) -> Dict[str, Any]: """ Run component with caching. @@ -173,11 +227,14 @@ def run(self, **kwargs: Any) -> Dict[str, Any]: single_output_items = self._extract_output_data(single_output) # Store: input → its outputs (handles both 1→1 and 1→N) + # Include FAIR compliance fields if single_output_items: self.outgate.store( input_data=input_item, output_data=single_output_items, component_config=component_config, + content_type=self.content_type, + run_id=self.run_id, ) except Exception as e: self.logger.warning(f"Failed to cache item: {e}") @@ -316,11 +373,14 @@ async def run_async(self, **kwargs: Any) -> Dict[str, Any]: single_output_items = self._extract_output_data(single_output) # Store: input → its outputs (handles both 1→1 and 1→N) - use async version + # Include FAIR compliance fields if single_output_items: await self.outgate.store_async( input_data=input_item, output_data=single_output_items, component_config=component_config, + content_type=self.content_type, + run_id=self.run_id, ) except Exception as e: self.logger.warning(f"Failed to cache item: {e}") diff --git a/agentic_rag/components/gates/ingate.py b/agentic_rag/components/gates/ingate.py index edd9c2a..d046558 100644 --- a/agentic_rag/components/gates/ingate.py +++ b/agentic_rag/components/gates/ingate.py @@ -8,7 +8,7 @@ import json from typing import Any, Dict, List, Optional -from ...utils.ipfs_client import LighthouseClient +from ...utils.akave_client import AkaveClient from ...utils.logger import get_logger from ..neo4j_manager import GraphStore @@ -35,8 +35,8 @@ def __init__( component_id: str, component_name: str, username: Optional[str] = None, - ipfs_client: Optional[LighthouseClient] = None, - retrieve_from_ipfs: bool = False, + storage_client: Optional[AkaveClient] = None, + retrieve_from_storage: bool = False, ): """ Initialize InGate. @@ -46,14 +46,14 @@ def __init__( component_id: Unique ID of the component this gate protects component_name: Human-readable component name username: Username for per-user logging - ipfs_client: Optional Lighthouse IPFS client (creates one if not provided) - retrieve_from_ipfs: If True, retrieves actual data from IPFS (default: False, returns metadata only) + storage_client: Optional Akave storage client (creates one if not provided) + retrieve_from_storage: If True, retrieves actual data from storage (default: False, returns metadata only) """ self.graph_store = graph_store self.component_id = component_id self.component_name = component_name - self.ipfs_client = ipfs_client or LighthouseClient() - self.retrieve_from_ipfs = retrieve_from_ipfs + self.storage_client = storage_client or AkaveClient() + self.retrieve_from_storage = retrieve_from_storage self.logger = get_logger(f"{__name__}.{component_name}", username=username) def check_cache_batch( @@ -77,8 +77,8 @@ def check_cache_batch( } cached_result format: - - If retrieve_from_ipfs=False: List[{fingerprint, ipfs_hash, data_type}] - - If retrieve_from_ipfs=True: List of actual data retrieved from IPFS + - If retrieve_from_storage=False: List[{fingerprint, object_key, data_type}] + - If retrieve_from_storage=True: List of actual data retrieved from storage Example (batch): >>> result = ingate.check_cache_batch( @@ -128,22 +128,24 @@ def check_cache_batch( # Found cached result cached_metadata = cache_map[fp] - if self.retrieve_from_ipfs: + if self.retrieve_from_storage: try: cached_data = [] for output_meta in cached_metadata: - ipfs_hash = output_meta["ipfs_hash"] + object_key = output_meta[ + "ipfs_hash" + ] # Field name kept for Neo4j compat data_type = output_meta.get("data_type") - data = self._retrieve_from_ipfs(ipfs_hash, data_type) + data = self._retrieve_from_storage(object_key, data_type) cached_data.append(data) cached.append((item, cached_data)) cache_hits += 1 except (ConnectionError, Exception) as e: - # IPFS retrieval failed - treat as cache miss - self.logger.warning(f"IPFS retrieval failed for {fp}: {e}") + # Storage retrieval failed - treat as cache miss + self.logger.warning(f"Storage retrieval failed for {fp}: {e}") uncached.append(item) else: - # Just return metadata (fingerprints + IPFS hashes) + # Just return metadata (fingerprints + object keys) cached.append((item, cached_metadata)) cache_hits += 1 else: @@ -213,24 +215,26 @@ async def check_cache_batch_async( # Found cached result cached_metadata = cache_map[fp] - if self.retrieve_from_ipfs: + if self.retrieve_from_storage: try: cached_data = [] for output_meta in cached_metadata: - ipfs_hash = output_meta["ipfs_hash"] + object_key = output_meta[ + "ipfs_hash" + ] # Field name kept for Neo4j compat data_type = output_meta.get("data_type") - data = await self._retrieve_from_ipfs_async( - ipfs_hash, data_type + data = await self._retrieve_from_storage_async( + object_key, data_type ) cached_data.append(data) cached.append((item, cached_data)) cache_hits += 1 except (ConnectionError, Exception) as e: - # IPFS retrieval failed - treat as cache miss - self.logger.warning(f"IPFS retrieval failed for {fp}: {e}") + # Storage retrieval failed - treat as cache miss + self.logger.warning(f"Storage retrieval failed for {fp}: {e}") uncached.append(item) else: - # Just return metadata (fingerprints + IPFS hashes) + # Just return metadata (fingerprints + object keys) cached.append((item, cached_metadata)) cache_hits += 1 else: @@ -308,22 +312,52 @@ def hash_config(self, config: Dict[str, Any]) -> str: hash_obj = hashlib.sha256(config_str.encode("utf-8")) return f"cfg_{hash_obj.hexdigest()[:16]}" - def _retrieve_from_ipfs( - self, ipfs_hash: str, data_type: Optional[str] = None + def _retrieve_from_storage( + self, object_key: str, data_type: Optional[str] = None ) -> Any: - """Retrieve from IPFS and reconstruct as Document.""" + """Retrieve from Akave storage and reconstruct as Document.""" from haystack import Document - text = self.ipfs_client.retrieve_text(ipfs_hash) + # Try to retrieve as JSON first (may contain embedding) + try: + data = self.storage_client.retrieve_json(object_key) + if isinstance(data, dict) and "content" in data: + # Reconstruct Document with all fields + return Document( + content=data.get("content", ""), + meta=data.get("meta", {}), + id=data.get("id"), + embedding=data.get("embedding"), # May be None or list + ) + except Exception: + pass + + # Fallback: retrieve as text + text = self.storage_client.retrieve_text(object_key) return Document(content=text) - async def _retrieve_from_ipfs_async( - self, ipfs_hash: str, data_type: Optional[str] = None + async def _retrieve_from_storage_async( + self, object_key: str, data_type: Optional[str] = None ) -> Any: - """Async version of _retrieve_from_ipfs.""" + """Async version of _retrieve_from_storage.""" from haystack import Document - text = await self.ipfs_client.retrieve_text_async(ipfs_hash) + # Try to retrieve as JSON first (may contain embedding) + try: + data = await self.storage_client.retrieve_json_async(object_key) + if isinstance(data, dict) and "content" in data: + # Reconstruct Document with all fields + return Document( + content=data.get("content", ""), + meta=data.get("meta", {}), + id=data.get("id"), + embedding=data.get("embedding"), # May be None or list + ) + except Exception: + pass + + # Fallback: retrieve as text + text = await self.storage_client.retrieve_text_async(object_key) return Document(content=text) def _serialize_for_fingerprint(self, data: Any) -> str: diff --git a/agentic_rag/components/gates/outgate.py b/agentic_rag/components/gates/outgate.py index 51cfbbc..bfca860 100644 --- a/agentic_rag/components/gates/outgate.py +++ b/agentic_rag/components/gates/outgate.py @@ -4,13 +4,19 @@ 1. Upload data to IPFS → get IPFS hash 2. Create DataPiece nodes in Neo4j 3. Create TRANSFORMED_BY edges connecting them + +FAIR Compliance: +- Supports content_type for semantic typing (Document, Chunk, Embedding) +- Supports run_id for provenance tracking (links to RunNode) +- Generates persistent URIs for each DataPiece """ import hashlib import json from typing import Any, Dict, List, Optional -from ...utils.ipfs_client import LighthouseClient +from ...types.node_types import ContentType, generate_uri +from ...utils.akave_client import AkaveClient from ...utils.logger import get_logger from ..neo4j_manager import GraphStore @@ -21,11 +27,11 @@ class OutGate: Simple responsibilities: - Upload data to IPFS (placeholder) - - Create DataPiece nodes with IPFS hashes + - Create DataPiece nodes with Akave object keys - Create TRANSFORMED_BY edges (input → output) Flow: - Output → [Upload to IPFS] → [Create DataPiece] → [Create TRANSFORMED_BY edge] + Output → [Upload to Akave] → [Create DataPiece] → [Create TRANSFORMED_BY edge] """ def __init__( @@ -34,7 +40,7 @@ def __init__( component_id: str, component_name: str, username: str, - ipfs_client: Optional[LighthouseClient] = None, + storage_client: Optional[AkaveClient] = None, ): """ Initialize OutGate. @@ -44,13 +50,13 @@ def __init__( component_id: Component doing the transformation component_name: Human-readable name username: Data owner - ipfs_client: Optional Lighthouse IPFS client (creates one if not provided) + storage_client: Optional Akave storage client (creates one if not provided) """ self.graph_store = graph_store self.component_id = component_id self.component_name = component_name self.username = username - self.ipfs_client = ipfs_client or LighthouseClient() + self.storage_client = storage_client or AkaveClient() self.logger = get_logger(f"{__name__}.{component_name}", username=username) def store( @@ -59,6 +65,8 @@ def store( output_data: List[Any], component_config: Dict[str, Any], processing_time_ms: Optional[int] = None, + content_type: Optional[ContentType] = None, + run_id: Optional[str] = None, ) -> Dict[str, Any]: """ Store transformation: input → component → outputs. @@ -76,6 +84,8 @@ def store( output_data: List of outputs (even if just one) component_config: Component configuration processing_time_ms: Processing time + content_type: FAIR semantic type (Document, Chunk, Embedding, DataNode) + run_id: FAIR provenance - links outputs to a Run node Returns: { @@ -85,8 +95,12 @@ def store( } Example: - >>> # Single output - >>> outgate.store(markdown, [chunk], config) + >>> # Single output with FAIR metadata + >>> outgate.store( + ... markdown, [chunk], config, + ... content_type=ContentType.CHUNK, + ... run_id="run_abc123" + ... ) >>> # Multiple outputs >>> outgate.store(markdown, [chunk1, chunk2, chunk3], config) @@ -97,26 +111,49 @@ def store( # Fingerprint input input_fingerprint = self.fingerprint_data(input_data) - # Process all outputs in batch + # Batch upload all outputs to Akave in parallel using ThreadPoolExecutor + from concurrent.futures import ThreadPoolExecutor, as_completed + output_records = [] - for output_item in output_data: - # Upload to IPFS - ipfs_hash = self._upload_to_ipfs(output_item) - # Fingerprint output + # Prepare data for parallel upload + upload_tasks = [] + for idx, output_item in enumerate(output_data): output_fingerprint = self.fingerprint_data(output_item) - - # Detect data type data_type = type(output_item).__name__ - - output_records.append( + upload_tasks.append( { + "idx": idx, + "item": output_item, "fingerprint": output_fingerprint, - "ipfs_hash": ipfs_hash, "data_type": data_type, } ) + # Upload in parallel (max 10 concurrent uploads) + with ThreadPoolExecutor(max_workers=min(10, len(upload_tasks))) as executor: + future_to_task = { + executor.submit(self._upload_to_storage, task["item"]): task + for task in upload_tasks + } + + for future in as_completed(future_to_task): + task = future_to_task[future] + object_key = future.result() + + record: Dict[str, Any] = { + "fingerprint": task["fingerprint"], + "ipfs_hash": object_key, + "data_type": task["data_type"], + } + + # Add FAIR compliance fields if provided + if content_type is not None: + record["content_type"] = content_type.value + record["uri"] = generate_uri(content_type, task["fingerprint"]) + + output_records.append(record) + # Store everything to Neo4j in one batch self.logger.info( f"Storing {len(output_records)} outputs to cache (component: {self.component_name})" @@ -127,7 +164,7 @@ def store( ) self.graph_store.store_transformation_batch( input_fingerprint=input_fingerprint, - input_ipfs_hash=self._upload_to_ipfs(input_data), + input_ipfs_hash=self._upload_to_storage(input_data), input_data_type=type(input_data).__name__, output_records=output_records, component_id=self.component_id, @@ -135,6 +172,7 @@ def store( config_hash=config_hash, username=self.username, processing_time_ms=processing_time_ms, + run_id=run_id, ) self.logger.debug( @@ -153,6 +191,8 @@ async def store_async( output_data: List[Any], component_config: Dict[str, Any], processing_time_ms: Optional[int] = None, + content_type: Optional[ContentType] = None, + run_id: Optional[str] = None, ) -> Dict[str, Any]: """ Async version of store. @@ -164,6 +204,8 @@ async def store_async( output_data: List of outputs (even if just one) component_config: Component configuration processing_time_ms: Processing time + content_type: FAIR semantic type (Document, Chunk, Embedding, DataNode) + run_id: FAIR provenance - links outputs to a Run node Returns: { @@ -172,32 +214,46 @@ async def store_async( "output_ipfs_hashes": List[str] } """ + import asyncio + # Hash config config_hash = self.hash_config(component_config) # Fingerprint input input_fingerprint = self.fingerprint_data(input_data) - # Process all outputs in batch - output_records = [] + # Prepare metadata for all outputs + output_metadata = [] for output_item in output_data: - # Upload to IPFS (async) - ipfs_hash = await self._upload_to_ipfs_async(output_item) - - # Fingerprint output output_fingerprint = self.fingerprint_data(output_item) - - # Detect data type data_type = type(output_item).__name__ - - output_records.append( + output_metadata.append( { + "item": output_item, "fingerprint": output_fingerprint, - "ipfs_hash": ipfs_hash, "data_type": data_type, } ) + # Batch upload all outputs to Akave in parallel using asyncio.gather + async def upload_with_metadata(meta: dict) -> dict: + object_key = await self._upload_to_storage_async(meta["item"]) + record: Dict[str, Any] = { + "fingerprint": meta["fingerprint"], + "ipfs_hash": object_key, + "data_type": meta["data_type"], + } + if content_type is not None: + record["content_type"] = content_type.value + record["uri"] = generate_uri(content_type, meta["fingerprint"]) + return record + + # Upload all in parallel + output_records = await asyncio.gather( + *[upload_with_metadata(meta) for meta in output_metadata] + ) + output_records = list(output_records) + # Store everything to Neo4j in one batch (async) self.logger.info( f"Storing {len(output_records)} outputs to cache (async, component: {self.component_name})" @@ -208,7 +264,7 @@ async def store_async( ) await self.graph_store.store_transformation_batch_async( input_fingerprint=input_fingerprint, - input_ipfs_hash=await self._upload_to_ipfs_async(input_data), + input_ipfs_hash=await self._upload_to_storage_async(input_data), input_data_type=type(input_data).__name__, output_records=output_records, component_id=self.component_id, @@ -216,6 +272,7 @@ async def store_async( config_hash=config_hash, username=self.username, processing_time_ms=processing_time_ms, + run_id=run_id, ) self.logger.debug( @@ -258,49 +315,49 @@ def hash_config(self, config: Dict[str, Any]) -> str: hash_obj = hashlib.sha256(config_str.encode("utf-8")) return f"cfg_{hash_obj.hexdigest()[:16]}" - def _upload_to_ipfs(self, data: Any) -> str: + def _upload_to_storage(self, data: Any) -> str: """ - Upload data to IPFS via Lighthouse and return CID. + Upload data to Akave and return object key. Args: data: Any data type (Document, dict, str, bytes, etc.) Returns: - IPFS CID (hash) + Object key (stored as ipfs_hash for Neo4j compatibility) """ try: - result = self.ipfs_client.upload_any(data) - ipfs_hash: str = result["Hash"] + result = self.storage_client.upload_any(data) + object_key: str = result["Hash"] self.logger.debug( - f"IPFS upload successful: CID={ipfs_hash}, Size={result.get('Size', 'unknown')}" + f"Akave upload successful: key={object_key}, Size={result.get('Size', 'unknown')}" ) - return ipfs_hash + return object_key except Exception as e: - self.logger.error(f"IPFS upload failed: {type(e).__name__}: {str(e)}") + self.logger.error(f"Akave upload failed: {type(e).__name__}: {str(e)}") raise - async def _upload_to_ipfs_async(self, data: Any) -> str: + async def _upload_to_storage_async(self, data: Any) -> str: """ - Async version of _upload_to_ipfs. + Async version of _upload_to_storage. - Upload data to IPFS via Lighthouse and return CID. + Upload data to Akave and return object key. Args: data: Any data type (Document, dict, str, bytes, etc.) Returns: - IPFS CID (hash) + Object key (stored as ipfs_hash for Neo4j compatibility) """ try: - result = await self.ipfs_client.upload_any_async(data) - ipfs_hash: str = result["Hash"] + result = await self.storage_client.upload_any_async(data) + object_key: str = result["Hash"] self.logger.debug( - f"IPFS upload successful (async): CID={ipfs_hash}, Size={result.get('Size', 'unknown')}" + f"Akave upload successful (async): key={object_key}, Size={result.get('Size', 'unknown')}" ) - return ipfs_hash + return object_key except Exception as e: self.logger.error( - f"IPFS upload failed (async): {type(e).__name__}: {str(e)}" + f"Akave upload failed (async): {type(e).__name__}: {str(e)}" ) raise diff --git a/agentic_rag/components/neo4j_manager.py b/agentic_rag/components/neo4j_manager.py index 97b612b..f191407 100644 --- a/agentic_rag/components/neo4j_manager.py +++ b/agentic_rag/components/neo4j_manager.py @@ -1,6 +1,7 @@ """Simple graph database store for batch nodes and edges.""" import ssl +import sys from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import certifi @@ -71,9 +72,12 @@ def __init__( " GraphStore(config=config)" ) - print(f"GraphStore connecting to: {self.uri} with user: {self.neo4j_username}") + print( + f"GraphStore connecting to: {self.uri} with user: {self.neo4j_username}", + file=sys.stderr, + ) if self.database: - print(f"Using database: {self.database}") + print(f"Using database: {self.database}", file=sys.stderr) # Use the same SSL setup as the working example ssl_ctx = ssl.create_default_context(cafile=certifi.where()) @@ -98,10 +102,10 @@ def __init__( # Verify connectivity like the working example try: self.driver.verify_connectivity() - print("GraphStore connected successfully!") + print("GraphStore connected successfully!", file=sys.stderr) self._initialized = True except Exception as e: - print(f"GraphStore connection failed: {e}") + print(f"GraphStore connection failed: {e}", file=sys.stderr) raise @classmethod @@ -556,6 +560,94 @@ async def _dfs_traversal_same_pipeline_async( return components + async def create_user_async(self, username: str) -> None: + """Create a new user node if it doesn't exist.""" + async with self.async_driver.session(database=self.database) as session: + # Use id as the merge key to match add_nodes_batch behavior + query = """ + MERGE (u:User {id: $username}) + SET u.username = $username + """ + await session.run(query, username=username) + + async def create_project_async(self, username: str, project_name: str) -> None: + """Create a new project for a user.""" + # Import here to avoid circular imports + import hashlib + + project_id = f"proj_{hashlib.sha256(f'{username}__{project_name}'.encode()).hexdigest()[:12]}" + + async with self.async_driver.session(database=self.database) as session: + query = """ + MERGE (u:User {id: $username}) + SET u.username = $username + MERGE (p:Project {id: $project_id}) + ON CREATE SET + p.name = $project_name, + p.username = $username, + p.created_at = datetime() + MERGE (u)-[:OWNS]->(p) + """ + await session.run( + query, + username=username, + project_name=project_name, + project_id=project_id, + ) + + async def get_user_projects_and_pipelines_async( + self, username: str + ) -> List[Dict[str, Any]]: + """Get all projects and pipelines for a user.""" + async with self.async_driver.session(database=self.database) as session: + query = """ + MATCH (u:User {username: $username})-[:OWNS]->(p:Project) + OPTIONAL MATCH (p)-[:FLOWS_TO]->(c:Component) + WHERE c.pipeline_name IS NOT NULL + WITH p, c.pipeline_name as pname, c.pipeline_type as ptype + WITH p, collect(DISTINCT CASE WHEN pname IS NOT NULL THEN {name: pname, type: ptype} END) as raw_pipelines + WITH p, [x IN raw_pipelines WHERE x IS NOT NULL] as pipelines + RETURN { + project: p.name, + pipelines: [x in pipelines | x.name], + details: pipelines + } as project_data + """ + result = await session.run(query, username=username) + results = await result.data() + return [record["project_data"] for record in results] + + async def project_exists_async(self, username: str, project_name: str) -> bool: + """Check if a project exists for a user.""" + async with self.async_driver.session(database=self.database) as session: + query = """ + MATCH (u:User {username: $username})-[:OWNS]->(p:Project {name: $project_name}) + RETURN p + """ + result = await session.run( + query, username=username, project_name=project_name + ) + return await result.single() is not None + + async def pipeline_exists_async( + self, username: str, project_name: str, pipeline_name: str + ) -> bool: + """Check if a pipeline exists in a project.""" + async with self.async_driver.session(database=self.database) as session: + query = """ + MATCH (u:User {username: $username})-[:OWNS]->(p:Project {name: $project_name}) + MATCH (p)-[:FLOWS_TO]->(c:Component {pipeline_name: $pipeline_name}) + RETURN c + LIMIT 1 + """ + result = await session.run( + query, + username=username, + project_name=project_name, + pipeline_name=pipeline_name, + ) + return await result.single() is not None + def lookup_cached_transformations_batch( self, input_fingerprints: List[str], component_id: str, config_hash: str ) -> Dict[str, List[Dict[str, Any]]]: @@ -680,6 +772,7 @@ def store_transformation_batch( config_hash: str, username: str, processing_time_ms: Optional[int] = None, + run_id: Optional[str] = None, ) -> None: """ Store a 1→N transformation in Neo4j. @@ -688,19 +781,23 @@ def store_transformation_batch( - Input DataPiece node (if not exists) - Output DataPiece nodes for all outputs - TRANSFORMED_BY edges from input to each output + - Optional GENERATED_BY edge to Run node (for FAIR compliance) Args: input_fingerprint: Fingerprint of input data - input_ipfs_hash: IPFS hash of input data + input_ipfs_hash: Storage object key (field name kept as ipfs_hash for Neo4j schema compat) input_data_type: Type of input data - output_records: List of {fingerprint, ipfs_hash, data_type} + output_records: List of {fingerprint, ipfs_hash, data_type, content_type?, uri?} + (ipfs_hash field stores Akave object key) component_id: ID of component that did transformation component_name: Name of component config_hash: Hash of component config username: User who owns this data processing_time_ms: Optional processing time + run_id: Optional run ID for FAIR provenance tracking """ with self.driver.session(database=self.database) as session: + # Main query to create DataPiece nodes and TRANSFORMED_BY edges query = """ // Create or get input DataPiece MERGE (input:DataPiece {fingerprint: $input_fingerprint}) @@ -719,6 +816,10 @@ def store_transformation_batch( out.data_type = output.data_type, out.username = $username, out.created_at = datetime() + // Set FAIR compliance fields if provided + SET out.content_type = COALESCE(output.content_type, out.content_type), + out.uri = COALESCE(output.uri, out.uri), + out.generated_by = COALESCE($run_id, out.generated_by) // Create TRANSFORMED_BY edge MERGE (input)-[t:TRANSFORMED_BY { @@ -742,8 +843,22 @@ def store_transformation_batch( config_hash=config_hash, username=username, processing_time_ms=processing_time_ms, + run_id=run_id, ).consume() + # Create GENERATED_BY relationship to Run node if run_id provided + if run_id: + run_query = """ + UNWIND $output_fingerprints AS fp + MATCH (d:DataPiece {fingerprint: fp}) + MATCH (r:Run {id: $run_id}) + MERGE (d)-[:GENERATED_BY]->(r) + """ + output_fps = [r["fingerprint"] for r in output_records] + session.run( + run_query, output_fingerprints=output_fps, run_id=run_id + ).consume() + async def store_transformation_batch_async( self, input_fingerprint: str, @@ -755,6 +870,7 @@ async def store_transformation_batch_async( config_hash: str, username: str, processing_time_ms: Optional[int] = None, + run_id: Optional[str] = None, ) -> None: """ Async version of store_transformation_batch. @@ -763,14 +879,16 @@ async def store_transformation_batch_async( Args: input_fingerprint: Fingerprint of input data - input_ipfs_hash: IPFS hash of input data + input_ipfs_hash: Storage object key (field name kept as ipfs_hash for Neo4j schema compat) input_data_type: Type of input data - output_records: List of {fingerprint, ipfs_hash, data_type} + output_records: List of {fingerprint, ipfs_hash, data_type, content_type?, uri?} + (ipfs_hash field stores Akave object key) component_id: ID of component that did transformation component_name: Name of component config_hash: Hash of component config username: User who owns this data processing_time_ms: Optional processing time + run_id: Optional run ID for FAIR provenance tracking """ async with self.async_driver.session(database=self.database) as session: query = """ @@ -791,6 +909,10 @@ async def store_transformation_batch_async( out.data_type = output.data_type, out.username = $username, out.created_at = datetime() + // Set FAIR compliance fields if provided + SET out.content_type = COALESCE(output.content_type, out.content_type), + out.uri = COALESCE(output.uri, out.uri), + out.generated_by = COALESCE($run_id, out.generated_by) // Create TRANSFORMED_BY edge MERGE (input)-[t:TRANSFORMED_BY { @@ -814,5 +936,241 @@ async def store_transformation_batch_async( config_hash=config_hash, username=username, processing_time_ms=processing_time_ms, + run_id=run_id, ) await result.consume() + + # Create GENERATED_BY relationship to Run node if run_id provided + if run_id: + run_query = """ + UNWIND $output_fingerprints AS fp + MATCH (d:DataPiece {fingerprint: fp}) + MATCH (r:Run {id: $run_id}) + MERGE (d)-[:GENERATED_BY]->(r) + """ + output_fps = [r["fingerprint"] for r in output_records] + result = await session.run( + run_query, output_fingerprints=output_fps, run_id=run_id + ) + await result.consume() + + # ========================================================================= + # FAIR Compliance: RunNode Methods + # ========================================================================= + + def store_run_node( + self, + run_id: str, + pipeline_name: str, + username: str, + project: str = "default", + pipeline_version: Optional[str] = None, + git_commit: Optional[str] = None, + model_names: Optional[List[str]] = None, + config_hash: Optional[str] = None, + started_at: Optional[str] = None, + uri: Optional[str] = None, + ) -> None: + """ + Store a RunNode in Neo4j for provenance tracking. + + Args: + run_id: Unique run identifier + pipeline_name: Name of the pipeline being run + username: User who initiated the run + project: Project name + pipeline_version: Optional pipeline version + git_commit: Optional git commit hash + model_names: Optional list of model names used + config_hash: Optional hash of pipeline config + started_at: ISO timestamp when run started + uri: Persistent URI for FAIR compliance + """ + with self.driver.session(database=self.database) as session: + query = """ + MERGE (r:Run {id: $run_id}) + ON CREATE SET + r.pipeline_name = $pipeline_name, + r.username = $username, + r.project = $project, + r.started_at = $started_at, + r.uri = $uri, + r.created_at = datetime() + ON MATCH SET + r.pipeline_name = $pipeline_name, + r.username = $username, + r.project = $project + """ + params: Dict[str, Any] = { + "run_id": run_id, + "pipeline_name": pipeline_name, + "username": username, + "project": project, + "started_at": started_at, + "uri": uri, + } + + session.run(query, **params).consume() + + # Set optional properties separately to avoid null issues + if pipeline_version or git_commit or model_names or config_hash: + set_clauses = [] + opt_params: Dict[str, Any] = {"run_id": run_id} + + if pipeline_version: + set_clauses.append("r.pipeline_version = $pipeline_version") + opt_params["pipeline_version"] = pipeline_version + if git_commit: + set_clauses.append("r.git_commit = $git_commit") + opt_params["git_commit"] = git_commit + if model_names: + set_clauses.append("r.model_names = $model_names") + opt_params["model_names"] = model_names + if config_hash: + set_clauses.append("r.config_hash = $config_hash") + opt_params["config_hash"] = config_hash + + if set_clauses: + update_query = f""" + MATCH (r:Run {{id: $run_id}}) + SET {', '.join(set_clauses)} + """ + session.run(update_query, **opt_params).consume() + + async def store_run_node_async( + self, + run_id: str, + pipeline_name: str, + username: str, + project: str = "default", + pipeline_version: Optional[str] = None, + git_commit: Optional[str] = None, + model_names: Optional[List[str]] = None, + config_hash: Optional[str] = None, + started_at: Optional[str] = None, + uri: Optional[str] = None, + ) -> None: + """Async version of store_run_node.""" + async with self.async_driver.session(database=self.database) as session: + query = """ + MERGE (r:Run {id: $run_id}) + ON CREATE SET + r.pipeline_name = $pipeline_name, + r.username = $username, + r.project = $project, + r.started_at = $started_at, + r.uri = $uri, + r.created_at = datetime() + ON MATCH SET + r.pipeline_name = $pipeline_name, + r.username = $username, + r.project = $project + """ + params: Dict[str, Any] = { + "run_id": run_id, + "pipeline_name": pipeline_name, + "username": username, + "project": project, + "started_at": started_at, + "uri": uri, + } + + result = await session.run(query, **params) + await result.consume() + + # Set optional properties + if pipeline_version or git_commit or model_names or config_hash: + set_clauses = [] + opt_params: Dict[str, Any] = {"run_id": run_id} + + if pipeline_version: + set_clauses.append("r.pipeline_version = $pipeline_version") + opt_params["pipeline_version"] = pipeline_version + if git_commit: + set_clauses.append("r.git_commit = $git_commit") + opt_params["git_commit"] = git_commit + if model_names: + set_clauses.append("r.model_names = $model_names") + opt_params["model_names"] = model_names + if config_hash: + set_clauses.append("r.config_hash = $config_hash") + opt_params["config_hash"] = config_hash + + if set_clauses: + update_query = f""" + MATCH (r:Run {{id: $run_id}}) + SET {', '.join(set_clauses)} + """ + result = await session.run(update_query, **opt_params) + await result.consume() + + def update_run_finished( + self, + run_id: str, + finished_at: str, + success: bool = True, + error: Optional[str] = None, + ) -> None: + """ + Update a RunNode with completion information. + + Args: + run_id: Run identifier + finished_at: ISO timestamp when run finished + success: Whether the run completed successfully + error: Optional error message if failed + """ + with self.driver.session(database=self.database) as session: + query = """ + MATCH (r:Run {id: $run_id}) + SET r.finished_at = $finished_at, + r.success = $success + """ + params: Dict[str, Any] = { + "run_id": run_id, + "finished_at": finished_at, + "success": success, + } + + if error: + query = """ + MATCH (r:Run {id: $run_id}) + SET r.finished_at = $finished_at, + r.success = $success, + r.error = $error + """ + params["error"] = error + + session.run(query, **params).consume() + + async def update_run_finished_async( + self, + run_id: str, + finished_at: str, + success: bool = True, + error: Optional[str] = None, + ) -> None: + """Async version of update_run_finished.""" + async with self.async_driver.session(database=self.database) as session: + query = """ + MATCH (r:Run {id: $run_id}) + SET r.finished_at = $finished_at, + r.success = $success + """ + params: Dict[str, Any] = { + "run_id": run_id, + "finished_at": finished_at, + "success": success, + } + + if error: + query = """ + MATCH (r:Run {id: $run_id}) + SET r.finished_at = $finished_at, + r.success = $success, + r.error = $error + """ + params["error"] = error + + result = await session.run(query, **params) + await result.consume() diff --git a/agentic_rag/components/registry.py b/agentic_rag/components/registry.py index 7a9ba08..9783c37 100644 --- a/agentic_rag/components/registry.py +++ b/agentic_rag/components/registry.py @@ -12,6 +12,7 @@ PipelineUsage, create_haystack_component, ) +from ..types.node_types import ContentType class ComponentRegistry: @@ -74,7 +75,7 @@ def get_component_instance(self, spec: ComponentSpec) -> Any: # Generate cache key based on component class and sorted configuration # This ensures identical configs share the same instance - config_str = json.dumps(spec.default_config, sort_keys=True) + config_str = json.dumps(spec.get_config(), sort_keys=True) config_hash = hashlib.md5(config_str.encode()).hexdigest() cache_key = f"{spec.haystack_class}_{config_hash}" @@ -104,7 +105,7 @@ def clear_cache(self) -> None: def _initialize_default_components(self) -> None: """Initialize default Haystack component mappings.""" - # Converters + # Converters - produce DOCUMENT self.register_component( ComponentSpec( name="pdf_converter", @@ -114,6 +115,7 @@ def _initialize_default_components(self) -> None: output_types=[DataType.LIST_DOCUMENT], pipeline_usage=PipelineUsage.INDEXING, default_config={}, + produces_content_type=ContentType.DOCUMENT.value, ) ) @@ -126,6 +128,7 @@ def _initialize_default_components(self) -> None: output_types=[DataType.LIST_DOCUMENT], pipeline_usage=PipelineUsage.INDEXING, default_config={}, + produces_content_type=ContentType.DOCUMENT.value, ) ) @@ -138,6 +141,7 @@ def _initialize_default_components(self) -> None: output_types=[DataType.LIST_DOCUMENT], pipeline_usage=PipelineUsage.INDEXING, default_config={}, + produces_content_type=ContentType.DOCUMENT.value, ) ) @@ -150,6 +154,7 @@ def _initialize_default_components(self) -> None: output_types=[DataType.LIST_DOCUMENT], pipeline_usage=PipelineUsage.INDEXING, default_config={}, + produces_content_type=ContentType.DOCUMENT.value, ) ) @@ -162,6 +167,7 @@ def _initialize_default_components(self) -> None: output_types=[DataType.LIST_DOCUMENT], pipeline_usage=PipelineUsage.INDEXING, default_config={}, + produces_content_type=ContentType.DOCUMENT.value, ) ) @@ -178,6 +184,7 @@ def _initialize_default_components(self) -> None: "output_format": "markdown", "store_full_path": False, }, + produces_content_type=ContentType.DOCUMENT.value, ) ) @@ -192,10 +199,11 @@ def _initialize_default_components(self) -> None: default_config={ "store_full_path": False, }, + produces_content_type=ContentType.DOCUMENT.value, ) ) - # Chunkers/Splitters - Only verified existing ones + # Chunkers/Splitters - produce CHUNK self.register_component( ComponentSpec( name="chunker", @@ -205,10 +213,11 @@ def _initialize_default_components(self) -> None: output_types=[DataType.LIST_DOCUMENT], pipeline_usage=PipelineUsage.INDEXING, default_config={"split_by": "sentence", "split_length": 512}, + produces_content_type=ContentType.CHUNK.value, ) ) - # Custom Chunkers + # Custom Chunkers - produce CHUNK self.register_component( ComponentSpec( name="markdown_aware_chunker", @@ -218,6 +227,7 @@ def _initialize_default_components(self) -> None: output_types=[DataType.LIST_DOCUMENT], pipeline_usage=PipelineUsage.INDEXING, default_config={"chunk_size": 1000, "chunk_overlap": 100}, + produces_content_type=ContentType.CHUNK.value, ) ) @@ -234,10 +244,11 @@ def _initialize_default_components(self) -> None: "max_chunk_size": 1000, "overlap_size": 50, }, + produces_content_type=ContentType.CHUNK.value, ) ) - # Embedders - Only verified existing ones + # Embedders - produce EMBEDDING self.register_component( ComponentSpec( name="embedder", @@ -247,6 +258,7 @@ def _initialize_default_components(self) -> None: output_types=[DataType.LIST_LIST_FLOAT], pipeline_usage=PipelineUsage.RETRIEVAL, default_config={"model": "sentence-transformers/all-MiniLM-L6-v2"}, + produces_content_type=ContentType.EMBEDDING.value, ) ) @@ -259,6 +271,7 @@ def _initialize_default_components(self) -> None: output_types=[DataType.LIST_DOCUMENT], # Documents with embeddings pipeline_usage=PipelineUsage.INDEXING, default_config={"model": "sentence-transformers/all-MiniLM-L6-v2"}, + produces_content_type=ContentType.EMBEDDING.value, ) ) @@ -324,7 +337,18 @@ def _initialize_default_components(self) -> None: output_types=[DataType.STRING], pipeline_usage=PipelineUsage.RETRIEVAL, dependencies=[], - default_config={}, + default_config={ + "template": ( + "Answer the question based only on the provided documents.\n\n" + "Question: {{query}}\n\n" + "Documents:\n" + "{% for doc in documents %}\n" + "{{ doc.content }}\n" + "---\n" + "{% endfor %}\n\n" + "Answer:" + ) + }, parallelizable=False, ) ) diff --git a/agentic_rag/config.py b/agentic_rag/config.py index c9f9f6d..97218fa 100644 --- a/agentic_rag/config.py +++ b/agentic_rag/config.py @@ -14,7 +14,9 @@ neo4j_username="neo4j", neo4j_password="password", openrouter_api_key="your-key", - lighthouse_api_key="your-key" + akave_access_key="your-access-key", + akave_secret_key="your-secret-key", + akave_bucket="your-bucket" ) # Use with PipelineFactory @@ -27,7 +29,10 @@ - NEO4J_PASSWORD - NEO4J_DATABASE - OPENROUTER_API_KEY - - LIGHTHOUSE_API_KEY + - AKAVE_ACCESS_KEY + - AKAVE_SECRET_KEY + - AKAVE_BUCKET + - AKAVE_ENDPOINT - AGENTIC_RAG_LOG_LEVEL - AGENTIC_RAG_LOG_FILE """ @@ -46,7 +51,10 @@ class Config: neo4j_password: Neo4j password neo4j_database: Neo4j database name openrouter_api_key: OpenRouter API key for LLM access - lighthouse_api_key: Lighthouse IPFS API key + akave_access_key: Akave S3 access key + akave_secret_key: Akave S3 secret key + akave_bucket: Akave S3 bucket name + akave_endpoint: Akave S3 endpoint URL log_level: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL) log_file: Optional log file path agentic_root_dir: Root directory for agentic-rag data storage @@ -59,7 +67,10 @@ def __init__( neo4j_password: Optional[str] = None, neo4j_database: Optional[str] = None, openrouter_api_key: Optional[str] = None, - lighthouse_api_key: Optional[str] = None, + akave_access_key: Optional[str] = None, + akave_secret_key: Optional[str] = None, + akave_bucket: Optional[str] = None, + akave_endpoint: Optional[str] = None, log_level: Optional[str] = None, log_file: Optional[str] = None, agentic_root_dir: Optional[str] = None, @@ -77,7 +88,10 @@ def __init__( neo4j_password: Neo4j password (defaults to NEO4J_PASSWORD env var) neo4j_database: Neo4j database name (defaults to NEO4J_DATABASE env var) openrouter_api_key: OpenRouter API key (defaults to OPENROUTER_API_KEY env var) - lighthouse_api_key: Lighthouse IPFS key (defaults to LIGHTHOUSE_API_KEY env var) + akave_access_key: Akave S3 access key (defaults to AKAVE_ACCESS_KEY env var) + akave_secret_key: Akave S3 secret key (defaults to AKAVE_SECRET_KEY env var) + akave_bucket: Akave bucket name (defaults to AKAVE_BUCKET env var or "agentic-rag") + akave_endpoint: Akave endpoint URL (defaults to AKAVE_ENDPOINT env var) log_level: Log level (defaults to AGENTIC_RAG_LOG_LEVEL env var or INFO) log_file: Log file path (defaults to AGENTIC_RAG_LOG_FILE env var) agentic_root_dir: Root directory for agentic-rag data (defaults to AGENTIC_ROOT_DIR env var or ./data) @@ -91,7 +105,14 @@ def __init__( # API keys self.openrouter_api_key = openrouter_api_key or os.getenv("OPENROUTER_API_KEY") - self.lighthouse_api_key = lighthouse_api_key or os.getenv("LIGHTHOUSE_API_KEY") + + # Akave S3-compatible storage configuration + self.akave_access_key = akave_access_key or os.getenv("AKAVE_ACCESS_KEY") + self.akave_secret_key = akave_secret_key or os.getenv("AKAVE_SECRET_KEY") + self.akave_bucket = akave_bucket or os.getenv("AKAVE_BUCKET", "agentic-rag") + self.akave_endpoint = akave_endpoint or os.getenv( + "AKAVE_ENDPOINT", "https://o3-rc3.akave.xyz" + ) # Logging configuration self.log_level = log_level or os.getenv("AGENTIC_RAG_LOG_LEVEL") or "INFO" @@ -144,14 +165,14 @@ def validate_openrouter(self) -> bool: """ return self.openrouter_api_key is not None - def validate_lighthouse(self) -> bool: + def validate_akave(self) -> bool: """ - Check if Lighthouse API key is configured. + Check if Akave S3 credentials are configured. Returns: - True if Lighthouse API key is present + True if Akave access key and secret key are present """ - return self.lighthouse_api_key is not None + return self.akave_access_key is not None and self.akave_secret_key is not None def get_project_path(self, username: str, project: str) -> str: """ @@ -181,7 +202,10 @@ def to_dict(self) -> dict: "neo4j_password": "***" if self.neo4j_password else None, "neo4j_database": self.neo4j_database, "openrouter_api_key": "***" if self.openrouter_api_key else None, - "lighthouse_api_key": "***" if self.lighthouse_api_key else None, + "akave_access_key": "***" if self.akave_access_key else None, + "akave_secret_key": "***" if self.akave_secret_key else None, + "akave_bucket": self.akave_bucket, + "akave_endpoint": self.akave_endpoint, "log_level": self.log_level, "log_file": self.log_file, "agentic_root_dir": self.agentic_root_dir, diff --git a/agentic_rag/export/__init__.py b/agentic_rag/export/__init__.py new file mode 100644 index 0000000..18a7fb1 --- /dev/null +++ b/agentic_rag/export/__init__.py @@ -0,0 +1,5 @@ +"""Export module for FAIR-compliant data serialization.""" + +from .jsonld_exporter import JSONLDExporter + +__all__ = ["JSONLDExporter"] diff --git a/agentic_rag/export/jsonld_exporter.py b/agentic_rag/export/jsonld_exporter.py new file mode 100644 index 0000000..cc5a1bf --- /dev/null +++ b/agentic_rag/export/jsonld_exporter.py @@ -0,0 +1,508 @@ +""" +JSON-LD Exporter for FAIR-compliant data serialization. + +Exports Neo4j graph data to JSON-LD format using standard vocabularies: +- schema.org: General metadata +- Dublin Core (dcterms): Document metadata +- PROV-O: Provenance information +- DCAT: Dataset catalog vocabulary + +This enables interoperability with other linked data systems and +compliance with FAIR principles I1-I3 (Interoperability). +""" + +import json +from datetime import datetime +from typing import Any, Dict, List, Optional + +from ..components.neo4j_manager import GraphStore +from ..types.node_types import ARKHAI_NAMESPACE + +# JSON-LD Context with standard vocabularies +JSONLD_CONTEXT = { + "@vocab": "https://schema.org/", + "schema": "https://schema.org/", + "dcterms": "http://purl.org/dc/terms/", + "prov": "http://www.w3.org/ns/prov#", + "dcat": "http://www.w3.org/ns/dcat#", + "xsd": "http://www.w3.org/2001/XMLSchema#", + "arkhai": ARKHAI_NAMESPACE + "/", + # Property mappings + "identifier": "dcterms:identifier", + "created": {"@id": "dcterms:created", "@type": "xsd:dateTime"}, + "modified": {"@id": "dcterms:modified", "@type": "xsd:dateTime"}, + "title": "dcterms:title", + "description": "dcterms:description", + "creator": "dcterms:creator", + "publisher": "dcterms:publisher", + "source": "dcterms:source", + "format": "dcterms:format", + # PROV-O mappings + "wasGeneratedBy": {"@id": "prov:wasGeneratedBy", "@type": "@id"}, + "wasDerivedFrom": {"@id": "prov:wasDerivedFrom", "@type": "@id"}, + "wasAttributedTo": {"@id": "prov:wasAttributedTo", "@type": "@id"}, + "generatedAtTime": {"@id": "prov:generatedAtTime", "@type": "xsd:dateTime"}, + "startedAtTime": {"@id": "prov:startedAtTime", "@type": "xsd:dateTime"}, + "endedAtTime": {"@id": "prov:endedAtTime", "@type": "xsd:dateTime"}, + "used": {"@id": "prov:used", "@type": "@id"}, + # DCAT mappings + "distribution": "dcat:distribution", + "accessURL": {"@id": "dcat:accessURL", "@type": "@id"}, + "downloadURL": {"@id": "dcat:downloadURL", "@type": "@id"}, + "byteSize": "dcat:byteSize", + "mediaType": "dcat:mediaType", + # Custom Arkhai properties + "fingerprint": "arkhai:fingerprint", + "contentType": "arkhai:contentType", + "storageKey": "arkhai:storageKey", + "configHash": "arkhai:configHash", + "processingTime": "arkhai:processingTimeMs", +} + + +class JSONLDExporter: + """ + Export Neo4j graph data to JSON-LD format. + + Supports exporting: + - Individual DataPieces as schema:DigitalDocument + - Runs as prov:Activity + - Pipelines as schema:SoftwareApplication + - Complete datasets as dcat:Dataset + + Example: + ```python + exporter = JSONLDExporter(graph_store) + + # Export a single DataPiece + doc = exporter.export_datapiece("fp_abc123") + + # Export a complete dataset for a user/project + dataset = exporter.export_dataset( + username="user1", + project="my_project", + title="My RAG Dataset", + description="Documents processed for RAG" + ) + + # Save to file + exporter.save_to_file(dataset, "dataset.jsonld") + ``` + """ + + def __init__(self, graph_store: GraphStore): + """ + Initialize the exporter. + + Args: + graph_store: Neo4j GraphStore instance + """ + self.graph_store = graph_store + + def export_datapiece(self, fingerprint: str) -> Optional[Dict[str, Any]]: + """ + Export a single DataPiece as JSON-LD. + + Args: + fingerprint: The fingerprint of the DataPiece + + Returns: + JSON-LD document or None if not found + """ + query = """ + MATCH (d:DataPiece {fingerprint: $fingerprint}) + OPTIONAL MATCH (d)-[:GENERATED_BY]->(r:Run) + OPTIONAL MATCH (parent:DataPiece)-[t:TRANSFORMED_BY]->(d) + RETURN d, r, parent, t + """ + + with self.graph_store.driver.session( + database=self.graph_store.database + ) as session: + result = session.run(query, fingerprint=fingerprint).single() + + if not result: + return None + + node = result["d"] + run = result["r"] + parent = result["parent"] + transform = result["t"] + + return self._datapiece_to_jsonld(node, run, parent, transform) + + def export_run(self, run_id: str) -> Optional[Dict[str, Any]]: + """ + Export a Run as JSON-LD prov:Activity. + + Args: + run_id: The run ID + + Returns: + JSON-LD document or None if not found + """ + query = """ + MATCH (r:Run {id: $run_id}) + OPTIONAL MATCH (d:DataPiece)-[:GENERATED_BY]->(r) + RETURN r, collect(d.fingerprint) as generated_fingerprints + """ + + with self.graph_store.driver.session( + database=self.graph_store.database + ) as session: + result = session.run(query, run_id=run_id).single() + + if not result: + return None + + run = result["r"] + generated = result["generated_fingerprints"] + + return self._run_to_jsonld(run, generated) + + def export_dataset( + self, + username: str, + project: str, + title: Optional[str] = None, + description: Optional[str] = None, + include_runs: bool = True, + include_lineage: bool = True, + ) -> Dict[str, Any]: + """ + Export a complete dataset as JSON-LD dcat:Dataset. + + Args: + username: Owner username + project: Project name + title: Dataset title (defaults to project name) + description: Dataset description + include_runs: Include Run provenance + include_lineage: Include transformation lineage + + Returns: + JSON-LD dataset document + """ + # Get all DataPieces for this user/project + datapieces_query = """ + MATCH (d:DataPiece {username: $username}) + OPTIONAL MATCH (d)-[:GENERATED_BY]->(r:Run) + WHERE r.pipeline_name CONTAINS $project OR r IS NULL + RETURN DISTINCT d, r + ORDER BY d.created_at + """ + + # Get lineage if requested + lineage_query = """ + MATCH (d1:DataPiece {username: $username})-[t:TRANSFORMED_BY]->(d2:DataPiece) + RETURN d1.fingerprint as from_fp, d2.fingerprint as to_fp, + t.component_name as component, t.config_hash as config_hash + """ + + # Get runs if requested + runs_query = """ + MATCH (r:Run) + WHERE r.pipeline_name CONTAINS $project + RETURN r + ORDER BY r.started_at + """ + + datapieces = [] + runs = [] + lineage = [] + + with self.graph_store.driver.session( + database=self.graph_store.database + ) as session: + # Fetch DataPieces + for record in session.run( + datapieces_query, username=username, project=project + ): + node = record["d"] + run = record["r"] + datapieces.append(self._datapiece_to_jsonld(node, run)) + + # Fetch lineage + if include_lineage: + for record in session.run(lineage_query, username=username): + lineage.append( + { + "from": f"{ARKHAI_NAMESPACE}/data/{record['from_fp']}", + "to": f"{ARKHAI_NAMESPACE}/data/{record['to_fp']}", + "component": record["component"], + "configHash": record["config_hash"], + } + ) + + # Fetch runs + if include_runs: + for record in session.run(runs_query, project=project): + run_node = record["r"] + runs.append(self._run_to_jsonld(run_node, [])) + + # Build dataset document + dataset_id = f"{ARKHAI_NAMESPACE}/dataset/{username}/{project}" + + dataset: Dict[str, Any] = { + "@context": JSONLD_CONTEXT, + "@type": "dcat:Dataset", + "@id": dataset_id, + "identifier": f"{username}/{project}", + "title": title or project, + "creator": username, + "created": datetime.utcnow().isoformat() + "Z", + } + + if description: + dataset["description"] = description + + # Add DataPieces as distributions + if datapieces: + dataset["distribution"] = datapieces + + # Add provenance activities + if runs: + dataset["prov:wasGeneratedBy"] = runs + + # Add lineage graph + if lineage: + dataset["arkhai:lineage"] = lineage + + # Add statistics + dataset["arkhai:statistics"] = { + "totalDataPieces": len(datapieces), + "totalRuns": len(runs), + "totalTransformations": len(lineage), + } + + return dataset + + def export_pipeline_definition( + self, username: str, project: str, pipeline_name: str + ) -> Optional[Dict[str, Any]]: + """ + Export a pipeline definition as JSON-LD. + + Args: + username: Owner username + project: Project name + pipeline_name: Pipeline name + + Returns: + JSON-LD document or None if not found + """ + query = """ + MATCH (c:Component) + WHERE c.author = $username + AND c.project = $project + AND c.pipeline_name = $pipeline_name + RETURN c + ORDER BY c.id + """ + + components = [] + with self.graph_store.driver.session( + database=self.graph_store.database + ) as session: + for record in session.run( + query, username=username, project=project, pipeline_name=pipeline_name + ): + comp = record["c"] + components.append(self._component_to_jsonld(comp)) + + if not components: + return None + + pipeline_id = ( + f"{ARKHAI_NAMESPACE}/pipeline/{username}/{project}/{pipeline_name}" + ) + + return { + "@context": JSONLD_CONTEXT, + "@type": "schema:SoftwareApplication", + "@id": pipeline_id, + "name": pipeline_name, + "creator": username, + "schema:hasPart": components, + } + + def _datapiece_to_jsonld( + self, + node: Any, + run: Optional[Any] = None, + parent: Optional[Any] = None, + transform: Optional[Any] = None, + ) -> Dict[str, Any]: + """Convert a DataPiece node to JSON-LD.""" + fingerprint = node.get("fingerprint", "") + uri = node.get("uri") or f"{ARKHAI_NAMESPACE}/data/{fingerprint}" + + doc: Dict[str, Any] = { + "@type": "schema:DigitalDocument", + "@id": uri, + "identifier": fingerprint, + "fingerprint": fingerprint, + } + + # Add optional fields + if node.get("ipfs_hash"): + doc["storageKey"] = node["ipfs_hash"] + + if node.get("data_type"): + doc["format"] = node["data_type"] + + if node.get("content_type"): + doc["contentType"] = node["content_type"] + + if node.get("created_at"): + created = node["created_at"] + if hasattr(created, "isoformat"): + doc["created"] = created.isoformat() + "Z" + else: + doc["created"] = str(created) + + if node.get("username"): + doc["creator"] = node["username"] + + # Add provenance + if run: + run_id = run.get("id", "") + doc["wasGeneratedBy"] = f"{ARKHAI_NAMESPACE}/run/{run_id}" + if run.get("started_at"): + started = run["started_at"] + if hasattr(started, "isoformat"): + doc["generatedAtTime"] = started.isoformat() + "Z" + + # Add derivation + if parent: + parent_fp = parent.get("fingerprint", "") + doc["wasDerivedFrom"] = f"{ARKHAI_NAMESPACE}/data/{parent_fp}" + + if transform: + doc["arkhai:transformedBy"] = { + "component": transform.get("component_name"), + "configHash": transform.get("config_hash"), + } + + return doc + + def _run_to_jsonld( + self, run: Any, generated_fingerprints: List[str] + ) -> Dict[str, Any]: + """Convert a Run node to JSON-LD prov:Activity.""" + run_id = run.get("id", "") + uri = f"{ARKHAI_NAMESPACE}/run/{run_id}" + + activity: Dict[str, Any] = { + "@type": "prov:Activity", + "@id": uri, + "identifier": run_id, + } + + if run.get("pipeline_name"): + activity["prov:used"] = ( + f"{ARKHAI_NAMESPACE}/pipeline/{run.get('pipeline_name')}" + ) + + if run.get("started_at"): + started = run["started_at"] + if hasattr(started, "isoformat"): + activity["startedAtTime"] = started.isoformat() + "Z" + else: + activity["startedAtTime"] = str(started) + + if run.get("ended_at"): + ended = run["ended_at"] + if hasattr(ended, "isoformat"): + activity["endedAtTime"] = ended.isoformat() + "Z" + else: + activity["endedAtTime"] = str(ended) + + if run.get("success") is not None: + activity["arkhai:success"] = run["success"] + + if run.get("input_count"): + activity["arkhai:inputCount"] = run["input_count"] + + if run.get("output_count"): + activity["arkhai:outputCount"] = run["output_count"] + + # Add generated entities + if generated_fingerprints: + activity["prov:generated"] = [ + f"{ARKHAI_NAMESPACE}/data/{fp}" for fp in generated_fingerprints + ] + + return activity + + def _component_to_jsonld(self, comp: Any) -> Dict[str, Any]: + """Convert a Component node to JSON-LD.""" + comp_id = comp.get("id", "") + + return { + "@type": "schema:SoftwareSourceCode", + "@id": f"{ARKHAI_NAMESPACE}/component/{comp_id}", + "identifier": comp_id, + "name": comp.get("component_name"), + "schema:programmingLanguage": "Python", + "arkhai:componentType": comp.get("component_type"), + "arkhai:cacheKey": comp.get("cache_key"), + "arkhai:version": comp.get("version"), + } + + def to_json(self, data: Dict[str, Any], indent: int = 2) -> str: + """ + Serialize JSON-LD to string. + + Args: + data: JSON-LD document + indent: Indentation level + + Returns: + JSON string + """ + return json.dumps(data, indent=indent, default=str) + + def save_to_file( + self, data: Dict[str, Any], filepath: str, indent: int = 2 + ) -> None: + """ + Save JSON-LD to file. + + Args: + data: JSON-LD document + filepath: Output file path + indent: Indentation level + """ + with open(filepath, "w") as f: + json.dump(data, f, indent=indent, default=str) + + +# Convenience function for quick exports +def export_user_dataset( + graph_store: GraphStore, + username: str, + project: str, + output_path: Optional[str] = None, + **kwargs: Any, +) -> Dict[str, Any]: + """ + Quick export of a user's dataset to JSON-LD. + + Args: + graph_store: Neo4j GraphStore instance + username: Owner username + project: Project name + output_path: Optional file path to save + **kwargs: Additional arguments for export_dataset + + Returns: + JSON-LD dataset document + """ + exporter = JSONLDExporter(graph_store) + dataset = exporter.export_dataset(username, project, **kwargs) + + if output_path: + exporter.save_to_file(dataset, output_path) + + return dataset diff --git a/agentic_rag/mcp/__init__.py b/agentic_rag/mcp/__init__.py new file mode 100644 index 0000000..5db55d5 --- /dev/null +++ b/agentic_rag/mcp/__init__.py @@ -0,0 +1,5 @@ +"""MCP integration for agentic-rag.""" + +from .server import create_mcp_server + +__all__ = ["create_mcp_server"] diff --git a/agentic_rag/mcp/server.py b/agentic_rag/mcp/server.py new file mode 100644 index 0000000..8ee5dcf --- /dev/null +++ b/agentic_rag/mcp/server.py @@ -0,0 +1,685 @@ +"""MCP server exposing agentic-rag pipeline operations as tools.""" + +from __future__ import annotations + +import contextlib +import sys +from dataclasses import asdict, is_dataclass +from typing import Any, Callable, Dict, List, Optional, TypeVar + +from agentic_rag import Config, PipelineFactory +from agentic_rag.components import get_default_registry +from agentic_rag.pipeline import PipelineRunner +from agentic_rag.types import ( + PipelineSpec, + get_component_value, + list_available_components, +) +from agentic_rag.types.component_enums import get_evaluator_mode, requires_gold_standard + +try: + from mcp.server.fastmcp import FastMCP +except ImportError as exc: # pragma: no cover - import-time dependency guard + raise ImportError( + "MCP dependencies are missing. Install with `pip install mcp` " + "or use `poetry install` after updating dependencies." + ) from exc + + +EMBEDDING_FIELD_NAMES = { + "embedding", + "embeddings", + "query_embedding", + "sparse_embedding", + "sparse_embeddings", +} + +T = TypeVar("T") + + +def _to_jsonable(value: Any, depth: int = 0, max_depth: int = 5) -> Any: + """Convert nested objects to JSON-friendly data.""" + if depth > max_depth: + return repr(value) + + if value is None or isinstance(value, (str, int, float, bool)): + return value + + if is_dataclass(value): + return _to_jsonable(asdict(value), depth + 1, max_depth) + + if isinstance(value, dict): + return { + str(key): _to_jsonable(item, depth + 1, max_depth) + for key, item in value.items() + if str(key).lower() not in EMBEDDING_FIELD_NAMES + } + + if isinstance(value, (list, tuple, set)): + return [_to_jsonable(item, depth + 1, max_depth) for item in value] + + if hasattr(value, "to_dict") and callable(value.to_dict): + try: + return _to_jsonable(value.to_dict(), depth + 1, max_depth) + except Exception: + return repr(value) + + if hasattr(value, "__dict__"): + raw = vars(value) + safe = { + key: _to_jsonable(item, depth + 1, max_depth) + for key, item in raw.items() + if not key.startswith("_") and key.lower() not in EMBEDDING_FIELD_NAMES + } + if safe: + return safe + + return repr(value) + + +def _run_with_stdout_redirect(func: Callable[..., T], *args: Any, **kwargs: Any) -> T: + """ + Run tool logic while redirecting accidental stdout to stderr. + + MCP stdio transport requires stdout to contain only JSON-RPC payloads. + """ + with contextlib.redirect_stdout(sys.stderr): + return func(*args, **kwargs) + + +def _extract_replies(component_result: Dict[str, Any]) -> List[str]: + """Extract text replies from a generator-like component output.""" + replies = component_result.get("replies") + if isinstance(replies, list): + return [str(reply) for reply in replies] + if isinstance(replies, str): + return [replies] + return [] + + +def _extract_eval_metrics(component_result: Dict[str, Any]) -> Dict[str, Any]: + """Extract evaluation metrics from evaluator output.""" + eval_data = component_result.get("eval_data") + if isinstance(eval_data, dict): + metrics = eval_data.get("eval_metrics") + if isinstance(metrics, dict): + return metrics + return {} + + +def _summarize_retrieval_result(result: Dict[str, Any]) -> Dict[str, Any]: + """Return a compact retrieval summary with answer and metric results.""" + branch_summaries: Dict[str, Any] = {} + first_answer: Optional[str] = None + + branches = result.get("branches", {}) + if not isinstance(branches, dict): + return {"query": result.get("query"), "branches": {}} + + for branch_id, branch_result in branches.items(): + if isinstance(branch_result, dict) and "error" in branch_result: + branch_summaries[branch_id] = {"error": branch_result["error"]} + continue + + branch_answer: Optional[str] = None + branch_metrics: Dict[str, Any] = {} + document_count = 0 + + if isinstance(branch_result, dict): + for component_output in branch_result.values(): + if not isinstance(component_output, dict): + continue + + replies = _extract_replies(component_output) + if replies and branch_answer is None: + branch_answer = replies[0] + + metrics = _extract_eval_metrics(component_output) + if metrics: + branch_metrics.update(metrics) + + documents = component_output.get("documents") + if isinstance(documents, list): + document_count = max(document_count, len(documents)) + + if first_answer is None and branch_answer: + first_answer = branch_answer + + branch_summaries[branch_id] = { + "answer": branch_answer, + "metrics": branch_metrics, + "document_count": document_count, + } + + return { + "query": result.get("query"), + "answer": first_answer, + "branches": branch_summaries, + "branches_count": result.get("branches_count", len(branch_summaries)), + "total_documents": result.get("total_documents"), + } + + +class PipelineMCPService: + """Thin service layer wrapping Factory/Runner/GraphStore operations.""" + + def __init__(self, config: Optional[Config] = None) -> None: + self.config = config or Config() + self.registry = get_default_registry() + self._factory: Optional[PipelineFactory] = None + self._runner: Optional[PipelineRunner] = None + + @property + def factory(self) -> PipelineFactory: + """Lazily create the pipeline factory.""" + if self._factory is None: + self._factory = PipelineFactory(config=self.config) + return self._factory + + @property + def runner(self) -> PipelineRunner: + """Lazily create the pipeline runner.""" + if self._runner is None: + self._runner = PipelineRunner(config=self.config, enable_caching=False) + return self._runner + + @property + def graph_store(self) -> Any: + """Access graph store via runner to avoid eager Neo4j connection on startup.""" + return self.runner.graph_store + + def list_projects_and_pipelines(self, username: str) -> List[Dict[str, Any]]: + """List all projects and pipeline names for a user.""" + with self.graph_store.driver.session( + database=self.graph_store.database + ) as session: + query = """ + MATCH (u:User {username: $username})-[:OWNS]->(p:Project) + OPTIONAL MATCH (p)-[:FLOWS_TO]->(c:Component) + WHERE c.pipeline_name IS NOT NULL + WITH p, c.pipeline_name as pname, c.pipeline_type as ptype + WITH p, collect(DISTINCT CASE WHEN pname IS NOT NULL THEN {name: pname, type: ptype} END) as raw_pipelines + WITH p, [x IN raw_pipelines WHERE x IS NOT NULL] as pipelines + RETURN { + project: p.name, + pipelines: [x in pipelines | x.name], + details: pipelines + } as project_data + """ + results = session.run(query, username=username).data() + return [record["project_data"] for record in results] + + def get_component_index(self) -> Dict[str, Any]: + """ + Return a full component index linking enum specs to concrete registry entries. + + Useful for agents to discover exact component specs they can compose in + `create_pipelines` calls. + """ + available_specs = list_available_components() + entries: List[Dict[str, Any]] = [] + + for category, member_names in available_specs.items(): + for member_name in member_names: + spec_string = f"{category}.{member_name}" + component_name = get_component_value(spec_string) + registered_spec = self.registry.get_component_spec(component_name) + + entry: Dict[str, Any] = { + "spec": spec_string, + "category": category, + "member": member_name, + "component_name": component_name, + "registered": registered_spec is not None, + } + + if registered_spec is not None: + entry.update( + { + "component_type": registered_spec.component_type.value, + "pipeline_usage": registered_spec.pipeline_usage.value, + "haystack_class": registered_spec.haystack_class, + "default_config_keys": sorted( + list(registered_spec.default_config.keys()) + ), + } + ) + if category == "EVALUATOR": + entry.update( + { + "requires_ground_truth": requires_gold_standard( + component_name + ), + "evaluation_mode": get_evaluator_mode( + component_name + ).value, + } + ) + + entries.append(entry) + + return { + "total_specs": len(entries), + "total_registered_components": len(self.registry.list_components()), + "component_specs": entries, + "registered_components": sorted(self.registry.list_components()), + } + + def create_pipelines( + self, + username: str, + project: str, + pipeline_specs: List[List[Dict[str, str]]], + configs: Optional[List[Dict[str, Any]]] = None, + pipeline_types: Optional[List[str]] = None, + ) -> Dict[str, Any]: + """Create one or more pipeline graphs in Neo4j.""" + created: List[PipelineSpec] = self.factory.build_pipeline_graphs_from_specs( + pipeline_specs=pipeline_specs, + username=username, + project=project, + configs=configs, + pipeline_types=pipeline_types, + ) + return { + "count": len(created), + "pipelines": [ + { + "name": pipe.name, + "pipeline_type": pipe.pipeline_type.value, + "component_count": len(pipe.components), + "components": [comp.name for comp in pipe.components], + "indexing_pipelines": pipe.indexing_pipelines, + } + for pipe in created + ], + } + + def load_pipelines( + self, username: str, project: str, pipeline_names: List[str] + ) -> Dict[str, Any]: + """Load one or more pipelines for execution.""" + self.runner.load_pipelines( + pipeline_names=pipeline_names, + username=username, + project=project, + ) + return { + "loaded": pipeline_names, + "loaded_count": len(pipeline_names), + } + + def run_pipeline( + self, + username: str, + project: str, + pipeline_name: str, + pipeline_type: str, + inputs: Dict[str, Any], + auto_load: bool = True, + ) -> Dict[str, Any]: + """Run a pipeline with provided input payload.""" + if auto_load: + self.runner.load_pipelines( + pipeline_names=[pipeline_name], + username=username, + project=project, + ) + + result = self.runner.run( + pipeline_name=pipeline_name, + username=username, + type=pipeline_type, + project=project, + **inputs, + ) + + summarized_result: Any = result + if pipeline_type == "retrieval" and isinstance(result, dict): + summarized_result = _summarize_retrieval_result(result) + + return { + "pipeline_name": pipeline_name, + "pipeline_type": pipeline_type, + "project": project, + "result": _to_jsonable(summarized_result), + } + + def list_pipeline_components( + self, username: str, project: str, pipeline_name: str + ) -> Dict[str, Any]: + """List persisted component nodes for a pipeline.""" + components = self.graph_store.get_components_by_pipeline( + pipeline_name=pipeline_name, + username=username, + project=project, + ) + return { + "pipeline_name": pipeline_name, + "component_count": len(components), + "components": _to_jsonable(components), + } + + +def create_mcp_server(config: Optional[Config] = None) -> FastMCP: + """Create and configure the MCP server instance.""" + mcp = FastMCP("agentic-rag") + service = PipelineMCPService(config=config) + + @mcp.tool() # type: ignore[misc] + def health() -> Dict[str, Any]: + """Simple connectivity and config status check.""" + return _run_with_stdout_redirect( + lambda: { + "service": "agentic-rag-mcp", + "neo4j_configured": service.config.validate_neo4j(), + "openrouter_configured": service.config.validate_openrouter(), + } + ) + + @mcp.tool() # type: ignore[misc] + def list_available_component_specs() -> Dict[str, List[str]]: + """List component spec enums grouped by category.""" + return _run_with_stdout_redirect(list_available_components) + + @mcp.tool() # type: ignore[misc] + def list_registered_components() -> List[str]: + """List concrete registered component names from registry.""" + return _run_with_stdout_redirect(service.registry.list_components) + + @mcp.tool() # type: ignore[misc] + def get_component_index() -> Dict[str, Any]: + """ + Return exact component index for MCP agents. + + Includes: + - enum spec strings like `CONVERTER.PDF` + - concrete registry names like `pdf_converter` + - registration status and metadata + + Example call: + { + "tool": "get_component_index", + "arguments": {} + } + """ + return _run_with_stdout_redirect(service.get_component_index) + + @mcp.tool() # type: ignore[misc] + def list_projects(username: str) -> List[Dict[str, Any]]: + """ + List projects and pipelines available for a username. + + Example call: + { + "tool": "list_projects", + "arguments": { + "username": "alice" + } + } + """ + return _run_with_stdout_redirect( + service.list_projects_and_pipelines, + username=username, + ) + + @mcp.tool() # type: ignore[misc] + def list_pipeline_components( + username: str, + project: str, + pipeline_name: str, + ) -> Dict[str, Any]: + """ + List components in a specific pipeline. + + Example call: + { + "tool": "list_pipeline_components", + "arguments": { + "username": "alice", + "project": "demo_rag_app", + "pipeline_name": "paper_index" + } + } + """ + return _run_with_stdout_redirect( + service.list_pipeline_components, + username=username, + project=project, + pipeline_name=pipeline_name, + ) + + @mcp.tool() # type: ignore[misc] + def create_pipelines( + username: str, + project: str, + pipeline_specs: List[List[Dict[str, str]]], + configs: Optional[List[Dict[str, Any]]] = None, + pipeline_types: Optional[List[str]] = None, + ) -> Dict[str, Any]: + """ + Create pipeline graphs. + + - `pipeline_specs` is a list of pipelines. + - Each pipeline is a list of component dicts with `type`. + - Example component: `{"type": "CONVERTER.PDF"}`. + - `pipeline_types` must contain one string per pipeline, usually + `"indexing"` or `"retrieval"`. + - For indexing pipelines, put component specs inside `pipeline_specs`, + not inside `pipeline_types`. + + Example indexing pipeline call: + { + "tool": "create_pipelines", + "arguments": { + "username": "alice", + "project": "demo_rag_app", + "pipeline_specs": [[ + {"type": "CONVERTER.MARKDOWN"}, + {"type": "CHUNKER.MARKDOWN_AWARE"}, + {"type": "EMBEDDER.SENTENCE_TRANSFORMERS_DOC"}, + {"type": "WRITER.CHROMA_DOCUMENT_WRITER"} + ]], + "configs": [{ + "_pipeline_name": "paper_index", + "markdown_aware_chunker": { + "chunk_size": 800, + "chunk_overlap": 100 + } + }], + "pipeline_types": ["indexing"] + } + } + + Example retrieval pipeline call: + { + "tool": "create_pipelines", + "arguments": { + "username": "alice", + "project": "demo_rag_app", + "pipeline_specs": [[ + {"type": "INDEX"}, + {"type": "GENERATOR.PROMPT_BUILDER"}, + {"type": "GENERATOR.OPENROUTER"} + ]], + "configs": [{ + "_pipeline_name": "paper_retrieval", + "_indexing_pipelines": ["paper_index"], + "prompt_builder": { + "template": "Answer the question based only on the provided documents.\n\nQuestion: {{query}}\n\nDocuments:\n{% for doc in documents %}\n{{ doc.content }}\n---\n{% endfor %}\n\nAnswer:" + } + }], + "pipeline_types": ["retrieval"] + } + } + """ + return _run_with_stdout_redirect( + service.create_pipelines, + username=username, + project=project, + pipeline_specs=pipeline_specs, + configs=configs, + pipeline_types=pipeline_types, + ) + + @mcp.tool() # type: ignore[misc] + def load_pipelines( + username: str, + project: str, + pipeline_names: List[str], + ) -> Dict[str, Any]: + """ + Load one or more pipeline graphs into the runtime runner. + + Example call: + { + "tool": "load_pipelines", + "arguments": { + "username": "alice", + "project": "demo_rag_app", + "pipeline_names": ["paper_index", "paper_retrieval"] + } + } + """ + return _run_with_stdout_redirect( + service.load_pipelines, + username=username, + project=project, + pipeline_names=pipeline_names, + ) + + @mcp.tool() # type: ignore[misc] + def run_indexing_pipeline( + username: str, + project: str, + pipeline_name: str, + data_path: str, + auto_load: bool = True, + ) -> Dict[str, Any]: + """ + Run an indexing pipeline. Input requires `data_path`. + + Example call: + { + "tool": "run_indexing_pipeline", + "arguments": { + "username": "alice", + "project": "demo_rag_app", + "pipeline_name": "paper_index", + "data_path": "/absolute/path/to/papers", + "auto_load": true + } + } + """ + return _run_with_stdout_redirect( + service.run_pipeline, + username=username, + project=project, + pipeline_name=pipeline_name, + pipeline_type="indexing", + inputs={"data_path": data_path}, + auto_load=auto_load, + ) + + @mcp.tool() # type: ignore[misc] + def run_retrieval_pipeline( + username: str, + project: str, + pipeline_name: str, + query: str, + ground_truth_answer: Optional[str] = None, + relevant_doc_ids: Optional[List[str]] = None, + auto_load: bool = True, + ) -> Dict[str, Any]: + """ + Run a retrieval pipeline with a query. + + Example call: + { + "tool": "run_retrieval_pipeline", + "arguments": { + "username": "alice", + "project": "demo_rag_app", + "pipeline_name": "paper_retrieval", + "query": "What chunking strategy is recommended?", + "auto_load": true + } + } + """ + return _run_with_stdout_redirect( + service.run_pipeline, + username=username, + project=project, + pipeline_name=pipeline_name, + pipeline_type="retrieval", + inputs={ + "query": query, + "ground_truth_answer": ground_truth_answer, + "relevant_doc_ids": relevant_doc_ids or [], + }, + auto_load=auto_load, + ) + + @mcp.tool() # type: ignore[misc] + def run_pipeline( + username: str, + project: str, + pipeline_name: str, + pipeline_type: str, + inputs: Dict[str, Any], + auto_load: bool = True, + ) -> Dict[str, Any]: + """ + Generic execution tool for indexing or retrieval pipelines. + + Example indexing call: + { + "tool": "run_pipeline", + "arguments": { + "username": "alice", + "project": "demo_rag_app", + "pipeline_name": "paper_index", + "pipeline_type": "indexing", + "inputs": { + "data_path": "/absolute/path/to/papers" + }, + "auto_load": true + } + } + + Example retrieval call: + { + "tool": "run_pipeline", + "arguments": { + "username": "alice", + "project": "demo_rag_app", + "pipeline_name": "paper_retrieval", + "pipeline_type": "retrieval", + "inputs": { + "query": "Summarize the paper's main contribution" + }, + "auto_load": true + } + } + """ + return _run_with_stdout_redirect( + service.run_pipeline, + username=username, + project=project, + pipeline_name=pipeline_name, + pipeline_type=pipeline_type, + inputs=inputs, + auto_load=auto_load, + ) + + return mcp + + +def main() -> None: + """Run MCP server over stdio transport.""" + server = create_mcp_server() + server.run() + + +if __name__ == "__main__": + main() diff --git a/agentic_rag/pipeline/factory.py b/agentic_rag/pipeline/factory.py index 3973105..0ddcd7a 100644 --- a/agentic_rag/pipeline/factory.py +++ b/agentic_rag/pipeline/factory.py @@ -418,41 +418,48 @@ async def _build_indexing_pipeline_async( self.logger.info( f"Building indexing pipeline (async): {pipeline_name} for user: {username}, project: {project}" ) + self.logger.info(f"Component specs: {component_specs}") # Parse component specifications and validate component_specs_list = [] for spec_item in component_specs: - component_name = self._parse_component_spec(spec_item) - - spec = self.registry.get_component_spec(component_name) - if spec is None: - raise ValueError(f"Unknown component: {component_name}") - - # Configure the spec directly with user config - user_config = config.get(component_name, {}) - - # Auto-generate root_dir for document writers if not provided - if ( - component_name in ["chroma_document_writer", "qdrant_document_writer"] - and "root_dir" not in user_config - ): - user_config = user_config.copy() # Don't modify original config - # Use agentic_root_dir from config if available - root_dir = self.config.agentic_root_dir if self.config else "./data" - user_config["root_dir"] = ( - f"{root_dir}/{username}/{project}/{pipeline_name}" - ) - self.logger.debug( - f"Auto-generated root_dir for {component_name}: {user_config['root_dir']}" - ) + try: + component_name = self._parse_component_spec(spec_item) + self.logger.info(f"Processing component: {component_name}") + + spec = self.registry.get_component_spec(component_name) + if spec is None: + raise ValueError(f"Unknown component: {component_name}") + + # Configure the spec directly with user config + user_config = config.get(component_name, {}) + + # Auto-generate root_dir for document writers if not provided + if ( + component_name in ["chroma_document_writer", "qdrant_document_writer"] + and "root_dir" not in user_config + ): + user_config = user_config.copy() # Don't modify original config + # Use agentic_root_dir from config if available + root_dir = self.config.agentic_root_dir if self.config else "./data" + user_config["root_dir"] = ( + f"{root_dir}/{username}/{project}/{pipeline_name}" + ) + self.logger.info( + f"Auto-generated root_dir for {component_name}: {user_config['root_dir']}" + ) - configured_spec = spec.configure(user_config) + configured_spec = spec.configure(user_config) - # Store the original full type string - configured_spec.full_type = spec_item.get("type", "") + # Store the original full type string + configured_spec.full_type = spec_item.get("type", "") - component_specs_list.append(configured_spec) + component_specs_list.append(configured_spec) + except Exception as e: + self.logger.error(f"Error processing component spec {spec_item}: {e}") + raise + self.logger.info("Finished processing components, creating PipelineSpec") # Create pipeline specification pipeline_spec = PipelineSpec( name=pipeline_name, diff --git a/agentic_rag/pipeline/runner.py b/agentic_rag/pipeline/runner.py index 43d5c7e..a7c4711 100644 --- a/agentic_rag/pipeline/runner.py +++ b/agentic_rag/pipeline/runner.py @@ -3,18 +3,34 @@ ARCHITECTURE: - Factory: Build pipelines once and store in Neo4j (creation time) - Runner: Load pipelines from Neo4j and execute (runtime) + +FAIR Compliance: +- Creates RunNode for each pipeline execution (provenance) +- Links all generated DataPieces to the Run via GENERATED_BY +- Stores run metadata: pipeline_name, started_at, finished_at, success """ +import hashlib import time +import uuid +from datetime import datetime from typing import Any, Dict, List, Optional from ..components import GraphStore from ..config import Config, get_config from ..types import PipelineUsage +from ..types.node_types import ARKHAI_NAMESPACE from ..utils.logger import configure_haystack_logging, get_logger from ..utils.metrics import MetricsCollector +def generate_run_id() -> str: + """Generate a unique run ID.""" + unique = f"{datetime.utcnow().isoformat()}_{uuid.uuid4().hex[:8]}" + hash_obj = hashlib.sha256(unique.encode()) + return f"run_{hash_obj.hexdigest()[:12]}" + + class PipelineRunner: """Executes pipelines with input data (Singleton).""" @@ -75,6 +91,106 @@ def __init__( self._initialized = True self.logger.info("PipelineRunner initialized (singleton)") + def _set_run_id_on_components(self, pipeline_name: str, run_id: str) -> None: + """ + Set run_id on all GatedComponents in a pipeline for FAIR provenance. + + Args: + pipeline_name: Name of the pipeline + run_id: Run ID to set on all components + """ + from ..components.gates import GatedComponent + + if pipeline_name not in self._haystack_components_by_pipeline: + return + + components = self._haystack_components_by_pipeline[pipeline_name] + for comp_id, component in components.items(): + if isinstance(component, GatedComponent): + component.set_run_id(run_id) + self.logger.debug(f"Set run_id={run_id} on component {comp_id}") + + def _create_run_node( + self, + run_id: str, + pipeline_name: str, + username: str, + project: str, + pipeline_type: str, + ) -> None: + """ + Create a RunNode in Neo4j for provenance tracking. + + Args: + run_id: Unique run identifier + pipeline_name: Name of pipeline being run + username: User who initiated the run + project: Project name + pipeline_type: "indexing" or "retrieval" + """ + started_at = datetime.utcnow().isoformat() + uri = f"{ARKHAI_NAMESPACE}/run/{run_id}" + + self.graph_store.store_run_node( + run_id=run_id, + pipeline_name=pipeline_name, + username=username, + project=project, + started_at=started_at, + uri=uri, + ) + self.logger.info(f"Created RunNode: {run_id} for pipeline {pipeline_name}") + + async def _create_run_node_async( + self, + run_id: str, + pipeline_name: str, + username: str, + project: str, + pipeline_type: str, + ) -> None: + """Async version of _create_run_node.""" + started_at = datetime.utcnow().isoformat() + uri = f"{ARKHAI_NAMESPACE}/run/{run_id}" + + await self.graph_store.store_run_node_async( + run_id=run_id, + pipeline_name=pipeline_name, + username=username, + project=project, + started_at=started_at, + uri=uri, + ) + self.logger.info( + f"Created RunNode (async): {run_id} for pipeline {pipeline_name}" + ) + + def _finalize_run_node( + self, run_id: str, success: bool, error: Optional[str] = None + ) -> None: + """Update RunNode with completion information.""" + finished_at = datetime.utcnow().isoformat() + self.graph_store.update_run_finished( + run_id=run_id, + finished_at=finished_at, + success=success, + error=error, + ) + self.logger.info(f"Finalized RunNode: {run_id} (success={success})") + + async def _finalize_run_node_async( + self, run_id: str, success: bool, error: Optional[str] = None + ) -> None: + """Async version of _finalize_run_node.""" + finished_at = datetime.utcnow().isoformat() + await self.graph_store.update_run_finished_async( + run_id=run_id, + finished_at=finished_at, + success=success, + error=error, + ) + self.logger.info(f"Finalized RunNode (async): {run_id} (success={success})") + @classmethod def reset_instance(cls) -> None: """Reset singleton instance (useful for testing).""" @@ -463,10 +579,11 @@ def build_haystack_components_from_graph( else: configured_spec = spec_copy - # Instantiate the actual Haystack component using registry cache - # This will return a cached instance if available (LRU) - # For writers/retrievers, it will handle document store creation - haystack_component = registry.get_component_instance(configured_spec) + # Instantiate a FRESH Haystack component (no caching) + # Haystack doesn't allow sharing component instances between pipelines + from ..components.registry import create_haystack_component + + haystack_component = create_haystack_component(configured_spec) # Optionally wrap with GatedComponent for caching if self.enable_caching: @@ -486,7 +603,7 @@ def build_haystack_components_from_graph( graph_store=self.graph_store, username=username, cache_key=cache_key, # Use pipeline-agnostic cache key - retrieve_from_ipfs=True, + retrieve_from_storage=True, ) else: self.logger.debug( @@ -1037,6 +1154,10 @@ def _run_indexing_pipeline( Returns: Indexing results + + FAIR Compliance: + - Creates a RunNode for provenance tracking + - Links all generated DataPieces to the Run """ from pathlib import Path @@ -1047,6 +1168,9 @@ def _run_indexing_pipeline( success = False error_msg = None + # FAIR: Generate run ID and create RunNode + run_id = generate_run_id() + try: # Get the pipeline if pipeline_name not in self._haystack_pipelines: @@ -1082,7 +1206,21 @@ def _run_indexing_pipeline( f"No supported files found in {data_path_obj}" ) - self.logger.info(f"Running indexing pipeline: {pipeline_name}") + # FAIR: Create RunNode before execution + if self.enable_caching: + self._create_run_node( + run_id=run_id, + pipeline_name=pipeline_name, + username=username, + project=project, + pipeline_type="indexing", + ) + # Set run_id on all GatedComponents + self._set_run_id_on_components(pipeline_name, run_id) + + self.logger.info( + f"Running indexing pipeline: {pipeline_name} (run_id={run_id})" + ) self.logger.info(f"Processing {len(input_files)} file(s)") # Run the pipeline with the file sources @@ -1108,15 +1246,24 @@ def _run_indexing_pipeline( "type": "indexing", "files_processed": len(input_files), "data_path": str(data_path_obj), + "run_id": run_id, }, ) + # FAIR: Finalize RunNode + if self.enable_caching: + self._finalize_run_node(run_id, success=True) + return result except Exception as e: error_msg = str(e) success = False + # FAIR: Finalize RunNode with error + if self.enable_caching: + self._finalize_run_node(run_id, success=False, error=error_msg) + # Log failed metrics end_time = time.time() component_count = len( @@ -1129,7 +1276,7 @@ def _run_indexing_pipeline( total_components=component_count, success=success, error=error_msg, - metadata={"type": "indexing"}, + metadata={"type": "indexing", "run_id": run_id}, ) raise @@ -1280,6 +1427,10 @@ async def _run_indexing_pipeline_async( Async version of _run_indexing_pipeline. Executes indexing pipeline using AsyncPipeline.run_async(). + + FAIR Compliance: + - Creates a RunNode for provenance tracking + - Links all generated DataPieces to the Run """ from pathlib import Path @@ -1290,6 +1441,9 @@ async def _run_indexing_pipeline_async( success = False error_msg = None + # FAIR: Generate run ID + run_id = generate_run_id() + try: # Get the pipeline if pipeline_name not in self._haystack_pipelines: @@ -1325,7 +1479,21 @@ async def _run_indexing_pipeline_async( f"No supported files found in {data_path_obj}" ) - self.logger.info(f"Running indexing pipeline (async): {pipeline_name}") + # FAIR: Create RunNode before execution + if self.enable_caching: + await self._create_run_node_async( + run_id=run_id, + pipeline_name=pipeline_name, + username=username, + project=project, + pipeline_type="indexing", + ) + # Set run_id on all GatedComponents + self._set_run_id_on_components(pipeline_name, run_id) + + self.logger.info( + f"Running indexing pipeline (async): {pipeline_name} (run_id={run_id})" + ) self.logger.info(f"Processing {len(input_files)} file(s)") # Run the pipeline with the file sources @@ -1353,15 +1521,26 @@ async def _run_indexing_pipeline_async( "files_processed": len(input_files), "data_path": str(data_path_obj), "mode": "async", + "run_id": run_id, }, ) + # FAIR: Finalize RunNode + if self.enable_caching: + await self._finalize_run_node_async(run_id, success=True) + return result except Exception as e: error_msg = str(e) success = False + # FAIR: Finalize RunNode with error + if self.enable_caching: + await self._finalize_run_node_async( + run_id, success=False, error=error_msg + ) + # Log failed metrics end_time = time.time() component_count = len( @@ -1374,7 +1553,7 @@ async def _run_indexing_pipeline_async( total_components=component_count, success=success, error=error_msg, - metadata={"type": "indexing", "mode": "async"}, + metadata={"type": "indexing", "mode": "async", "run_id": run_id}, ) raise diff --git a/agentic_rag/pipeline/storage.py b/agentic_rag/pipeline/storage.py index 4101fac..2ccd92f 100644 --- a/agentic_rag/pipeline/storage.py +++ b/agentic_rag/pipeline/storage.py @@ -1,5 +1,6 @@ """Graph storage for creating and managing graph representations of pipelines.""" +import sys from typing import Any, Dict, List, Optional, Tuple from haystack import Pipeline @@ -99,7 +100,7 @@ def create_pipeline_graph( self.logger.info( f"Creating pipeline graph for user '{username}', project '{project}', pipeline '{spec.name}'" ) - user_node = UserNode(username=username, display_name=username.title()) + user_node = UserNode(username=username, display_name=username) user_dict = user_node.to_dict() self.graph_store.add_nodes_batch([user_dict], "User") @@ -209,7 +210,7 @@ async def create_pipeline_graph_async( self.logger.info( f"Creating pipeline graph (async) for user '{username}', project '{project}', pipeline '{spec.name}'" ) - user_node = UserNode(username=username, display_name=username.title()) + user_node = UserNode(username=username, display_name=username) user_dict = user_node.to_dict() await self.graph_store.add_nodes_batch_async([user_dict], "User") @@ -343,7 +344,10 @@ def build_haystack_pipeline(self, spec: PipelineSpec) -> Any: try: pipeline.connect(source, target) except Exception as e: - print(f"Warning: Could not connect {source} -> {target}: {e}") + print( + f"Warning: Could not connect {source} -> {target}: {e}", + file=sys.stderr, + ) return pipeline @@ -383,7 +387,10 @@ def load_pipeline_by_hashes( # 2. Fetch components for each pipeline hash separately for pipeline_hash in pipeline_hashes: - print(f"\nFetching components for pipeline: {pipeline_hash}") + print( + f"\nFetching components for pipeline: {pipeline_hash}", + file=sys.stderr, + ) # Call Neo4j for this specific pipeline hash (single hash method) component_data_list = self.graph_store.get_pipeline_components_by_hash( @@ -392,13 +399,13 @@ def load_pipeline_by_hashes( project, # Single pipeline hash with project filter ) - print(f" Found {len(component_data_list)} components") + print(f" Found {len(component_data_list)} components", file=sys.stderr) if component_data_list: pipelines_data[pipeline_hash] = component_data_list # Print details for this pipeline - print(component_data_list) + print(component_data_list, file=sys.stderr) return pipelines_data @@ -426,7 +433,10 @@ async def load_pipeline_by_hashes_async( # 2. Fetch components for each pipeline hash separately for pipeline_hash in pipeline_hashes: - print(f"\nFetching components for pipeline (async): {pipeline_hash}") + print( + f"\nFetching components for pipeline (async): {pipeline_hash}", + file=sys.stderr, + ) # Call Neo4j for this specific pipeline hash (async) component_data_list = ( @@ -437,7 +447,7 @@ async def load_pipeline_by_hashes_async( ) ) - print(f" Found {len(component_data_list)} components") + print(f" Found {len(component_data_list)} components", file=sys.stderr) if component_data_list: pipelines_data[pipeline_hash] = component_data_list diff --git a/agentic_rag/types/__init__.py b/agentic_rag/types/__init__.py index fccb518..58159b4 100644 --- a/agentic_rag/types/__init__.py +++ b/agentic_rag/types/__init__.py @@ -26,14 +26,18 @@ get_relationship_name, get_safe_relationship_name, ) -from .node_types import ( +from .node_types import ( # Core node types; FAIR compliance types; FAIR compliance utilities + ARKHAI_NAMESPACE, ComponentNode, ComponentRelationship, + ContentType, DataPiece, ProcessedByRelationship, ProjectNode, + RunNode, TransformedByRelationship, UserNode, + generate_uri, ) from .pipeline_spec import PipelineSpec, PipelineType @@ -54,6 +58,11 @@ "DataPiece", "TransformedByRelationship", "ProcessedByRelationship", + # FAIR compliance types + "ContentType", + "RunNode", + "ARKHAI_NAMESPACE", + "generate_uri", # Component enums "CONVERTER", "CHUNKER", diff --git a/agentic_rag/types/component_spec.py b/agentic_rag/types/component_spec.py index e4c9a31..51babb0 100644 --- a/agentic_rag/types/component_spec.py +++ b/agentic_rag/types/component_spec.py @@ -27,6 +27,10 @@ class ComponentSpec: # Store the full type string (e.g., "EMBEDDER.SENTENCE_TRANSFORMERS_DOC") full_type: str = "" + # FAIR Compliance: What content type this component produces + # e.g., CONVERTER produces DOCUMENT, CHUNKER produces CHUNK, EMBEDDER produces EMBEDDING + produces_content_type: Optional[str] = None # String value of ContentType enum + def is_compatible_input(self, data_type: DataType) -> bool: """Check if a data type is compatible with this component's inputs.""" return data_type in self.input_types @@ -139,9 +143,21 @@ def create_haystack_component(spec: ComponentSpec) -> Any: # Get the final config from ComponentSpec config = spec.get_config().copy() # Copy to avoid modifying original + # Check if this is a Chroma component by name + is_chroma_component = "chroma" in spec.name.lower() + + # Check if this is a Qdrant component by name + is_qdrant_component = "qdrant" in spec.name.lower() + + # Check config for Qdrant-specific params (for remote Qdrant) + has_qdrant_config = any( + k in config + for k in ["qdrant_host", "qdrant_port", "qdrant_collection", "embedding_dim"] + ) + # Special handling for Chroma components - inject document store if "ChromaEmbeddingRetriever" in spec.haystack_class or ( - "DocumentWriter" in spec.haystack_class and "Chroma" in spec.haystack_class + "DocumentWriter" in spec.haystack_class and is_chroma_component ): # Get root directory from config or use current directory config.pop("model", None) # Remove model if present for these components @@ -150,6 +166,16 @@ def create_haystack_component(spec: ComponentSpec) -> Any: # Check for remote ChromaDB configuration (for async support) chroma_host = config.pop("chroma_host", None) chroma_port = config.pop("chroma_port", None) + + # Fallback to env vars if not in config (Safety net) + if not chroma_host: + chroma_host = os.getenv("CHROMA_HOST") + if not chroma_port: + chroma_port = os.getenv("CHROMA_PORT") + + if chroma_port: + chroma_port = int(chroma_port) + chroma_collection = config.pop("chroma_collection", None) document_store = _create_chroma_document_store( @@ -161,44 +187,32 @@ def create_haystack_component(spec: ComponentSpec) -> Any: config["document_store"] = document_store # Special handling for Qdrant components - inject document store - elif ( - "QdrantEmbeddingRetriever" in spec.haystack_class - or "DocumentWriter" in spec.haystack_class + elif "QdrantEmbeddingRetriever" in spec.haystack_class or ( + "DocumentWriter" in spec.haystack_class + and (is_qdrant_component or has_qdrant_config) ): - # Check config for Qdrant-specific params - has_qdrant_config = any( - k in config - for k in [ - "qdrant_host", - "qdrant_port", - "qdrant_collection", - "embedding_dim", - ] + # Get root directory from config or use current directory + config.pop("model", None) # Remove model if present for these components + root_dir = config.pop("root_dir", ".") + + # Check for remote Qdrant configuration + qdrant_host = config.pop("qdrant_host", None) + qdrant_port = config.pop("qdrant_port", None) + qdrant_collection = config.pop("qdrant_collection", None) + embedding_dim = config.pop("embedding_dim", 768) # Default to 768 + + document_store = _create_qdrant_document_store( + root_dir=root_dir, + host=qdrant_host, + port=qdrant_port, + collection_name=qdrant_collection, + embedding_dim=embedding_dim, ) + config["document_store"] = document_store - if has_qdrant_config: - # Get root directory from config or use current directory - config.pop("model", None) # Remove model if present for these components - root_dir = config.pop("root_dir", ".") - - # Check for remote Qdrant configuration - qdrant_host = config.pop("qdrant_host", None) - qdrant_port = config.pop("qdrant_port", None) - qdrant_collection = config.pop("qdrant_collection", None) - embedding_dim = config.pop("embedding_dim", 768) # Default to 768 - - document_store = _create_qdrant_document_store( - root_dir=root_dir, - host=qdrant_host, - port=qdrant_port, - collection_name=qdrant_collection, - embedding_dim=embedding_dim, - ) - config["document_store"] = document_store - elif "DocumentWriter" in spec.haystack_class: - # DocumentWriter without Qdrant config - remove root_dir if present - # (This handles the case where root_dir was auto-generated but no Qdrant config) - config.pop("root_dir", None) + elif "DocumentWriter" in spec.haystack_class: + # DocumentWriter without specific store config - remove root_dir if present + config.pop("root_dir", None) # Dynamic import and instantiation module_path, class_name = spec.haystack_class.rsplit(".", 1) diff --git a/agentic_rag/types/graph_relationships.py b/agentic_rag/types/graph_relationships.py index 8a63fc6..96e5e45 100644 --- a/agentic_rag/types/graph_relationships.py +++ b/agentic_rag/types/graph_relationships.py @@ -21,6 +21,10 @@ class GraphRelationship(Enum): TRANSFORMED_BY = "TRANSFORMED_BY" # DataPiece transformed to another DataPiece PROCESSED_BY = "PROCESSED_BY" # DataPiece processed by Component + # Provenance relationships (FAIR compliance - PROV-O) + GENERATED_BY = "GENERATED_BY" # DataPiece generated by Run (prov:wasGeneratedBy) + GENERATED = "GENERATED" # Run generated DataPiece (prov:generated) + def get_relationship_name(relationship: GraphRelationship) -> str: """Get the string name of a relationship for Neo4j queries.""" diff --git a/agentic_rag/types/node_types.py b/agentic_rag/types/node_types.py index 514cf81..4e37f01 100644 --- a/agentic_rag/types/node_types.py +++ b/agentic_rag/types/node_types.py @@ -2,7 +2,54 @@ from dataclasses import dataclass from datetime import datetime -from typing import Any, Dict, Optional, Tuple +from enum import Enum +from typing import Any, Dict, List, Optional, Tuple + +# ============================================================================= +# FAIR Compliance: URI Namespace and Content Types +# ============================================================================= + +# Persistent URI namespace for FAIR compliance (F1-F3) +ARKHAI_NAMESPACE = "https://w3id.org/arkhai" + + +class ContentType(Enum): + """ + Semantic content type for FAIR compliance. + + Maps to JSON-LD @type during export. Used to distinguish between + different kinds of data flowing through pipelines. + """ + + DOCUMENT = "Document" # Original/converted documents + CHUNK = "Chunk" # Text chunks from chunking + EMBEDDING = "Embedding" # Vector embeddings + DATA_NODE = "DataNode" # Generic data (fallback) + + +def generate_uri(content_type: ContentType, identifier: str) -> str: + """ + Generate a persistent URI for FAIR compliance (F1-F3). + + Args: + content_type: The semantic type of the content + identifier: Unique identifier (e.g., fingerprint) + + Returns: + URI like https://w3id.org/arkhai/doc/fp_abc123 + + Example: + >>> generate_uri(ContentType.DOCUMENT, "fp_abc123") + 'https://w3id.org/arkhai/doc/fp_abc123' + """ + type_to_path = { + ContentType.DOCUMENT: "doc", + ContentType.CHUNK: "chunk", + ContentType.EMBEDDING: "embedding", + ContentType.DATA_NODE: "data", + } + path = type_to_path.get(content_type, "data") + return f"{ARKHAI_NAMESPACE}/{path}/{identifier}" @dataclass @@ -191,6 +238,116 @@ def from_dict(cls, data: Dict[str, Any]) -> "ProjectNode": ) +# ============================================================================= +# FAIR Compliance: Run/Provenance Tracking +# ============================================================================= + + +@dataclass +class RunNode: + """ + Represents a pipeline execution run for provenance tracking. + + Links all data transformations from a single pipeline execution. + Required for FAIR compliance (PROV-O: prov:Activity). + + Graph pattern: + Run -[:GENERATED]-> DataPiece + DataPiece -[:GENERATED_BY]-> Run + + Example: + >>> run = RunNode( + ... id="run_abc123", + ... pipeline_name="indexing_pipeline", + ... username="alice", + ... ) + >>> run.uri + 'https://w3id.org/arkhai/run/run_abc123' + """ + + # Identity (required) + id: str # Run identifier (e.g., "run_abc123") + pipeline_name: str # Name of the pipeline that was executed + username: str # Who initiated the run + + # Pipeline info (optional) + pipeline_version: Optional[str] = None # Version of the pipeline + project: str = "default" # Project this run belongs to + + # Provenance metadata (optional) + git_commit: Optional[str] = None # Git commit hash at run time + model_names: Optional[List[str]] = None # Models used in this run + config_hash: Optional[str] = None # Hash of pipeline configuration + + # Timing (optional) + started_at: Optional[datetime] = None # When run started + finished_at: Optional[datetime] = None # When run completed + + # FAIR compliance + uri: Optional[str] = None # Persistent URI + + def __post_init__(self) -> None: + """Generate URI if not provided.""" + if self.uri is None: + self.uri = f"{ARKHAI_NAMESPACE}/run/{self.id}" + + def to_neo4j_properties(self) -> Dict[str, Any]: + """Convert to Neo4j node properties.""" + props: Dict[str, Any] = { + "id": self.id, + "uri": self.uri, + "pipeline_name": self.pipeline_name, + "username": self.username, + "project": self.project, + } + + if self.pipeline_version: + props["pipeline_version"] = self.pipeline_version + if self.git_commit: + props["git_commit"] = self.git_commit + if self.model_names: + # Neo4j supports lists natively + props["model_names"] = self.model_names + if self.config_hash: + props["config_hash"] = self.config_hash + if self.started_at: + props["started_at"] = self.started_at.isoformat() + if self.finished_at: + props["finished_at"] = self.finished_at.isoformat() + + return props + + @classmethod + def from_neo4j_node(cls, node: dict) -> "RunNode": + """Create RunNode from Neo4j node properties.""" + started_at = None + if node.get("started_at"): + started_at = datetime.fromisoformat(node["started_at"]) + + finished_at = None + if node.get("finished_at"): + finished_at = datetime.fromisoformat(node["finished_at"]) + + return cls( + id=node["id"], + pipeline_name=node["pipeline_name"], + username=node.get("username", ""), + pipeline_version=node.get("pipeline_version"), + project=node.get("project", "default"), + git_commit=node.get("git_commit"), + model_names=node.get("model_names"), + config_hash=node.get("config_hash"), + started_at=started_at, + finished_at=finished_at, + uri=node.get("uri"), + ) + + +# ============================================================================= +# Data Content Nodes +# ============================================================================= + + @dataclass class DataPiece: """ @@ -198,34 +355,77 @@ class DataPiece: Used by InGate/OutGate to track data transformations and enable caching. Each unique piece of content gets one DataPiece node (deduplicated by fingerprint). - Content is stored on IPFS, only hash is stored in Neo4j. + Content is stored on Akave (S3-compatible), object key is stored in Neo4j. + + Note: The field `ipfs_hash` is kept for backward compatibility with existing Neo4j + schemas but now stores Akave object keys instead of IPFS CIDs. + + FAIR Compliance: + - `content_type`: Semantic type (Document, Chunk, Embedding, DataNode) + - `generated_by`: Run ID linking to provenance (PROV-O: prov:wasGeneratedBy) + - `uri`: Persistent URI for identification (F1-F3) + + Example: + >>> piece = DataPiece( + ... fingerprint="fp_abc123", + ... ipfs_hash="ak_abc123...", # Akave object key + ... data_type="Document", + ... username="alice", + ... content_type=ContentType.DOCUMENT, + ... generated_by="run_xyz789", + ... ) + >>> piece.uri + 'https://w3id.org/arkhai/doc/fp_abc123' """ # Identity (required) fingerprint: str # SHA256 hash of content (PRIMARY KEY) # Content storage (required) - ipfs_hash: str # IPFS hash where actual content is stored - data_type: str # "Document", "ByteStream", "List[Document]", etc. + ipfs_hash: str # Akave object key (field name kept for Neo4j schema compatibility) + data_type: str # Python type: "Document", "ByteStream", "List[Document]", etc. # Authorship (required) username: str # Who created/owns this data + # FAIR compliance fields (optional, for semantic interoperability) + content_type: Optional[ContentType] = ( + None # Semantic type: Document | Chunk | Embedding | DataNode + ) + generated_by: Optional[str] = None # Run ID for provenance tracking + uri: Optional[str] = ( + None # Persistent URI: https://w3id.org/arkhai/doc/{fingerprint} + ) + # Metadata (optional) content_preview: Optional[str] = None # First 200 chars for quick viewing size_bytes: Optional[int] = None # Size of content created_at: Optional[datetime] = None # When first created source: Optional[str] = None # Original source (file path, URL, etc.) + def __post_init__(self) -> None: + """Generate URI if content_type is set but URI is not.""" + if self.content_type is not None and self.uri is None: + self.uri = generate_uri(self.content_type, self.fingerprint) + def to_neo4j_properties(self) -> Dict[str, Any]: """Convert to Neo4j node properties.""" props: Dict[str, Any] = { "fingerprint": self.fingerprint, "ipfs_hash": self.ipfs_hash, - "type": self.data_type, + "type": self.data_type, # Keep 'type' for backward compatibility "username": self.username, } + # FAIR compliance fields + if self.content_type is not None: + props["content_type"] = self.content_type.value + if self.generated_by: + props["generated_by"] = self.generated_by + if self.uri: + props["uri"] = self.uri + + # Metadata if self.content_preview: props["content_preview"] = self.content_preview if self.size_bytes is not None: @@ -238,11 +438,22 @@ def to_neo4j_properties(self) -> Dict[str, Any]: @classmethod def from_neo4j_node(cls, node: dict) -> "DataPiece": """Create DataPiece from Neo4j node properties.""" + # Parse content_type enum if present + content_type = None + if node.get("content_type"): + try: + content_type = ContentType(node["content_type"]) + except ValueError: + pass # Unknown content type, leave as None + return cls( fingerprint=node["fingerprint"], ipfs_hash=node["ipfs_hash"], data_type=node["type"], username=node["username"], + content_type=content_type, + generated_by=node.get("generated_by"), + uri=node.get("uri"), content_preview=node.get("content_preview"), size_bytes=node.get("size_bytes"), source=node.get("source"), diff --git a/agentic_rag/utils/__init__.py b/agentic_rag/utils/__init__.py index 08b6a7e..8e634d1 100644 --- a/agentic_rag/utils/__init__.py +++ b/agentic_rag/utils/__init__.py @@ -1,11 +1,15 @@ """Utilities for agentic-rag.""" -from .ipfs_client import LighthouseClient +from .akave_client import AkaveClient from .logger import configure_haystack_logging, get_logger, get_system_logger from .metrics import MetricsCollector, TimedExecution +# Alias for backward compatibility +StorageClient = AkaveClient + __all__ = [ - "LighthouseClient", + "AkaveClient", + "StorageClient", "get_logger", "get_system_logger", "configure_haystack_logging", diff --git a/agentic_rag/utils/akave_client.py b/agentic_rag/utils/akave_client.py new file mode 100644 index 0000000..1663558 --- /dev/null +++ b/agentic_rag/utils/akave_client.py @@ -0,0 +1,456 @@ +"""Akave S3-compatible storage client for uploading and retrieving data. + +Akave provides decentralized S3-compatible storage with encryption and +erasure coding. This client uses boto3 to interact with the Akave O3 API. + +Endpoint: https://o3-rc3.akave.xyz +""" + +import hashlib +import io +import json +from typing import TYPE_CHECKING, Any, Dict, Optional, Union + +import aioboto3 +import boto3 +from botocore.config import Config as BotoConfig +from dotenv import load_dotenv + +if TYPE_CHECKING: + from ..config import Config + +load_dotenv() + +# Default Akave O3 endpoint +AKAVE_ENDPOINT = "https://o3-rc3.akave.xyz" +AKAVE_REGION = "us-east-1" # S3 requires a region, use default + + +class AkaveClient: + """ + Client for interacting with Akave S3-compatible storage. + + Supports: + - Upload files/text/JSON as S3 objects + - Upload raw bytes/buffers + - Retrieve data by object key + - Content-addressed storage (hash-based keys) + + The client maintains compatibility with the LighthouseClient interface + by returning similar response formats. + """ + + def __init__( + self, + access_key: Optional[str] = None, + secret_key: Optional[str] = None, + bucket: Optional[str] = None, + endpoint: Optional[str] = None, + config: Optional["Config"] = None, + timeout: float = 300.0, + ): + """ + Initialize Akave client. + + Args: + access_key: Akave S3 access key (overrides config) + secret_key: Akave S3 secret key (overrides config) + bucket: S3 bucket name for storage (overrides config) + endpoint: Akave endpoint URL (defaults to o3-rc3.akave.xyz) + config: Config object with Akave credentials + timeout: Timeout for API requests in seconds (default: 300.0) + """ + # Priority: explicit params > config object > env vars + if config is not None: + self.access_key = access_key or config.akave_access_key + self.secret_key = secret_key or config.akave_secret_key + self.bucket = bucket or config.akave_bucket + self.endpoint = endpoint or config.akave_endpoint or AKAVE_ENDPOINT + else: + import os + + self.access_key = access_key or os.getenv("AKAVE_ACCESS_KEY") or "" + self.secret_key = secret_key or os.getenv("AKAVE_SECRET_KEY") or "" + self.bucket = bucket or os.getenv("AKAVE_BUCKET", "agentic-rag") + self.endpoint = endpoint or os.getenv("AKAVE_ENDPOINT") or AKAVE_ENDPOINT + + if not self.access_key or not self.secret_key: + raise ValueError( + "Akave credentials required. Provide via config parameter:\n" + " config = Config(akave_access_key='...', akave_secret_key='...')\n" + " AkaveClient(config=config)\n" + "Or set AKAVE_ACCESS_KEY and AKAVE_SECRET_KEY environment variables." + ) + + self.timeout = timeout + + # Configure boto3 client + boto_config = BotoConfig( + connect_timeout=30, + read_timeout=int(timeout), + retries={"max_attempts": 3}, + ) + + # Create sync S3 client + self.s3_client = boto3.client( + "s3", + endpoint_url=self.endpoint, + aws_access_key_id=self.access_key, + aws_secret_access_key=self.secret_key, + region_name=AKAVE_REGION, + config=boto_config, + ) + + # Store session for async client creation + self._boto_config = boto_config + + # Ensure bucket exists + self._ensure_bucket_exists() + + def _ensure_bucket_exists(self) -> None: + """Create bucket if it doesn't exist.""" + try: + self.s3_client.head_bucket(Bucket=self.bucket) + except Exception: + try: + self.s3_client.create_bucket(Bucket=self.bucket) + except Exception: + # Bucket might already exist or we don't have permissions + pass + + def _generate_key(self, data: bytes, prefix: str = "") -> str: + """ + Generate content-addressed key from data hash. + + Args: + data: Raw bytes to hash + prefix: Optional prefix for organization + + Returns: + Key in format: {prefix}/{sha256_hash[:16]} + """ + hash_obj = hashlib.sha256(data) + key = f"ak_{hash_obj.hexdigest()[:16]}" + if prefix: + key = f"{prefix}/{key}" + return key + + def upload_text(self, text: str, name: Optional[str] = None) -> Dict[str, Any]: + """ + Upload text to Akave. + + Args: + text: Text content to upload + name: Optional name (used as prefix) + + Returns: + { + "Name": str, + "Hash": str (object key), + "Size": str + } + """ + text_bytes = text.encode("utf-8") + return self.upload_buffer(text_bytes, name=name) + + async def upload_text_async( + self, text: str, name: Optional[str] = None + ) -> Dict[str, Any]: + """Async version of upload_text.""" + text_bytes = text.encode("utf-8") + return await self.upload_buffer_async(text_bytes, name=name) + + def upload_buffer( + self, data: Union[bytes, io.BytesIO], name: Optional[str] = None + ) -> Dict[str, Any]: + """ + Upload raw bytes/buffer to Akave. + + Args: + data: Bytes or BytesIO buffer + name: Optional name/prefix + + Returns: + { + "Name": str, + "Hash": str (object key), + "Size": str + } + """ + # Convert BytesIO to bytes if needed + if isinstance(data, io.BytesIO): + data = data.getvalue() + + # Generate content-addressed key + prefix = name.replace(".", "_") if name else "data" + key = self._generate_key(data, prefix) + + # Upload to S3 + self.s3_client.put_object( + Bucket=self.bucket, + Key=key, + Body=data, + ) + + return { + "Name": name or key, + "Hash": key, # Maintain compatibility with IPFS client interface + "Size": str(len(data)), + } + + async def upload_buffer_async( + self, data: Union[bytes, io.BytesIO], name: Optional[str] = None + ) -> Dict[str, Any]: + """Async version of upload_buffer.""" + # Convert BytesIO to bytes if needed + if isinstance(data, io.BytesIO): + data = data.getvalue() + + # Generate content-addressed key + prefix = name.replace(".", "_") if name else "data" + key = self._generate_key(data, prefix) + + # Upload using aioboto3 + session = aioboto3.Session() + async with session.client( + "s3", + endpoint_url=self.endpoint, + aws_access_key_id=self.access_key, + aws_secret_access_key=self.secret_key, + region_name=AKAVE_REGION, + ) as s3: + await s3.put_object( + Bucket=self.bucket, + Key=key, + Body=data, + ) + + return { + "Name": name or key, + "Hash": key, + "Size": str(len(data)), + } + + def upload_json( + self, data: Dict[str, Any], name: Optional[str] = None + ) -> Dict[str, Any]: + """ + Upload JSON object to Akave. + + Args: + data: Dictionary to upload as JSON + name: Optional name + + Returns: + { + "Name": str, + "Hash": str (object key), + "Size": str + } + """ + json_str = json.dumps(data, sort_keys=True) + return self.upload_text(json_str, name=name) + + async def upload_json_async( + self, data: Dict[str, Any], name: Optional[str] = None + ) -> Dict[str, Any]: + """Async version of upload_json.""" + json_str = json.dumps(data, sort_keys=True) + return await self.upload_text_async(json_str, name=name) + + def retrieve(self, key: str) -> bytes: + """ + Retrieve raw bytes from Akave. + + Args: + key: Object key (returned as "Hash" from upload methods) + + Returns: + Raw bytes + """ + response = self.s3_client.get_object(Bucket=self.bucket, Key=key) + result: bytes = response["Body"].read() + return result + + async def retrieve_async(self, key: str) -> bytes: + """Async version of retrieve.""" + session = aioboto3.Session() + async with session.client( + "s3", + endpoint_url=self.endpoint, + aws_access_key_id=self.access_key, + aws_secret_access_key=self.secret_key, + region_name=AKAVE_REGION, + ) as s3: + response = await s3.get_object(Bucket=self.bucket, Key=key) + async with response["Body"] as stream: + result: bytes = await stream.read() + return result + + def retrieve_text(self, key: str) -> str: + """ + Retrieve text data from Akave. + + Args: + key: Object key + + Returns: + Text content as string + """ + data = self.retrieve(key) + return data.decode("utf-8") + + async def retrieve_text_async(self, key: str) -> str: + """Async version of retrieve_text.""" + data = await self.retrieve_async(key) + return data.decode("utf-8") + + def retrieve_json(self, key: str) -> Dict[str, Any]: + """ + Retrieve JSON data from Akave. + + Args: + key: Object key + + Returns: + Parsed JSON as dictionary + """ + text = self.retrieve_text(key) + result: Dict[str, Any] = json.loads(text) + return result + + async def retrieve_json_async(self, key: str) -> Dict[str, Any]: + """Async version of retrieve_json.""" + text = await self.retrieve_text_async(key) + result: Dict[str, Any] = json.loads(text) + return result + + def upload_haystack_document( + self, document: Any, as_text: bool = True + ) -> Dict[str, Any]: + """ + Upload Haystack Document to Akave. + + Args: + document: Haystack Document object + as_text: If True, upload content as text (more readable). + If False, upload as JSON with full metadata and embedding. + + Returns: + Upload response with object key + """ + # If document has embedding, always use JSON to preserve it + has_embedding = ( + hasattr(document, "embedding") and document.embedding is not None + ) + + if as_text and not has_embedding: + content = ( + document.content if hasattr(document, "content") else str(document) + ) + doc_id = document.id if hasattr(document, "id") else "unknown" + return self.upload_text(content, name=f"document_{doc_id}") + else: + # Upload as JSON with full metadata and embedding + doc_dict = { + "content": ( + document.content if hasattr(document, "content") else str(document) + ), + "meta": document.meta if hasattr(document, "meta") else {}, + "id": document.id if hasattr(document, "id") else None, + } + # Include embedding if present (convert to list for JSON serialization) + if has_embedding: + embedding = document.embedding + # Handle numpy arrays + if hasattr(embedding, "tolist"): + embedding = embedding.tolist() + doc_dict["embedding"] = embedding + return self.upload_json( + doc_dict, name=f"document_{doc_dict.get('id', 'unknown')}" + ) + + def upload_any( + self, data: Any, name: Optional[str] = None, as_text: bool = True + ) -> Dict[str, Any]: + """ + Smart upload that handles any data type. + + Args: + data: Any data (Document, dict, str, bytes, etc.) + name: Optional name + as_text: For Documents, upload as readable text (default True) + + Returns: + Upload response with object key + """ + # Handle Haystack Document + if hasattr(data, "content"): + return self.upload_haystack_document(data, as_text=as_text) + + # Handle bytes - skip large binary files + if isinstance(data, bytes): + if len(data) > 5 * 1024 * 1024: + return {"Hash": "skipped_large_binary", "Size": str(len(data))} + return self.upload_buffer(data, name=name) + + # Handle dict/list (JSON) + if isinstance(data, dict): + return self.upload_json(data, name=name) + if isinstance(data, list): + return self.upload_json({"data": data}, name=name) + + # Handle string + if isinstance(data, str): + return self.upload_text(data, name=name) + + # Fallback: convert to string + return self.upload_text(str(data), name=name) + + async def upload_any_async( + self, data: Any, name: Optional[str] = None, as_text: bool = True + ) -> Dict[str, Any]: + """Async version of upload_any.""" + # Handle Haystack Document + if hasattr(data, "content"): + return self.upload_haystack_document(data, as_text=as_text) + + # Handle bytes - skip large binary files + if isinstance(data, bytes): + if len(data) > 5 * 1024 * 1024: + return {"Hash": "skipped_large_binary", "Size": str(len(data))} + return await self.upload_buffer_async(data, name=name) + + # Handle dict/list (JSON) + if isinstance(data, dict): + return await self.upload_json_async(data, name=name) + if isinstance(data, list): + return await self.upload_json_async({"data": data}, name=name) + + # Handle string + if isinstance(data, str): + return await self.upload_text_async(data, name=name) + + # Fallback: convert to string + return await self.upload_text_async(str(data), name=name) + + def delete(self, key: str) -> None: + """ + Delete an object from Akave. + + Args: + key: Object key to delete + """ + self.s3_client.delete_object(Bucket=self.bucket, Key=key) + + async def delete_async(self, key: str) -> None: + """Async version of delete.""" + session = aioboto3.Session() + async with session.client( + "s3", + endpoint_url=self.endpoint, + aws_access_key_id=self.access_key, + aws_secret_access_key=self.secret_key, + region_name=AKAVE_REGION, + ) as s3: + await s3.delete_object(Bucket=self.bucket, Key=key) diff --git a/agentic_rag/utils/ipfs_client.py b/agentic_rag/utils/ipfs_client.py deleted file mode 100644 index 0022039..0000000 --- a/agentic_rag/utils/ipfs_client.py +++ /dev/null @@ -1,437 +0,0 @@ -"""Lighthouse IPFS client for uploading and retrieving data.""" - -import io -import json -from typing import TYPE_CHECKING, Any, Dict, Optional, Union - -import httpx -from dotenv import load_dotenv - -if TYPE_CHECKING: - from ..config import Config - -load_dotenv() - - -class LighthouseClient: - """ - Client for interacting with Lighthouse IPFS storage. - - Supports: - - Upload files/text/JSON - - Upload raw bytes/buffers - - Retrieve data by CID - """ - - BASE_URL = "https://upload.lighthouse.storage/api/v0" - GATEWAY_URL = "https://gateway.lighthouse.storage/ipfs" - - def __init__( - self, - api_key: Optional[str] = None, - config: Optional["Config"] = None, - timeout: float = 300.0, - ): - """ - Initialize Lighthouse client. - - Args: - api_key: Lighthouse API key (overrides config) - config: Config object with API key (required if api_key not provided) - timeout: Timeout for API requests in seconds (default: 300.0 for large files) - """ - # Priority: explicit api_key > config object - if config is not None: - self.api_key = api_key or config.lighthouse_api_key - else: - self.api_key = api_key - - if not self.api_key: - raise ValueError( - "Lighthouse API key required. Provide via config parameter:\n" - " config = Config(lighthouse_api_key='your-key')\n" - " LighthouseClient(config=config)" - ) - - self.timeout = timeout - - # Create sync and async HTTP clients - self.client = httpx.Client(timeout=timeout) - self.async_client = httpx.AsyncClient(timeout=timeout) - - def upload_text(self, text: str, name: Optional[str] = None) -> Dict[str, Any]: - """ - Upload text/JSON to Lighthouse. - - Args: - text: Text content to upload (use json.dumps() for JSON) - name: Optional name for the uploaded content - - Returns: - { - "Name": str, - "Hash": str (IPFS CID), - "Size": str - } - """ - # Convert text to bytes - text_bytes = text.encode("utf-8") - - # Upload as buffer - return self.upload_buffer(text_bytes, name=name) - - async def upload_text_async( - self, text: str, name: Optional[str] = None - ) -> Dict[str, Any]: - """ - Async version of upload_text. - - Upload text/JSON to Lighthouse. - - Args: - text: Text content to upload - name: Optional name for the uploaded content - - Returns: - { - "Name": str, - "Hash": str (IPFS CID), - "Size": str - } - """ - # Convert text to bytes - text_bytes = text.encode("utf-8") - - # Upload as buffer - return await self.upload_buffer_async(text_bytes, name=name) - - def upload_buffer( - self, data: Union[bytes, io.BytesIO], name: Optional[str] = None - ) -> Dict[str, Any]: - """ - Upload raw bytes/buffer to Lighthouse. - - Args: - data: Bytes or BytesIO buffer - name: Optional filename - - Returns: - { - "Name": str, - "Hash": str (IPFS CID), - "Size": str - } - """ - # Convert BytesIO to bytes if needed - if isinstance(data, io.BytesIO): - data = data.getvalue() - - # Create file-like object - if name is None: - name = "data" - - files = {"file": (name, data)} - headers = {"Authorization": f"Bearer {self.api_key}"} - - response = self.client.post( - f"{self.BASE_URL}/add", - files=files, - headers=headers, - ) - response.raise_for_status() - - result = response.json() - return { - "Name": result.get("Name"), - "Hash": result.get("Hash"), - "Size": result.get("Size"), - } - - async def upload_buffer_async( - self, data: Union[bytes, io.BytesIO], name: Optional[str] = None - ) -> Dict[str, Any]: - """ - Async version of upload_buffer. - - Upload raw bytes/buffer to Lighthouse. - - Args: - data: Bytes or BytesIO buffer - name: Optional filename - - Returns: - { - "Name": str, - "Hash": str (IPFS CID), - "Size": str - } - """ - # Convert BytesIO to bytes if needed - if isinstance(data, io.BytesIO): - data = data.getvalue() - - # Create file-like object - if name is None: - name = "data" - - files = {"file": (name, data)} - headers = {"Authorization": f"Bearer {self.api_key}"} - - response = await self.async_client.post( - f"{self.BASE_URL}/add", - files=files, - headers=headers, - ) - response.raise_for_status() - - result = response.json() - return { - "Name": result.get("Name"), - "Hash": result.get("Hash"), - "Size": result.get("Size"), - } - - def upload_json( - self, data: Dict[str, Any], name: Optional[str] = None - ) -> Dict[str, Any]: - """ - Upload JSON object to Lighthouse. - - Args: - data: Dictionary to upload as JSON - name: Optional name - - Returns: - { - "Name": str, - "Hash": str (IPFS CID), - "Size": str - } - """ - json_str = json.dumps(data, sort_keys=True) - return self.upload_text(json_str, name=name) - - async def upload_json_async( - self, data: Dict[str, Any], name: Optional[str] = None - ) -> Dict[str, Any]: - """ - Async version of upload_json. - - Upload JSON object to Lighthouse. - - Args: - data: Dictionary to upload as JSON - name: Optional name - - Returns: - { - "Name": str, - "Hash": str (IPFS CID), - "Size": str - } - """ - json_str = json.dumps(data, sort_keys=True) - return await self.upload_text_async(json_str, name=name) - - def retrieve(self, cid: str) -> bytes: - """ - Retrieve raw bytes from IPFS via Lighthouse API. - - """ - - if not self.api_key: - raise ValueError("Lighthouse API key required.") - - headers = {"Authorization": f"Bearer {self.api_key}"} - url = f"https://gateway.lighthouse.storage/ipfs/{cid}" - response = self.client.get(url, headers=headers, timeout=10) - response.raise_for_status() - return bytes(response.content) - - async def retrieve_async(self, cid: str) -> bytes: - """ - Async version of retrieve. - - Retrieve raw bytes from IPFS via Lighthouse API. - - Uses Lighthouse API first (authenticated), falls back to public gateways. - """ - if not self.api_key: - raise ValueError("Lighthouse API key required.") - - headers = {"Authorization": f"Bearer {self.api_key}"} - url = f"https://gateway.lighthouse.storage/ipfs/{cid}" - response = await self.async_client.get(url, headers=headers, timeout=10) - response.raise_for_status() - return bytes(response.content) - - def retrieve_text(self, cid: str) -> str: - """ - Retrieve text data from IPFS. - - Args: - cid: IPFS content identifier - - Returns: - Text content as string - """ - data = self.retrieve(cid) - return data.decode("utf-8") - - async def retrieve_text_async(self, cid: str) -> str: - """ - Async version of retrieve_text. - - Retrieve text data from IPFS. - - Args: - cid: IPFS content identifier - - Returns: - Text content as string - """ - data = await self.retrieve_async(cid) - return data.decode("utf-8") - - def retrieve_json(self, cid: str) -> Dict[str, Any]: - """ - Retrieve JSON data from IPFS. - - Args: - cid: IPFS content identifier - - Returns: - Parsed JSON as dictionary - """ - text = self.retrieve_text(cid) - result: Dict[str, Any] = json.loads(text) - return result - - async def retrieve_json_async(self, cid: str) -> Dict[str, Any]: - """ - Async version of retrieve_json. - - Retrieve JSON data from IPFS. - - Args: - cid: IPFS content identifier - - Returns: - Parsed JSON as dictionary - """ - text = await self.retrieve_text_async(cid) - result: Dict[str, Any] = json.loads(text) - return result - - def upload_haystack_document( - self, document: Any, as_text: bool = True - ) -> Dict[str, Any]: - """ - Upload Haystack Document to IPFS. - - Args: - document: Haystack Document object - as_text: If True, upload content as .txt file (more readable). - If False, upload as JSON with full metadata. - - Returns: - Upload response with CID - """ - if as_text: - # Upload just the content as readable .txt - content = ( - document.content if hasattr(document, "content") else str(document) - ) - doc_id = document.id if hasattr(document, "id") else "unknown" - return self.upload_text(content, name=f"document_{doc_id}.txt") - else: - # Upload full document as JSON (includes metadata) - doc_dict = { - "content": ( - document.content if hasattr(document, "content") else str(document) - ), - "meta": document.meta if hasattr(document, "meta") else {}, - "id": document.id if hasattr(document, "id") else None, - } - return self.upload_json( - doc_dict, name=f"document_{doc_dict.get('id', 'unknown')}.json" - ) - - def upload_any( - self, data: Any, name: Optional[str] = None, as_text: bool = True - ) -> Dict[str, Any]: - """ - Smart upload that handles any data type. - - Args: - data: Any data (Document, dict, str, bytes, etc.) - name: Optional name - as_text: For Documents, upload as readable .txt (default True) - - Returns: - Upload response with CID - """ - # Handle Haystack Document - upload as .txt for readability - if hasattr(data, "content"): - return self.upload_haystack_document(data, as_text=as_text) - - # Handle bytes (PDFs, etc.) - avoid uploading large binaries - if isinstance(data, bytes): - # Skip large binary files (>5MB) to save IPFS space - if len(data) > 5 * 1024 * 1024: - return {"Hash": "skipped_large_binary", "Size": str(len(data))} - return self.upload_buffer(data, name=name) - - # Handle dict/list (JSON) - if isinstance(data, dict): - return self.upload_json(data, name=name) - if isinstance(data, list): - return self.upload_json({"data": data}, name=name) - - # Handle string - if isinstance(data, str): - return self.upload_text(data, name=name) - - # Fallback: convert to string - return self.upload_text(str(data), name=name) - - async def upload_any_async( - self, data: Any, name: Optional[str] = None, as_text: bool = True - ) -> Dict[str, Any]: - """ - Async version of upload_any. - - Smart upload that handles any data type. - - Args: - data: Any data (Document, dict, str, bytes, etc.) - name: Optional name - as_text: For Documents, upload as readable .txt (default True) - - Returns: - Upload response with CID - """ - # Handle Haystack Document - upload as .txt for readability - if hasattr(data, "content"): - # Use sync version for now (could be made async if needed) - return self.upload_haystack_document(data, as_text=as_text) - - # Handle bytes (PDFs, etc.) - avoid uploading large binaries - if isinstance(data, bytes): - # Skip large binary files (>5MB) to save IPFS space - if len(data) > 5 * 1024 * 1024: - return {"Hash": "skipped_large_binary", "Size": str(len(data))} - return await self.upload_buffer_async(data, name=name) - - # Handle dict/list (JSON) - if isinstance(data, dict): - return await self.upload_json_async(data, name=name) - if isinstance(data, list): - return await self.upload_json_async({"data": data}, name=name) - - # Handle string - if isinstance(data, str): - return await self.upload_text_async(data, name=name) - - # Fallback: convert to string - return await self.upload_text_async(str(data), name=name) diff --git a/examples/README.md b/examples/README.md index 4c0b5d4..76f44fd 100644 --- a/examples/README.md +++ b/examples/README.md @@ -25,6 +25,9 @@ poetry run python examples/indexing_pipeline_example.py # Query indexed documents poetry run python examples/retrieval_pipeline_example.py + +# Start MCP server (stdio transport) +poetry run python examples/mcp_server_example.py ``` ## Examples @@ -45,8 +48,16 @@ Multi-source retrieval with evaluation: - Evaluates with BLEU, ROUGE, coherence, readability - **Important**: Must use same `project` as indexing pipelines +### `mcp_server_example.py` + +Runs an MCP server that exposes `agentic-rag` operations as tools: +- create pipelines (indexing and retrieval) +- load pipelines for execution +- run indexing and retrieval pipelines +- list projects, pipelines, and components + ## Storage & Organization - **Neo4j**: Pipeline graphs with `User → Project → Pipelines` hierarchy - **ChromaDB**: Vector embeddings at `./data/{username}/{project}/{pipeline_name}/` -- **IPFS**: Document content with user-based access control +- **Akave**: Document content on decentralized S3-compatible storage with user-based access control diff --git a/examples/async_indexing_pipeline_example.py b/examples/async_indexing_pipeline_example.py index 88a4368..600161e 100644 --- a/examples/async_indexing_pipeline_example.py +++ b/examples/async_indexing_pipeline_example.py @@ -13,7 +13,7 @@ - Generates embeddings using different models - Stores them in separate ChromaDB collections -The pipelines automatically store metadata in Neo4j and content in IPFS. +The pipelines automatically store metadata in Neo4j and content in Akave. """ import asyncio @@ -43,7 +43,9 @@ neo4j_uri=os.getenv("NEO4J_URI", "bolt://localhost:7687"), neo4j_username=os.getenv("NEO4J_USERNAME", "neo4j"), neo4j_password=os.getenv("NEO4J_PASSWORD"), - lighthouse_api_key=os.getenv("LIGHTHOUSE_API_KEY"), # Optional: For IPFS storage + akave_access_key=os.getenv("AKAVE_ACCESS_KEY"), # For Akave S3-compatible storage + akave_secret_key=os.getenv("AKAVE_SECRET_KEY"), + akave_bucket=os.getenv("AKAVE_BUCKET", "agentic-rag"), log_level=os.getenv("AGENTIC_RAG_LOG_LEVEL", "INFO"), ) @@ -160,9 +162,10 @@ async def run_indexing_pipelines_async(data_directory: str) -> Dict[str, Any]: graph_store = GraphStore(config=config) # Initialize runner (singleton - no username needed at init) + # enable_caching=True stores content in Akave and enables transformation caching runner = PipelineRunner( graph_store=graph_store, - enable_caching=False, + enable_caching=True, # Store data in Akave config=config, ) diff --git a/examples/async_qdrant_indexing_example.py b/examples/async_qdrant_indexing_example.py index ad88b22..3628129 100644 --- a/examples/async_qdrant_indexing_example.py +++ b/examples/async_qdrant_indexing_example.py @@ -41,7 +41,9 @@ neo4j_uri=os.getenv("NEO4J_URI", "bolt://localhost:7687"), neo4j_username=os.getenv("NEO4J_USERNAME", "neo4j"), neo4j_password=os.getenv("NEO4J_PASSWORD"), - lighthouse_api_key=os.getenv("LIGHTHOUSE_API_KEY"), + akave_access_key=os.getenv("AKAVE_ACCESS_KEY"), + akave_secret_key=os.getenv("AKAVE_SECRET_KEY"), + akave_bucket=os.getenv("AKAVE_BUCKET", "agentic-rag"), log_level=os.getenv("AGENTIC_RAG_LOG_LEVEL", "INFO"), ) @@ -132,9 +134,10 @@ async def run_qdrant_indexing_pipelines_async(data_directory: str) -> Dict[str, """ graph_store = GraphStore(config=config) + # enable_caching=True stores content in Akave and enables transformation caching runner = PipelineRunner( graph_store=graph_store, - enable_caching=False, + enable_caching=True, # Store data in Akave config=config, ) diff --git a/examples/indexing_pipeline_example.py b/examples/indexing_pipeline_example.py index e22570a..9ab7e3b 100644 --- a/examples/indexing_pipeline_example.py +++ b/examples/indexing_pipeline_example.py @@ -13,7 +13,7 @@ - Generates embeddings using different models - Stores them in separate ChromaDB collections -The pipelines automatically store metadata in Neo4j and content in IPFS. +The pipelines automatically store metadata in Neo4j and content in Akave. """ import os @@ -42,7 +42,9 @@ neo4j_uri=os.getenv("NEO4J_URI", "bolt://localhost:7687"), neo4j_username=os.getenv("NEO4J_USERNAME", "neo4j"), neo4j_password=os.getenv("NEO4J_PASSWORD"), - lighthouse_api_key=os.getenv("LIGHTHOUSE_API_KEY"), # Optional: For IPFS storage + akave_access_key=os.getenv("AKAVE_ACCESS_KEY"), # For Akave S3-compatible storage + akave_secret_key=os.getenv("AKAVE_SECRET_KEY"), + akave_bucket=os.getenv("AKAVE_BUCKET", "agentic-rag"), log_level=os.getenv("AGENTIC_RAG_LOG_LEVEL", "INFO"), ) @@ -153,9 +155,10 @@ def run_indexing_pipelines(data_directory: str) -> Dict[str, Any]: graph_store = GraphStore(config=config) # Initialize runner (singleton - no username needed at init) + # enable_caching=True stores content in Akave and enables transformation caching runner = PipelineRunner( graph_store=graph_store, - enable_caching=False, + enable_caching=True, # Store data in Akave config=config, ) diff --git a/examples/mcp_server_example.py b/examples/mcp_server_example.py new file mode 100644 index 0000000..082c773 --- /dev/null +++ b/examples/mcp_server_example.py @@ -0,0 +1,6 @@ +"""Run the agentic-rag MCP server over stdio.""" + +from agentic_rag.mcp.server import main + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index c4f3377..10b4475 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,8 +37,17 @@ dependencies = [ "scipy (>=1.16.3,<2.0.0)", "types-requests>=2.31.0", "httpx>=0.27.0", + # Required for Akave S3-compatible storage + "boto3>=1.34.0", + "aioboto3>=12.0.0", + "mcp>=1.6.0", + "markdown-it-py (>=4.0.0,<5.0.0)", + "mdit-plain (>=1.0.1,<2.0.0)", ] +[project.scripts] +agentic-rag-mcp = "agentic_rag.mcp.server:main" + [project.optional-dependencies] dev = [ "black>=23.0.0", diff --git a/tests/components/gates/test_gated_component.py b/tests/components/gates/test_gated_component.py index d03121c..3fc45da 100644 --- a/tests/components/gates/test_gated_component.py +++ b/tests/components/gates/test_gated_component.py @@ -44,10 +44,10 @@ def mock_graph_store(): @pytest.fixture -def mock_lighthouse_client(): - """Create a mock Lighthouse IPFS client.""" +def mock_storage_client(): + """Create a mock Akave storage client.""" client = Mock() - client.upload_any.return_value = {"Hash": "QmTestHash123", "Size": "100"} + client.upload_any.return_value = {"Hash": "ak_test_hash_123", "Size": "100"} client.retrieve_text.return_value = "cached content" return client @@ -79,19 +79,19 @@ def mock_outgate(): @pytest.mark.asyncio -@patch("agentic_rag.components.gates.ingate.LighthouseClient") -@patch("agentic_rag.components.gates.outgate.LighthouseClient") +@patch("agentic_rag.components.gates.ingate.AkaveClient") +@patch("agentic_rag.components.gates.outgate.AkaveClient") async def test_run_async_with_async_component( - mock_outgate_lighthouse, - mock_ingate_lighthouse, + mock_outgate_storage, + mock_ingate_storage, mock_graph_store, mock_ingate, mock_outgate, - mock_lighthouse_client, + mock_storage_client, ): """Test run_async with a component that has run_async method.""" - mock_ingate_lighthouse.return_value = mock_lighthouse_client - mock_outgate_lighthouse.return_value = mock_lighthouse_client + mock_ingate_storage.return_value = mock_storage_client + mock_outgate_storage.return_value = mock_storage_client component = MockAsyncComponent() gated = GatedComponent( @@ -118,19 +118,19 @@ async def test_run_async_with_async_component( @pytest.mark.asyncio -@patch("agentic_rag.components.gates.ingate.LighthouseClient") -@patch("agentic_rag.components.gates.outgate.LighthouseClient") +@patch("agentic_rag.components.gates.ingate.AkaveClient") +@patch("agentic_rag.components.gates.outgate.AkaveClient") async def test_run_async_with_sync_component( - mock_outgate_lighthouse, - mock_ingate_lighthouse, + mock_outgate_storage, + mock_ingate_storage, mock_graph_store, mock_ingate, mock_outgate, - mock_lighthouse_client, + mock_storage_client, ): """Test run_async with a component that only has sync run method.""" - mock_ingate_lighthouse.return_value = mock_lighthouse_client - mock_outgate_lighthouse.return_value = mock_lighthouse_client + mock_ingate_storage.return_value = mock_storage_client + mock_outgate_storage.return_value = mock_storage_client component = MockComponent() gated = GatedComponent( @@ -157,17 +157,17 @@ async def test_run_async_with_sync_component( @pytest.mark.asyncio -@patch("agentic_rag.components.gates.ingate.LighthouseClient") -@patch("agentic_rag.components.gates.outgate.LighthouseClient") +@patch("agentic_rag.components.gates.ingate.AkaveClient") +@patch("agentic_rag.components.gates.outgate.AkaveClient") async def test_run_async_with_cached_items( - mock_outgate_lighthouse, - mock_ingate_lighthouse, + mock_outgate_storage, + mock_ingate_storage, mock_graph_store, - mock_lighthouse_client, + mock_storage_client, ): """Test run_async when all items are in cache.""" - mock_ingate_lighthouse.return_value = mock_lighthouse_client - mock_outgate_lighthouse.return_value = mock_lighthouse_client + mock_ingate_storage.return_value = mock_storage_client + mock_outgate_storage.return_value = mock_storage_client component = MockAsyncComponent() gated = GatedComponent( @@ -200,18 +200,18 @@ async def test_run_async_with_cached_items( @pytest.mark.asyncio -@patch("agentic_rag.components.gates.ingate.LighthouseClient") -@patch("agentic_rag.components.gates.outgate.LighthouseClient") +@patch("agentic_rag.components.gates.ingate.AkaveClient") +@patch("agentic_rag.components.gates.outgate.AkaveClient") async def test_run_async_partial_cache( - mock_outgate_lighthouse, - mock_ingate_lighthouse, + mock_outgate_storage, + mock_ingate_storage, mock_graph_store, mock_outgate, - mock_lighthouse_client, + mock_storage_client, ): """Test run_async with partial cache hits.""" - mock_ingate_lighthouse.return_value = mock_lighthouse_client - mock_outgate_lighthouse.return_value = mock_lighthouse_client + mock_ingate_storage.return_value = mock_storage_client + mock_outgate_storage.return_value = mock_storage_client component = MockAsyncComponent() gated = GatedComponent( @@ -242,17 +242,17 @@ async def test_run_async_partial_cache( @pytest.mark.asyncio -@patch("agentic_rag.components.gates.ingate.LighthouseClient") -@patch("agentic_rag.components.gates.outgate.LighthouseClient") +@patch("agentic_rag.components.gates.ingate.AkaveClient") +@patch("agentic_rag.components.gates.outgate.AkaveClient") async def test_run_async_no_cacheable_inputs( - mock_outgate_lighthouse, - mock_ingate_lighthouse, + mock_outgate_storage, + mock_ingate_storage, mock_graph_store, - mock_lighthouse_client, + mock_storage_client, ): """Test run_async with no cacheable inputs.""" - mock_ingate_lighthouse.return_value = mock_lighthouse_client - mock_outgate_lighthouse.return_value = mock_lighthouse_client + mock_ingate_storage.return_value = mock_storage_client + mock_outgate_storage.return_value = mock_storage_client component = MockAsyncComponent() gated = GatedComponent( @@ -271,18 +271,18 @@ async def test_run_async_no_cacheable_inputs( @pytest.mark.asyncio -@patch("agentic_rag.components.gates.ingate.LighthouseClient") -@patch("agentic_rag.components.gates.outgate.LighthouseClient") +@patch("agentic_rag.components.gates.ingate.AkaveClient") +@patch("agentic_rag.components.gates.outgate.AkaveClient") async def test_run_async_error_handling( - mock_outgate_lighthouse, - mock_ingate_lighthouse, + mock_outgate_storage, + mock_ingate_storage, mock_graph_store, mock_ingate, - mock_lighthouse_client, + mock_storage_client, ): """Test run_async handles errors gracefully.""" - mock_ingate_lighthouse.return_value = mock_lighthouse_client - mock_outgate_lighthouse.return_value = mock_lighthouse_client + mock_ingate_storage.return_value = mock_storage_client + mock_outgate_storage.return_value = mock_storage_client class FailingComponent: __haystack_supports_async__ = True @@ -307,19 +307,19 @@ async def run_async(self, **kwargs): await gated.run_async(documents=["doc1"]) -@patch("agentic_rag.components.gates.ingate.LighthouseClient") -@patch("agentic_rag.components.gates.outgate.LighthouseClient") +@patch("agentic_rag.components.gates.ingate.AkaveClient") +@patch("agentic_rag.components.gates.outgate.AkaveClient") def test_sync_run_still_works( - mock_outgate_lighthouse, - mock_ingate_lighthouse, + mock_outgate_storage, + mock_ingate_storage, mock_graph_store, mock_ingate, mock_outgate, - mock_lighthouse_client, + mock_storage_client, ): """Test that sync run method still works after adding async.""" - mock_ingate_lighthouse.return_value = mock_lighthouse_client - mock_outgate_lighthouse.return_value = mock_lighthouse_client + mock_ingate_storage.return_value = mock_storage_client + mock_outgate_storage.return_value = mock_storage_client component = MockComponent() gated = GatedComponent( diff --git a/tests/conftest.py b/tests/conftest.py index cf689bd..da9147f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -29,7 +29,9 @@ def test_config(): neo4j_password=os.getenv("NEO4J_PASSWORD", "test_password"), neo4j_database=os.getenv("NEO4J_DATABASE"), openrouter_api_key=os.getenv("OPENROUTER_API_KEY", "test_openrouter_key"), - lighthouse_api_key=os.getenv("LIGHTHOUSE_API_KEY", "test_lighthouse_key"), + akave_access_key=os.getenv("AKAVE_ACCESS_KEY", "test_akave_access_key"), + akave_secret_key=os.getenv("AKAVE_SECRET_KEY", "test_akave_secret_key"), + akave_bucket=os.getenv("AKAVE_BUCKET", "test-bucket"), log_level="DEBUG", )