chuk-ai-session-manager 0.1.1__py3-none-any.whl

chuk_ai_session_manager/models/token_usage.py

@@ -0,0 +1,316 @@
+# chuk_ai_session_manager/models/token_usage.py
+"""
+Token usage tracking models for the chuk session manager.
+
+This module provides models for tracking token usage in LLM interactions
+with proper async support.
+"""
+from __future__ import annotations
+from datetime import datetime, timezone
+from typing import Dict, Optional, Union, List, Any
+from uuid import uuid4
+from pydantic import BaseModel, Field, ConfigDict
+import asyncio
+
+# Try to import tiktoken, but make it optional
+try:
+    import tiktoken
+    TIKTOKEN_AVAILABLE = True
+except ImportError:
+    TIKTOKEN_AVAILABLE = False
+
+
+class TokenUsage(BaseModel):
+    """
+    Tracks token usage for LLM interactions.
+    
+    Attributes:
+        prompt_tokens: Number of tokens in the prompt/input
+        completion_tokens: Number of tokens in the completion/output
+        total_tokens: Total tokens (prompt + completion)
+        model: The model used for the interaction (helps with pricing calculations)
+        estimated_cost_usd: Estimated cost in USD (if pricing info is available)
+    """
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = Field(default=0)
+    model: str = ""
+    estimated_cost_usd: Optional[float] = None
+    
+    def __init__(self, **data):
+        super().__init__(**data)
+        # Auto-calculate total tokens if not explicitly provided
+        if self.total_tokens == 0 and (self.prompt_tokens > 0 or self.completion_tokens > 0):
+            self.total_tokens = self.prompt_tokens + self.completion_tokens
+            
+        # Auto-calculate estimated cost if model is provided
+        if self.model and self.estimated_cost_usd is None:
+            self.estimated_cost_usd = self._calculate_cost_sync()
+    
+    def _calculate_cost_sync(self) -> float:
+        """
+        Synchronous implementation of calculate_cost.
+        
+        Returns:
+            Estimated cost in USD
+        """
+        # Model pricing per 1000 tokens (approximate as of May 2025)
+        pricing = {
+            # OpenAI models
+            "gpt-4": {"input": 0.03, "output": 0.06},
+            "gpt-4-turbo": {"input": 0.01, "output": 0.03},
+            "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015},
+            
+            # Claude models
+            "claude-3-opus": {"input": 0.015, "output": 0.075},
+            "claude-3-sonnet": {"input": 0.003, "output": 0.015},
+            "claude-3-haiku": {"input": 0.00025, "output": 0.00125},
+            
+            # Fallback for unknown models
+            "default": {"input": 0.001, "output": 0.002}
+        }
+        
+        # Get pricing for this model or use default
+        model_pricing = pricing.get(self.model.lower(), pricing["default"])
+        
+        # Calculate cost
+        input_cost = (self.prompt_tokens / 1000) * model_pricing["input"]
+        output_cost = (self.completion_tokens / 1000) * model_pricing["output"]
+        
+        return round(input_cost + output_cost, 6)
+    
+    async def calculate_cost(self) -> float:
+        """
+        Async version of calculate_cost.
+        
+        Returns:
+            Estimated cost in USD
+        """
+        # Token calculation is CPU-bound, so run in executor
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(None, self._calculate_cost_sync)
+    
+    def _update_sync(self, prompt_tokens: int = 0, completion_tokens: int = 0) -> None:
+        """
+        Synchronous implementation of update.
+        
+        Args:
+            prompt_tokens: Additional prompt tokens to add
+            completion_tokens: Additional completion tokens to add
+        """
+        self.prompt_tokens += prompt_tokens
+        self.completion_tokens += completion_tokens
+        self.total_tokens = self.prompt_tokens + self.completion_tokens
+        
+        if self.model:
+            self.estimated_cost_usd = self._calculate_cost_sync()
+    
+    async def update(self, prompt_tokens: int = 0, completion_tokens: int = 0) -> None:
+        """
+        Async version of update.
+        
+        Args:
+            prompt_tokens: Additional prompt tokens to add
+            completion_tokens: Additional completion tokens to add
+        """
+        self.prompt_tokens += prompt_tokens
+        self.completion_tokens += completion_tokens
+        self.total_tokens = self.prompt_tokens + self.completion_tokens
+        
+        if self.model:
+            self.estimated_cost_usd = await self.calculate_cost()
+    
+    @classmethod
+    def _from_text_sync(
+        cls, 
+        prompt: str, 
+        completion: Optional[str] = None, 
+        model: str = "gpt-3.5-turbo"
+    ) -> TokenUsage:
+        """
+        Synchronous implementation of from_text.
+        
+        Args:
+            prompt: The prompt/input text
+            completion: The completion/output text (optional)
+            model: The model name to use for counting and pricing
+            
+        Returns:
+            A TokenUsage instance with token counts
+        """
+        prompt_tokens = cls._count_tokens_sync(prompt, model)
+        completion_tokens = cls._count_tokens_sync(completion, model) if completion else 0
+        
+        return cls(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            model=model
+        )
+    
+    @classmethod
+    async def from_text(
+        cls, 
+        prompt: str, 
+        completion: Optional[str] = None, 
+        model: str = "gpt-3.5-turbo"
+    ) -> TokenUsage:
+        """
+        Async version of from_text.
+        
+        Args:
+            prompt: The prompt/input text
+            completion: The completion/output text (optional)
+            model: The model name to use for counting and pricing
+            
+        Returns:
+            A TokenUsage instance with token counts
+        """
+        # Run token counting in executor since it's CPU-bound
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(
+            None, 
+            lambda: cls._from_text_sync(prompt, completion, model)
+        )
+    
+    @staticmethod
+    def _count_tokens_sync(text: Optional[str], model: str = "gpt-3.5-turbo") -> int:
+        """
+        Synchronous implementation of count_tokens.
+        
+        Args:
+            text: The text to count tokens for
+            model: The model name to use for counting
+            
+        Returns:
+            The number of tokens
+        """
+        if text is None:
+            return 0
+            
+        if TIKTOKEN_AVAILABLE:
+            try:
+                encoding = tiktoken.encoding_for_model(model)
+                return len(encoding.encode(text))
+            except (KeyError, ValueError):
+                # Fall back to cl100k_base encoding if the model is not found
+                try:
+                    encoding = tiktoken.get_encoding("cl100k_base")
+                    return len(encoding.encode(text))
+                except Exception:
+                    # If all else fails, use the approximation
+                    pass
+        
+        # Simple approximation: ~4 chars per token for English text
+        return int(len(text) / 4)
+    
+    @staticmethod
+    async def count_tokens(text: Optional[str], model: str = "gpt-3.5-turbo") -> int:
+        """
+        Async version of count_tokens.
+        
+        Args:
+            text: The text to count tokens for
+            model: The model name to use for counting
+            
+        Returns:
+            The number of tokens
+        """
+        # Run in executor since token counting is CPU-bound
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(
+            None,
+            lambda: TokenUsage._count_tokens_sync(text, model)
+        )
+    
+    def __add__(self, other: TokenUsage) -> TokenUsage:
+        """
+        Add two TokenUsage instances together.
+        
+        Args:
+            other: Another TokenUsage instance
+            
+        Returns:
+            A new TokenUsage instance with combined counts
+        """
+        # Use the model from self if it exists, otherwise use the other's model
+        model = self.model if self.model else other.model
+        
+        return TokenUsage(
+            prompt_tokens=self.prompt_tokens + other.prompt_tokens,
+            completion_tokens=self.completion_tokens + other.completion_tokens,
+            model=model
+        )
+
+
+class TokenSummary(BaseModel):
+    """
+    Summarizes token usage across multiple interactions.
+    
+    Attributes:
+        total_prompt_tokens: Total tokens used in prompts
+        total_completion_tokens: Total tokens used in completions
+        total_tokens: Total tokens overall
+        usage_by_model: Breakdown of usage by model
+        total_estimated_cost_usd: Total estimated cost across all models
+    """
+    total_prompt_tokens: int = 0
+    total_completion_tokens: int = 0
+    total_tokens: int = 0
+    usage_by_model: Dict[str, TokenUsage] = Field(default_factory=dict)
+    total_estimated_cost_usd: float = 0.0
+    
+    def _add_usage_sync(self, usage: TokenUsage) -> None:
+        """
+        Synchronous implementation of add_usage.
+        
+        Args:
+            usage: The TokenUsage to add
+        """
+        self.total_prompt_tokens += usage.prompt_tokens
+        self.total_completion_tokens += usage.completion_tokens
+        self.total_tokens += usage.total_tokens
+        
+        if usage.estimated_cost_usd is not None:
+            self.total_estimated_cost_usd += usage.estimated_cost_usd
+        
+        if usage.model:
+            if usage.model in self.usage_by_model:
+                self.usage_by_model[usage.model]._update_sync(
+                    prompt_tokens=usage.prompt_tokens,
+                    completion_tokens=usage.completion_tokens
+                )
+            else:
+                self.usage_by_model[usage.model] = TokenUsage(
+                    prompt_tokens=usage.prompt_tokens,
+                    completion_tokens=usage.completion_tokens,
+                    model=usage.model
+                )
+    
+    async def add_usage(self, usage: TokenUsage) -> None:
+        """
+        Async version of add_usage.
+        
+        Args:
+            usage: The TokenUsage to add
+        """
+        self.total_prompt_tokens += usage.prompt_tokens
+        self.total_completion_tokens += usage.completion_tokens
+        self.total_tokens += usage.total_tokens
+        
+        if usage.estimated_cost_usd is not None:
+            self.total_estimated_cost_usd += usage.estimated_cost_usd
+        
+        if usage.model:
+            if usage.model in self.usage_by_model:
+                await self.usage_by_model[usage.model].update(
+                    prompt_tokens=usage.prompt_tokens,
+                    completion_tokens=usage.completion_tokens
+                )
+            else:
+                self.usage_by_model[usage.model] = TokenUsage(
+                    prompt_tokens=usage.prompt_tokens,
+                    completion_tokens=usage.completion_tokens,
+                    model=usage.model
+                )

chuk_ai_session_manager/sample_tools.py

@@ -0,0 +1,194 @@
+# sample_tools.py
+"""
+Sample tools for chuk session manager demos - corrected version following registry example
+"""
+
+import asyncio
+import random
+from datetime import datetime
+from typing import Dict, Any
+
+from chuk_tool_processor.registry import register_tool
+
+
+@register_tool(name="calculator", namespace="default", description="Perform basic arithmetic operations")
+class CalculatorTool:
+    """Calculator tool for basic arithmetic."""
+    
+    async def execute(self, operation: str, a: float, b: float) -> Dict[str, Any]:
+        """
+        Perform a basic arithmetic operation.
+        
+        Args:
+            operation: One of "add", "subtract", "multiply", "divide"
+            a: First operand
+            b: Second operand
+            
+        Returns:
+            Dictionary with the result
+        """
+        print(f"🧮 Calculator executing: {a} {operation} {b}")
+        
+        if operation == "add":
+            result = a + b
+        elif operation == "subtract":
+            result = a - b
+        elif operation == "multiply":
+            result = a * b
+        elif operation == "divide":
+            if b == 0:
+                raise ValueError("Cannot divide by zero")
+            result = a / b
+        else:
+            raise ValueError(f"Unknown operation: {operation}")
+            
+        return {
+            "operation": operation,
+            "a": a,
+            "b": b,
+            "result": result,
+            "timestamp": datetime.now().isoformat()
+        }
+
+
+@register_tool(name="weather", namespace="default", description="Get current weather information for a location")
+class WeatherTool:
+    """Weather tool that returns mock weather data."""
+    
+    async def execute(self, location: str) -> Dict[str, Any]:
+        """
+        Get weather information for a specific location.
+        
+        Args:
+            location: The city or location to get weather for
+            
+        Returns:
+            Dictionary with weather information
+        """
+        print(f"🌤️ Weather tool executing for: {location}")
+        await asyncio.sleep(0.1)  # Simulate API delay
+        
+        # Mock realistic weather based on location
+        base_temp = 15  # Default moderate temperature
+        if any(city in location.lower() for city in ["miami", "phoenix", "dubai", "singapore"]):
+            base_temp = 28
+        elif any(city in location.lower() for city in ["moscow", "montreal", "oslo", "anchorage"]):
+            base_temp = -5
+        elif any(city in location.lower() for city in ["london", "seattle", "vancouver"]):
+            base_temp = 12
+        elif any(city in location.lower() for city in ["tokyo", "new york", "paris", "berlin"]):
+            base_temp = 18
+        
+        # Add some randomness
+        temperature = base_temp + random.randint(-8, 12)
+        conditions = ["Sunny", "Partly Cloudy", "Cloudy", "Light Rain", "Heavy Rain", "Snow", "Thunderstorm", "Foggy"]
+        condition = random.choice(conditions)
+        humidity = random.randint(35, 85)
+        wind_speed = random.uniform(2.0, 25.0)
+        feels_like = temperature + random.randint(-3, 3)
+        
+        # Adjust conditions based on temperature
+        if temperature < 0:
+            condition = random.choice(["Snow", "Cloudy", "Partly Cloudy"])
+        elif temperature > 30:
+            condition = random.choice(["Sunny", "Partly Cloudy", "Hot"])
+        
+        description = f"Current weather in {location} is {condition.lower()} with temperature {temperature}°C"
+        
+        return {
+            "location": location,
+            "temperature": float(temperature),
+            "condition": condition,
+            "humidity": humidity,
+            "wind_speed": round(wind_speed, 1),
+            "description": description,
+            "feels_like": float(feels_like),
+            "timestamp": datetime.now().isoformat()
+        }
+
+
+@register_tool(name="search", namespace="default", description="Search for information on the internet")  
+class SearchTool:
+    """Search tool that returns mock search results."""
+    
+    async def execute(self, query: str, max_results: int = 3) -> Dict[str, Any]:
+        """
+        Search for information on the internet.
+        
+        Args:
+            query: Search query
+            max_results: Maximum number of results to return
+            
+        Returns:
+            Dictionary with search results
+        """
+        print(f"🔍 Search tool executing for: {query}")
+        await asyncio.sleep(0.2)  # Simulate API delay
+        
+        results = []
+        query_lower = query.lower()
+        
+        # Generate contextually relevant mock results based on query
+        if "climate" in query_lower or "environment" in query_lower:
+            result_templates = [
+                {
+                    "title": "Climate Change Adaptation Strategies - IPCC Report",
+                    "url": "https://www.ipcc.ch/adaptation-strategies",
+                    "snippet": "Comprehensive guide to climate change adaptation strategies for communities, businesses, and governments. Includes resilience planning and risk assessment."
+                },
+                {
+                    "title": "Environmental Adaptation Solutions | Climate.gov", 
+                    "url": "https://www.climate.gov/adaptation-solutions",
+                    "snippet": "Evidence-based climate adaptation solutions including infrastructure improvements, ecosystem restoration, and community planning approaches."
+                },
+                {
+                    "title": "Building Climate Resilience: A Practical Guide",
+                    "url": "https://www.resilience.org/climate-guide", 
+                    "snippet": "Practical steps for building climate resilience in your community. Covers early warning systems, green infrastructure, and adaptation planning."
+                }
+            ]
+        elif "weather" in query_lower:
+            result_templates = [
+                {
+                    "title": "Weather Forecast and Current Conditions",
+                    "url": "https://weather.com/forecast",
+                    "snippet": "Get accurate weather forecasts, current conditions, and severe weather alerts for your location."
+                },
+                {
+                    "title": "Climate and Weather Patterns Explained",
+                    "url": "https://www.weatherpatterns.org",
+                    "snippet": "Understanding weather patterns, climate systems, and meteorological phenomena that affect daily weather."
+                }
+            ]
+        else:
+            # Generic results for other queries
+            result_templates = [
+                {
+                    "title": f"Everything You Need to Know About {query.title()}",
+                    "url": f"https://encyclopedia.com/{query.lower().replace(' ', '-')}",
+                    "snippet": f"Comprehensive information and resources about {query}. Expert insights, latest research, and practical applications."
+                },
+                {
+                    "title": f"{query.title()} - Latest News and Updates",
+                    "url": f"https://news.example.com/{query.lower().replace(' ', '-')}",
+                    "snippet": f"Stay up to date with the latest news, trends, and developments related to {query}."
+                },
+                {
+                    "title": f"Guide to {query.title()} - Best Practices",
+                    "url": f"https://guides.com/{query.lower().replace(' ', '-')}",
+                    "snippet": f"Expert guide covering best practices, tips, and strategies for {query}. Includes real-world examples and case studies."
+                }
+            ]
+        
+        # Select results up to max_results
+        selected_results = result_templates[:max_results]
+        
+        return {
+            "query": query,
+            "results_count": len(selected_results),
+            "results": selected_results,
+            "timestamp": datetime.now().isoformat()
+        }
+
+
+print("✅ sample_tools.py: 3 tools defined with @register_tool decorator (corrected version)")

chuk_ai_session_manager/session_aware_tool_processor.py

@@ -0,0 +1,178 @@
+# chuk_ai_session_manager/session_aware_tool_processor.py
+#!/usr/bin/env python3
+"""Session-aware Tool-processor for chuk_tool_processor 0.1.x.
+
+* Converts OpenAI `tool_calls` → `ToolCall` objects.
+* Executes them with **ToolProcessor().executor.execute**.
+* Adds caching / retry.
+* Logs every call into the session tree, storing the **string-repr**
+  of the result (this is what the prompt-builder currently expects)."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+import logging
+from typing import Any, Dict, List
+
+from chuk_tool_processor.core.processor import ToolProcessor
+from chuk_tool_processor.models.tool_call import ToolCall
+from chuk_tool_processor.models.tool_result import ToolResult
+
+from chuk_ai_session_manager.models.event_source import EventSource
+from chuk_ai_session_manager.models.event_type import EventType
+from chuk_ai_session_manager.models.session_event import SessionEvent
+from chuk_ai_session_manager.storage import SessionStoreProvider
+
+logger = logging.getLogger(__name__)
+
+
+class SessionAwareToolProcessor:
+    """Run tool-calls, add retry/caching, and log into a session."""
+
+    # ─────────────────────────── construction ──────────────────────────
+    def __init__(
+        self,
+        session_id: str,
+        *,
+        enable_caching: bool = True,
+        max_retries: int = 2,
+        retry_delay: float = 1.0,
+    ) -> None:
+        self.session_id     = session_id
+        self.enable_caching = enable_caching
+        self.max_retries    = max_retries
+        self.retry_delay    = retry_delay
+        self.cache: Dict[str, ToolResult] = {}
+
+        self._tp = ToolProcessor()
+        if not hasattr(self._tp, "executor"):
+            raise AttributeError("Installed chuk_tool_processor is too old - missing `.executor`")
+
+    @classmethod
+    async def create(cls, session_id: str, **kwargs):
+        store = SessionStoreProvider.get_store()
+        if not await store.get(session_id):
+            raise ValueError(f"Session {session_id} not found")
+        return cls(session_id=session_id, **kwargs)
+
+    # ─────────────────────────── internals ─────────────────────────────
+    async def _maybe_await(self, val: Any) -> Any:
+        return await val if asyncio.iscoroutine(val) else val
+
+    async def _exec_calls(self, calls: List[Dict[str, Any]]) -> List[ToolResult]:
+        """Convert dicts → ToolCall and drive the executor."""
+        tool_calls: list[ToolCall] = []
+        for c in calls:
+            fn   = c.get("function", {})
+            name = fn.get("name", "tool")
+            try:
+                args = json.loads(fn.get("arguments", "{}"))
+            except json.JSONDecodeError:
+                args = {"raw": fn.get("arguments")}
+            tool_calls.append(ToolCall(tool=name, arguments=args))
+
+        results = await self._tp.executor.execute(tool_calls)
+        for r in results:
+            r.result = await self._maybe_await(r.result)
+        return results
+
+    async def _log_event(
+        self,
+        session,
+        parent_id: str,
+        res: ToolResult,
+        attempt: int,
+        *,
+        cached: bool,
+        failed: bool = False,
+    ) -> None:
+        """Persist TOOL_CALL with *string* result (prompt-friendly)."""
+        result_str = str(res.result) if res.result is not None else "null"
+
+        ev = await SessionEvent.create_with_tokens(
+            message={
+                "tool":      res.tool,
+                "arguments": getattr(res, "arguments", None),
+                "result":    result_str,
+                "error":     res.error,
+                "cached":    cached,
+            },
+            prompt=f"{res.tool}({json.dumps(getattr(res, 'arguments', None), default=str)})",
+            completion=result_str,
+            model="tool-execution",
+            source=EventSource.SYSTEM,
+            type=EventType.TOOL_CALL,
+        )
+        await ev.update_metadata("parent_event_id", parent_id)
+        await ev.update_metadata("call_id", getattr(res, "id", "cid"))
+        await ev.update_metadata("attempt", attempt)
+        if failed:
+            await ev.update_metadata("failed", True)
+        await session.add_event_and_save(ev)
+
+    # ─────────────────────────── public API ────────────────────────────
+    async def process_llm_message(self, llm_msg: Dict[str, Any], _) -> List[ToolResult]:
+        store   = SessionStoreProvider.get_store()
+        session = await store.get(self.session_id)
+        if not session:
+            raise ValueError(f"Session {self.session_id} not found")
+
+        parent_evt = await SessionEvent.create_with_tokens(
+            message=llm_msg,
+            prompt="",
+            completion=json.dumps(llm_msg, ensure_ascii=False),
+            model="gpt-4o-mini",
+            source=EventSource.LLM,
+            type=EventType.MESSAGE,
+        )
+        await session.add_event_and_save(parent_evt)
+
+        calls = llm_msg.get("tool_calls", [])
+        if not calls:
+            return []
+
+        out: list[ToolResult] = []
+        for call in calls:
+            fn   = call.get("function", {})
+            name = fn.get("name", "tool")
+            try:
+                args = json.loads(fn.get("arguments", "{}"))
+            except json.JSONDecodeError:
+                args = {"raw": fn.get("arguments")}
+
+            cache_key = (
+                hashlib.md5(f"{name}:{json.dumps(args, sort_keys=True)}".encode()).hexdigest()
+                if self.enable_caching else None
+            )
+
+            # 1) cache hit --------------------------------------------------
+            if cache_key and (cached := self.cache.get(cache_key)):
+                await self._log_event(session, parent_evt.id, cached, 1, cached=True)
+                out.append(cached)
+                continue
+
+            # 2) execute with retry ----------------------------------------
+            last_err: str | None = None
+            for attempt in range(1, self.max_retries + 2):
+                try:
+                    res = (await self._exec_calls([call]))[0]
+                    if cache_key:
+                        self.cache[cache_key] = res
+                    await self._log_event(session, parent_evt.id, res, attempt, cached=False)
+                    out.append(res)
+                    break
+                except Exception as exc:  # noqa: BLE001
+                    last_err = str(exc)
+                    if attempt <= self.max_retries:
+                        await asyncio.sleep(self.retry_delay)
+                        continue
+                    err_res = ToolResult(tool=name, result=None, error=last_err)  # type: ignore[arg-type]
+                    await self._log_event(
+                        session, parent_evt.id, err_res, attempt,
+                        cached=False, failed=True
+                    )
+                    out.append(err_res)
+
+        return out