├── README.md ├── nexus ├── __init__.py ├── agents │ ├── __init__.py │ ├── evaluator_agent.py │ ├── code_agent.py │ ├── data_agent.py │ ├── audio_agent.py │ └── image_agent.py ├── security.py ├── vector_store.py ├── monitoring.py ├── evaluator.py ├── caching.py ├── workers.py ├── chains.py └── orchestrator.py ├── examples ├── sample_data.csv ├── tts_example.py ├── audio_example.py ├── sample_code.py ├── basic_usage.py ├── ocr_example.py ├── security_example.py ├── parallel_eval_example.py ├── vector_store_example.py ├── evaluator_example.py ├── workers_example.py ├── chain_example.py ├── monitoring_example.py ├── image_example.py ├── caching_example.py └── full_demo.py ├── .gitignore ├── requirements.txt ├── pyproject.toml ├── .cursorrules └── action_plan.md /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/benjaminegger/nexus-agents/HEAD/README.md -------------------------------------------------------------------------------- /nexus/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Nexus Agents - A Multi-Modal Agentic Platform 3 | """ 4 | 5 | __version__ = "0.1.0" -------------------------------------------------------------------------------- /nexus/agents/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Nexus Agents module. 3 | """ 4 | from .audio_agent import AudioAgent, AudioConfig 5 | from .image_agent import ImageAgent, ImageConfig, GenerationConfig 6 | 7 | __all__ = [ 8 | 'AudioAgent', 9 | 'AudioConfig', 10 | 'ImageAgent', 11 | 'ImageConfig', 12 | 'GenerationConfig', 13 | ] -------------------------------------------------------------------------------- /examples/sample_data.csv: -------------------------------------------------------------------------------- 1 | id,name,age,salary,department 2 | 1,John Doe,35,75000,Engineering 3 | 2,Jane Smith,28,65000,Marketing 4 | 3,Bob Johnson,42,85000,Engineering 5 | 4,Alice Brown,31,70000,Sales 6 | 5,Charlie Wilson,39,80000,Engineering 7 | 6,Diana Lee,33,72000,Marketing 8 | 7,Edward Chen,45,90000,Sales 9 | 8,Fiona Wright,29,68000,Marketing 10 | 9,George Kim,37,78000,Engineering 11 | 10,Helen Park,34,71000,Sales -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Python 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | *.so 6 | .Python 7 | build/ 8 | develop-eggs/ 9 | dist/ 10 | downloads/ 11 | eggs/ 12 | .eggs/ 13 | lib/ 14 | lib64/ 15 | parts/ 16 | sdist/ 17 | var/ 18 | wheels/ 19 | *.egg-info/ 20 | .installed.cfg 21 | *.egg 22 | 23 | # Virtual Environment 24 | .env 25 | .venv 26 | env/ 27 | venv/ 28 | ENV/ 29 | 30 | # IDE 31 | .idea/ 32 | .vscode/ 33 | *.swp 34 | *.swo 35 | 36 | # Project specific 37 | logs/ 38 | data/ 39 | *.log 40 | .uv/ 41 | 42 | # OS specific 43 | .DS_Store 44 | Thumbs.db -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # Core dependencies 2 | cryptography>=41.0.0 3 | backoff>=2.2.0 4 | pydantic>=2.6.3 5 | loguru>=0.7.2 6 | python-dotenv>=1.0.1 7 | typing-extensions>=4.9.0 8 | 9 | # Data processing 10 | pandas>=2.2.1 11 | numpy>=1.26.4 12 | openpyxl>=3.1.2 # For Excel support 13 | 14 | # Async support 15 | aiohttp>=3.9.3 16 | asyncio>=3.4.3 17 | 18 | # Monitoring & metrics 19 | prometheus-client>=0.20.0 20 | psutil>=5.9.8 # For system metrics 21 | 22 | # Web UI & API 23 | gradio>=4.19.2 24 | fastapi>=0.115.6 25 | uvicorn>=0.27.1 26 | starlette>=0.41.3 27 | python-multipart>=0.0.20 # For file uploads 28 | 29 | # Development tools 30 | ruff>=0.9.1 # For code linting 31 | huggingface-hub>=0.27.1 # For model integration 32 | safehttpx>=0.1.6 # For safe HTTP requests 33 | semantic-version>=2.10.0 # For version parsing 34 | tomlkit>=0.13.2 # For TOML handling 35 | ffmpy>=0.5.0 # For media processing -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "nexus-agents" 3 | version = "0.1.0" 4 | description = "Multi-modal agentic platform with LLM routing capabilities" 5 | dependencies = [ 6 | "phidata>=2.1.0", 7 | "openai>=1.35.8", 8 | "anthropic>=0.30.1", 9 | "aisuite[all]>=0.1.7", 10 | "faster-whisper>=0.10.0", 11 | "soundfile==0.12.1", 12 | "numpy>=1.24.3", 13 | "torch>=2.0.0", 14 | "torchaudio>=2.0.0", 15 | "TTS>=0.22.0", 16 | "noisereduce>=3.0.0", 17 | "easyocr>=1.7.1", 18 | "pillow>=10.0.0", 19 | "opencv-python>=4.8.0", 20 | "diffusers>=0.25.0", 21 | "transformers>=4.36.0", 22 | "accelerate>=0.27.0", 23 | "safetensors>=0.4.1", 24 | "compel>=2.0.2", 25 | "controlnet-aux>=0.0.7", 26 | "python-dotenv==1.0.0", 27 | "loguru==0.7.2", 28 | "rich==13.7.0", 29 | "pydantic>=2.3.0", 30 | ] 31 | requires-python = ">=3.9" 32 | readme = "README.md" 33 | license = { text = "MIT" } 34 | 35 | [project.optional-dependencies] 36 | dev = [ 37 | "ruff==0.1.8", 38 | "black==23.12.1", 39 | "pytest==7.4.3", 40 | ] 41 | 42 | [build-system] 43 | requires = ["hatchling"] 44 | build-backend = "hatchling.build" 45 | 46 | [tool.hatch.build.targets.wheel] 47 | packages = ["mnemosyne"] 48 | 49 | [tool.ruff] 50 | line-length = 88 51 | target-version = "py39" 52 | 53 | [tool.black] 54 | line-length = 88 55 | target-version = ["py39"] -------------------------------------------------------------------------------- /examples/tts_example.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example demonstrating AudioAgent's text-to-speech capabilities. 3 | """ 4 | from nexus.agents.audio_agent import AudioAgent, AudioConfig 5 | from loguru import logger 6 | import sys 7 | 8 | 9 | def main(): 10 | # Initialize audio agent with TTS configuration 11 | audio_agent = AudioAgent( 12 | config=AudioConfig( 13 | # TTS settings 14 | tts_model="tts_models/en/ljspeech/tacotron2-DDC", 15 | tts_language="en", 16 | # Audio processing settings 17 | noise_reduction=True, 18 | volume_normalize=True 19 | ) 20 | ) 21 | 22 | try: 23 | # List available voices and languages 24 | voices = audio_agent.list_available_voices() 25 | languages = audio_agent.list_available_languages() 26 | 27 | logger.info("Available voices: {}", voices) 28 | logger.info("Available languages: {}", languages) 29 | 30 | # Example text to synthesize 31 | text = "Hello! This is a test of the text-to-speech system." 32 | 33 | # Synthesize speech and save to file 34 | output_path = "test_output.wav" 35 | logger.info("Synthesizing speech to: {}", output_path) 36 | 37 | audio_path = audio_agent.synthesize( 38 | text=text, 39 | output_path=output_path 40 | ) 41 | 42 | logger.info("Speech synthesized successfully to: {}", audio_path) 43 | 44 | # Now let's try transcribing the synthesized audio 45 | logger.info("Transcribing the synthesized audio...") 46 | 47 | transcription = audio_agent.transcribe(audio_path) 48 | logger.info("Transcription: {}", transcription["text"]) 49 | 50 | except Exception as e: 51 | logger.exception("Error in TTS example") 52 | sys.exit(1) 53 | 54 | 55 | if __name__ == "__main__": 56 | main() -------------------------------------------------------------------------------- /examples/audio_example.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example demonstrating AudioAgent usage for speech-to-text conversion. 3 | """ 4 | from nexus.agents.audio_agent import AudioAgent, AudioConfig 5 | from nexus.orchestrator import Orchestrator, OrchestratorConfig, LLMProviderConfig 6 | from loguru import logger 7 | from pathlib import Path 8 | import sys 9 | 10 | 11 | def main(): 12 | # Initialize audio agent with base model 13 | audio_agent = AudioAgent( 14 | config=AudioConfig( 15 | model_type="base", # Use base model for faster processing 16 | device="cpu", # Use CPU for inference 17 | language="en" # Default to English 18 | ) 19 | ) 20 | 21 | # Initialize orchestrator for processing transcribed text 22 | orchestrator = Orchestrator( 23 | config=OrchestratorConfig( 24 | debug=True, 25 | primary_provider=LLMProviderConfig( 26 | provider="openai", 27 | model="gpt-4" 28 | ) 29 | ) 30 | ) 31 | 32 | try: 33 | # Path to your audio file 34 | audio_path = "path/to/your/audio.wav" # Replace with actual path 35 | 36 | if not Path(audio_path).exists(): 37 | logger.error("Audio file not found: {}", audio_path) 38 | sys.exit(1) 39 | 40 | # First, detect the language 41 | detected_lang = audio_agent.detect_language(audio_path) 42 | logger.info("Detected language: {}", detected_lang) 43 | 44 | # Transcribe the audio 45 | transcription = audio_agent.transcribe(audio_path) 46 | logger.info("Transcription: {}", transcription["text"]) 47 | 48 | # Process the transcribed text with the orchestrator 49 | prompt = f"Please analyze this transcribed text and provide a summary: {transcription['text']}" 50 | response = orchestrator.process_input(prompt) 51 | 52 | logger.info("Analysis: {}", response["response"]) 53 | 54 | except Exception as e: 55 | logger.exception("Error processing audio") 56 | sys.exit(1) 57 | 58 | 59 | if __name__ == "__main__": 60 | main() -------------------------------------------------------------------------------- /examples/sample_code.py: -------------------------------------------------------------------------------- 1 | """ 2 | Sample code for testing code analysis agent. 3 | """ 4 | from typing import List, Dict, Optional 5 | import os 6 | import sys 7 | from pathlib import Path 8 | 9 | 10 | def calculate_statistics(numbers: List[float]) -> Dict[str, float]: 11 | """ 12 | Calculate basic statistics for a list of numbers. 13 | 14 | Args: 15 | numbers: List of numbers to analyze 16 | 17 | Returns: 18 | Dictionary with statistics 19 | """ 20 | if not numbers: 21 | return { 22 | "mean": 0.0, 23 | "min": 0.0, 24 | "max": 0.0, 25 | "sum": 0.0 26 | } 27 | 28 | total = sum(numbers) 29 | count = len(numbers) 30 | 31 | return { 32 | "mean": total / count, 33 | "min": min(numbers), 34 | "max": max(numbers), 35 | "sum": total 36 | } 37 | 38 | 39 | class FileProcessor: 40 | """Process files with size validation.""" 41 | 42 | def __init__(self, max_size: int = 1024 * 1024): 43 | """Initialize processor with max file size.""" 44 | self.max_size = max_size 45 | 46 | def process_file(self, path: str) -> Optional[str]: 47 | """ 48 | Process a file if it's within size limit. 49 | 50 | Args: 51 | path: Path to file 52 | 53 | Returns: 54 | File contents or None if too large 55 | """ 56 | file_path = Path(path) 57 | 58 | if not file_path.exists(): 59 | raise FileNotFoundError(f"File not found: {path}") 60 | 61 | if file_path.stat().st_size > self.max_size: 62 | return None 63 | 64 | return file_path.read_text() 65 | 66 | 67 | def main(): 68 | """Main function.""" 69 | # Example usage 70 | numbers = [1.5, 2.7, 3.2, 4.8, 5.1] 71 | stats = calculate_statistics(numbers) 72 | print("Statistics:", stats) 73 | 74 | processor = FileProcessor(max_size=1024) # 1KB limit 75 | try: 76 | content = processor.process_file("test.txt") 77 | if content: 78 | print("File contents:", content) 79 | else: 80 | print("File too large") 81 | except FileNotFoundError as e: 82 | print("Error:", e) 83 | 84 | 85 | if __name__ == "__main__": 86 | main() -------------------------------------------------------------------------------- /examples/basic_usage.py: -------------------------------------------------------------------------------- 1 | """ 2 | Basic example demonstrating Orchestrator usage with multiple LLM providers. 3 | Shows enhanced logging and debugging features. 4 | """ 5 | from loguru import logger 6 | from nexus.orchestrator import Orchestrator, OrchestratorConfig, LLMProviderConfig 7 | import os 8 | from dotenv import load_dotenv 9 | import sys 10 | 11 | def main(): 12 | # Load environment variables from .env file 13 | load_dotenv() 14 | 15 | # Initialize the orchestrator with multiple providers and debug mode 16 | orchestrator = Orchestrator( 17 | config=OrchestratorConfig( 18 | debug=True, # Enable detailed logging 19 | history_length=5, 20 | # Configure primary provider 21 | primary_provider=LLMProviderConfig( 22 | provider="openai", 23 | model="gpt-4" 24 | ), 25 | # Configure fallback providers 26 | fallback_providers=[ 27 | LLMProviderConfig( 28 | provider="anthropic", 29 | model="claude-3-sonnet-20240229" 30 | ), 31 | LLMProviderConfig( 32 | provider="google", 33 | model="gemini-pro" 34 | ) 35 | ] 36 | ) 37 | ) 38 | 39 | logger.info("Starting conversation with enhanced logging...") 40 | 41 | try: 42 | # First question - will try primary provider first 43 | with logger.contextualize(question_number=1): 44 | question1 = "What are the three laws of robotics?" 45 | logger.info("Asking question: {}", question1) 46 | response1 = orchestrator.process_input(question1) 47 | logger.info("Received response from {}: {}", 48 | response1["provider_used"], 49 | response1["response"]) 50 | 51 | # Follow-up question (demonstrates memory and potentially fallback providers) 52 | with logger.contextualize(question_number=2): 53 | question2 = "Who created these laws?" 54 | logger.info("Asking follow-up: {}", question2) 55 | response2 = orchestrator.process_input(question2) 56 | logger.info("Received response from {}: {}", 57 | response2["provider_used"], 58 | response2["response"]) 59 | 60 | # Log conversation summary 61 | logger.info("Conversation completed successfully") 62 | logger.debug("Final history length: {}", response2["history_length"]) 63 | 64 | except Exception as e: 65 | logger.exception("Error during conversation") 66 | sys.exit(1) 67 | 68 | if __name__ == "__main__": 69 | main() -------------------------------------------------------------------------------- /examples/ocr_example.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example demonstrating ImageAgent's OCR capabilities. 3 | """ 4 | from nexus.agents.image_agent import ImageAgent, ImageConfig 5 | from nexus.orchestrator import Orchestrator, OrchestratorConfig, LLMProviderConfig 6 | from loguru import logger 7 | from pathlib import Path 8 | import sys 9 | 10 | 11 | def main(): 12 | # Initialize image agent with OCR configuration 13 | image_agent = ImageAgent( 14 | config=ImageConfig( 15 | languages=["en"], # English OCR 16 | device="cpu", # Use CPU for inference 17 | # Image processing settings 18 | preprocessing=True, 19 | contrast_enhance=True, 20 | denoise=True, 21 | min_confidence=0.5 22 | ) 23 | ) 24 | 25 | # Initialize orchestrator for processing extracted text 26 | orchestrator = Orchestrator( 27 | config=OrchestratorConfig( 28 | debug=True, 29 | primary_provider=LLMProviderConfig( 30 | provider="openai", 31 | model="gpt-4" 32 | ) 33 | ) 34 | ) 35 | 36 | try: 37 | # Path to your image file 38 | image_path = "path/to/your/image.png" # Replace with actual path 39 | 40 | if not Path(image_path).exists(): 41 | logger.error("Image file not found: {}", image_path) 42 | sys.exit(1) 43 | 44 | # List supported languages 45 | supported_langs = image_agent.supported_languages() 46 | logger.info("Supported OCR languages: {}", supported_langs) 47 | 48 | # First, detect text regions 49 | logger.info("Detecting text regions...") 50 | regions = image_agent.detect_text_regions(image_path) 51 | logger.info("Found {} text regions", len(regions)) 52 | 53 | # Extract text from the image 54 | logger.info("Extracting text...") 55 | result = image_agent.extract_text(image_path) 56 | 57 | logger.info("Extracted text: {}", result["text"]) 58 | if result["blocks"]: 59 | logger.info("Text blocks found: {}", len(result["blocks"])) 60 | for i, block in enumerate(result["blocks"], 1): 61 | logger.info("Block {}: '{}' (confidence: {:.2f})", 62 | i, block["text"], block["confidence"]) 63 | 64 | # Process the extracted text with the orchestrator 65 | prompt = f"Please analyze this text extracted from an image and provide a summary: {result['text']}" 66 | response = orchestrator.process_input(prompt) 67 | 68 | logger.info("Analysis: {}", response["response"]) 69 | 70 | except Exception as e: 71 | logger.exception("Error in OCR example") 72 | sys.exit(1) 73 | 74 | 75 | if __name__ == "__main__": 76 | main() -------------------------------------------------------------------------------- /examples/security_example.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example demonstrating security features including encryption and retry mechanisms. 3 | """ 4 | from nexus.security import Security, SecurityConfig 5 | from nexus.orchestrator import ( 6 | Orchestrator, 7 | OrchestratorConfig, 8 | LLMProviderConfig 9 | ) 10 | from loguru import logger 11 | import asyncio 12 | import sys 13 | from datetime import datetime 14 | 15 | 16 | class SimulatedFailure(Exception): 17 | """Simulated failure for testing retries.""" 18 | pass 19 | 20 | 21 | class ExampleAgent: 22 | """Example agent to demonstrate retry mechanism.""" 23 | 24 | def __init__(self, security: Security): 25 | self.security = security 26 | self.fail_count = 0 27 | 28 | @Security.with_retries(max_tries=3, initial_wait=1.0, max_wait=5.0) 29 | async def flaky_operation(self) -> str: 30 | """Simulated operation that sometimes fails.""" 31 | self.fail_count += 1 32 | 33 | if self.fail_count <= 2: # Fail first two attempts 34 | logger.warning("Operation failed, attempt {}", self.fail_count) 35 | raise SimulatedFailure("Simulated failure") 36 | 37 | return "Operation succeeded on attempt 3!" 38 | 39 | 40 | async def main(): 41 | # Initialize security with custom config 42 | security = Security( 43 | config=SecurityConfig( 44 | max_retries=3, 45 | initial_wait=1.0, 46 | max_wait=30.0 47 | ) 48 | ) 49 | 50 | try: 51 | # Example 1: Data Encryption 52 | logger.info("\nTesting data encryption...") 53 | 54 | # Sample data to encrypt 55 | sensitive_data = { 56 | "user_id": "12345", 57 | "timestamp": datetime.now().isoformat(), 58 | "ocr_results": "Confidential document text...", 59 | "audio_transcript": "Sensitive conversation content..." 60 | } 61 | 62 | # Encrypt data 63 | encrypted = security.encrypt_data(sensitive_data) 64 | logger.info("Data encrypted: {} bytes", len(encrypted)) 65 | 66 | # Decrypt data 67 | decrypted = security.decrypt_data(encrypted) 68 | logger.info("Data decrypted successfully") 69 | logger.debug("Decrypted content: {}", decrypted) 70 | 71 | # Verify data integrity 72 | assert decrypted == sensitive_data 73 | logger.info("Data integrity verified") 74 | 75 | # Example 2: Retry Mechanism 76 | logger.info("\nTesting retry mechanism...") 77 | 78 | # Create example agent 79 | agent = ExampleAgent(security) 80 | 81 | # Try flaky operation 82 | result = await agent.flaky_operation() 83 | logger.info("Final result: {}", result) 84 | 85 | except Exception as e: 86 | logger.exception("Error during example") 87 | sys.exit(1) 88 | 89 | 90 | if __name__ == "__main__": 91 | # Run the async main function 92 | asyncio.run(main()) -------------------------------------------------------------------------------- /.cursorrules: -------------------------------------------------------------------------------- 1 | # .cursorrules 2 | 3 | 4 | Always start with "YOOO mi amigo!!" 5 | 6 | Important rules you HAVE TO FOLLOW 7 | -You will create a documentation file and write all tasks you do in that file. 8 | When I ask you to write a new task, check the documentation first and remember our project requirements and steps. 9 | -Always add debug logs and comments in the code for easier debugging & readability. 10 | -Everytime you are asked to do something, you MUST ask for clarification first. 11 | -Everytime you choose to apply a rule(s), explicitly state which rule(s) in the output. You can abbreviate the rule description 12 | to a single word or phrase. 13 | -After implementing any new feature, update action_plan.md with: 14 | * Mark [X] for completed items 15 | * Add detailed explanation of the feature 16 | * Document how to use it 17 | * Explain its purpose and benefits 18 | * Describe relationships with other features 19 | * Include example usage if applicable 20 | 21 | # Multi-Modal Platform Development Rules 22 | 23 | ## Code Quality Rules 24 | -Follow Python best practices (3.9+) 25 | -Maintain modular code structure with clear separation of concerns 26 | -Add comprehensive docstrings and type hints 27 | -Implement proper error handling and logging 28 | -Use loguru for structured logging with appropriate debug levels 29 | 30 | ## Architecture Rules 31 | -Keep each agent (Text, Audio, Image, Video) independent and modular 32 | -Ensure the Orchestrator remains the central control point 33 | -Implement proper memory management using Vector DB 34 | -Follow the defined workflow patterns (Prompt Chaining, Routing, Parallelization, etc.) 35 | 36 | ## Development Process Rules 37 | -Break tasks into small, independently testable units 38 | -Add tests for each new feature or modification 39 | -Update documentation with each significant change 40 | -Follow phased development approach as outlined in action plan 41 | -Keep PRs small and focused 42 | 43 | ## Security Rules 44 | -Never expose API keys or sensitive data in logs 45 | -Implement proper encryption for sensitive data 46 | -Add appropriate authentication mechanisms 47 | -Follow security best practices for external API calls 48 | 49 | ## Performance Rules 50 | -Optimize for response time (<5s for text-based tasks) 51 | -Implement caching where appropriate 52 | -Use parallelization for independent tasks 53 | -Monitor and log performance metrics 54 | 55 | ## Multi-Modal Rules 56 | -Handle text, audio, image, and video inputs appropriately 57 | -Implement proper format validation 58 | -Use appropriate tools for each modality (Whisper, Tesseract, etc.) 59 | -Ensure seamless integration between different modalities 60 | 61 | ## Memory & Context Rules 62 | -Maintain conversation history appropriately 63 | -Use Vector DB for long-term memory storage 64 | -Implement proper context management 65 | -Handle context windows efficiently 66 | 67 | ## Debugging Rules 68 | -Use structured logging with loguru 69 | -Implement comprehensive error messages 70 | -Add debug mode with detailed logging 71 | -Maintain clean and readable error traces 72 | 73 | 74 | -------------------------------------------------------------------------------- /examples/parallel_eval_example.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example demonstrating parallel task execution with evaluation. 3 | Shows how to run OCR and content moderation in parallel, then evaluate results. 4 | """ 5 | from nexus.orchestrator import Orchestrator, OrchestratorConfig, LLMProviderConfig 6 | from nexus.agents.evaluator_agent import EvaluatorAgent, EvaluatorConfig 7 | from nexus.agents.image_agent import ImageAgent 8 | from loguru import logger 9 | import asyncio 10 | from pathlib import Path 11 | import sys 12 | 13 | 14 | async def main(): 15 | # Initialize agents 16 | orchestrator = Orchestrator( 17 | config=OrchestratorConfig( 18 | debug=True, 19 | primary_provider=LLMProviderConfig( 20 | provider="openai", 21 | model="gpt-4" 22 | ) 23 | ) 24 | ) 25 | 26 | evaluator = EvaluatorAgent( 27 | config=EvaluatorConfig( 28 | confidence_threshold=0.7, 29 | check_types=["factual_coherence"] 30 | ) 31 | ) 32 | 33 | image_agent = ImageAgent() # Using default config 34 | 35 | try: 36 | # Example: Process image and evaluate text in parallel 37 | image_path = "path/to/your/image.jpg" # Replace with actual path 38 | 39 | if not Path(image_path).exists(): 40 | logger.error("Image file not found: {}", image_path) 41 | sys.exit(1) 42 | 43 | logger.info("Starting parallel processing") 44 | 45 | # Define tasks to run in parallel 46 | tasks = [ 47 | # Task 1: OCR the image 48 | image_agent.extract_text(image_path), 49 | 50 | # Task 2: Get moderation check from LLM 51 | orchestrator.process_input_async( 52 | "Please analyze this image for any concerning content." 53 | ) 54 | ] 55 | 56 | # Run tasks in parallel with timeout 57 | results = await orchestrator.run_parallel_tasks(tasks, timeout=30.0) 58 | 59 | # Unpack results 60 | ocr_text, moderation_result = results 61 | 62 | logger.info("OCR Text: {}", ocr_text) 63 | logger.info("Moderation Result: {}", moderation_result) 64 | 65 | # Evaluate the results 66 | evaluation = await evaluator.evaluate_output( 67 | output=moderation_result["response"], 68 | context={"ocr_text": ocr_text} 69 | ) 70 | 71 | logger.info("Evaluation Results:") 72 | logger.info("Passed: {}", evaluation["passed"]) 73 | logger.info("Confidence: {:.2f}", evaluation["confidence"]) 74 | 75 | if not evaluation["passed"]: 76 | logger.warning("Evaluation failed. Suggestions:") 77 | for suggestion in evaluation["suggestions"]: 78 | logger.warning("- {}", suggestion) 79 | 80 | except asyncio.TimeoutError: 81 | logger.error("Processing timed out") 82 | sys.exit(1) 83 | except Exception as e: 84 | logger.exception("Error during processing") 85 | sys.exit(1) 86 | 87 | 88 | if __name__ == "__main__": 89 | # Run the async main function 90 | asyncio.run(main()) -------------------------------------------------------------------------------- /examples/vector_store_example.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example demonstrating vector store integration for persistent memory. 3 | Shows conversation history storage and semantic search capabilities. 4 | """ 5 | from nexus.orchestrator import ( 6 | Orchestrator, 7 | OrchestratorConfig, 8 | LLMProviderConfig 9 | ) 10 | from nexus.vector_store import VectorStoreConfig 11 | from loguru import logger 12 | import asyncio 13 | import sys 14 | 15 | 16 | async def main(): 17 | # Initialize orchestrator with vector store 18 | orchestrator = Orchestrator( 19 | config=OrchestratorConfig( 20 | debug=True, 21 | primary_provider=LLMProviderConfig( 22 | provider="openai", 23 | model="gpt-4" 24 | ), 25 | vector_store=VectorStoreConfig( 26 | db_path="data/example_store", 27 | collection_name="example_history" 28 | ) 29 | ) 30 | ) 31 | 32 | try: 33 | # Example conversation 34 | logger.info("Starting conversation...") 35 | 36 | # First message about Python 37 | response1 = await orchestrator.process_input_async( 38 | "What are the key features of Python?" 39 | ) 40 | logger.info("Response 1: {}", response1["response"]) 41 | 42 | # Second message about a different topic 43 | response2 = await orchestrator.process_input_async( 44 | "Tell me about machine learning." 45 | ) 46 | logger.info("Response 2: {}", response2["response"]) 47 | 48 | # Third message about Python again 49 | response3 = await orchestrator.process_input_async( 50 | "How does Python handle memory management?" 51 | ) 52 | logger.info("Response 3: {}", response3["response"]) 53 | 54 | # Demonstrate semantic search 55 | logger.info("\nSearching for messages about Python...") 56 | python_messages = await orchestrator.search_similar_messages( 57 | query="Python programming language features", 58 | limit=2 59 | ) 60 | 61 | logger.info("Found {} relevant messages:", len(python_messages)) 62 | for msg in python_messages: 63 | logger.info("- Role: {}, Content: {}", msg["role"], msg["content"]) 64 | 65 | # Get recent history 66 | logger.info("\nGetting recent history...") 67 | history = await orchestrator.get_conversation_history(limit=5) 68 | 69 | logger.info("Recent conversation:") 70 | for msg in history: 71 | logger.info("- {}: {}", msg["role"], msg["content"]) 72 | 73 | # Filter by role 74 | logger.info("\nGetting only assistant responses...") 75 | assistant_msgs = await orchestrator.get_conversation_history( 76 | role_filter="assistant" 77 | ) 78 | 79 | logger.info("Assistant messages:") 80 | for msg in assistant_msgs: 81 | logger.info("- {}", msg["content"]) 82 | 83 | # Clear history 84 | logger.info("\nClearing conversation history...") 85 | await orchestrator.clear_history() 86 | 87 | # Verify it's cleared 88 | empty_history = await orchestrator.get_conversation_history() 89 | logger.info("History after clearing: {} messages", len(empty_history)) 90 | 91 | except Exception as e: 92 | logger.exception("Error during example") 93 | sys.exit(1) 94 | 95 | 96 | if __name__ == "__main__": 97 | # Run the async main function 98 | asyncio.run(main()) -------------------------------------------------------------------------------- /examples/evaluator_example.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example demonstrating the enhanced evaluator features. 3 | """ 4 | import asyncio 5 | from loguru import logger 6 | from nexus.evaluator import ( 7 | ContentEvaluator, 8 | EvaluationConfig, 9 | EvaluationCriteria 10 | ) 11 | 12 | 13 | async def main(): 14 | # Create evaluator config 15 | config = EvaluationConfig( 16 | criteria=[ 17 | EvaluationCriteria.FACTUAL_ACCURACY, 18 | EvaluationCriteria.STYLE_CONSISTENCY, 19 | EvaluationCriteria.POLICY_COMPLIANCE, 20 | EvaluationCriteria.RELEVANCE 21 | ], 22 | weights={ 23 | EvaluationCriteria.FACTUAL_ACCURACY: 0.4, 24 | EvaluationCriteria.STYLE_CONSISTENCY: 0.2, 25 | EvaluationCriteria.POLICY_COMPLIANCE: 0.2, 26 | EvaluationCriteria.RELEVANCE: 0.2 27 | }, 28 | thresholds={ 29 | EvaluationCriteria.FACTUAL_ACCURACY: 0.8, 30 | EvaluationCriteria.STYLE_CONSISTENCY: 0.7, 31 | EvaluationCriteria.POLICY_COMPLIANCE: 0.9, 32 | EvaluationCriteria.RELEVANCE: 0.7 33 | }, 34 | refinement_threshold=0.75 35 | ) 36 | 37 | # Initialize evaluator 38 | evaluator = ContentEvaluator(config) 39 | 40 | # Example 1: Single Evaluation 41 | logger.info("Running single evaluation example...") 42 | 43 | content = """ 44 | The Earth orbits the Sun at an average distance of 93 million miles. 45 | This journey takes approximately 365.25 days to complete. 46 | The Earth's atmosphere is composed primarily of nitrogen and oxygen. 47 | """ 48 | 49 | result = await evaluator.evaluate(content) 50 | 51 | logger.info("Evaluation scores:") 52 | for score in result.scores: 53 | logger.info( 54 | "{}: {:.2f} - {}", 55 | score.criterion.value, 56 | score.score, 57 | score.feedback 58 | ) 59 | logger.info("Overall score: {:.2f}", result.overall_score) 60 | 61 | if result.needs_refinement: 62 | logger.info("Content needs refinement") 63 | logger.info("Refinement prompt: {}", result.refinement_prompt) 64 | 65 | # Example 2: Evaluation with Context 66 | logger.info("\nRunning evaluation with context example...") 67 | 68 | content_with_context = """ 69 | Our new product features advanced AI capabilities. 70 | It can process natural language and generate responses. 71 | The system is built on cutting-edge technology. 72 | """ 73 | 74 | context = { 75 | "target_audience": "technical professionals", 76 | "style_guide": { 77 | "tone": "professional", 78 | "formality": "high" 79 | }, 80 | "domain": "artificial intelligence" 81 | } 82 | 83 | result = await evaluator.evaluate(content_with_context, context) 84 | 85 | logger.info("Evaluation with context scores:") 86 | for score in result.scores: 87 | logger.info( 88 | "{}: {:.2f} - {}", 89 | score.criterion.value, 90 | score.score, 91 | score.feedback 92 | ) 93 | 94 | # Example 3: Automatic Refinement 95 | logger.info("\nRunning automatic refinement example...") 96 | 97 | content_to_refine = """ 98 | AI is really good at doing stuff. 99 | It helps people work better and faster. 100 | Everyone should use AI because it's amazing. 101 | """ 102 | 103 | refined_content, refinement_results = await evaluator.evaluate_and_refine( 104 | content_to_refine, 105 | max_iterations=3 106 | ) 107 | 108 | logger.info("Refinement history:") 109 | for i, result in enumerate(refinement_results): 110 | logger.info( 111 | "Iteration {}: Score {:.2f}", 112 | i + 1, 113 | result.overall_score 114 | ) 115 | 116 | logger.info("Final content:\n{}", refined_content) 117 | 118 | 119 | if __name__ == "__main__": 120 | asyncio.run(main()) -------------------------------------------------------------------------------- /examples/workers_example.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example demonstrating the Orchestrator-Workers pattern. 3 | """ 4 | import asyncio 5 | from loguru import logger 6 | from nexus.workers import ( 7 | WorkerOrchestrator, 8 | WorkerType, 9 | WorkerStatus 10 | ) 11 | 12 | 13 | async def text_processor(content: dict) -> str: 14 | """Example text processing function.""" 15 | # Simulate processing time 16 | await asyncio.sleep(0.5) 17 | text = content.get("text", "") 18 | return f"Processed text (len={len(text)}): {text[:50]}..." 19 | 20 | 21 | async def image_processor(content: dict) -> dict: 22 | """Example image processing function.""" 23 | # Simulate processing time 24 | await asyncio.sleep(1.0) 25 | image_path = content.get("path", "") 26 | return { 27 | "path": image_path, 28 | "dimensions": "800x600", 29 | "format": "JPEG" 30 | } 31 | 32 | 33 | async def main(): 34 | # Initialize orchestrator 35 | orchestrator = WorkerOrchestrator() 36 | 37 | # Example 1: Parallel Text Processing 38 | logger.info("Running parallel text processing example...") 39 | 40 | # Create text processing tasks 41 | text_tasks = [ 42 | {"text": f"Sample text {i} for processing"} 43 | for i in range(5) 44 | ] 45 | 46 | # Process in parallel with 3 workers 47 | text_results = await orchestrator.process_parallel( 48 | tasks=text_tasks, 49 | worker_type=WorkerType.TEXT, 50 | processor=text_processor, 51 | num_workers=3 52 | ) 53 | 54 | logger.info("Text processing results: {}", text_results) 55 | 56 | # Example 2: Sequential Image Processing 57 | logger.info("Running sequential image processing example...") 58 | 59 | # Create image processing tasks 60 | image_tasks = [ 61 | {"path": f"image_{i}.jpg"} 62 | for i in range(3) 63 | ] 64 | 65 | # Process sequentially 66 | image_results = await orchestrator.process_sequential( 67 | tasks=image_tasks, 68 | worker_type=WorkerType.IMAGE, 69 | processor=image_processor 70 | ) 71 | 72 | logger.info("Image processing results: {}", image_results) 73 | 74 | # Example 3: Mixed Workload with Dependencies 75 | logger.info("Running mixed workload example...") 76 | 77 | # Create a worker pool with different types 78 | orchestrator.create_worker_group( 79 | type=WorkerType.TEXT, 80 | count=2, 81 | capabilities=["summarize", "analyze"] 82 | ) 83 | orchestrator.create_worker_group( 84 | type=WorkerType.IMAGE, 85 | count=2, 86 | capabilities=["resize", "format"] 87 | ) 88 | 89 | # Add tasks with dependencies 90 | text_task_id = orchestrator.pool.add_task( 91 | type=WorkerType.TEXT, 92 | content={"text": "Text to process before image"} 93 | ) 94 | 95 | # Image task depends on text task 96 | image_task_id = orchestrator.pool.add_task( 97 | type=WorkerType.IMAGE, 98 | content={"path": "dependent_image.jpg"}, 99 | dependencies=[text_task_id] 100 | ) 101 | 102 | # Process text task 103 | text_result = await orchestrator.pool.process_task( 104 | worker_id=next( 105 | w_id for w_id, w in orchestrator.pool.workers.items() 106 | if w.type == WorkerType.TEXT and w.status == WorkerStatus.IDLE 107 | ), 108 | task_id=text_task_id, 109 | processor=text_processor 110 | ) 111 | logger.info("Text task result: {}", text_result) 112 | 113 | # Process image task (will wait for text task) 114 | image_result = await orchestrator.pool.process_task( 115 | worker_id=next( 116 | w_id for w_id, w in orchestrator.pool.workers.items() 117 | if w.type == WorkerType.IMAGE and w.status == WorkerStatus.IDLE 118 | ), 119 | task_id=image_task_id, 120 | processor=image_processor 121 | ) 122 | logger.info("Image task result: {}", image_result) 123 | 124 | 125 | if __name__ == "__main__": 126 | asyncio.run(main()) -------------------------------------------------------------------------------- /examples/chain_example.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example demonstrating the prompt chaining library features. 3 | """ 4 | import asyncio 5 | from loguru import logger 6 | from nexus.chains import ( 7 | PromptChain, 8 | ChainNodeType, 9 | create_linear_chain, 10 | create_branching_chain 11 | ) 12 | 13 | 14 | async def example_tool(text: str) -> dict: 15 | """Example tool that processes text.""" 16 | return { 17 | "processed_text": f"Processed: {text}", 18 | "length": len(text) 19 | } 20 | 21 | 22 | async def main(): 23 | # Example 1: Linear Chain 24 | logger.info("Creating linear chain example...") 25 | linear_chain = create_linear_chain( 26 | "text_processor", 27 | prompts=[ 28 | "Analyze this text: {input_text}", 29 | "Summarize the analysis: {processed_text}" 30 | ], 31 | tools=[{ 32 | "name": "text_processor", 33 | "params": {"text": "{input_text}"} 34 | }] 35 | ) 36 | 37 | # Execute linear chain 38 | linear_results = await linear_chain.execute( 39 | initial_context={"input_text": "Hello, World!"}, 40 | tools={"text_processor": example_tool} 41 | ) 42 | logger.info("Linear chain results: {}", linear_results) 43 | 44 | # Example 2: Branching Chain 45 | logger.info("Creating branching chain example...") 46 | branching_chain = create_branching_chain( 47 | "text_classifier", 48 | condition_prompt="Is this text a question: {input_text}", 49 | true_branch=[ 50 | "Answer the question: {input_text}", 51 | "Verify the answer's accuracy" 52 | ], 53 | false_branch=[ 54 | "Process the statement: {input_text}", 55 | "Generate related insights" 56 | ] 57 | ) 58 | 59 | # Add dynamic modification during execution 60 | question_node_id = branching_chain.add_node( 61 | ChainNodeType.PROMPT, 62 | "dynamic_question", 63 | "Additional question about: {input_text}" 64 | ) 65 | 66 | # We'll connect this dynamically during execution 67 | # based on some condition 68 | 69 | # Execute branching chain 70 | branching_results = await branching_chain.execute( 71 | initial_context={"input_text": "What is the meaning of life?"} 72 | ) 73 | logger.info("Branching chain results: {}", branching_results) 74 | 75 | # Example 3: Dynamic Chain Modification 76 | logger.info("Demonstrating dynamic chain modification...") 77 | dynamic_chain = PromptChain("dynamic_example") 78 | 79 | # Add initial nodes 80 | start_id = dynamic_chain.add_node( 81 | ChainNodeType.PROMPT, 82 | "start", 83 | "Initial analysis of: {input_text}" 84 | ) 85 | 86 | process_id = dynamic_chain.add_node( 87 | ChainNodeType.TOOL, 88 | "process", 89 | { 90 | "name": "text_processor", 91 | "params": {"text": "{input_text}"} 92 | } 93 | ) 94 | 95 | dynamic_chain.connect(start_id, process_id) 96 | 97 | # Execute first part 98 | context = { 99 | "input_text": "This is a dynamic chain example" 100 | } 101 | 102 | initial_results = await dynamic_chain.execute( 103 | initial_context=context, 104 | tools={"text_processor": example_tool} 105 | ) 106 | 107 | # Based on results, add new nodes 108 | if len(context["input_text"]) > 10: 109 | logger.info("Adding detailed analysis branch...") 110 | detail_id = dynamic_chain.add_node( 111 | ChainNodeType.PROMPT, 112 | "detailed_analysis", 113 | "Detailed analysis of: {processed_text}" 114 | ) 115 | dynamic_chain.connect(process_id, detail_id) 116 | 117 | # Execute the modified chain 118 | final_results = await dynamic_chain.execute( 119 | initial_context=context, 120 | tools={"text_processor": example_tool} 121 | ) 122 | logger.info("Final results after modification: {}", final_results) 123 | 124 | 125 | if __name__ == "__main__": 126 | asyncio.run(main()) -------------------------------------------------------------------------------- /examples/monitoring_example.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example demonstrating performance monitoring features. 3 | """ 4 | import asyncio 5 | from loguru import logger 6 | from nexus.monitoring import ( 7 | PerformanceMetrics, 8 | MonitoredComponent, 9 | MetricType 10 | ) 11 | import random 12 | import time 13 | 14 | 15 | class ExampleProcessor(MonitoredComponent): 16 | """Example component that processes tasks.""" 17 | 18 | def __init__(self, metrics: PerformanceMetrics): 19 | """Initialize processor.""" 20 | super().__init__(metrics) 21 | self.queue = [] 22 | 23 | async def process_task(self, task_id: int): 24 | """Process a single task with monitoring.""" 25 | async with self.track_operation( 26 | operation="process_task", 27 | agent_type="processor" 28 | ): 29 | # Simulate processing 30 | await asyncio.sleep(random.uniform(0.1, 0.5)) 31 | 32 | # Randomly fail some tasks 33 | if random.random() < 0.2: 34 | raise ValueError("Random task failure") 35 | 36 | return f"Processed task {task_id}" 37 | 38 | async def process_batch(self, batch: list): 39 | """Process a batch of tasks with worker tracking.""" 40 | async with self.track_worker(agent_type="processor"): 41 | results = [] 42 | for item in batch: 43 | try: 44 | result = await self.process_task(item) 45 | results.append(result) 46 | except Exception as e: 47 | self.record_error( 48 | error_type=type(e).__name__, 49 | agent_type="processor" 50 | ) 51 | return results 52 | 53 | def update_queue(self, items: list): 54 | """Update queue with monitoring.""" 55 | self.queue.extend(items) 56 | self.update_queue_size( 57 | size=len(self.queue), 58 | agent_type="processor" 59 | ) 60 | 61 | 62 | async def main(): 63 | # Initialize metrics 64 | metrics = PerformanceMetrics(port=8000) 65 | logger.info("Started metrics server on :8000") 66 | 67 | # Create processor 68 | processor = ExampleProcessor(metrics) 69 | 70 | # Example 1: Process individual tasks 71 | logger.info("Running individual task processing example...") 72 | 73 | for i in range(5): 74 | try: 75 | result = await processor.process_task(i) 76 | logger.info("Task {} result: {}", i, result) 77 | except Exception as e: 78 | logger.error("Task {} failed: {}", i, str(e)) 79 | 80 | # Example 2: Process batches with workers 81 | logger.info("\nRunning batch processing example...") 82 | 83 | batches = [ 84 | list(range(5)), 85 | list(range(5, 10)), 86 | list(range(10, 15)) 87 | ] 88 | 89 | for i, batch in enumerate(batches): 90 | processor.update_queue(batch) 91 | logger.info("Processing batch {} (queue size: {})", i + 1, len(batch)) 92 | 93 | results = await processor.process_batch(batch) 94 | logger.info("Batch {} results: {}", i + 1, results) 95 | 96 | # Clear processed items from queue 97 | processor.queue = processor.queue[len(batch):] 98 | processor.update_queue_size( 99 | size=len(processor.queue), 100 | agent_type="processor" 101 | ) 102 | 103 | # Example 3: Get metrics summary 104 | logger.info("\nGetting metrics summary...") 105 | 106 | summary = metrics.get_summary() 107 | logger.info("Metrics summary:\n{}", summary) 108 | 109 | # Keep server running to allow viewing metrics 110 | logger.info("\nMetrics server running on :8000") 111 | logger.info("Visit http://localhost:8000/metrics to view Prometheus metrics") 112 | logger.info("Press Ctrl+C to exit") 113 | 114 | try: 115 | while True: 116 | await asyncio.sleep(1) 117 | except KeyboardInterrupt: 118 | logger.info("Shutting down...") 119 | 120 | 121 | if __name__ == "__main__": 122 | asyncio.run(main()) -------------------------------------------------------------------------------- /examples/image_example.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example demonstrating the usage of the ImageAgent for OCR and image generation. 3 | """ 4 | import os 5 | from pathlib import Path 6 | from dotenv import load_dotenv 7 | from nexus.agents import ImageAgent, ImageConfig, GenerationConfig 8 | from nexus.agents.image_agent import GENERATION_AVAILABLE 9 | from loguru import logger 10 | import torch 11 | 12 | # Load environment variables 13 | load_dotenv() 14 | 15 | def main(): 16 | # Configure the image agent 17 | config = ImageConfig( 18 | languages=["en"], 19 | device="cuda" if torch.cuda.is_available() else "cpu", 20 | return_bboxes=True, # Get bounding boxes for visualization 21 | return_confidence=True, # Get confidence scores 22 | min_confidence=0.5, # Minimum confidence threshold 23 | ) 24 | 25 | # Add generation config if available 26 | if GENERATION_AVAILABLE: 27 | config.generation = GenerationConfig( 28 | model_id="runwayml/stable-diffusion-v1-5", 29 | device="cuda" if torch.cuda.is_available() else "cpu", 30 | num_inference_steps=50, 31 | guidance_scale=7.5, 32 | width=512, 33 | height=512 34 | ) 35 | 36 | # Initialize the agent 37 | agent = ImageAgent(config) 38 | 39 | # Example 1: OCR on an image 40 | image_path = "path/to/your/image.png" # Replace with actual image path 41 | if os.path.exists(image_path): 42 | logger.info("Performing OCR on image") 43 | try: 44 | # Get detailed OCR results 45 | results = agent.extract_text(image_path) 46 | 47 | # Process results 48 | if isinstance(results, str): 49 | # Simple text output 50 | logger.info("Extracted text: {}", results) 51 | else: 52 | # Detailed results with bounding boxes and confidence 53 | logger.info("Found {} text regions:", len(results)) 54 | for i, result in enumerate(results, 1): 55 | logger.info( 56 | "Region {}: '{}' (confidence: {:.2f})", 57 | i, result["text"], result["confidence"] 58 | ) 59 | if "bbox" in result: 60 | logger.debug("Bounding box: {}", result["bbox"]) 61 | except Exception as e: 62 | logger.error("OCR failed: {}", str(e)) 63 | 64 | # Only run generation examples if available 65 | if GENERATION_AVAILABLE: 66 | # Example 2: Generate an image from text 67 | logger.info("Generating image from text prompt") 68 | try: 69 | prompt = "A serene landscape with mountains and a lake at sunset, digital art style" 70 | image = agent.generate_image( 71 | prompt, 72 | negative_prompt="blurry, low quality, distorted" 73 | ) 74 | 75 | # Save the generated image 76 | output_path = "generated_landscape.png" 77 | image.save(output_path) 78 | logger.info("Image saved to: {}", output_path) 79 | except Exception as e: 80 | logger.error("Image generation failed: {}", str(e)) 81 | 82 | # Example 3: Generate a diagram 83 | logger.info("Generating a flowchart diagram") 84 | try: 85 | description = "Software development lifecycle with stages: Planning, Development, Testing, Deployment, Maintenance" 86 | diagram = agent.generate_diagram( 87 | description, 88 | style="flowchart", 89 | guidance_scale=9.0 # Higher guidance for more precise diagram 90 | ) 91 | 92 | # Save the generated diagram 93 | output_path = "generated_diagram.png" 94 | diagram.save(output_path) 95 | logger.info("Diagram saved to: {}", output_path) 96 | except Exception as e: 97 | logger.error("Diagram generation failed: {}", str(e)) 98 | else: 99 | logger.warning("Image generation features are not available. Install required dependencies to enable them.") 100 | 101 | if __name__ == "__main__": 102 | main() -------------------------------------------------------------------------------- /nexus/agents/evaluator_agent.py: -------------------------------------------------------------------------------- 1 | """ 2 | Evaluator agent for checking output quality and factual coherence. 3 | """ 4 | from typing import Dict, Any, Optional, List 5 | from pydantic import BaseModel 6 | from loguru import logger 7 | import asyncio 8 | from datetime import datetime 9 | 10 | 11 | class EvaluatorConfig(BaseModel): 12 | """Configuration for the EvaluatorAgent.""" 13 | confidence_threshold: float = 0.7 14 | max_retries: int = 3 15 | check_types: List[str] = ["factual_coherence"] # Can add more types later 16 | 17 | 18 | class EvaluatorAgent: 19 | """ 20 | Agent responsible for evaluating LLM outputs for quality and factual coherence. 21 | Can be extended with additional evaluation criteria. 22 | """ 23 | 24 | def __init__(self, config: Optional[EvaluatorConfig] = None): 25 | """Initialize the evaluator with optional configuration.""" 26 | self.config = config or EvaluatorConfig() 27 | logger.info("EvaluatorAgent initialized with config: {}", self.config) 28 | 29 | async def evaluate_output(self, 30 | output: str, 31 | context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: 32 | """ 33 | Evaluate the given output for quality and factual coherence. 34 | 35 | Args: 36 | output: The text to evaluate 37 | context: Optional context for evaluation (e.g., source material) 38 | 39 | Returns: 40 | Dictionary containing evaluation results 41 | """ 42 | logger.debug("Starting evaluation of output") 43 | start_time = datetime.now() 44 | 45 | try: 46 | # Initialize results 47 | results = { 48 | "passed": False, 49 | "confidence": 0.0, 50 | "checks": {}, 51 | "suggestions": [] 52 | } 53 | 54 | # Run enabled checks 55 | if "factual_coherence" in self.config.check_types: 56 | coherence_result = await self._check_factual_coherence(output, context) 57 | results["checks"]["factual_coherence"] = coherence_result 58 | results["confidence"] = coherence_result.get("confidence", 0.0) 59 | 60 | # Determine overall pass/fail 61 | results["passed"] = results["confidence"] >= self.config.confidence_threshold 62 | 63 | duration = (datetime.now() - start_time).total_seconds() 64 | logger.info( 65 | "Evaluation completed in {:.2f}s. Passed: {}", 66 | duration, 67 | results["passed"] 68 | ) 69 | 70 | return results 71 | 72 | except Exception as e: 73 | logger.exception("Error during evaluation") 74 | raise RuntimeError(f"Evaluation failed: {str(e)}") 75 | 76 | async def _check_factual_coherence(self, 77 | output: str, 78 | context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: 79 | """ 80 | Check the factual coherence of the output against available context. 81 | 82 | Args: 83 | output: Text to check 84 | context: Optional reference material or facts 85 | 86 | Returns: 87 | Dictionary with coherence check results 88 | """ 89 | logger.debug("Checking factual coherence") 90 | 91 | try: 92 | # Initialize basic check result 93 | result = { 94 | "confidence": 0.8, # Placeholder - would use actual LLM evaluation 95 | "issues": [], 96 | "suggestions": [] 97 | } 98 | 99 | # TODO: Implement actual coherence checking logic 100 | # This could involve: 101 | # 1. Comparing against known facts in context 102 | # 2. Checking for internal consistency 103 | # 3. Validating against external knowledge base 104 | 105 | return result 106 | 107 | except Exception as e: 108 | logger.error("Factual coherence check failed: {}", str(e)) 109 | return { 110 | "confidence": 0.0, 111 | "issues": [str(e)], 112 | "suggestions": ["Unable to complete coherence check"] 113 | } -------------------------------------------------------------------------------- /nexus/security.py: -------------------------------------------------------------------------------- 1 | """ 2 | Security module for encryption and retry mechanisms. 3 | """ 4 | from typing import TypeVar, Callable, Any, Optional 5 | from pydantic import BaseModel 6 | from loguru import logger 7 | import asyncio 8 | from datetime import datetime 9 | from cryptography.fernet import Fernet 10 | import base64 11 | import os 12 | import json 13 | from functools import wraps 14 | import backoff 15 | 16 | 17 | T = TypeVar('T') # Generic type for retry decorator 18 | 19 | 20 | class SecurityConfig(BaseModel): 21 | """Configuration for security features.""" 22 | encryption_key: Optional[str] = None # If not provided, will generate one 23 | max_retries: int = 3 24 | initial_wait: float = 1.0 # Initial wait time in seconds 25 | max_wait: float = 30.0 # Maximum wait time in seconds 26 | timeout: float = 10.0 # Default timeout in seconds 27 | 28 | 29 | class Security: 30 | """ 31 | Handles encryption and retry mechanisms. 32 | """ 33 | 34 | def __init__(self, config: Optional[SecurityConfig] = None): 35 | """Initialize security with optional configuration.""" 36 | self.config = config or SecurityConfig() 37 | 38 | # Initialize encryption key 39 | if not self.config.encryption_key: 40 | self.config.encryption_key = base64.urlsafe_b64encode(os.urandom(32)).decode() 41 | logger.info("Generated new encryption key") 42 | 43 | self.fernet = Fernet(self.config.encryption_key.encode()) 44 | logger.info("Security module initialized") 45 | 46 | def encrypt_data(self, data: Any) -> bytes: 47 | """ 48 | Encrypt any serializable data. 49 | 50 | Args: 51 | data: Data to encrypt (must be JSON serializable) 52 | 53 | Returns: 54 | Encrypted bytes 55 | """ 56 | try: 57 | # Convert data to JSON string 58 | json_data = json.dumps(data) 59 | 60 | # Encrypt 61 | encrypted = self.fernet.encrypt(json_data.encode()) 62 | 63 | logger.debug("Data encrypted successfully") 64 | return encrypted 65 | 66 | except Exception as e: 67 | logger.exception("Encryption failed") 68 | raise RuntimeError(f"Failed to encrypt data: {str(e)}") 69 | 70 | def decrypt_data(self, encrypted_data: bytes) -> Any: 71 | """ 72 | Decrypt data back to its original form. 73 | 74 | Args: 75 | encrypted_data: Data to decrypt 76 | 77 | Returns: 78 | Decrypted data in its original form 79 | """ 80 | try: 81 | # Decrypt 82 | decrypted = self.fernet.decrypt(encrypted_data) 83 | 84 | # Parse JSON 85 | data = json.loads(decrypted.decode()) 86 | 87 | logger.debug("Data decrypted successfully") 88 | return data 89 | 90 | except Exception as e: 91 | logger.exception("Decryption failed") 92 | raise RuntimeError(f"Failed to decrypt data: {str(e)}") 93 | 94 | @staticmethod 95 | def with_retries( 96 | max_tries: Optional[int] = None, 97 | initial_wait: Optional[float] = None, 98 | max_wait: Optional[float] = None 99 | ): 100 | """ 101 | Decorator for adding exponential backoff retry logic to functions. 102 | 103 | Args: 104 | max_tries: Maximum number of retry attempts 105 | initial_wait: Initial wait time between retries in seconds 106 | max_wait: Maximum wait time between retries in seconds 107 | """ 108 | def decorator(func: Callable[..., T]) -> Callable[..., T]: 109 | @wraps(func) 110 | @backoff.on_exception( 111 | backoff.expo, 112 | Exception, 113 | max_tries=max_tries, 114 | max_time=max_wait, 115 | base=initial_wait 116 | ) 117 | async def wrapper(*args, **kwargs) -> T: 118 | try: 119 | return await func(*args, **kwargs) 120 | except Exception as e: 121 | logger.warning( 122 | "Attempt failed for {}: {}. Retrying...", 123 | func.__name__, 124 | str(e) 125 | ) 126 | raise 127 | return wrapper 128 | return decorator 129 | 130 | 131 | # Convenience function to get a configured security instance 132 | def get_security(config: Optional[SecurityConfig] = None) -> Security: 133 | """Get a configured security instance.""" 134 | return Security(config) -------------------------------------------------------------------------------- /examples/caching_example.py: -------------------------------------------------------------------------------- 1 | """ 2 | Example demonstrating caching and rate limiting features. 3 | """ 4 | from nexus.caching import ( 5 | CacheManager, 6 | CacheConfig, 7 | RateLimitConfig, 8 | cache_embedding, 9 | cache_llm_response 10 | ) 11 | from loguru import logger 12 | import asyncio 13 | import numpy as np 14 | from datetime import datetime 15 | import sys 16 | 17 | 18 | class ExampleProcessor: 19 | """Example processor to demonstrate caching and rate limiting.""" 20 | 21 | def __init__(self, cache_manager: CacheManager): 22 | self.cache_manager = cache_manager 23 | 24 | @cache_embedding(cache_manager=lambda self: self.cache_manager) 25 | async def generate_embedding(self, text: str) -> np.ndarray: 26 | """Simulate embedding generation.""" 27 | # Simulate processing time 28 | await asyncio.sleep(0.5) 29 | 30 | # Generate fake embedding 31 | return np.random.rand(384) # Same dimension as our vector store 32 | 33 | @cache_llm_response(cache_manager=lambda self: self.cache_manager) 34 | async def generate_response(self, prompt: str) -> str: 35 | """Simulate LLM response generation.""" 36 | # Simulate processing time 37 | await asyncio.sleep(1.0) 38 | 39 | # Generate fake response 40 | return f"Response to: {prompt} at {datetime.now().isoformat()}" 41 | 42 | 43 | async def main(): 44 | # Initialize cache manager with custom config 45 | cache_manager = CacheManager( 46 | cache_config=CacheConfig( 47 | embedding_ttl=60, # 1 minute for testing 48 | llm_response_ttl=30, # 30 seconds for testing 49 | max_embedding_size=5, 50 | max_llm_size=3 51 | ), 52 | rate_limit_config=RateLimitConfig( 53 | user_limit=5, # 5 requests per user per minute 54 | global_limit=10, # 10 total requests per minute 55 | embedding_limit=8, # 8 embeddings per minute 56 | llm_limit=6 # 6 LLM calls per minute 57 | ) 58 | ) 59 | 60 | # Initialize processor 61 | processor = ExampleProcessor(cache_manager) 62 | 63 | try: 64 | # Example 1: Embedding Caching 65 | logger.info("\nTesting embedding caching...") 66 | 67 | # Generate embedding for same text multiple times 68 | text = "This is a test text for embedding." 69 | 70 | logger.info("First call - should generate new embedding") 71 | embedding1 = await processor.generate_embedding(text) 72 | 73 | logger.info("Second call - should use cached embedding") 74 | embedding2 = await processor.generate_embedding(text) 75 | 76 | # Verify embeddings are identical 77 | assert np.array_equal(embedding1, embedding2) 78 | logger.info("Embedding cache working correctly") 79 | 80 | # Example 2: LLM Response Caching 81 | logger.info("\nTesting LLM response caching...") 82 | 83 | # Generate response for same prompt multiple times 84 | prompt = "What is the meaning of life?" 85 | 86 | logger.info("First call - should generate new response") 87 | response1 = await processor.generate_response(prompt) 88 | 89 | logger.info("Second call - should use cached response") 90 | response2 = await processor.generate_response(prompt) 91 | 92 | # Verify responses are identical 93 | assert response1 == response2 94 | logger.info("LLM cache working correctly") 95 | 96 | # Example 3: Rate Limiting 97 | logger.info("\nTesting rate limiting...") 98 | 99 | # Test user rate limiting 100 | user_id = "test_user" 101 | logger.info("Testing user rate limiting...") 102 | 103 | for i in range(7): # Try more than the limit 104 | try: 105 | if await cache_manager.check_rate_limit(user_id=user_id): 106 | logger.info("Request {} allowed", i + 1) 107 | else: 108 | logger.warning("Request {} rate limited", i + 1) 109 | except Exception as e: 110 | logger.error("Request {} failed: {}", i + 1, str(e)) 111 | 112 | # Test embedding rate limiting 113 | logger.info("\nTesting embedding rate limiting...") 114 | texts = [f"Text {i}" for i in range(10)] 115 | 116 | for i, text in enumerate(texts): 117 | try: 118 | embedding = await processor.generate_embedding(text) 119 | logger.info("Embedding {} generated", i + 1) 120 | except RuntimeError as e: 121 | logger.warning("Embedding {} rate limited: {}", i + 1, str(e)) 122 | 123 | logger.info("Rate limiting tests completed") 124 | 125 | except Exception as e: 126 | logger.exception("Error during example") 127 | sys.exit(1) 128 | 129 | 130 | if __name__ == "__main__": 131 | # Run the async main function 132 | asyncio.run(main()) -------------------------------------------------------------------------------- /examples/full_demo.py: -------------------------------------------------------------------------------- 1 | """ 2 | Comprehensive demo of Nexus Agents functionality. 3 | """ 4 | import asyncio 5 | import gradio as gr 6 | import pandas as pd 7 | from pathlib import Path 8 | from loguru import logger 9 | from nexus.orchestrator import Orchestrator, OrchestratorConfig 10 | from nexus.security import Security, SecurityConfig 11 | from nexus.caching import CacheManager, CacheConfig, RateLimitConfig 12 | from nexus.monitoring import PerformanceMetrics 13 | from nexus.agents.code_agent import CodeAgent, CodeSandbox 14 | from nexus.agents.data_agent import DataAgent, DataConfig 15 | 16 | 17 | async def initialize_components(): 18 | """Initialize all components with proper configuration.""" 19 | # Set up monitoring 20 | metrics = PerformanceMetrics(port=8000) 21 | logger.info("Initialized performance metrics") 22 | 23 | # Set up security 24 | security = Security( 25 | config=SecurityConfig( 26 | encryption_key="your-secure-key-here", 27 | max_retries=3, 28 | retry_delay=1.0 29 | ) 30 | ) 31 | logger.info("Initialized security") 32 | 33 | # Set up caching 34 | cache_manager = CacheManager( 35 | cache_config=CacheConfig( 36 | ttl_seconds=3600, 37 | max_size_bytes=1024 * 1024 * 100 # 100MB 38 | ), 39 | rate_limit_config=RateLimitConfig( 40 | requests_per_minute=60, 41 | burst_limit=10 42 | ) 43 | ) 44 | logger.info("Initialized cache manager") 45 | 46 | # Set up domain agents 47 | code_agent = CodeAgent( 48 | metrics=metrics, 49 | security=security, 50 | sandbox=CodeSandbox( 51 | allowed_paths=[Path("./examples")], 52 | max_file_size=1024 * 1024, # 1MB 53 | allowed_imports=["os", "sys", "pathlib", "typing"] 54 | ) 55 | ) 56 | logger.info("Initialized code agent") 57 | 58 | data_agent = DataAgent( 59 | metrics=metrics, 60 | security=security, 61 | config=DataConfig( 62 | max_rows=1000000, 63 | max_size=100 * 1024 * 1024, # 100MB 64 | allowed_file_types=[".csv", ".json", ".xlsx"] 65 | ) 66 | ) 67 | logger.info("Initialized data agent") 68 | 69 | # Set up orchestrator 70 | orchestrator = Orchestrator( 71 | config=OrchestratorConfig( 72 | debug=True, 73 | history_length=10 74 | ), 75 | metrics=metrics, 76 | security=security, 77 | cache_manager=cache_manager, 78 | agents={ 79 | "code": code_agent, 80 | "data": data_agent 81 | } 82 | ) 83 | logger.info("Initialized orchestrator") 84 | 85 | return orchestrator, metrics 86 | 87 | 88 | class DemoUI: 89 | """Gradio UI for demonstrating functionality.""" 90 | 91 | def __init__(self): 92 | """Initialize demo UI.""" 93 | self.orchestrator = None 94 | self.metrics = None 95 | 96 | async def startup(self): 97 | """Initialize components on startup.""" 98 | self.orchestrator, self.metrics = await initialize_components() 99 | 100 | async def process_code(self, code: str) -> dict: 101 | """Process code with code agent.""" 102 | try: 103 | agent = self.orchestrator.agents["code"] 104 | analysis = await agent.analyze_code(code) 105 | return { 106 | "status": "success", 107 | "analysis": analysis 108 | } 109 | except Exception as e: 110 | logger.exception("Error processing code") 111 | return { 112 | "status": "error", 113 | "message": str(e) 114 | } 115 | 116 | async def process_data(self, file: str) -> dict: 117 | """Process data with data agent.""" 118 | try: 119 | agent = self.orchestrator.agents["data"] 120 | data = await agent.load_data(file) 121 | analysis = await agent.analyze_data(data, "summary") 122 | return { 123 | "status": "success", 124 | "analysis": analysis 125 | } 126 | except Exception as e: 127 | logger.exception("Error processing data") 128 | return { 129 | "status": "error", 130 | "message": str(e) 131 | } 132 | 133 | async def get_metrics(self) -> dict: 134 | """Get current performance metrics.""" 135 | try: 136 | return self.metrics.get_summary() 137 | except Exception as e: 138 | logger.exception("Error getting metrics") 139 | return { 140 | "status": "error", 141 | "message": str(e) 142 | } 143 | 144 | def launch(self): 145 | """Launch Gradio interface.""" 146 | # Initialize components 147 | asyncio.run(self.startup()) 148 | 149 | # Create interface 150 | with gr.Blocks(title="Mnemosyne Agents Demo") as demo: 151 | gr.Markdown("# Mnemosyne Agents Demo") 152 | 153 | with gr.Tab("Code Analysis"): 154 | code_input = gr.Code( 155 | label="Enter Python Code", 156 | language="python" 157 | ) 158 | code_button = gr.Button("Analyze Code") 159 | code_output = gr.JSON(label="Analysis Results") 160 | 161 | code_button.click( 162 | fn=self.process_code, 163 | inputs=[code_input], 164 | outputs=[code_output] 165 | ) 166 | 167 | with gr.Tab("Data Analysis"): 168 | file_input = gr.File(label="Upload Data File") 169 | data_button = gr.Button("Analyze Data") 170 | data_output = gr.JSON(label="Analysis Results") 171 | 172 | data_button.click( 173 | fn=self.process_data, 174 | inputs=[file_input], 175 | outputs=[data_output] 176 | ) 177 | 178 | with gr.Tab("Monitoring"): 179 | metrics_button = gr.Button("Get Metrics") 180 | metrics_output = gr.JSON(label="Current Metrics") 181 | 182 | metrics_button.click( 183 | fn=self.get_metrics, 184 | inputs=[], 185 | outputs=[metrics_output] 186 | ) 187 | 188 | # Launch interface 189 | demo.launch(server_name="0.0.0.0", server_port=7860) 190 | 191 | 192 | if __name__ == "__main__": 193 | # Create and launch demo 194 | demo = DemoUI() 195 | demo.launch() -------------------------------------------------------------------------------- /nexus/agents/code_agent.py: -------------------------------------------------------------------------------- 1 | """ 2 | Code generation agent with sandbox restrictions. 3 | """ 4 | from typing import Any, Dict, List, Optional 5 | from pydantic import BaseModel, Field 6 | from loguru import logger 7 | import ast 8 | import os 9 | from pathlib import Path 10 | import tempfile 11 | from nexus.monitoring import MonitoredComponent 12 | from nexus.security import Security 13 | 14 | 15 | class CodeSandbox(BaseModel): 16 | """Sandbox configuration for code operations.""" 17 | allowed_paths: List[Path] 18 | max_file_size: int = 1024 * 1024 # 1MB 19 | allowed_imports: List[str] = ["os", "sys", "pathlib", "typing"] 20 | 21 | 22 | class CodeAgent(MonitoredComponent): 23 | """Agent for code generation and manipulation.""" 24 | 25 | def __init__(self, metrics, security: Security, sandbox: CodeSandbox): 26 | """Initialize code agent.""" 27 | super().__init__(metrics) 28 | self.security = security 29 | self.sandbox = sandbox 30 | logger.info("Initialized code agent with sandbox") 31 | 32 | def _validate_path(self, path: Path) -> bool: 33 | """Check if path is allowed in sandbox.""" 34 | try: 35 | path = Path(path).resolve() 36 | return any( 37 | str(path).startswith(str(allowed.resolve())) 38 | for allowed in self.sandbox.allowed_paths 39 | ) 40 | except Exception: 41 | return False 42 | 43 | def _validate_code(self, code: str) -> bool: 44 | """Validate code for security.""" 45 | try: 46 | # Parse code to AST 47 | tree = ast.parse(code) 48 | 49 | # Check imports 50 | for node in ast.walk(tree): 51 | if isinstance(node, ast.Import): 52 | for name in node.names: 53 | if name.name not in self.sandbox.allowed_imports: 54 | logger.warning( 55 | "Disallowed import: {}", 56 | name.name 57 | ) 58 | return False 59 | elif isinstance(node, ast.ImportFrom): 60 | if node.module not in self.sandbox.allowed_imports: 61 | logger.warning( 62 | "Disallowed import from: {}", 63 | node.module 64 | ) 65 | return False 66 | 67 | return True 68 | 69 | except SyntaxError: 70 | logger.warning("Invalid Python syntax") 71 | return False 72 | except Exception as e: 73 | logger.exception("Error validating code") 74 | return False 75 | 76 | async def generate_code(self, 77 | prompt: str, 78 | context: Optional[Dict[str, Any]] = None) -> str: 79 | """ 80 | Generate code based on prompt. 81 | 82 | Args: 83 | prompt: Description of code to generate 84 | context: Additional context (e.g. existing code) 85 | 86 | Returns: 87 | Generated code 88 | """ 89 | async with self.track_operation( 90 | operation="generate_code", 91 | agent_type="code" 92 | ): 93 | # TODO: Call LLM to generate code 94 | # For now, return dummy code 95 | code = """ 96 | def hello_world(): 97 | print("Hello, World!") 98 | return 42 99 | """ 100 | 101 | # Validate generated code 102 | if not self._validate_code(code): 103 | raise ValueError("Generated code failed validation") 104 | 105 | return code 106 | 107 | async def read_file(self, path: str) -> str: 108 | """Read file contents with sandbox validation.""" 109 | async with self.track_operation( 110 | operation="read_file", 111 | agent_type="code" 112 | ): 113 | path = Path(path) 114 | if not self._validate_path(path): 115 | raise ValueError(f"Path not allowed: {path}") 116 | 117 | if not path.exists(): 118 | raise FileNotFoundError(f"File not found: {path}") 119 | 120 | if path.stat().st_size > self.sandbox.max_file_size: 121 | raise ValueError("File too large") 122 | 123 | return path.read_text() 124 | 125 | async def write_file(self, path: str, content: str): 126 | """Write file contents with sandbox validation.""" 127 | async with self.track_operation( 128 | operation="write_file", 129 | agent_type="code" 130 | ): 131 | path = Path(path) 132 | if not self._validate_path(path): 133 | raise ValueError(f"Path not allowed: {path}") 134 | 135 | if len(content.encode()) > self.sandbox.max_file_size: 136 | raise ValueError("Content too large") 137 | 138 | # Validate code if it's a Python file 139 | if path.suffix == ".py": 140 | if not self._validate_code(content): 141 | raise ValueError("Code failed validation") 142 | 143 | # Write to temp file first 144 | with tempfile.NamedTemporaryFile(mode="w", delete=False) as tmp: 145 | tmp.write(content) 146 | 147 | # Move temp file to target 148 | os.replace(tmp.name, path) 149 | 150 | async def analyze_code(self, 151 | code: str, 152 | context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]: 153 | """ 154 | Analyze code for quality and issues. 155 | 156 | Args: 157 | code: Code to analyze 158 | context: Additional context 159 | 160 | Returns: 161 | Analysis results 162 | """ 163 | async with self.track_operation( 164 | operation="analyze_code", 165 | agent_type="code" 166 | ): 167 | # Validate code first 168 | if not self._validate_code(code): 169 | raise ValueError("Code failed validation") 170 | 171 | # TODO: Call LLM to analyze code 172 | # For now, return dummy analysis 173 | return { 174 | "complexity": "low", 175 | "maintainability": "good", 176 | "issues": [], 177 | "suggestions": [ 178 | "Add docstrings", 179 | "Add type hints" 180 | ] 181 | } -------------------------------------------------------------------------------- /nexus/agents/data_agent.py: -------------------------------------------------------------------------------- 1 | """ 2 | Data analysis agent for processing and analyzing data. 3 | """ 4 | from typing import Any, Dict, List, Optional 5 | from pydantic import BaseModel 6 | from loguru import logger 7 | import pandas as pd 8 | import numpy as np 9 | from nexus.monitoring import MonitoredComponent 10 | from nexus.security import Security 11 | 12 | 13 | class DataConfig(BaseModel): 14 | """Configuration for data operations.""" 15 | max_rows: int = 1000000 # 1M rows 16 | max_size: int = 100 * 1024 * 1024 # 100MB 17 | allowed_file_types: List[str] = [".csv", ".json", ".xlsx"] 18 | 19 | 20 | class DataAgent(MonitoredComponent): 21 | """Agent for data analysis and processing.""" 22 | 23 | def __init__(self, metrics, security: Security, config: DataConfig): 24 | """Initialize data agent.""" 25 | super().__init__(metrics) 26 | self.security = security 27 | self.config = config 28 | logger.info("Initialized data agent") 29 | 30 | def _validate_data(self, data: pd.DataFrame) -> bool: 31 | """Validate data size and content.""" 32 | if len(data) > self.config.max_rows: 33 | logger.warning("Data exceeds max rows") 34 | return False 35 | 36 | # Check memory usage 37 | memory_usage = data.memory_usage(deep=True).sum() 38 | if memory_usage > self.config.max_size: 39 | logger.warning("Data exceeds max size") 40 | return False 41 | 42 | return True 43 | 44 | async def load_data(self, path: str) -> pd.DataFrame: 45 | """Load data from file with validation.""" 46 | async with self.track_operation( 47 | operation="load_data", 48 | agent_type="data" 49 | ): 50 | # Check file type 51 | if not any(path.endswith(ext) for ext in self.config.allowed_file_types): 52 | raise ValueError("Unsupported file type") 53 | 54 | # Load data based on type 55 | if path.endswith(".csv"): 56 | data = pd.read_csv(path) 57 | elif path.endswith(".json"): 58 | data = pd.read_json(path) 59 | elif path.endswith(".xlsx"): 60 | data = pd.read_excel(path) 61 | 62 | # Validate loaded data 63 | if not self._validate_data(data): 64 | raise ValueError("Data validation failed") 65 | 66 | return data 67 | 68 | async def analyze_data(self, 69 | data: pd.DataFrame, 70 | analysis_type: str) -> Dict[str, Any]: 71 | """ 72 | Analyze data based on specified type. 73 | 74 | Args: 75 | data: DataFrame to analyze 76 | analysis_type: Type of analysis to perform 77 | 78 | Returns: 79 | Analysis results 80 | """ 81 | async with self.track_operation( 82 | operation="analyze_data", 83 | agent_type="data" 84 | ): 85 | # Validate data first 86 | if not self._validate_data(data): 87 | raise ValueError("Data validation failed") 88 | 89 | results = {} 90 | 91 | if analysis_type == "summary": 92 | results["summary"] = data.describe().to_dict() 93 | results["missing"] = data.isnull().sum().to_dict() 94 | results["dtypes"] = data.dtypes.astype(str).to_dict() 95 | 96 | elif analysis_type == "correlation": 97 | # Only numeric columns 98 | numeric_data = data.select_dtypes(include=[np.number]) 99 | results["correlation"] = numeric_data.corr().to_dict() 100 | 101 | elif analysis_type == "distribution": 102 | results["distribution"] = { 103 | col: data[col].value_counts().to_dict() 104 | for col in data.columns 105 | } 106 | 107 | else: 108 | raise ValueError(f"Unknown analysis type: {analysis_type}") 109 | 110 | return results 111 | 112 | async def process_data(self, 113 | data: pd.DataFrame, 114 | operations: List[Dict[str, Any]]) -> pd.DataFrame: 115 | """ 116 | Process data with specified operations. 117 | 118 | Args: 119 | data: DataFrame to process 120 | operations: List of operations to apply 121 | 122 | Returns: 123 | Processed DataFrame 124 | """ 125 | async with self.track_operation( 126 | operation="process_data", 127 | agent_type="data" 128 | ): 129 | # Validate input data 130 | if not self._validate_data(data): 131 | raise ValueError("Input data validation failed") 132 | 133 | result = data.copy() 134 | 135 | for op in operations: 136 | op_type = op["type"] 137 | 138 | if op_type == "filter": 139 | col = op["column"] 140 | condition = op["condition"] 141 | value = op["value"] 142 | 143 | if condition == "equals": 144 | result = result[result[col] == value] 145 | elif condition == "greater_than": 146 | result = result[result[col] > value] 147 | elif condition == "less_than": 148 | result = result[result[col] < value] 149 | else: 150 | raise ValueError(f"Unknown condition: {condition}") 151 | 152 | elif op_type == "transform": 153 | col = op["column"] 154 | transform = op["transform"] 155 | 156 | if transform == "normalize": 157 | result[col] = (result[col] - result[col].mean()) / result[col].std() 158 | elif transform == "fillna": 159 | value = op.get("value", result[col].mean()) 160 | result[col] = result[col].fillna(value) 161 | else: 162 | raise ValueError(f"Unknown transform: {transform}") 163 | 164 | else: 165 | raise ValueError(f"Unknown operation type: {op_type}") 166 | 167 | # Validate intermediate result 168 | if not self._validate_data(result): 169 | raise ValueError("Intermediate result validation failed") 170 | 171 | return result -------------------------------------------------------------------------------- /nexus/vector_store.py: -------------------------------------------------------------------------------- 1 | """ 2 | Vector store module for persistent memory using LanceDB. 3 | """ 4 | from typing import List, Dict, Any, Optional 5 | from datetime import datetime 6 | import lancedb 7 | import os 8 | from pathlib import Path 9 | from loguru import logger 10 | from pydantic import BaseModel 11 | import numpy as np 12 | from sentence_transformers import SentenceTransformer 13 | 14 | 15 | class VectorStoreConfig(BaseModel): 16 | """Configuration for the vector store.""" 17 | db_path: str = "data/vector_store" 18 | collection_name: str = "conversation_history" 19 | embedding_model: str = "all-MiniLM-L6-v2" 20 | dimension: int = 384 # Default for all-MiniLM-L6-v2 21 | 22 | 23 | class ConversationEntry(BaseModel): 24 | """Schema for conversation entries in the vector store.""" 25 | id: str 26 | role: str 27 | content: str 28 | embedding: List[float] 29 | timestamp: datetime 30 | metadata: Dict[str, Any] = {} 31 | 32 | 33 | class VectorStore: 34 | """ 35 | Vector store for persistent memory using LanceDB. 36 | Handles conversation history storage and retrieval. 37 | """ 38 | 39 | def __init__(self, config: Optional[VectorStoreConfig] = None): 40 | """Initialize the vector store with optional configuration.""" 41 | self.config = config or VectorStoreConfig() 42 | 43 | # Create db directory if it doesn't exist 44 | os.makedirs(os.path.dirname(self.config.db_path), exist_ok=True) 45 | 46 | # Initialize LanceDB 47 | self.db = lancedb.connect(self.config.db_path) 48 | 49 | # Load embedding model 50 | self.embedding_model = SentenceTransformer(self.config.embedding_model) 51 | 52 | # Initialize or get collection 53 | self._init_collection() 54 | 55 | logger.info( 56 | "VectorStore initialized with path: {} and collection: {}", 57 | self.config.db_path, 58 | self.config.collection_name 59 | ) 60 | 61 | def _init_collection(self): 62 | """Initialize or get the conversation history collection.""" 63 | try: 64 | self.collection = self.db.open_table(self.config.collection_name) 65 | logger.debug("Opened existing collection: {}", self.config.collection_name) 66 | except Exception: 67 | logger.info("Creating new collection: {}", self.config.collection_name) 68 | schema = { 69 | "id": "string", 70 | "role": "string", 71 | "content": "string", 72 | "embedding": f"vector({self.config.dimension})", 73 | "timestamp": "timestamp", 74 | "metadata": "json" 75 | } 76 | self.collection = self.db.create_table( 77 | self.config.collection_name, 78 | schema=schema 79 | ) 80 | 81 | def _get_embedding(self, text: str) -> List[float]: 82 | """Get embedding vector for text using sentence-transformers.""" 83 | return self.embedding_model.encode(text).tolist() 84 | 85 | async def add_message(self, 86 | role: str, 87 | content: str, 88 | metadata: Optional[Dict[str, Any]] = None) -> str: 89 | """ 90 | Add a message to the conversation history. 91 | 92 | Args: 93 | role: Message role (user/assistant/system) 94 | content: Message content 95 | metadata: Optional metadata about the message 96 | 97 | Returns: 98 | ID of the added message 99 | """ 100 | try: 101 | # Generate unique ID based on timestamp 102 | message_id = f"{role}_{datetime.now().isoformat()}" 103 | 104 | # Get embedding 105 | embedding = self._get_embedding(content) 106 | 107 | # Create entry 108 | entry = ConversationEntry( 109 | id=message_id, 110 | role=role, 111 | content=content, 112 | embedding=embedding, 113 | timestamp=datetime.now(), 114 | metadata=metadata or {} 115 | ) 116 | 117 | # Add to collection 118 | self.collection.add([entry.dict()]) 119 | 120 | logger.debug("Added message to vector store. ID: {}", message_id) 121 | return message_id 122 | 123 | except Exception as e: 124 | logger.exception("Error adding message to vector store") 125 | raise RuntimeError(f"Failed to add message: {str(e)}") 126 | 127 | async def search_similar(self, 128 | query: str, 129 | limit: int = 5, 130 | role_filter: Optional[str] = None) -> List[Dict[str, Any]]: 131 | """ 132 | Search for similar messages in the conversation history. 133 | 134 | Args: 135 | query: Text to search for 136 | limit: Maximum number of results 137 | role_filter: Optional filter by role 138 | 139 | Returns: 140 | List of similar messages with scores 141 | """ 142 | try: 143 | # Get query embedding 144 | query_embedding = self._get_embedding(query) 145 | 146 | # Build search query 147 | search = self.collection.search(query_embedding) 148 | if role_filter: 149 | search = search.where(f"role = '{role_filter}'") 150 | 151 | # Execute search 152 | results = search.limit(limit).to_list() 153 | 154 | logger.debug("Found {} similar messages", len(results)) 155 | return results 156 | 157 | except Exception as e: 158 | logger.exception("Error searching vector store") 159 | raise RuntimeError(f"Search failed: {str(e)}") 160 | 161 | async def get_recent_history(self, 162 | limit: int = 10, 163 | role_filter: Optional[str] = None) -> List[Dict[str, Any]]: 164 | """ 165 | Get the most recent conversation history. 166 | 167 | Args: 168 | limit: Maximum number of messages to return 169 | role_filter: Optional filter by role 170 | 171 | Returns: 172 | List of recent messages 173 | """ 174 | try: 175 | # Build query 176 | query = self.collection 177 | if role_filter: 178 | query = query.where(f"role = '{role_filter}'") 179 | 180 | # Execute query 181 | results = query.order_by("timestamp", "desc").limit(limit).to_list() 182 | 183 | logger.debug("Retrieved {} recent messages", len(results)) 184 | return results 185 | 186 | except Exception as e: 187 | logger.exception("Error retrieving history") 188 | raise RuntimeError(f"History retrieval failed: {str(e)}") 189 | 190 | async def clear_history(self): 191 | """Clear all conversation history.""" 192 | try: 193 | self.db.drop_table(self.config.collection_name) 194 | self._init_collection() 195 | logger.info("Cleared conversation history") 196 | 197 | except Exception as e: 198 | logger.exception("Error clearing history") 199 | raise RuntimeError(f"Failed to clear history: {str(e)}") -------------------------------------------------------------------------------- /nexus/monitoring.py: -------------------------------------------------------------------------------- 1 | """ 2 | Performance monitoring module with Prometheus integration. 3 | """ 4 | from typing import Any, Dict, List, Optional, Union, Callable 5 | from pydantic import BaseModel, Field 6 | from loguru import logger 7 | import time 8 | import asyncio 9 | from datetime import datetime, timedelta 10 | from enum import Enum 11 | from collections import defaultdict 12 | import threading 13 | from prometheus_client import ( 14 | Counter, 15 | Histogram, 16 | Gauge, 17 | start_http_server, 18 | REGISTRY, 19 | CollectorRegistry 20 | ) 21 | 22 | 23 | class MetricType(str, Enum): 24 | """Types of metrics to track.""" 25 | COUNTER = "counter" # Monotonically increasing counter 26 | HISTOGRAM = "histogram" # Distribution of values 27 | GAUGE = "gauge" # Point-in-time value 28 | 29 | 30 | class MetricLabel(str, Enum): 31 | """Common metric labels.""" 32 | AGENT_TYPE = "agent_type" 33 | OPERATION = "operation" 34 | STATUS = "status" 35 | ERROR_TYPE = "error_type" 36 | 37 | 38 | class PerformanceMetrics: 39 | """ 40 | Manages performance metrics collection and reporting. 41 | """ 42 | 43 | def __init__(self, port: int = 8000): 44 | """Initialize metrics with Prometheus registry.""" 45 | # Create metrics 46 | self.request_count = Counter( 47 | "mnemosyne_requests_total", 48 | "Total number of requests", 49 | ["agent_type", "operation", "status"] 50 | ) 51 | 52 | self.request_latency = Histogram( 53 | "mnemosyne_request_duration_seconds", 54 | "Request duration in seconds", 55 | ["agent_type", "operation"], 56 | buckets=[0.1, 0.5, 1.0, 2.0, 5.0, 10.0] 57 | ) 58 | 59 | self.error_count = Counter( 60 | "mnemosyne_errors_total", 61 | "Total number of errors", 62 | ["agent_type", "error_type"] 63 | ) 64 | 65 | self.active_workers = Gauge( 66 | "mnemosyne_active_workers", 67 | "Number of active workers", 68 | ["agent_type"] 69 | ) 70 | 71 | self.queue_size = Gauge( 72 | "mnemosyne_queue_size", 73 | "Number of items in queue", 74 | ["agent_type"] 75 | ) 76 | 77 | self.memory_usage = Gauge( 78 | "mnemosyne_memory_mb", 79 | "Memory usage in MB", 80 | ["component"] 81 | ) 82 | 83 | # Start Prometheus HTTP server 84 | start_http_server(port) 85 | logger.info("Started Prometheus metrics server on port {}", port) 86 | 87 | # Initialize aggregation state 88 | self._reset_aggregation() 89 | 90 | # Start background tasks 91 | self._start_background_tasks() 92 | 93 | def _reset_aggregation(self): 94 | """Reset aggregation state.""" 95 | self._agg_lock = threading.Lock() 96 | self._agg_data = defaultdict(lambda: { 97 | "count": 0, 98 | "total_time": 0.0, 99 | "errors": defaultdict(int) 100 | }) 101 | 102 | def _start_background_tasks(self): 103 | """Start background monitoring tasks.""" 104 | async def update_memory_metrics(): 105 | while True: 106 | try: 107 | # TODO: Get actual memory usage 108 | # For now, use dummy values 109 | self.memory_usage.labels("orchestrator").set(100) 110 | self.memory_usage.labels("workers").set(200) 111 | self.memory_usage.labels("cache").set(50) 112 | except Exception as e: 113 | logger.exception("Error updating memory metrics") 114 | await asyncio.sleep(60) # Update every minute 115 | 116 | # Start memory metrics task 117 | asyncio.create_task(update_memory_metrics()) 118 | 119 | def track_request(self, 120 | agent_type: str, 121 | operation: str) -> "RequestTracker": 122 | """ 123 | Create a context manager to track request metrics. 124 | 125 | Args: 126 | agent_type: Type of agent making request 127 | operation: Operation being performed 128 | 129 | Returns: 130 | Context manager for tracking request 131 | """ 132 | return RequestTracker( 133 | metrics=self, 134 | agent_type=agent_type, 135 | operation=operation 136 | ) 137 | 138 | def track_worker(self, 139 | agent_type: str) -> "WorkerTracker": 140 | """ 141 | Create a context manager to track worker metrics. 142 | 143 | Args: 144 | agent_type: Type of worker 145 | 146 | Returns: 147 | Context manager for tracking worker 148 | """ 149 | return WorkerTracker( 150 | metrics=self, 151 | agent_type=agent_type 152 | ) 153 | 154 | def record_error(self, 155 | agent_type: str, 156 | error_type: str): 157 | """Record an error occurrence.""" 158 | self.error_count.labels( 159 | agent_type=agent_type, 160 | error_type=error_type 161 | ).inc() 162 | 163 | with self._agg_lock: 164 | self._agg_data[agent_type]["errors"][error_type] += 1 165 | 166 | def update_queue_size(self, 167 | agent_type: str, 168 | size: int): 169 | """Update queue size metric.""" 170 | self.queue_size.labels(agent_type=agent_type).set(size) 171 | 172 | def get_summary(self, 173 | agent_type: Optional[str] = None, 174 | window: timedelta = timedelta(minutes=5)) -> Dict[str, Any]: 175 | """ 176 | Get summary metrics for the specified window. 177 | 178 | Args: 179 | agent_type: Optional agent type to filter by 180 | window: Time window to summarize 181 | 182 | Returns: 183 | Dictionary of summary metrics 184 | """ 185 | with self._agg_lock: 186 | if agent_type: 187 | data = { 188 | agent_type: self._agg_data[agent_type] 189 | } 190 | else: 191 | data = dict(self._agg_data) 192 | 193 | summary = {} 194 | for agent, metrics in data.items(): 195 | if metrics["count"] > 0: 196 | avg_time = metrics["total_time"] / metrics["count"] 197 | else: 198 | avg_time = 0.0 199 | 200 | summary[agent] = { 201 | "requests": metrics["count"], 202 | "avg_latency": avg_time, 203 | "error_counts": dict(metrics["errors"]) 204 | } 205 | 206 | return summary 207 | 208 | 209 | class RequestTracker: 210 | """Context manager for tracking request metrics.""" 211 | 212 | def __init__(self, 213 | metrics: PerformanceMetrics, 214 | agent_type: str, 215 | operation: str): 216 | """Initialize the tracker.""" 217 | self.metrics = metrics 218 | self.agent_type = agent_type 219 | self.operation = operation 220 | self.start_time = None 221 | self.error = None 222 | 223 | async def __aenter__(self): 224 | """Start tracking request.""" 225 | self.start_time = time.time() 226 | return self 227 | 228 | async def __aexit__(self, exc_type, exc_val, exc_tb): 229 | """Stop tracking request and record metrics.""" 230 | duration = time.time() - self.start_time 231 | 232 | # Record latency 233 | self.metrics.request_latency.labels( 234 | agent_type=self.agent_type, 235 | operation=self.operation 236 | ).observe(duration) 237 | 238 | # Record request count 239 | status = "error" if exc_type else "success" 240 | self.metrics.request_count.labels( 241 | agent_type=self.agent_type, 242 | operation=self.operation, 243 | status=status 244 | ).inc() 245 | 246 | # Record error if any 247 | if exc_type: 248 | self.metrics.record_error( 249 | self.agent_type, 250 | exc_type.__name__ 251 | ) 252 | 253 | # Update aggregation 254 | with self.metrics._agg_lock: 255 | agg = self.metrics._agg_data[self.agent_type] 256 | agg["count"] += 1 257 | agg["total_time"] += duration 258 | 259 | 260 | class WorkerTracker: 261 | """Context manager for tracking worker metrics.""" 262 | 263 | def __init__(self, 264 | metrics: PerformanceMetrics, 265 | agent_type: str): 266 | """Initialize the tracker.""" 267 | self.metrics = metrics 268 | self.agent_type = agent_type 269 | 270 | async def __aenter__(self): 271 | """Start tracking worker.""" 272 | self.metrics.active_workers.labels( 273 | agent_type=self.agent_type 274 | ).inc() 275 | return self 276 | 277 | async def __aexit__(self, exc_type, exc_val, exc_tb): 278 | """Stop tracking worker.""" 279 | self.metrics.active_workers.labels( 280 | agent_type=self.agent_type 281 | ).dec() 282 | 283 | 284 | class MonitoredComponent: 285 | """Base class for components with performance monitoring.""" 286 | 287 | def __init__(self, metrics: PerformanceMetrics): 288 | """Initialize with metrics.""" 289 | self.metrics = metrics 290 | 291 | async def track_operation(self, 292 | operation: str, 293 | agent_type: str = "general") -> "RequestTracker": 294 | """Track an operation with metrics.""" 295 | return self.metrics.track_request(agent_type, operation) 296 | 297 | def track_worker(self, agent_type: str = "general") -> "WorkerTracker": 298 | """Track a worker with metrics.""" 299 | return self.metrics.track_worker(agent_type) 300 | 301 | def record_error(self, 302 | error_type: str, 303 | agent_type: str = "general"): 304 | """Record an error.""" 305 | self.metrics.record_error(agent_type, error_type) 306 | 307 | def update_queue_size(self, 308 | size: int, 309 | agent_type: str = "general"): 310 | """Update queue size metric.""" 311 | self.metrics.update_queue_size(agent_type, size) -------------------------------------------------------------------------------- /nexus/evaluator.py: -------------------------------------------------------------------------------- 1 | """ 2 | Enhanced evaluator module for multi-criteria output evaluation and refinement. 3 | """ 4 | from typing import Any, Dict, List, Optional, Union 5 | from pydantic import BaseModel, Field, validator 6 | from loguru import logger 7 | import asyncio 8 | from datetime import datetime 9 | from enum import Enum 10 | import json 11 | 12 | 13 | class EvaluationCriteria(str, Enum): 14 | """Types of evaluation criteria.""" 15 | FACTUAL_ACCURACY = "factual_accuracy" 16 | STYLE_CONSISTENCY = "style_consistency" 17 | POLICY_COMPLIANCE = "policy_compliance" 18 | RELEVANCE = "relevance" 19 | COHERENCE = "coherence" 20 | CREATIVITY = "creativity" 21 | SAFETY = "safety" 22 | 23 | 24 | class EvaluationScore(BaseModel): 25 | """Score for a single evaluation criterion.""" 26 | criterion: EvaluationCriteria 27 | score: float # 0.0 to 1.0 28 | feedback: str 29 | suggestions: List[str] = [] 30 | 31 | 32 | class EvaluationResult(BaseModel): 33 | """Complete evaluation result for a piece of content.""" 34 | scores: List[EvaluationScore] 35 | overall_score: float 36 | needs_refinement: bool 37 | refinement_prompt: Optional[str] = None 38 | metadata: Dict[str, Any] = {} 39 | timestamp: datetime = Field(default_factory=datetime.now) 40 | 41 | @validator("overall_score") 42 | def validate_score(cls, v): 43 | """Validate score is between 0 and 1.""" 44 | if not 0 <= v <= 1: 45 | raise ValueError("Score must be between 0 and 1") 46 | return v 47 | 48 | 49 | class EvaluationConfig(BaseModel): 50 | """Configuration for the evaluator.""" 51 | criteria: List[EvaluationCriteria] 52 | weights: Dict[EvaluationCriteria, float] 53 | thresholds: Dict[EvaluationCriteria, float] 54 | refinement_threshold: float = 0.7 # Overall score below this triggers refinement 55 | 56 | @validator("weights") 57 | def validate_weights(cls, v, values): 58 | """Validate weights match criteria and sum to 1.""" 59 | if "criteria" in values: 60 | if set(v.keys()) != set(values["criteria"]): 61 | raise ValueError("Weights must match criteria exactly") 62 | if abs(sum(v.values()) - 1.0) > 0.001: 63 | raise ValueError("Weights must sum to 1") 64 | return v 65 | 66 | @validator("thresholds") 67 | def validate_thresholds(cls, v, values): 68 | """Validate thresholds match criteria and are between 0 and 1.""" 69 | if "criteria" in values: 70 | if set(v.keys()) != set(values["criteria"]): 71 | raise ValueError("Thresholds must match criteria exactly") 72 | if not all(0 <= t <= 1 for t in v.values()): 73 | raise ValueError("Thresholds must be between 0 and 1") 74 | return v 75 | 76 | 77 | class ContentEvaluator: 78 | """ 79 | Evaluates content based on multiple criteria and suggests refinements. 80 | """ 81 | 82 | def __init__(self, config: EvaluationConfig): 83 | """Initialize the evaluator with config.""" 84 | self.config = config 85 | logger.info( 86 | "Initialized evaluator with {} criteria", 87 | len(config.criteria) 88 | ) 89 | 90 | async def evaluate_criterion(self, 91 | content: str, 92 | criterion: EvaluationCriteria, 93 | context: Optional[Dict[str, Any]] = None) -> EvaluationScore: 94 | """ 95 | Evaluate content for a single criterion. 96 | 97 | Args: 98 | content: Content to evaluate 99 | criterion: Criterion to evaluate against 100 | context: Additional context for evaluation 101 | 102 | Returns: 103 | Evaluation score with feedback 104 | """ 105 | # Create prompt based on criterion 106 | prompt = self._create_evaluation_prompt(content, criterion, context) 107 | 108 | # TODO: Call LLM to evaluate 109 | # For now, simulate evaluation 110 | await asyncio.sleep(0.5) 111 | 112 | # Parse response into score and feedback 113 | if criterion == EvaluationCriteria.FACTUAL_ACCURACY: 114 | score = 0.85 115 | feedback = "Content appears mostly accurate" 116 | suggestions = ["Consider adding sources for key claims"] 117 | else: 118 | score = 0.75 119 | feedback = f"Acceptable {criterion.value}" 120 | suggestions = [f"Could improve {criterion.value}"] 121 | 122 | return EvaluationScore( 123 | criterion=criterion, 124 | score=score, 125 | feedback=feedback, 126 | suggestions=suggestions 127 | ) 128 | 129 | def _create_evaluation_prompt(self, 130 | content: str, 131 | criterion: EvaluationCriteria, 132 | context: Optional[Dict[str, Any]] = None) -> str: 133 | """Create an evaluation prompt for the LLM.""" 134 | prompts = { 135 | EvaluationCriteria.FACTUAL_ACCURACY: """ 136 | Evaluate the factual accuracy of this content: 137 | --- 138 | {content} 139 | --- 140 | Consider: 141 | 1. Are all statements supported by evidence? 142 | 2. Are there any contradictions? 143 | 3. Are statistics and numbers accurate? 144 | 145 | Provide a score from 0.0 to 1.0 and explain your reasoning. 146 | """, 147 | EvaluationCriteria.STYLE_CONSISTENCY: """ 148 | Evaluate the style consistency of this content: 149 | --- 150 | {content} 151 | --- 152 | Consider: 153 | 1. Is the tone consistent? 154 | 2. Is the formality level maintained? 155 | 3. Are transitions smooth? 156 | 157 | Provide a score from 0.0 to 1.0 and explain your reasoning. 158 | """, 159 | # Add prompts for other criteria... 160 | } 161 | 162 | base_prompt = prompts.get( 163 | criterion, 164 | "Evaluate the {criterion} of this content: {content}" 165 | ) 166 | 167 | return base_prompt.format( 168 | content=content, 169 | criterion=criterion.value, 170 | **(context or {}) 171 | ) 172 | 173 | def _create_refinement_prompt(self, 174 | content: str, 175 | scores: List[EvaluationScore], 176 | context: Optional[Dict[str, Any]] = None) -> str: 177 | """Create a refinement prompt based on evaluation scores.""" 178 | # Identify areas needing improvement 179 | improvements = [] 180 | for score in scores: 181 | if score.score < self.config.thresholds[score.criterion]: 182 | improvements.extend(score.suggestions) 183 | 184 | prompt = f""" 185 | Please improve this content: 186 | --- 187 | {content} 188 | --- 189 | Focus on these aspects: 190 | {json.dumps(improvements, indent=2)} 191 | 192 | Maintain any correct and high-quality portions while addressing the issues. 193 | """ 194 | 195 | return prompt 196 | 197 | async def evaluate(self, 198 | content: str, 199 | context: Optional[Dict[str, Any]] = None) -> EvaluationResult: 200 | """ 201 | Evaluate content against all configured criteria. 202 | 203 | Args: 204 | content: Content to evaluate 205 | context: Additional context for evaluation 206 | 207 | Returns: 208 | Complete evaluation result 209 | """ 210 | # Evaluate each criterion 211 | scores = [] 212 | for criterion in self.config.criteria: 213 | score = await self.evaluate_criterion( 214 | content, 215 | criterion, 216 | context 217 | ) 218 | scores.append(score) 219 | 220 | # Calculate overall score 221 | overall_score = sum( 222 | score.score * self.config.weights[score.criterion] 223 | for score in scores 224 | ) 225 | 226 | # Check if refinement is needed 227 | needs_refinement = overall_score < self.config.refinement_threshold 228 | refinement_prompt = None 229 | 230 | if needs_refinement: 231 | refinement_prompt = self._create_refinement_prompt( 232 | content, 233 | scores, 234 | context 235 | ) 236 | 237 | return EvaluationResult( 238 | scores=scores, 239 | overall_score=overall_score, 240 | needs_refinement=needs_refinement, 241 | refinement_prompt=refinement_prompt 242 | ) 243 | 244 | async def evaluate_and_refine(self, 245 | content: str, 246 | context: Optional[Dict[str, Any]] = None, 247 | max_iterations: int = 3) -> tuple[str, List[EvaluationResult]]: 248 | """ 249 | Evaluate content and automatically refine until it meets thresholds. 250 | 251 | Args: 252 | content: Initial content to evaluate and refine 253 | context: Additional context for evaluation 254 | max_iterations: Maximum number of refinement attempts 255 | 256 | Returns: 257 | Tuple of (final content, list of evaluation results) 258 | """ 259 | current_content = content 260 | results = [] 261 | 262 | for i in range(max_iterations): 263 | # Evaluate current content 264 | result = await self.evaluate(current_content, context) 265 | results.append(result) 266 | 267 | # Check if refinement is needed 268 | if not result.needs_refinement: 269 | logger.info("Content meets all criteria after {} iterations", i + 1) 270 | break 271 | 272 | # Refine content 273 | if result.refinement_prompt: 274 | # TODO: Call LLM with refinement prompt 275 | # For now, simulate refinement 276 | await asyncio.sleep(1.0) 277 | current_content = f"Refined ({i + 1}): {current_content}" 278 | 279 | logger.info( 280 | "Completed refinement iteration {} with score {:.2f}", 281 | i + 1, 282 | result.overall_score 283 | ) 284 | 285 | return current_content, results -------------------------------------------------------------------------------- /nexus/caching.py: -------------------------------------------------------------------------------- 1 | """ 2 | Caching module for embeddings and LLM responses. 3 | 4 | Key Components: 5 | 1. Configuration Classes: 6 | - CacheConfig: Defines caching parameters such as TTL and maximum sizes. 7 | - RateLimitConfig: Sets up rate limiting for user and global requests. 8 | 9 | 2. CacheEntry Class: Represents a single cache entry with data and expiration time. 10 | 11 | 3. Cache Class: Implements the caching mechanism, allowing for storing and retrieving data based on generated keys, 12 | enforcing size limits, and removing the oldest entries when necessary. 13 | 14 | 4. RateLimiter Class: Manages request limits for users and globally, ensuring the system is not overwhelmed by too many requests. 15 | 16 | 5. CacheManager Class: A higher-level interface for managing caching and rate limiting, integrating both functionalities 17 | into the overall workflow of the platform. 18 | 19 | 6. Caching Decorators: Provides decorators for caching embeddings and LLM responses automatically, enhancing performance 20 | without modifying core function logic. 21 | 22 | Usage: 23 | - Create an instance of CacheManager to manage caching and rate limiting. 24 | - Decorate functions with caching decorators to automatically handle caching. 25 | 26 | Capabilities: 27 | - Reduces redundant computations and API calls, improving efficiency. 28 | - Provides control over request limits, ensuring fair usage. 29 | - Essential for maintaining responsiveness and reliability under heavy load. 30 | 31 | This module is integral to building efficient and responsive AI applications across different modalities. 32 | """ 33 | from typing import Any, Optional, Dict, List 34 | from pydantic import BaseModel 35 | from loguru import logger 36 | import hashlib 37 | import json 38 | import time 39 | from datetime import datetime, timedelta 40 | import asyncio 41 | from functools import wraps 42 | import numpy as np 43 | 44 | 45 | class CacheConfig(BaseModel): 46 | """Configuration for caching.""" 47 | embedding_ttl: int = 3600 # Time to live for embeddings in seconds 48 | llm_response_ttl: int = 1800 # Time to live for LLM responses in seconds 49 | max_embedding_size: int = 10000 # Maximum number of embeddings to cache 50 | max_llm_size: int = 1000 # Maximum number of LLM responses to cache 51 | 52 | 53 | class RateLimitConfig(BaseModel): 54 | """Configuration for rate limiting.""" 55 | user_limit: int = 100 # Requests per user per minute 56 | global_limit: int = 1000 # Total requests per minute 57 | embedding_limit: int = 500 # Embedding generations per minute 58 | llm_limit: int = 200 # LLM calls per minute 59 | 60 | 61 | class CacheEntry(BaseModel): 62 | """A single cache entry.""" 63 | data: Any 64 | expires_at: datetime 65 | 66 | 67 | class Cache: 68 | """ 69 | Cache implementation with TTL and size limits. 70 | """ 71 | def __init__(self, max_size: int, ttl: int): 72 | """Initialize cache with size and TTL limits.""" 73 | self.max_size = max_size 74 | self.ttl = ttl 75 | self._cache: Dict[str, CacheEntry] = {} 76 | logger.info("Cache initialized with max_size={}, ttl={}s", max_size, ttl) 77 | 78 | def _generate_key(self, data: Any) -> str: 79 | """Generate cache key from data.""" 80 | if isinstance(data, np.ndarray): 81 | data = data.tobytes() 82 | elif not isinstance(data, (str, bytes)): 83 | data = json.dumps(data, sort_keys=True).encode() 84 | 85 | return hashlib.sha256(data).hexdigest() 86 | 87 | def get(self, key: str) -> Optional[Any]: 88 | """Get item from cache if not expired.""" 89 | entry = self._cache.get(key) 90 | if not entry: 91 | return None 92 | 93 | if datetime.now() > entry.expires_at: 94 | del self._cache[key] 95 | return None 96 | 97 | return entry.data 98 | 99 | def set(self, key: str, value: Any): 100 | """Set item in cache with TTL.""" 101 | # Remove oldest items if cache is full 102 | while len(self._cache) >= self.max_size: 103 | oldest_key = min(self._cache.keys(), 104 | key=lambda k: self._cache[k].expires_at) 105 | del self._cache[oldest_key] 106 | 107 | self._cache[key] = CacheEntry( 108 | data=value, 109 | expires_at=datetime.now() + timedelta(seconds=self.ttl) 110 | ) 111 | 112 | 113 | class RateLimiter: 114 | """ 115 | Rate limiter implementation with per-user and global limits. 116 | """ 117 | def __init__(self, config: RateLimitConfig): 118 | """Initialize rate limiter with config.""" 119 | self.config = config 120 | self._user_requests: Dict[str, List[datetime]] = {} 121 | self._global_requests: List[datetime] = [] 122 | self._embedding_requests: List[datetime] = [] 123 | self._llm_requests: List[datetime] = [] 124 | logger.info("Rate limiter initialized with config: {}", config) 125 | 126 | def _clean_old_requests(self, requests: List[datetime]): 127 | """Remove requests older than 1 minute.""" 128 | cutoff = datetime.now() - timedelta(minutes=1) 129 | while requests and requests[0] < cutoff: 130 | requests.pop(0) 131 | 132 | async def check_rate_limit(self, 133 | user_id: Optional[str] = None, 134 | limit_type: str = "global") -> bool: 135 | """ 136 | Check if request is within rate limits. 137 | 138 | Args: 139 | user_id: Optional user ID for per-user limits 140 | limit_type: Type of limit to check (global/embedding/llm) 141 | 142 | Returns: 143 | True if request is allowed, False if rate limited 144 | """ 145 | now = datetime.now() 146 | 147 | # Check user limit if user_id provided 148 | if user_id: 149 | if user_id not in self._user_requests: 150 | self._user_requests[user_id] = [] 151 | 152 | user_requests = self._user_requests[user_id] 153 | self._clean_old_requests(user_requests) 154 | 155 | if len(user_requests) >= self.config.user_limit: 156 | logger.warning("User {} rate limited", user_id) 157 | return False 158 | 159 | user_requests.append(now) 160 | 161 | # Check type-specific limit 162 | if limit_type == "embedding": 163 | self._clean_old_requests(self._embedding_requests) 164 | if len(self._embedding_requests) >= self.config.embedding_limit: 165 | logger.warning("Embedding rate limit reached") 166 | return False 167 | self._embedding_requests.append(now) 168 | 169 | elif limit_type == "llm": 170 | self._clean_old_requests(self._llm_requests) 171 | if len(self._llm_requests) >= self.config.llm_limit: 172 | logger.warning("LLM rate limit reached") 173 | return False 174 | self._llm_requests.append(now) 175 | 176 | # Check global limit 177 | self._clean_old_requests(self._global_requests) 178 | if len(self._global_requests) >= self.config.global_limit: 179 | logger.warning("Global rate limit reached") 180 | return False 181 | 182 | self._global_requests.append(now) 183 | return True 184 | 185 | 186 | class CacheManager: 187 | """ 188 | Manages caching and rate limiting for the platform. 189 | """ 190 | def __init__(self, 191 | cache_config: Optional[CacheConfig] = None, 192 | rate_limit_config: Optional[RateLimitConfig] = None): 193 | """Initialize cache manager.""" 194 | self.cache_config = cache_config or CacheConfig() 195 | self.rate_limit_config = rate_limit_config or RateLimitConfig() 196 | 197 | # Initialize caches 198 | self.embedding_cache = Cache( 199 | max_size=self.cache_config.max_embedding_size, 200 | ttl=self.cache_config.embedding_ttl 201 | ) 202 | self.llm_cache = Cache( 203 | max_size=self.cache_config.max_llm_size, 204 | ttl=self.cache_config.llm_response_ttl 205 | ) 206 | 207 | # Initialize rate limiter 208 | self.rate_limiter = RateLimiter(self.rate_limit_config) 209 | 210 | logger.info("Cache manager initialized") 211 | 212 | def cache_embedding(self, text: str, embedding: np.ndarray) -> None: 213 | """Cache an embedding vector.""" 214 | key = self.embedding_cache._generate_key(text) 215 | self.embedding_cache.set(key, embedding) 216 | 217 | def get_cached_embedding(self, text: str) -> Optional[np.ndarray]: 218 | """Get cached embedding if available.""" 219 | key = self.embedding_cache._generate_key(text) 220 | return self.embedding_cache.get(key) 221 | 222 | def cache_llm_response(self, prompt: str, response: str) -> None: 223 | """Cache an LLM response.""" 224 | key = self.llm_cache._generate_key(prompt) 225 | self.llm_cache.set(key, response) 226 | 227 | def get_cached_llm_response(self, prompt: str) -> Optional[str]: 228 | """Get cached LLM response if available.""" 229 | key = self.llm_cache._generate_key(prompt) 230 | return self.llm_cache.get(key) 231 | 232 | async def check_rate_limit(self, 233 | user_id: Optional[str] = None, 234 | limit_type: str = "global") -> bool: 235 | """Check if request is within rate limits.""" 236 | return await self.rate_limiter.check_rate_limit(user_id, limit_type) 237 | 238 | 239 | # Decorator for caching embeddings 240 | def cache_embedding(cache_manager: CacheManager): 241 | """Decorator to cache embedding results.""" 242 | def decorator(func): 243 | @wraps(func) 244 | async def wrapper(text: str, *args, **kwargs): 245 | # Check cache first 246 | cached = cache_manager.get_cached_embedding(text) 247 | if cached is not None: 248 | logger.debug("Using cached embedding for: {}", text[:50]) 249 | return cached 250 | 251 | # Check rate limit 252 | if not await cache_manager.check_rate_limit(limit_type="embedding"): 253 | logger.warning("Rate limit reached for embedding generation") 254 | raise RuntimeError("Rate limit exceeded for embedding generation") 255 | 256 | # Generate and cache embedding 257 | embedding = await func(text, *args, **kwargs) 258 | cache_manager.cache_embedding(text, embedding) 259 | return embedding 260 | return wrapper 261 | return decorator 262 | 263 | 264 | # Decorator for caching LLM responses 265 | def cache_llm_response(cache_manager: CacheManager): 266 | """Decorator to cache LLM responses.""" 267 | def decorator(func): 268 | @wraps(func) 269 | async def wrapper(prompt: str, *args, **kwargs): 270 | # Check cache first 271 | cached = cache_manager.get_cached_llm_response(prompt) 272 | if cached is not None: 273 | logger.debug("Using cached LLM response for: {}", prompt[:50]) 274 | return cached 275 | 276 | # Check rate limit 277 | if not await cache_manager.check_rate_limit(limit_type="llm"): 278 | logger.warning("Rate limit reached for LLM calls") 279 | raise RuntimeError("Rate limit exceeded for LLM calls") 280 | 281 | # Generate and cache response 282 | response = await func(prompt, *args, **kwargs) 283 | cache_manager.cache_llm_response(prompt, response) 284 | return response 285 | return wrapper 286 | return decorator -------------------------------------------------------------------------------- /nexus/workers.py: -------------------------------------------------------------------------------- 1 | """ 2 | Workers module for handling distributed task processing. 3 | """ 4 | from typing import Any, Dict, List, Optional, Union, Callable 5 | from pydantic import BaseModel, Field 6 | from loguru import logger 7 | import asyncio 8 | from datetime import datetime 9 | import uuid 10 | from enum import Enum 11 | 12 | 13 | class WorkerType(str, Enum): 14 | """Types of worker agents.""" 15 | TEXT = "text" 16 | AUDIO = "audio" 17 | IMAGE = "image" 18 | VIDEO = "video" 19 | GENERAL = "general" 20 | 21 | 22 | class WorkerStatus(str, Enum): 23 | """Status of a worker agent.""" 24 | IDLE = "idle" 25 | BUSY = "busy" 26 | ERROR = "error" 27 | 28 | 29 | class Task(BaseModel): 30 | """A task to be processed by a worker.""" 31 | id: str = Field(default_factory=lambda: str(uuid.uuid4())) 32 | type: WorkerType 33 | content: Dict[str, Any] # Task-specific data 34 | priority: int = 0 # Higher number = higher priority 35 | dependencies: List[str] = [] # IDs of tasks that must complete first 36 | metadata: Dict[str, Any] = {} 37 | created_at: datetime = Field(default_factory=datetime.now) 38 | started_at: Optional[datetime] = None 39 | completed_at: Optional[datetime] = None 40 | error: Optional[str] = None 41 | 42 | 43 | class WorkerAgent(BaseModel): 44 | """A worker agent that processes tasks.""" 45 | id: str = Field(default_factory=lambda: str(uuid.uuid4())) 46 | type: WorkerType 47 | status: WorkerStatus = WorkerStatus.IDLE 48 | current_task: Optional[str] = None # ID of current task 49 | capabilities: List[str] = [] # Specific operations this worker can handle 50 | metadata: Dict[str, Any] = {} 51 | 52 | 53 | class WorkerPool: 54 | """ 55 | Manages a pool of worker agents and task distribution. 56 | """ 57 | 58 | def __init__(self): 59 | """Initialize an empty worker pool.""" 60 | self.workers: Dict[str, WorkerAgent] = {} 61 | self.tasks: Dict[str, Task] = {} 62 | self.task_queue: List[str] = [] # IDs of pending tasks 63 | self.results: Dict[str, Any] = {} # Task results 64 | logger.info("Initialized worker pool") 65 | 66 | def add_worker(self, 67 | type: WorkerType, 68 | capabilities: Optional[List[str]] = None) -> str: 69 | """ 70 | Add a new worker to the pool. 71 | 72 | Args: 73 | type: Type of worker 74 | capabilities: List of specific operations this worker can handle 75 | 76 | Returns: 77 | ID of the new worker 78 | """ 79 | worker = WorkerAgent( 80 | type=type, 81 | capabilities=capabilities or [] 82 | ) 83 | self.workers[worker.id] = worker 84 | logger.info("Added {} worker {} with capabilities: {}", 85 | type, worker.id, capabilities) 86 | return worker.id 87 | 88 | def add_task(self, 89 | type: WorkerType, 90 | content: Dict[str, Any], 91 | priority: int = 0, 92 | dependencies: Optional[List[str]] = None) -> str: 93 | """ 94 | Add a new task to the pool. 95 | 96 | Args: 97 | type: Type of task 98 | content: Task-specific data 99 | priority: Task priority (higher = more important) 100 | dependencies: IDs of tasks that must complete first 101 | 102 | Returns: 103 | ID of the new task 104 | """ 105 | task = Task( 106 | type=type, 107 | content=content, 108 | priority=priority, 109 | dependencies=dependencies or [] 110 | ) 111 | self.tasks[task.id] = task 112 | self._update_queue() 113 | logger.info("Added {} task {} with priority {}", 114 | type, task.id, priority) 115 | return task.id 116 | 117 | def _update_queue(self): 118 | """Update the task queue based on priorities and dependencies.""" 119 | # Get all pending tasks 120 | pending = [ 121 | task for task in self.tasks.values() 122 | if not task.completed_at and not task.started_at 123 | ] 124 | 125 | # Filter out tasks with incomplete dependencies 126 | ready = [ 127 | task for task in pending 128 | if all(dep in self.results for dep in task.dependencies) 129 | ] 130 | 131 | # Sort by priority (descending) and creation time 132 | sorted_tasks = sorted( 133 | ready, 134 | key=lambda t: (-t.priority, t.created_at) 135 | ) 136 | 137 | # Update queue with task IDs 138 | self.task_queue = [task.id for task in sorted_tasks] 139 | 140 | def get_next_task(self, worker_id: str) -> Optional[Task]: 141 | """ 142 | Get the next task for a worker to process. 143 | 144 | Args: 145 | worker_id: ID of the worker requesting a task 146 | 147 | Returns: 148 | Next task to process, or None if no suitable tasks 149 | """ 150 | if worker_id not in self.workers: 151 | raise ValueError(f"Unknown worker: {worker_id}") 152 | 153 | worker = self.workers[worker_id] 154 | if worker.status != WorkerStatus.IDLE: 155 | return None 156 | 157 | # Find the first task matching worker type and capabilities 158 | for task_id in self.task_queue: 159 | task = self.tasks[task_id] 160 | if task.type == worker.type: 161 | return task 162 | 163 | return None 164 | 165 | async def process_task(self, 166 | worker_id: str, 167 | task_id: str, 168 | processor: Callable) -> Any: 169 | """ 170 | Process a task using the specified worker. 171 | 172 | Args: 173 | worker_id: ID of worker processing the task 174 | task_id: ID of task to process 175 | processor: Async function to process the task 176 | 177 | Returns: 178 | Task result 179 | """ 180 | if worker_id not in self.workers: 181 | raise ValueError(f"Unknown worker: {worker_id}") 182 | if task_id not in self.tasks: 183 | raise ValueError(f"Unknown task: {task_id}") 184 | 185 | worker = self.workers[worker_id] 186 | task = self.tasks[task_id] 187 | 188 | try: 189 | # Update status 190 | worker.status = WorkerStatus.BUSY 191 | worker.current_task = task_id 192 | task.started_at = datetime.now() 193 | 194 | # Process task 195 | logger.info("Worker {} processing task {}", worker_id, task_id) 196 | result = await processor(task.content) 197 | 198 | # Store result 199 | self.results[task_id] = result 200 | task.completed_at = datetime.now() 201 | 202 | # Update queue 203 | self._update_queue() 204 | 205 | return result 206 | 207 | except Exception as e: 208 | logger.exception("Error processing task {}", task_id) 209 | task.error = str(e) 210 | worker.status = WorkerStatus.ERROR 211 | raise 212 | 213 | finally: 214 | worker.status = WorkerStatus.IDLE 215 | worker.current_task = None 216 | 217 | 218 | class WorkerOrchestrator: 219 | """ 220 | Orchestrates task distribution and parallel processing across workers. 221 | """ 222 | 223 | def __init__(self): 224 | """Initialize the orchestrator.""" 225 | self.pool = WorkerPool() 226 | logger.info("Initialized worker orchestrator") 227 | 228 | def create_worker_group(self, 229 | type: WorkerType, 230 | count: int, 231 | capabilities: Optional[List[str]] = None) -> List[str]: 232 | """ 233 | Create a group of workers with the same type and capabilities. 234 | 235 | Args: 236 | type: Type of workers to create 237 | count: Number of workers to create 238 | capabilities: List of capabilities for the workers 239 | 240 | Returns: 241 | List of worker IDs 242 | """ 243 | worker_ids = [] 244 | for _ in range(count): 245 | worker_id = self.pool.add_worker(type, capabilities) 246 | worker_ids.append(worker_id) 247 | return worker_ids 248 | 249 | async def process_parallel(self, 250 | tasks: List[Dict[str, Any]], 251 | worker_type: WorkerType, 252 | processor: Callable, 253 | num_workers: int = 3) -> Dict[str, Any]: 254 | """ 255 | Process multiple tasks in parallel using a pool of workers. 256 | 257 | Args: 258 | tasks: List of task content dictionaries 259 | worker_type: Type of workers to use 260 | processor: Async function to process each task 261 | num_workers: Number of workers to create 262 | 263 | Returns: 264 | Dictionary mapping task IDs to results 265 | """ 266 | # Create workers if needed 267 | if not any(w.type == worker_type for w in self.pool.workers.values()): 268 | self.create_worker_group(worker_type, num_workers) 269 | 270 | # Add tasks to pool 271 | task_ids = [] 272 | for task_content in tasks: 273 | task_id = self.pool.add_task( 274 | type=worker_type, 275 | content=task_content 276 | ) 277 | task_ids.append(task_id) 278 | 279 | # Process tasks in parallel 280 | async def worker_loop(worker_id: str): 281 | while True: 282 | task = self.pool.get_next_task(worker_id) 283 | if not task: 284 | # No more tasks for this worker 285 | break 286 | 287 | await self.pool.process_task( 288 | worker_id, 289 | task.id, 290 | processor 291 | ) 292 | 293 | # Start worker tasks 294 | worker_tasks = [] 295 | for worker_id in self.pool.workers: 296 | if self.pool.workers[worker_id].type == worker_type: 297 | worker_tasks.append( 298 | asyncio.create_task(worker_loop(worker_id)) 299 | ) 300 | 301 | # Wait for all tasks to complete 302 | await asyncio.gather(*worker_tasks) 303 | 304 | # Return results for requested tasks 305 | return { 306 | task_id: self.pool.results[task_id] 307 | for task_id in task_ids 308 | if task_id in self.pool.results 309 | } 310 | 311 | async def process_sequential(self, 312 | tasks: List[Dict[str, Any]], 313 | worker_type: WorkerType, 314 | processor: Callable) -> List[Any]: 315 | """ 316 | Process tasks sequentially using a single worker. 317 | 318 | Args: 319 | tasks: List of task content dictionaries 320 | worker_type: Type of worker to use 321 | processor: Async function to process each task 322 | 323 | Returns: 324 | List of task results in order 325 | """ 326 | # Create a worker 327 | worker_id = self.pool.add_worker(worker_type) 328 | 329 | results = [] 330 | for task_content in tasks: 331 | # Add and process task 332 | task_id = self.pool.add_task( 333 | type=worker_type, 334 | content=task_content 335 | ) 336 | 337 | result = await self.pool.process_task( 338 | worker_id, 339 | task_id, 340 | processor 341 | ) 342 | results.append(result) 343 | 344 | return results -------------------------------------------------------------------------------- /nexus/agents/audio_agent.py: -------------------------------------------------------------------------------- 1 | """ 2 | Audio processing agent with speech-to-text and text-to-speech capabilities. 3 | """ 4 | from typing import Optional, Union, BinaryIO, List 5 | from pathlib import Path 6 | from faster_whisper import WhisperModel 7 | from TTS.api import TTS 8 | import soundfile as sf 9 | import numpy as np 10 | from loguru import logger 11 | from pydantic import BaseModel 12 | import tempfile 13 | import os 14 | import noisereduce as nr 15 | 16 | 17 | class AudioConfig(BaseModel): 18 | """Configuration for audio processing.""" 19 | # STT Configuration 20 | model_type: str = "base" # Whisper model type: tiny, base, small, medium, large 21 | device: str = "cpu" # Device to run inference on: cpu or cuda 22 | compute_type: str = "int8" # Compute type: int8, int8_float16, float16, float32 23 | language: Optional[str] = None # Language code (e.g., 'en' for English) 24 | 25 | # TTS Configuration 26 | tts_model: str = "tts_models/en/ljspeech/tacotron2-DDC" # Default English TTS model 27 | tts_language: str = "en" 28 | tts_voice: Optional[str] = None 29 | 30 | # Audio Processing 31 | sample_rate: int = 16000 32 | noise_reduction: bool = False 33 | volume_normalize: bool = True 34 | 35 | 36 | class AudioAgent: 37 | """ 38 | Agent for processing audio inputs and outputs. 39 | Handles speech-to-text conversion and text-to-speech synthesis. 40 | """ 41 | 42 | def __init__(self, config: Optional[AudioConfig] = None): 43 | """Initialize the audio agent with optional configuration.""" 44 | self.config = config or AudioConfig() 45 | 46 | # Initialize STT (Whisper) 47 | with logger.contextualize(model_type=self.config.model_type): 48 | logger.info("Initializing Whisper model") 49 | self.stt_model = WhisperModel( 50 | model_size_or_path=self.config.model_type, 51 | device=self.config.device, 52 | compute_type=self.config.compute_type 53 | ) 54 | logger.info("Whisper model loaded successfully") 55 | 56 | # Initialize TTS 57 | with logger.contextualize(tts_model=self.config.tts_model): 58 | logger.info("Initializing TTS model") 59 | self.tts_model = TTS(self.config.tts_model) 60 | logger.info("TTS model loaded successfully") 61 | 62 | def _preprocess_audio(self, audio_data: np.ndarray, sample_rate: int) -> np.ndarray: 63 | """ 64 | Apply preprocessing steps to the audio data. 65 | """ 66 | try: 67 | # Ensure correct sample rate 68 | if sample_rate != self.config.sample_rate: 69 | # TODO: Add resampling logic if needed 70 | pass 71 | 72 | # Apply noise reduction if enabled 73 | if self.config.noise_reduction: 74 | logger.debug("Applying noise reduction") 75 | audio_data = nr.reduce_noise(y=audio_data, sr=sample_rate) 76 | 77 | # Normalize volume if enabled 78 | if self.config.volume_normalize: 79 | logger.debug("Normalizing volume") 80 | audio_data = audio_data / np.max(np.abs(audio_data)) 81 | 82 | return audio_data 83 | 84 | except Exception as e: 85 | logger.error("Error in audio preprocessing: {}", str(e)) 86 | raise ValueError(f"Failed to preprocess audio: {str(e)}") 87 | 88 | def _load_audio(self, audio_input: Union[str, Path, BinaryIO, np.ndarray]) -> str: 89 | """ 90 | Load audio from various input types and return a path to the audio file. 91 | Creates a temporary file if needed. 92 | """ 93 | try: 94 | if isinstance(audio_input, (str, Path)): 95 | # Load, preprocess, and save back if preprocessing is enabled 96 | if self.config.noise_reduction or self.config.volume_normalize: 97 | audio_data, sample_rate = sf.read(str(audio_input)) 98 | audio_data = self._preprocess_audio(audio_data, sample_rate) 99 | with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file: 100 | sf.write(temp_file.name, audio_data, sample_rate) 101 | return temp_file.name 102 | return str(audio_input) 103 | elif isinstance(audio_input, np.ndarray): 104 | # Preprocess and save numpy array to temporary file 105 | audio_data = self._preprocess_audio(audio_input, self.config.sample_rate) 106 | with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file: 107 | sf.write(temp_file.name, audio_data, self.config.sample_rate) 108 | return temp_file.name 109 | else: 110 | # Save file-like object to temporary file 111 | with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file: 112 | temp_file.write(audio_input.read()) 113 | # Load, preprocess, and save back if preprocessing is enabled 114 | if self.config.noise_reduction or self.config.volume_normalize: 115 | audio_data, sample_rate = sf.read(temp_file.name) 116 | audio_data = self._preprocess_audio(audio_data, sample_rate) 117 | sf.write(temp_file.name, audio_data, sample_rate) 118 | return temp_file.name 119 | 120 | except Exception as e: 121 | logger.error("Error loading audio: {}", str(e)) 122 | raise ValueError(f"Failed to load audio: {str(e)}") 123 | 124 | def synthesize(self, 125 | text: str, 126 | output_path: Optional[str] = None, 127 | language: Optional[str] = None, 128 | voice: Optional[str] = None) -> Optional[str]: 129 | """ 130 | Convert text to speech using the TTS model. 131 | 132 | Args: 133 | text: Text to convert to speech 134 | output_path: Optional path to save the audio file 135 | language: Optional language override 136 | voice: Optional voice override 137 | 138 | Returns: 139 | Path to the generated audio file if output_path is provided, 140 | otherwise plays the audio and returns None 141 | """ 142 | try: 143 | with logger.contextualize(operation="synthesize"): 144 | logger.info("Starting text-to-speech synthesis") 145 | 146 | # Use provided values or fall back to config 147 | language = language or self.config.tts_language 148 | voice = voice or self.config.tts_voice 149 | 150 | # Generate temporary path if none provided 151 | if output_path is None: 152 | with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file: 153 | output_path = temp_file.name 154 | 155 | # Synthesize speech 156 | logger.debug("Running TTS inference") 157 | self.tts_model.tts_to_file( 158 | text=text, 159 | file_path=output_path, 160 | language=language, 161 | speaker=voice 162 | ) 163 | 164 | logger.info("Speech synthesis completed successfully") 165 | return output_path 166 | 167 | except Exception as e: 168 | logger.exception("Speech synthesis failed") 169 | raise RuntimeError(f"Speech synthesis failed: {str(e)}") 170 | 171 | def list_available_voices(self) -> List[str]: 172 | """List available voices for the current TTS model.""" 173 | try: 174 | voices = self.tts_model.speakers 175 | return voices if voices else [] 176 | except Exception: 177 | return [] 178 | 179 | def list_available_languages(self) -> List[str]: 180 | """List available languages for the current TTS model.""" 181 | try: 182 | languages = self.tts_model.languages 183 | return languages if languages else [] 184 | except Exception: 185 | return [] 186 | 187 | def transcribe(self, 188 | audio_input: Union[str, Path, BinaryIO, np.ndarray], 189 | language: Optional[str] = None) -> dict: 190 | """ 191 | Transcribe speech from audio input to text. 192 | 193 | Args: 194 | audio_input: Audio file path, binary data, or numpy array 195 | language: Optional language code to override config 196 | 197 | Returns: 198 | Dictionary containing: 199 | - text: The transcribed text 200 | - language: Detected or specified language 201 | - segments: List of transcribed segments with timestamps 202 | """ 203 | try: 204 | with logger.contextualize(operation="transcribe"): 205 | logger.info("Starting transcription") 206 | 207 | # Load and preprocess audio 208 | audio_path = self._load_audio(audio_input) 209 | temp_file_created = not isinstance(audio_input, (str, Path)) 210 | 211 | try: 212 | # Run transcription 213 | logger.debug("Running Whisper inference") 214 | segments, info = self.stt_model.transcribe( 215 | audio_path, 216 | language=language or self.config.language, 217 | task="transcribe" 218 | ) 219 | 220 | # Convert segments to list and join text 221 | segments_list = list(segments) # Convert generator to list 222 | full_text = " ".join(seg.text for seg in segments_list) 223 | 224 | # Format segments for output 225 | formatted_segments = [ 226 | { 227 | "text": seg.text, 228 | "start": seg.start, 229 | "end": seg.end, 230 | "words": [{"word": w.word, "probability": w.probability} 231 | for w in (seg.words or [])] 232 | } 233 | for seg in segments_list 234 | ] 235 | 236 | logger.info("Transcription completed successfully") 237 | 238 | return { 239 | "text": full_text, 240 | "language": info.language, 241 | "segments": formatted_segments, 242 | "language_probability": info.language_probability 243 | } 244 | 245 | finally: 246 | # Clean up temporary file if we created one 247 | if temp_file_created and os.path.exists(audio_path): 248 | os.unlink(audio_path) 249 | 250 | except Exception as e: 251 | logger.exception("Transcription failed") 252 | raise RuntimeError(f"Transcription failed: {str(e)}") 253 | 254 | def detect_language(self, audio_input: Union[str, Path, BinaryIO, np.ndarray]) -> str: 255 | """ 256 | Detect the language of speech in the audio. 257 | 258 | Args: 259 | audio_input: Audio file path, binary data, or numpy array 260 | 261 | Returns: 262 | Detected language code 263 | """ 264 | try: 265 | with logger.contextualize(operation="detect_language"): 266 | logger.info("Starting language detection") 267 | 268 | # Load and preprocess audio 269 | audio_path = self._load_audio(audio_input) 270 | temp_file_created = not isinstance(audio_input, (str, Path)) 271 | 272 | try: 273 | # Run language detection 274 | _, info = self.stt_model.transcribe( 275 | audio_path, 276 | task="transcribe", 277 | language=None # Force language detection 278 | ) 279 | 280 | detected_lang = info.language 281 | logger.info("Detected language: {} (probability: {:.2f})", 282 | detected_lang, info.language_probability) 283 | 284 | return detected_lang 285 | 286 | finally: 287 | # Clean up temporary file if we created one 288 | if temp_file_created and os.path.exists(audio_path): 289 | os.unlink(audio_path) 290 | 291 | except Exception as e: 292 | logger.exception("Language detection failed") 293 | raise RuntimeError(f"Language detection failed: {str(e)}") -------------------------------------------------------------------------------- /nexus/chains.py: -------------------------------------------------------------------------------- 1 | """ 2 | Prompt chaining module for defining and executing multi-step workflows. 3 | """ 4 | from typing import Any, Dict, List, Optional, Union, Callable, Set 5 | from pydantic import BaseModel, Field 6 | from loguru import logger 7 | import asyncio 8 | from datetime import datetime 9 | from enum import Enum 10 | import uuid 11 | 12 | 13 | class ChainNodeType(str, Enum): 14 | """Types of chain nodes.""" 15 | PROMPT = "prompt" # LLM prompt node 16 | BRANCH = "branch" # Branching decision node 17 | MERGE = "merge" # Merge multiple paths 18 | TOOL = "tool" # External tool/function call 19 | 20 | 21 | class ChainNode(BaseModel): 22 | """A single node in the prompt chain.""" 23 | id: str = Field(default_factory=lambda: str(uuid.uuid4())) 24 | type: ChainNodeType 25 | name: str 26 | content: Union[str, Dict[str, Any]] # Prompt text or tool config 27 | next_nodes: List[str] = [] # IDs of next nodes 28 | prev_nodes: List[str] = [] # IDs of previous nodes 29 | requires_all_prev: bool = False # For merge nodes: require all prev nodes to complete 30 | metadata: Dict[str, Any] = {} 31 | 32 | 33 | class ChainContext(BaseModel): 34 | """Context passed between chain nodes during execution.""" 35 | variables: Dict[str, Any] = {} 36 | results: Dict[str, Any] = {} 37 | current_path: List[str] = [] # Node IDs in current execution path 38 | completed_nodes: Set[str] = set() # IDs of completed nodes 39 | 40 | 41 | class ChainValidationError(Exception): 42 | """Raised when chain validation fails.""" 43 | pass 44 | 45 | 46 | class PromptChain: 47 | """ 48 | Manages a chain of prompts and tools with support for branching. 49 | """ 50 | 51 | def __init__(self, name: str): 52 | """Initialize an empty chain.""" 53 | self.name = name 54 | self.nodes: Dict[str, ChainNode] = {} 55 | self.start_nodes: List[str] = [] # Nodes with no predecessors 56 | self.end_nodes: List[str] = [] # Nodes with no successors 57 | logger.info("Created new prompt chain: {}", name) 58 | 59 | def add_node(self, 60 | node_type: ChainNodeType, 61 | name: str, 62 | content: Union[str, Dict[str, Any]], 63 | requires_all_prev: bool = False) -> str: 64 | """ 65 | Add a new node to the chain. 66 | 67 | Args: 68 | node_type: Type of node 69 | name: Node name 70 | content: Node content (prompt text or tool config) 71 | requires_all_prev: For merge nodes, require all prev nodes 72 | 73 | Returns: 74 | ID of the new node 75 | """ 76 | node = ChainNode( 77 | type=node_type, 78 | name=name, 79 | content=content, 80 | requires_all_prev=requires_all_prev 81 | ) 82 | 83 | self.nodes[node.id] = node 84 | logger.debug("Added {} node '{}' ({})", node_type, name, node.id) 85 | 86 | # Update start/end nodes 87 | self._update_topology() 88 | return node.id 89 | 90 | def connect(self, from_id: str, to_id: str): 91 | """Connect two nodes in the chain.""" 92 | if from_id not in self.nodes or to_id not in self.nodes: 93 | raise ValueError("Invalid node IDs") 94 | 95 | from_node = self.nodes[from_id] 96 | to_node = self.nodes[to_id] 97 | 98 | if to_id not in from_node.next_nodes: 99 | from_node.next_nodes.append(to_id) 100 | if from_id not in to_node.prev_nodes: 101 | to_node.prev_nodes.append(from_id) 102 | 103 | self._update_topology() 104 | logger.debug("Connected node {} to {}", from_id, to_id) 105 | 106 | def _update_topology(self): 107 | """Update start and end node lists.""" 108 | self.start_nodes = [ 109 | node.id for node in self.nodes.values() 110 | if not node.prev_nodes 111 | ] 112 | self.end_nodes = [ 113 | node.id for node in self.nodes.values() 114 | if not node.next_nodes 115 | ] 116 | 117 | async def validate_node(self, 118 | node: ChainNode, 119 | context: ChainContext) -> bool: 120 | """ 121 | Validate a node before execution. 122 | 123 | Args: 124 | node: Node to validate 125 | context: Current execution context 126 | 127 | Returns: 128 | True if validation passes 129 | """ 130 | try: 131 | # Check if node can be executed 132 | if node.requires_all_prev: 133 | # All previous nodes must be completed 134 | if not all(n in context.completed_nodes for n in node.prev_nodes): 135 | logger.warning( 136 | "Node {} requires all prev nodes to complete", 137 | node.id 138 | ) 139 | return False 140 | else: 141 | # At least one previous node must be completed 142 | if node.prev_nodes and not any( 143 | n in context.completed_nodes for n in node.prev_nodes 144 | ): 145 | logger.warning( 146 | "Node {} requires at least one prev node", 147 | node.id 148 | ) 149 | return False 150 | 151 | # Validate based on node type 152 | if node.type == ChainNodeType.PROMPT: 153 | # Check if prompt can be formatted with context 154 | try: 155 | node.content.format(**context.variables) 156 | except KeyError as e: 157 | logger.warning( 158 | "Missing variable {} for node {}", 159 | str(e), 160 | node.id 161 | ) 162 | return False 163 | 164 | elif node.type == ChainNodeType.TOOL: 165 | # Check if required tool configs are present 166 | required_configs = {"name", "params"} 167 | if not all(k in node.content for k in required_configs): 168 | logger.warning( 169 | "Missing tool configs for node {}: {}", 170 | node.id, 171 | required_configs - node.content.keys() 172 | ) 173 | return False 174 | 175 | return True 176 | 177 | except Exception as e: 178 | logger.exception("Error validating node {}", node.id) 179 | return False 180 | 181 | async def execute_node(self, 182 | node: ChainNode, 183 | context: ChainContext, 184 | tools: Dict[str, Callable] = None) -> Any: 185 | """ 186 | Execute a single node. 187 | 188 | Args: 189 | node: Node to execute 190 | context: Current execution context 191 | tools: Available tool functions 192 | 193 | Returns: 194 | Node execution result 195 | """ 196 | try: 197 | if node.type == ChainNodeType.PROMPT: 198 | # Format and execute prompt 199 | prompt = node.content.format(**context.variables) 200 | # TODO: Call LLM with prompt 201 | result = f"LLM response to: {prompt}" 202 | 203 | elif node.type == ChainNodeType.BRANCH: 204 | # Evaluate branch condition 205 | condition = node.content.get("condition", "true") 206 | result = eval(condition, {"context": context.variables}) 207 | 208 | elif node.type == ChainNodeType.TOOL: 209 | # Execute tool function 210 | tool_name = node.content["name"] 211 | if tool_name not in tools: 212 | raise ValueError(f"Tool {tool_name} not found") 213 | 214 | tool_func = tools[tool_name] 215 | result = await tool_func(**node.content["params"]) 216 | 217 | else: # MERGE 218 | # Combine results from previous nodes 219 | merge_strategy = node.content.get("strategy", "last") 220 | prev_results = [ 221 | context.results[n] 222 | for n in node.prev_nodes 223 | if n in context.completed_nodes 224 | ] 225 | 226 | if merge_strategy == "last": 227 | result = prev_results[-1] 228 | elif merge_strategy == "list": 229 | result = prev_results 230 | else: 231 | raise ValueError(f"Unknown merge strategy: {merge_strategy}") 232 | 233 | return result 234 | 235 | except Exception as e: 236 | logger.exception("Error executing node {}", node.id) 237 | raise 238 | 239 | async def execute(self, 240 | initial_context: Optional[Dict[str, Any]] = None, 241 | tools: Optional[Dict[str, Callable]] = None) -> Dict[str, Any]: 242 | """ 243 | Execute the entire chain. 244 | 245 | Args: 246 | initial_context: Initial variables 247 | tools: Available tool functions 248 | 249 | Returns: 250 | Dictionary of execution results 251 | """ 252 | context = ChainContext( 253 | variables=initial_context or {}, 254 | results={}, 255 | current_path=[], 256 | completed_nodes=set() 257 | ) 258 | 259 | # Start with all start nodes 260 | next_nodes = self.start_nodes.copy() 261 | 262 | while next_nodes: 263 | current_id = next_nodes.pop(0) 264 | current_node = self.nodes[current_id] 265 | 266 | # Validate node 267 | if not await self.validate_node(current_node, context): 268 | logger.warning("Validation failed for node {}", current_id) 269 | continue 270 | 271 | # Execute node 272 | try: 273 | result = await self.execute_node( 274 | current_node, 275 | context, 276 | tools 277 | ) 278 | 279 | # Update context 280 | context.results[current_id] = result 281 | context.completed_nodes.add(current_id) 282 | context.current_path.append(current_id) 283 | 284 | # Update variables if node provides output 285 | if isinstance(result, dict): 286 | context.variables.update(result) 287 | 288 | # Add next nodes to queue 289 | for next_id in current_node.next_nodes: 290 | if next_id not in next_nodes: 291 | next_nodes.append(next_id) 292 | 293 | except Exception as e: 294 | logger.exception("Error in node {}", current_id) 295 | raise 296 | 297 | return context.results 298 | 299 | 300 | # Helper functions to create common chain patterns 301 | def create_linear_chain( 302 | name: str, 303 | prompts: List[str], 304 | tools: Optional[List[Dict[str, Any]]] = None 305 | ) -> PromptChain: 306 | """Create a linear chain of prompts and tools.""" 307 | chain = PromptChain(name) 308 | prev_id = None 309 | 310 | # Add prompts 311 | for i, prompt in enumerate(prompts): 312 | node_id = chain.add_node( 313 | ChainNodeType.PROMPT, 314 | f"prompt_{i}", 315 | prompt 316 | ) 317 | if prev_id: 318 | chain.connect(prev_id, node_id) 319 | prev_id = node_id 320 | 321 | # Add tools 322 | if tools: 323 | for i, tool in enumerate(tools): 324 | node_id = chain.add_node( 325 | ChainNodeType.TOOL, 326 | f"tool_{i}", 327 | tool 328 | ) 329 | if prev_id: 330 | chain.connect(prev_id, node_id) 331 | prev_id = node_id 332 | 333 | return chain 334 | 335 | 336 | def create_branching_chain( 337 | name: str, 338 | condition_prompt: str, 339 | true_branch: List[str], 340 | false_branch: List[str] 341 | ) -> PromptChain: 342 | """Create a chain with conditional branching.""" 343 | chain = PromptChain(name) 344 | 345 | # Add condition node 346 | condition_id = chain.add_node( 347 | ChainNodeType.PROMPT, 348 | "condition", 349 | condition_prompt 350 | ) 351 | 352 | # Add branch node 353 | branch_id = chain.add_node( 354 | ChainNodeType.BRANCH, 355 | "branch", 356 | {"condition": "context.get('condition', False)"} 357 | ) 358 | chain.connect(condition_id, branch_id) 359 | 360 | # Add true branch 361 | prev_id = branch_id 362 | for i, prompt in enumerate(true_branch): 363 | node_id = chain.add_node( 364 | ChainNodeType.PROMPT, 365 | f"true_{i}", 366 | prompt 367 | ) 368 | chain.connect(prev_id, node_id) 369 | prev_id = node_id 370 | true_end_id = prev_id 371 | 372 | # Add false branch 373 | prev_id = branch_id 374 | for i, prompt in enumerate(false_branch): 375 | node_id = chain.add_node( 376 | ChainNodeType.PROMPT, 377 | f"false_{i}", 378 | prompt 379 | ) 380 | chain.connect(prev_id, node_id) 381 | prev_id = node_id 382 | false_end_id = prev_id 383 | 384 | # Add merge node 385 | merge_id = chain.add_node( 386 | ChainNodeType.MERGE, 387 | "merge", 388 | {"strategy": "last"}, 389 | requires_all_prev=False 390 | ) 391 | chain.connect(true_end_id, merge_id) 392 | chain.connect(false_end_id, merge_id) 393 | 394 | return chain -------------------------------------------------------------------------------- /nexus/orchestrator.py: -------------------------------------------------------------------------------- 1 | """ 2 | Core Orchestrator module for coordinating agent interactions. 3 | """ 4 | from typing import Optional, Any, List, Dict, Coroutine 5 | from collections import deque 6 | from loguru import logger 7 | from pydantic import BaseModel 8 | import aisuite as ai 9 | import os 10 | import sys 11 | from datetime import datetime 12 | import asyncio 13 | from concurrent.futures import ThreadPoolExecutor 14 | from .vector_store import VectorStore, VectorStoreConfig 15 | 16 | 17 | # Configure loguru logger 18 | logger.remove() # Remove default handler 19 | logger.add( 20 | sys.stderr, 21 | format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}", 22 | level="INFO" 23 | ) 24 | logger.add( 25 | "logs/mnemosyne_{time}.log", 26 | rotation="500 MB", 27 | retention="10 days", 28 | format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}", 29 | level="DEBUG" 30 | ) 31 | 32 | 33 | class LLMProviderConfig(BaseModel): 34 | """Configuration for an LLM provider.""" 35 | provider: str # e.g., 'openai', 'anthropic', 'google' 36 | model: str # e.g., 'gpt-4', 'claude-3' 37 | api_key: Optional[str] = None 38 | 39 | 40 | class OrchestratorConfig(BaseModel): 41 | """Configuration for the Orchestrator.""" 42 | debug: bool = False 43 | primary_provider: LLMProviderConfig = LLMProviderConfig( 44 | provider="openai", 45 | model="gpt-4" 46 | ) 47 | fallback_providers: List[LLMProviderConfig] = [] 48 | history_length: int = 10 # Number of conversation turns to remember in memory 49 | vector_store: Optional[VectorStoreConfig] = None # Vector store config for persistent memory 50 | 51 | 52 | class Message(BaseModel): 53 | """A single message in the conversation.""" 54 | role: str # 'user', 'assistant', or 'system' 55 | content: str 56 | timestamp: datetime = datetime.now() 57 | 58 | 59 | class Orchestrator: 60 | """ 61 | Central orchestrator for managing agent workflows and interactions. 62 | Uses aisuite for flexible LLM provider routing. 63 | """ 64 | 65 | def __init__(self, config: Optional[OrchestratorConfig] = None): 66 | """Initialize the orchestrator with optional configuration.""" 67 | self.config = config or OrchestratorConfig() 68 | 69 | # Set up logging based on debug mode 70 | if self.config.debug: 71 | logger.configure(handlers=[{"sink": sys.stderr, "level": "DEBUG"}]) 72 | 73 | # Log initialization with context 74 | with logger.contextualize( 75 | config=self.config.dict(), 76 | session_id=id(self) 77 | ): 78 | logger.info("Orchestrator initialized") 79 | logger.debug("Full configuration: {}", self.config) 80 | 81 | # Initialize aisuite client 82 | self.client = ai.Client() 83 | 84 | # Initialize conversation history 85 | self.conversation_history: deque[Message] = deque(maxlen=self.config.history_length) 86 | 87 | # Initialize thread pool for parallel execution 88 | self._thread_pool = ThreadPoolExecutor() 89 | 90 | # Initialize vector store if configured 91 | self.vector_store = VectorStore(self.config.vector_store) if self.config.vector_store else None 92 | if self.vector_store: 93 | logger.info("Vector store initialized for persistent memory") 94 | 95 | def _get_model_string(self, provider_config: LLMProviderConfig) -> str: 96 | """Convert provider config to aisuite model string format.""" 97 | return f"{provider_config.provider}:{provider_config.model}" 98 | 99 | def _log_llm_request(self, messages: List[Dict[str, str]], provider_config: LLMProviderConfig): 100 | """Log LLM request details in debug mode.""" 101 | if self.config.debug: 102 | with logger.contextualize( 103 | provider=provider_config.provider, 104 | model=provider_config.model, 105 | message_count=len(messages) 106 | ): 107 | logger.debug("LLM Request:") 108 | for idx, msg in enumerate(messages): 109 | logger.debug(f"Message {idx}: {msg}") 110 | 111 | def _call_llm(self, messages: List[Dict[str, str]], provider_config: Optional[LLMProviderConfig] = None) -> str: 112 | """ 113 | Make a call to the LLM and return its response. 114 | Attempts fallback providers if primary fails. 115 | """ 116 | # Use primary provider if none specified 117 | provider_config = provider_config or self.config.primary_provider 118 | 119 | # Log request details in debug mode 120 | self._log_llm_request(messages, provider_config) 121 | 122 | try: 123 | with logger.contextualize( 124 | provider=provider_config.provider, 125 | model=provider_config.model 126 | ): 127 | # Try primary provider 128 | start_time = datetime.now() 129 | response = self.client.chat.completions.create( 130 | model=self._get_model_string(provider_config), 131 | messages=messages 132 | ) 133 | duration = (datetime.now() - start_time).total_seconds() 134 | 135 | logger.info("LLM call successful. Duration: {:.2f}s", duration) 136 | if self.config.debug: 137 | logger.debug("Raw response: {}", response) 138 | 139 | return response.choices[0].message.content 140 | 141 | except Exception as e: 142 | logger.error("Error with provider {}: {}", provider_config.provider, str(e)) 143 | 144 | # Try fallback providers if available 145 | for fallback in self.config.fallback_providers: 146 | try: 147 | with logger.contextualize( 148 | provider=fallback.provider, 149 | model=fallback.model, 150 | is_fallback=True 151 | ): 152 | logger.info("Attempting fallback provider") 153 | start_time = datetime.now() 154 | response = self.client.chat.completions.create( 155 | model=self._get_model_string(fallback), 156 | messages=messages 157 | ) 158 | duration = (datetime.now() - start_time).total_seconds() 159 | 160 | logger.info("Fallback successful. Duration: {:.2f}s", duration) 161 | return response.choices[0].message.content 162 | 163 | except Exception as fallback_error: 164 | logger.error("Fallback provider failed: {}", str(fallback_error)) 165 | 166 | # If all providers fail, raise the original error 167 | raise 168 | 169 | async def run_parallel_tasks(self, 170 | tasks: List[Coroutine], 171 | timeout: Optional[float] = None) -> List[Any]: 172 | """ 173 | Run multiple async tasks in parallel and return their results. 174 | 175 | Args: 176 | tasks: List of coroutines to execute 177 | timeout: Optional timeout in seconds 178 | 179 | Returns: 180 | List of results from completed tasks 181 | """ 182 | logger.debug("Running {} tasks in parallel", len(tasks)) 183 | start_time = datetime.now() 184 | 185 | try: 186 | # Run tasks with optional timeout 187 | if timeout: 188 | results = await asyncio.gather(*tasks, timeout=timeout) 189 | else: 190 | results = await asyncio.gather(*tasks) 191 | 192 | duration = (datetime.now() - start_time).total_seconds() 193 | logger.info("Parallel execution completed in {:.2f}s", duration) 194 | 195 | return results 196 | 197 | except asyncio.TimeoutError: 198 | logger.warning("Parallel execution timed out after {}s", timeout) 199 | raise 200 | except Exception as e: 201 | logger.exception("Error during parallel execution") 202 | raise RuntimeError(f"Parallel execution failed: {str(e)}") 203 | 204 | async def process_input_async(self, user_input: str) -> dict[str, Any]: 205 | """ 206 | Async version of process_input that supports parallel operations. 207 | Maintains conversation history and handles provider routing. 208 | """ 209 | with logger.contextualize( 210 | session_id=id(self), 211 | history_length=len(self.conversation_history) 212 | ): 213 | logger.info("Processing new input (async)") 214 | logger.debug("User input: {}", user_input) 215 | 216 | try: 217 | # Add user message to history 218 | user_message = Message(role="user", content=user_input) 219 | self.conversation_history.append(user_message) 220 | 221 | # Store in vector store if available 222 | if self.vector_store: 223 | await self.vector_store.add_message( 224 | role="user", 225 | content=user_input, 226 | metadata={"session_id": id(self)} 227 | ) 228 | 229 | # Prepare messages for LLM including history 230 | messages = [{"role": msg.role, "content": msg.content} 231 | for msg in self.conversation_history] 232 | 233 | # Get response from LLM (with potential fallbacks) 234 | llm_response = await self._call_llm_async(messages) 235 | 236 | # Add assistant response to history 237 | assistant_message = Message( 238 | role="assistant", 239 | content=llm_response 240 | ) 241 | self.conversation_history.append(assistant_message) 242 | 243 | # Store assistant response in vector store 244 | if self.vector_store: 245 | await self.vector_store.add_message( 246 | role="assistant", 247 | content=llm_response, 248 | metadata={ 249 | "session_id": id(self), 250 | "provider": self.config.primary_provider.provider, 251 | "model": self.config.primary_provider.model 252 | } 253 | ) 254 | 255 | response = { 256 | "status": "success", 257 | "input_received": user_input, 258 | "response": llm_response, 259 | "provider_used": self.config.primary_provider.provider, 260 | "model_used": self.config.primary_provider.model, 261 | "history_length": len(self.conversation_history), 262 | "timestamp": datetime.now().isoformat() 263 | } 264 | 265 | logger.info("Successfully processed input (async)") 266 | if self.config.debug: 267 | logger.debug("Full response: {}", response) 268 | 269 | return response 270 | 271 | except Exception as e: 272 | logger.exception("Error processing input (async)") 273 | return { 274 | "status": "error", 275 | "input_received": user_input, 276 | "error": str(e), 277 | "provider_used": self.config.primary_provider.provider, 278 | "model_used": self.config.primary_provider.model, 279 | "timestamp": datetime.now().isoformat() 280 | } 281 | 282 | async def _call_llm_async(self, messages: List[Dict[str, str]], 283 | provider_config: Optional[LLMProviderConfig] = None) -> str: 284 | """ 285 | Async version of _call_llm. 286 | Makes a call to the LLM and returns its response. 287 | Attempts fallback providers if primary fails. 288 | """ 289 | # Use primary provider if none specified 290 | provider_config = provider_config or self.config.primary_provider 291 | 292 | # Log request details in debug mode 293 | self._log_llm_request(messages, provider_config) 294 | 295 | try: 296 | with logger.contextualize( 297 | provider=provider_config.provider, 298 | model=provider_config.model 299 | ): 300 | # Try primary provider 301 | start_time = datetime.now() 302 | 303 | # Run LLM call in thread pool to avoid blocking 304 | response = await asyncio.get_event_loop().run_in_executor( 305 | self._thread_pool, 306 | lambda: self.client.chat.completions.create( 307 | model=self._get_model_string(provider_config), 308 | messages=messages 309 | ) 310 | ) 311 | 312 | duration = (datetime.now() - start_time).total_seconds() 313 | 314 | logger.info("LLM call successful (async). Duration: {:.2f}s", duration) 315 | if self.config.debug: 316 | logger.debug("Raw response: {}", response) 317 | 318 | return response.choices[0].message.content 319 | 320 | except Exception as e: 321 | logger.error("Error with provider {} (async): {}", provider_config.provider, str(e)) 322 | 323 | # Try fallback providers if available 324 | for fallback in self.config.fallback_providers: 325 | try: 326 | with logger.contextualize( 327 | provider=fallback.provider, 328 | model=fallback.model, 329 | is_fallback=True 330 | ): 331 | logger.info("Attempting fallback provider (async)") 332 | start_time = datetime.now() 333 | 334 | # Run fallback in thread pool 335 | response = await asyncio.get_event_loop().run_in_executor( 336 | self._thread_pool, 337 | lambda: self.client.chat.completions.create( 338 | model=self._get_model_string(fallback), 339 | messages=messages 340 | ) 341 | ) 342 | 343 | duration = (datetime.now() - start_time).total_seconds() 344 | 345 | logger.info("Fallback successful (async). Duration: {:.2f}s", duration) 346 | return response.choices[0].message.content 347 | 348 | except Exception as fallback_error: 349 | logger.error("Fallback provider failed (async): {}", str(fallback_error)) 350 | 351 | # If all providers fail, raise the original error 352 | raise 353 | 354 | async def get_conversation_history(self, 355 | limit: Optional[int] = None, 356 | role_filter: Optional[str] = None) -> List[Dict[str, Any]]: 357 | """ 358 | Get conversation history from vector store if available, otherwise from memory. 359 | 360 | Args: 361 | limit: Maximum number of messages to return 362 | role_filter: Optional filter by role 363 | 364 | Returns: 365 | List of conversation messages 366 | """ 367 | if self.vector_store: 368 | return await self.vector_store.get_recent_history( 369 | limit=limit or self.config.history_length, 370 | role_filter=role_filter 371 | ) 372 | else: 373 | # Return from in-memory history 374 | history = list(self.conversation_history) 375 | if role_filter: 376 | history = [msg for msg in history if msg.role == role_filter] 377 | if limit: 378 | history = history[-limit:] 379 | return [msg.dict() for msg in history] 380 | 381 | async def search_similar_messages(self, 382 | query: str, 383 | limit: int = 5, 384 | role_filter: Optional[str] = None) -> List[Dict[str, Any]]: 385 | """ 386 | Search for similar messages in conversation history. 387 | Only available if vector store is configured. 388 | 389 | Args: 390 | query: Text to search for 391 | limit: Maximum number of results 392 | role_filter: Optional filter by role 393 | 394 | Returns: 395 | List of similar messages with scores 396 | """ 397 | if not self.vector_store: 398 | logger.warning("Vector store not configured, semantic search unavailable") 399 | return [] 400 | 401 | return await self.vector_store.search_similar( 402 | query=query, 403 | limit=limit, 404 | role_filter=role_filter 405 | ) 406 | 407 | async def clear_history(self): 408 | """Clear conversation history from both memory and vector store.""" 409 | self.conversation_history.clear() 410 | if self.vector_store: 411 | await self.vector_store.clear_history() 412 | logger.info("Cleared conversation history") -------------------------------------------------------------------------------- /nexus/agents/image_agent.py: -------------------------------------------------------------------------------- 1 | """ 2 | Image processing agent with OCR, image analysis, and generation capabilities. 3 | """ 4 | from typing import Optional, Union, BinaryIO, List, Dict, Any, Tuple 5 | from pathlib import Path 6 | import easyocr 7 | import cv2 8 | import numpy as np 9 | from PIL import Image 10 | from loguru import logger 11 | from pydantic import BaseModel, Field 12 | import tempfile 13 | import os 14 | import torch 15 | 16 | try: 17 | from diffusers import ( 18 | StableDiffusionPipeline, 19 | StableDiffusionControlNetPipeline, 20 | ControlNetModel, 21 | ) 22 | from compel import Compel 23 | from transformers import CLIPTokenizer 24 | GENERATION_AVAILABLE = True 25 | except ImportError: 26 | logger.warning("Image generation dependencies not available. Some features will be disabled.") 27 | GENERATION_AVAILABLE = False 28 | 29 | 30 | class GenerationConfig(BaseModel): 31 | """Configuration for image generation.""" 32 | model_id: str = "runwayml/stable-diffusion-v1-5" 33 | device: str = "cuda" if torch.cuda.is_available() else "cpu" 34 | dtype: torch.dtype = torch.float16 if torch.cuda.is_available() else torch.float32 35 | 36 | # Generation parameters 37 | num_inference_steps: int = 50 38 | guidance_scale: float = 7.5 39 | negative_prompt: str = "" 40 | width: int = 512 41 | height: int = 512 42 | num_images: int = 1 43 | seed: Optional[int] = None 44 | 45 | # ControlNet settings 46 | use_controlnet: bool = False 47 | controlnet_model: Optional[str] = None 48 | controlnet_conditioning_scale: float = 1.0 49 | 50 | # Advanced settings 51 | use_compel: bool = True # Use compel for advanced prompt weighting 52 | use_safety_checker: bool = True 53 | enable_attention_slicing: bool = True 54 | enable_vae_slicing: bool = True 55 | 56 | 57 | class ImageConfig(BaseModel): 58 | """Configuration for image processing.""" 59 | # OCR Configuration 60 | languages: List[str] = ["en"] # List of language codes for OCR 61 | device: str = "cpu" # Device to run inference on: cpu or cuda 62 | 63 | # Image Processing 64 | min_confidence: float = 0.5 # Minimum confidence for OCR detection 65 | preprocessing: bool = True # Whether to apply preprocessing 66 | contrast_enhance: bool = True # Whether to enhance contrast 67 | denoise: bool = True # Whether to apply denoising 68 | 69 | # Output Configuration 70 | return_bboxes: bool = True # Whether to return bounding boxes 71 | return_confidence: bool = True # Whether to return confidence scores 72 | 73 | # Generation Configuration 74 | generation: GenerationConfig = Field(default_factory=GenerationConfig) 75 | 76 | 77 | class ImageAgent: 78 | """ 79 | Agent for processing images, handling OCR, image analysis, and generation. 80 | Uses EasyOCR for text detection and Stable Diffusion for image generation. 81 | """ 82 | 83 | def __init__(self, config: Optional[ImageConfig] = None): 84 | """Initialize the image agent with optional configuration.""" 85 | self.config = config or ImageConfig() 86 | 87 | # Initialize OCR 88 | with logger.contextualize(languages=self.config.languages): 89 | logger.info("Initializing OCR model") 90 | self.reader = easyocr.Reader( 91 | lang_list=self.config.languages, 92 | gpu=self.config.device == "cuda" 93 | ) 94 | logger.info("OCR model loaded successfully") 95 | 96 | # Initialize generation pipeline 97 | self._init_generation_pipeline() 98 | 99 | def _init_generation_pipeline(self): 100 | """Initialize the image generation pipeline.""" 101 | if not GENERATION_AVAILABLE: 102 | logger.warning("Image generation is disabled due to missing dependencies") 103 | self.pipeline = None 104 | self.compel = None 105 | return 106 | 107 | try: 108 | with logger.contextualize( 109 | model=self.config.generation.model_id, 110 | device=self.config.generation.device 111 | ): 112 | logger.info("Initializing generation pipeline") 113 | 114 | # Basic pipeline 115 | if not self.config.generation.use_controlnet: 116 | self.pipeline = StableDiffusionPipeline.from_pretrained( 117 | self.config.generation.model_id, 118 | torch_dtype=self.config.generation.dtype, 119 | safety_checker=None if not self.config.generation.use_safety_checker else "default" 120 | ) 121 | else: 122 | # ControlNet pipeline 123 | controlnet = ControlNetModel.from_pretrained( 124 | self.config.generation.controlnet_model, 125 | torch_dtype=self.config.generation.dtype 126 | ) 127 | self.pipeline = StableDiffusionControlNetPipeline.from_pretrained( 128 | self.config.generation.model_id, 129 | controlnet=controlnet, 130 | torch_dtype=self.config.generation.dtype, 131 | safety_checker=None if not self.config.generation.use_safety_checker else "default" 132 | ) 133 | 134 | # Move to device 135 | self.pipeline = self.pipeline.to(self.config.generation.device) 136 | 137 | # Optimizations 138 | if self.config.generation.enable_attention_slicing: 139 | self.pipeline.enable_attention_slicing() 140 | if self.config.generation.enable_vae_slicing: 141 | self.pipeline.enable_vae_slicing() 142 | 143 | # Initialize Compel for advanced prompt weighting 144 | if self.config.generation.use_compel: 145 | self.compel = Compel( 146 | tokenizer=self.pipeline.tokenizer, 147 | text_encoder=self.pipeline.text_encoder, 148 | truncate_long_prompts=False 149 | ) 150 | 151 | logger.info("Generation pipeline initialized successfully") 152 | 153 | except Exception as e: 154 | logger.exception("Failed to initialize generation pipeline") 155 | raise RuntimeError(f"Generation pipeline initialization failed: {str(e)}") 156 | 157 | def generate_image( 158 | self, 159 | prompt: str, 160 | negative_prompt: Optional[str] = None, 161 | control_image: Optional[Union[str, Path, Image.Image]] = None, 162 | **kwargs 163 | ) -> Union[Image.Image, List[Image.Image]]: 164 | """ 165 | Generate an image from a text prompt using Stable Diffusion. 166 | 167 | Args: 168 | prompt: Text description of the desired image 169 | negative_prompt: Optional text description of what to avoid 170 | control_image: Optional control image for ControlNet 171 | **kwargs: Additional generation parameters to override config 172 | 173 | Returns: 174 | One or more PIL Images depending on num_images setting 175 | """ 176 | if not GENERATION_AVAILABLE: 177 | raise RuntimeError("Image generation is not available. Please install required dependencies.") 178 | 179 | try: 180 | with logger.contextualize(operation="generate"): 181 | logger.info("Starting image generation") 182 | logger.debug("Prompt: {}", prompt) 183 | 184 | # Process prompt with Compel if enabled 185 | if self.config.generation.use_compel: 186 | prompt_embeds = self.compel(prompt) 187 | neg_prompt_embeds = self.compel( 188 | negative_prompt or self.config.generation.negative_prompt 189 | ) 190 | else: 191 | prompt_embeds = None 192 | neg_prompt_embeds = None 193 | 194 | # Prepare generation parameters 195 | generator = None 196 | if self.config.generation.seed is not None: 197 | generator = torch.Generator(device=self.config.generation.device) 198 | generator.manual_seed(self.config.generation.seed) 199 | 200 | # Generate image(s) 201 | if self.config.generation.use_controlnet and control_image: 202 | # Load and preprocess control image 203 | if isinstance(control_image, (str, Path)): 204 | control_image = Image.open(str(control_image)) 205 | control_image = control_image.resize( 206 | (self.config.generation.width, self.config.generation.height) 207 | ) 208 | 209 | result = self.pipeline( 210 | prompt=prompt if not prompt_embeds else None, 211 | prompt_embeds=prompt_embeds, 212 | negative_prompt=negative_prompt if not neg_prompt_embeds else None, 213 | negative_prompt_embeds=neg_prompt_embeds, 214 | image=control_image, 215 | num_inference_steps=self.config.generation.num_inference_steps, 216 | guidance_scale=self.config.generation.guidance_scale, 217 | width=self.config.generation.width, 218 | height=self.config.generation.height, 219 | num_images_per_prompt=self.config.generation.num_images, 220 | generator=generator, 221 | controlnet_conditioning_scale=self.config.generation.controlnet_conditioning_scale, 222 | **kwargs 223 | ).images 224 | else: 225 | result = self.pipeline( 226 | prompt=prompt if not prompt_embeds else None, 227 | prompt_embeds=prompt_embeds, 228 | negative_prompt=negative_prompt if not neg_prompt_embeds else None, 229 | negative_prompt_embeds=neg_prompt_embeds, 230 | num_inference_steps=self.config.generation.num_inference_steps, 231 | guidance_scale=self.config.generation.guidance_scale, 232 | width=self.config.generation.width, 233 | height=self.config.generation.height, 234 | num_images_per_prompt=self.config.generation.num_images, 235 | generator=generator, 236 | **kwargs 237 | ).images 238 | 239 | logger.info("Image generation completed successfully") 240 | return result[0] if len(result) == 1 else result 241 | 242 | except Exception as e: 243 | logger.exception("Image generation failed") 244 | raise RuntimeError(f"Image generation failed: {str(e)}") 245 | 246 | def generate_diagram( 247 | self, 248 | description: str, 249 | style: str = "flowchart", 250 | **kwargs 251 | ) -> Image.Image: 252 | """ 253 | Generate a diagram from a text description. 254 | 255 | Args: 256 | description: Text description of the desired diagram 257 | style: Type of diagram (flowchart, mindmap, etc.) 258 | **kwargs: Additional generation parameters 259 | 260 | Returns: 261 | PIL Image of the generated diagram 262 | """ 263 | if not GENERATION_AVAILABLE: 264 | raise RuntimeError("Image generation is not available. Please install required dependencies.") 265 | 266 | # Construct a specialized prompt for diagram generation 267 | diagram_prompt = f"A clean, professional {style} diagram showing {description}. " \ 268 | f"Minimalist design, high contrast, clear text and arrows." 269 | 270 | # Add style-specific negative prompts 271 | negative_prompt = "blurry, photo-realistic, complex background, " \ 272 | "natural scene, artistic, painterly" 273 | 274 | # Generate with specific settings for diagrams 275 | return self.generate_image( 276 | prompt=diagram_prompt, 277 | negative_prompt=negative_prompt, 278 | width=1024, # Larger size for diagrams 279 | height=768, 280 | guidance_scale=8.5, # Higher guidance for more precise output 281 | num_inference_steps=60, # More steps for detail 282 | **kwargs 283 | ) 284 | 285 | def _preprocess_image(self, image: np.ndarray) -> np.ndarray: 286 | """ 287 | Apply preprocessing steps to the image. 288 | """ 289 | try: 290 | processed = image.copy() 291 | 292 | if self.config.preprocessing: 293 | # Convert to grayscale if not already 294 | if len(processed.shape) == 3: 295 | processed = cv2.cvtColor(processed, cv2.COLOR_BGR2GRAY) 296 | 297 | # Enhance contrast if enabled 298 | if self.config.contrast_enhance: 299 | logger.debug("Enhancing contrast") 300 | processed = cv2.equalizeHist(processed) 301 | 302 | # Apply denoising if enabled 303 | if self.config.denoise: 304 | logger.debug("Applying denoising") 305 | processed = cv2.fastNlMeansDenoising(processed) 306 | 307 | return processed 308 | 309 | except Exception as e: 310 | logger.error("Error in image preprocessing: {}", str(e)) 311 | raise ValueError(f"Failed to preprocess image: {str(e)}") 312 | 313 | def _load_image(self, image_input: Union[str, Path, BinaryIO, np.ndarray]) -> np.ndarray: 314 | """ 315 | Load image from various input types and convert to numpy array. 316 | """ 317 | try: 318 | if isinstance(image_input, (str, Path)): 319 | # Load from file path 320 | image = cv2.imread(str(image_input)) 321 | if image is None: 322 | raise ValueError(f"Failed to load image from {image_input}") 323 | return image 324 | elif isinstance(image_input, np.ndarray): 325 | # Already a numpy array 326 | return image_input.copy() 327 | else: 328 | # Assume it's a file-like object, save to temp file first 329 | with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as temp_file: 330 | temp_file.write(image_input.read()) 331 | temp_path = temp_file.name 332 | 333 | try: 334 | image = cv2.imread(temp_path) 335 | if image is None: 336 | raise ValueError(f"Failed to load image from file-like object") 337 | return image 338 | finally: 339 | os.unlink(temp_path) # Clean up temp file 340 | 341 | except Exception as e: 342 | logger.error("Error loading image: {}", str(e)) 343 | raise ValueError(f"Failed to load image: {str(e)}") 344 | 345 | def extract_text( 346 | self, 347 | image: Union[str, Path, Image.Image, np.ndarray, BinaryIO], 348 | **kwargs 349 | ) -> Union[str, List[Dict[str, Any]]]: 350 | """ 351 | Extract text from an image using OCR. 352 | 353 | Args: 354 | image: Input image as file path, PIL Image, numpy array, or file object 355 | **kwargs: Additional parameters to pass to the OCR reader 356 | 357 | Returns: 358 | If return_bboxes and return_confidence are False, returns the extracted text as a string. 359 | Otherwise, returns a list of dictionaries containing text, bounding boxes, and confidence scores. 360 | """ 361 | try: 362 | with logger.contextualize(operation="ocr"): 363 | logger.info("Starting OCR text extraction") 364 | 365 | # Load and preprocess image 366 | if isinstance(image, (str, Path)): 367 | logger.debug("Loading image from path: {}", image) 368 | image = cv2.imread(str(image)) 369 | elif isinstance(image, Image.Image): 370 | logger.debug("Converting PIL Image to numpy array") 371 | image = np.array(image) 372 | elif isinstance(image, BinaryIO): 373 | logger.debug("Loading image from file object") 374 | image = cv2.imdecode( 375 | np.frombuffer(image.read(), np.uint8), 376 | cv2.IMREAD_COLOR 377 | ) 378 | 379 | if image is None: 380 | raise ValueError("Failed to load image") 381 | 382 | # Preprocess image 383 | processed_image = self._preprocess_image(image) 384 | 385 | # Perform OCR 386 | logger.debug("Running OCR") 387 | results = self.reader.readtext( 388 | processed_image, 389 | detail=self.config.return_bboxes or self.config.return_confidence, 390 | **kwargs 391 | ) 392 | 393 | # Format results 394 | if not (self.config.return_bboxes or self.config.return_confidence): 395 | # Return simple text string 396 | text = " ".join([result[1] for result in results]) 397 | logger.info("OCR completed successfully") 398 | return text 399 | else: 400 | # Return detailed results 401 | formatted_results = [] 402 | for bbox, text, conf in results: 403 | if conf >= self.config.min_confidence: 404 | result = {"text": text} 405 | if self.config.return_bboxes: 406 | result["bbox"] = bbox 407 | if self.config.return_confidence: 408 | result["confidence"] = conf 409 | formatted_results.append(result) 410 | 411 | logger.info("OCR completed successfully") 412 | return formatted_results 413 | 414 | except Exception as e: 415 | logger.exception("OCR text extraction failed") 416 | raise RuntimeError(f"OCR text extraction failed: {str(e)}") 417 | 418 | def detect_text_regions(self, 419 | image_input: Union[str, Path, BinaryIO, np.ndarray]) -> List[Dict[str, Any]]: 420 | """ 421 | Detect regions containing text in the image. 422 | 423 | Args: 424 | image_input: Image file path, binary data, or numpy array 425 | 426 | Returns: 427 | List of detected regions with coordinates and confidence scores 428 | """ 429 | try: 430 | with logger.contextualize(operation="detect_regions"): 431 | logger.info("Starting text region detection") 432 | 433 | # Load and preprocess image 434 | image = self._load_image(image_input) 435 | processed_image = self._preprocess_image(image) 436 | 437 | # Run detection 438 | results = self.reader.detect( 439 | processed_image, 440 | min_size=10, 441 | text_threshold=self.config.min_confidence 442 | ) 443 | 444 | # Format regions 445 | regions = [] 446 | if results is not None and len(results) == 2: 447 | boxes, scores = results 448 | for box, score in zip(boxes, scores): 449 | if score >= self.config.min_confidence: 450 | regions.append({ 451 | "bbox": box.tolist(), 452 | "confidence": float(score) 453 | }) 454 | 455 | logger.info("Detected {} text regions", len(regions)) 456 | return regions 457 | 458 | except Exception as e: 459 | logger.exception("Region detection failed") 460 | raise RuntimeError(f"Region detection failed: {str(e)}") 461 | 462 | def supported_languages(self) -> List[str]: 463 | """List supported OCR languages.""" 464 | return self.reader.lang_list -------------------------------------------------------------------------------- /action_plan.md: -------------------------------------------------------------------------------- 1 | # Updated Action Plan for Multi-Modal Agentic Platform 2 | 3 | Below is an **updated action plan** that factors in the new Orchestrator code you've added. This roadmap continues to follow the **phased** approach, but now explicitly acknowledges the **Orchestrator** foundation already in place. The steps are broken into **bite-sized, independently codable tasks**, ensuring each one yields a meaningful improvement while maintaining **Separation of Concerns** and a **Layered Architecture**. 4 | 5 | --- 6 | 7 | ## Phase 1: Foundational Orchestrator & Basic Text Interaction (Updated) 8 | 9 | ### 1.1 [X] Minimal Orchestrator 10 | - **What We've Done** 11 | - Created `Orchestrator` and `OrchestratorConfig` classes. 12 | - Implemented a `process_input` method that logs user input and returns a placeholder response. 13 | - Added a basic usage script (`basic_usage.py`) to demonstrate how to instantiate and call the Orchestrator. 14 | - Confirmed logging works as intended (via `loguru`). 15 | 16 | **Usage & Integration**: 17 | ```python 18 | from mnemosyne.orchestrator import Orchestrator, OrchestratorConfig 19 | 20 | orchestrator = Orchestrator( 21 | config=OrchestratorConfig(debug=True) 22 | ) 23 | response = orchestrator.process_input("Your query here") 24 | ``` 25 | 26 | **Purpose**: Provides the foundation for all agent interactions and workflow management. 27 | **Related Features**: Core component that all other features build upon. 28 | 29 | ### 1.2 [X] Integrate a Basic Text LLM 30 | 1. **Add LLM Client** [X] 31 | - Implemented `_call_llm()` with support for multiple providers (OpenAI, Anthropic, Google) 32 | - Added fallback provider support for reliability 33 | 2. **Refine `process_input`** [X] 34 | - Now uses actual LLM calls with comprehensive logging 35 | - Handles errors gracefully with fallback providers 36 | 3. **Testing & Validation** [X] 37 | - Added example in `basic_usage.py` 38 | - Includes error handling for API issues 39 | 40 | **Usage**: 41 | ```python 42 | orchestrator = Orchestrator( 43 | config=OrchestratorConfig( 44 | primary_provider=LLMProviderConfig( 45 | provider="openai", 46 | model="gpt-4" 47 | ), 48 | fallback_providers=[...] 49 | ) 50 | ) 51 | ``` 52 | 53 | **Purpose**: Enables intelligent responses using state-of-the-art LLMs. 54 | **Related Features**: Foundation for all LLM-based operations. 55 | 56 | ### 1.3 [X] Basic Memory (In-Memory) 57 | 1. **Conversation Buffer** [X] 58 | - Implemented using `deque` for efficient history management 59 | - Automatically includes context in LLM prompts 60 | 2. **Configuration** [X] 61 | - Added `history_length` to `OrchestratorConfig` 62 | 3. **Testing** [X] 63 | - Demonstrated in `basic_usage.py` 64 | 65 | **Usage**: 66 | ```python 67 | orchestrator = Orchestrator( 68 | config=OrchestratorConfig( 69 | history_length=10 # Store last 10 messages 70 | ) 71 | ) 72 | ``` 73 | 74 | **Purpose**: Enables context-aware conversations. 75 | **Related Features**: Foundation for Vector DB integration. 76 | 77 | ### 1.4 [X] Enhanced Logging & Debugging 78 | 1. **Enhanced Debug Mode** [X] 79 | - Implemented comprehensive logging with loguru 80 | - Added structured context for better filtering 81 | 2. **Structured Logging** [X] 82 | - Added file rotation and retention policies 83 | - Implemented context-aware logging 84 | 85 | **Usage**: Debug logs are automatically managed based on config: 86 | ```python 87 | orchestrator = Orchestrator( 88 | config=OrchestratorConfig(debug=True) 89 | ) 90 | ``` 91 | 92 | **Purpose**: Facilitates debugging and monitoring. 93 | **Related Features**: Supports all components with standardized logging. 94 | 95 | --- 96 | 97 | ## Phase 2: Multi-Modal & Workflow Patterns 98 | 99 | ### 2.1 [X] Audio Agent (STT/TTS) 100 | 1. **Speech-to-Text Integration** [X] 101 | - Created `AudioAgent` with Whisper integration 102 | - Added language detection and transcription 103 | 2. **Optional: TTS** [X] 104 | - Implemented voice selection and synthesis 105 | 106 | **Usage**: 107 | ```python 108 | from mnemosyne.agents.audio_agent import AudioAgent, AudioConfig 109 | 110 | agent = AudioAgent( 111 | config=AudioConfig( 112 | model_type="base", 113 | device="cpu", 114 | language="en" 115 | ) 116 | ) 117 | text = await agent.transcribe("audio.wav") 118 | ``` 119 | 120 | **Purpose**: Enables audio processing capabilities. 121 | **Related Features**: Integrates with Orchestrator for multi-modal input. 122 | 123 | ### 2.2 [X] Image Agent (OCR, Generation) 124 | 1. **OCR** [X] 125 | - Implemented `ImageAgent` with Tesseract/EasyOCR 126 | - Added text extraction and processing 127 | 2. **Image Generation** [X] 128 | - Added support for image generation from text 129 | 130 | **Usage**: 131 | ```python 132 | from mnemosyne.agents.image_agent import ImageAgent 133 | 134 | agent = ImageAgent() 135 | text = await agent.extract_text("image.jpg") 136 | ``` 137 | 138 | **Purpose**: Enables image processing and generation. 139 | **Related Features**: Works with Orchestrator and EvaluatorAgent. 140 | 141 | ### 2.3 [X] Parallelization & Basic Evaluator 142 | 1. **Parallelization Example** [X] 143 | - Added async support to Orchestrator 144 | - Implemented `run_parallel_tasks` 145 | 2. **Evaluator-Optimizer** [X] 146 | - Created `EvaluatorAgent` for output validation 147 | - Added configurable evaluation criteria 148 | 149 | **Usage**: 150 | ```python 151 | # Parallel execution 152 | results = await orchestrator.run_parallel_tasks([ 153 | task1, 154 | task2 155 | ], timeout=30.0) 156 | 157 | # Evaluation 158 | evaluator = EvaluatorAgent() 159 | result = await evaluator.evaluate_output(text) 160 | ``` 161 | 162 | **Purpose**: Enables efficient parallel processing and output validation. 163 | **Related Features**: Enhances all agents with parallel capabilities and quality checks. 164 | 165 | --- 166 | 167 | ## Phase 3: Refinement & Memory 168 | 169 | ### 3.1 Vector Database Integration 170 | 1. **Vector Store Setup** 171 | - Choose a vector DB (e.g. Pinecone, Weaviate, or local LanceDB). 172 | - Implement a small module (`vector_store.py`) for storing and retrieving text embeddings. 173 | 2. **Augmented Retrieval** 174 | - For each user query, retrieve relevant context from the vector DB to inject into the LLM prompt (RAG approach). 175 | 3. **Persisted Memory** 176 | - Optionally store conversation turns in the vector DB so that memory can survive restarts. 177 | 178 | **Outcome**: 179 | A more robust memory mechanism and retrieval-augmented generation to reduce hallucinations. 180 | 181 | --- 182 | 183 | ### 3.2 [X] Security & Hardening 184 | 1. **Auth & Encryption** [X] 185 | - Implemented Fernet-based encryption for sensitive data 186 | - Added configurable encryption key management 187 | 2. **Error Handling** [X] 188 | - Added exponential backoff retry mechanism 189 | - Implemented configurable timeouts and retry limits 190 | 191 | **Usage**: 192 | ```python 193 | # Encryption 194 | from mnemosyne.security import Security, SecurityConfig 195 | 196 | security = Security( 197 | config=SecurityConfig( 198 | max_retries=3, 199 | initial_wait=1.0, 200 | max_wait=30.0 201 | ) 202 | ) 203 | 204 | # Encrypt sensitive data 205 | encrypted = security.encrypt_data(sensitive_data) 206 | decrypted = security.decrypt_data(encrypted) 207 | 208 | # Retry mechanism 209 | @Security.with_retries(max_tries=3, initial_wait=1.0, max_wait=5.0) 210 | async def flaky_operation(): 211 | # Your code here 212 | pass 213 | ``` 214 | 215 | **Purpose**: Ensures data security and system reliability. 216 | **Related Features**: 217 | - Works with all agents to protect sensitive data (OCR results, audio transcripts) 218 | - Enhances reliability of external API calls with retry mechanism 219 | 220 | --- 221 | 222 | ### 3.3 [X] Performance & Monitoring 223 | 1. **Caching** [X] 224 | - Implemented TTL-based caching for embeddings and LLM responses 225 | - Added size limits and automatic cleanup of old entries 226 | 2. **Rate Limiting** [X] 227 | - Added per-user and global rate limits 228 | - Implemented type-specific limits (embedding/LLM) 229 | - Added exponential backoff for retries 230 | 231 | **Usage**: 232 | ```python 233 | from mnemosyne.caching import CacheManager, CacheConfig, RateLimitConfig 234 | 235 | # Initialize cache manager 236 | cache_manager = CacheManager( 237 | cache_config=CacheConfig( 238 | embedding_ttl=3600, # 1 hour 239 | llm_response_ttl=1800, # 30 minutes 240 | max_embedding_size=10000, 241 | max_llm_size=1000 242 | ), 243 | rate_limit_config=RateLimitConfig( 244 | user_limit=100, # Per minute 245 | global_limit=1000, # Per minute 246 | embedding_limit=500, # Per minute 247 | llm_limit=200 # Per minute 248 | ) 249 | ) 250 | 251 | # Use decorators for automatic caching and rate limiting 252 | @cache_embedding(cache_manager=cache_manager) 253 | async def generate_embedding(text: str): 254 | # Your embedding code here 255 | pass 256 | 257 | @cache_llm_response(cache_manager=cache_manager) 258 | async def generate_response(prompt: str): 259 | # Your LLM code here 260 | pass 261 | ``` 262 | 263 | **Purpose**: Improves performance and reliability while controlling resource usage. 264 | **Related Features**: 265 | - Works with VectorStore for efficient embedding management 266 | - Enhances Orchestrator with response caching 267 | - Protects external APIs with rate limiting 268 | 269 | --- 270 | 271 | ### 3.4 Advanced Workflows [X] 272 | 273 | #### Prompt Chaining Library [X] 274 | 1. **Core Implementation** [X] 275 | - Created flexible `PromptChain` class with support for: 276 | - Linear chains (A → B → C) 277 | - Branching chains (A → [B1, B2] → C) 278 | - Dynamic chain modification during execution 279 | - Per-node validation before execution 280 | - Added comprehensive logging and error handling 281 | 282 | 2. **Node Types** [X] 283 | - PROMPT: LLM prompt nodes 284 | - BRANCH: Conditional branching nodes 285 | - MERGE: Combine results from multiple paths 286 | - TOOL: External tool/function calls 287 | 288 | 3. **Features** [X] 289 | - Context Management: 290 | - Variables passed between nodes 291 | - Results tracking 292 | - Execution path history 293 | - Dynamic Modification: 294 | - Add/remove nodes during execution 295 | - Modify connections between nodes 296 | - Conditional path selection 297 | 298 | **Usage**: 299 | ```python 300 | from mnemosyne.chains import ( 301 | PromptChain, 302 | ChainNodeType, 303 | create_linear_chain, 304 | create_branching_chain 305 | ) 306 | 307 | # Create a linear chain 308 | chain = create_linear_chain( 309 | "text_processor", 310 | prompts=[ 311 | "Analyze this text: {input_text}", 312 | "Summarize the analysis: {processed_text}" 313 | ], 314 | tools=[{ 315 | "name": "text_processor", 316 | "params": {"text": "{input_text}"} 317 | }] 318 | ) 319 | 320 | # Execute with context 321 | results = await chain.execute( 322 | initial_context={"input_text": "Hello, World!"}, 323 | tools={"text_processor": my_tool} 324 | ) 325 | 326 | # Create a branching chain 327 | branch_chain = create_branching_chain( 328 | "classifier", 329 | condition_prompt="Is this a question: {input_text}", 330 | true_branch=["Answer: {input_text}"], 331 | false_branch=["Process: {input_text}"] 332 | ) 333 | 334 | # Dynamic modification 335 | chain.add_node(...) 336 | chain.connect(from_id, to_id) 337 | ``` 338 | 339 | **Purpose**: Enables creation of complex, dynamic prompt workflows. 340 | 341 | **Benefits**: 342 | - Modular and reusable prompt patterns 343 | - Flexible workflow construction 344 | - Dynamic adaptation to results 345 | - Built-in validation and error handling 346 | - Comprehensive logging for debugging 347 | 348 | **Related Features**: 349 | - Integrates with caching system for LLM responses 350 | - Works with rate limiting for external API calls 351 | - Supports all agent types (Text, Audio, Image, Video) 352 | - Compatible with security module for sensitive data 353 | 354 | --- 355 | 356 | ## Phase 4: Advanced Workflows & Domain-Specific Agents 357 | 358 | ### 4.1 Prompt Chaining Library 359 | 1. **Chain Utility** 360 | - Implement a reusable "PromptChain" class or decorator to define multi-step tasks (Outline → Validate → Expand). 361 | - Let the Orchestrator load or execute these chains dynamically based on user requests. 362 | 2. **Testing** 363 | - Provide small chain examples (like a "blog post writer" flow). 364 | 365 | --- 366 | 367 | ### 4.2 Orchestrator-Workers Pattern [X] 368 | 369 | 1. **Worker Pool Implementation** [X] 370 | - Created flexible `WorkerPool` class with support for: 371 | - Multiple worker types (TEXT, AUDIO, IMAGE, VIDEO) 372 | - Task prioritization and dependencies 373 | - Worker status tracking and error handling 374 | - Added comprehensive logging for debugging 375 | 376 | 2. **Task Management** [X] 377 | - Implemented task queue with priority-based scheduling 378 | - Added support for task dependencies 379 | - Created task status tracking and result storage 380 | 381 | 3. **Worker Orchestration** [X] 382 | - Created `WorkerOrchestrator` for managing worker pools 383 | - Added support for parallel and sequential processing 384 | - Implemented worker group creation with capabilities 385 | 386 | **Usage**: 387 | ```python 388 | from mnemosyne.workers import WorkerOrchestrator, WorkerType 389 | 390 | # Initialize orchestrator 391 | orchestrator = WorkerOrchestrator() 392 | 393 | # Create worker groups 394 | orchestrator.create_worker_group( 395 | type=WorkerType.TEXT, 396 | count=3, 397 | capabilities=["summarize", "analyze"] 398 | ) 399 | 400 | # Process tasks in parallel 401 | results = await orchestrator.process_parallel( 402 | tasks=[{"text": "Task 1"}, {"text": "Task 2"}], 403 | worker_type=WorkerType.TEXT, 404 | processor=text_processor, 405 | num_workers=3 406 | ) 407 | 408 | # Process tasks sequentially 409 | results = await orchestrator.process_sequential( 410 | tasks=[{"path": "image1.jpg"}, {"path": "image2.jpg"}], 411 | worker_type=WorkerType.IMAGE, 412 | processor=image_processor 413 | ) 414 | 415 | # Add tasks with dependencies 416 | task1_id = orchestrator.pool.add_task( 417 | type=WorkerType.TEXT, 418 | content={"text": "Process first"} 419 | ) 420 | task2_id = orchestrator.pool.add_task( 421 | type=WorkerType.IMAGE, 422 | content={"path": "process_after.jpg"}, 423 | dependencies=[task1_id] 424 | ) 425 | ``` 426 | 427 | **Purpose**: Enables efficient parallel processing and task orchestration. 428 | 429 | **Benefits**: 430 | - Flexible worker pool management 431 | - Priority-based task scheduling 432 | - Support for task dependencies 433 | - Parallel and sequential processing 434 | - Comprehensive logging and error handling 435 | 436 | **Related Features**: 437 | - Works with all agent types (Text, Audio, Image, Video) 438 | - Integrates with logging system for debugging 439 | - Compatible with caching and rate limiting 440 | - Supports the security module for sensitive data 441 | 442 | --- 443 | 444 | ### 4.3 Enhanced Evaluator-Optimizer [X] 445 | 446 | 1. **Multi-Criteria Evaluation** [X] 447 | - Created flexible `ContentEvaluator` with support for: 448 | - Multiple evaluation criteria (factual accuracy, style, policy, etc.) 449 | - Weighted scoring system 450 | - Configurable thresholds per criterion 451 | - Comprehensive feedback and suggestions 452 | - Added validation for weights and thresholds 453 | 454 | 2. **Automatic Refinement** [X] 455 | - Implemented iterative refinement loop 456 | - Added smart prompt generation based on failing criteria 457 | - Created configurable refinement thresholds 458 | - Added maximum iteration limits 459 | 460 | 3. **Context-Aware Evaluation** [X] 461 | - Added support for evaluation context (audience, style guide, domain) 462 | - Created criterion-specific prompt templates 463 | - Implemented context-aware refinement suggestions 464 | 465 | **Usage**: 466 | ```python 467 | from mnemosyne.evaluator import ( 468 | ContentEvaluator, 469 | EvaluationConfig, 470 | EvaluationCriteria 471 | ) 472 | 473 | # Configure evaluator 474 | config = EvaluationConfig( 475 | criteria=[ 476 | EvaluationCriteria.FACTUAL_ACCURACY, 477 | EvaluationCriteria.STYLE_CONSISTENCY 478 | ], 479 | weights={ 480 | EvaluationCriteria.FACTUAL_ACCURACY: 0.6, 481 | EvaluationCriteria.STYLE_CONSISTENCY: 0.4 482 | }, 483 | thresholds={ 484 | EvaluationCriteria.FACTUAL_ACCURACY: 0.8, 485 | EvaluationCriteria.STYLE_CONSISTENCY: 0.7 486 | } 487 | ) 488 | 489 | evaluator = ContentEvaluator(config) 490 | 491 | # Single evaluation 492 | result = await evaluator.evaluate( 493 | content="Your content here", 494 | context={"audience": "technical"} 495 | ) 496 | 497 | # Print scores 498 | for score in result.scores: 499 | print(f"{score.criterion}: {score.score} - {score.feedback}") 500 | 501 | # Automatic refinement 502 | refined_content, history = await evaluator.evaluate_and_refine( 503 | content="Content to improve", 504 | max_iterations=3 505 | ) 506 | ``` 507 | 508 | **Purpose**: Ensures high-quality output through systematic evaluation and refinement. 509 | 510 | **Benefits**: 511 | - Comprehensive quality assessment 512 | - Automated content improvement 513 | - Flexible evaluation criteria 514 | - Context-aware evaluation 515 | - Detailed feedback and suggestions 516 | 517 | **Related Features**: 518 | - Works with all content types (text, code, etc.) 519 | - Integrates with logging system for debugging 520 | - Compatible with worker pool for parallel evaluation 521 | - Supports the security module for sensitive data 522 | 523 | --- 524 | 525 | ## Phase 5: Performance, Monitoring, & CI/CD 526 | 527 | ### 5.1 Performance Monitoring [X] 528 | 529 | 1. **Metrics Collection** [X] 530 | - Created flexible `PerformanceMetrics` class with support for: 531 | - Request counts and latencies 532 | - Error tracking 533 | - Worker and queue metrics 534 | - Memory usage monitoring 535 | - Added Prometheus integration for metrics export 536 | 537 | 2. **Monitoring Tools** [X] 538 | - Implemented context managers for tracking: 539 | - Request durations and outcomes 540 | - Worker activity 541 | - Queue sizes 542 | - Added comprehensive error tracking 543 | - Created metrics aggregation and summaries 544 | 545 | 3. **Component Integration** [X] 546 | - Created `MonitoredComponent` base class 547 | - Added monitoring decorators and utilities 548 | - Implemented background metric collection 549 | 550 | **Usage**: 551 | ```python 552 | from mnemosyne.monitoring import ( 553 | PerformanceMetrics, 554 | MonitoredComponent 555 | ) 556 | 557 | # Initialize metrics with Prometheus 558 | metrics = PerformanceMetrics(port=8000) 559 | 560 | # Create monitored component 561 | class MyComponent(MonitoredComponent): 562 | def __init__(self, metrics): 563 | super().__init__(metrics) 564 | 565 | async def process(self, data): 566 | # Track operation metrics 567 | async with self.track_operation( 568 | operation="process", 569 | agent_type="processor" 570 | ): 571 | # Your code here 572 | pass 573 | 574 | # Track worker metrics 575 | async with self.track_worker( 576 | agent_type="processor" 577 | ): 578 | # Worker code here 579 | pass 580 | 581 | # Record errors 582 | self.record_error( 583 | error_type="ValidationError", 584 | agent_type="processor" 585 | ) 586 | 587 | # Update queue metrics 588 | self.update_queue_size( 589 | size=len(queue), 590 | agent_type="processor" 591 | ) 592 | 593 | # Get metrics summary 594 | summary = metrics.get_summary( 595 | agent_type="processor", 596 | window=timedelta(minutes=5) 597 | ) 598 | ``` 599 | 600 | **Purpose**: Enables comprehensive performance monitoring and debugging. 601 | 602 | **Benefits**: 603 | - Real-time performance tracking 604 | - Prometheus integration for visualization 605 | - Flexible metric collection 606 | - Error tracking and analysis 607 | - Queue and worker monitoring 608 | 609 | **Related Features**: 610 | - Works with all components and agents 611 | - Integrates with logging system 612 | - Compatible with worker pool 613 | - Supports the security module 614 | 615 | ### 5.2 CI/CD Pipeline 616 | 617 | 1. **Continuous Integration** 618 | - Set up tests to run on each pull request (e.g., GitHub Actions). 619 | - Automatically build Docker images on merges to `main`. 620 | 621 | --- 622 | 623 | ## Phase 6: Extensions & Community Contributions 624 | 625 | 1. **Domain-Specific Agents** 626 | - E.g. a "Code Generation Agent" that can read/write local files in a restricted sandbox. 627 | - A "Legal Document Agent" that references known legal data from the vector store. 628 | 2. **UI Integrations** 629 | - Optionally create a minimal web front-end or connect to no-code platforms (Rivet, Vellum, etc.). 630 | 631 | --- 632 | 633 | # Where We Stand 634 | 635 | - **Currently**: 636 | - You have a **minimal Orchestrator** (`Orchestrator` & `OrchestratorConfig`) with basic logging and a placeholder response. 637 | - A usage demo (`basic_usage.py`) confirms that the Orchestrator can handle text input and log the interaction. 638 | 639 | - **Next Immediate Step**: 640 | - **Tie in a live LLM** (see Phase 1.2) so that `process_input` produces real AI-driven responses. 641 | - Then add **short-term in-memory conversation history** (Phase 1.3) for basic multi-turn interactions. 642 | 643 | --- 644 | 645 | ## Tips & Best Practices 646 | 647 | - **Keep PRs Small**: Implement each sub-step in its own branch and PR so code reviews remain manageable. 648 | - **Document as You Go**: Update `README.md` and your docstrings each time you add a new phase or feature. 649 | - **Test Often**: Each new feature (e.g., `AudioAgent`, `ImageAgent`) should come with at least a small test script or unit test. 650 | 651 | With this updated plan, you'll **incrementally build** a powerful multi-modal agentic platform—starting with the **Orchestrator** you've already set up, then layering in specialized agents, memory, retrieval, and advanced workflows. 652 | ``` --------------------------------------------------------------------------------