paygent-sdk 1.0.0__py3-none-any.whl → 3.0.0__py3-none-any.whl

paygent_sdk/wrappers/gemini_wrapper.py

@@ -0,0 +1,334 @@
+"""
+Google Gemini wrapper for automatic usage tracking with Paygent.
+This wrapper intercepts Gemini API calls and automatically sends usage data to Paygent.
+
+Usage is identical to the standard google-genai SDK, with only the addition of tracking parameters.
+"""
+
+import json
+from typing import Any
+
+try:
+    from google import genai
+except ImportError:
+    raise ImportError(
+        "google-genai package is a peer-dependency. To use the Paygent wrapper around google-genai "
+        "you're assumed to already have google-genai package installed."
+    )
+
+from ..client import Client
+from ..models import UsageData, UsageDataWithStrings
+
+
+class PaygentGemini:
+    """Main wrapper class for Google Gemini that provides automatic usage tracking."""
+
+    def __init__(self, gemini_client: genai.Client, paygent_client: Client):
+        """
+        Create a new PaygentGemini wrapper.
+        
+        Args:
+            gemini_client: The GoogleGenAI client instance from google-genai
+            paygent_client: The Paygent client instance for usage tracking
+        """
+        self.gemini = gemini_client
+        self.paygent_client = paygent_client
+
+    @property
+    def models(self) -> 'ModelsWrapper':
+        """Access to models API with automatic usage tracking."""
+        return ModelsWrapper(self.gemini, self.paygent_client)
+
+
+class ModelsWrapper:
+    """Wrapper for Gemini models API."""
+
+    def __init__(self, gemini_client: genai.Client, paygent_client: Client):
+        self.gemini = gemini_client
+        self.paygent_client = paygent_client
+
+    def generate_content(
+        self,
+        *,
+        model: str,
+        contents: Any,
+        indicator: str,
+        external_agent_id: str,
+        external_customer_id: str,
+        **kwargs
+    ) -> Any:
+        """
+        Generate content with automatic usage tracking.
+        
+        Args:
+            model: The model to use
+            contents: The content to generate from
+            indicator: Indicator for the usage event
+            external_agent_id: External agent identifier
+            external_customer_id: External customer identifier
+            **kwargs: Additional Gemini parameters
+            
+        Returns:
+            The generation response from Gemini
+        """
+        # Make the Gemini API call using the standard SDK method
+        response = self.gemini.models.generate_content(
+            model=model,
+            contents=contents,
+            **kwargs
+        )
+
+        # Extract usage data from response with robust fallback mechanism
+        has_valid_usage = (
+            hasattr(response, 'usage_metadata') and
+            response.usage_metadata and
+            (response.usage_metadata.prompt_token_count > 0 or
+             response.usage_metadata.candidates_token_count > 0)
+        )
+
+        if has_valid_usage:
+            # Primary path: Use usage metadata from API response
+            usage_data = UsageData(
+                service_provider=model,
+                model=model,
+                prompt_tokens=response.usage_metadata.prompt_token_count or 0,
+                completion_tokens=response.usage_metadata.candidates_token_count or 0,
+                total_tokens=response.usage_metadata.total_token_count or 0
+            )
+
+            self.paygent_client.send_usage(
+                external_agent_id,
+                external_customer_id,
+                indicator,
+                usage_data
+            )
+        else:
+            # Fallback path: Calculate tokens from actual strings
+            prompt_string = json.dumps(contents) if not isinstance(contents, str) else contents
+            output_string = ''
+            
+            if hasattr(response, 'candidates') and response.candidates:
+                if len(response.candidates) > 0:
+                    candidate = response.candidates[0]
+                    if hasattr(candidate, 'content') and hasattr(candidate.content, 'parts'):
+                        if len(candidate.content.parts) > 0:
+                            part = candidate.content.parts[0]
+                            if hasattr(part, 'text'):
+                                output_string = part.text or ''
+
+            usage_data_with_strings = UsageDataWithStrings(
+                service_provider=model,
+                model=model,
+                prompt_string=prompt_string,
+                output_string=output_string
+            )
+
+            self.paygent_client.send_usage_with_token_string(
+                external_agent_id,
+                external_customer_id,
+                indicator,
+                usage_data_with_strings
+            )
+
+        return response
+
+    def start_chat(
+        self,
+        *,
+        model: str,
+        indicator: str,
+        external_agent_id: str,
+        external_customer_id: str,
+        **kwargs
+    ) -> 'ChatSessionWrapper':
+        """
+        Start a chat session with automatic usage tracking.
+        
+        Args:
+            model: The model to use
+            indicator: Indicator for the usage event
+            external_agent_id: External agent identifier
+            external_customer_id: External customer identifier
+            **kwargs: Additional chat configuration
+            
+        Returns:
+            A wrapped chat session with automatic tracking
+        """
+        # Start chat session using the standard SDK
+        chat_session = self.gemini.models.start_chat(model=model, **kwargs)
+
+        return ChatSessionWrapper(
+            chat_session,
+            self.paygent_client,
+            model,
+            indicator,
+            external_agent_id,
+            external_customer_id
+        )
+
+    def generate_image(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        indicator: str,
+        external_agent_id: str,
+        external_customer_id: str,
+        **kwargs
+    ) -> Any:
+        """
+        Generate images with automatic usage tracking.
+        
+        Args:
+            model: The model to use (e.g., "imagen-3.0-generate-001")
+            prompt: The prompt for image generation
+            indicator: Indicator for the usage event
+            external_agent_id: External agent identifier
+            external_customer_id: External customer identifier
+            **kwargs: Additional image generation parameters
+            
+        Returns:
+            The image generation response
+        """
+        # Make the image generation API call
+        response = self.gemini.models.generate_content(
+            model=model,
+            contents=prompt,
+            **kwargs
+        )
+
+        # Extract usage data from response with robust fallback mechanism
+        has_valid_usage = (
+            hasattr(response, 'usage_metadata') and
+            response.usage_metadata and
+            (response.usage_metadata.prompt_token_count > 0 or
+             response.usage_metadata.candidates_token_count > 0)
+        )
+
+        if has_valid_usage:
+            # Primary path: Use usage metadata from API response
+            usage_data = UsageData(
+                service_provider=model,
+                model=model,
+                prompt_tokens=response.usage_metadata.prompt_token_count or 0,
+                completion_tokens=response.usage_metadata.candidates_token_count or 0,
+                total_tokens=response.usage_metadata.total_token_count or 0
+            )
+
+            self.paygent_client.send_usage(
+                external_agent_id,
+                external_customer_id,
+                indicator,
+                usage_data
+            )
+        else:
+            # Fallback path: Calculate tokens from prompt string
+            usage_data_with_strings = UsageDataWithStrings(
+                service_provider=model,
+                model=model,
+                prompt_string=prompt,
+                output_string=''  # Image generation doesn't have text output
+            )
+
+            self.paygent_client.send_usage_with_token_string(
+                external_agent_id,
+                external_customer_id,
+                indicator,
+                usage_data_with_strings
+            )
+
+        return response
+
+
+class ChatSessionWrapper:
+    """Wrapper for Gemini ChatSession."""
+
+    def __init__(
+        self,
+        chat_session: Any,
+        paygent_client: Client,
+        model_name: str,
+        indicator: str,
+        external_agent_id: str,
+        external_customer_id: str
+    ):
+        self.chat_session = chat_session
+        self.paygent_client = paygent_client
+        self.model_name = model_name
+        self.indicator = indicator
+        self.external_agent_id = external_agent_id
+        self.external_customer_id = external_customer_id
+
+    def send_message(self, message: str) -> Any:
+        """
+        Send a message in the chat with automatic usage tracking.
+        
+        Args:
+            message: The message to send
+            
+        Returns:
+            The chat response from Gemini
+        """
+        # Make the Gemini API call
+        response = self.chat_session.send_message(message)
+
+        # Extract usage data from response with robust fallback mechanism
+        has_valid_usage = (
+            hasattr(response, 'usage_metadata') and
+            response.usage_metadata and
+            (response.usage_metadata.prompt_token_count > 0 or
+             response.usage_metadata.candidates_token_count > 0)
+        )
+
+        if has_valid_usage:
+            # Primary path: Use usage metadata from API response
+            usage_data = UsageData(
+                service_provider=self.model_name,
+                model=self.model_name,
+                prompt_tokens=response.usage_metadata.prompt_token_count or 0,
+                completion_tokens=response.usage_metadata.candidates_token_count or 0,
+                total_tokens=response.usage_metadata.total_token_count or 0
+            )
+
+            self.paygent_client.send_usage(
+                self.external_agent_id,
+                self.external_customer_id,
+                self.indicator,
+                usage_data
+            )
+        else:
+            # Fallback path: Calculate tokens from message and response
+            output_string = ''
+            if hasattr(response, 'candidates') and response.candidates:
+                if len(response.candidates) > 0:
+                    candidate = response.candidates[0]
+                    if hasattr(candidate, 'content') and hasattr(candidate.content, 'parts'):
+                        if len(candidate.content.parts) > 0:
+                            part = candidate.content.parts[0]
+                            if hasattr(part, 'text'):
+                                output_string = part.text or ''
+
+            usage_data_with_strings = UsageDataWithStrings(
+                service_provider=self.model_name,
+                model=self.model_name,
+                prompt_string=message,
+                output_string=output_string
+            )
+
+            self.paygent_client.send_usage_with_token_string(
+                self.external_agent_id,
+                self.external_customer_id,
+                self.indicator,
+                usage_data_with_strings
+            )
+
+        return response
+
+    def get_history(self) -> Any:
+        """
+        Get the chat history.
+        
+        Returns:
+            The chat history
+        """
+        return self.chat_session.get_history()

paygent_sdk/wrappers/langchain_wrapper.py

@@ -0,0 +1,257 @@
+"""
+LangChain callback handler for automatic usage tracking with Paygent.
+This callback intercepts LangChain LLM calls and automatically sends usage data to Paygent.
+"""
+
+from typing import Any, Dict, List, Optional
+from uuid import UUID
+
+try:
+    from langchain_core.callbacks import BaseCallbackHandler
+    from langchain_core.outputs import LLMResult
+except ImportError:
+    raise ImportError(
+        "LangChain is required for this integration. "
+        "Install it with: pip install langchain-core"
+    )
+
+from ..client import Client
+from ..models import UsageData, UsageDataWithStrings
+
+
+class PaygentLangChainCallback(BaseCallbackHandler):
+    """
+    LangChain callback handler that automatically tracks usage with Paygent.
+    
+    Usage example:
+    ```python
+    from paygent_sdk import Client, PaygentLangChainCallback
+    from langchain_openai import ChatOpenAI
+    
+    paygent_client = Client("your-api-key")
+    callback = PaygentLangChainCallback(
+        paygent_client,
+        indicator="chat_completion",
+        external_agent_id="agent-123",
+        external_customer_id="customer-456"
+    )
+    
+    llm = ChatOpenAI(callbacks=[callback])
+    response = llm.invoke("Hello!")
+    # Usage automatically tracked!
+    ```
+    """
+    
+    def __init__(
+        self,
+        paygent_client: Client,
+        indicator: str,
+        external_agent_id: str,
+        external_customer_id: str
+    ):
+        """
+        Create a new PaygentLangChainCallback.
+        
+        Args:
+            paygent_client: The Paygent client instance for usage tracking
+            indicator: Indicator for the usage event (e.g., "chat_completion")
+            external_agent_id: External agent identifier
+            external_customer_id: External customer identifier
+        """
+        super().__init__()
+        self.paygent_client = paygent_client
+        self.indicator = indicator
+        self.external_agent_id = external_agent_id
+        self.external_customer_id = external_customer_id
+        self.run_info: Dict[UUID, Dict[str, str]] = {}
+    
+    def _extract_provider(self, serialized: Dict[str, Any]) -> str:
+        """Extract the service provider from the serialized LLM data."""
+        if not serialized:
+            return "unknown"
+        
+        # Check id field
+        if "id" in serialized:
+            id_list = serialized["id"]
+            if isinstance(id_list, list):
+                id_str = " ".join(id_list).lower()
+            else:
+                id_str = str(id_list).lower()
+            
+            if "openai" in id_str:
+                return "openai"
+            if "anthropic" in id_str:
+                return "anthropic"
+            if "mistral" in id_str:
+                return "mistral"
+            if "cohere" in id_str:
+                return "cohere"
+            if "google" in id_str or "gemini" in id_str:
+                return "gemini"
+            if "huggingface" in id_str:
+                return "huggingface"
+            if "azure" in id_str:
+                return "azure"
+        
+        # Check name field
+        if "name" in serialized:
+            name = str(serialized["name"]).lower()
+            if "openai" in name:
+                return "openai"
+            if "anthropic" in name:
+                return "anthropic"
+            if "mistral" in name:
+                return "mistral"
+            if "gemini" in name or "google" in name:
+                return "gemini"
+        
+        return "unknown"
+    
+    def on_llm_start(
+        self,
+        serialized: Dict[str, Any],
+        prompts: List[str],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when an LLM starts running."""
+        try:
+            provider = self._extract_provider(serialized)
+            model_name = metadata.get("ls_model_name", "unknown") if metadata else "unknown"
+            
+            # Store the run info for use in on_llm_end
+            self.run_info[run_id] = {
+                "provider": provider,
+                "model_name": model_name
+            }
+        except Exception as e:
+            print(f"Error in on_llm_start: {e}")
+    
+    def on_llm_end(
+        self,
+        response: LLMResult,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when an LLM ends running."""
+        try:
+            # Get the stored run info
+            info = self.run_info.get(run_id, {})
+            provider = info.get("provider", "unknown")
+            model_name = info.get("model_name", "unknown")
+            
+            # Extract usage information from LangChain's LLMResult
+            llm_output = response.llm_output or {}
+            
+            prompt_tokens = 0
+            completion_tokens = 0
+            total_tokens = 0
+            
+            # OpenAI-style usage (in llm_output["token_usage"])
+            if "token_usage" in llm_output:
+                token_usage = llm_output["token_usage"]
+                prompt_tokens = token_usage.get("prompt_tokens", 0)
+                completion_tokens = token_usage.get("completion_tokens", 0)
+                total_tokens = token_usage.get("total_tokens", 0)
+            
+            # Anthropic-style usage (in llm_output["usage"])
+            elif "usage" in llm_output:
+                usage = llm_output["usage"]
+                prompt_tokens = usage.get("input_tokens", usage.get("prompt_tokens", 0))
+                completion_tokens = usage.get("output_tokens", usage.get("completion_tokens", 0))
+                total_tokens = prompt_tokens + completion_tokens
+            
+            # Gemini-style usage (in generations[0].message.usage_metadata)
+            elif (response.generations and 
+                  len(response.generations) > 0 and 
+                  len(response.generations[0]) > 0):
+                gen = response.generations[0][0]
+                if hasattr(gen, "message") and hasattr(gen.message, "usage_metadata"):
+                    usage_metadata = gen.message.usage_metadata
+                    prompt_tokens = getattr(usage_metadata, "input_tokens", 0)
+                    completion_tokens = getattr(usage_metadata, "output_tokens", 0)
+                    total_tokens = getattr(usage_metadata, "total_tokens", 0)
+            
+            # Try to extract model name from response if not already set
+            if model_name == "unknown":
+                if "model_name" in llm_output:
+                    model_name = llm_output["model_name"]
+                elif (response.generations and 
+                      len(response.generations) > 0 and 
+                      len(response.generations[0]) > 0):
+                    gen = response.generations[0][0]
+                    if hasattr(gen, "message") and hasattr(gen.message, "response_metadata"):
+                        model_name = gen.message.response_metadata.get("model_name", "unknown")
+            
+            # Send usage data if we have token information
+            if total_tokens > 0:
+                usage_data = UsageData(
+                    service_provider=provider,
+                    model=model_name,
+                    prompt_tokens=prompt_tokens,
+                    completion_tokens=completion_tokens,
+                    total_tokens=total_tokens
+                )
+                
+                self.paygent_client.send_usage(
+                    self.external_agent_id,
+                    self.external_customer_id,
+                    self.indicator,
+                    usage_data
+                )
+            else:
+                # Fallback: use string-based tracking
+                # Extract prompt and output strings
+                prompt_string = ""
+                if response.generations and len(response.generations) > 0:
+                    # Get the first generation's text
+                    if len(response.generations[0]) > 0:
+                        prompt_string = str(response.generations[0][0].text)
+                
+                output_string = ""
+                if response.generations and len(response.generations) > 0:
+                    outputs = [gen.text for gen in response.generations[0]]
+                    output_string = " ".join(outputs)
+                
+                usage_data_with_strings = UsageDataWithStrings(
+                    service_provider=provider,
+                    model=model_name,
+                    prompt_string=prompt_string,
+                    output_string=output_string
+                )
+                
+                self.paygent_client.send_usage_with_token_string(
+                    self.external_agent_id,
+                    self.external_customer_id,
+                    self.indicator,
+                    usage_data_with_strings
+                )
+            
+            # Clean up the run info
+            if run_id in self.run_info:
+                del self.run_info[run_id]
+                
+        except Exception as e:
+            print(f"Error tracking LangChain usage with Paygent: {e}")
+            if run_id in self.run_info:
+                del self.run_info[run_id]
+    
+    def on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when an LLM encounters an error."""
+        print(f"LLM error: {error}")
+        # Clean up the run info
+        if run_id in self.run_info:
+            del self.run_info[run_id]