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 # Attempt to import Whisper try: import whisper WHISPER_AVAILABLE = True except ImportError: logging.warning("openai-whisper not installed. Audio transcription tool will be unavailable.") WHISPER_AVAILABLE = False # 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 # --- Helper function to load Whisper model --- def _load_whisper_model(model_size: str = "small") -> Optional[object]: """Loads the Whisper model instance, lazy loading.""" global _whisper_model if not WHISPER_AVAILABLE: logger.error("Whisper library not available, cannot load model.") return None if _whisper_model is None: try: logger.info(f"Loading Whisper model: {model_size}...") # Allow model size selection via env var, default to "base" selected_model_size = os.getenv("WHISPER_MODEL_SIZE", model_size) print(f"Available Whisper models: {whisper.available_models()}") _whisper_model = whisper.load_model(selected_model_size) logger.info(f"Whisper model {selected_model_size} loaded successfully.") except Exception as e: logger.error(f"Failed to load Whisper model {selected_model_size}: {e}", exc_info=True) _whisper_model = None # Ensure it remains None on failure return _whisper_model # --- 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=summarizer_llm_model) 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)] def transcribe_audio(audio_file_path: str, language: Optional[str] = None) -> str: """Transcribes an audio file using the OpenAI Whisper model. Args: audio_file_path (str): The path to the audio file (e.g., mp3, wav, m4a). language (Optional[str]): The language code (e.g., "en", "es") or full name ("English", "Spanish"). If None, Whisper will detect the language. Returns: str: The transcribed text or an error message. """ logger.info(f"Attempting to transcribe audio file: {audio_file_path}, Language: {language}") # Check if Whisper is available if not WHISPER_AVAILABLE: return "Error: openai-whisper library is required but not installed." # Check if file exists if not os.path.exists(audio_file_path): logger.error(f"Audio file not found: {audio_file_path}") return f"Error: Audio file not found at {audio_file_path}" # Load the Whisper model (lazy loading) model = _load_whisper_model() # Uses default size "base" or WHISPER_MODEL_SIZE env var if model is None: return "Error: Failed to load Whisper model." try: # Perform transcription # The transcribe function handles various audio formats via ffmpeg result = model.transcribe(audio_file_path, language=language) transcribed_text = result["text"] detected_language = result.get("language", "unknown") # Get detected language if available logger.info(f"Audio transcription successful. Detected language: {detected_language}. Text length: {len(transcribed_text)}") return transcribed_text except Exception as e: # Check if it might be an ffmpeg issue if "ffmpeg" in str(e).lower(): logger.error(f"Error during transcription, possibly ffmpeg issue: {e}", exc_info=True) # Check if ffmpeg is installed using shell command try: subprocess.run(["ffmpeg", "-version"], check=True, capture_output=True) # If ffmpeg is installed, the error is likely something else return f"Error during transcription (ffmpeg seems installed): {e}" except (FileNotFoundError, subprocess.CalledProcessError): logger.error("ffmpeg command not found or failed. Please ensure ffmpeg is installed and in PATH.") return "Error: ffmpeg not found or not working. Please install ffmpeg." else: logger.error(f"Unexpected error during transcription: {e}", exc_info=True) return f"Error during transcription: {e}" # --- 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." ), ) # Conditionally create transcribe_audio_tool transcribe_audio_tool = None if WHISPER_AVAILABLE: transcribe_audio_tool = FunctionTool.from_defaults( fn=transcribe_audio, name="transcribe_audio_file", description=( "Transcribes speech from an audio file (e.g., mp3, wav, m4a) into text using Whisper. " "Input: audio_file_path (str), Optional: language (str - e.g., \"en\", \"Spanish\"). " "Output: transcribed text (str) or error message." ), ) logger.info("Audio transcription tool created.") else: logger.warning("Audio transcription tool disabled because openai-whisper is not installed.") # --- 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", "models/gemini-1.5-pro") 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=agent_llm_model) 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] if transcribe_audio_tool: tools.append(transcribe_audio_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`)." ) if transcribe_audio_tool: agent_description += " Can also transcribe audio files to text (`transcribe_audio_file`)." 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]}") # Test audio transcription (if available) if WHISPER_AVAILABLE: print("\nTesting audio transcription...") # Create a dummy audio file for testing (requires ffmpeg) dummy_file = "dummy_audio.mp3" try: # Generate a 1-second silent MP3 using ffmpeg subprocess.run(["ffmpeg", "-f", "lavfi", "-i", "anullsrc=r=44100:cl=mono", "-t", "1", "-q:a", "9", "-y", dummy_file], check=True, capture_output=True) print(f"Created dummy audio file: {dummy_file}") transcript = transcribe_audio(dummy_file) print(f"Transcription Result: '{transcript}' (Expected: empty or silence markers)") os.remove(dummy_file) # Clean up dummy file except Exception as ffmpeg_err: print(f"Could not create/test dummy audio file (ffmpeg required): {ffmpeg_err}") else: print("\nSkipping audio transcription test as openai-whisper is not available.") # 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}")