diff --git a/robosystems_client/api/query/execute_cypher_query.py b/robosystems_client/api/query/execute_cypher_query.py index c34a26d4..25d2e29b 100644 --- a/robosystems_client/api/query/execute_cypher_query.py +++ b/robosystems_client/api/query/execute_cypher_query.py @@ -67,6 +67,11 @@ def _parse_response( *, client: Union[AuthenticatedClient, Client], response: httpx.Response ) -> Optional[Union[Any, HTTPValidationError]]: if response.status_code == 200: + # Check if this is NDJSON - if so, skip parsing (will be handled by client) + content_type = response.headers.get("content-type", "") + stream_format = response.headers.get("x-stream-format", "") + if "application/x-ndjson" in content_type or stream_format == "ndjson": + return None # Skip parsing, client will handle NDJSON response_200 = response.json() return response_200 diff --git a/robosystems_client/extensions/__init__.py b/robosystems_client/extensions/__init__.py index f8cc4086..718f86d8 100644 --- a/robosystems_client/extensions/__init__.py +++ b/robosystems_client/extensions/__init__.py @@ -27,6 +27,12 @@ UploadResult, TableInfo, ) +from .graph_client import ( + GraphClient, + GraphMetadata, + InitialEntityData, + GraphInfo, +) from .extensions import ( RoboSystemsExtensions, RoboSystemsExtensionConfig, @@ -122,6 +128,11 @@ "IngestOptions", "UploadResult", "TableInfo", + # Graph Client + "GraphClient", + "GraphMetadata", + "InitialEntityData", + "GraphInfo", # Utilities "QueryBuilder", "ResultProcessor", diff --git a/robosystems_client/extensions/extensions.py b/robosystems_client/extensions/extensions.py index c25f78f5..98943e46 100644 --- a/robosystems_client/extensions/extensions.py +++ b/robosystems_client/extensions/extensions.py @@ -9,6 +9,7 @@ from .query_client import QueryClient from .operation_client import OperationClient from .table_ingest_client import TableIngestClient +from .graph_client import GraphClient from .sse_client import SSEClient @@ -59,6 +60,7 @@ def __init__(self, config: RoboSystemsExtensionConfig = None): self.query = QueryClient(self.config) self.operations = OperationClient(self.config) self.tables = TableIngestClient(self.config) + self.graphs = GraphClient(self.config) def monitor_operation( self, operation_id: str, on_progress: Optional[Callable] = None @@ -88,6 +90,7 @@ def close(self): self.query.close() self.operations.close_all() self.tables.close() + self.graphs.close() # Convenience methods that delegate to the appropriate clients def execute_query(self, graph_id: str, query: str, parameters: Dict[str, Any] = None): diff --git a/robosystems_client/extensions/graph_client.py b/robosystems_client/extensions/graph_client.py new file mode 100644 index 00000000..a40a2b99 --- /dev/null +++ b/robosystems_client/extensions/graph_client.py @@ -0,0 +1,321 @@ +"""Graph Management Client + +Provides high-level graph management operations with automatic operation monitoring. +""" + +from dataclasses import dataclass +from typing import Dict, Any, Optional, Callable +import time +import logging + +logger = logging.getLogger(__name__) + + +@dataclass +class GraphMetadata: + """Graph metadata for creation""" + + graph_name: str + description: Optional[str] = None + schema_extensions: Optional[list] = None + tags: Optional[list] = None + + +@dataclass +class InitialEntityData: + """Initial entity data for graph creation""" + + name: str + uri: str + category: Optional[str] = None + sic: Optional[str] = None + sic_description: Optional[str] = None + + +@dataclass +class GraphInfo: + """Information about a graph""" + + graph_id: str + graph_name: str + description: Optional[str] = None + schema_extensions: Optional[list] = None + tags: Optional[list] = None + created_at: Optional[str] = None + status: Optional[str] = None + + +class GraphClient: + """Client for graph management operations""" + + def __init__(self, config: Dict[str, Any]): + self.config = config + self.base_url = config["base_url"] + self.headers = config.get("headers", {}) + self.token = config.get("token") + + def create_graph_and_wait( + self, + metadata: GraphMetadata, + initial_entity: Optional[InitialEntityData] = None, + timeout: int = 60, + poll_interval: int = 2, + on_progress: Optional[Callable[[str], None]] = None, + ) -> str: + """ + Create a graph and wait for completion. + + Args: + metadata: Graph metadata + initial_entity: Optional initial entity data + timeout: Maximum time to wait in seconds + poll_interval: Time between status checks in seconds + on_progress: Callback for progress updates + + Returns: + graph_id when creation completes + + Raises: + Exception: If creation fails or times out + """ + from ..client import AuthenticatedClient + from ..api.graphs.create_graph import sync_detailed as create_graph + from ..api.operations.get_operation_status import sync_detailed as get_status + from ..models.create_graph_request import CreateGraphRequest + from ..models.graph_metadata import GraphMetadata as APIGraphMetadata + + if not self.token: + raise ValueError("No API key provided. Set X-API-Key in headers.") + + client = AuthenticatedClient( + base_url=self.base_url, + token=self.token, + prefix="", + auth_header_name="X-API-Key", + headers=self.headers, + ) + + # Build API metadata + api_metadata = APIGraphMetadata( + graph_name=metadata.graph_name, + description=metadata.description, + schema_extensions=metadata.schema_extensions or [], + tags=metadata.tags or [], + ) + + # Build initial entity if provided + initial_entity_dict = None + if initial_entity: + initial_entity_dict = { + "name": initial_entity.name, + "uri": initial_entity.uri, + } + if initial_entity.category: + initial_entity_dict["category"] = initial_entity.category + if initial_entity.sic: + initial_entity_dict["sic"] = initial_entity.sic + if initial_entity.sic_description: + initial_entity_dict["sic_description"] = initial_entity.sic_description + + # Create graph request + graph_create = CreateGraphRequest( + metadata=api_metadata, + initial_entity=initial_entity_dict, + ) + + if on_progress: + on_progress(f"Creating graph: {metadata.graph_name}") + + # Execute create request + response = create_graph(client=client, body=graph_create) + + if not response.parsed: + raise RuntimeError(f"Failed to create graph: {response.status_code}") + + # Extract graph_id or operation_id + if isinstance(response.parsed, dict): + graph_id = response.parsed.get("graph_id") + operation_id = response.parsed.get("operation_id") + else: + graph_id = getattr(response.parsed, "graph_id", None) + operation_id = getattr(response.parsed, "operation_id", None) + + # If graph_id returned immediately, we're done + if graph_id: + if on_progress: + on_progress(f"Graph created: {graph_id}") + return graph_id + + # Otherwise, poll operation until complete + if not operation_id: + raise RuntimeError("No graph_id or operation_id in response") + + if on_progress: + on_progress(f"Graph creation queued (operation: {operation_id})") + + max_attempts = timeout // poll_interval + for attempt in range(max_attempts): + time.sleep(poll_interval) + + status_response = get_status(operation_id=operation_id, client=client) + + if not status_response.parsed: + continue + + # Handle both dict and object responses + status_data = status_response.parsed + if isinstance(status_data, dict): + status = status_data.get("status") + else: + # Check for additional_properties first (common in generated clients) + if hasattr(status_data, "additional_properties"): + status = status_data.additional_properties.get("status") + else: + status = getattr(status_data, "status", None) + + if on_progress: + on_progress(f"Status: {status} (attempt {attempt + 1}/{max_attempts})") + + if status == "completed": + # Extract graph_id from result + if isinstance(status_data, dict): + result = status_data.get("result", {}) + elif hasattr(status_data, "additional_properties"): + result = status_data.additional_properties.get("result", {}) + else: + result = getattr(status_data, "result", {}) + + if isinstance(result, dict): + graph_id = result.get("graph_id") + else: + graph_id = getattr(result, "graph_id", None) + + if graph_id: + if on_progress: + on_progress(f"Graph created: {graph_id}") + return graph_id + else: + raise RuntimeError("Operation completed but no graph_id in result") + + elif status == "failed": + # Extract error message + if isinstance(status_data, dict): + error = ( + status_data.get("error") or status_data.get("message") or "Unknown error" + ) + elif hasattr(status_data, "additional_properties"): + props = status_data.additional_properties + error = props.get("error") or props.get("message") or "Unknown error" + else: + error = getattr(status_data, "message", "Unknown error") + raise RuntimeError(f"Graph creation failed: {error}") + + raise TimeoutError(f"Graph creation timed out after {timeout}s") + + def get_graph_info(self, graph_id: str) -> GraphInfo: + """ + Get information about a graph. + + Args: + graph_id: The graph ID + + Returns: + GraphInfo with graph details + + Raises: + ValueError: If graph not found + """ + from ..client import AuthenticatedClient + from ..api.graphs.get_graphs import sync_detailed as get_graphs + + if not self.token: + raise ValueError("No API key provided. Set X-API-Key in headers.") + + client = AuthenticatedClient( + base_url=self.base_url, + token=self.token, + prefix="", + auth_header_name="X-API-Key", + headers=self.headers, + ) + + # Use get_graphs and filter for the specific graph + response = get_graphs(client=client) + + if not response.parsed: + raise RuntimeError(f"Failed to get graphs: {response.status_code}") + + data = response.parsed + graphs = None + + # Extract graphs list from response + if isinstance(data, dict): + graphs = data.get("graphs", []) + elif hasattr(data, "additional_properties"): + graphs = data.additional_properties.get("graphs", []) + elif hasattr(data, "graphs"): + graphs = data.graphs + else: + raise RuntimeError("Unexpected response format from get_graphs") + + # Find the specific graph by ID + graph_data = None + for graph in graphs: + if isinstance(graph, dict): + if graph.get("graph_id") == graph_id or graph.get("id") == graph_id: + graph_data = graph + break + elif hasattr(graph, "graph_id"): + if graph.graph_id == graph_id or getattr(graph, "id", None) == graph_id: + graph_data = graph + break + + if not graph_data: + raise ValueError(f"Graph not found: {graph_id}") + + # Build GraphInfo from the found graph + if isinstance(graph_data, dict): + return GraphInfo( + graph_id=graph_data.get("graph_id") or graph_data.get("id", graph_id), + graph_name=graph_data.get("graph_name") or graph_data.get("name", ""), + description=graph_data.get("description"), + schema_extensions=graph_data.get("schema_extensions"), + tags=graph_data.get("tags"), + created_at=graph_data.get("created_at"), + status=graph_data.get("status"), + ) + else: + return GraphInfo( + graph_id=getattr(graph_data, "graph_id", None) + or getattr(graph_data, "id", graph_id), + graph_name=getattr(graph_data, "graph_name", None) + or getattr(graph_data, "name", ""), + description=getattr(graph_data, "description", None), + schema_extensions=getattr(graph_data, "schema_extensions", None), + tags=getattr(graph_data, "tags", None), + created_at=getattr(graph_data, "created_at", None), + status=getattr(graph_data, "status", None), + ) + + def delete_graph(self, graph_id: str) -> None: + """ + Delete a graph. + + Note: This method is not yet available as the delete_graph endpoint + is not included in the generated SDK. This will be implemented when + the endpoint is added to the API specification. + + Args: + graph_id: The graph ID to delete + + Raises: + NotImplementedError: This feature is not yet available + """ + raise NotImplementedError( + "Graph deletion is not yet available. " + "The delete_graph endpoint needs to be added to the API specification." + ) + + def close(self): + """Clean up resources (placeholder for consistency)""" + pass diff --git a/robosystems_client/extensions/query_client.py b/robosystems_client/extensions/query_client.py index 47a52805..a0642505 100644 --- a/robosystems_client/extensions/query_client.py +++ b/robosystems_client/extensions/query_client.py @@ -113,9 +113,30 @@ def execute_query( ) try: - kwargs = {"graph_id": graph_id, "client": client, "body": query_request} + kwargs = { + "graph_id": graph_id, + "client": client, + "body": query_request, + "mode": options.mode if options.mode else None, + "chunk_size": options.chunk_size if options.chunk_size else 1000, + "test_mode": options.test_mode if options.test_mode else False, + } response = execute_cypher_query(**kwargs) + # Check if this is an NDJSON streaming response (parsed will be None for NDJSON) + if ( + hasattr(response, "headers") + and ( + "application/x-ndjson" in response.headers.get("content-type", "") + or response.headers.get("x-stream-format") == "ndjson" + ) + ) or ( + hasattr(response, "parsed") + and response.parsed is None + and response.status_code == 200 + ): + return self._parse_ndjson_response(response, graph_id) + # Check response type and handle accordingly if hasattr(response, "parsed") and response.parsed: response_data = response.parsed @@ -187,6 +208,58 @@ def execute_query( # Unexpected response format raise Exception("Unexpected response format from query endpoint") + def _parse_ndjson_response(self, response, graph_id: str) -> QueryResult: + """Parse NDJSON streaming response and aggregate into QueryResult""" + import json + + all_data = [] + columns = None + total_rows = 0 + execution_time_ms = 0 + + # Parse NDJSON line by line + content = ( + response.content.decode("utf-8") + if isinstance(response.content, bytes) + else response.content + ) + + for line in content.strip().split("\n"): + if not line.strip(): + continue + + try: + chunk = json.loads(line) + + # Extract columns from first chunk + if columns is None and "columns" in chunk: + columns = chunk["columns"] + + # Aggregate data rows (NDJSON uses "rows", regular JSON uses "data") + if "rows" in chunk: + all_data.extend(chunk["rows"]) + total_rows += len(chunk["rows"]) + elif "data" in chunk: + all_data.extend(chunk["data"]) + total_rows += len(chunk["data"]) + + # Track execution time (use max from all chunks) + if "execution_time_ms" in chunk: + execution_time_ms = max(execution_time_ms, chunk["execution_time_ms"]) + + except json.JSONDecodeError as e: + raise Exception(f"Failed to parse NDJSON line: {e}") + + # Return aggregated result + return QueryResult( + data=all_data, + columns=columns or [], + row_count=total_rows, + execution_time_ms=execution_time_ms, + graph_id=graph_id, + timestamp=datetime.now().isoformat(), + ) + def _stream_query_results( self, operation_id: str, options: QueryOptions ) -> Iterator[Any]: diff --git a/tests/test_extensions.py b/tests/test_extensions.py index 1678f7bb..c43ba203 100644 --- a/tests/test_extensions.py +++ b/tests/test_extensions.py @@ -7,6 +7,7 @@ TableIngestClient, QueryClient, OperationClient, + GraphClient, ) @@ -22,6 +23,7 @@ def test_extensions_initialization_default(self): assert isinstance(extensions.query, QueryClient) assert isinstance(extensions.operations, OperationClient) assert isinstance(extensions.tables, TableIngestClient) + assert isinstance(extensions.graphs, GraphClient) # Cleanup extensions.close() @@ -115,6 +117,20 @@ def test_table_ingest_client_receives_config(self): extensions.close() + def test_graph_client_receives_config(self): + """Test that GraphClient receives proper config.""" + config = RoboSystemsExtensionConfig( + base_url="https://api.robosystems.ai", + headers={"X-API-Key": "test-token"}, + ) + + extensions = RoboSystemsExtensions(config) + + assert extensions.graphs.base_url == "https://api.robosystems.ai" + assert "X-API-Key" in extensions.graphs.headers + + extensions.close() + def test_multiple_extensions_instances(self): """Test that multiple extension instances can coexist.""" ext1 = RoboSystemsExtensions( diff --git a/tests/test_graph_client.py b/tests/test_graph_client.py new file mode 100644 index 00000000..9b3f18fb --- /dev/null +++ b/tests/test_graph_client.py @@ -0,0 +1,110 @@ +"""Unit tests for GraphClient.""" + +import pytest +from robosystems_client.extensions.graph_client import ( + GraphClient, + GraphMetadata, + InitialEntityData, + GraphInfo, +) + + +@pytest.mark.unit +class TestGraphClient: + """Test suite for GraphClient.""" + + def test_client_initialization(self, mock_config): + """Test that client initializes correctly with config.""" + client = GraphClient(mock_config) + + assert client.base_url == "http://localhost:8000" + assert client.token == "test-api-key" + assert client.headers == {"X-API-Key": "test-api-key"} + + def test_graph_metadata_dataclass(self): + """Test GraphMetadata dataclass.""" + metadata = GraphMetadata( + graph_name="Test Graph", + description="A test graph", + schema_extensions=["custom_prop"], + tags=["test", "demo"], + ) + + assert metadata.graph_name == "Test Graph" + assert metadata.description == "A test graph" + assert metadata.schema_extensions == ["custom_prop"] + assert metadata.tags == ["test", "demo"] + + def test_graph_metadata_defaults(self): + """Test GraphMetadata default values.""" + metadata = GraphMetadata(graph_name="Simple Graph") + + assert metadata.graph_name == "Simple Graph" + assert metadata.description is None + assert metadata.schema_extensions is None + assert metadata.tags is None + + def test_initial_entity_data_dataclass(self): + """Test InitialEntityData dataclass.""" + entity = InitialEntityData( + name="ACME Corp", + uri="https://example.com/acme", + category="Technology", + sic="7372", + sic_description="Prepackaged Software", + ) + + assert entity.name == "ACME Corp" + assert entity.uri == "https://example.com/acme" + assert entity.category == "Technology" + assert entity.sic == "7372" + assert entity.sic_description == "Prepackaged Software" + + def test_initial_entity_data_defaults(self): + """Test InitialEntityData default values.""" + entity = InitialEntityData(name="Basic Entity", uri="https://example.com") + + assert entity.name == "Basic Entity" + assert entity.uri == "https://example.com" + assert entity.category is None + assert entity.sic is None + assert entity.sic_description is None + + def test_graph_info_dataclass(self): + """Test GraphInfo dataclass.""" + info = GraphInfo( + graph_id="graph-123", + graph_name="Production Graph", + description="Production knowledge graph", + schema_extensions=["prop1", "prop2"], + tags=["prod"], + created_at="2024-01-15T10:30:00Z", + status="active", + ) + + assert info.graph_id == "graph-123" + assert info.graph_name == "Production Graph" + assert info.description == "Production knowledge graph" + assert info.schema_extensions == ["prop1", "prop2"] + assert info.tags == ["prod"] + assert info.created_at == "2024-01-15T10:30:00Z" + assert info.status == "active" + + def test_graph_info_minimal(self): + """Test GraphInfo with minimal data.""" + info = GraphInfo(graph_id="graph-456", graph_name="Minimal Graph") + + assert info.graph_id == "graph-456" + assert info.graph_name == "Minimal Graph" + assert info.description is None + assert info.schema_extensions is None + assert info.tags is None + assert info.created_at is None + assert info.status is None + + def test_close_method(self, mock_config): + """Test that close method exists and can be called.""" + client = GraphClient(mock_config) + + # Should not raise any exceptions + client.close()