667 lines
24 KiB
Python
667 lines
24 KiB
Python
"""
|
|
HCFS Asynchronous Client
|
|
|
|
High-level asynchronous client for HCFS API operations with WebSocket support.
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import time
|
|
from typing import List, Optional, Dict, Any, AsyncIterator, Callable
|
|
from datetime import datetime
|
|
|
|
import httpx
|
|
import websockets
|
|
from websockets.exceptions import ConnectionClosed, WebSocketException
|
|
|
|
from .models import (
|
|
Context, SearchResult, ContextFilter, PaginationOptions,
|
|
SearchOptions, ClientConfig, AnalyticsData, BatchResult, StreamEvent
|
|
)
|
|
from .exceptions import (
|
|
HCFSError, HCFSConnectionError, HCFSAuthenticationError,
|
|
HCFSNotFoundError, HCFSValidationError, HCFSStreamError, handle_api_error
|
|
)
|
|
from .utils import MemoryCache, validate_path, normalize_path
|
|
from .decorators import cached_context, retry_on_failure, rate_limited
|
|
|
|
|
|
class HCFSAsyncClient:
|
|
"""
|
|
Asynchronous HCFS API client with WebSocket streaming capabilities.
|
|
|
|
This client provides async/await support for all operations and includes
|
|
real-time streaming capabilities through WebSocket connections.
|
|
|
|
Example:
|
|
>>> import asyncio
|
|
>>> from hcfs.sdk import HCFSAsyncClient, Context
|
|
>>>
|
|
>>> async def main():
|
|
... async with HCFSAsyncClient(
|
|
... base_url="https://api.hcfs.example.com",
|
|
... api_key="your-api-key"
|
|
... ) as client:
|
|
... # Create a context
|
|
... context = Context(
|
|
... path="/docs/async_readme",
|
|
... content="Async README content",
|
|
... summary="Async documentation"
|
|
... )
|
|
... created = await client.create_context(context)
|
|
...
|
|
... # Search with async
|
|
... results = await client.search_contexts("async README")
|
|
... async for result in results:
|
|
... print(f"Found: {result.context.path}")
|
|
>>>
|
|
>>> asyncio.run(main())
|
|
"""
|
|
|
|
def __init__(self, config: Optional[ClientConfig] = None, **kwargs):
|
|
"""
|
|
Initialize async HCFS client.
|
|
|
|
Args:
|
|
config: Client configuration object
|
|
**kwargs: Configuration overrides
|
|
"""
|
|
# Merge configuration
|
|
if config:
|
|
self.config = config
|
|
else:
|
|
self.config = ClientConfig(**kwargs)
|
|
|
|
# HTTP client will be initialized in __aenter__
|
|
self.http_client: Optional[httpx.AsyncClient] = None
|
|
self.websocket: Optional[websockets.WebSocketServerProtocol] = None
|
|
self._websocket_listeners: List[Callable[[StreamEvent], None]] = []
|
|
self._websocket_task: Optional[asyncio.Task] = None
|
|
|
|
# Initialize cache
|
|
self._cache = MemoryCache(
|
|
max_size=self.config.cache.max_size,
|
|
strategy=self.config.cache.strategy,
|
|
ttl_seconds=self.config.cache.ttl_seconds
|
|
) if self.config.cache.enabled else None
|
|
|
|
# Analytics
|
|
self.analytics = AnalyticsData()
|
|
|
|
async def __aenter__(self):
|
|
"""Async context manager entry."""
|
|
await self._initialize_http_client()
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
|
"""Async context manager exit."""
|
|
await self.close()
|
|
|
|
async def _initialize_http_client(self):
|
|
"""Initialize the HTTP client with proper configuration."""
|
|
headers = {
|
|
"User-Agent": self.config.user_agent,
|
|
"Content-Type": "application/json"
|
|
}
|
|
|
|
if self.config.api_key:
|
|
headers["X-API-Key"] = self.config.api_key
|
|
elif self.config.jwt_token:
|
|
headers["Authorization"] = f"Bearer {self.config.jwt_token}"
|
|
|
|
# Configure timeouts
|
|
timeout = httpx.Timeout(
|
|
connect=self.config.timeout,
|
|
read=self.config.timeout,
|
|
write=self.config.timeout,
|
|
pool=self.config.timeout * 2
|
|
)
|
|
|
|
# Configure connection limits
|
|
limits = httpx.Limits(
|
|
max_connections=self.config.max_connections,
|
|
max_keepalive_connections=self.config.max_keepalive_connections
|
|
)
|
|
|
|
self.http_client = httpx.AsyncClient(
|
|
base_url=self.config.base_url,
|
|
headers=headers,
|
|
timeout=timeout,
|
|
limits=limits,
|
|
follow_redirects=True
|
|
)
|
|
|
|
async def health_check(self) -> Dict[str, Any]:
|
|
"""
|
|
Check API health status asynchronously.
|
|
|
|
Returns:
|
|
Health status information
|
|
|
|
Raises:
|
|
HCFSConnectionError: If health check fails
|
|
"""
|
|
try:
|
|
response = await self.http_client.get("/health")
|
|
|
|
if response.status_code == 200:
|
|
self._update_analytics("health_check", success=True)
|
|
return response.json()
|
|
else:
|
|
self._update_analytics("health_check", success=False)
|
|
handle_api_error(response)
|
|
|
|
except httpx.RequestError as e:
|
|
self._update_analytics("health_check", success=False, error=str(e))
|
|
raise HCFSConnectionError(f"Health check failed: {str(e)}")
|
|
|
|
@cached_context()
|
|
@retry_on_failure()
|
|
async def create_context(self, context: Context) -> Context:
|
|
"""
|
|
Create a new context asynchronously.
|
|
|
|
Args:
|
|
context: Context object to create
|
|
|
|
Returns:
|
|
Created context with assigned ID
|
|
|
|
Raises:
|
|
HCFSValidationError: If context data is invalid
|
|
HCFSError: If creation fails
|
|
"""
|
|
if not validate_path(context.path):
|
|
raise HCFSValidationError(f"Invalid context path: {context.path}")
|
|
|
|
context.path = normalize_path(context.path)
|
|
|
|
try:
|
|
response = await self.http_client.post(
|
|
"/api/v1/contexts",
|
|
json=context.to_create_dict()
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
data = response.json()["data"]
|
|
created_context = Context(**data)
|
|
self._update_analytics("create_context", success=True)
|
|
return created_context
|
|
else:
|
|
self._update_analytics("create_context", success=False)
|
|
handle_api_error(response)
|
|
|
|
except httpx.RequestError as e:
|
|
self._update_analytics("create_context", success=False, error=str(e))
|
|
raise HCFSConnectionError(f"Failed to create context: {str(e)}")
|
|
|
|
@cached_context()
|
|
async def get_context(self, context_id: int) -> Context:
|
|
"""
|
|
Retrieve a context by ID asynchronously.
|
|
|
|
Args:
|
|
context_id: Context identifier
|
|
|
|
Returns:
|
|
Context object
|
|
|
|
Raises:
|
|
HCFSNotFoundError: If context doesn't exist
|
|
"""
|
|
try:
|
|
response = await self.http_client.get(f"/api/v1/contexts/{context_id}")
|
|
|
|
if response.status_code == 200:
|
|
data = response.json()["data"]
|
|
context = Context(**data)
|
|
self._update_analytics("get_context", success=True)
|
|
return context
|
|
else:
|
|
self._update_analytics("get_context", success=False)
|
|
handle_api_error(response)
|
|
|
|
except httpx.RequestError as e:
|
|
self._update_analytics("get_context", success=False, error=str(e))
|
|
raise HCFSConnectionError(f"Failed to get context: {str(e)}")
|
|
|
|
async def list_contexts(self,
|
|
filter_opts: Optional[ContextFilter] = None,
|
|
pagination: Optional[PaginationOptions] = None) -> List[Context]:
|
|
"""
|
|
List contexts with filtering and pagination asynchronously.
|
|
|
|
Args:
|
|
filter_opts: Context filtering options
|
|
pagination: Pagination configuration
|
|
|
|
Returns:
|
|
List of contexts
|
|
"""
|
|
params = {}
|
|
|
|
if filter_opts:
|
|
params.update(filter_opts.to_query_params())
|
|
|
|
if pagination:
|
|
params.update(pagination.to_query_params())
|
|
|
|
try:
|
|
response = await self.http_client.get("/api/v1/contexts", params=params)
|
|
|
|
if response.status_code == 200:
|
|
data = response.json()["data"]
|
|
contexts = [Context(**ctx_data) for ctx_data in data]
|
|
self._update_analytics("list_contexts", success=True)
|
|
return contexts
|
|
else:
|
|
self._update_analytics("list_contexts", success=False)
|
|
handle_api_error(response)
|
|
|
|
except httpx.RequestError as e:
|
|
self._update_analytics("list_contexts", success=False, error=str(e))
|
|
raise HCFSConnectionError(f"Failed to list contexts: {str(e)}")
|
|
|
|
async def update_context(self, context_id: int, updates: Dict[str, Any]) -> Context:
|
|
"""
|
|
Update an existing context asynchronously.
|
|
|
|
Args:
|
|
context_id: Context identifier
|
|
updates: Fields to update
|
|
|
|
Returns:
|
|
Updated context
|
|
|
|
Raises:
|
|
HCFSNotFoundError: If context doesn't exist
|
|
HCFSValidationError: If update data is invalid
|
|
"""
|
|
try:
|
|
response = await self.http_client.put(
|
|
f"/api/v1/contexts/{context_id}",
|
|
json=updates
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
data = response.json()["data"]
|
|
updated_context = Context(**data)
|
|
self._update_analytics("update_context", success=True)
|
|
|
|
# Invalidate cache
|
|
if self._cache:
|
|
cache_key = f"get_context:{context_id}"
|
|
self._cache.remove(cache_key)
|
|
|
|
return updated_context
|
|
else:
|
|
self._update_analytics("update_context", success=False)
|
|
handle_api_error(response)
|
|
|
|
except httpx.RequestError as e:
|
|
self._update_analytics("update_context", success=False, error=str(e))
|
|
raise HCFSConnectionError(f"Failed to update context: {str(e)}")
|
|
|
|
async def delete_context(self, context_id: int) -> bool:
|
|
"""
|
|
Delete a context asynchronously.
|
|
|
|
Args:
|
|
context_id: Context identifier
|
|
|
|
Returns:
|
|
True if deletion was successful
|
|
|
|
Raises:
|
|
HCFSNotFoundError: If context doesn't exist
|
|
"""
|
|
try:
|
|
response = await self.http_client.delete(f"/api/v1/contexts/{context_id}")
|
|
|
|
if response.status_code == 200:
|
|
self._update_analytics("delete_context", success=True)
|
|
|
|
# Invalidate cache
|
|
if self._cache:
|
|
cache_key = f"get_context:{context_id}"
|
|
self._cache.remove(cache_key)
|
|
|
|
return True
|
|
else:
|
|
self._update_analytics("delete_context", success=False)
|
|
handle_api_error(response)
|
|
|
|
except httpx.RequestError as e:
|
|
self._update_analytics("delete_context", success=False, error=str(e))
|
|
raise HCFSConnectionError(f"Failed to delete context: {str(e)}")
|
|
|
|
@rate_limited(requests_per_second=10.0)
|
|
async def search_contexts(self,
|
|
query: str,
|
|
options: Optional[SearchOptions] = None) -> List[SearchResult]:
|
|
"""
|
|
Search contexts asynchronously using various search methods.
|
|
|
|
Args:
|
|
query: Search query string
|
|
options: Search configuration options
|
|
|
|
Returns:
|
|
List of search results ordered by relevance
|
|
"""
|
|
search_opts = options or SearchOptions()
|
|
|
|
request_data = {
|
|
"query": query,
|
|
**search_opts.to_request_dict()
|
|
}
|
|
|
|
try:
|
|
response = await self.http_client.post(
|
|
"/api/v1/search",
|
|
json=request_data
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
data = response.json()["data"]
|
|
results = []
|
|
|
|
for result_data in data:
|
|
context = Context(**result_data["context"])
|
|
search_result = SearchResult(
|
|
context=context,
|
|
score=result_data["score"],
|
|
explanation=result_data.get("explanation"),
|
|
highlights=result_data.get("highlights", [])
|
|
)
|
|
results.append(search_result)
|
|
|
|
self._update_analytics("search_contexts", success=True)
|
|
return sorted(results, key=lambda x: x.score, reverse=True)
|
|
else:
|
|
self._update_analytics("search_contexts", success=False)
|
|
handle_api_error(response)
|
|
|
|
except httpx.RequestError as e:
|
|
self._update_analytics("search_contexts", success=False, error=str(e))
|
|
raise HCFSConnectionError(f"Search failed: {str(e)}")
|
|
|
|
async def batch_create_contexts(self, contexts: List[Context]) -> BatchResult:
|
|
"""
|
|
Create multiple contexts in a single batch operation asynchronously.
|
|
|
|
Args:
|
|
contexts: List of contexts to create
|
|
|
|
Returns:
|
|
Batch operation results
|
|
"""
|
|
request_data = {
|
|
"contexts": [ctx.to_create_dict() for ctx in contexts]
|
|
}
|
|
|
|
start_time = time.time()
|
|
|
|
try:
|
|
response = await self.http_client.post(
|
|
"/api/v1/contexts/batch",
|
|
json=request_data,
|
|
timeout=self.config.timeout * 3 # Extended timeout for batch ops
|
|
)
|
|
|
|
execution_time = time.time() - start_time
|
|
|
|
if response.status_code == 200:
|
|
data = response.json()["data"]
|
|
|
|
result = BatchResult(
|
|
success_count=data["success_count"],
|
|
error_count=data["error_count"],
|
|
total_items=data["total_items"],
|
|
successful_items=data.get("created_ids", []),
|
|
failed_items=data.get("errors", []),
|
|
execution_time=execution_time
|
|
)
|
|
|
|
self._update_analytics("batch_create", success=True)
|
|
return result
|
|
else:
|
|
self._update_analytics("batch_create", success=False)
|
|
handle_api_error(response)
|
|
|
|
except httpx.RequestError as e:
|
|
execution_time = time.time() - start_time
|
|
self._update_analytics("batch_create", success=False, error=str(e))
|
|
|
|
return BatchResult(
|
|
success_count=0,
|
|
error_count=len(contexts),
|
|
total_items=len(contexts),
|
|
successful_items=[],
|
|
failed_items=[{"error": str(e)}],
|
|
execution_time=execution_time
|
|
)
|
|
|
|
async def get_statistics(self) -> Dict[str, Any]:
|
|
"""
|
|
Get comprehensive system statistics asynchronously.
|
|
|
|
Returns:
|
|
System statistics and metrics
|
|
"""
|
|
try:
|
|
response = await self.http_client.get("/api/v1/stats")
|
|
|
|
if response.status_code == 200:
|
|
self._update_analytics("get_statistics", success=True)
|
|
return response.json()
|
|
else:
|
|
self._update_analytics("get_statistics", success=False)
|
|
handle_api_error(response)
|
|
|
|
except httpx.RequestError as e:
|
|
self._update_analytics("get_statistics", success=False, error=str(e))
|
|
raise HCFSConnectionError(f"Failed to get statistics: {str(e)}")
|
|
|
|
async def iterate_contexts(self,
|
|
filter_opts: Optional[ContextFilter] = None,
|
|
page_size: int = 100) -> AsyncIterator[Context]:
|
|
"""
|
|
Asynchronously iterate through all contexts with automatic pagination.
|
|
|
|
Args:
|
|
filter_opts: Context filtering options
|
|
page_size: Number of contexts per page
|
|
|
|
Yields:
|
|
Context objects
|
|
"""
|
|
page = 1
|
|
|
|
while True:
|
|
pagination = PaginationOptions(page=page, page_size=page_size)
|
|
contexts = await self.list_contexts(filter_opts, pagination)
|
|
|
|
if not contexts:
|
|
break
|
|
|
|
for context in contexts:
|
|
yield context
|
|
|
|
# If we got fewer contexts than requested, we've reached the end
|
|
if len(contexts) < page_size:
|
|
break
|
|
|
|
page += 1
|
|
|
|
async def connect_websocket(self,
|
|
path_prefix: Optional[str] = None,
|
|
event_types: Optional[List[str]] = None) -> None:
|
|
"""
|
|
Connect to WebSocket for real-time updates.
|
|
|
|
Args:
|
|
path_prefix: Filter events by path prefix
|
|
event_types: List of event types to subscribe to
|
|
|
|
Raises:
|
|
HCFSStreamError: If WebSocket connection fails
|
|
"""
|
|
if self.websocket and not self.websocket.closed:
|
|
return # Already connected
|
|
|
|
# Convert HTTP URL to WebSocket URL
|
|
ws_url = self.config.base_url.replace("http://", "ws://").replace("https://", "wss://")
|
|
ws_url += "/ws"
|
|
|
|
# Add authentication headers
|
|
headers = {}
|
|
if self.config.api_key:
|
|
headers["X-API-Key"] = self.config.api_key
|
|
elif self.config.jwt_token:
|
|
headers["Authorization"] = f"Bearer {self.config.jwt_token}"
|
|
|
|
try:
|
|
self.websocket = await websockets.connect(
|
|
ws_url,
|
|
extra_headers=headers,
|
|
ping_interval=self.config.websocket.ping_interval,
|
|
ping_timeout=self.config.websocket.ping_timeout
|
|
)
|
|
|
|
# Send subscription request
|
|
subscription = {
|
|
"type": "subscribe",
|
|
"data": {
|
|
"path_prefix": path_prefix,
|
|
"event_types": event_types or ["created", "updated", "deleted"],
|
|
"filters": {}
|
|
}
|
|
}
|
|
|
|
await self.websocket.send(json.dumps(subscription))
|
|
|
|
# Start listening task
|
|
self._websocket_task = asyncio.create_task(self._websocket_listener())
|
|
|
|
except (WebSocketException, ConnectionClosed) as e:
|
|
raise HCFSStreamError(f"Failed to connect to WebSocket: {str(e)}")
|
|
|
|
async def disconnect_websocket(self) -> None:
|
|
"""Disconnect from WebSocket."""
|
|
if self._websocket_task:
|
|
self._websocket_task.cancel()
|
|
try:
|
|
await self._websocket_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
self._websocket_task = None
|
|
|
|
if self.websocket:
|
|
await self.websocket.close()
|
|
self.websocket = None
|
|
|
|
def add_event_listener(self, listener: Callable[[StreamEvent], None]) -> None:
|
|
"""
|
|
Add an event listener for WebSocket events.
|
|
|
|
Args:
|
|
listener: Function to call when events are received
|
|
"""
|
|
self._websocket_listeners.append(listener)
|
|
|
|
def remove_event_listener(self, listener: Callable[[StreamEvent], None]) -> None:
|
|
"""
|
|
Remove an event listener.
|
|
|
|
Args:
|
|
listener: Function to remove
|
|
"""
|
|
if listener in self._websocket_listeners:
|
|
self._websocket_listeners.remove(listener)
|
|
|
|
async def _websocket_listener(self) -> None:
|
|
"""Internal WebSocket message listener."""
|
|
try:
|
|
async for message in self.websocket:
|
|
try:
|
|
data = json.loads(message)
|
|
event = StreamEvent(
|
|
event_type=data.get("type", "unknown"),
|
|
data=data.get("data", {}),
|
|
timestamp=datetime.fromisoformat(data.get("timestamp", datetime.utcnow().isoformat())),
|
|
context_id=data.get("context_id"),
|
|
path=data.get("path")
|
|
)
|
|
|
|
# Notify all listeners
|
|
for listener in self._websocket_listeners:
|
|
try:
|
|
if asyncio.iscoroutinefunction(listener):
|
|
await listener(event)
|
|
else:
|
|
listener(event)
|
|
except Exception:
|
|
pass # Don't let listener errors break the connection
|
|
|
|
except json.JSONDecodeError:
|
|
pass # Ignore malformed messages
|
|
|
|
except (WebSocketException, ConnectionClosed):
|
|
# Connection was closed, attempt reconnection if configured
|
|
if self.config.websocket.auto_reconnect:
|
|
await self._attempt_websocket_reconnection()
|
|
|
|
async def _attempt_websocket_reconnection(self) -> None:
|
|
"""Attempt to reconnect WebSocket with backoff."""
|
|
for attempt in range(self.config.websocket.max_reconnect_attempts):
|
|
try:
|
|
await asyncio.sleep(self.config.websocket.reconnect_interval)
|
|
await self.connect_websocket()
|
|
return # Successfully reconnected
|
|
except Exception:
|
|
continue # Try again
|
|
|
|
# All reconnection attempts failed
|
|
raise HCFSStreamError("Failed to reconnect WebSocket after multiple attempts")
|
|
|
|
def clear_cache(self) -> None:
|
|
"""Clear all cached data."""
|
|
if self._cache:
|
|
self._cache.clear()
|
|
|
|
def get_cache_stats(self) -> Dict[str, Any]:
|
|
"""Get cache statistics."""
|
|
if self._cache:
|
|
stats = self._cache.stats()
|
|
self.analytics.cache_stats = stats
|
|
return stats
|
|
return {}
|
|
|
|
def get_analytics(self) -> AnalyticsData:
|
|
"""
|
|
Get client analytics and usage statistics.
|
|
|
|
Returns:
|
|
Analytics data including operation counts and performance metrics
|
|
"""
|
|
# Update cache stats
|
|
if self._cache:
|
|
self.analytics.cache_stats = self._cache.stats()
|
|
|
|
return self.analytics
|
|
|
|
def _update_analytics(self, operation: str, success: bool, error: Optional[str] = None):
|
|
"""Update internal analytics tracking."""
|
|
self.analytics.operation_count[operation] = self.analytics.operation_count.get(operation, 0) + 1
|
|
|
|
if not success:
|
|
error_key = error or "unknown_error"
|
|
self.analytics.error_stats[error_key] = self.analytics.error_stats.get(error_key, 0) + 1
|
|
|
|
async def close(self):
|
|
"""Close the client and cleanup resources."""
|
|
await self.disconnect_websocket()
|
|
|
|
if self.http_client:
|
|
await self.http_client.aclose() |