feat: Add Fireworks Python SDK support

agentops/client.py

@@ -15,7 +15,7 @@
 import traceback
 from decimal import Decimal
 from functools import cached_property
-from typing import List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
 from uuid import UUID, uuid4
 
 from termcolor import colored
@@ -247,7 +247,7 @@
         end_state_reason: Optional[str] = None,
         video: Optional[str] = None,
         is_auto_end: Optional[bool] = None,
-    ) -> Optional[Decimal]:
+    ) -> Optional[Dict[str, Any]]:
         """
         End the current session with the AgentOps service.
 
@@ -258,17 +258,15 @@
             is_auto_end (bool, optional): is this an automatic use of end_session and should be skipped with skip_auto_end_session
 
         Returns:
-            Decimal: The token cost of the session. Returns 0 if the cost is unknown.
+            Dict[str, Any]: Session statistics including duration, cost, and event counts.
         """
         session = self._safe_get_session()
         if session is None:
-            return
+            return None
         if is_auto_end and self._config.skip_auto_end_session:
-            return
-
-        token_cost = session.end_session(end_state=end_state, end_state_reason=end_state_reason, video=video)
+            return None
 
-        return token_cost
+        return session.end_session(end_state=end_state, end_state_reason=end_state_reason, video=video)
 
     def create_agent(
         self,

agentops/llms/providers/fireworks.py

@@ -0,0 +1,196 @@
+import logging
+from typing import Optional, AsyncGenerator
+import pprint
+from agentops.session import Session
+from agentops.helpers import get_ISO_time
+from agentops.event import LLMEvent
+from agentops.enums import EventType
+from .instrumented_provider import InstrumentedProvider
+
+logger = logging.getLogger(__name__)
+
+
+class FireworksProvider(InstrumentedProvider):
+    """Provider for Fireworks.ai API."""
+
+    def __init__(self, client):
+        super().__init__(client)
+        self._provider_name = "Fireworks"
+        self._original_completion = None
+        self._original_async_completion = None
+        self._session = None  # Initialize session attribute
+        self._accumulated_content = ""  # Track accumulated content for streaming
+        self._init_timestamp = None  # Track stream start time
+        logger.info(f"Initializing {self._provider_name} provider")
+
+    def set_session(self, session: Session):
+        """Set the session for event tracking."""
+        self._session = session
+        logger.debug(f"Set session {session.session_id} for {self._provider_name} provider")
+
+    def handle_response(self, response, kwargs, init_timestamp, session: Optional[Session] = None) -> dict:
+        """Handle the response from the Fireworks API."""
+        if session:
+            self._session = session
+            logger.debug(f"Updated session to {session.session_id} for {self._provider_name} provider")
+
+        try:
+            # Pass ChatML messages directly to LLMEvent
+            messages = kwargs.get("messages", [])
+            logger.debug(f"Using ChatML messages: {messages}")
+
+            # Handle streaming response
+            if kwargs.get("stream", False):
+                # Create single LLMEvent for streaming response
+                stream_event = LLMEvent(
+                    event_type=EventType.LLM.value,
+                    init_timestamp=init_timestamp,
+                    model=kwargs.get("model", "unknown"),
+                    prompt=messages,  # Pass ChatML directly
+                    prompt_tokens=0,
+                    completion_tokens=0,
+                    cost=0.0,
+                )
+
+                async def async_generator(stream):
+                    """Handle async streaming response."""
+                    self._accumulated_content = ""
+                    async for chunk in stream:
+                        try:
+                            if hasattr(chunk, "choices") and chunk.choices:
+                                content = (
+                                    chunk.choices[0].delta.content
+                                    if hasattr(chunk.choices[0].delta, "content")
+                                    else None
+                                )
+                            else:
+                                content = chunk
+
+                            if content:
+                                self._accumulated_content += content
+                                # Only create event when stream is finished
+                                if hasattr(chunk.choices[0], "finish_reason") and chunk.choices[0].finish_reason:
+                                    stream_event.completion = self._accumulated_content
+                                    stream_event.end_timestamp = get_ISO_time()
+                                    if self._session:
+                                        self._session.record(stream_event)
+                                        logger.debug(f"Recorded complete streaming response for session {self._session.session_id}")
+                                    self._accumulated_content = ""  # Reset for next stream
+                                yield content
+                        except Exception as e:
+                            logger.error(f"Error in async streaming: {str(e)}")
+                            raise
+
+                def generator(stream):
+                    """Handle synchronous streaming response."""
+                    self._accumulated_content = ""
+                    for chunk in stream:
+                        try:
+                            if hasattr(chunk, "choices") and chunk.choices:
+                                content = (
+                                    chunk.choices[0].delta.content
+                                    if hasattr(chunk.choices[0].delta, "content")
+                                    else None
+                                )
+                            else:
+                                content = chunk
+
+                            if content:
+                                self._accumulated_content += content
+                                # Only create event when stream is finished
+                                if hasattr(chunk.choices[0], "finish_reason") and chunk.choices[0].finish_reason:
+                                    stream_event.completion = self._accumulated_content
+                                    stream_event.end_timestamp = get_ISO_time()
+                                    if self._session:
+                                        self._session.record(stream_event)
+                                        logger.debug(f"Recorded complete streaming response for session {self._session.session_id}")
+                                    self._accumulated_content = ""  # Reset for next stream
+                                yield content
+                        except Exception as e:
+                            logger.error(f"Error in sync streaming: {str(e)}")
+                            raise
+
+                if hasattr(response, "__aiter__"):
+                    return async_generator(response)
+                else:
+                    return generator(response)
+
+            # Handle non-streaming response
+            if hasattr(response, "choices") and response.choices:
+                content = response.choices[0].message.content if hasattr(response.choices[0], "message") else ""
+
+                # Create LLMEvent for non-streaming response
+                non_stream_event = LLMEvent(
+                    event_type=EventType.LLM.value,
+                    init_timestamp=init_timestamp,
+                    end_timestamp=get_ISO_time(),
+                    model=kwargs.get("model", "unknown"),
+                    prompt=messages,  # Pass ChatML directly
+                    completion=content,
+                    prompt_tokens=0,
+                    completion_tokens=0,
+                    cost=0.0,
+                )
+
+                if self._session:
+                    self._session.record(non_stream_event)
+                    logger.debug(f"Recorded non-streaming response for session {self._session.session_id}")
+
+            return response
+
+        except Exception as e:
+            logger.error(f"Error handling Fireworks response: {str(e)}")
+            raise
+
+    def override(self):
+        """Override Fireworks API methods with instrumented versions."""
+        logger.info(f"Overriding {self._provider_name} provider methods")
+
+        # Store original methods
+        self._original_completion = self.client.chat.completions.create
+        self._original_async_completion = getattr(self.client.chat.completions, "acreate", None)
+
+        # Override methods
+        self._override_fireworks_completion()
+        if self._original_async_completion:
+            self._override_fireworks_async_completion()
+
+    def _override_fireworks_completion(self):
+        """Override synchronous completion method."""
+        original_create = self._original_completion
+        provider = self
+
+        def patched_function(*args, **kwargs):
+            try:
+                init_timestamp = get_ISO_time()
+                response = original_create(*args, **kwargs)
+                return provider.handle_response(response, kwargs, init_timestamp, provider._session)
+            except Exception as e:
+                logger.error(f"Error in Fireworks completion: {str(e)}")
+                raise
+
+        self.client.chat.completions.create = patched_function
+
+    def _override_fireworks_async_completion(self):
+        """Override asynchronous completion method."""
+        original_acreate = self._original_async_completion
+        provider = self
+
+        async def patched_function(*args, **kwargs):
+            try:
+                init_timestamp = get_ISO_time()
+                response = await original_acreate(*args, **kwargs)
+                return provider.handle_response(response, kwargs, init_timestamp, provider._session)
+            except Exception as e:
+                logger.error(f"Error in Fireworks async completion: {str(e)}")
+                raise
+
+        self.client.chat.completions.acreate = patched_function
+
+    def undo_override(self):
+        """Restore original Fireworks API methods."""
+        logger.info(f"Restoring original {self._provider_name} provider methods")
+        if self._original_completion:
+            self.client.chat.completions.create = self._original_completion
+        if self._original_async_completion:
+            self.client.chat.completions.acreate = self._original_async_completion

agentops/llms/tracker.py

@@ -16,6 +16,7 @@
 from .providers.anthropic import AnthropicProvider
 from .providers.mistral import MistralProvider
 from .providers.ai21 import AI21Provider
+from .providers.fireworks import FireworksProvider
 
 original_func = {}
 original_create = None
@@ -48,6 +49,9 @@
         "mistralai": {
             "1.0.1": ("chat.complete", "chat.stream"),
         },
+        "fireworks-ai": {
+            "0.1.0": ("chat.completions.create",),
+        },
         "ai21": {
             "2.0.0": (
                 "chat.completions.create",
@@ -155,6 +159,15 @@
                     else:
                         logger.warning(f"Only AI21>=2.0.0 supported. v{module_version} found.")
 
+                if api == "fireworks-ai":
+                    module_version = version(api)
+
+                    if Version(module_version) >= parse("0.1.0"):
+                        provider = FireworksProvider(self.client)
+                        provider.override()
+                    else:
+                        logger.warning(f"Only Fireworks>=0.1.0 supported. v{module_version} found.")
+
                 if api == "llama_stack_client":
                     module_version = version(api)
 
@@ -174,3 +187,4 @@
         MistralProvider(self.client).undo_override()
         AI21Provider(self.client).undo_override()
         LlamaStackClientProvider(self.client).undo_override()
+        FireworksProvider(self.client).undo_override()

agentops/session.py

@@ -284,7 +284,7 @@
         end_state: str = "Indeterminate",
         end_state_reason: Optional[str] = None,
         video: Optional[str] = None,
-    ) -> Union[Decimal, None]:
+    ) -> Union[Dict[str, Any], None]:
         with self._end_session_lock:
             if not self.is_running:
                 return None
@@ -324,8 +324,9 @@
                     finally:
                         del self._span_processor
 
-                # 5. Final session update
-                if not (analytics_stats := self.get_analytics()):
+                # 5. Final session update and get analytics
+                analytics_stats = self.get_analytics()
+                if not analytics_stats:
                     return None
 
                 analytics = (
@@ -341,6 +342,7 @@
 
             except Exception as e:
                 logger.exception(f"Error during session end: {e}")
+                return None
             finally:
                 active_sessions.remove(self)  # First thing, get rid of the session
 
@@ -350,7 +352,7 @@
                         "blue",
                     )
                 )
-            return self.token_cost
+            return analytics_stats
 
     def add_tags(self, tags: List[str]) -> None:
         """