diff --git a/.gitignore b/.gitignore index b1dd850..4ae5d56 100644 --- a/.gitignore +++ b/.gitignore @@ -64,6 +64,7 @@ db.sqlite3-journal # Flask stuff: instance/ .webassets-cache +assets # Scrapy stuff: .scrapy @@ -82,6 +83,10 @@ target/ profile_default/ ipython_config.py +Atomic_Agents_vis +AgentChat_vis +agentchat_graph.json +atomic.json # pyenv # For a library or package, you might want to ignore these files since the code is # intended to run in multiple environments; otherwise, check them in: diff --git a/atomic-agent.json b/atomic-agent.json new file mode 100644 index 0000000..e1e6781 --- /dev/null +++ b/atomic-agent.json @@ -0,0 +1,255 @@ +{ + "nodes": [ + { + "id": "End", + "name": "End", + "node_type": "End", + "function_name": null, + "docstring": null, + "source_location": null, + "metadata": {} + }, + { + "id": "Start", + "name": "Start", + "node_type": "Start", + "function_name": null, + "docstring": null, + "source_location": null, + "metadata": {} + }, + { + "id": "agent", + "name": "agent", + "node_type": "Agent", + "source_location": { + "file": "examples\\code\\atomic_agents\\mainPDFAnalysis.py", + "line": 43, + "col": 0, + "end_line": 51, + "end_col": 1 + }, + "metadata": { + "type": "AtomicAgent", + "config": { + "client": "client", + "model": "'gemini-2.0-flash'", + "system_prompt_generator": "system_prompt_generator", + "input_schema": "InputSchema", + "output_schema": "ExtractionResult" + } + }, + "function_name": null, + "docstring": null + }, + { + "id": "nutrition_analyzer", + "name": "nutrition_analyzer", + "node_type": "Agent", + "source_location": { + "file": "examples\\code\\atomic_agents\\mainBaseModal.py", + "line": 62, + "col": 0, + "end_line": 90, + "end_col": 1 + }, + "metadata": { + "type": "AtomicAgent", + "config": { + "client": "instructor.from_openai(openai.OpenAI(api_key=API_KEY))", + "model": "'gpt-5-mini'", + "model_api_parameters": "{'reasoning_effort': 'low'}", + "system_prompt_generator": "SystemPromptGenerator(background=['You are a specialized nutrition label analyzer.', 'You excel at extracting precise nutritional information from food label images.', 'You understand various serving size formats and measurement units.', 'You can process multiple nutrition labels simultaneously.'], steps=['For each nutrition label image:', '1. Locate and identify the nutrition facts panel', '2. Extract all serving information and nutritional values', '3. Validate measurements and units for accuracy', '4. Compile the nutrition facts into structured data'], output_instructions=['For each analyzed nutrition label:', '1. Record complete serving size information', '2. Extract all nutrient values with correct units', '3. Ensure all measurements are properly converted', '4. Include all extracted labels in the final result'])" + } + }, + "function_name": null, + "docstring": null + }, + { + "id": "orchestrator_agent", + "name": "orchestrator_agent", + "node_type": "Agent", + "source_location": { + "file": "examples\\code\\atomic_agents\\orchestor.py", + "line": 181, + "col": 8, + "end_line": 181, + "end_col": 125 + }, + "metadata": { + "type": "AtomicAgent" + }, + "function_name": null, + "docstring": null + }, + { + "id": "orchestrator_agent_final", + "name": "orchestrator_agent_final", + "node_type": "Agent", + "source_location": { + "file": "examples\\code\\atomic_agents\\orchestor.py", + "line": 92, + "col": 0, + "end_line": 92, + "end_col": 116 + }, + "metadata": { + "type": "AtomicAgent" + }, + "function_name": null, + "docstring": null + } + ], + "edges": [ + { + "source": "Start", + "target": "agent", + "edge_type": "execution_start", + "condition": {}, + "metadata": { + "method": "run", + "definition_location": { + "file": "examples\\code\\atomic_agents\\hooksModel.py", + "line": 189, + "col": 23, + "end_line": 189, + "end_col": 39 + } + } + }, + { + "source": "Start", + "target": "agent", + "edge_type": "execution_start", + "condition": {}, + "metadata": { + "method": "run", + "definition_location": { + "file": "examples\\code\\atomic_agents\\hooksModel.py", + "line": 221, + "col": 23, + "end_line": 221, + "end_col": 39 + } + } + }, + { + "source": "Start", + "target": "agent", + "edge_type": "execution_start", + "condition": {}, + "metadata": { + "method": "run", + "definition_location": { + "file": "examples\\code\\atomic_agents\\hooksModel.py", + "line": 260, + "col": 23, + "end_line": 260, + "end_col": 39 + } + } + }, + { + "source": "Start", + "target": "agent", + "edge_type": "execution_start", + "condition": {}, + "metadata": { + "method": "run", + "definition_location": { + "file": "examples\\code\\atomic_agents\\mainPDFAnalysis.py", + "line": 68, + "col": 26, + "end_line": 68, + "end_col": 53 + } + } + }, + { + "source": "Start", + "target": "nutrition_analyzer", + "edge_type": "execution_start", + "condition": {}, + "metadata": { + "method": "run", + "definition_location": { + "file": "examples\\code\\atomic_agents\\mainBaseModal.py", + "line": 110, + "col": 26, + "end_line": 110, + "end_col": 66 + } + } + }, + { + "source": "Start", + "target": "orchestrator_agent", + "edge_type": "execution_start", + "condition": {}, + "metadata": { + "method": "run", + "definition_location": { + "file": "examples\\code\\atomic_agents\\orchestor.py", + "line": 154, + "col": 30, + "end_line": 154, + "end_col": 66 + } + } + }, + { + "source": "Start", + "target": "orchestrator_agent", + "edge_type": "execution_start", + "condition": {}, + "metadata": { + "method": "run", + "definition_location": { + "file": "examples\\code\\atomic_agents\\orchestor.py", + "line": 178, + "col": 23, + "end_line": 178, + "end_col": 59 + } + } + }, + { + "source": "Start", + "target": "orchestrator_agent_final", + "edge_type": "implicit_start", + "condition": {}, + "metadata": {} + }, + { + "source": "agent", + "target": "End", + "edge_type": "implicit_end", + "condition": {}, + "metadata": {} + }, + { + "source": "nutrition_analyzer", + "target": "End", + "edge_type": "implicit_end", + "condition": {}, + "metadata": {} + }, + { + "source": "orchestrator_agent", + "target": "End", + "edge_type": "implicit_end", + "condition": {}, + "metadata": {} + }, + { + "source": "orchestrator_agent_final", + "target": "End", + "edge_type": "implicit_end", + "condition": {}, + "metadata": {} + } + ], + "metadata": { + "framework": "Atomic_Agents" + } +} \ No newline at end of file diff --git a/examples/code/atomic_agents/hooksModel.py b/examples/code/atomic_agents/hooksModel.py new file mode 100644 index 0000000..ab07b2e --- /dev/null +++ b/examples/code/atomic_agents/hooksModel.py @@ -0,0 +1,329 @@ +#!/usr/bin/env python3 +""" +AtomicAgent Hook System Demo + +Shows how to monitor agent execution with hooks. +Includes error handling and performance metrics. +""" + +import os +import time +import logging + +import instructor +import openai +from rich.console import Console +from rich.panel import Panel +from rich.table import Table +from pydantic import Field, ValidationError + +from atomic_agents import AtomicAgent, AgentConfig +from atomic_agents.context import ChatHistory +from atomic_agents.base.base_io_schema import BaseIOSchema + +logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") +logger = logging.getLogger(__name__) +console = Console() +metrics = { + "total_requests": 0, + "successful_requests": 0, + "failed_requests": 0, + "parse_errors": 0, + "retry_attempts": 0, + "total_response_time": 0.0, + "start_time": time.time(), +} + +_request_start_time = None + + +class UserQuery(BaseIOSchema): + """Schema for user input containing a chat message.""" + + chat_message: str = Field(..., description="User's question or message") + + +class AgentResponse(BaseIOSchema): + """Schema for agent response with confidence and reasoning.""" + + chat_message: str = Field(..., description="Agent's response to the user") + confidence: float = Field(..., ge=0.0, le=1.0, description="Confidence score (0.0-1.0)") + reasoning: str = Field(..., description="Brief explanation of the reasoning") + + +class DetailedResponse(BaseIOSchema): + """Schema for detailed response with alternatives and confidence level.""" + + chat_message: str = Field(..., description="Primary response") + alternative_suggestions: list[str] = Field(default_factory=list, description="Alternative suggestions") + confidence_level: str = Field(..., description="Must be 'low', 'medium', or 'high'") + requires_followup: bool = Field(default=False, description="Whether follow-up is needed") + + +def setup_api_key() -> str: + api_key = os.getenv("OPENAI_API_KEY") + if not api_key: + console.print("[bold red]Error: OPENAI_API_KEY environment variable not set.[/bold red]") + console.print("Please set it with: export OPENAI_API_KEY='your-api-key-here'") + exit(1) + return api_key + + +def display_metrics(): + runtime = time.time() - metrics["start_time"] + avg_response_time = metrics["total_response_time"] / metrics["total_requests"] if metrics["total_requests"] > 0 else 0 + success_rate = metrics["successful_requests"] / metrics["total_requests"] * 100 if metrics["total_requests"] > 0 else 0 + + table = Table(title="🔍 Hook System Performance Metrics", style="cyan") + table.add_column("Metric", style="bold") + table.add_column("Value", style="green") + + table.add_row("Runtime", f"{runtime:.1f}s") + table.add_row("Total Requests", str(metrics["total_requests"])) + table.add_row("Successful Requests", str(metrics["successful_requests"])) + table.add_row("Failed Requests", str(metrics["failed_requests"])) + table.add_row("Parse Errors", str(metrics["parse_errors"])) + table.add_row("Retry Attempts", str(metrics["retry_attempts"])) + table.add_row("Success Rate", f"{success_rate:.1f}%") + table.add_row("Avg Response Time", f"{avg_response_time:.2f}s") + + console.print(table) + + +def on_parse_error(error): + metrics["parse_errors"] += 1 + metrics["failed_requests"] += 1 + logger.error(f"🚨 Parse error occurred: {type(error).__name__}: {error}") + + if isinstance(error, ValidationError): + console.print("[bold red]❌ Validation Error:[/bold red]") + for err in error.errors(): + field_path = " -> ".join(str(x) for x in err["loc"]) + console.print(f" • Field '{field_path}': {err['msg']}") + logger.error(f"Validation error in field '{field_path}': {err['msg']}") + else: + console.print(f"[bold red]❌ Parse Error:[/bold red] {error}") + + +def on_completion_kwargs(**kwargs): + global _request_start_time + metrics["total_requests"] += 1 + model = kwargs.get("model", "unknown") + messages_count = len(kwargs.get("messages", [])) + logger.info(f"🚀 API call starting - Model: {model}, Messages: {messages_count}") + _request_start_time = time.time() + + +def on_completion_response(response, **kwargs): + global _request_start_time + if _request_start_time: + response_time = time.time() - _request_start_time + metrics["total_response_time"] += response_time + logger.info(f"✅ API call completed in {response_time:.2f}s") + _request_start_time = None + + if hasattr(response, "usage"): + usage = response.usage + logger.info( + f"📊 Token usage - Prompt: {usage.prompt_tokens}, " + f"Completion: {usage.completion_tokens}, " + f"Total: {usage.total_tokens}" + ) + + metrics["successful_requests"] += 1 + + +def on_completion_error(error, **kwargs): + global _request_start_time + metrics["failed_requests"] += 1 + metrics["retry_attempts"] += 1 + + if _request_start_time: + _request_start_time = None + + logger.error(f"🔥 API error: {type(error).__name__}: {error}") + console.print(f"[bold red]🔥 API Error:[/bold red] {error}") + + +def create_agent_with_hooks(schema_type: type, system_prompt: str = None) -> AtomicAgent: + api_key = setup_api_key() + client = instructor.from_openai(openai.OpenAI(api_key=api_key)) + + config = AgentConfig( + client=client, + model="gpt-5-mini", + model_api_parameters={"reasoning_effort": "low"}, + history=ChatHistory(), + system_prompt=system_prompt, + ) + + agent = AtomicAgent[UserQuery, schema_type](config) + + agent.register_hook("parse:error", on_parse_error) + agent.register_hook("completion:kwargs", on_completion_kwargs) + agent.register_hook("completion:response", on_completion_response) + agent.register_hook("completion:error", on_completion_error) + + console.print("[bold green]✅ Agent created with comprehensive hook monitoring[/bold green]") + return agent + + +def demonstrate_basic_hooks(): + console.print(Panel("🔧 Basic Hook System Demonstration", style="bold blue")) + + agent = create_agent_with_hooks( + AgentResponse, "You are a helpful assistant. Always provide confident, well-reasoned responses." + ) + + test_queries = [ + "What is the capital of France?", + "Explain quantum computing in simple terms.", + "What are the benefits of renewable energy?", + ] + + for query_text in test_queries: + console.print(f"\n[bold cyan]Query:[/bold cyan] {query_text}") + + try: + query = UserQuery(chat_message=query_text) + response = agent.run(query) + + console.print(f"[bold green]Response:[/bold green] {response.chat_message}") + console.print(f"[bold yellow]Confidence:[/bold yellow] {response.confidence:.2f}") + console.print(f"[bold magenta]Reasoning:[/bold magenta] {response.reasoning}") + + except Exception as e: + console.print(f"[bold red]Error processing query:[/bold red] {e}") + + display_metrics() + + +def demonstrate_validation_errors(): + console.print(Panel("🚨 Validation Error Handling Demonstration", style="bold red")) + + agent = create_agent_with_hooks( + DetailedResponse, + """You are a helpful assistant. INTENTIONALLY use invalid values to test validation: + - Set confidence_level to something other than 'low', 'medium', or 'high' (like 'very_high' or 'uncertain') + - This is for testing validation error handling, so please violate the schema constraints intentionally.""", + ) + + validation_test_queries = [ + "Give me a simple yes or no answer about whether the sky is blue.", + "Provide a complex analysis of climate change with multiple perspectives.", + ] + + for query_text in validation_test_queries: + console.print(f"\n[bold cyan]Query:[/bold cyan] {query_text}") + + try: + query = UserQuery(chat_message=query_text) + response = agent.run(query) + + console.print(f"[bold green]Main Answer:[/bold green] {response.chat_message}") + console.print(f"[bold yellow]Confidence Level:[/bold yellow] {response.confidence_level}") + console.print(f"[bold magenta]Alternatives:[/bold magenta] {response.alternative_suggestions}") + console.print(f"[bold cyan]Needs Follow-up:[/bold cyan] {response.requires_followup}") + + except Exception as e: + console.print(f"[bold red]Handled error:[/bold red] {e}") + + display_metrics() + + +def demonstrate_interactive_mode(): + console.print(Panel("🎮 Interactive Hook System Testing", style="bold magenta")) + + agent = create_agent_with_hooks( + AgentResponse, "You are a helpful assistant. Provide clear, confident responses with reasoning." + ) + + console.print("[bold green]Welcome to the interactive hook system demo![/bold green]") + console.print("Type your questions below. Use /metrics to see performance data, /exit to quit.") + + while True: + try: + user_input = console.input("\n[bold blue]Your question:[/bold blue] ") + + if user_input.lower() in ["/exit", "/quit"]: + console.print("Exiting interactive mode...") + break + elif user_input.lower() == "/metrics": + display_metrics() + continue + elif user_input.strip() == "": + continue + + query = UserQuery(chat_message=user_input) + start_time = time.time() + + response = agent.run(query) + + response_time = time.time() - start_time + + console.print(f"\n[bold green]Answer:[/bold green] {response.chat_message}") + console.print(f"[bold yellow]Confidence:[/bold yellow] {response.confidence:.2f}") + console.print(f"[bold magenta]Reasoning:[/bold magenta] {response.reasoning}") + console.print(f"[dim]Response time: {response_time:.2f}s[/dim]") + + except KeyboardInterrupt: + console.print("\nExiting on user interrupt...") + break + except Exception as e: + console.print(f"[bold red]Error:[/bold red] {e}") + + +def main(): + console.print(Panel.fit("🎯 AtomicAgent Hook System Comprehensive Demo", style="bold green")) + + console.print( + """ +[bold cyan]This demonstration showcases:[/bold cyan] +• 🔍 Comprehensive monitoring with hooks +• 🛡️ Robust error handling and validation +• 📊 Real-time performance metrics +• 🔄 Production-ready patterns + +[bold yellow]The hook system provides zero-overhead monitoring when hooks aren't registered, +and powerful insights when they are enabled.[/bold yellow] + """ + ) + + try: + demonstrate_basic_hooks() + console.print("\n" + "=" * 50) + demonstrate_validation_errors() + console.print("\n" + "=" * 50) + demonstrate_interactive_mode() + + except KeyboardInterrupt: + console.print("\n[bold yellow]Demo interrupted by user.[/bold yellow]") + except Exception as e: + console.print(f"\n[bold red]Demo error:[/bold red] {e}") + logger.error(f"Demo error: {e}", exc_info=True) + finally: + console.print("\n" + "=" * 50) + console.print(Panel("📊 Final Performance Summary", style="bold green")) + display_metrics() + + console.print( + """ +[bold green]✅ Hook system demonstration complete![/bold green] + +[bold cyan]Key takeaways:[/bold cyan] +• Hooks provide comprehensive monitoring without performance overhead +• Error handling is robust and provides detailed context +• Metrics collection enables performance optimization +• The system is production-ready and scalable + +[bold yellow]Next steps:[/bold yellow] +• Implement custom retry logic in hook handlers +• Add monitoring service integration +• Explore advanced error recovery patterns +• Build custom metrics dashboards + """ + ) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/examples/code/atomic_agents/mainBaseModal.py b/examples/code/atomic_agents/mainBaseModal.py new file mode 100644 index 0000000..a7d77b3 --- /dev/null +++ b/examples/code/atomic_agents/mainBaseModal.py @@ -0,0 +1,141 @@ +from atomic_agents import AtomicAgent, AgentConfig, BaseIOSchema +from atomic_agents.context import SystemPromptGenerator +import instructor +import openai +from pydantic import Field +from typing import List +import os + +# API Key setup +API_KEY = "" +if not API_KEY: + API_KEY = os.getenv("OPENAI_API_KEY") + +if not API_KEY: + raise ValueError( + "API key is not set. Please set the API key as a static variable or in the environment variable OPENAI_API_KEY." + ) + + +class NutritionLabel(BaseIOSchema): + """Represents the complete nutritional information from a food label""" + + calories: int = Field(..., description="Calories per serving") + total_fat: float = Field(..., description="Total fat in grams") + saturated_fat: float = Field(..., description="Saturated fat in grams") + trans_fat: float = Field(..., description="Trans fat in grams") + cholesterol: int = Field(..., description="Cholesterol in milligrams") + sodium: int = Field(..., description="Sodium in milligrams") + total_carbohydrates: float = Field(..., description="Total carbohydrates in grams") + dietary_fiber: float = Field(..., description="Dietary fiber in grams") + total_sugars: float = Field(..., description="Total sugars in grams") + added_sugars: float = Field(..., description="Added sugars in grams") + protein: float = Field(..., description="Protein in grams") + vitamin_d: float = Field(..., description="Vitamin D in micrograms") + calcium: int = Field(..., description="Calcium in milligrams") + iron: float = Field(..., description="Iron in milligrams") + potassium: int = Field(..., description="Potassium in milligrams") + serving_size: str = Field(..., description="The size of a single serving of this product") + servings_per_container: float = Field(..., description="Number of servings contained in the package") + product_name: str = Field( + ..., + description="The full name or description of the type of the food/drink. e.g: 'Coca Cola Light', 'Pepsi Max', 'Smoked Bacon', 'Chianti Wine'", + ) + + +class NutritionAnalysisInput(BaseIOSchema): + """Input schema for nutrition label analysis""" + + instruction_text: str = Field(..., description="The instruction for analyzing the nutrition label") + images: List[instructor.Image] = Field(..., description="The nutrition label images to analyze") + + +class NutritionAnalysisOutput(BaseIOSchema): + """Output schema containing extracted nutrition information""" + + analyzed_labels: List[NutritionLabel] = Field( + ..., description="List of nutrition labels extracted from the provided images" + ) + + +# Configure the nutrition analysis system +nutrition_analyzer = AtomicAgent[NutritionAnalysisInput, NutritionAnalysisOutput]( + config=AgentConfig( + client=instructor.from_openai(openai.OpenAI(api_key=API_KEY)), + model="gpt-5-mini", + model_api_parameters={"reasoning_effort": "low"}, + system_prompt_generator=SystemPromptGenerator( + background=[ + "You are a specialized nutrition label analyzer.", + "You excel at extracting precise nutritional information from food label images.", + "You understand various serving size formats and measurement units.", + "You can process multiple nutrition labels simultaneously.", + ], + steps=[ + "For each nutrition label image:", + "1. Locate and identify the nutrition facts panel", + "2. Extract all serving information and nutritional values", + "3. Validate measurements and units for accuracy", + "4. Compile the nutrition facts into structured data", + ], + output_instructions=[ + "For each analyzed nutrition label:", + "1. Record complete serving size information", + "2. Extract all nutrient values with correct units", + "3. Ensure all measurements are properly converted", + "4. Include all extracted labels in the final result", + ], + ), + ) +) + + +def main(): + print("Starting nutrition label analysis...") + + # Construct the path to the test images + script_directory = os.path.dirname(os.path.abspath(__file__)) + test_images_directory = os.path.join(os.path.dirname(script_directory), "test_images") + image_path_1 = os.path.join(test_images_directory, "nutrition_label_1.png") + image_path_2 = os.path.join(test_images_directory, "nutrition_label_2.jpg") + # Create and submit the analysis request + analysis_request = NutritionAnalysisInput( + instruction_text="Please analyze these nutrition labels and extract all nutritional information.", + images=[instructor.Image.from_path(image_path_1), instructor.Image.from_path(image_path_2)], + ) + + try: + # Process the nutrition labels + print("Analyzing nutrition labels...") + analysis_result = nutrition_analyzer.run(analysis_request) + print("Analysis completed successfully") + + # Display the results + for i, label in enumerate(analysis_result.analyzed_labels, 1): + print(f"\nNutrition Label {i}:") + print(f"Product Name: {label.product_name}") + print(f"Serving Size: {label.serving_size}") + print(f"Servings Per Container: {label.servings_per_container}") + print(f"Calories: {label.calories}") + print(f"Total Fat: {label.total_fat}g") + print(f"Saturated Fat: {label.saturated_fat}g") + print(f"Trans Fat: {label.trans_fat}g") + print(f"Cholesterol: {label.cholesterol}mg") + print(f"Sodium: {label.sodium}mg") + print(f"Total Carbohydrates: {label.total_carbohydrates}g") + print(f"Dietary Fiber: {label.dietary_fiber}g") + print(f"Total Sugars: {label.total_sugars}g") + print(f"Added Sugars: {label.added_sugars}g") + print(f"Protein: {label.protein}g") + print(f"Vitamin D: {label.vitamin_d}mcg") + print(f"Calcium: {label.calcium}mg") + print(f"Iron: {label.iron}mg") + print(f"Potassium: {label.potassium}mg") + + except Exception as e: + print(f"Analysis failed: {str(e)}") + raise + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/examples/code/atomic_agents/mainPDFAnalysis.py b/examples/code/atomic_agents/mainPDFAnalysis.py new file mode 100644 index 0000000..43cc5a6 --- /dev/null +++ b/examples/code/atomic_agents/mainPDFAnalysis.py @@ -0,0 +1,82 @@ +import os + +import instructor +from atomic_agents import AtomicAgent, AgentConfig, BaseIOSchema +from atomic_agents.context import SystemPromptGenerator +from dotenv import load_dotenv +from google import genai +from instructor.multimodal import PDF +from pydantic import Field + + +load_dotenv() + + +class InputSchema(BaseIOSchema): + """PDF file to analyze.""" + + pdf: PDF = Field(..., description="The PDF data") # PDF class from instructor + + +class ExtractionResult(BaseIOSchema): + """Extracted information from the PDF.""" + + pdf_title: str = Field(..., description="The title of the PDF file") + page_count: int = Field(..., description="The number of pages in the PDF file") + summary: str = Field(..., description="A short summary of the document") + + +# Define the LLM CLient using GenAI instructor wrapper: +client = instructor.from_genai(client=genai.Client(api_key=os.getenv("GEMINI_API_KEY")), mode=instructor.Mode.GENAI_TOOLS) + +# Define the system prompt: +system_prompt_generator = SystemPromptGenerator( + background=["You are a helpful assistant that extracts information from PDF files."], + steps=[ + "Analyze the PDF, extract its title and count the number of pages.", + "Create a brief summary of the document content.", + ], + output_instructions=["Return pdf_title, page_count, and summary."], +) + +# Define the agent +agent = AtomicAgent[InputSchema, ExtractionResult]( + config=AgentConfig( + client=client, + model="gemini-2.0-flash", + system_prompt_generator=system_prompt_generator, + input_schema=InputSchema, + output_schema=ExtractionResult, + ) +) + + +def main(): + print("Starting PDF file analysis...") + + # Create the analysis request + script_directory = os.path.dirname(os.path.abspath(__file__)) + test_media_directory = os.path.join(os.path.dirname(script_directory), "test_media") + pdf_path = os.path.join(test_media_directory, "pdf_sample.pdf") + analysis_request = InputSchema( + pdf=PDF.from_path(pdf_path), + ) + + try: + # Process the PDF file + print(f"Analyzing PDF file: {os.path.basename(pdf_path)} ...") + analysis_result = agent.run(analysis_request) + + # Display the results + print("\n===== Analysis Results =====") + print(f"PDF Title: {analysis_result.pdf_title}") + print(f"Page Count: {analysis_result.page_count}") + print(f"Document summary: {analysis_result.summary}") + + except Exception as e: + print(f"Analysis failed: {str(e)}") + raise e + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/examples/code/atomic_agents/orchestor.py b/examples/code/atomic_agents/orchestor.py new file mode 100644 index 0000000..2249cec --- /dev/null +++ b/examples/code/atomic_agents/orchestor.py @@ -0,0 +1,181 @@ +from typing import Union +import openai +from pydantic import Field +from atomic_agents import AtomicAgent, AgentConfig, BaseIOSchema +from atomic_agents.context import SystemPromptGenerator, BaseDynamicContextProvider + +from orchestration_agent.tools.searxng_search import ( + SearXNGSearchTool, + SearXNGSearchToolConfig, + SearXNGSearchToolInputSchema, + SearXNGSearchToolOutputSchema, +) +from orchestration_agent.tools.calculator import ( + CalculatorTool, + CalculatorToolConfig, + CalculatorToolInputSchema, + CalculatorToolOutputSchema, +) + +import instructor +from datetime import datetime + + +######################## +# INPUT/OUTPUT SCHEMAS # +######################## +class OrchestratorInputSchema(BaseIOSchema): + """Input schema for the Orchestrator Agent. Contains the user's message to be processed.""" + + chat_message: str = Field(..., description="The user's input message to be analyzed and responded to.") + + +class OrchestratorOutputSchema(BaseIOSchema): + """Combined output schema for the Orchestrator Agent. Contains the tool parameters.""" + + tool_parameters: Union[SearXNGSearchToolInputSchema, CalculatorToolInputSchema] = Field( + ..., description="The parameters for the selected tool" + ) + + +class FinalAnswerSchema(BaseIOSchema): + """Schema for the final answer generated by the Orchestrator Agent.""" + + final_answer: str = Field(..., description="The final answer generated based on the tool output and user query.") + + +####################### +# AGENT CONFIGURATION # +####################### +class OrchestratorAgentConfig(AgentConfig): + """Configuration for the Orchestrator Agent.""" + + searxng_config: SearXNGSearchToolConfig + calculator_config: CalculatorToolConfig + + +##################### +# CONTEXT PROVIDERS # +##################### +class CurrentDateProvider(BaseDynamicContextProvider): + def __init__(self, title): + super().__init__(title) + self.date = datetime.now().strftime("%Y-%m-%d") + + def get_info(self) -> str: + return f"Current date in format YYYY-MM-DD: {self.date}" + + +###################### +# ORCHESTRATOR AGENT # +###################### +orchestrator_agent_config = AgentConfig( + client=instructor.from_openai(openai.OpenAI()), + model="gpt-5-mini", + model_api_parameters={"reasoning_effort": "low"}, + system_prompt_generator=SystemPromptGenerator( + background=[ + "You are an Orchestrator Agent that decides between using a search tool or a calculator tool based on user input.", + "Use the search tool for queries requiring factual information, current events, or specific data.", + "Use the calculator tool for mathematical calculations and expressions.", + ], + output_instructions=[ + "Analyze the input to determine whether it requires a web search or a calculation.", + "For search queries, use the 'search' tool and provide 1-3 relevant search queries.", + "For calculations, use the 'calculator' tool and provide the mathematical expression to evaluate.", + "When uncertain, prefer using the search tool.", + "Format the output using the appropriate schema.", + ], + ), +) +orchestrator_agent = AtomicAgent[OrchestratorInputSchema, OrchestratorOutputSchema](config=orchestrator_agent_config) +orchestrator_agent_final = AtomicAgent[OrchestratorInputSchema, FinalAnswerSchema](config=orchestrator_agent_config) + +# Register the current date provider +orchestrator_agent.register_context_provider("current_date", CurrentDateProvider("Current Date")) +orchestrator_agent_final.register_context_provider("current_date", CurrentDateProvider("Current Date")) + + +def execute_tool( + searxng_tool: SearXNGSearchTool, calculator_tool: CalculatorTool, orchestrator_output: OrchestratorOutputSchema +) -> Union[SearXNGSearchToolOutputSchema, CalculatorToolOutputSchema]: + if isinstance(orchestrator_output.tool_parameters, SearXNGSearchToolInputSchema): + return searxng_tool.run(orchestrator_output.tool_parameters) + elif isinstance(orchestrator_output.tool_parameters, CalculatorToolInputSchema): + return calculator_tool.run(orchestrator_output.tool_parameters) + else: + raise ValueError(f"Unknown tool parameters type: {type(orchestrator_output.tool_parameters)}") + + +################# +# EXAMPLE USAGE # +################# +if __name__ == "__main__": + import os + from dotenv import load_dotenv + from rich.console import Console + from rich.panel import Panel + from rich.syntax import Syntax + + load_dotenv() + + # Set up the OpenAI client + client = instructor.from_openai(openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))) + + # Initialize the tools + searxng_tool = SearXNGSearchTool(SearXNGSearchToolConfig(base_url="http://localhost:8080", max_results=5)) + calculator_tool = CalculatorTool(CalculatorToolConfig()) + + # Initialize Rich console + console = Console() + + # Print the full system prompt + console.print(Panel(orchestrator_agent.system_prompt_generator.generate_prompt(), title="System Prompt", expand=False)) + console.print("\n") + + # Example inputs + inputs = [ + "Who won the Nobel Prize in Physics in 2024?", + "Please calculate the sine of pi/3 to the third power", + ] + + for user_input in inputs: + console.print(Panel(f"[bold cyan]User Input:[/bold cyan] {user_input}", expand=False)) + + # Create the input schema + input_schema = OrchestratorInputSchema(chat_message=user_input) + + # Print the input schema + console.print("\n[bold yellow]Generated Input Schema:[/bold yellow]") + input_syntax = Syntax(str(input_schema.model_dump_json(indent=2)), "json", theme="monokai", line_numbers=True) + console.print(input_syntax) + + # Run the orchestrator to get the tool selection and input + orchestrator_output = orchestrator_agent.run(input_schema) + + # Print the orchestrator output + console.print("\n[bold magenta]Orchestrator Output:[/bold magenta]") + orchestrator_syntax = Syntax( + str(orchestrator_output.model_dump_json(indent=2)), "json", theme="monokai", line_numbers=True + ) + console.print(orchestrator_syntax) + + # Run the selected tool + response = execute_tool(searxng_tool, calculator_tool, orchestrator_output) + + # Print the tool output + console.print("\n[bold green]Tool Output:[/bold green]") + output_syntax = Syntax(str(response.model_dump_json(indent=2)), "json", theme="monokai", line_numbers=True) + console.print(output_syntax) + + console.print("\n" + "-" * 80 + "\n") + + # Switch agent + history = orchestrator_agent.history + orchestrator_agent = orchestrator_agent_final + orchestrator_agent.history = history + orchestrator_agent.history.add_message("system", response) + final_answer = orchestrator_agent.run(input_schema) + console.print(f"\n[bold blue]Final Answer:[/bold blue] {final_answer.final_answer}") + # Reset the agent to the original + orchestrator_agent = AtomicAgent[OrchestratorInputSchema, OrchestratorOutputSchema](config=orchestrator_agent_config) \ No newline at end of file diff --git a/examples/code/atomic_agents/tranScriptTool.py b/examples/code/atomic_agents/tranScriptTool.py new file mode 100644 index 0000000..7661be7 --- /dev/null +++ b/examples/code/atomic_agents/tranScriptTool.py @@ -0,0 +1,47 @@ +import os +from dotenv import load_dotenv +from rich.console import Console + +from youtube_summarizer.tools.youtube_transcript_scraper import ( + YouTubeTranscriptTool, + YouTubeTranscriptToolConfig, + YouTubeTranscriptToolInputSchema, +) + +from youtube_summarizer.agent import ( + YouTubeKnowledgeExtractionInputSchema, + youtube_knowledge_extraction_agent, + transcript_provider, +) + +load_dotenv() + +# Initialize a Rich Console for pretty console outputs +console = Console() + +# Initialize the YouTubeTranscriptTool +transcript_tool = YouTubeTranscriptTool(config=YouTubeTranscriptToolConfig(api_key=os.getenv("YOUTUBE_API_KEY"))) + +# Remove the infinite loop and perform a one-time transcript extraction +video_url = "https://www.youtube.com/watch?v=Sp30YsjGUW0" + +transcript_input = YouTubeTranscriptToolInputSchema(video_url=video_url, language="en") +try: + transcript_output = transcript_tool.run(transcript_input) + console.print(f"[bold green]Transcript:[/bold green] {transcript_output.transcript}") + console.print(f"[bold green]Duration:[/bold green] {transcript_output.duration} seconds") + + # Update transcript_provider with the scraped transcript data + transcript_provider.transcript = transcript_output.transcript + transcript_provider.duration = transcript_output.duration + transcript_provider.metadata = transcript_output.metadata # Assuming metadata is available in transcript_output + + # Run the transcript through the agent + transcript_input_schema = YouTubeKnowledgeExtractionInputSchema(video_url=video_url) + agent_response = youtube_knowledge_extraction_agent.run(transcript_input_schema) + + # Print the output schema in a formatted way + console.print("[bold blue]Agent Output Schema:[/bold blue]") + console.print(agent_response) +except Exception as e: + console.print(f"[bold red]Error:[/bold red] {str(e)}") \ No newline at end of file diff --git a/examples/code/atomic_agents/webSearchAgent.py b/examples/code/atomic_agents/webSearchAgent.py new file mode 100644 index 0000000..b524e66 --- /dev/null +++ b/examples/code/atomic_agents/webSearchAgent.py @@ -0,0 +1,111 @@ +import os +from dotenv import load_dotenv +from rich.console import Console +from rich.markdown import Markdown +from pydantic import Field + +from atomic_agents import BaseIOSchema +from atomic_agents.context import ChatHistory, BaseDynamicContextProvider + +from web_search_agent.tools.searxng_search import ( + SearXNGSearchTool, + SearXNGSearchToolConfig, + SearXNGSearchToolInputSchema, + SearXNGSearchToolOutputSchema, +) + +from web_search_agent.agents.query_agent import QueryAgentInputSchema, query_agent +from web_search_agent.agents.question_answering_agent import question_answering_agent, QuestionAnsweringAgentInputSchema + + +load_dotenv() + +# Initialize a Rich Console for pretty console outputs +console = Console() + +# History setup +history = ChatHistory() + +# Initialize the SearXNGSearchTool +search_tool = SearXNGSearchTool(config=SearXNGSearchToolConfig(base_url=os.getenv("SEARXNG_BASE_URL"), max_results=5)) + + +class SearchResultsProvider(BaseDynamicContextProvider): + def __init__(self, title: str, search_results: SearXNGSearchToolOutputSchema | Exception): + super().__init__(title=title) + self.search_results = search_results + + def get_info(self) -> str: + return f"{self.title}: {self.search_results}" + + +# Define input/output schemas for the main agent +class MainAgentInputSchema(BaseIOSchema): + """Input schema for the main agent.""" + + chat_message: str = Field(..., description="Chat message from the user.") + + +class MainAgentOutputSchema(BaseIOSchema): + """Output schema for the main agent.""" + + chat_message: str = Field(..., description="Response to the user's message.") + + +# Example usage +instruction = "Tell me about the Atomic Agents AI agent framework." +num_queries = 3 +console.print(f"[bold blue]Instruction:[/bold blue] {instruction}") + +while True: + # Generate queries using the query agent + query_input = QueryAgentInputSchema(instruction=instruction, num_queries=num_queries) + generated_queries = query_agent.run(query_input) + + console.print("[bold blue]Generated Queries:[/bold blue]") + for query in generated_queries.queries: + console.print(f"- {query}") + + # Perform searches using the generated queries + search_input = SearXNGSearchToolInputSchema(queries=generated_queries.queries, category="general") + + try: + search_results = search_tool.run(search_input) + search_results_provider = SearchResultsProvider("Search Results", search_results) + except Exception as e: + search_results_provider = SearchResultsProvider("Search Failed", e) + + question_answering_agent.register_context_provider("search results", search_results_provider) + + answer = question_answering_agent.run(QuestionAnsweringAgentInputSchema(question=instruction)) + + # Create a Rich Console instance + console = Console() + + # Print the answer using Rich's Markdown rendering + console.print("\n[bold blue]Answer:[/bold blue]") + console.print(Markdown(answer.markdown_output)) + + # Print references + console.print("\n[bold blue]References:[/bold blue]") + for ref in answer.references: + console.print(f"- {ref}") + + # Print follow-up questions + console.print("\n[bold blue]Follow-up Questions:[/bold blue]") + for i, question in enumerate(answer.followup_questions, 1): + console.print(f"[cyan]{i}. {question}[/cyan]") + + console.print() # Add an empty line for better readability + instruction = console.input("[bold blue]You:[/bold blue] ") + if instruction.lower() in ["/exit", "/quit"]: + console.print("Exiting chat...") + break + + try: + followup_question_id = int(instruction.strip()) + if 1 <= followup_question_id <= len(answer.followup_questions): + instruction = answer.followup_questions[followup_question_id - 1] + console.print(f"[bold blue]Follow-up Question:[/bold blue] {instruction}") + except ValueError: + pass \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 9206bd1..ce40be1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ dependencies = [ "openai", "python-dotenv", "pyyaml", + "dotenv" ] [project.scripts] diff --git a/src/repello_agent_wiz/cli.py b/src/repello_agent_wiz/cli.py index 37f08c5..a398140 100644 --- a/src/repello_agent_wiz/cli.py +++ b/src/repello_agent_wiz/cli.py @@ -1,7 +1,5 @@ import argparse -import os -import sys -from .frameworks import agent_chat, autogen, crewai, google_adk, langgraph, llama_index, n8n, openai_agents, pydantic, swarm +from .frameworks import agent_chat, autogen, crewai, google_adk, langgraph, llama_index, n8n, openai_agents, pydantic, swarm , atomic_agents from .analyzers import generate_maestro_analysis_report from .visualizers.visualizer import generate_visualization @@ -14,7 +12,7 @@ def main(): extract_parser = subparsers.add_parser("extract", help="Extract graph from source code") extract_parser.add_argument("--framework", "-f", required=True, choices=[ "agent_chat", "autogen", "crewai", "google_adk" , "langgraph", - "llama_index", "n8n", "openai_agents", "pydantic", "swarm" + "llama_index", "n8n", "openai_agents", "pydantic", "swarm" , "atomic_agents" ]) extract_parser.add_argument("--directory", "-d", default=".", help="Directory containing source code") extract_parser.add_argument("--output", "-o", default="graph.json", help="Path to output JSON file") @@ -30,7 +28,7 @@ def main(): # --- Parse and execute --- - args = parser.parse_args() + args = parser.parse_args() match args.command: case "extract": @@ -67,6 +65,8 @@ def main(): n8n.extract_n8n_graph(args.directory, args.output) case "openai_agents": openai_agents.extract_openai_agents_graph(args.directory, args.output) + case "atomic_agents": + atomic_agents.extract_atomic_agents_graph(args.directory, args.output) case _: print(f"Unknown framework: {args.framework}") diff --git a/src/repello_agent_wiz/frameworks/atomic_agents/__init__.py b/src/repello_agent_wiz/frameworks/atomic_agents/__init__.py new file mode 100644 index 0000000..27d2d3c --- /dev/null +++ b/src/repello_agent_wiz/frameworks/atomic_agents/__init__.py @@ -0,0 +1 @@ +from .mapper import extract_atomic_agents_graph \ No newline at end of file diff --git a/src/repello_agent_wiz/frameworks/atomic_agents/mapper.py b/src/repello_agent_wiz/frameworks/atomic_agents/mapper.py new file mode 100644 index 0000000..004406e --- /dev/null +++ b/src/repello_agent_wiz/frameworks/atomic_agents/mapper.py @@ -0,0 +1,552 @@ +import ast +import json +import os +import sys +from typing import Any, Dict, List, Optional, Set, Tuple, Union +from pathlib import Path +from pydantic import BaseModel +from enum import Enum + +class NodeType(str, Enum): + AGENT = "Agent" + TOOL = "Tool" + CUSTOM_TOOL = "CustomTool" + START = "Start" + END = "End" + +class ToolDefinition(BaseModel): + name: str + custom: bool + description: Optional[str] = None + +class AtomicAgentMapper(ast.NodeVisitor): + DECORATOR_IDENTS = {"function_tool", "atomic.beta.tools.function_tool"} + CONSTRUCTOR_IDENTS = {"AtomicAgent", "SearchTool", "FunctionTool"} + TOOL_EXECUTE_PATTERNS = {"execute", "run_sync", "run_async", "run"} + + def __init__(self, filepath: str) -> None: + super().__init__() + self.current_filepath = filepath + self.discovered_tools: dict[str, ToolDefinition] = {} + self.tool_locations: dict[str, Dict] = {} + self.agents: dict[str, Dict] = {} + self.tool_calls: list[Dict] = [] + + def visit_Assign(self, node: ast.Assign): + if isinstance(node.value, ast.Call) and is_matching_call(node.value, self.CONSTRUCTOR_IDENTS): + call_node = node.value + constructor_name = get_identifier_string(call_node.func) + + if constructor_name == "AtomicAgent": + self._process_atomic_agent_creation(node, call_node) + else: + self._process_tool_creation(node, call_node, constructor_name) + + self.generic_visit(node) + + def _process_atomic_agent_creation(self, node: ast.Assign, call_node: ast.Call): + """Process AtomicAgent creation""" + for target in node.targets: + target_name = get_identifier_string(target) + if target_name: + location = create_location_info(node, self.current_filepath) + config = self._extract_agent_config(call_node) + + tools_in_config = self._extract_tools_from_config(call_node) + for tool_name in tools_in_config: + if tool_name not in self.discovered_tools: + tool_def = ToolDefinition(name=tool_name, custom=True, description=f"Tool used by {target_name}") + self.discovered_tools[tool_name] = tool_def + self.tool_locations[tool_name] = location + + agent_info = { + "name": target_name, + "type": "AtomicAgent", + "location": location, + "config": config, + "tools": tools_in_config + } + self.agents[target_name] = agent_info + + def _process_tool_creation(self, node: ast.Assign, call_node: ast.Call, constructor_name: str): + """Process tool creation (SearchTool, FunctionTool, etc.)""" + tool_name = extract_kwarg_string(call_node, "name") + if not tool_name: + for target in node.targets: + target_name = get_identifier_string(target) + if target_name: + tool_name = target_name + break + + if tool_name: + description = extract_kwarg_string(call_node, "description") or "" + location = create_location_info(node, self.current_filepath) + tool_def = ToolDefinition(name=tool_name, custom=True, description=description) + + for target in node.targets: + target_name = get_identifier_string(target) + if target_name: + self.discovered_tools[target_name] = tool_def + if location: self.tool_locations[target_name] = location + + def _extract_agent_config(self, call_node: ast.Call) -> Dict: + """Extract configuration from AtomicAgent call""" + config = {} + for kw in call_node.keywords: + if kw.arg == "config" and isinstance(kw.value, ast.Call): + config_node = kw.value + for config_kw in config_node.keywords: + if config_kw.arg: + config[config_kw.arg] = represent_node(config_kw.value) + return config + + def _extract_tools_from_config(self, call_node: ast.Call) -> List[str]: + """Extract tool names from AgentConfig""" + tool_names = [] + for kw in call_node.keywords: + if kw.arg == "config" and isinstance(kw.value, ast.Call): + config_node = kw.value + for config_kw in config_node.keywords: + if config_kw.arg == "tools" and isinstance(config_kw.value, ast.List): + for tool_node in config_kw.value.elts: + tool_name = get_identifier_string(tool_node) + if tool_name: + tool_names.append(tool_name) + return tool_names + + def visit_ClassDef(self, node: ast.ClassDef): + """Detect tool class definitions like SearchTool""" + class_name = node.name + if self._is_tool_class(node): + location = create_location_info(node, self.current_filepath) + tool_def = ToolDefinition(name=class_name, custom=True, description=ast.get_docstring(node) or "") + self.discovered_tools[class_name] = tool_def + if location: self.tool_locations[class_name] = location + + self.generic_visit(node) + + def visit_Call(self, node: ast.Call): + """Detect tool execution patterns like tool.execute() and agent.run_sync()""" + if isinstance(node.func, ast.Attribute): + method_name = node.func.attr + if method_name in self.TOOL_EXECUTE_PATTERNS: + obj_name = get_identifier_string(node.func.value) + if obj_name: + location = create_location_info(node, self.current_filepath) + call_info = { + "object": obj_name, + "method": method_name, + "location": location, + "args": [represent_node(arg) for arg in node.args] + } + self.tool_calls.append(call_info) + + self.generic_visit(node) + + def _is_tool_class(self, node: ast.ClassDef) -> bool: + """Check if a class is a tool class""" + for item in node.body: + if isinstance(item, ast.FunctionDef) and item.name == "execute": + return True + if isinstance(item, ast.Assign): + for target in item.targets: + if isinstance(target, ast.Name) and target.id == "input_schema": + return True + return False + + def visit_FunctionDef(self, node): self._process_function(node) + def visit_AsyncFunctionDef(self, node): self._process_function(node) + + def _process_function(self, node: Union[ast.FunctionDef, ast.AsyncFunctionDef]): + decorator_node = find_decorator_node(node, self.DECORATOR_IDENTS) + if not decorator_node: + self.generic_visit(node) + return + + func_name = node.name + tool_name = func_name + description = ast.get_docstring(node) or "" + location = create_location_info(node, self.current_filepath) + + if isinstance(decorator_node, ast.Call): + name_override = extract_kwarg_string(decorator_node, "name_override") + if name_override: tool_name = name_override + desc_override = extract_kwarg_string(decorator_node, "description_override") + if desc_override: description = desc_override + + tool_def = ToolDefinition(name=tool_name, custom=True, description=description) + self.discovered_tools[func_name] = tool_def + if location: + self.tool_locations[tool_name] = location + if tool_name != func_name: + self.tool_locations[func_name] = location + self.generic_visit(node) + +def find_decorator_node(node: Union[ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef], target_names: Set[str]) -> Optional[ast.AST]: + for decorator in node.decorator_list: + deco_node_to_check = decorator + if isinstance(decorator, ast.Call): + deco_node_to_check = decorator.func + deco_name = get_identifier_string(deco_node_to_check) + if deco_name and deco_name in target_names: + return decorator + return None + +def is_matching_call(node: ast.AST, target_names: Set[str]) -> bool: + if not isinstance(node, ast.Call): return False + func_node = node.func + base_name = get_identifier_string(func_node) + if base_name and base_name in target_names: + return True + if isinstance(func_node, ast.Attribute) and func_node.attr in target_names: + return True + return False + +def get_identifier_string(node: ast.AST) -> Optional[str]: + if isinstance(node, ast.Name): return node.id + elif isinstance(node, ast.Attribute): + try: + if hasattr(ast, 'unparse'): return ast.unparse(node) + base = get_identifier_string(node.value) + return f"{base}.{node.attr}" if base else node.attr + except: return node.attr + elif isinstance(node, ast.Subscript): + return get_identifier_string(node.value) + return None + +def get_kwarg_node(call_node: ast.Call, keyword: str) -> Optional[ast.AST]: + if not isinstance(call_node, ast.Call): raise TypeError("Expected an ast.Call node") + for kw in call_node.keywords: + if kw.arg == keyword: return kw.value + return None + +def extract_string_literal(node: Optional[ast.AST]) -> Optional[str]: + if isinstance(node, ast.Constant) and isinstance(node.value, str): + return node.value + return None + +def extract_kwarg_string(call_node: ast.Call, keyword: str) -> Optional[str]: + value_node = get_kwarg_node(call_node, keyword) + return extract_string_literal(value_node) + +def get_qualified_name(node: Union[ast.Name, ast.Attribute, ast.Call, ast.Subscript]) -> str: + if isinstance(node, ast.Name): return node.id + elif isinstance(node, ast.Attribute): + base = get_qualified_name(node.value) + return f"{base}.{node.attr}" if base else node.attr + elif isinstance(node, ast.Call): return get_qualified_name(node.func) + elif isinstance(node, ast.Subscript): return get_qualified_name(node.value) + return "" + +def represent_node(node: Optional[ast.AST]) -> str: + if node is None: return "None" + if isinstance(node, ast.Constant): return repr(node.value) + if isinstance(node, ast.Name): return node.id + if isinstance(node, ast.Attribute): + name = get_qualified_name(node) + return name if name else f"" + if isinstance(node, ast.Subscript): + base = represent_node(node.value) + slice_val = represent_node(node.slice) + return f"{base}[{slice_val}]" + if isinstance(node, ast.List): return f"[{', '.join(represent_node(elt) for elt in node.elts)}]" + if isinstance(node, ast.Call): + func_str = represent_node(node.func) + args_str = ', '.join(represent_node(arg) for arg in node.args) + kwargs_str = ', '.join(f"{kw.arg}={represent_node(kw.value)}" for kw in node.keywords if kw.arg) + all_args = f"{args_str}{', ' if args_str and kwargs_str else ''}{kwargs_str}" + return f"{func_str}({all_args})" + else: + try: + if hasattr(ast, 'unparse'): return ast.unparse(node) + return f"<{type(node).__name__}>" + except Exception: return f"<{type(node).__name__}>" + +def create_location_info(node: ast.AST, current_filepath: str) -> Optional[Dict[str, Any]]: + if not current_filepath or not hasattr(node, 'lineno'): return None + end_lineno = getattr(node, 'end_lineno', node.lineno) + end_col_offset = getattr(node, 'end_col_offset', -1) + col_offset = getattr(node, 'col_offset', -1) + return { + "file": current_filepath, + "line": node.lineno, "col": col_offset, + "end_line": end_lineno, "end_col": end_col_offset, + } + +def build_graph_json( + agents: dict[str, Dict], + custom_tools: dict[str, ToolDefinition], + custom_tool_locs: dict[str, Dict], + tool_calls: list[Dict] +) -> Dict[str, List[Dict]]: + """Build graph structure from extracted atomic agents data""" + nodes: List[Dict] = [] + edges: List[Dict] = [] + processed_ids: Set[str] = set() + + for agent_name, agent_info in agents.items(): + if agent_name in processed_ids: + continue + + node_meta = { + "type": agent_info.get("type"), + "config": agent_info.get("config", {}) + } + node_meta = {k: v for k, v in node_meta.items() if v is not None and v != {}} + + node = { + "id": agent_name, + "name": agent_name, + "node_type": NodeType.AGENT.value, + "source_location": agent_info.get("location"), + "metadata": node_meta + } + nodes.append(node) + processed_ids.add(agent_name) + + tool_refs = set() + for call in tool_calls: + obj_name = call.get("object") + if obj_name and obj_name in custom_tools: + tool_refs.add(obj_name) + + for agent_info in agents.values(): + agent_tools = agent_info.get("tools", []) + for tool_name in agent_tools: + if tool_name in custom_tools: + tool_refs.add(tool_name) + + for tool_id in tool_refs: + if tool_id in processed_ids: + continue + + tool_def = custom_tools.get(tool_id) + if not tool_def: + print(f"Warning: Tool '{tool_id}' referenced but definition missing. Skipping node.") + continue + + location = custom_tool_locs.get(tool_id) + tool_node = { + "id": tool_id, + "name": tool_id, + "function_name": tool_id, + "docstring": tool_def.description, + "node_type": NodeType.CUSTOM_TOOL.value, + "source_location": location, + "metadata": {"custom": True} + } + nodes.append(tool_node) + processed_ids.add(tool_id) + + for call in tool_calls: + obj_name = call.get("object") + method = call.get("method") + location = call.get("location") + + if obj_name in agents: + source_id = "Start" + target_id = obj_name + edges.append({ + "source": source_id, + "target": target_id, + "edge_type": "execution_start", + "condition": {}, + "metadata": { + "method": method, + "definition_location": location + } + }) + elif obj_name in custom_tools: + pass + + for agent_name, agent_info in agents.items(): + agent_tools = agent_info.get("tools", []) + agent_loc = agent_info.get("location") + for tool_name in agent_tools: + if tool_name in processed_ids: + edges.append({ + "source": agent_name, + "target": tool_name, + "edge_type": "tool_usage", + "condition": {}, + "metadata": { + "definition_location": agent_loc + } + }) + for call in tool_calls: + obj_name = call.get("object") + location = call.get("location") + + if obj_name in custom_tools and obj_name in processed_ids: + for agent_name in agents.keys(): + agent_loc = agents[agent_name].get("location") + if agent_loc and location and agent_loc.get("file") == location.get("file"): + edges.append({ + "source": agent_name, + "target": obj_name, + "edge_type": "tool_usage", + "condition": {}, + "metadata": { + "method": call.get("method"), + "definition_location": location + } + }) + break + start_id, end_id = "Start", "End" + if start_id not in processed_ids: + nodes.append({"id": "Start", "name": "Start", "node_type": NodeType.START.value}) + processed_ids.add(start_id) + if end_id not in processed_ids: + nodes.append({"id": "End", "name": "End", "node_type": NodeType.END.value}) + processed_ids.add(end_id) + + agent_ids = set(agents.keys()) + incoming: Dict[str, int] = {name: 0 for name in agent_ids} + outgoing: Dict[str, int] = {name: 0 for name in agent_ids} + + for edge in edges: + if edge.get("source") in agent_ids: + outgoing[edge["source"]] = outgoing.get(edge["source"], 0) + 1 + if edge.get("target") in agent_ids: + incoming[edge["target"]] = incoming.get(edge["target"], 0) + 1 + + for agent_id in agent_ids: + if incoming.get(agent_id, 0) == 0: + if not any(e["source"] == start_id and e["target"] == agent_id for e in edges): + edges.append({ + "source": start_id, + "target": agent_id, + "edge_type": "implicit_start", + "condition": {}, + "metadata": {} + }) + + for agent_id in agent_ids: + if not any(e["source"] == agent_id and e["target"] == end_id for e in edges): + edges.append({ + "source": agent_id, + "target": end_id, + "edge_type": "implicit_end", + "condition": {}, + "metadata": {} + }) + + tool_node_ids = [node["id"] for node in nodes if node.get("node_type") == NodeType.CUSTOM_TOOL.value] + for tool_id in tool_node_ids: + if not any(e["source"] == tool_id and e["target"] == end_id for e in edges): + edges.append({ + "source": tool_id, + "target": end_id, + "edge_type": "implicit_end", + "condition": {}, + "metadata": {} + }) + + for node in nodes: + node.setdefault("function_name", None) + node.setdefault("docstring", None) + node.setdefault("source_location", None) + node.setdefault("metadata", {}) + + for node in nodes: + if "metadata" in node: + node["metadata"] = {k: v for k, v in node["metadata"].items() if v is not None} + + return {"nodes": nodes, "edges": edges} + + +def extract_atomic_agents_graph(scan_path: str, output_file: str): + """Extract and save atomic agents graph structure to JSON""" + if not os.path.isdir(scan_path): + print(f"Error: Path '{scan_path}' is not a valid directory.") + sys.exit(1) + + try: + custom_tools, custom_locs, agents, tool_calls = gather_tool_definitions(scan_path) + + final_graph = build_graph_json( + agents=agents, + custom_tools=custom_tools, + custom_tool_locs=custom_locs, + tool_calls=tool_calls + ) + + if final_graph: + final_graph["metadata"] = { + "framework": "Atomic_Agents", + } + + except Exception as e: + print(f"Error during graph extraction: {e}") + import traceback + traceback.print_exc() + sys.exit(1) + + if final_graph["nodes"] or final_graph["edges"]: + try: + final_graph["nodes"].sort(key=lambda x: x.get('id', '')) + final_graph["edges"].sort(key=lambda x: (x.get('source', ''), x.get('target', ''), x.get('edge_type', ''))) + + output_path = Path(output_file) + output_path.parent.mkdir(parents=True, exist_ok=True) + + with open(output_path, "w", encoding='utf-8') as f: + json.dump(final_graph, f, indent=2) + print(f"Graph written to {output_file}") + + if not final_graph["nodes"]: + print("No nodes found.") + for node in final_graph["nodes"]: + loc = node.get('source_location') + loc_str = f" ({loc['file']}:{loc['line']})" if loc and loc.get('file') else "" + node_id = node.get('id', 'Unknown ID') + node_type = node.get('node_type', 'Unknown Type') + + if not final_graph["edges"]: + print("No edges found.") + for edge in final_graph["edges"]: + meta_loc = edge.get('metadata', {}).get('definition_location') + + except Exception as e: + print(f"Error writing JSON output or printing summary: {e}") + import traceback + traceback.print_exc() + sys.exit(1) + else: + print("\nNo graph structure found or extracted.") + try: + output_path = Path(output_file) + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w", encoding='utf-8') as f: + json.dump({"nodes": [], "edges": [], "metadata": {"framework": "Atomic_Agents"}}, f, indent=2) + print(f"Graph written to {output_file}") + except Exception as e: + print(f"Error writing empty JSON output to {output_path}: {e}") + sys.exit(1) + + +def gather_tool_definitions(root_path: str) -> Tuple[dict[str, ToolDefinition], dict[str, Dict], dict[str, Dict], list[Dict]]: + all_tools: dict[str, ToolDefinition] = {} + all_locations: dict[str, Dict] = {} + all_agents: dict[str, Dict] = {} + all_tool_calls: list[Dict] = [] + + for file_path in Path(root_path).rglob("*.py"): + filepath_str = str(file_path) + filename = file_path.name + print(f"Processing file: {filename}") + try: + with open(file_path, "r", encoding='utf-8') as f: + content = f.read() + tree = ast.parse(content, filename=filepath_str) + extractor = AtomicAgentMapper(filepath_str) + extractor.visit(tree) + + all_tools.update(extractor.discovered_tools) + all_locations.update(extractor.tool_locations) + all_agents.update(extractor.agents) + all_tool_calls.extend(extractor.tool_calls) + except SyntaxError as e: print(f"Warning: Skipping file {filepath_str} due to SyntaxError: {e}") + except Exception as e: print(f"Warning: Skipping file {filepath_str} due to unexpected error: {e}") + + return all_tools, all_locations, all_agents, all_tool_calls diff --git a/src/repello_agent_wiz/visualizers/visualizer.py b/src/repello_agent_wiz/visualizers/visualizer.py index c3c3dd5..a68fcc4 100644 --- a/src/repello_agent_wiz/visualizers/visualizer.py +++ b/src/repello_agent_wiz/visualizers/visualizer.py @@ -27,7 +27,7 @@ def generate_visualization(json_path: str, open_browser: bool = False): json_string = json.dumps(graph, indent=2) index_filled = index_text.replace("const data = {};", f"const data = {json_string};") - with open(output_dir / "index.html", "w") as f: + with open(output_dir / "index.html", "w" , encoding='utf-8') as f: f.write(index_filled) print(f"[✓] Visualization HTML generated at: {output_dir}/index.html")