Add cost tracking functionality across various components, including Executor, Planner, and FieldMapperAgent. Integrate CostTracker to monitor LLM and document intelligence costs, enhancing logging for cost-related metrics and providing detailed cost breakdowns in the user interface.

src/agents/field_mapper_agent.py

@@ -35,7 +35,25 @@ class FieldMapperAgent(BaseAgent):
         try:
             self.logger.info("Inferring document context...")
             self.logger.debug(f"Using text preview: {text[:500]}...")
-            context = self.llm.responses(prompt, temperature=0.0)
+            
+            # Get cost tracker from context
+            cost_tracker = self.ctx.get("cost_tracker") if hasattr(self, 'ctx') else None
+            if cost_tracker:
+                self.logger.info("Cost tracker found in context")
+            else:
+                self.logger.warning("No cost tracker found in context")
+            
+            context = self.llm.responses(
+                prompt, temperature=0.0,
+                ctx={"cost_tracker": cost_tracker} if cost_tracker else None,
+                description="Document Context Inference"
+            )
+            
+            # Log cost tracking results if available
+            if cost_tracker:
+                self.logger.info(f"Context inference costs - Input tokens: {cost_tracker.llm_input_tokens}, Output tokens: {cost_tracker.llm_output_tokens}")
+                self.logger.info(f"Context inference cost: ${cost_tracker.calculate_current_file_costs()['openai']['total_cost']:.4f}")
+            
             self.logger.info(f"Inferred context: {context}")
             return context
         except Exception as e:
@@ -142,7 +160,25 @@ class FieldMapperAgent(BaseAgent):
         try:
             self.logger.info(f"Calling LLM to extract value for field '{field}'")
             self.logger.debug(f"Using prompt: {prompt}")
-            value = self.llm.responses(prompt, temperature=0.0)
+            
+            # Get cost tracker from context
+            cost_tracker = self.ctx.get("cost_tracker") if hasattr(self, 'ctx') else None
+            if cost_tracker:
+                self.logger.info("Cost tracker found in context")
+            else:
+                self.logger.warning("No cost tracker found in context")
+            
+            value = self.llm.responses(
+                prompt, temperature=0.0,
+                ctx={"cost_tracker": cost_tracker} if cost_tracker else None,
+                description=f"Field Extraction - {field} (Search)"
+            )
+            
+            # Log cost tracking results if available
+            if cost_tracker:
+                self.logger.info(f"Field extraction costs - Input tokens: {cost_tracker.llm_input_tokens}, Output tokens: {cost_tracker.llm_output_tokens}")
+                self.logger.info(f"Field extraction cost: ${cost_tracker.calculate_current_file_costs()['openai']['total_cost']:.4f}")
+            
             self.logger.debug(f"Raw LLM response: {value}")
             
             if value and value.lower() not in ["none", "null", "n/a"]:
@@ -191,7 +227,25 @@ class FieldMapperAgent(BaseAgent):
         
         try:
             self.logger.info(f"Calling LLM to extract value for field '{field}' from page")
-            value = self.llm.responses(prompt, temperature=0.0)
+            
+            # Get cost tracker from context
+            cost_tracker = self.ctx.get("cost_tracker") if hasattr(self, 'ctx') else None
+            if cost_tracker:
+                self.logger.info("Cost tracker found in context")
+            else:
+                self.logger.warning("No cost tracker found in context")
+            
+            value = self.llm.responses(
+                prompt, temperature=0.0,
+                ctx={"cost_tracker": cost_tracker} if cost_tracker else None,
+                description=f"Field Extraction - {field} (Page)"
+            )
+            
+            # Log cost tracking results if available
+            if cost_tracker:
+                self.logger.info(f"Page extraction costs - Input tokens: {cost_tracker.llm_input_tokens}, Output tokens: {cost_tracker.llm_output_tokens}")
+                self.logger.info(f"Page extraction cost: ${cost_tracker.calculate_current_file_costs()['openai']['total_cost']:.4f}")
+            
             self.logger.debug(f"Raw LLM response: {value}")
             
             if value and value.lower() not in ["none", "null", "n/a"]:

src/agents/pdf_agent.py

@@ -25,4 +25,10 @@ class PDFAgent(BaseAgent):
         pdf_bytes = pdf_file.read()
         text = self._extract_text(pdf_bytes)
         ctx["text"] = text
+
+        # After extracting pages
+        num_pages = len(fitz.open(stream=pdf_bytes, filetype="pdf"))  # type: ignore[arg-type]
+        if "cost_tracker" in ctx:
+            ctx["cost_tracker"].add_di_pages(num_pages)
+
         return text

src/app.py

@@ -13,6 +13,7 @@ from datetime import datetime
 import io
 import sys
 from io import StringIO
+from services.cost_tracker import CostTracker
 
 # Create a custom stream handler to capture logs
 class LogCaptureHandler(logging.StreamHandler):
@@ -317,7 +318,10 @@ else:  # page == "Execution"
                 doc = fitz.open(stream=pdf_file.getvalue(), filetype="pdf")  # type: ignore[arg-type]
                 preview = "\n".join(page.get_text() for page in doc[:10])[:20000]  # first 2 pages, 2k chars
 
-                planner = Planner()
+                # Create a cost tracker for this run
+                cost_tracker = CostTracker()
+
+                planner = Planner(cost_tracker=cost_tracker)
                 plan = planner.build_plan(
                     pdf_meta={"filename": pdf_file.name},
                     doc_preview=preview,
@@ -329,9 +333,28 @@ else:  # page == "Execution"
                 st.markdown("---")
 
             with st.spinner("Executing …"):
-                executor = Executor(settings=settings)
+                executor = Executor(settings=settings, cost_tracker=cost_tracker)
                 results, logs = executor.run(plan, pdf_file)
 
+                # Get detailed costs
+                costs = executor.cost_tracker.calculate_current_file_costs()
+                model_cost = costs["openai"]["total_cost"]
+                di_cost = costs["document_intelligence"]["total_cost"]
+
+                # Display detailed costs table
+                st.subheader("Detailed Costs")
+                costs_df = executor.cost_tracker.get_detailed_costs_table()
+                st.dataframe(costs_df, use_container_width=True)
+
+                st.info(
+                    f"LLM input tokens: {executor.cost_tracker.llm_input_tokens}, "
+                    f"LLM output tokens: {executor.cost_tracker.llm_output_tokens}, "
+                    f"DI pages: {executor.cost_tracker.di_pages}, "
+                    f"Model cost: ${model_cost:.4f}, "
+                    f"DI cost: ${di_cost:.4f}, "
+                    f"Total cost: ${model_cost + di_cost:.4f}"
+                )
+
                 # Add detailed logging about what executor returned
                 logger.info(f"Executor returned results of type: {type(results)}")
                 logger.info(f"Results content: {results}")

src/orchestrator/executor.py

@@ -16,9 +16,12 @@ from agents.semantic_reasoner import SemanticReasonerAgent
 from agents.confidence_scorer import ConfidenceScorer
 from agents.query_generator import QueryGenerator
 
+# Add import for CostTracker
+from services.cost_tracker import CostTracker
+
 
 class Executor:
-    def __init__(self, settings):
+    def __init__(self, settings, cost_tracker=None):
         # map name → instance
         self.tools = {
             "PDFAgent": PDFAgent(),
@@ -32,6 +35,7 @@ class Executor:
 
         self.logs: List[Dict[str, Any]] = []
         self.logger = logging.getLogger(__name__)
+        self.cost_tracker = cost_tracker or CostTracker()
 
     # ---------------------------------------------------------
     def run(self, plan: Dict[str, Any], pdf_file) -> tuple[pd.DataFrame, List[Dict[str, Any]]]:  # noqa: D401
@@ -47,6 +51,7 @@ class Executor:
             "results": [],
             "conf": 1.0,
             "pdf_meta": plan.get("pdf_meta", {}),  # Include the plan's metadata
+            "cost_tracker": self.cost_tracker,
         }
         
         try:
@@ -112,6 +117,12 @@ class Executor:
         else:
             df = pd.DataFrame()
             
+        # At the end, log the costs
+        self.logger.info(f"Total LLM input tokens: {self.cost_tracker.llm_input_tokens}")
+        self.logger.info(f"Total LLM output tokens: {self.cost_tracker.llm_output_tokens}")
+        self.logger.info(f"Total DI pages: {self.cost_tracker.di_pages}")
+        self.logger.info(f"Total cost: ${self.cost_tracker.total_cost():.4f}")
+        
         return df, self.logs
 
     # ---------------------------------------------------------

src/orchestrator/planner.py

@@ -11,6 +11,7 @@ import yaml
 
 from services.llm_client import LLMClient
 from config.settings import settings
+from services.cost_tracker import CostTracker
 
 
 _PROMPTS_FILE = Path(__file__).parent.parent / "config" / "prompts.yaml"
@@ -23,9 +24,10 @@ logger = logging.getLogger(__name__)
 class Planner:
     """Generate a plan with the Responses API; fall back to a static template if parsing fails."""
 
-    def __init__(self) -> None:
+    def __init__(self, cost_tracker=None) -> None:
         self.prompt_template = self._load_prompt("planner")
         self.llm = LLMClient(settings)
+        self.cost_tracker = cost_tracker or CostTracker()
         logger.info("Planner initialized with prompt template")
 
     # --------------------------------------------------
@@ -53,7 +55,12 @@ class Planner:
 
         try:
             logger.info("Calling LLM to generate plan")
-            raw = self.llm.responses(prompt, temperature=0.0)
+            raw = self.llm.responses(
+                prompt,
+                temperature=0.0,
+                ctx={"cost_tracker": self.cost_tracker},
+                description="Execution Plan Generation"
+            )
             logger.debug(f"Raw LLM response: {raw}")
             
             try:

src/services/cost_tracker.py

@@ -0,0 +1,221 @@
+"""Cost tracking service for Azure OpenAI and Document Intelligence."""
+
+from dataclasses import dataclass
+from typing import Dict, Optional, List
+import logging
+from datetime import datetime
+import pandas as pd
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class TokenCosts:
+    """Costs per 1M tokens for different models."""
+    GPT41: float = 2.0  # $10 per 1M input tokens for GPT-4.1
+    GPT41_OUTPUT: float = 8.0  # $30 per 1M output tokens for GPT-4.1
+
+@dataclass
+class DocumentIntelligenceCosts:
+    """Costs for Document Intelligence."""
+    READ: float = 1.5  # $1.5 per 1,000 pages
+    LAYOUT: float = 10.0  # $10 per 1,000 pages
+    CUSTOM: float = 30.0  # $30 per 1,000 pages
+
+@dataclass
+class LLMCall:
+    """Represents a single LLM call with its details."""
+    description: str
+    input_tokens: int
+    output_tokens: int
+    timestamp: datetime = datetime.now()
+
+    @property
+    def input_cost(self) -> float:
+        return (self.input_tokens / 1_000_000) * TokenCosts.GPT41
+
+    @property
+    def output_cost(self) -> float:
+        return (self.output_tokens / 1_000_000) * TokenCosts.GPT41_OUTPUT
+
+    @property
+    def total_cost(self) -> float:
+        return self.input_cost + self.output_cost
+
+class CostTracker:
+    """Tracks costs for Azure OpenAI and Document Intelligence usage."""
+    
+    def __init__(self):
+        self.token_costs = TokenCosts()
+        self.di_costs = DocumentIntelligenceCosts()
+        self.current_file_costs = {
+            "input_tokens": 0,
+            "output_tokens": 0,
+            "di_pages": 0,
+            "di_operations": {}  # Track different DI operations
+        }
+        self.total_costs = {
+            "input_tokens": 0,
+            "output_tokens": 0,
+            "di_pages": 0,
+            "di_operations": {}
+        }
+        self.llm_input_tokens = 0
+        self.llm_output_tokens = 0
+        self.di_pages = 0
+        self.llm_calls: List[LLMCall] = []  # Track individual LLM calls
+    
+    def reset_current_file(self):
+        """Reset costs for current file."""
+        self.current_file_costs = {
+            "input_tokens": 0,
+            "output_tokens": 0,
+            "di_pages": 0,
+            "di_operations": {}
+        }
+        self.llm_calls = []  # Reset LLM calls for new file
+    
+    def add_tokens(self, input_tokens: int, output_tokens: int, model: str = "GPT41"):
+        """Add tokens for current file and total."""
+        self.current_file_costs["input_tokens"] += input_tokens
+        self.current_file_costs["output_tokens"] += output_tokens
+        self.total_costs["input_tokens"] += input_tokens
+        self.total_costs["output_tokens"] += output_tokens
+        
+        logger.info(f"Added tokens - Input: {input_tokens}, Output: {output_tokens} for model {model}")
+    
+    def add_di_operation(self, operation: str, pages: int):
+        """Add Document Intelligence operation costs."""
+        if operation not in self.current_file_costs["di_operations"]:
+            self.current_file_costs["di_operations"][operation] = 0
+            self.total_costs["di_operations"][operation] = 0
+            
+        self.current_file_costs["di_operations"][operation] += pages
+        self.current_file_costs["di_pages"] += pages
+        self.total_costs["di_operations"][operation] += pages
+        self.total_costs["di_pages"] += pages
+        
+        logger.info(f"Added DI operation - {operation}: {pages} pages")
+    
+    def calculate_current_file_costs(self) -> Dict:
+        """Calculate costs for current file."""
+        costs = {
+            "openai": {
+                "input_tokens": self.current_file_costs["input_tokens"],
+                "output_tokens": self.current_file_costs["output_tokens"],
+                "input_cost": (self.current_file_costs["input_tokens"] / 1_000_000) * self.token_costs.GPT41,
+                "output_cost": (self.current_file_costs["output_tokens"] / 1_000_000) * self.token_costs.GPT41_OUTPUT,
+                "total_cost": 0,
+                "calls": [call.__dict__ for call in self.llm_calls]  # Include detailed call information
+            },
+            "document_intelligence": {
+                "total_pages": self.current_file_costs["di_pages"],
+                "operations": {},
+                "total_cost": 0
+            }
+        }
+        
+        # Calculate OpenAI total cost
+        costs["openai"]["total_cost"] = costs["openai"]["input_cost"] + costs["openai"]["output_cost"]
+        
+        # Calculate Document Intelligence costs
+        for operation, pages in self.current_file_costs["di_operations"].items():
+            cost = (pages / 1000) * getattr(self.di_costs, operation.upper(), self.di_costs.READ)
+            costs["document_intelligence"]["operations"][operation] = {
+                "pages": pages,
+                "cost": cost
+            }
+            costs["document_intelligence"]["total_cost"] += cost
+        
+        return costs
+    
+    def get_total_costs(self) -> Dict:
+        """Get total costs across all files."""
+        return {
+            "openai": {
+                "input_tokens": self.total_costs["input_tokens"],
+                "output_tokens": self.total_costs["output_tokens"],
+                "input_cost": (self.total_costs["input_tokens"] / 1_000_000) * self.token_costs.GPT41,
+                "output_cost": (self.total_costs["output_tokens"] / 1_000_000) * self.token_costs.GPT41_OUTPUT,
+                "total_cost": 0
+            },
+            "document_intelligence": {
+                "total_pages": self.total_costs["di_pages"],
+                "operations": {},
+                "total_cost": 0
+            }
+        }
+
+    def add_llm_tokens(self, input_tokens, output_tokens, description: str = "LLM Call"):
+        """Add tokens for an LLM call with a description."""
+        self.llm_input_tokens += input_tokens
+        self.llm_output_tokens += output_tokens
+        # Also update the current file costs
+        self.current_file_costs["input_tokens"] += input_tokens
+        self.current_file_costs["output_tokens"] += output_tokens
+        self.total_costs["input_tokens"] += input_tokens
+        self.total_costs["output_tokens"] += output_tokens
+        # Add to LLM calls list
+        self.llm_calls.append(LLMCall(
+            description=description,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens
+        ))
+
+    def add_di_pages(self, num_pages):
+        self.di_pages += num_pages
+        # Also update the current file costs
+        self.current_file_costs["di_pages"] += num_pages
+        self.total_costs["di_pages"] += num_pages
+        # Add to LAYOUT operation by default since we're using layout analysis
+        if "LAYOUT" not in self.current_file_costs["di_operations"]:
+            self.current_file_costs["di_operations"]["LAYOUT"] = 0
+            self.total_costs["di_operations"]["LAYOUT"] = 0
+        self.current_file_costs["di_operations"]["LAYOUT"] += num_pages
+        self.total_costs["di_operations"]["LAYOUT"] += num_pages
+
+    def total_cost(self):
+        """Calculate total cost using the same pricing as calculate_current_file_costs."""
+        costs = self.calculate_current_file_costs()
+        return costs["openai"]["total_cost"] + costs["document_intelligence"]["total_cost"]
+
+    def get_detailed_costs_table(self) -> pd.DataFrame:
+        """Return a DataFrame of detailed costs."""
+        if not self.llm_calls:
+            return pd.DataFrame()
+
+        # Create list of dictionaries for DataFrame
+        rows = []
+        for call in self.llm_calls:
+            rows.append({
+                'Description': call.description,
+                'Input Tokens': call.input_tokens,
+                'Output Tokens': call.output_tokens,
+                'Input Cost': f"${call.input_cost:.4f}",
+                'Output Cost': f"${call.output_cost:.4f}",
+                'Total Cost': f"${call.total_cost:.4f}"
+            })
+        
+        # Calculate totals
+        total_input = sum(call.input_tokens for call in self.llm_calls)
+        total_output = sum(call.output_tokens for call in self.llm_calls)
+        total_input_cost = sum(call.input_cost for call in self.llm_calls)
+        total_output_cost = sum(call.output_cost for call in self.llm_calls)
+        total_cost = total_input_cost + total_output_cost
+        
+        # Add total row
+        rows.append({
+            'Description': 'TOTAL',
+            'Input Tokens': total_input,
+            'Output Tokens': total_output,
+            'Input Cost': f"${total_input_cost:.4f}",
+            'Output Cost': f"${total_output_cost:.4f}",
+            'Total Cost': f"${total_cost:.4f}"
+        })
+        
+        # Create DataFrame
+        df = pd.DataFrame(rows)
+        
+        # Set column order
+        df = df[['Description', 'Input Tokens', 'Output Tokens', 'Input Cost', 'Output Cost', 'Total Cost']]
+        
+        return df 

src/services/llm_client.py

@@ -35,12 +35,15 @@ class LLMClient:
         logger.info(f"API Key length: {len(openai.api_key) if openai.api_key else 0}")
 
     # --------------------------------------------------
-    def responses(self, prompt: str, tools: List[dict] | None = None, **kwargs: Any) -> str:
+    def responses(self, prompt: str, tools: List[dict] | None = None, description: str = "LLM Call", **kwargs: Any) -> str:
         """Call the Responses API and return the assistant content as string."""
         logger = logging.getLogger(__name__)
         logger.info(f"Making request with API version: {openai.api_version}")
         logger.info(f"Request URL will be: {openai.api_base}/openai/responses?api-version={openai.api_version}")
         
+        # Remove ctx from kwargs before passing to openai
+        ctx = kwargs.pop("ctx", None)
+
         resp = openai.responses.create(
             input=prompt,
             model=self._deployment,
@@ -49,7 +52,31 @@ class LLMClient:
         )
         # Log the raw response for debugging
         logging.debug(f"LLM raw response: {resp}")
-        
+
+        # --- Cost tracking: must be BEFORE any return! ---
+        logger.info(f"LLMClient.responses: ctx is {ctx}")
+        if ctx and "cost_tracker" in ctx:
+            logger.info(f"LLMClient.responses: cost_tracker is {ctx['cost_tracker']}")
+            usage = getattr(resp, "usage", None)
+            if usage:
+                logger.info(f"LLMClient.responses: usage is {usage}")
+                ctx["cost_tracker"].add_llm_tokens(
+                    input_tokens=getattr(usage, "input_tokens", 0),
+                    output_tokens=getattr(usage, "output_tokens", 0),
+                    description=description
+                )
+                logger.info(f"LLMClient.responses: prompt: {prompt[:200]}...")  # Log first 200 chars
+                logger.info(f"LLMClient.responses: resp: {str(resp)[:200]}...")  # Log first 200 chars
+                if usage:
+                    logger.info(f"LLMClient.responses: usage.input_tokens={getattr(usage, 'input_tokens', None)}, usage.output_tokens={getattr(usage, 'output_tokens', None)}, usage.total_tokens={getattr(usage, 'total_tokens', None)}")
+            else:
+                # Fallback: estimate tokens (very rough)
+                ctx["cost_tracker"].add_llm_tokens(
+                    input_tokens=len(prompt.split()),
+                    output_tokens=len(str(resp).split()),
+                    description=description
+                )
+
         # Extract the text content from the response
         if hasattr(resp, "output") and isinstance(resp.output, list):
             # Handle list of ResponseOutputMessage objects