nora-lib-impl 1.0.2__tar.gz → 1.2.0__tar.gz

{nora_lib_impl-1.0.2 → nora_lib_impl-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nora_lib-impl
-Version: 1.0.2
+Version: 1.2.0
 Summary: For making and coordinating agents and tools
 Home-page: https://github.com/allenai/nora_lib/impl
 Requires-Python: >=3.9

{nora_lib_impl-1.0.2 → nora_lib_impl-1.2.0}/src/nora_lib/impl/context/agent_context.py

@@ -1,17 +1,24 @@
 from typing import Optional
+from uuid import UUID
 
 from nora_lib.impl.interactions.models import Surface
 from pydantic import BaseModel
 
+from nora_lib.serializers import UuidWithSerializer
+
 
 class MessageAgentContext(BaseModel):
     """
-    Identifiers for the triggering user message
+    Identifiers for the triggering message
     """
 
     message_id: str
     thread_id: str
     channel_id: str
+    # Pydantic models don't actually need this, but some codebases pass model_dump() objects
+    # to other JSON serializers (e.g. FastAPI) which don't handle UUIDs and other python built-ins.
+    # So we use this special UUID wrapper for convenience.
+    actor_id: UuidWithSerializer
     surface: Surface
 
 

{nora_lib_impl-1.0.2 → nora_lib_impl-1.2.0}/src/nora_lib/impl/interactions/models.py

@@ -5,10 +5,18 @@ Model for interactions to be sent to the interactions service.
 import json
 from datetime import datetime
 from enum import Enum
-from typing import Dict, List, Optional, Tuple
+from typing import Annotated, Dict, List, Literal, Optional, Tuple, Union
 from uuid import UUID
 
-from pydantic import BaseModel, Field, field_serializer, ConfigDict, model_validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Discriminator,
+    Field,
+    field_serializer,
+    field_validator,
+    model_validator,
+)
 
 
 class Surface(str, Enum):
@@ -300,25 +308,51 @@ class VirtualThread:
 class CostDetail(BaseModel):
     """
     Base class to store details of cost to service a request by an agent.
-    If an agent has different cost details,
-    it should create another class inheriting this class and add those fields.
-    See LLMCost and LLMTokenBreakdown below for examples.
+    If an agent has different cost details (e.g. for non-llm costs), it should:
+
+    - create another class inheriting this class and add any additional fields
+    - give the class a unique detail_type
+    - add the class to CostDetailType below
+
+    See LLMCost below as an example.
+
+    To add new details about the cost of an LLM call, add to LLMCost below
+    rather than creating a new detail.  A ServiceCost can have details of
+    multiple LLM calls, so if details are split into multiple instances then
+    they can't easily be connected.
+
+    Any new fields added to an existing subclass must have a default value
+    specified for backward-compatibility.
     """
 
-    model_config = ConfigDict(protected_namespaces=())
-    pass
+    detail_type: str = "unknown"
 
+    model_config = ConfigDict(protected_namespaces=(), extra="allow")
 
-class LLMCost(CostDetail):
-    """LLM cost detail"""
+    def try_subclass_conversion(self):
+        """For events with no detail_type, attempt to convert to an appropriate
+        subclass based on the fields.  This is useful for handling legacy
+        events but should not be needed for new events that have detail_type."""
+        # Already a subclass
+        if type(self) is not CostDetail:
+            return self
 
-    token_count: int
-    model_name: str
+        d = self.dict()
+        del d["detail_type"]
+        if "token_count" in d and "model_name" in d:
+            return LLMCost(**d)
+        if "prompt_tokens" in d and "completion_tokens" in d:
+            return LLMTokenBreakdown(**d)
+        if "run_id" in d:
+            return LangChainRun(**d)
+        return self
 
 
 class LLMTokenBreakdown(CostDetail):
     """Token usage breakdown"""
 
+    detail_type: Literal["llm_token_breakdown"] = "llm_token_breakdown"
+
     prompt_tokens: int
     completion_tokens: int
 
@@ -326,6 +360,8 @@ class LLMTokenBreakdown(CostDetail):
 class LangChainRun(CostDetail):
     """LangChain Run"""
 
+    detail_type: Literal["langchain_run"] = "langchain_run"
+
     # Subset of run fields which allow future lookup of run details.
     run_id: UUID
     run_name: Optional[str] = None
@@ -336,15 +372,51 @@ class LangChainRun(CostDetail):
     # Serialize the UUIDs as strings
     @field_serializer("run_id")
     def serialize_id(self, run_id: UUID):
-        return str(run_id)
+        return str(run_id) if run_id is not None else None
 
     @field_serializer("trace_id")
     def serialize_trace_id(self, trace_id: UUID):
-        return str(trace_id)
+        return str(trace_id) if trace_id is not None else None
 
     @field_serializer("session_id")
     def serialize_session_id(self, session_id: UUID):
-        return str(session_id)
+        return str(session_id) if session_id is not None else None
+
+    # Validators to handle legacy "None" strings
+    @field_validator("trace_id", "session_id", mode="before")
+    @classmethod
+    def validate_optional_uuid(cls, value):
+        if value == "None":
+            return None
+        return value
+
+
+class LLMCost(CostDetail):
+    """Details for the cost/usage of an LLM call."""
+
+    detail_type: Literal["llm_cost"] = "llm_cost"
+
+    model_name: str
+    token_count: int
+
+    token_breakdown: Optional[LLMTokenBreakdown] = None
+
+
+# Note: CostDetailType is a Union of all the subclasses of CostDetail, with
+# a discriminator for pydantic deserialization
+CostDetailType = Union[
+    Annotated[
+        Union[
+            LLMCost,
+            LLMTokenBreakdown,
+            LangChainRun,
+        ],
+        Discriminator("detail_type"),
+    ],
+    # We fall back to the base class if the discriminator is not found (legacy
+    # events or custom ones written by other apps may be missing detail_type)
+    CostDetail,
+]
 
 
 class ServiceCost(BaseModel):
@@ -363,10 +435,32 @@ class ServiceCost(BaseModel):
         description="Agent generated task_id used to track nora assigned tasks",
     )
     tool_call_id: Optional[str] = None
-    details: List[CostDetail] = []
+    details: list[CostDetailType] = Field(default_factory=list)
     env: Optional[str] = None
     git_sha: Optional[str] = None
 
+    def with_unified_llm_costs(self) -> "ServiceCost":
+        """
+        Creates a new ServiceCost object with unified LLMCost details.
+
+        This method converts old-style separate LLMCost and LLMTokenBreakdown details
+        into unified LLMCost objects that include token breakdown information.
+
+        If the details are already in the unified format (no separate LLMTokenBreakdown
+        objects), the method will return a copy without modifying the details.
+
+        Returns:
+            A new ServiceCost object with unified LLMCost details
+
+        Raises:
+            ValueError: If there is ambiguity in matching LLMCost with LLMTokenBreakdown
+        """
+        unified_details = unify_llm_cost_details(self.details)
+
+        data = self.model_dump()
+        data["details"] = unified_details
+        return ServiceCost.model_validate(data)
+
 
 class StepCost(BaseModel):
     """Wrapping service cost with event metadata so that it can be converted to an Event object."""
@@ -402,6 +496,112 @@ class StepCost(BaseModel):
         )
 
 
+def unify_llm_cost_details(details: List[CostDetailType]) -> List[CostDetailType]:
+    """
+    Convert a list of old-style LLMCost and LLMTokenBreakdown details into a list
+    where these pairs are combined into the new unified LLMCost format.
+
+    This is useful for migrating old ServiceCost events which might have separate
+    LLMCost and LLMTokenBreakdown details for the same LLM call.
+
+    If the function encounters ambiguity (can't confidently pair LLMCost with
+    its corresponding LLMTokenBreakdown), it will raise a ValueError.
+
+    Args:
+        details: List of CostDetail objects that might contain old-style
+                LLMCost and LLMTokenBreakdown details
+
+    Returns:
+        A new list where matching LLMCost and LLMTokenBreakdown pairs are combined
+        into unified LLMCost objects, and other details are left unchanged.
+
+    Raises:
+        ValueError: If there is ambiguity in matching LLMCost with LLMTokenBreakdown
+    """
+    result: List[CostDetailType] = []
+    llm_costs: List[LLMCost] = []
+    token_breakdowns: List[LLMTokenBreakdown] = []
+    other_details: List[CostDetailType] = []
+
+    if any(
+        isinstance(detail, LLMCost) and detail.token_breakdown is not None
+        for detail in details
+    ) and any(isinstance(detail, LLMTokenBreakdown) for detail in details):
+        raise ValueError(
+            "Cannot mix LLMCost with token breakdowns with unified LLMCost details"
+        )
+
+    # Separate details by type
+    for detail in details:
+        if isinstance(detail, LLMCost):
+            llm_costs.append(detail)
+        elif isinstance(detail, LLMTokenBreakdown):
+            token_breakdowns.append(detail)
+        else:
+            other_details.append(detail)
+
+    # If there are no token breakdowns, just return the original details
+    if not token_breakdowns:
+        return details
+
+    # If there's a mismatch in count, we can't be confident about pairing
+    if len(llm_costs) != len(token_breakdowns):
+        raise ValueError(
+            f"Cannot confidently pair LLMCost and LLMTokenBreakdown details: "
+            f"Found {len(llm_costs)} LLMCost and {len(token_breakdowns)} LLMTokenBreakdown details."
+        )
+
+    # Try to pair costs with breakdowns based on token counts
+    matched_breakdowns = set()
+
+    for cost in llm_costs:
+        match_found = False
+        matching_breakdown = None
+        matched_index = -1
+
+        # Look for a token breakdown where total tokens matches the cost token count
+        for i, breakdown in enumerate(token_breakdowns):
+            if i in matched_breakdowns:
+                continue
+
+            total_tokens = breakdown.prompt_tokens + breakdown.completion_tokens
+            if total_tokens == cost.token_count:
+                if match_found:
+                    # If we already found a matching breakdown, we have ambiguity
+                    raise ValueError(
+                        f"Ambiguity in matching LLMCost with token count {cost.token_count} "
+                        f"to LLMTokenBreakdown - multiple matches found."
+                    )
+                match_found = True
+                matching_breakdown = breakdown
+                matched_index = i
+
+        if match_found and matching_breakdown is not None:
+            # Create a new unified LLMCost object
+            new_cost = LLMCost(
+                model_name=cost.model_name,
+                token_count=cost.token_count,
+                token_breakdown=matching_breakdown,
+            )
+            result.append(new_cost)
+            matched_breakdowns.add(matched_index)
+        else:
+            # If no matching breakdown was found based on token count
+            raise ValueError(
+                f"Could not find matching LLMTokenBreakdown for LLMCost with token count {cost.token_count}"
+            )
+
+    # Check if all breakdowns were matched
+    if len(matched_breakdowns) != len(token_breakdowns):
+        raise ValueError("Not all LLMTokenBreakdown details were matched to an LLMCost")
+
+    # Add the other details
+    for detail in other_details:
+        result.append(detail)
+
+    return result
+
+
 def thread_message_lookup_request(message_id: str, event_type: str) -> dict:
     """retrieve messages and events for the thread associated with a message"""
     return {

{nora_lib_impl-1.0.2 → nora_lib_impl-1.2.0}/src/nora_lib_impl.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nora_lib-impl
-Version: 1.0.2
+Version: 1.2.0
 Summary: For making and coordinating agents and tools
 Home-page: https://github.com/allenai/nora_lib/impl
 Requires-Python: >=3.9