import os import certifi import logging import subprocess # For calling ffmpeg if needed from typing import List, Dict, Optional from llama_index.core.agent.workflow import ReActAgent from llama_index.core.tools import FunctionTool from llama_index.llms.google_genai import GoogleGenAI from llama_index.core.node_parser import SentenceSplitter from llama_index.core import Document # Setup logging logger = logging.getLogger(__name__) # Global Whisper model instance (lazy loaded) _whisper_model = None os.environ["SSL_CERT_FILE"] = certifi.where() # Helper function to load prompt from file def load_prompt_from_file(filename: str, default_prompt: str) -> str: """Loads a prompt from a text file.""" try: script_dir = os.path.dirname(__file__) prompt_path = os.path.join(script_dir, filename) with open(prompt_path, "r") as f: prompt = f.read() logger.info(f"Successfully loaded prompt from {prompt_path}") return prompt except FileNotFoundError: logger.warning(f"Prompt file {filename} not found at {prompt_path}. Using default.") return default_prompt except Exception as e: logger.error(f"Error loading prompt file {filename}: {e}", exc_info=True) return default_prompt # --- Tool Functions --- def summarize_text(text: str, max_length: int = 150, min_length: int = 30) -> str: """Summarize the provided text using an LLM.""" logger.info(f"Summarizing text (length: {len(text)} chars). Max/Min length: {max_length}/{min_length}") # Configuration for summarization LLM summarizer_llm_model = os.getenv("SUMMARIZER_LLM_MODEL", "models/gemini-1.5-flash") # Use flash for speed gemini_api_key = os.getenv("GEMINI_API_KEY") if not gemini_api_key: logger.error("GEMINI_API_KEY not found for summarization tool LLM.") return "Error: GEMINI_API_KEY not set for summarization." # Truncate input text if excessively long to avoid API limits/costs max_input_chars = 30000 # Example limit, adjust as needed if len(text) > max_input_chars: logger.warning(f"Input text truncated to {max_input_chars} chars for summarization.") text = text[:max_input_chars] prompt = ( f"Summarize the following text concisely. Aim for a length between {min_length} and {max_length} words. " f"Focus on the main points and key information.\n\n" f"TEXT:\n{text}\n\nSUMMARY:" ) try: llm = GoogleGenAI(api_key=gemini_api_key, model="gemini-2.5-pro-preview-03-25", temperature=0.05) logger.info(f"Using summarization LLM: {summarizer_llm_model}") response = llm.complete(prompt) summary = response.text.strip() logger.info(f"Summarization successful (output length: {len(summary.split())} words).") return summary except Exception as e: logger.error(f"LLM call failed during summarization: {e}", exc_info=True) return f"Error during summarization: {e}" def extract_entities(text: str, entity_types: List[str] = ["PERSON", "ORG", "GPE", "DATE", "EVENT"]) -> Dict[str, List[str]]: """Extract named entities (like people, organizations, locations, dates) from the text using an LLM.""" logger.info(f"Extracting entities (types: {entity_types}) from text (length: {len(text)} chars).") # Configuration for entity extraction LLM entity_llm_model = os.getenv("ENTITY_LLM_MODEL", "models/gemini-1.5-flash") # Use flash for speed gemini_api_key = os.getenv("GEMINI_API_KEY") if not gemini_api_key: logger.error("GEMINI_API_KEY not found for entity extraction tool LLM.") return {"error": "GEMINI_API_KEY not set for entity extraction."} # Truncate input text if excessively long max_input_chars = 30000 # Example limit if len(text) > max_input_chars: logger.warning(f"Input text truncated to {max_input_chars} chars for entity extraction.") text = text[:max_input_chars] # Define the desired output format clearly in the prompt prompt = ( f"Extract named entities from the following text. Identify entities of these types: {', '.join(entity_types)}. " f"Format the output as a JSON object where keys are the entity types (uppercase) and values are lists of unique strings found for that type. " f"If no entities of a type are found, include the key with an empty list.\n\n" f"TEXT:\n{text}\n\nJSON_OUTPUT:" ) try: llm = GoogleGenAI(api_key=gemini_api_key, model=entity_llm_model, response_mime_type="application/json") # Request JSON output logger.info(f"Using entity extraction LLM: {entity_llm_model}") response = llm.complete(prompt) # Attempt to parse the JSON response import json try: # The response might be wrapped in ```json ... ```, try to extract it json_str = response.text.strip() if json_str.startswith("```json"): json_str = json_str[7:] if json_str.endswith("```"): json_str = json_str[:-3] entities = json.loads(json_str.strip()) # Validate structure (optional but good practice) if not isinstance(entities, dict): raise ValueError("LLM response is not a JSON object.") # Ensure all requested types are present, even if empty for entity_type in entity_types: if entity_type not in entities: entities[entity_type] = [] elif not isinstance(entities[entity_type], list): logger.warning(f"Entity type {entity_type} value is not a list, converting.") entities[entity_type] = [str(entities[entity_type])] # Attempt conversion logger.info(f"Entity extraction successful. Found entities: { {k: len(v) for k, v in entities.items()} }") return entities except json.JSONDecodeError as json_err: logger.error(f"Failed to parse JSON response from LLM: {json_err}. Response text: {response.text}") return {"error": f"Failed to parse LLM JSON response: {json_err}"} except ValueError as val_err: logger.error(f"Invalid JSON structure from LLM: {val_err}. Response text: {response.text}") return {"error": f"Invalid JSON structure from LLM: {val_err}"} except Exception as e: logger.error(f"LLM call failed during entity extraction: {e}", exc_info=True) return {"error": f"Error during entity extraction: {e}"} def split_text_into_chunks(text: str, chunk_size: int = 1000, chunk_overlap: int = 200) -> List[str]: """Split a long text into smaller chunks suitable for processing.""" logger.info(f"Splitting text (length: {len(text)} chars) into chunks (size: {chunk_size}, overlap: {chunk_overlap}).") if not text: return [] try: splitter = SentenceSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) document = Document(text=text) nodes = splitter.get_nodes_from_documents([document]) chunks = [node.get_content() for node in nodes] logger.info(f"Text split into {len(chunks)} chunks.") return chunks except Exception as e: logger.error(f"Error splitting text: {e}", exc_info=True) # Fallback to simple splitting if SentenceSplitter fails logger.warning("Falling back to simple text splitting.") return [text[i:i + chunk_size] for i in range(0, len(text), chunk_size - chunk_overlap)] # --- Tool Definitions --- summarize_tool = FunctionTool.from_defaults( fn=summarize_text, name="summarize_text", description=( "Summarizes a given block of text. Useful for condensing long documents or articles. " "Input: text (str), Optional: max_length (int), min_length (int). Output: summary (str) or error." ), ) extract_entities_tool = FunctionTool.from_defaults( fn=extract_entities, name="extract_entities", description=( "Extracts named entities (people, organizations, locations, dates, events) from text. " "Input: text (str), Optional: entity_types (List[str]). Output: Dict[str, List[str]] or error dict." ), ) split_text_tool = FunctionTool.from_defaults( fn=split_text_into_chunks, name="split_text_into_chunks", description=( "Splits a long text document into smaller, overlapping chunks. " "Input: text (str), Optional: chunk_size (int), chunk_overlap (int). Output: List[str] of chunks." ), ) # --- Agent Initialization --- def initialize_text_analyzer_agent() -> ReActAgent: """Initializes the Text Analyzer Agent.""" logger.info("Initializing TextAnalyzerAgent...") # Configuration for the agent's main LLM agent_llm_model = os.getenv("TEXT_ANALYZER_AGENT_LLM_MODEL", "gemini-2.5-pro-preview-03-25") gemini_api_key = os.getenv("GEMINI_API_KEY") if not gemini_api_key: logger.error("GEMINI_API_KEY not found for TextAnalyzerAgent.") raise ValueError("GEMINI_API_KEY must be set for TextAnalyzerAgent") try: llm = GoogleGenAI(api_key=gemini_api_key, model="gemini-2.5-pro-preview-03-25", temperature=0.05) logger.info(f"Using agent LLM: {agent_llm_model}") # Load system prompt default_system_prompt = ("You are TextAnalyzerAgent... [Default prompt content - replace with actual]" # Placeholder ) system_prompt = load_prompt_from_file("../prompts/text_analyzer_prompt.txt", default_system_prompt) if system_prompt == default_system_prompt: logger.warning("Using default/fallback system prompt for TextAnalyzerAgent.") # Define available tools, including the audio tool if available tools = [summarize_tool, extract_entities_tool, split_text_tool] # Update agent description based on available tools agent_description = ( "Analyzes text content. Can summarize text (`summarize_text`), extract named entities (`extract_entities`), " "and split long texts (`split_text_into_chunks`)." ) agent = ReActAgent( name="text_analyzer_agent", description=agent_description, tools=tools, llm=llm, system_prompt=system_prompt, can_handoff_to=["planner_agent", "research_agent", "reasoning_agent", "verifier_agent", "advanced_validation_agent"], # Example handoffs ) logger.info("TextAnalyzerAgent initialized successfully.") return agent except Exception as e: logger.error(f"Error during TextAnalyzerAgent initialization: {e}", exc_info=True) raise # Example usage (for testing if run directly) if __name__ == "__main__": logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger.info("Running text_analyzer_agent.py directly for testing...") # Check required keys required_keys = ["GEMINI_API_KEY"] missing_keys = [key for key in required_keys if not os.getenv(key)] if missing_keys: print(f"Error: Required environment variable(s) not set: {', '.join(missing_keys)}. Cannot run test.") else: try: # Test summarization print("\nTesting summarization...") long_text = """The Industrial Revolution, now also known as the First Industrial Revolution, was a period of global transition of the human economy towards more efficient and stable manufacturing processes that succeeded the Agricultural Revolution, starting from Great Britain, continental Europe and the United States, that occurred during the period from around 1760 to about 1820–1840. This transition included going from hand production methods to machines; new chemical manufacturing and iron production processes; the increasing use of water power and steam power; the development of machine tools; and the rise of the mechanized factory system. The Revolution also saw an unprecedented rise in the rate of population growth.""" summary = summarize_text(long_text, max_length=50) print(f"Summary:\n{summary}") # Test entity extraction print("\nTesting entity extraction...") entities = extract_entities(long_text, entity_types=["EVENT", "GPE", "DATE"]) print(f"Extracted Entities:\n{entities}") # Test text splitting print("\nTesting text splitting...") chunks = split_text_into_chunks(long_text * 3, chunk_size=150, chunk_overlap=30) # Make text longer print(f"Split into {len(chunks)} chunks. First chunk:\n{chunks[0]}") # Initialize the agent (optional) # test_agent = initialize_text_analyzer_agent() # print("\nText Analyzer Agent initialized successfully for testing.") except Exception as e: print(f"Error during testing: {e}")