Spaces:
Running
Running
| """Map a single field to a candidate value using page-by-page analysis and LLM-based extraction.""" | |
| from typing import Dict, Any, Optional, List | |
| import logging | |
| import re | |
| import json | |
| from .base_agent import BaseAgent | |
| from services.llm_client import LLMClient | |
| from services.embedding_client import EmbeddingClient | |
| from config.settings import settings | |
| # Configure logging to disable verbose Azure HTTP logs | |
| logging.getLogger('azure.core.pipeline.policies.http_logging_policy').setLevel(logging.WARNING) | |
| logging.getLogger('azure.core.pipeline').setLevel(logging.WARNING) | |
| logging.getLogger('azure').setLevel(logging.WARNING) | |
| class FieldMapperAgent(BaseAgent): | |
| def __init__(self): | |
| self.logger = logging.getLogger(__name__) | |
| self.llm = LLMClient(settings) | |
| self.embedding_client = EmbeddingClient() | |
| def _infer_document_context(self, text: str) -> str: | |
| """Use LLM to infer document context and user profile.""" | |
| prompt = f"""Given this document text, describe the document type and typical user profile in 1-2 sentences. | |
| Focus on the domain, purpose, and who would use this document. | |
| Document text: | |
| {text[:2000]} # First 2000 chars for context | |
| Response format: | |
| Document type: [type] | |
| User profile: [profile] | |
| """ | |
| try: | |
| self.logger.info("Inferring document context...") | |
| self.logger.debug(f"Using text preview: {text[:500]}...") | |
| # 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: | |
| self.logger.error(f"Error inferring context: {str(e)}") | |
| return "Generic document user" | |
| def _find_similar_chunks_search(self, query: str, index: Dict[str, Any], top_k: int = 3) -> List[Dict[str, Any]]: | |
| """Find chunks semantically similar to the query using cosine similarity.""" | |
| try: | |
| self.logger.info(f"Finding similar chunks for query: {query}") | |
| self.logger.debug(f"Index contains {len(index['chunks'])} chunks and {len(index['embeddings'])} embeddings") | |
| # Get query embedding | |
| self.logger.debug("Generating embedding for query...") | |
| query_embedding = self.embedding_client.embed([query])[0] | |
| self.logger.debug(f"Query embedding generated, length: {len(query_embedding)}") | |
| # Calculate similarities | |
| similarities = [] | |
| for i, (chunk, embedding) in enumerate(zip(index["chunks"], index["embeddings"])): | |
| similarity = self._cosine_similarity(query_embedding, embedding) | |
| similarities.append((similarity, chunk)) | |
| self.logger.debug(f"Chunk {i} similarity: {similarity:.3f}") | |
| self.logger.debug(f"Chunk {i} preview: {chunk['text'][:100]}...") | |
| # Sort by similarity and return top k | |
| similarities.sort(reverse=True) | |
| results = [chunk for _, chunk in similarities[:top_k]] | |
| # Log top results | |
| self.logger.info(f"Found {len(results)} similar chunks") | |
| for i, (sim, chunk) in enumerate(similarities[:top_k]): | |
| self.logger.info(f"Top {i+1} match (similarity: {sim:.3f}): {chunk['text'][:200]}...") | |
| return results | |
| except Exception as e: | |
| self.logger.error(f"Error finding similar chunks: {str(e)}", exc_info=True) | |
| return [] | |
| def _cosine_similarity(self, a: List[float], b: List[float]) -> float: | |
| """Calculate cosine similarity between two vectors.""" | |
| import numpy as np | |
| try: | |
| # Check for zero vectors | |
| if not a or not b or all(x == 0 for x in a) or all(x == 0 for x in b): | |
| self.logger.warning("Zero vector detected in cosine similarity calculation") | |
| return 0.0 | |
| # Convert to numpy arrays | |
| a_np = np.array(a) | |
| b_np = np.array(b) | |
| # Calculate norms | |
| norm_a = np.linalg.norm(a_np) | |
| norm_b = np.linalg.norm(b_np) | |
| # Check for zero norms | |
| if norm_a == 0 or norm_b == 0: | |
| self.logger.warning("Zero norm detected in cosine similarity calculation") | |
| return 0.0 | |
| # Calculate similarity | |
| similarity = np.dot(a_np, b_np) / (norm_a * norm_b) | |
| # Check for NaN | |
| if np.isnan(similarity): | |
| self.logger.warning("NaN detected in cosine similarity calculation") | |
| return 0.0 | |
| return float(similarity) | |
| except Exception as e: | |
| self.logger.error(f"Error calculating cosine similarity: {str(e)}") | |
| return 0.0 | |
| def _extract_field_value_search(self, field: str, chunks: List[Dict[str, Any]], context: str) -> Optional[str]: | |
| """Use LLM to extract field value from relevant chunks.""" | |
| # Combine chunks into context | |
| chunk_texts = [chunk["text"] for chunk in chunks] | |
| combined_context = "\n".join(chunk_texts) | |
| self.logger.info(f"Extracting value for field '{field}' from {len(chunks)} chunks") | |
| self.logger.debug(f"Combined context preview: {combined_context[:500]}...") | |
| # Get filename from context if available | |
| filename = self.ctx.get("pdf_meta", {}).get("filename", "") | |
| filename_context = f"\nDocument filename: {filename}" if filename else "" | |
| # Get field descriptions from context if available | |
| field_descriptions = self.ctx.get("field_descriptions", {}) | |
| # Format field descriptions for the prompt | |
| field_descriptions_text = "" | |
| if field_descriptions and field in field_descriptions: | |
| desc_info = field_descriptions[field] | |
| if isinstance(desc_info, dict): | |
| description = desc_info.get('description', '') | |
| format_info = desc_info.get('format', '') | |
| examples = desc_info.get('examples', '') | |
| possible_values = desc_info.get('possible_values', '') | |
| field_descriptions_text = f"\nField information for '{field}':" | |
| if description: | |
| field_descriptions_text += f"\nDescription: {description}" | |
| if format_info: | |
| field_descriptions_text += f"\nFormat: {format_info}" | |
| if examples: | |
| field_descriptions_text += f"\nExamples: {examples}" | |
| if possible_values: | |
| field_descriptions_text += f"\nPossible Values: {possible_values}" | |
| field_descriptions_text += "\n" | |
| prompt = f"""You are an expert in {context} | |
| Your task is to extract the value for the field: {field}{filename_context}{field_descriptions_text} | |
| Consider the following context from the document: | |
| {combined_context} | |
| Instructions: | |
| 1. Look for the field value in the context | |
| 2. If you find multiple potential values, choose the most relevant one | |
| 3. If you're not sure, return None | |
| 4. Return ONLY the value, no explanations | |
| Field value:""" | |
| try: | |
| self.logger.info(f"Calling LLM to extract value for field '{field}'") | |
| self.logger.debug(f"Using prompt: {prompt}") | |
| # 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"]: | |
| self.logger.info(f"Successfully extracted value: {value}") | |
| return value.strip() | |
| else: | |
| self.logger.warning(f"LLM returned no valid value for field '{field}'") | |
| return None | |
| except Exception as e: | |
| self.logger.error(f"Error extracting field value: {str(e)}", exc_info=True) | |
| return None | |
| def _extract_field_value_from_page(self, field: str, page_text: str, context: str) -> Optional[str]: | |
| """Use LLM to extract field value from a single page.""" | |
| self.logger.info(f"Extracting value for field '{field}' from page") | |
| self.logger.debug(f"Page text preview: {page_text[:500]}...") | |
| # Get filename from context if available | |
| filename = self.ctx.get("pdf_meta", {}).get("filename", "") | |
| filename_context = f"\nDocument filename: {filename}" if filename else "" | |
| # Get field descriptions from context if available | |
| field_descriptions = self.ctx.get("field_descriptions", {}) | |
| # Format field descriptions for the prompt | |
| field_descriptions_text = "" | |
| if field_descriptions and field in field_descriptions: | |
| desc_info = field_descriptions[field] | |
| if isinstance(desc_info, dict): | |
| description = desc_info.get('description', '') | |
| format_info = desc_info.get('format', '') | |
| examples = desc_info.get('examples', '') | |
| possible_values = desc_info.get('possible_values', '') | |
| field_descriptions_text = f"\nField information for '{field}':" | |
| if description: | |
| field_descriptions_text += f"\nDescription: {description}" | |
| if format_info: | |
| field_descriptions_text += f"\nFormat: {format_info}" | |
| if examples: | |
| field_descriptions_text += f"\nExamples: {examples}" | |
| if possible_values: | |
| field_descriptions_text += f"\nPossible Values: {possible_values}" | |
| field_descriptions_text += "\n" | |
| prompt = f"""You are an expert in {context} | |
| Your task is to extract the value for the field: {field}{filename_context}{field_descriptions_text} | |
| Consider the following page from the document: | |
| {page_text} | |
| Instructions: | |
| 1. Look for the field values in this page | |
| 2. Return the data in a tabular format where each field is a column | |
| 3. Each field should have an array of values | |
| 4. The arrays must be aligned (same length) to represent rows | |
| 5. Return ONLY the JSON value, no explanations | |
| 6. Format the response as a valid JSON object with field names as keys | |
| 7. Keep the structure flat - do not nest values under 'details' or other keys | |
| Example response format: | |
| {{ | |
| "field1": ["value1", "value2", "value3"], | |
| "field2": ["value4", "value5", "value6"], | |
| "field3": ["value7", "value8", "value9"] | |
| }} | |
| Field value:""" | |
| try: | |
| self.logger.info(f"Calling LLM to extract value for field '{field}' from page") | |
| # 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"]: | |
| # Try to parse as JSON to ensure it's valid | |
| try: | |
| json_value = json.loads(value) | |
| self.logger.info(f"Successfully extracted value: {json.dumps(json_value, indent=2)}") | |
| return json.dumps(json_value, indent=2) | |
| except json.JSONDecodeError: | |
| # If not valid JSON, wrap it in a JSON object | |
| json_value = {field: value.strip()} | |
| self.logger.info(f"Wrapped non-JSON value in JSON object: {json.dumps(json_value, indent=2)}") | |
| return json.dumps(json_value, indent=2) | |
| else: | |
| self.logger.warning(f"LLM returned no valid value for field '{field}'") | |
| return None | |
| except Exception as e: | |
| self.logger.error(f"Error extracting field value from page: {str(e)}", exc_info=True) | |
| return None | |
| def _extract_with_unique_indices(self, text: str, context: str, unique_indices: List[str], fields_to_extract: List[str]) -> Optional[str]: | |
| """Extract values using unique indices strategy.""" | |
| self.logger.info(f"Using unique indices strategy with indices: {unique_indices}") | |
| self.logger.info(f"Fields to extract: {fields_to_extract}") | |
| # Get filename from context if available | |
| filename = self.ctx.get("pdf_meta", {}).get("filename", "") | |
| filename_context = f"\nDocument filename: {filename}" if filename else "" | |
| # Get field descriptions from context if available | |
| field_descriptions = self.ctx.get("field_descriptions", {}) | |
| unique_indices_descriptions = self.ctx.get("unique_indices_descriptions", {}) | |
| # Format field descriptions for the prompt | |
| field_descriptions_text = "" | |
| if field_descriptions: | |
| field_descriptions_text = "\nField descriptions:\n" | |
| for field, desc_info in field_descriptions.items(): | |
| if isinstance(desc_info, dict): | |
| description = desc_info.get('description', '') | |
| format_info = desc_info.get('format', '') | |
| examples = desc_info.get('examples', '') | |
| possible_values = desc_info.get('possible_values', '') | |
| desc_line = f" {field}:" | |
| if description: | |
| desc_line += f" {description}" | |
| if format_info: | |
| desc_line += f" (Format: {format_info})" | |
| if examples: | |
| desc_line += f" (Examples: {examples})" | |
| if possible_values: | |
| desc_line += f" (Possible Values: {possible_values})" | |
| field_descriptions_text += desc_line + "\n" | |
| else: | |
| field_descriptions_text += f" {field}: {desc_info}\n" | |
| # Format unique indices descriptions for the prompt | |
| unique_indices_text = "" | |
| if unique_indices_descriptions: | |
| unique_indices_text = "\nUnique indices descriptions:\n" | |
| for index, desc_info in unique_indices_descriptions.items(): | |
| if isinstance(desc_info, dict): | |
| description = desc_info.get('description', '') | |
| format_info = desc_info.get('format', '') | |
| examples = desc_info.get('examples', '') | |
| possible_values = desc_info.get('possible_values', '') | |
| desc_line = f" {index}:" | |
| if description: | |
| desc_line += f" {description}" | |
| if format_info: | |
| desc_line += f" (Format: {format_info})" | |
| if examples: | |
| desc_line += f" (Examples: {examples})" | |
| if possible_values: | |
| desc_line += f" (Possible Values: {possible_values})" | |
| unique_indices_text += desc_line + "\n" | |
| else: | |
| unique_indices_text += f" {index}: {desc_info}\n" | |
| prompt = f"""You are an expert in {context} | |
| Your task is to extract information from the document based on unique combinations of indices and their corresponding fields. | |
| Unique Indices to look for: {', '.join(unique_indices)} | |
| Fields to extract for each combination: {', '.join(fields_to_extract)}{filename_context}{field_descriptions_text}{unique_indices_text} | |
| Consider the following document: | |
| {text} | |
| Instructions: | |
| 1. First, identify all unique combinations of the specified indices ({', '.join(unique_indices)}) in the document | |
| 2. For each unique combination found, extract the values for all specified fields ({', '.join(fields_to_extract)}) | |
| 3. Return the data in a tabular format where: | |
| - Each row represents a unique combination | |
| - Each column represents a field value | |
| 4. Return ONLY the JSON value, no explanations | |
| 5. Format the response as a valid JSON object with field names as keys | |
| 6. Keep the structure flat - do not nest values under 'details' or other keys | |
| Example response format: | |
| {{ | |
| "index1": ["value1", "value2", "value3"], | |
| "index2": ["value4", "value5", "value6"], | |
| "field1": ["value7", "value8", "value9"], | |
| "field2": ["value10", "value11", "value12"] | |
| }} | |
| Note: Each array in the response must have the same length, representing aligned rows of data. | |
| For example, if there are 3 unique combinations found, each array should have 3 values. | |
| Field values:""" | |
| try: | |
| self.logger.info("Calling LLM for unique indices extraction") | |
| # Get cost tracker from context | |
| cost_tracker = self.ctx.get("cost_tracker") if hasattr(self, 'ctx') else None | |
| value = self.llm.responses( | |
| prompt, temperature=0.0, | |
| ctx={"cost_tracker": cost_tracker} if cost_tracker else None, | |
| description="Unique Indices Field Extraction" | |
| ) | |
| # Log cost tracking results if available | |
| if cost_tracker: | |
| self.logger.info(f"Unique indices extraction costs - Input tokens: {cost_tracker.llm_input_tokens}, Output tokens: {cost_tracker.llm_output_tokens}") | |
| self.logger.info(f"Unique indices 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"]: | |
| # Try to parse as JSON to ensure it's valid | |
| try: | |
| json_value = json.loads(value) | |
| self.logger.info(f"Successfully extracted values: {json.dumps(json_value, indent=2)}") | |
| return json.dumps(json_value, indent=2) | |
| except json.JSONDecodeError: | |
| # If not valid JSON, wrap it in a JSON object | |
| json_value = {", ".join(unique_indices + fields_to_extract): value.strip()} | |
| self.logger.info(f"Wrapped non-JSON value in JSON object: {json.dumps(json_value, indent=2)}") | |
| return json.dumps(json_value, indent=2) | |
| else: | |
| self.logger.warning("LLM returned no valid value") | |
| return None | |
| except Exception as e: | |
| self.logger.error(f"Error in unique indices extraction: {str(e)}", exc_info=True) | |
| return None | |
| def execute(self, ctx: Dict[str, Any]): # noqa: D401 | |
| field = ctx.get("current_field") | |
| strategy = ctx.get("strategy", "original") # Default to original strategy | |
| self.logger.info(f"Starting field mapping for: {field} using strategy: {strategy}") | |
| # Store context for use in extraction methods | |
| self.ctx = ctx | |
| # Get text and index | |
| text = "" | |
| index = None | |
| if "index" in ctx and isinstance(ctx["index"], dict): | |
| index = ctx["index"] | |
| text = index.get("text", "") | |
| self.logger.info(f"Using text from index (length: {len(text)})") | |
| self.logger.debug(f"Index contains {len(index.get('chunks', []))} chunks") | |
| self.logger.debug(f"Index contains {len(index.get('embeddings', []))} embeddings") | |
| elif "text" in ctx: | |
| text = ctx["text"] | |
| self.logger.info(f"Using text from direct context (length: {len(text)})") | |
| if not text: | |
| self.logger.warning("No text content found in context or index") | |
| return None | |
| # Infer document context if not already present | |
| if "document_context" not in ctx: | |
| ctx["document_context"] = self._infer_document_context(text) | |
| self.logger.info(f"Using document context: {ctx['document_context']}") | |
| # Process based on selected strategy | |
| if strategy == "unique_indices": | |
| unique_indices = ctx.get("unique_indices", []) | |
| fields_to_extract = ctx.get("fields_to_extract", []) | |
| if not unique_indices or not fields_to_extract: | |
| self.logger.warning("Missing unique indices or fields to extract") | |
| return None | |
| return self._extract_with_unique_indices(text, ctx["document_context"], unique_indices, fields_to_extract) | |
| else: | |
| # Original strategy | |
| if not field: | |
| self.logger.warning("No field provided in context") | |
| return None | |
| self.logger.info(f"Processing field: {field}") | |
| self.logger.info("Processing entire document...") | |
| value = self._extract_field_value_from_page(field, text, ctx["document_context"]) | |
| if value: | |
| return value | |
| # If no value found, try the search-based approach as fallback | |
| self.logger.warning("No value found in document analysis, falling back to search-based approach") | |
| if index and "embeddings" in index: | |
| self.logger.info("Using semantic search with embeddings") | |
| search_query = f"{field} in {ctx['document_context']}" | |
| similar_chunks = self._find_similar_chunks_search(search_query, index) | |
| if similar_chunks: | |
| self.logger.info(f"Found {len(similar_chunks)} relevant chunks, attempting value extraction") | |
| value = self._extract_field_value_search(field, similar_chunks, ctx["document_context"]) | |
| if value: | |
| return value | |
| self.logger.warning(f"No candidate found for field: {field}") | |
| return f"<no candidate for {field}>" |