robosystems-client 0.2.22__py3-none-any.whl → 0.2.24__py3-none-any.whl

robosystems_client/api/agent/auto_select_agent.py

@@ -134,7 +134,7 @@ def sync_detailed(
   **Execution Strategies (automatic):**
   - Fast operations (<5s): Immediate synchronous response
   - Medium operations (5-30s): SSE streaming with progress updates
-  - Long operations (>30s): Async Celery worker with operation tracking
+  - Long operations (>30s): Background queue with operation tracking
 
   **Response Mode Override:**
   Use query parameter `?mode=sync|async` to override automatic strategy selection.
@@ -228,7 +228,7 @@ def sync(
   **Execution Strategies (automatic):**
   - Fast operations (<5s): Immediate synchronous response
   - Medium operations (5-30s): SSE streaming with progress updates
-  - Long operations (>30s): Async Celery worker with operation tracking
+  - Long operations (>30s): Background queue with operation tracking
 
   **Response Mode Override:**
   Use query parameter `?mode=sync|async` to override automatic strategy selection.
@@ -317,7 +317,7 @@ async def asyncio_detailed(
   **Execution Strategies (automatic):**
   - Fast operations (<5s): Immediate synchronous response
   - Medium operations (5-30s): SSE streaming with progress updates
-  - Long operations (>30s): Async Celery worker with operation tracking
+  - Long operations (>30s): Background queue with operation tracking
 
   **Response Mode Override:**
   Use query parameter `?mode=sync|async` to override automatic strategy selection.
@@ -409,7 +409,7 @@ async def asyncio(
   **Execution Strategies (automatic):**
   - Fast operations (<5s): Immediate synchronous response
   - Medium operations (5-30s): SSE streaming with progress updates
-  - Long operations (>30s): Async Celery worker with operation tracking
+  - Long operations (>30s): Background queue with operation tracking
 
   **Response Mode Override:**
   Use query parameter `?mode=sync|async` to override automatic strategy selection.

robosystems_client/api/agent/execute_specific_agent.py

@@ -124,7 +124,7 @@ def sync_detailed(
   **Execution Strategies (automatic):**
   - Fast operations (<5s): Immediate synchronous response
   - Medium operations (5-30s): SSE streaming with progress updates
-  - Long operations (>30s): Async Celery worker with operation tracking
+  - Long operations (>30s): Background queue with operation tracking
 
   **Response Mode Override:**
   Use query parameter `?mode=sync|async` to override automatic strategy selection.
@@ -179,7 +179,7 @@ def sync(
   **Execution Strategies (automatic):**
   - Fast operations (<5s): Immediate synchronous response
   - Medium operations (5-30s): SSE streaming with progress updates
-  - Long operations (>30s): Async Celery worker with operation tracking
+  - Long operations (>30s): Background queue with operation tracking
 
   **Response Mode Override:**
   Use query parameter `?mode=sync|async` to override automatic strategy selection.
@@ -229,7 +229,7 @@ async def asyncio_detailed(
   **Execution Strategies (automatic):**
   - Fast operations (<5s): Immediate synchronous response
   - Medium operations (5-30s): SSE streaming with progress updates
-  - Long operations (>30s): Async Celery worker with operation tracking
+  - Long operations (>30s): Background queue with operation tracking
 
   **Response Mode Override:**
   Use query parameter `?mode=sync|async` to override automatic strategy selection.
@@ -282,7 +282,7 @@ async def asyncio(
   **Execution Strategies (automatic):**
   - Fast operations (<5s): Immediate synchronous response
   - Medium operations (5-30s): SSE streaming with progress updates
-  - Long operations (>30s): Async Celery worker with operation tracking
+  - Long operations (>30s): Background queue with operation tracking
 
   **Response Mode Override:**
   Use query parameter `?mode=sync|async` to override automatic strategy selection.

robosystems_client/api/files/update_file.py

@@ -103,7 +103,7 @@ def sync_detailed(
   **What Happens (status='uploaded'):**
   1. File validated in S3
   2. Row count calculated
-  3. DuckDB staging triggered immediately (Celery task)
+  3. DuckDB staging triggered immediately (background task)
   4. If ingest_to_graph=true, graph ingestion queued
   5. File queryable in DuckDB within seconds
 
@@ -165,7 +165,7 @@ def sync(
   **What Happens (status='uploaded'):**
   1. File validated in S3
   2. Row count calculated
-  3. DuckDB staging triggered immediately (Celery task)
+  3. DuckDB staging triggered immediately (background task)
   4. If ingest_to_graph=true, graph ingestion queued
   5. File queryable in DuckDB within seconds
 
@@ -222,7 +222,7 @@ async def asyncio_detailed(
   **What Happens (status='uploaded'):**
   1. File validated in S3
   2. Row count calculated
-  3. DuckDB staging triggered immediately (Celery task)
+  3. DuckDB staging triggered immediately (background task)
   4. If ingest_to_graph=true, graph ingestion queued
   5. File queryable in DuckDB within seconds
 
@@ -282,7 +282,7 @@ async def asyncio(
   **What Happens (status='uploaded'):**
   1. File validated in S3
   2. Row count calculated
-  3. DuckDB staging triggered immediately (Celery task)
+  3. DuckDB staging triggered immediately (background task)
   4. If ingest_to_graph=true, graph ingestion queued
   5. File queryable in DuckDB within seconds
 

robosystems_client/extensions/README.md

@@ -1,6 +1,6 @@
 # RoboSystems Python Client Extensions
 
-🚀 **Production-Ready Extensions** for the RoboSystems Financial Knowledge Graph API
+**Production-Ready Extensions** for the RoboSystems Financial Knowledge Graph API
 
 [![Python 3.7+](https://img.shields.io/badge/python-3.7+-blue.svg)](https://www.python.org/downloads/)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
@@ -17,7 +17,7 @@ The RoboSystems Python Client Extensions provide enhanced functionality for the
 - **Caching** with TTL and LRU eviction
 - **Full Async/Await Support** throughout
 
-## 🚀 Quick Start
+## Quick Start
 
 ### Installation
 
@@ -84,7 +84,7 @@ async def main():
 asyncio.run(main())
 ```
 
-## 🔐 Authentication
+## Authentication
 
 ### API Key Authentication (Recommended)
 
@@ -139,7 +139,7 @@ dev_ext = create_extensions(
 )
 ```
 
-## 🛠 Advanced Features
+## Advanced Features
 
 ### Query Builder
 
@@ -182,7 +182,7 @@ print(f"Complexity: {cost['complexity_category']} (score: {cost['complexity_scor
 
 # Get optimization recommendations  
 for rec in cost['recommendations']:
-    print(f"💡 {rec}")
+    print(f"Tip: {rec}")
 ```
 
 ### Result Processing
@@ -271,7 +271,7 @@ client.connect("operation_id")
 client.close()
 ```
 
-## 📊 Examples
+## Examples
 
 ### Financial Data Analysis
 
@@ -332,7 +332,7 @@ if final_batch:
     process_transaction_batch(final_batch)
     total_processed += len(final_batch)
 
-print(f"✅ Processed {total_processed:,} transactions total")
+print(f"Processed {total_processed:,} transactions total")
 ```
 
 ### Error Handling
@@ -361,7 +361,7 @@ except Exception as e:
     print(f"Query failed: {e}")
 ```
 
-## ⚡ Performance Optimization
+## Performance Optimization
 
 ### Connection Pooling
 
@@ -391,21 +391,21 @@ query = "MATCH (c:Company) WHERE c.revenue > 1000000 RETURN c"
 # Check syntax
 validation = validate_cypher_query(query)
 if not validation['valid']:
-    print("❌ Query has syntax errors:", validation['issues'])
+    print("Query has syntax errors:", validation['issues'])
     
 # Estimate cost
 cost = estimate_query_cost(query)
-print(f"📊 Query complexity: {cost['complexity_category']}")
+print(f"Query complexity: {cost['complexity_category']}")
 
 # Follow recommendations
 for rec in cost['recommendations']:
-    print(f"💡 Optimization tip: {rec}")
+    print(f"Optimization tip: {rec}")
 
 # Execute only if reasonable complexity
 if cost['complexity_category'] in ['low', 'medium']:
     result = extensions.execute_query("graph_id", query)
 else:
-    print("⚠️ Query may be too expensive - consider optimization")
+    print("Query may be too expensive - consider optimization")
 ```
 
 ### Caching Strategy
@@ -425,7 +425,7 @@ else:
     cache = results_cache
 ```
 
-## 🧪 Testing
+## Testing
 
 Run the test suite:
 
@@ -457,7 +457,7 @@ def test_query_execution():
         assert result["data"] == [{"count": 100}]
 ```
 
-## 🔧 Configuration
+## Configuration
 
 ### Environment Variables
 
@@ -492,7 +492,7 @@ config = RoboSystemsExtensionConfig(
 extensions = RoboSystemsExtensions(config)
 ```
 
-## 📚 API Reference
+## API Reference
 
 ### Core Classes
 
@@ -528,7 +528,7 @@ extensions = RoboSystemsExtensions(config)
 - **`format_duration(milliseconds)`** - Human-readable time formatting
 - **`create_extensions(method, **kwargs)`** - Extensions factory
 
-## 🐛 Troubleshooting
+## Troubleshooting
 
 ### Common Issues
 
@@ -545,9 +545,9 @@ pip install pandas # For DataFrame conversion (optional)
 extensions = AuthenticatedExtensions("your-api-key")
 try:
     result = extensions.execute_query("graph_id", "MATCH (n) RETURN count(n) LIMIT 1")
-    print("✅ Authentication successful")
+    print("Authentication successful")
 except Exception as e:
-    print(f"❌ Auth failed: {e}")
+    print(f"Auth failed: {e}")
 ```
 
 **Connection Issues**
@@ -585,11 +585,7 @@ logging.basicConfig(level=logging.DEBUG)
 extensions = AuthenticatedExtensions("your-key")
 ```
 
-## 📄 License
-
-MIT License - see [LICENSE](LICENSE) file for details.
-
-## 🤝 Contributing
+## Contributing
 
 1. Fork the repository
 2. Create a feature branch
@@ -597,7 +593,7 @@ MIT License - see [LICENSE](LICENSE) file for details.
 4. Run the test suite: `python run_tests.py`
 5. Submit a pull request
 
-## 📞 Support
+## Support
 
 - **API Reference**: [api.robosystems.ai](https://api.robosystems.ai)
 - **Issues**: [GitHub Issues](https://github.com/RoboFinSystems/robosystems-python-client/issues)

robosystems_client/extensions/agent_client.py

@@ -177,7 +177,7 @@ class AgentClient:
               else datetime.now().isoformat(),
             )
 
-        # Check if this is a queued response (async Celery execution)
+        # Check if this is a queued response (async background task execution)
         is_queued = False
         queued_response = None
 

robosystems_client/extensions/graph_client.py

@@ -1,13 +1,17 @@
 """Graph Management Client
 
 Provides high-level graph management operations with automatic operation monitoring.
+Supports both SSE (Server-Sent Events) for real-time updates and polling fallback.
 """
 
 from dataclasses import dataclass
 from typing import Dict, Any, Optional, Callable
 import time
+import json
 import logging
 
+import httpx
+
 logger = logging.getLogger(__name__)
 
 
@@ -62,10 +66,14 @@ class GraphClient:
     timeout: int = 60,
     poll_interval: int = 2,
     on_progress: Optional[Callable[[str], None]] = None,
+    use_sse: bool = True,
   ) -> str:
     """
     Create a graph and wait for completion.
 
+    Uses SSE (Server-Sent Events) for real-time progress updates with
+    automatic fallback to polling if SSE connection fails.
+
     Args:
         metadata: Graph metadata
         initial_entity: Optional initial entity data
@@ -73,8 +81,9 @@ class GraphClient:
             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
+        poll_interval: Time between status checks in seconds (for polling fallback)
         on_progress: Callback for progress updates
+        use_sse: Whether to try SSE first (default True). Falls back to polling on failure.
 
     Returns:
         graph_id when creation completes
@@ -84,7 +93,6 @@ class GraphClient:
     """
     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
 
@@ -151,13 +159,165 @@ class GraphClient:
         on_progress(f"Graph created: {graph_id}")
       return graph_id
 
-    # Otherwise, poll operation until complete
+    # Otherwise, wait for operation to 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})")
 
+    # Try SSE first, fall back to polling
+    if use_sse:
+      try:
+        return self._wait_with_sse(operation_id, timeout, on_progress)
+      except Exception as e:
+        logger.debug(f"SSE connection failed, falling back to polling: {e}")
+        if on_progress:
+          on_progress("SSE unavailable, using polling...")
+
+    # Fallback to polling
+    return self._wait_with_polling(
+      operation_id, timeout, poll_interval, on_progress, client
+    )
+
+  def _wait_with_sse(
+    self,
+    operation_id: str,
+    timeout: int,
+    on_progress: Optional[Callable[[str], None]],
+  ) -> str:
+    """
+    Wait for operation completion using SSE stream.
+
+    Args:
+        operation_id: Operation ID to monitor
+        timeout: Maximum time to wait in seconds
+        on_progress: Callback for progress updates
+
+    Returns:
+        graph_id when operation completes
+
+    Raises:
+        RuntimeError: If operation fails
+        TimeoutError: If operation times out
+    """
+    stream_url = f"{self.base_url}/v1/operations/{operation_id}/stream"
+    headers = {"X-API-Key": self.token, "Accept": "text/event-stream"}
+
+    with httpx.Client(timeout=httpx.Timeout(timeout + 5.0)) as http_client:
+      with http_client.stream("GET", stream_url, headers=headers) as response:
+        if response.status_code != 200:
+          raise RuntimeError(f"SSE connection failed: {response.status_code}")
+
+        start_time = time.time()
+        event_type = None
+        event_data = ""
+
+        for line in response.iter_lines():
+          # Check timeout
+          if time.time() - start_time > timeout:
+            raise TimeoutError(f"Graph creation timed out after {timeout}s")
+
+          line = line.strip()
+
+          if not line:
+            # Empty line = end of event, process it
+            if event_type and event_data:
+              result = self._process_sse_event(event_type, event_data, on_progress)
+              if result is not None:
+                return result
+            event_type = None
+            event_data = ""
+            continue
+
+          if line.startswith("event:"):
+            event_type = line[6:].strip()
+          elif line.startswith("data:"):
+            event_data = line[5:].strip()
+          # Ignore other lines (comments, id, retry, etc.)
+
+    raise TimeoutError(f"SSE stream ended without completion after {timeout}s")
+
+  def _process_sse_event(
+    self,
+    event_type: str,
+    event_data: str,
+    on_progress: Optional[Callable[[str], None]],
+  ) -> Optional[str]:
+    """
+    Process a single SSE event.
+
+    Returns:
+        graph_id if operation completed, None to continue waiting
+
+    Raises:
+        RuntimeError: If operation failed
+    """
+    try:
+      data = json.loads(event_data)
+    except json.JSONDecodeError:
+      logger.debug(f"Failed to parse SSE event data: {event_data}")
+      return None
+
+    if event_type == "operation_progress":
+      if on_progress:
+        message = data.get("message", "Processing...")
+        percent = data.get("progress_percent")
+        if percent is not None:
+          on_progress(f"{message} ({percent:.0f}%)")
+        else:
+          on_progress(message)
+      return None
+
+    elif event_type == "operation_completed":
+      result = data.get("result", {})
+      graph_id = result.get("graph_id") if isinstance(result, dict) else 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 event_type == "operation_error":
+      error = data.get("error", "Unknown error")
+      raise RuntimeError(f"Graph creation failed: {error}")
+
+    elif event_type == "operation_cancelled":
+      reason = data.get("reason", "Operation was cancelled")
+      raise RuntimeError(f"Graph creation cancelled: {reason}")
+
+    # Ignore other event types (keepalive, etc.)
+    return None
+
+  def _wait_with_polling(
+    self,
+    operation_id: str,
+    timeout: int,
+    poll_interval: int,
+    on_progress: Optional[Callable[[str], None]],
+    client: Any,
+  ) -> str:
+    """
+    Wait for operation completion using polling.
+
+    Args:
+        operation_id: Operation ID to monitor
+        timeout: Maximum time to wait in seconds
+        poll_interval: Time between status checks
+        on_progress: Callback for progress updates
+        client: Authenticated HTTP client
+
+    Returns:
+        graph_id when operation completes
+
+    Raises:
+        RuntimeError: If operation fails
+        TimeoutError: If operation times out
+    """
+    from ..api.operations.get_operation_status import sync_detailed as get_status
+
     max_attempts = timeout // poll_interval
     for attempt in range(max_attempts):
       time.sleep(poll_interval)

robosystems_client/extensions/materialization_client.py

@@ -15,6 +15,7 @@ from ..api.materialize.get_materialization_status import (
   sync_detailed as get_materialization_status,
 )
 from ..models.materialize_request import MaterializeRequest
+from .operation_client import OperationClient, OperationProgress, MonitorOptions
 
 logger = logging.getLogger(__name__)
 
@@ -27,6 +28,7 @@ class MaterializationOptions:
   rebuild: bool = False
   force: bool = False
   on_progress: Optional[Callable[[str], None]] = None
+  timeout: Optional[int] = 600  # 10 minute default timeout
 
 
 @dataclass
@@ -66,6 +68,14 @@ class MaterializationClient:
     self.base_url = config["base_url"]
     self.headers = config.get("headers", {})
     self.token = config.get("token")
+    self._operation_client = None
+
+  @property
+  def operation_client(self) -> OperationClient:
+    """Get or create the operation client for SSE monitoring."""
+    if self._operation_client is None:
+      self._operation_client = OperationClient(self.config)
+    return self._operation_client
 
   def materialize(
     self,
@@ -75,13 +85,13 @@ class MaterializationClient:
     """
     Materialize graph from DuckDB staging tables.
 
-    Rebuilds the complete graph database from the current state of DuckDB
-    staging tables. Automatically discovers all tables, materializes them in
-    the correct order (nodes before relationships), and clears the staleness flag.
+    Submits a materialization job to Dagster and monitors progress via SSE.
+    The operation runs asynchronously on the server but this method waits
+    for completion and returns the final result.
 
     Args:
         graph_id: Graph database identifier
-        options: Materialization options (ignore_errors, rebuild, force)
+        options: Materialization options (ignore_errors, rebuild, force, timeout)
 
     Returns:
         MaterializationResult with detailed execution information
@@ -96,7 +106,7 @@ class MaterializationClient:
 
     try:
       if options.on_progress:
-        options.on_progress("Starting graph materialization...")
+        options.on_progress("Submitting materialization job...")
 
       request = MaterializeRequest(
         ignore_errors=options.ignore_errors,
@@ -125,6 +135,7 @@ class MaterializationClient:
 
       response = materialize_graph(**kwargs)
 
+      # Handle non-200 status codes
       if response.status_code != 200 or not response.parsed:
         error_msg = f"Materialization failed: {response.status_code}"
         if hasattr(response, "content"):
@@ -148,25 +159,68 @@ class MaterializationClient:
           error=error_msg,
         )
 
+      # Get the operation_id from the queued response
       result_data = response.parsed
+      operation_id = result_data.operation_id
 
       if options.on_progress:
-        options.on_progress(
-          f"✅ Materialization complete: {len(result_data.tables_materialized)} tables, "
-          f"{result_data.total_rows:,} rows in {result_data.execution_time_ms:.2f}ms"
-        )
-
-      return MaterializationResult(
-        status=result_data.status,
-        was_stale=result_data.was_stale,
-        stale_reason=result_data.stale_reason,
-        tables_materialized=result_data.tables_materialized,
-        total_rows=result_data.total_rows,
-        execution_time_ms=result_data.execution_time_ms,
-        message=result_data.message,
-        success=True,
+        options.on_progress(f"Materialization queued (operation: {operation_id})")
+
+      # Monitor the operation via SSE until completion
+      def on_sse_progress(progress: OperationProgress):
+        if options.on_progress:
+          msg = progress.message
+          if progress.percentage is not None:
+            msg += f" ({progress.percentage:.0f}%)"
+          options.on_progress(msg)
+
+      monitor_options = MonitorOptions(
+        on_progress=on_sse_progress,
+        timeout=options.timeout,
       )
 
+      op_result = self.operation_client.monitor_operation(operation_id, monitor_options)
+
+      # Convert operation result to materialization result
+      if op_result.status.value == "completed":
+        # Extract details from SSE completion event result
+        sse_result = op_result.result or {}
+
+        if options.on_progress:
+          tables = sse_result.get("tables_materialized", [])
+          rows = sse_result.get("total_rows", 0)
+          time_ms = sse_result.get("execution_time_ms", 0)
+          options.on_progress(
+            f"✅ Materialization complete: {len(tables)} tables, "
+            f"{rows:,} rows in {time_ms:.2f}ms"
+          )
+
+        return MaterializationResult(
+          status="success",
+          was_stale=sse_result.get("was_stale", False),
+          stale_reason=sse_result.get("stale_reason"),
+          tables_materialized=sse_result.get("tables_materialized", []),
+          total_rows=sse_result.get("total_rows", 0),
+          execution_time_ms=sse_result.get(
+            "execution_time_ms", op_result.execution_time_ms or 0
+          ),
+          message=sse_result.get("message", "Graph materialized successfully"),
+          success=True,
+        )
+      else:
+        # Operation failed or was cancelled
+        return MaterializationResult(
+          status=op_result.status.value,
+          was_stale=False,
+          stale_reason=None,
+          tables_materialized=[],
+          total_rows=0,
+          execution_time_ms=op_result.execution_time_ms or 0,
+          message=op_result.error or f"Operation {op_result.status.value}",
+          success=False,
+          error=op_result.error,
+        )
+
     except Exception as e:
       logger.error(f"Materialization failed: {e}")
       return MaterializationResult(

robosystems_client/extensions/sse_client.py

@@ -104,8 +104,10 @@ class SSEClient:
 
     try:
       self.client = httpx.Client(timeout=self.config.timeout)
-      self._response = self.client.stream("GET", url, params=params, headers=headers)
-      self._response.__enter__()
+      self._context_manager = self.client.stream(
+        "GET", url, params=params, headers=headers
+      )
+      self._response = self._context_manager.__enter__()
 
       self.reconnect_attempts = 0
       self.emit("connected", None)
@@ -124,11 +126,9 @@ class SSEClient:
 
     try:
       event_buffer = {"event": None, "data": [], "id": None, "retry": None}
-      print("[SSE DEBUG] Starting to process events...")
 
       for line in self._response.iter_lines():
         if self.closed:
-          print("[SSE DEBUG] Stream closed, breaking out of loop")
           break
 
         line = line.strip()
@@ -136,7 +136,6 @@ class SSEClient:
         # Empty line indicates end of event
         if not line:
           if event_buffer["data"] or event_buffer["event"]:
-            print(f"[SSE DEBUG] Dispatching event: {event_buffer.get('event')}")
             self._dispatch_event(event_buffer)
           event_buffer = {"event": None, "data": [], "id": None, "retry": None}
           continue
@@ -172,13 +171,9 @@ class SSEClient:
 
       # Handle final event if stream ends without empty line
       if event_buffer["data"] or event_buffer["event"]:
-        print("[SSE DEBUG] Dispatching final event after stream end")
         self._dispatch_event(event_buffer)
 
-      print("[SSE DEBUG] Event processing loop ended")
-
     except Exception as error:
-      print(f"[SSE DEBUG] Exception in event processing: {error}")
       if not self.closed:
         self.emit("error", error)
 
@@ -214,13 +209,14 @@ class SSEClient:
     # Emit typed event
     self.emit(event_type, parsed_data)
 
-    # Check for completion events
+    # Check for completion events - just set flag, don't close from within loop
+    # The loop will break on next iteration and close() will be called in finally
     if event_type in [
       EventType.OPERATION_COMPLETED.value,
       EventType.OPERATION_ERROR.value,
       EventType.OPERATION_CANCELLED.value,
     ]:
-      self.close()
+      self.closed = True
 
   def _handle_error(
     self, error: Exception, operation_id: str, from_sequence: int
@@ -285,12 +281,13 @@ class SSEClient:
     """Close the SSE connection"""
     self.closed = True
 
-    if self._response:
+    if hasattr(self, "_context_manager") and self._context_manager:
       try:
-        self._response.__exit__(None, None, None)
+        self._context_manager.__exit__(None, None, None)
       except Exception:
         pass
-      self._response = None
+      self._context_manager = None
+    self._response = None
 
     if self.client:
       self.client.close()
@@ -334,10 +331,10 @@ class AsyncSSEClient:
 
     try:
       self.client = httpx.AsyncClient(timeout=self.config.timeout)
-      self._response = await self.client.stream(
+      self._context_manager = self.client.stream(
         "GET", url, params=params, headers=headers
       )
-      await self._response.__aenter__()
+      self._response = await self._context_manager.__aenter__()
 
       self.reconnect_attempts = 0
       self.emit("connected", None)
@@ -401,13 +398,9 @@ class AsyncSSEClient:
 
       # Handle final event if stream ends without empty line
       if event_buffer["data"] or event_buffer["event"]:
-        print("[SSE DEBUG] Dispatching final event after stream end")
         self._dispatch_event(event_buffer)
 
-      print("[SSE DEBUG] Event processing loop ended")
-
     except Exception as error:
-      print(f"[SSE DEBUG] Exception in event processing: {error}")
       if not self.closed:
         self.emit("error", error)
 
@@ -443,13 +436,14 @@ class AsyncSSEClient:
     # Emit typed event
     self.emit(event_type, parsed_data)
 
-    # Check for completion events
+    # Check for completion events - just set flag, don't close from within loop
+    # The loop will break on next iteration and close() will be called in finally
     if event_type in [
       EventType.OPERATION_COMPLETED.value,
       EventType.OPERATION_ERROR.value,
       EventType.OPERATION_CANCELLED.value,
     ]:
-      asyncio.create_task(self.close())
+      self.closed = True
 
   async def _handle_error(
     self, error: Exception, operation_id: str, from_sequence: int
@@ -512,12 +506,13 @@ class AsyncSSEClient:
     """Close the SSE connection (async)"""
     self.closed = True
 
-    if self._response:
+    if hasattr(self, "_context_manager") and self._context_manager:
       try:
-        await self._response.__aexit__(None, None, None)
+        await self._context_manager.__aexit__(None, None, None)
       except Exception:
         pass
-      self._response = None
+      self._context_manager = None
+    self._response = None
 
     if self.client:
       await self.client.aclose()

robosystems_client/extensions/subgraph_workspace_client.py

@@ -197,7 +197,9 @@ class SubgraphWorkspaceClient:
         sse_url = f"{self.api._base_url}/v1/operations/{operation_id}/stream"
         headers = {"X-API-Key": self.api.token}
 
-        async with httpx.AsyncClient() as client:
+        # Use longer timeout for SSE streaming (Dagster jobs can take time)
+        timeout = httpx.Timeout(connect=30.0, read=120.0, write=30.0, pool=30.0)
+        async with httpx.AsyncClient(timeout=timeout) as client:
           async with client.stream("GET", sse_url, headers=headers) as sse_response:
             async for line in sse_response.aiter_lines():
               if line.startswith("data: "):

robosystems_client/models/materialize_response.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Mapping
-from typing import Any, TypeVar, cast
+from typing import Any, TypeVar
 
 from attrs import define as _attrs_define
 from attrs import field as _attrs_field
@@ -13,102 +13,64 @@ T = TypeVar("T", bound="MaterializeResponse")
 
 @_attrs_define
 class MaterializeResponse:
-  """
+  """Response for queued materialization operation.
+
+  Example:
+      {'graph_id': 'kg_abc123', 'message': 'Materialization queued. Monitor via SSE stream.', 'operation_id':
+          '550e8400-e29b-41d4-a716-446655440000', 'status': 'queued'}
+
   Attributes:
-      status (str): Materialization status
       graph_id (str): Graph database identifier
-      was_stale (bool): Whether graph was stale before materialization
-      tables_materialized (list[str]): List of tables successfully materialized
-      total_rows (int): Total rows materialized across all tables
-      execution_time_ms (float): Total materialization time
+      operation_id (str): SSE operation ID for progress tracking
       message (str): Human-readable status message
-      stale_reason (None | str | Unset): Reason graph was stale
+      status (str | Unset): Operation status Default: 'queued'.
   """
 
-  status: str
   graph_id: str
-  was_stale: bool
-  tables_materialized: list[str]
-  total_rows: int
-  execution_time_ms: float
+  operation_id: str
   message: str
-  stale_reason: None | str | Unset = UNSET
+  status: str | Unset = "queued"
   additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
 
   def to_dict(self) -> dict[str, Any]:
-    status = self.status
-
     graph_id = self.graph_id
 
-    was_stale = self.was_stale
-
-    tables_materialized = self.tables_materialized
-
-    total_rows = self.total_rows
-
-    execution_time_ms = self.execution_time_ms
+    operation_id = self.operation_id
 
     message = self.message
 
-    stale_reason: None | str | Unset
-    if isinstance(self.stale_reason, Unset):
-      stale_reason = UNSET
-    else:
-      stale_reason = self.stale_reason
+    status = self.status
 
     field_dict: dict[str, Any] = {}
     field_dict.update(self.additional_properties)
     field_dict.update(
       {
-        "status": status,
         "graph_id": graph_id,
-        "was_stale": was_stale,
-        "tables_materialized": tables_materialized,
-        "total_rows": total_rows,
-        "execution_time_ms": execution_time_ms,
+        "operation_id": operation_id,
         "message": message,
       }
     )
-    if stale_reason is not UNSET:
-      field_dict["stale_reason"] = stale_reason
+    if status is not UNSET:
+      field_dict["status"] = status
 
     return field_dict
 
   @classmethod
   def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
     d = dict(src_dict)
-    status = d.pop("status")
-
     graph_id = d.pop("graph_id")
 
-    was_stale = d.pop("was_stale")
-
-    tables_materialized = cast(list[str], d.pop("tables_materialized"))
-
-    total_rows = d.pop("total_rows")
-
-    execution_time_ms = d.pop("execution_time_ms")
+    operation_id = d.pop("operation_id")
 
     message = d.pop("message")
 
-    def _parse_stale_reason(data: object) -> None | str | Unset:
-      if data is None:
-        return data
-      if isinstance(data, Unset):
-        return data
-      return cast(None | str | Unset, data)
-
-    stale_reason = _parse_stale_reason(d.pop("stale_reason", UNSET))
+    status = d.pop("status", UNSET)
 
     materialize_response = cls(
-      status=status,
       graph_id=graph_id,
-      was_stale=was_stale,
-      tables_materialized=tables_materialized,
-      total_rows=total_rows,
-      execution_time_ms=execution_time_ms,
+      operation_id=operation_id,
       message=message,
-      stale_reason=stale_reason,
+      status=status,
     )
 
     materialize_response.additional_properties = d

robosystems_client/models/token_pricing.py

@@ -14,12 +14,12 @@ class TokenPricing:
   """AI token pricing for a specific model.
 
   Attributes:
-      input_per_1k_tokens (int): Credits per 1K input tokens
-      output_per_1k_tokens (int): Credits per 1K output tokens
+      input_per_1k_tokens (float): Credits per 1K input tokens
+      output_per_1k_tokens (float): Credits per 1K output tokens
   """
 
-  input_per_1k_tokens: int
-  output_per_1k_tokens: int
+  input_per_1k_tokens: float
+  output_per_1k_tokens: float
   additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
 
   def to_dict(self) -> dict[str, Any]:

{robosystems_client-0.2.22.dist-info → robosystems_client-0.2.24.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: robosystems-client
-Version: 0.2.22
+Version: 0.2.24
 Summary: Python Client for RoboSystems financial graph database API
 Author: RFS LLC
 License: MIT
@@ -33,8 +33,8 @@ Requires-Dist: pydantic>=2.11.7
 Requires-Dist: python-dateutil>=2.8.0
 Requires-Dist: typing-extensions>=4.0.0
 Provides-Extra: all
-Requires-Dist: pandas>=1.5.0; extra == 'all'
-Requires-Dist: pyarrow>=10.0.0; extra == 'all'
+Requires-Dist: pandas>=2.1.0; extra == 'all'
+Requires-Dist: pyarrow>=17.0.0; extra == 'all'
 Provides-Extra: dev
 Requires-Dist: basedpyright>=1.21.0; extra == 'dev'
 Requires-Dist: build>=1.0.0; extra == 'dev'
@@ -45,11 +45,11 @@ Requires-Dist: pytest>=8.3.5; extra == 'dev'
 Requires-Dist: ruff>=0.12; extra == 'dev'
 Requires-Dist: twine>=5.0.0; extra == 'dev'
 Provides-Extra: extensions
-Requires-Dist: pandas>=1.5.0; extra == 'extensions'
-Requires-Dist: pyarrow>=10.0.0; extra == 'extensions'
+Requires-Dist: pandas>=2.1.0; extra == 'extensions'
+Requires-Dist: pyarrow>=17.0.0; extra == 'extensions'
 Provides-Extra: tables
-Requires-Dist: pandas>=1.5.0; extra == 'tables'
-Requires-Dist: pyarrow>=10.0.0; extra == 'tables'
+Requires-Dist: pandas>=2.1.0; extra == 'tables'
+Requires-Dist: pyarrow>=17.0.0; extra == 'tables'
 Description-Content-Type: text/markdown
 
 # RoboSystems Python Client

{robosystems_client-0.2.22.dist-info → robosystems_client-0.2.24.dist-info}/RECORD

@@ -6,9 +6,9 @@ robosystems_client/sdk-config.yaml,sha256=Y_A8qSC2zHLYy6d443Rlgdkw2GleOSFjYvq_Qm
 robosystems_client/types.py,sha256=p37iLHBX-cpNh7jZmbG4Xl1UI6hexvRS-HkY_QBlZWE,1301
 robosystems_client/api/__init__.py,sha256=zTSiG_ujSjAqWPyc435YXaX9XTlpMjiJWBbV-f-YtdA,45
 robosystems_client/api/agent/__init__.py,sha256=5vd9uJWAjRqa9xzxzYkLD1yoZ12Ld_bAaNB5WX4fbE8,56
-robosystems_client/api/agent/auto_select_agent.py,sha256=yeav6I1-DAL7g-E0f3AhClva_BEsKRXL0HQzKvnDroU,16408
+robosystems_client/api/agent/auto_select_agent.py,sha256=uHoIjrKoA6e7GyQu1DMAIPGZ1JAx-xA8sYodyLTFUfU,16396
 robosystems_client/api/agent/batch_process_queries.py,sha256=d0htQmRZoTig8XnSQHINhPDf10w_kgPfTSBNGDUrg1c,6697
-robosystems_client/api/agent/execute_specific_agent.py,sha256=70Y9dVy84z-xuQlwDFUUULi-mb-tBCmDO67Avmny6eo,9440
+robosystems_client/api/agent/execute_specific_agent.py,sha256=pZWf6dsWtMjJ2qkV0tOLEhWB8s9HlsEXUtu2DZKv_RE,9428
 robosystems_client/api/agent/get_agent_metadata.py,sha256=MS4MCA5fQE2F654iVTTnWGxckDx3YlPvYScMtsyKkAk,5730
 robosystems_client/api/agent/list_agents.py,sha256=BZ0H9ZfpwyXY57TPvkMVSsm72d6JcOY9qDUU4LxTvyE,6183
 robosystems_client/api/agent/recommend_agent.py,sha256=fy4cm-DA_FHiDsX9gtrHfITOmHJFGNekOkxAhPWz6I8,6601
@@ -67,7 +67,7 @@ robosystems_client/api/files/create_file_upload.py,sha256=OVMbGuJ004lqYrHmUFzr2g
 robosystems_client/api/files/delete_file.py,sha256=wAt0wHIYgX1D57HwQvUUb1GMr2GQzXNq3wz0MuCSI3I,9815
 robosystems_client/api/files/get_file.py,sha256=cZv2YurkjF8bhBzC_pprblPUguunV0ciOjuB4ygQS1Q,10558
 robosystems_client/api/files/list_files.py,sha256=3JnJ4gwZS7k4KdjnHp7fv0jq8JuTGzV0xRhZcMEfBcA,10541
-robosystems_client/api/files/update_file.py,sha256=PTFUwT8qsI8tUFo7Cno4iG7ouicIBGYWme0BiLy6Ogc,9027
+robosystems_client/api/files/update_file.py,sha256=X80_TbYmQvL7mLdQejCy-UzUiyVx96SpCE2EHNJJSr8,9043
 robosystems_client/api/graph_health/__init__.py,sha256=5vd9uJWAjRqa9xzxzYkLD1yoZ12Ld_bAaNB5WX4fbE8,56
 robosystems_client/api/graph_health/get_database_health.py,sha256=rGMjLUjD5DfIGv1Qu8BARWRafZhTqyrsb65YCIFe1k8,8738
 robosystems_client/api/graph_info/__init__.py,sha256=5vd9uJWAjRqa9xzxzYkLD1yoZ12Ld_bAaNB5WX4fbE8,56
@@ -141,20 +141,20 @@ robosystems_client/api/user/update_user_password.py,sha256=0couyI63dlspqvtO62cI4
 robosystems_client/api/views/__init__.py,sha256=5vd9uJWAjRqa9xzxzYkLD1yoZ12Ld_bAaNB5WX4fbE8,56
 robosystems_client/api/views/create_view.py,sha256=aBjk0FsAj6A3K2iqjE7DcTLKjDRsxegFbs5RC0hmbqA,6316
 robosystems_client/api/views/save_view.py,sha256=T7PUUNvp0JIeOYWG2n8Nm4Y9hhEWV7i-Gl2Vl8Oz4Ls,9224
-robosystems_client/extensions/README.md,sha256=qfHFjdgA_J-zNXziNZE6M1MKJiwVkocBi01w_HhvzEk,16136
+robosystems_client/extensions/README.md,sha256=NUkPte5GM8rJkT-tVrVXbujbKVxP-sw-QocYaLZfAL0,15971
 robosystems_client/extensions/__init__.py,sha256=eTuJQGygQTOWC51YVhJOWUWFUMLcPo7MpZ0H3GaxoR0,7076
-robosystems_client/extensions/agent_client.py,sha256=Db2C4hrakVsf6ScnBcNk6rte3Kwn4cQBEHsR_joWMTs,17750
+robosystems_client/extensions/agent_client.py,sha256=kvC5UO5JjTpsVB_1ARFN2H7MSDZYjw016zJrnNTMZUM,17759
 robosystems_client/extensions/auth_integration.py,sha256=ABOJ8aVjfHehNGNzim1iR9-Cdh7Mr22ce-WgWWeqJt0,6535
 robosystems_client/extensions/dataframe_utils.py,sha256=gK1bgkVqBF0TvWVdGQvqWrt-ur_Rw11j8uNtMoulLWE,12312
 robosystems_client/extensions/element_mapping_client.py,sha256=yuh0QPQBPM33E7r6QWWDiKm3T4TfCdbn2kvO3Jlw4Cs,18516
 robosystems_client/extensions/extensions.py,sha256=IqTuVJNt6vvO5LhnH-9SexCF5k9Abqig1sBdY4SdD1I,6858
 robosystems_client/extensions/file_client.py,sha256=bknqmbYrIqxc2PnWskeDvcWYP2bemBQaxY_4A6zr5LY,12326
-robosystems_client/extensions/graph_client.py,sha256=OBi0xj0SLIRKLeSu_DiGt2ZakCmhggvNrMP3jdRfEgQ,10326
-robosystems_client/extensions/materialization_client.py,sha256=Pvq4Gz5fS4B2eTxFrpCGTj7-diOZwOtWnNo7wKT9mVE,6811
+robosystems_client/extensions/graph_client.py,sha256=2Y9DaIH3pkzofS7ZujrDxzgBGSov0_uBUsM2B0jWrLw,15294
+robosystems_client/extensions/materialization_client.py,sha256=3PzzEYn5NqENbb4PggGb1Z31SWmcW9D2mJ83uem22_c,8923
 robosystems_client/extensions/operation_client.py,sha256=B1qju-wWQrnrnVJixKGgsA_KEInviwJwdlJxzm_i7P0,13359
 robosystems_client/extensions/query_client.py,sha256=cX3e8EBoTeg4Lwm6edJYRULM2UmGpfqNX3f48S8TQbE,19430
-robosystems_client/extensions/sse_client.py,sha256=XvQIq3JQ0Yiax11E7cwclhupShYOpEMURM2cYQodiz8,15058
-robosystems_client/extensions/subgraph_workspace_client.py,sha256=Ioc7FNJEKaD_kAJBeymwtFlVI-U9t47RouD5ibUHv4g,24036
+robosystems_client/extensions/sse_client.py,sha256=rqYs3mfnLhE3kwdhEWcbeiAFfT_6AFZRY5rpXpV5KHg,14956
+robosystems_client/extensions/subgraph_workspace_client.py,sha256=ISMo7xkOB1Py35njIV_qjw94RJJKEZk8Ag4RGkII3-o,24208
 robosystems_client/extensions/table_client.py,sha256=hegYiPlR-HMacyinDDACZaMeNn0bamkYtbkzDkVawVI,4521
 robosystems_client/extensions/token_utils.py,sha256=qCK_s1vBzRnSYwtgncPZRLJVIw3WXmzqNTWjdEEpdgs,10899
 robosystems_client/extensions/utils.py,sha256=vhmUnEsq-UEAMgNhmkqlbJg4oJj096QPiHALEHJ-y4A,16207
@@ -300,7 +300,7 @@ robosystems_client/models/list_table_files_response.py,sha256=D4SuwlQy11oOhEq9k2
 robosystems_client/models/login_request.py,sha256=8LbzTIfoUWAKzkEYs7E6AXr4TEPAeejoX8kIMBRm7XQ,1579
 robosystems_client/models/logout_user_response_logoutuser.py,sha256=dpvqyIqCDQdXVyx-gDyl3qex2KuyjJUE0OuldQcL4k0,1257
 robosystems_client/models/materialize_request.py,sha256=t3TXLcPJabN_lSx6BwEeku2v2zGceHFRVSS6TpK59Iw,1483
-robosystems_client/models/materialize_response.py,sha256=RQSN9CqXKL1HvRlZP7DGj_unMWMp4vTcsFv_BBjsIMg,3540
+robosystems_client/models/materialize_response.py,sha256=Luuf3FjUj1cZQ0QmcONR5ulMbKKvjhL5SkkMx7lBp-Q,2386
 robosystems_client/models/materialize_status_response.py,sha256=eXNmPgudtCneoVb2rbSW9DGGTG10DFoGN-oaiHuYjLo,5402
 robosystems_client/models/mcp_tool_call.py,sha256=dz--5OucQNcbFT7GaUR1uw-ZgSxHI0-zVYmz0J9fixE,2215
 robosystems_client/models/mcp_tool_call_arguments.py,sha256=FX9fVbI_8tzM1XF4oKaH7Tf2WFChsSuIHn03yZY0Q5w,1245
@@ -395,7 +395,7 @@ robosystems_client/models/table_info.py,sha256=vjOwl4srG9ktgg_ExG58kruAOh727Ithc
 robosystems_client/models/table_list_response.py,sha256=Yq6CMAEpWFGsLVGSXGOH4Q5ylqycqbFb-NfIWk8HPKU,2035
 robosystems_client/models/table_query_request.py,sha256=jcrXMey2DEebEgalsvVxJLtTNMMQYkBKIYPMDshd0i8,1918
 robosystems_client/models/table_query_response.py,sha256=HZcsK0XjPN7XukPStVyH3DDAsYVKuY1so3ft4dDU7-c,2304
-robosystems_client/models/token_pricing.py,sha256=r2ai68qdKuquv0cAE-1uB1P3g2kX66tayy7AkeFKswY,1877
+robosystems_client/models/token_pricing.py,sha256=VR5CI6czu1b0amw4W0B1FxwTtIiztVvZqHmfho9Ac-Q,1885
 robosystems_client/models/transaction_summary_response.py,sha256=zP6d-iypgggzkRTrdK35MTIKugWl_UVt4bKFmQx6tB0,3693
 robosystems_client/models/upcoming_invoice.py,sha256=n2RyBnypJpkWp1jWEt73BB80YIlO9tBQCZXHNpGKOnw,3518
 robosystems_client/models/update_api_key_request.py,sha256=JEGm8NzYAh_5AzxBMZDCpYak9QbDbb85lPQ6qZlzwRE,2506
@@ -414,7 +414,7 @@ robosystems_client/models/view_axis_config_member_labels_type_0.py,sha256=kkzpHx
 robosystems_client/models/view_config.py,sha256=HQnqYjLMXRhjZLOc5ypwILriMFKuvPzu0hPQi2vyNoM,3795
 robosystems_client/models/view_source.py,sha256=h66cASj-P_-qOptKv26uAIe9PtIewU2nTs42Ls-lFFk,4098
 robosystems_client/models/view_source_type.py,sha256=KpgczHUeOinV01jdLvytZ2URKwcsRcp1doPx2D3USyw,169
-robosystems_client-0.2.22.dist-info/METADATA,sha256=yEG5FHqaUwnCVaVxGPrcXOAMY3DMYelvRAvJK3p5ikE,3950
-robosystems_client-0.2.22.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-robosystems_client-0.2.22.dist-info/licenses/LICENSE,sha256=LjFqQPU4eQh7jAQ04SmE9eC0j74HCdXvzbo0hjW4mWo,1063
-robosystems_client-0.2.22.dist-info/RECORD,,
+robosystems_client-0.2.24.dist-info/METADATA,sha256=H_jYC_RHMKkZjsZfrf7Tq3ohH0OaOK3xVQDQV7ieX4A,3950
+robosystems_client-0.2.24.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+robosystems_client-0.2.24.dist-info/licenses/LICENSE,sha256=LjFqQPU4eQh7jAQ04SmE9eC0j74HCdXvzbo0hjW4mWo,1063
+robosystems_client-0.2.24.dist-info/RECORD,,