raindrop-ai 0.0.28__py3-none-any.whl → 0.0.30__py3-none-any.whl

raindrop/analytics.py

@@ -1,7 +1,9 @@
 import sys
 import time
 import threading
-from typing import Union, List, Dict, Optional, Literal, Any
+import os
+from contextlib import contextmanager
+from typing import Callable, Union, List, Dict, Optional, Literal, Any
 import requests
 from datetime import datetime, timezone
 import logging
@@ -23,6 +25,39 @@ from raindrop.models import (
 )
 from raindrop.interaction import Interaction
 from raindrop.redact import perform_pii_redaction
+import weakref
+import urllib.parse
+
+from traceloop.sdk import Traceloop
+from traceloop.sdk.tracing.tracing import TracerWrapper
+from opentelemetry.trace import get_current_span
+from traceloop.sdk.decorators import (
+    task as tlp_task,
+    workflow as tlp_workflow,
+    TraceloopSpanKindValues,
+    F,
+)
+
+__all__ = [
+    # Configuration functions
+    "set_debug_logs",
+    "set_redact_pii", 
+    "init",
+    
+    "identify",
+    "track_ai",
+    "track_signal",
+    "begin",
+    "resume_interaction",
+    
+    "interaction",
+    "task", 
+    "tool",
+    "set_span_properties",
+    
+    "flush",
+    "shutdown",
+]
 
 
 # Configure logging
@@ -40,6 +75,7 @@ buffer = []
 flush_lock = threading.Lock()
 debug_logs = False
 redact_pii = False
+_tracing_enabled = False
 flush_thread = None
 shutdown_event = threading.Event()
 max_ingest_size_bytes = 1 * 1024 * 1024  # 1 MB
@@ -236,6 +272,11 @@ def shutdown():
     if flush_thread:
         flush_thread.join(timeout=10)
     flush()  # Final flush to ensure all events are sent
+    if _tracing_enabled:
+        try:
+            TracerWrapper().flush()
+        except Exception as e:
+            logger.debug(f"Could not flush TracerWrapper during shutdown: {e}")
 
 
 def _check_write_key():
@@ -365,6 +406,14 @@ def track_signal(
     save_to_buffer({"type": "signals/track", "data": data})
 
 
+INTERACTION_TRACE_ID_REGISTRY: weakref.WeakValueDictionary[int, Interaction] = (
+    weakref.WeakValueDictionary()
+)
+INTERACTION_EVENT_ID_REGISTRY: weakref.WeakValueDictionary[str, Interaction] = (
+    weakref.WeakValueDictionary()
+)
+
+
 def begin(
     user_id: str,
     event: str,
@@ -387,6 +436,10 @@ def begin(
     # Combine properties with initial_fields, giving precedence to initial_fields if keys clash
     final_properties = (properties or {}).copy()
 
+    current_trace_id = _safe_current_trace_id()
+    if current_trace_id is not None:
+        final_properties["trace_id"] = f"{current_trace_id:032x}"
+
     partial_event = PartialTrackAIEvent(
         event_id=eid,
         user_id=user_id,
@@ -397,13 +450,145 @@ def begin(
         attachments=attachments,
     )
 
+    span_attributes = {
+        "user_id": user_id,
+        "convo_id": convo_id,
+        "event": event,
+        "event_id": eid,
+    }
+    if _tracing_enabled:
+        Traceloop.set_association_properties(
+            {k: v for k, v in span_attributes.items() if v is not None}
+        )
+
+    interaction = Interaction(eid)
+    INTERACTION_EVENT_ID_REGISTRY[eid] = interaction
+    if current_trace_id is not None and current_trace_id != 0:
+        INTERACTION_TRACE_ID_REGISTRY[current_trace_id] = interaction
+
     _track_ai_partial(partial_event)
-    return Interaction(eid)
+    return interaction
+
+
+@contextmanager
+def _temp_env(key: str, value: str):
+    """Temporarily sets an environment variable. Hacky helper to deal with traceloop's BS"""
+    orig = os.environ.get(key)
+    os.environ[key] = value
+    try:
+        yield
+    finally:
+        if orig is None:
+            del os.environ[key]
+        else:
+            os.environ[key] = orig
+
+
+def init(
+    api_key: str,
+    tracing_enabled: bool = False,
+    **traceloop_kwargs,
+):
+    """Initialize Raindrop with Traceloop integration."""
+    global write_key
+    write_key = api_key
+
+    global _tracing_enabled
+    _tracing_enabled = tracing_enabled
+
+    if not _tracing_enabled:
+        return
+
+    parsed_url = urllib.parse.urlparse(api_url)
+    api_endpoint = f"{parsed_url.scheme}://{parsed_url.netloc}"
+
+    with _temp_env("TRACELOOP_METRICS_ENABLED", "false"):
+        Traceloop.init(
+            api_endpoint=api_endpoint,
+            api_key=api_key,
+            telemetry_enabled=False,
+            **traceloop_kwargs,
+        )
+
+
+def _safe_current_trace_id() -> int | None:
+    """Return current trace id or None if unavailable."""
+    try:
+        trace_id = get_current_span().get_span_context().trace_id
+    except Exception:
+        return None
+    return trace_id if trace_id else None
+
+
+def interaction(
+    name: Optional[str] = None,
+    version: Optional[int] = None,
+    method_name: Optional[str] = None,
+) -> Callable[[F], F]:
+    return tlp_workflow(
+        name=name,
+        version=version,
+        method_name=method_name,
+        tlp_span_kind=TraceloopSpanKindValues.WORKFLOW,
+    )
+
+
+def task(
+    name: Optional[str] = None,
+    version: Optional[int] = None,
+    method_name: Optional[str] = None,
+    tlp_span_kind: Optional[TraceloopSpanKindValues] = TraceloopSpanKindValues.TASK,
+) -> Callable[[F], F]:
+    return tlp_task(
+        name=name,
+        version=version,
+        method_name=method_name,
+        tlp_span_kind=tlp_span_kind,
+    )
+
+
+def tool(
+    name: Optional[str] = None,
+    version: Optional[int] = None,
+    method_name: Optional[str] = None,
+) -> Callable[[F], F]:
+    return tlp_task(
+        name=name,
+        version=version,
+        method_name=method_name,
+        tlp_span_kind=TraceloopSpanKindValues.TOOL,
+    )
+
+
+def set_span_properties(properties: Dict[str, Any]) -> None:
+    """
+    Set association properties on the current span for tracing.
+
+    Args:
+        properties: Dictionary of properties to associate with the current span
+    """
+    if not _tracing_enabled:
+        return
+
+    Traceloop.set_association_properties(properties)
+
+
+def resume_interaction(event_id: str | None = None) -> Interaction:
+    """Return an Interaction associated with the current trace or given event_id."""
+
+    if event_id is not None:
+        if (interaction := INTERACTION_EVENT_ID_REGISTRY.get(event_id)) is not None:
+            return interaction
+        return Interaction(event_id)
 
+    if (trace_id := _safe_current_trace_id()) is not None:
+        if (interaction := INTERACTION_TRACE_ID_REGISTRY.get(trace_id)) is not None:
+            return interaction
 
-def resume_interaction(event_id: str) -> Interaction:
-    """Return an Interaction wrapper for an existing eventId."""
-    return Interaction(event_id)
+    # Fallback: create a fresh Interaction when no identifiers are available
+    # TODO: Return No-Op interaction if event_id is None
+    logger.debug("No interaction found, creating a new one")
+    return Interaction()
 
 
 def _track_ai_partial(event: PartialTrackAIEvent) -> None:

raindrop/interaction.py

@@ -1,9 +1,18 @@
 from __future__ import annotations
-from typing import Any, Dict, List, Optional, Callable
+from typing import (
+    Any,
+    Dict,
+    List,
+    Optional,
+    Union,
+    Iterator,
+)
 from uuid import uuid4
+from dataclasses import dataclass
 
 from .models import Attachment, PartialTrackAIEvent
 from . import analytics as _core
+from opentelemetry import context as context_api
 
 
 class Interaction:
@@ -11,12 +20,19 @@ class Interaction:
     Thin helper returned by analytics.begin().
     Each mutator just relays a partial update back to Analytics.
     """
-    __slots__ = ("_event_id", "_analytics")
 
-    def __init__(self, event_id: Optional[str] = None):
+    __slots__ = (
+        "_event_id",
+        "_analytics",
+        "__weakref__",
+    )
+
+    def __init__(
+        self,
+        event_id: Optional[str] = None,
+    ):
         self._event_id = event_id or str(uuid4())
         self._analytics = _core
-        
 
     # -- mutators ----------------------------------------------------------- #
     def set_input(self, text: str) -> None:
@@ -34,7 +50,11 @@ class Interaction:
             PartialTrackAIEvent(event_id=self._event_id, properties=props)
         )
 
+    def set_property(self, key: str, value: Any) -> None:
+        self.set_properties({key: value})
+
     def finish(self, *, output: str | None = None, **extra) -> None:
+
         payload = PartialTrackAIEvent(
             event_id=self._event_id,
             ai_data={"output": output} if output is not None else None,
@@ -46,4 +66,4 @@ class Interaction:
     # convenience
     @property
     def id(self) -> str:
-        return self._event_id
+        return self._event_id

raindrop/models.py

@@ -1,6 +1,8 @@
 from pydantic import BaseModel, Field, ValidationError, model_validator, field_validator
 from typing import Any, Optional, Dict, Literal, List, Union
 from datetime import datetime, timezone
+from dataclasses import dataclass, field
+
 
 class _Base(BaseModel):
     model_config = dict(extra="forbid", validate_default=True)
@@ -8,10 +10,10 @@ class _Base(BaseModel):
 
 class Attachment(BaseModel):
     type: Literal["code", "text", "image", "iframe"]
-    value: str                       # URL, raw code, etc.
-    name: Optional[str] = None       # e.g. "Generated SQL"
+    value: str  # URL, raw code, etc.
+    name: Optional[str] = None  # e.g. "Generated SQL"
     role: Optional[Literal["input", "output", "context"]] = None
-    language: Optional[str] = None   # for code snippets
+    language: Optional[str] = None  # for code snippets
 
     @model_validator(mode="after")
     def _require_value(cls, values):
@@ -19,6 +21,7 @@ class Attachment(BaseModel):
             raise ValueError("value must be non-empty.")
         return values
 
+
 class AIData(_Base):
     model: Optional[str]
     input: Optional[str]
@@ -30,7 +33,8 @@ class AIData(_Base):
         if not (values.input or values.output):
             raise ValueError("Either 'input' or 'output' must be non-empty.")
         return values
-    
+
+
 class TrackAIEvent(_Base):
     event_id: Optional[str] = None
     user_id: str
@@ -38,7 +42,9 @@ class TrackAIEvent(_Base):
     ai_data: AIData
     properties: Dict[str, Any] = Field(default_factory=dict)
     timestamp: datetime = Field(
-        default_factory=lambda: datetime.now(timezone.utc).replace(microsecond=0).isoformat()
+        default_factory=lambda: datetime.now(timezone.utc)
+        .replace(microsecond=0)
+        .isoformat()
     )
     attachments: Optional[List[Attachment]] = None
 
@@ -55,8 +61,10 @@ class TrackAIEvent(_Base):
 
 # --- Signal Tracking Models --- #
 
+
 class BaseSignal(_Base):
     """Base model for signal events, containing common fields."""
+
     event_id: str
     signal_name: str
     timestamp: datetime = Field(
@@ -73,41 +81,52 @@ class BaseSignal(_Base):
             raise ValueError(f"'{info.field_name}' must be a non-empty string.")
         return v
 
+
 class DefaultSignal(BaseSignal):
     """Model for default signal events."""
+
     signal_type: Literal["default"] = "default"
 
+
 class FeedbackSignal(BaseSignal):
     """Model for feedback signal events, requiring a comment."""
+
     signal_type: Literal["feedback"]
 
     @model_validator(mode="after")
     def _check_comment_in_properties(cls, values):
         # Check properties safely after potential initialization
         # Use getattr to safely access properties, returning None if not present
-        props = getattr(values, 'properties', None) 
+        props = getattr(values, "properties", None)
         if not isinstance(props, dict):
-             raise ValueError("'properties' must be a dictionary for feedback signals.")
+            raise ValueError("'properties' must be a dictionary for feedback signals.")
         comment = props.get("comment")
         if not comment or not isinstance(comment, str) or not comment.strip():
-            raise ValueError("'properties' must contain a non-empty string 'comment' for feedback signals.")
+            raise ValueError(
+                "'properties' must contain a non-empty string 'comment' for feedback signals."
+            )
         return values
 
+
 class EditSignal(BaseSignal):
     """Model for edit signal events, requiring after content."""
+
     signal_type: Literal["edit"]
 
     @model_validator(mode="after")
     def _check_after_in_properties(cls, values):
         # Check properties safely after potential initialization
-        props = getattr(values, 'properties', None)
+        props = getattr(values, "properties", None)
         if not isinstance(props, dict):
-             raise ValueError("'properties' must be a dictionary for edit signals.")
+            raise ValueError("'properties' must be a dictionary for edit signals.")
         after = props.get("after")
         if not after or not isinstance(after, str) or not after.strip():
-             raise ValueError("'properties' must contain a non-empty string 'after' for edit signals.")
+            raise ValueError(
+                "'properties' must contain a non-empty string 'after' for edit signals."
+            )
         return values
 
+
 # Discriminated Union for Signal Events
 # Pydantic will automatically use the 'signal_type' field to determine which model to use.
 SignalEvent = Union[DefaultSignal, FeedbackSignal, EditSignal]
@@ -117,18 +136,21 @@ SignalEvent = Union[DefaultSignal, FeedbackSignal, EditSignal]
 
 class PartialAIData(_Base):
     """Looser version for incremental updates."""
+
     model: Optional[str] = None
     input: Optional[str] = None
     output: Optional[str] = None
-    convo_id: Optional[str] = None 
+    convo_id: Optional[str] = None
+
 
 class PartialTrackAIEvent(_Base):
     """Accepts *any subset* of TrackAIEvent fields."""
-    event_id: str                                 # always required for merge-key
+
+    event_id: str  # always required for merge-key
     user_id: Optional[str] = None
     event: Optional[str] = None
     ai_data: Optional[PartialAIData] = None
     timestamp: Optional[datetime] = None
     properties: Optional[Dict[str, Any]] = None
     attachments: Optional[List[Attachment]] = None
-    is_pending: Optional[bool] = True
+    is_pending: Optional[bool] = True

raindrop/version.py

@@ -1 +1 @@
-VERSION = "0.0.28"
+VERSION = "0.0.30"

{raindrop_ai-0.0.28.dist-info → raindrop_ai-0.0.30.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: raindrop-ai
-Version: 0.0.28
+Version: 0.0.30
 Summary: Raindrop AI (Python SDK)
 License: MIT
 Author: Raindrop AI
@@ -13,6 +13,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: pydantic (>=2.09,<3)
 Requires-Dist: requests (>=2.32.3,<3.0.0)
+Requires-Dist: traceloop-sdk (==0.45.6)
 Description-Content-Type: text/markdown
 
 # Raindrop Python SDK
@@ -31,6 +32,25 @@ poetry install
 
 ## Run tests
 
+### Using pytest (recommended)
+```bash
+# Run all tests
+poetry run pytest
+
+# Run tests with verbose output
+poetry run pytest -v
+
+# Run specific test file
+poetry run pytest tests/test_trace_attributes.py
+
+# Run specific test
+poetry run pytest tests/test_trace_attributes.py::TestTraceAttributes::test_user_id_attribute_direct
+
+# Run tests with coverage
+poetry run pytest --cov=raindrop
+```
+
+### Using green (legacy)
 ```bash
 poetry run green -vv
 ```

raindrop_ai-0.0.30.dist-info/RECORD

@@ -0,0 +1,10 @@
+raindrop/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+raindrop/analytics.py,sha256=-n9vKdsTkiSid_y11bxSUl0sOLZ7drqgD9ahGOpO68E,19548
+raindrop/interaction.py,sha256=1gYFj2-oEobLtG56RZw5de3RjYr0pCXg_NqYDGx32Uc,1896
+raindrop/models.py,sha256=dEhwTtBEgwOOW_qaPEYugCLeOqvDXuBpEfi7vUANVK0,5361
+raindrop/redact.py,sha256=rMNUoI90KxOY3d_zcHAr0TFD2yQ_CDgpDz-1XJLVmHs,7658
+raindrop/version.py,sha256=3Utuv7Oo1v_YaWRMBX5-yd895qkotgnwqFPSZDHaL1g,19
+raindrop/well-known-names.json,sha256=9giJF6u6W1R0APW-Pf1dvNUU32OXQEoQ9CBQXSnA3ks,144403
+raindrop_ai-0.0.30.dist-info/METADATA,sha256=uM1mTZbxJ6ZlVjRGpsmnrWiAEXKZ4PQ3mJv8KZQSNGM,1269
+raindrop_ai-0.0.30.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+raindrop_ai-0.0.30.dist-info/RECORD,,

raindrop_ai-0.0.28.dist-info/RECORD

@@ -1,10 +0,0 @@
-raindrop/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-raindrop/analytics.py,sha256=T7oUrbtZmB0CkC2Xlp9mm_ue7UvqmnTe8w5SMhr3kKc,14618
-raindrop/interaction.py,sha256=U8OapN5tZGSMVkn-dTLum5sIg_AgSYKX9lrtwgXJOuI,1613
-raindrop/models.py,sha256=8YKuyB34P7sfawPFL3jOt5RgYlNXzf9dEsj8_2rjZKU,5293
-raindrop/redact.py,sha256=rMNUoI90KxOY3d_zcHAr0TFD2yQ_CDgpDz-1XJLVmHs,7658
-raindrop/version.py,sha256=1s8fJcG0T-ks-jhi7sPF29aUmHv4p8qr2zQ-eevL20k,19
-raindrop/well-known-names.json,sha256=9giJF6u6W1R0APW-Pf1dvNUU32OXQEoQ9CBQXSnA3ks,144403
-raindrop_ai-0.0.28.dist-info/METADATA,sha256=rGpUHqWn6m_cueg4pdBvznT8fQCX1DMu2CtWT8AzsXM,815
-raindrop_ai-0.0.28.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-raindrop_ai-0.0.28.dist-info/RECORD,,