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

robosystems_client/extensions/graph_client.py

@@ -0,0 +1,326 @@
+"""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,
+    create_entity: bool = True,
+    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
+        create_entity: Whether to create the entity node and upload initial data.
+            Only applies when initial_entity is provided. Set to False to create
+            graph without populating entity data (useful for file-based ingestion).
+        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,
+      create_entity=create_entity,
+    )
+
+    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/extensions/table_ingest_client.py

@@ -11,20 +11,20 @@ import json
 import logging
 import httpx
 
-from ..api.tables.get_upload_url_v1_graphs_graph_id_tables_table_name_files_post import (
+from ..api.tables.get_upload_url import (
   sync_detailed as get_upload_url,
 )
-from ..api.tables.update_file_v1_graphs_graph_id_tables_files_file_id_patch import (
-  sync_detailed as update_file,
+from ..api.tables.update_file_status import (
+  sync_detailed as update_file_status,
 )
-from ..api.tables.list_tables_v1_graphs_graph_id_tables_get import (
+from ..api.tables.list_tables import (
   sync_detailed as list_tables,
 )
-from ..api.tables.ingest_tables_v1_graphs_graph_id_tables_ingest_post import (
+from ..api.tables.ingest_tables import (
   sync_detailed as ingest_tables,
 )
 from ..models.file_upload_request import FileUploadRequest
-from ..models.file_update_request import FileUpdateRequest
+from ..models.file_status_update import FileStatusUpdate
 from ..models.bulk_ingest_request import BulkIngestRequest
 
 logger = logging.getLogger(__name__)
@@ -95,7 +95,7 @@ class TableIngestClient:
     This method handles the complete 3-step upload process:
     1. Get presigned upload URL
     2. Upload file to S3
-    3. Update file metadata
+    3. Mark file as 'uploaded' (backend validates, calculates size/row count)
 
     Args:
         graph_id: The graph ID
@@ -104,7 +104,7 @@ class TableIngestClient:
         options: Upload options
 
     Returns:
-        UploadResult with upload details
+        UploadResult with upload details (size/row count calculated by backend)
     """
     if options is None:
       options = UploadOptions()
@@ -216,12 +216,10 @@ class TableIngestClient:
           # BinaryIO or file-like object
           file_or_buffer.seek(0)
           file_content = file_or_buffer.read()
-        file_size = len(file_content)
       else:
         # Read from file path
         with open(file_path, "rb") as f:
           file_content = f.read()
-        file_size = len(file_content)
 
       s3_response = self._http_client.put(
         upload_url,
@@ -230,54 +228,47 @@ class TableIngestClient:
       )
       s3_response.raise_for_status()
 
-      # Step 3: Get row count and update file metadata
+      # Step 3: Mark file as uploaded (backend validates and calculates size/row count)
       if options.on_progress:
-        options.on_progress(f"Updating file metadata for {file_name}...")
+        options.on_progress(f"Marking {file_name} as uploaded...")
 
-      try:
-        import pyarrow.parquet as pq
-
-        if is_buffer:
-          # Read from buffer for row count
-          if hasattr(file_or_buffer, "seek"):
-            file_or_buffer.seek(0)
-          parquet_table = pq.read_table(file_or_buffer)
-        else:
-          # Read from file path
-          parquet_table = pq.read_table(file_path)
-
-        row_count = parquet_table.num_rows
-      except ImportError:
-        logger.warning(
-          "pyarrow not installed, row count will be estimated from file size"
-        )
-        # Rough estimate: ~100 bytes per row for typical data
-        row_count = file_size // 100
-
-      metadata_update = FileUpdateRequest(
-        file_size_bytes=file_size, row_count=row_count
-      )
+      status_update = FileStatusUpdate(status="uploaded")
 
       kwargs = {
         "graph_id": graph_id,
         "file_id": file_id,
         "client": client,
-        "body": metadata_update,
+        "body": status_update,
       }
 
-      update_response = update_file(**kwargs)
+      update_response = update_file_status(**kwargs)
 
       if not update_response.parsed:
+        logger.error(
+          f"No parsed response from update_file_status. Status code: {update_response.status_code}"
+        )
         return UploadResult(
           file_id=file_id,
-          file_size=file_size,
-          row_count=row_count,
+          file_size=0,
+          row_count=0,
           table_name=table_name,
           file_name=file_name,
           success=False,
-          error="Failed to update file metadata",
+          error="Failed to complete file upload",
         )
 
+      response_data = update_response.parsed
+
+      if isinstance(response_data, dict):
+        file_size = response_data.get("file_size_bytes", 0)
+        row_count = response_data.get("row_count", 0)
+      elif hasattr(response_data, "additional_properties"):
+        file_size = response_data.additional_properties.get("file_size_bytes", 0)
+        row_count = response_data.additional_properties.get("row_count", 0)
+      else:
+        file_size = getattr(response_data, "file_size_bytes", 0)
+        row_count = getattr(response_data, "row_count", 0)
+
       if options.on_progress:
         options.on_progress(
           f"✅ Uploaded {file_name} ({file_size:,} bytes, {row_count:,} rows)"

robosystems_client/models/__init__.py

@@ -75,9 +75,7 @@ from .cypher_query_request import CypherQueryRequest
 from .cypher_query_request_parameters_type_0 import CypherQueryRequestParametersType0
 from .database_health_response import DatabaseHealthResponse
 from .database_info_response import DatabaseInfoResponse
-from .delete_file_v1_graphs_graph_id_tables_files_file_id_delete_response_delete_file_v1_graphs_graph_id_tables_files_file_id_delete import (
-  DeleteFileV1GraphsGraphIdTablesFilesFileIdDeleteResponseDeleteFileV1GraphsGraphIdTablesFilesFileIdDelete,
-)
+from .delete_file_response import DeleteFileResponse
 from .delete_subgraph_request import DeleteSubgraphRequest
 from .delete_subgraph_response import DeleteSubgraphResponse
 from .detailed_transactions_response import DetailedTransactionsResponse
@@ -93,7 +91,8 @@ from .enhanced_credit_transaction_response_metadata import (
 from .error_response import ErrorResponse
 from .exchange_token_request import ExchangeTokenRequest
 from .exchange_token_request_metadata_type_0 import ExchangeTokenRequestMetadataType0
-from .file_update_request import FileUpdateRequest
+from .file_info import FileInfo
+from .file_status_update import FileStatusUpdate
 from .file_upload_request import FileUploadRequest
 from .file_upload_response import FileUploadResponse
 from .forgot_password_request import ForgotPasswordRequest
@@ -115,9 +114,7 @@ from .get_current_auth_user_response_getcurrentauthuser import (
 from .get_current_graph_bill_response_getcurrentgraphbill import (
   GetCurrentGraphBillResponseGetcurrentgraphbill,
 )
-from .get_file_info_v1_graphs_graph_id_tables_files_file_id_get_response_get_file_info_v1_graphs_graph_id_tables_files_file_id_get import (
-  GetFileInfoV1GraphsGraphIdTablesFilesFileIdGetResponseGetFileInfoV1GraphsGraphIdTablesFilesFileIdGet,
-)
+from .get_file_info_response import GetFileInfoResponse
 from .get_graph_billing_history_response_getgraphbillinghistory import (
   GetGraphBillingHistoryResponseGetgraphbillinghistory,
 )
@@ -164,9 +161,7 @@ from .link_token_request_options_type_0 import LinkTokenRequestOptionsType0
 from .link_token_request_provider_type_0 import LinkTokenRequestProviderType0
 from .list_connections_provider_type_0 import ListConnectionsProviderType0
 from .list_subgraphs_response import ListSubgraphsResponse
-from .list_table_files_v1_graphs_graph_id_tables_table_name_files_get_response_list_table_files_v1_graphs_graph_id_tables_table_name_files_get import (
-  ListTableFilesV1GraphsGraphIdTablesTableNameFilesGetResponseListTableFilesV1GraphsGraphIdTablesTableNameFilesGet,
-)
+from .list_table_files_response import ListTableFilesResponse
 from .login_request import LoginRequest
 from .logout_user_response_logoutuser import LogoutUserResponseLogoutuser
 from .mcp_tool_call import MCPToolCall
@@ -249,8 +244,8 @@ from .table_query_response import TableQueryResponse
 from .tier_upgrade_request import TierUpgradeRequest
 from .transaction_summary_response import TransactionSummaryResponse
 from .update_api_key_request import UpdateAPIKeyRequest
-from .update_file_v1_graphs_graph_id_tables_files_file_id_patch_response_update_file_v1_graphs_graph_id_tables_files_file_id_patch import (
-  UpdateFileV1GraphsGraphIdTablesFilesFileIdPatchResponseUpdateFileV1GraphsGraphIdTablesFilesFileIdPatch,
+from .update_file_status_response_updatefilestatus import (
+  UpdateFileStatusResponseUpdatefilestatus,
 )
 from .update_password_request import UpdatePasswordRequest
 from .update_user_request import UpdateUserRequest
@@ -339,7 +334,7 @@ __all__ = (
   "CypherQueryRequestParametersType0",
   "DatabaseHealthResponse",
   "DatabaseInfoResponse",
-  "DeleteFileV1GraphsGraphIdTablesFilesFileIdDeleteResponseDeleteFileV1GraphsGraphIdTablesFilesFileIdDelete",
+  "DeleteFileResponse",
   "DeleteSubgraphRequest",
   "DeleteSubgraphResponse",
   "DetailedTransactionsResponse",
@@ -351,7 +346,8 @@ __all__ = (
   "ErrorResponse",
   "ExchangeTokenRequest",
   "ExchangeTokenRequestMetadataType0",
-  "FileUpdateRequest",
+  "FileInfo",
+  "FileStatusUpdate",
   "FileUploadRequest",
   "FileUploadResponse",
   "ForgotPasswordRequest",
@@ -361,7 +357,7 @@ __all__ = (
   "GetBackupDownloadUrlResponseGetbackupdownloadurl",
   "GetCurrentAuthUserResponseGetcurrentauthuser",
   "GetCurrentGraphBillResponseGetcurrentgraphbill",
-  "GetFileInfoV1GraphsGraphIdTablesFilesFileIdGetResponseGetFileInfoV1GraphsGraphIdTablesFilesFileIdGet",
+  "GetFileInfoResponse",
   "GetGraphBillingHistoryResponseGetgraphbillinghistory",
   "GetGraphLimitsResponseGetgraphlimits",
   "GetGraphMonthlyBillResponseGetgraphmonthlybill",
@@ -390,7 +386,7 @@ __all__ = (
   "LinkTokenRequestProviderType0",
   "ListConnectionsProviderType0",
   "ListSubgraphsResponse",
-  "ListTableFilesV1GraphsGraphIdTablesTableNameFilesGetResponseListTableFilesV1GraphsGraphIdTablesTableNameFilesGet",
+  "ListTableFilesResponse",
   "LoginRequest",
   "LogoutUserResponseLogoutuser",
   "MCPToolCall",
@@ -455,7 +451,7 @@ __all__ = (
   "TierUpgradeRequest",
   "TransactionSummaryResponse",
   "UpdateAPIKeyRequest",
-  "UpdateFileV1GraphsGraphIdTablesFilesFileIdPatchResponseUpdateFileV1GraphsGraphIdTablesFilesFileIdPatch",
+  "UpdateFileStatusResponseUpdatefilestatus",
   "UpdatePasswordRequest",
   "UpdateUserRequest",
   "UserAnalyticsResponse",