droidrun 0.3.5__py3-none-any.whl → 0.3.7__py3-none-any.whl

droidrun/agent/usage.py

@@ -0,0 +1,213 @@
+import contextlib
+from llama_index.core.callbacks import CallbackManager
+from llama_index.core.callbacks.base_handler import BaseCallbackHandler
+from llama_index.core.callbacks.schema import CBEventType, EventPayload
+from llama_index.core.llms import LLM, ChatResponse
+from pydantic import BaseModel
+from typing import Any, Dict, List, Optional
+from uuid import uuid4
+import logging
+
+logger = logging.getLogger("droidrun")
+SUPPORTED_PROVIDERS = [
+    "Gemini",
+    "GoogleGenAI",
+    "OpenAI",
+    "Anthropic",
+    "Ollama",
+    "DeepSeek",
+]
+
+
+class UsageResult(BaseModel):
+    request_tokens: int
+    response_tokens: int
+    total_tokens: int
+    requests: int
+
+def get_usage_from_response(provider: str, chat_rsp: ChatResponse) -> UsageResult:
+    rsp = chat_rsp.raw
+    if not rsp:
+        raise ValueError("No raw response in chat response")
+
+    print(f"rsp: {rsp.__class__.__name__}")
+
+    if provider == "Gemini" or provider == "GoogleGenAI":
+        return UsageResult(
+            request_tokens=rsp["usage_metadata"]["prompt_token_count"],
+            response_tokens=rsp["usage_metadata"]["candidates_token_count"],
+            total_tokens=rsp["usage_metadata"]["total_token_count"],
+            requests=1,
+        )
+    elif provider == "OpenAI":
+        from openai.types import CompletionUsage as OpenAIUsage
+
+        usage: OpenAIUsage = rsp.usage
+        return UsageResult(
+            request_tokens=usage.prompt_tokens,
+            response_tokens=usage.completion_tokens,
+            total_tokens=usage.total_tokens,
+            requests=1,
+        )
+    elif provider == "Anthropic":
+        from anthropic.types import Usage as AnthropicUsage
+
+        usage: AnthropicUsage = rsp["usage"]
+        return UsageResult(
+            request_tokens=usage.input_tokens,
+            response_tokens=usage.output_tokens,
+            total_tokens=usage.input_tokens + usage.output_tokens,
+            requests=1,
+        )
+    elif provider == "Ollama":
+        # Ollama response format uses different field names
+        prompt_eval_count = rsp.get("prompt_eval_count", 0)
+        eval_count = rsp.get("eval_count", 0)
+        return UsageResult(
+            request_tokens=prompt_eval_count,
+            response_tokens=eval_count,
+            total_tokens=prompt_eval_count + eval_count,
+            requests=1,
+        )
+    elif provider == "DeepSeek":
+        # DeepSeek follows OpenAI-compatible format
+        usage = rsp.usage
+        if not usage:
+            usage = {}
+        return UsageResult(
+            request_tokens=usage.prompt_tokens or 0,
+            response_tokens=usage.completion_tokens or 0,
+            total_tokens=usage.total_tokens or 0,
+            requests=1,
+        )
+
+    raise ValueError(f"Unsupported provider: {provider}")
+
+class TokenCountingHandler(BaseCallbackHandler):
+    """Token counting handler for LLamaIndex LLM calls."""
+
+    def __init__(self, provider: str):
+        super().__init__(event_starts_to_ignore=[], event_ends_to_ignore=[])
+        self.provider = provider
+        self.request_tokens: int = 0
+        self.response_tokens: int = 0
+        self.total_tokens: int = 0
+        self.requests: int = 0
+
+    @classmethod
+    def class_name(cls) -> str:
+        """Class name."""
+        return "TokenCountingHandler"
+
+    @property
+    def usage(self) -> UsageResult:
+        return UsageResult(
+            request_tokens=self.request_tokens,
+            response_tokens=self.response_tokens,
+            total_tokens=self.total_tokens,
+            requests=self.requests,
+        )
+
+    def _get_event_usage(self, payload: Dict[str, Any]) -> UsageResult:
+        if not EventPayload.RESPONSE in payload:
+            raise ValueError("No response in payload")
+
+        chat_rsp: ChatResponse = payload.get(EventPayload.RESPONSE)
+        return get_usage_from_response(self.provider, chat_rsp)
+
+    def on_event_start(
+        self,
+        event_type: CBEventType,
+        payload: Optional[Dict[str, Any]] = None,
+        event_id: str = "",
+        parent_id: str = "",
+        **kwargs: Any,
+    ) -> str:
+        """Run when an event starts and return id of event."""
+        return event_id or str(uuid4())
+
+    def on_event_end(
+        self,
+        event_type: CBEventType,
+        payload: Optional[Dict[str, Any]] = None,
+        event_id: str = "",
+        **kwargs: Any,
+    ) -> None:
+        """Run when an event ends."""
+        try:
+            usage = self._get_event_usage(payload)
+
+            self.request_tokens += usage.request_tokens
+            self.response_tokens += usage.response_tokens
+            self.total_tokens += usage.total_tokens
+            self.requests += usage.requests
+        except Exception as e:
+            self.requests += 1
+            logger.warning(
+                f"Error tracking usage for provider {self.provider}: {e}",
+                extra={"provider": self.provider},
+            )
+
+    def start_trace(self, trace_id: Optional[str] = None) -> None:
+        """Run when an overall trace is launched."""
+        pass
+
+    def end_trace(
+        self,
+        trace_id: Optional[str] = None,
+        trace_map: Optional[Dict[str, List[str]]] = None,
+    ) -> None:
+        """Run when an overall trace is exited."""
+        pass
+
+@contextlib.contextmanager
+def llm_callback(llm: LLM, *args: List[BaseCallbackHandler]):
+    for arg in args:
+        llm.callback_manager.add_handler(arg)
+    yield
+    for arg in args:
+        llm.callback_manager.remove_handler(arg)
+
+def create_tracker(llm: LLM) -> TokenCountingHandler:
+    provider = llm.__class__.__name__
+    if provider not in SUPPORTED_PROVIDERS:
+        raise ValueError(f"Tracking not yet supported for provider: {provider}")
+
+    return TokenCountingHandler(provider)
+
+
+def track_usage(llm: LLM) -> TokenCountingHandler:
+    """Track token usage for an LLM instance across all requests.
+    
+    This function:
+    - Creates a new TokenCountingHandler for the LLM provider
+    - Registers that handler as an LLM callback to monitor all requests
+    - Returns the handler for accessing cumulative usage statistics
+    
+    The handler counts tokens for total LLM usage across all requests. For fine-grained
+    per-request counting, use either:
+    - `create_tracker()` with `llm_callback()` context manager for temporary tracking
+    - `get_usage_from_response()` to extract usage from individual responses
+    
+    Args:
+        llm: The LLamaIndex LLM instance to track usage for
+        
+    Returns:
+        TokenCountingHandler: The registered handler that accumulates usage statistics
+        
+    Raises:
+        ValueError: If the LLM provider is not supported for tracking
+        
+    Example:
+        >>> llm = OpenAI()
+        >>> tracker = track_usage(llm)
+        >>> # ... make LLM calls ...
+        >>> print(f"Total tokens used: {tracker.usage.total_tokens}")
+    """
+    provider = llm.__class__.__name__
+    if provider not in SUPPORTED_PROVIDERS:
+        raise ValueError(f"Tracking not yet supported for provider: {provider}")
+
+    tracker = TokenCountingHandler(provider)
+    llm.callback_manager.add_handler(tracker)
+    return tracker

droidrun/agent/utils/executer.py

@@ -143,6 +143,6 @@ class SimpleCodeExecutor:
         result = {
             'output': output,
             'screenshots': self.globals['step_screenshots'],
-            'ui_states': self.globals['step_ui_states']
+            'ui_states': self.globals['step_ui_states'],
         }
         return result

droidrun/agent/utils/llm_picker.py

@@ -2,9 +2,12 @@ import importlib
 import logging
 from typing import Any
 from llama_index.core.llms.llm import LLM
+from droidrun.agent.usage import track_usage
+
 # Configure logging
 logger = logging.getLogger("droidrun")
 
+
 def load_llm(provider_name: str, **kwargs: Any) -> LLM:
     """
     Dynamically loads and initializes a LlamaIndex LLM.
@@ -51,29 +54,39 @@ def load_llm(provider_name: str, **kwargs: Any) -> LLM:
         logger.debug(f"Successfully imported module: {module_path}")
 
     except ModuleNotFoundError:
-        logger.error(f"Module '{module_path}' not found. Try: pip install {install_package_name}")
+        logger.error(
+            f"Module '{module_path}' not found. Try: pip install {install_package_name}"
+        )
         raise ModuleNotFoundError(
             f"Could not import '{module_path}'. Is '{install_package_name}' installed?"
         ) from None
 
     try:
-        logger.debug(f"Attempting to get class '{provider_name}' from module {module_path}")
+        logger.debug(
+            f"Attempting to get class '{provider_name}' from module {module_path}"
+        )
         llm_class = getattr(llm_module, provider_name)
         logger.debug(f"Found class: {llm_class.__name__}")
 
         # Verify the class is a subclass of LLM
         if not isinstance(llm_class, type) or not issubclass(llm_class, LLM):
-            raise TypeError(f"Class '{provider_name}' found in '{module_path}' is not a valid LLM subclass.")
+            raise TypeError(
+                f"Class '{provider_name}' found in '{module_path}' is not a valid LLM subclass."
+            )
 
         # Filter out None values from kwargs
         filtered_kwargs = {k: v for k, v in kwargs.items() if v is not None}
-        
+
         # Initialize
-        logger.debug(f"Initializing {llm_class.__name__} with kwargs: {list(filtered_kwargs.keys())}")
+        logger.debug(
+            f"Initializing {llm_class.__name__} with kwargs: {list(filtered_kwargs.keys())}"
+        )
         llm_instance = llm_class(**filtered_kwargs)
         logger.debug(f"Successfully loaded and initialized LLM: {provider_name}")
         if not llm_instance:
-            raise RuntimeError(f"Failed to initialize LLM instance for {provider_name}.")
+            raise RuntimeError(
+                f"Failed to initialize LLM instance for {provider_name}."
+            )
         return llm_instance
 
     except AttributeError:
@@ -83,11 +96,12 @@ def load_llm(provider_name: str, **kwargs: Any) -> LLM:
         ) from None
     except TypeError as e:
         logger.error(f"Error initializing {provider_name}: {e}")
-        raise # Re-raise TypeError (could be from issubclass check or __init__)
+        raise  # Re-raise TypeError (could be from issubclass check or __init__)
     except Exception as e:
         logger.error(f"An unexpected error occurred initializing {provider_name}: {e}")
         raise e
-    
+
+
 # --- Example Usage ---
 if __name__ == "__main__":
     # Install the specific LLM integrations you want to test:
@@ -97,52 +111,75 @@ if __name__ == "__main__":
     #   llama-index-llms-gemini \
     #   llama-index-llms-openai
 
-    # Example 1: Load Anthropic (requires ANTHROPIC_API_KEY env var or kwarg)
-    print("\n--- Loading Anthropic ---")
-    try:
-        anthropic_llm = load_llm(
-            "Anthropic",
-            model="claude-3-7-sonnet-latest",
-        )
-        print(f"Loaded LLM: {type(anthropic_llm)}")
-        print(f"Model: {anthropic_llm.metadata}")
-    except Exception as e:
-        print(f"Failed to load Anthropic: {e}")
+    from llama_index.core.base.llms.types import ChatMessage
 
-    # Example 2: Load DeepSeek (requires DEEPSEEK_API_KEY env var or kwarg)
-    print("\n--- Loading DeepSeek ---")
-    try:
-        deepseek_llm = load_llm(
-            "DeepSeek",
-            model="deepseek-reasoner",
-            api_key="your api",  # or set DEEPSEEK_API_KEY
-        )
-        print(f"Loaded LLM: {type(deepseek_llm)}")
-        print(f"Model: {deepseek_llm.metadata}")
-    except Exception as e:
-        print(f"Failed to load DeepSeek: {e}")
+    providers = [
+        {
+            "name": "Anthropic",
+            "model": "claude-3-7-sonnet-latest",
+        },
+        {
+            "name": "DeepSeek",
+            "model": "deepseek-reasoner",
+        },
+        {
+            "name": "GoogleGenAI",
+            "model": "gemini-2.5-flash",
+        },
+        {
+            "name": "OpenAI",
+            "model": "gpt-4",
+        },
+        {
+            "name": "Ollama",
+            "model": "llama3.2:1b",
+            "base_url": "http://localhost:11434",
+        },
+    ]
 
-    # Example 3: Load Gemini (requires GOOGLE_APPLICATION_CREDENTIALS or kwarg)
-    print("\n--- Loading Gemini ---")
-    try:
-        gemini_llm = load_llm(
-            "Gemini",
-            model="gemini-2.0-fash",
-        )
-        print(f"Loaded LLM: {type(gemini_llm)}")
-        print(f"Model: {gemini_llm.metadata}")
-    except Exception as e:
-        print(f"Failed to load Gemini: {e}")
+    system_prompt = ChatMessage(
+        role="system",
+        content="You are a personal health and food coach. You are given a user's health and food preferences and you need to recommend a meal plan for them. only output the meal plan, no other text.",
+    )
 
-    # Example 4: Load OpenAI (requires OPENAI_API_KEY env var or kwarg)
-    print("\n--- Loading OpenAI ---")
-    try:
-        openai_llm = load_llm(
-            "OpenAI",
-            model="gp-4o",
-            temperature=0.5,
-        )
-        print(f"Loaded LLM: {type(openai_llm)}")
-        print(f"Model: {openai_llm.metadata}")
-    except Exception as e:
-        print(f"Failed to load OpenAI: {e}")
+    user_prompt = ChatMessage(
+        role="user",
+        content="I am a 25 year old male. I am 5'10 and 180 pounds. I am a vegetarian. I am allergic to peanuts and tree nuts. I am allergic to shellfish. I am allergic to eggs. I am allergic to dairy. I am allergic to soy. I am allergic to wheat. I am allergic to corn. I am allergic to oats. I am allergic to rice. I am allergic to barley. I am allergic to rye. I am allergic to oats. I am allergic to rice. I am allergic to barley. I am allergic to rye.",
+    )
+
+    messages = [system_prompt, user_prompt]
+
+    for provider in providers:
+        print(f"\n{'#' * 35} Loading {provider['name']} {'#' * 35}")
+        print("-" * 100)
+
+        try:
+            provider_name = provider.pop("name")
+            llm = load_llm(provider_name, **provider)
+            provider["name"] = provider_name
+            print(f"Loaded LLM: {type(llm)}")
+            print(f"Model: {llm.metadata}")
+            print("-" * 100)
+
+            tracker = track_usage(llm)
+            print(f"Tracker: {type(tracker)}")
+            print(f"Usage: {tracker.usage}")
+            print("-" * 100)
+
+            assert tracker.usage.requests == 0
+            assert tracker.usage.request_tokens == 0
+            assert tracker.usage.response_tokens == 0
+            assert tracker.usage.total_tokens == 0
+
+            res = llm.chat(messages)
+            print(f"Response: {res.message.content}")
+            print("-" * 100)
+            print(f"Usage: {tracker.usage}")
+
+            assert tracker.usage.requests == 1
+            assert tracker.usage.request_tokens > 0
+            assert tracker.usage.response_tokens > 0
+            assert tracker.usage.total_tokens > tracker.usage.request_tokens
+            assert tracker.usage.total_tokens > tracker.usage.response_tokens
+        except Exception as e:
+            print(f"Failed to load and track usage for {provider['name']}: {e}")