polos-sdk 0.1.0__py3-none-any.whl

polos/utils/serializer.py

@@ -0,0 +1,167 @@
+"""JSON serialization utilities for handling non-serializable types."""
+
+import json
+from typing import Any
+
+from pydantic import BaseModel
+
+from ..types.types import AgentResult
+
+
+def is_json_serializable(obj: Any) -> bool:
+    """Check if an object is JSON serializable by attempting json.dumps.
+
+    Args:
+        obj: Object to check
+
+    Returns:
+        True if the object is JSON serializable, False otherwise
+    """
+    try:
+        json.dumps(obj)
+        return True
+    except (TypeError, ValueError):
+        return False
+
+
+def serialize(obj: Any) -> Any:
+    """Serialize an object to a JSON-serializable object.
+
+    If the input is a Pydantic BaseModel, uses model_dump(mode="json") to serialize it to a dict.
+    Otherwise, checks if it's JSON serializable via json.dumps and raises TypeError if not.
+
+    Args:
+        obj: Object to serialize
+
+    Returns:
+        JSON-serializable value (dict, list, str, int, float, bool, None)
+
+    Raises:
+        TypeError: If the object is not a Pydantic model and not JSON serializable
+    """
+    # Handle Pydantic models
+    if isinstance(obj, BaseModel):
+        return obj.model_dump(mode="json")
+
+    # Check if it's JSON serializable
+    if not is_json_serializable(obj):
+        raise TypeError(
+            f"Object of type {type(obj).__name__} is not JSON serializable. "
+            f"If it's a Pydantic model, ensure it inherits from BaseModel."
+        )
+
+    return obj
+
+
+def json_serialize(obj: Any) -> str:
+    """Serialize an object to a JSON string.
+
+    If the input is a Pydantic BaseModel, uses model_dump_json() to serialize it to a JSON string.
+    Otherwise, uses json.dumps and raises TypeError if not.
+
+    Args:
+        obj: Object to serialize
+
+    Returns:
+        JSON string
+
+    Raises:
+        TypeError: If the object is not a Pydantic model and not JSON serializable
+    """
+    # Handle Pydantic models
+    if isinstance(obj, BaseModel):
+        return obj.model_dump_json()
+
+    # Check if it's JSON serializable
+    try:
+        return json.dumps(obj)
+    except (TypeError, ValueError) as e:
+        raise TypeError(
+            f"Object of type {type(obj).__name__} is not JSON serializable. "
+            f"If it's a Pydantic model, ensure it inherits from BaseModel."
+        ) from e
+
+
+async def deserialize(obj: Any, output_schema_name: str | None = None) -> Any:
+    """Deserialize an object from a JSON string.
+
+    Args:
+        obj: Object to deserialize
+        output_schema_name: The name of the output schema
+
+    Returns:
+        Deserialized object
+    """
+    # If output_schema_name is present, try to reconstruct the Pydantic model
+    if output_schema_name and isinstance(obj, dict):
+        try:
+            # Dynamically import the Pydantic class
+            module_path, class_name = output_schema_name.rsplit(".", 1)
+            module = __import__(module_path, fromlist=[class_name])
+            model_class = getattr(module, class_name)
+
+            # Validate the dict back to the Pydantic model
+            if issubclass(model_class, BaseModel):
+                obj = model_class.model_validate(obj)
+        except (ImportError, AttributeError, ValueError, TypeError) as e:
+            raise Exception(
+                f"Failed to reconstruct Pydantic model from output_schema_name: "
+                f"{output_schema_name}. Error: {str(e)}"
+            ) from e
+    return obj
+
+
+async def deserialize_agent_result(result: AgentResult) -> AgentResult:
+    """Deserialize an agent result to a Pydantic model."""
+
+    # Convert result to structured output schema
+    if result.result_schema and result.result is not None:
+        result.result = await _deserialize_agent_result_schema(result.result, result.result_schema)
+
+    # Convert tool results to structured output schema
+    for tool_result in result.tool_results:
+        if tool_result.result_schema and tool_result.result is not None:
+            tool_result.result = await _deserialize_agent_result_schema(
+                tool_result.result, tool_result.result_schema
+            )
+
+    return result
+
+
+async def _deserialize_agent_result_schema(result, schema: str):
+    try:
+        # Dynamically import the Pydantic class
+        module_path, class_name = schema.rsplit(".", 1)
+        module = __import__(module_path, fromlist=[class_name])
+        model_class = getattr(module, class_name)
+
+        # Validate that it's a Pydantic BaseModel
+        if issubclass(model_class, BaseModel):
+            if isinstance(result, str):
+                return model_class.model_validate_json(result)
+            elif isinstance(result, dict):
+                return model_class.model_validate(result)
+            else:
+                return result
+    except (ImportError, AttributeError, ValueError, TypeError) as e:
+        # If reconstruction fails, log warning but return dict
+        # This allows backward compatibility if the class is not available
+        import warnings
+
+        warnings.warn(
+            f"Failed to reconstruct Pydantic model '{schema}': {e}. Returning dict instead.",
+            UserWarning,
+            stacklevel=2,
+        )
+    return result
+
+
+def safe_serialize(value):
+    """Serialize with fallback for non-serializable values."""
+    try:
+        return serialize(value)
+    except (TypeError, ValueError):
+        # Fallback representations
+        if hasattr(value, "__name__"):
+            return f"<{value.__name__}>"
+        return f"<{type(value).__name__}>"

polos/utils/tracing.py

@@ -0,0 +1,27 @@
+from opentelemetry.trace import NonRecordingSpan, set_span_in_context
+
+
+def get_parent_span_context_from_execution_context(exec_context):
+    parent_context = None
+    if exec_context:
+        parent_span_context = exec_context.get("_otel_span_context")
+        if parent_span_context:
+            # Convert SpanContext to Context by wrapping in NonRecordingSpan
+            parent_span = NonRecordingSpan(parent_span_context)
+            parent_context = set_span_in_context(parent_span)
+
+    return parent_context
+
+
+def get_span_context_from_execution_context(exec_context):
+    if exec_context:
+        return exec_context.get("_otel_span_context")
+    return None
+
+
+def set_span_context_in_execution_context(exec_context, span_context):
+    if exec_context:
+        exec_context["_otel_span_context"] = span_context
+        if span_context is not None:
+            exec_context["_otel_trace_id"] = format(span_context.trace_id, "032x")
+            exec_context["_otel_span_id"] = format(span_context.span_id, "016x")

polos/utils/worker_singleton.py

@@ -0,0 +1,40 @@
+"""Singleton worker instance management for HTTP client reuse."""
+
+from typing import Any
+
+import httpx
+
+# Singleton worker instance (set by Worker when it starts)
+_current_worker: Any | None = None
+
+
+def get_worker_client() -> httpx.AsyncClient | None:
+    """Get the HTTP client from the current worker instance if available.
+
+    Returns:
+        The worker's HTTP client if a worker is running, None otherwise
+    """
+    global _current_worker
+    if _current_worker is not None and _current_worker.client is not None:
+        return _current_worker.client
+    return None
+
+
+def set_current_worker(worker: Any | None) -> None:
+    """Set the current worker instance (called by Worker when it starts/stops).
+
+    Args:
+        worker: The Worker instance, or None to clear
+    """
+    global _current_worker
+    _current_worker = worker
+
+
+def get_current_worker() -> Any | None:
+    """Get the current worker instance if available.
+
+    Returns:
+        The current Worker instance, or None if no worker is running
+    """
+    global _current_worker
+    return _current_worker