robosystems-client 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

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
 

robosystems_client/extensions/__init__.py

@@ -27,6 +27,12 @@ from .table_ingest_client import (
   UploadResult,
   TableInfo,
 )
+from .graph_client import (
+  GraphClient,
+  GraphMetadata,
+  InitialEntityData,
+  GraphInfo,
+)
 from .extensions import (
   RoboSystemsExtensions,
   RoboSystemsExtensionConfig,
@@ -122,6 +128,11 @@ __all__ = [
   "IngestOptions",
   "UploadResult",
   "TableInfo",
+  # Graph Client
+  "GraphClient",
+  "GraphMetadata",
+  "InitialEntityData",
+  "GraphInfo",
   # Utilities
   "QueryBuilder",
   "ResultProcessor",

robosystems_client/extensions/extensions.py

@@ -9,6 +9,7 @@ from typing import Dict, Any, Optional, Callable
 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 @@ class RoboSystemsExtensions:
     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 @@ class RoboSystemsExtensions:
     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):

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

robosystems_client/extensions/query_client.py

@@ -113,9 +113,30 @@ class QueryClient:
     )
 
     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 @@ class QueryClient:
     # 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]:

{robosystems_client-0.2.1.dist-info → robosystems_client-0.2.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: robosystems-client
-Version: 0.2.1
+Version: 0.2.2
 Summary: Python Client for RoboSystems financial graph database API
 Author: RFS LLC
 License: MIT

{robosystems_client-0.2.1.dist-info → robosystems_client-0.2.2.dist-info}/RECORD

@@ -79,7 +79,7 @@ robosystems_client/api/operations/cancel_operation.py,sha256=3IiKfMfSANQS_inSg9a
 robosystems_client/api/operations/get_operation_status.py,sha256=r77igqYIVW9x34tlNFiWrueXGYzfTrNJ3Bsb-qb7urI,8513
 robosystems_client/api/operations/stream_operation_events.py,sha256=1pl-BNZSaNynIKwaT03U8WXEbL_B5KOPSZuEb1vS4CM,14042
 robosystems_client/api/query/__init__.py,sha256=5vd9uJWAjRqa9xzxzYkLD1yoZ12Ld_bAaNB5WX4fbE8,56
-robosystems_client/api/query/execute_cypher_query.py,sha256=CZrjL8jYsd_OCYF8zDJfZlI_5ERD13mhXtIrW1reEOY,18069
+robosystems_client/api/query/execute_cypher_query.py,sha256=FgXjAjj9_sXFxpNuflQ9mmKVSJzdQr509AlsEdp41hk,18410
 robosystems_client/api/schema/__init__.py,sha256=5vd9uJWAjRqa9xzxzYkLD1yoZ12Ld_bAaNB5WX4fbE8,56
 robosystems_client/api/schema/export_graph_schema.py,sha256=rcx9UzQHTaTQCUEUI2HIqJ5fj_n8L096-8awrt4qfa8,7589
 robosystems_client/api/schema/get_graph_schema.py,sha256=a86EAFzysQdjowI4eJF25MeRVicc3IEQSZi-3Dl16jA,8503
@@ -128,12 +128,13 @@ robosystems_client/api/user_subscriptions/get_user_shared_subscriptions.py,sha25
 robosystems_client/api/user_subscriptions/subscribe_to_shared_repository.py,sha256=GBr_AahYctNKUVg5O-wqju3sdL-DYFZLPtYnjfR9c5c,6606
 robosystems_client/api/user_subscriptions/upgrade_shared_repository_subscription.py,sha256=-kfP0zkXo9-OojY8-EvxK-hFiJBCdR-MV4puqCcHtrA,6742
 robosystems_client/extensions/README.md,sha256=qfHFjdgA_J-zNXziNZE6M1MKJiwVkocBi01w_HhvzEk,16136
-robosystems_client/extensions/__init__.py,sha256=Jj2K6iNTkymAXuQAEhWAOFc0Y54lBXx2orfm83x5TpE,4806
+robosystems_client/extensions/__init__.py,sha256=BH2VcEDBL8-HcT-nmwmEPUveRxoWVTEnTog6X3oJZAY,4993
 robosystems_client/extensions/auth_integration.py,sha256=9dBuLP-dWeZecaf7hXZcJA-L6eiSCeTHIPFIUMmASkU,6610
 robosystems_client/extensions/dataframe_utils.py,sha256=gK1bgkVqBF0TvWVdGQvqWrt-ur_Rw11j8uNtMoulLWE,12312
-robosystems_client/extensions/extensions.py,sha256=JPxy154A_NZLZBB3iB5BUmh59MUuXg0jp0e6sX69JKU,6152
+robosystems_client/extensions/extensions.py,sha256=ROnCobUek4Dke9dVx2sTzNKhz309NOG40EDSYHtNmWs,6257
+robosystems_client/extensions/graph_client.py,sha256=GwsFwETti6xxrQbqk7C_ckxZAVhMO05DrS-qACacLfk,10012
 robosystems_client/extensions/operation_client.py,sha256=B1qju-wWQrnrnVJixKGgsA_KEInviwJwdlJxzm_i7P0,13359
-robosystems_client/extensions/query_client.py,sha256=lwkv1ySm6kiVgnke_K9BulXRx4WXU4EFQZb_UJAxx_M,14630
+robosystems_client/extensions/query_client.py,sha256=NVaoTrYzSdR_pzrSwi8AsZ7WaUi7yTAcN9eJ5HFeFa8,16885
 robosystems_client/extensions/sse_client.py,sha256=fWO6EPiLmS3LJ8iph5wup57ldtHSLBkDoiAPfBNNCyk,15048
 robosystems_client/extensions/table_ingest_client.py,sha256=MyCnBAijFlkWY88dY4CJkouY0w5d4D3cxftw_wmSt60,12790
 robosystems_client/extensions/token_utils.py,sha256=qCK_s1vBzRnSYwtgncPZRLJVIw3WXmzqNTWjdEEpdgs,10899
@@ -342,7 +343,7 @@ robosystems_client/models/user_usage_response_graphs.py,sha256=xAH-ZnhaUfWQ_2EpZ
 robosystems_client/models/user_usage_summary_response.py,sha256=4hthwTH7bXyzdYlHoekDYOgDLI-stGRH507Bl2rUjYA,3655
 robosystems_client/models/user_usage_summary_response_usage_vs_limits.py,sha256=XrZnRcy1nD3xtKX4svbww7QfEHrN7_XIfeL9j5ZMbyQ,1298
 robosystems_client/models/validation_error.py,sha256=R77OuQG2nJ3WDFfY--xbEhg6x1D7gAAp_1UdnG8Ka2A,1949
-robosystems_client-0.2.1.dist-info/METADATA,sha256=Lz6GHnT549mGN8HpaOYLqfb_VGjXx1n4IEXHaYyA9Ic,3959
-robosystems_client-0.2.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-robosystems_client-0.2.1.dist-info/licenses/LICENSE,sha256=LjFqQPU4eQh7jAQ04SmE9eC0j74HCdXvzbo0hjW4mWo,1063
-robosystems_client-0.2.1.dist-info/RECORD,,
+robosystems_client-0.2.2.dist-info/METADATA,sha256=qMZLxDrWEKUYi1o1EDIr1zPQVuMIzkw61hBupaUEP-M,3959
+robosystems_client-0.2.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+robosystems_client-0.2.2.dist-info/licenses/LICENSE,sha256=LjFqQPU4eQh7jAQ04SmE9eC0j74HCdXvzbo0hjW4mWo,1063
+robosystems_client-0.2.2.dist-info/RECORD,,