Kye Gomez
2 months ago
parent
146402e149
commit
982601614e
@ -0,0 +1,359 @@
|
||||
import os
|
||||
import re
|
||||
|
||||
import litellm
|
||||
|
||||
from swarms import Agent
|
||||
|
||||
litellm.drop_params = True
|
||||
|
||||
|
||||
def extract_code_blocks(text: str) -> str:
|
||||
"""
|
||||
Extracts code blocks enclosed in triple backticks from the given text.
|
||||
|
||||
Args:
|
||||
text (str): The input text containing code blocks.
|
||||
|
||||
Returns:
|
||||
str: The extracted code blocks joined together as a single string.
|
||||
"""
|
||||
# Regular expression to match code blocks enclosed in triple backticks
|
||||
code_block_pattern = re.compile(
|
||||
r"```(?:[a-zA-Z]*)\n(.*?)```", re.DOTALL
|
||||
)
|
||||
|
||||
# Find all matches and join them into a single string
|
||||
matches = code_block_pattern.findall(text)
|
||||
return "\n".join(matches)
|
||||
|
||||
|
||||
# Prompt #1: Translate PyTorch code into CUDA (extensive instructions)
|
||||
translate_pytorch_to_cuda_prompt = """
|
||||
You are an AI agent specialized in converting PyTorch code into efficient,
|
||||
well-structured, and properly functioning CUDA code. Your role is to transform
|
||||
the given PyTorch code (which may include Python-based tensor operations,
|
||||
layers, and training loops) into a CUDA codebase capable of running directly on
|
||||
NVIDIA GPUs without relying on the high-level abstractions of PyTorch.
|
||||
|
||||
Detailed Instructions:
|
||||
1. Read and thoroughly understand every line of the provided PyTorch code,
|
||||
including model definitions, forward passes, loss functions, backward passes,
|
||||
and optimization steps.
|
||||
|
||||
2. Rewrite the code in pure CUDA C/C++:
|
||||
- Convert high-level tensor operations to raw GPU kernel operations and
|
||||
memory management.
|
||||
- Allocate and deallocate memory on the GPU using CUDA APIs like cudaMalloc
|
||||
and cudaFree.
|
||||
- Handle data transfers between the host (CPU) and the device (GPU)
|
||||
appropriately (e.g., cudaMemcpy).
|
||||
- Convert PyTorch's autograd-based gradient calculations into manual CUDA
|
||||
kernel operations or unify them within your own backward pass
|
||||
implementation if needed.
|
||||
- Replace Pythonic loops and array operations with explicit CUDA kernels.
|
||||
- Implement or port any custom CUDA kernels for special functionality.
|
||||
|
||||
3. Optimize the CUDA code for parallel execution:
|
||||
- Use grid and block dimensions that take advantage of the target GPU's
|
||||
compute capabilities and memory bandwidth.
|
||||
- Leverage shared memory, constant memory, or registers when beneficial.
|
||||
- Unroll loops and reduce warp divergence where possible.
|
||||
- Use efficient memory access patterns (e.g., coalesced memory access).
|
||||
|
||||
4. Preserve the overall logic and functionality:
|
||||
- The final CUDA code must produce equivalent numerical results as the
|
||||
original PyTorch code for the same inputs and initialization.
|
||||
- Keep the data types, precision, and numerical stability in mind while
|
||||
performing computations in CUDA.
|
||||
|
||||
5. Ensure readability and maintainability:
|
||||
- Add comments where necessary to explain tricky parts of the CUDA kernels
|
||||
and memory management.
|
||||
- Use clear function or file organization to separate different components
|
||||
(e.g., forward pass, backward pass, kernel definitions, helper utilities).
|
||||
|
||||
6. Remember that the goal is to create an extensive and fully functional .cu
|
||||
file or files that can be compiled with NVCC or integrated into a larger
|
||||
project.
|
||||
|
||||
7. Ensure PyTorch Integration:
|
||||
- Implement proper PyTorch C++ extension interfaces using torch::extension
|
||||
- Include necessary PyTorch headers and macros for tensor operations
|
||||
- Create Python bindings using pybind11 to expose CUDA functions to PyTorch
|
||||
- Handle PyTorch tensor types and conversions appropriately
|
||||
- Ensure compatibility with PyTorch's autograd system
|
||||
- Follow PyTorch's extension conventions for error handling and memory management
|
||||
- Make the code loadable as a PyTorch extension module
|
||||
|
||||
Output Requirements:
|
||||
- Provide only the final CUDA code.
|
||||
- Do not include any explanatory text outside of comments in the code itself.
|
||||
- The CUDA code should stand on its own as a complete solution that can be
|
||||
compiled or integrated without additional references.
|
||||
- The code must be fully compatible with PyTorch's extension system and
|
||||
able to be imported and used seamlessly within PyTorch models.
|
||||
- Ensure the code is fully compatible with PyTorch's extension system.
|
||||
- And, the code should be very long and extensive
|
||||
"""
|
||||
|
||||
# Prompt #2: Make the CUDA code super reliable and fast (extensive instructions)
|
||||
make_cuda_super_reliable_and_fast_prompt = """
|
||||
You are an advanced AI agent whose goal is to take a fully functional CUDA codebase
|
||||
and optimize it to be extraordinarily robust, reliable, and efficient. You must
|
||||
focus on performance improvements at both the kernel level and architectural
|
||||
level, ensuring that the code is streamlined and built to handle potential edge
|
||||
cases gracefully.
|
||||
|
||||
Detailed Instructions:
|
||||
1. Dive deeply into the CUDA kernels and identify performance bottlenecks:
|
||||
- Look for any uncoalesced memory accesses, unnecessary global memory reads,
|
||||
or writes that could be optimized.
|
||||
- Identify any opportunities to reduce shared memory usage, if that memory
|
||||
usage does not yield significant benefit, or to increase it if it can help
|
||||
reduce global memory accesses.
|
||||
|
||||
2. Implement advanced optimization techniques:
|
||||
- Use loop unrolling where beneficial to reduce overhead.
|
||||
- Employ occupancy analysis to find the ideal block size and grid size for
|
||||
maximum parallelism.
|
||||
- Consider the use of constant memory for frequently accessed read-only data.
|
||||
- Evaluate the benefits of pinned (page-locked) memory for host-device
|
||||
transfers.
|
||||
|
||||
3. Improve reliability and error handling:
|
||||
- Insert meaningful checks for CUDA API function calls (e.g., cudaMalloc,
|
||||
cudaMemcpy, cudaFree, kernel launches). Ensure proper cleanup if any call
|
||||
fails or returns an error code.
|
||||
- Handle edge cases where input sizes might be too large or too small,
|
||||
preventing invalid memory accesses or out-of-bounds kernel launches.
|
||||
- Ensure that any macros, compilation flags, or library calls align with
|
||||
best practices for portability across different GPU architectures.
|
||||
|
||||
4. Document any advanced tricks or strategies used:
|
||||
- In-line commentary is crucial. Briefly explain why a specific optimization
|
||||
was chosen and how it impacts performance.
|
||||
|
||||
5. Maintain numerical equivalence:
|
||||
- All transformations must preserve the original numerical outcomes within
|
||||
reasonable floating-point precision limits. Do not alter the algorithmic
|
||||
results unexpectedly.
|
||||
|
||||
6. Provide a clean final output:
|
||||
- The output must be only the improved and optimized CUDA source code.
|
||||
- All extraneous explanation beyond code comments should be avoided.
|
||||
- The code must be very long and extensive, containing all necessary functionality.
|
||||
- Output only the complete code file with no additional text.
|
||||
|
||||
Goal:
|
||||
- Deliver a high-performance CUDA code that is not only efficient in running
|
||||
time but is also resilient, capable of handling unexpected conditions, and
|
||||
thoroughly commented to allow easy maintenance and future modifications.
|
||||
- Ensure the output is a complete, extensive codebase with all functionality included.
|
||||
"""
|
||||
|
||||
# Prompt #3: Cleanup errors and add extensive documentation (extensive instructions)
|
||||
cleanup_and_document_cuda_code_prompt = """
|
||||
You are a specialized AI agent focused on thoroughly debugging, cleaning up, and
|
||||
documenting a CUDA codebase. Your mission is to ensure the code compiles cleanly,
|
||||
handles errors gracefully, follows best coding practices, and includes detailed
|
||||
documentation for maintainability and educational purposes.
|
||||
|
||||
Detailed Instructions:
|
||||
1. Debug and error cleanup:
|
||||
- Identify any compilation or runtime errors and fix them.
|
||||
- Check for mismatched types, undeclared variables, improper memory usage, or
|
||||
incorrect kernel configurations.
|
||||
- Make sure that each kernel launch has correct grid and block dimensions and
|
||||
that indexing inside kernels is handled properly (avoid out-of-bounds
|
||||
threads).
|
||||
|
||||
2. Strengthen error handling:
|
||||
- Wrap all CUDA library calls (cudaMalloc, cudaMemcpy, kernel launches, etc.)
|
||||
with macros or functions that check for and report errors in a consistent
|
||||
manner (e.g., using cudaGetErrorString).
|
||||
- Ensure resources are deallocated in case of failure and that the program
|
||||
can exit safely without memory leaks or GPU lockups.
|
||||
|
||||
3. Add thorough documentation:
|
||||
- Provide high-level explanations near the top of the file describing the
|
||||
overall code structure and flow.
|
||||
- Within each function, write docstrings or block comments to explain the
|
||||
function's purpose, inputs, outputs, and major steps.
|
||||
- For each kernel, add comments describing how threads are mapped to data,
|
||||
how memory is accessed, and what the main loop or computational logic
|
||||
accomplishes.
|
||||
|
||||
4. Check performance remains robust:
|
||||
- Ensure that none of the debugging or cleanup processes introduce unnecessary
|
||||
slowdowns. Where possible, maintain or improve previous optimizations.
|
||||
|
||||
5. Provide final cleaned, well-documented CUDA source code:
|
||||
- The output must contain only the updated source code, with no additional
|
||||
explanation outside of code comments and docstrings.
|
||||
- The code must be ready to compile, fully functional, and in a polished
|
||||
state that one can confidently integrate into a production environment.
|
||||
- The code must be very long and extensive, containing all necessary functionality.
|
||||
- Output only the complete code file with no additional text.
|
||||
|
||||
Goal:
|
||||
- Deliver a thoroughly cleaned, expertly documented CUDA code file. The
|
||||
readability, reliability, and educational clarity of the code should be
|
||||
enhanced without sacrificing computational performance.
|
||||
- Ensure the output is a complete, extensive codebase with all functionality included.
|
||||
"""
|
||||
|
||||
# Prompt #4: Produce a final, extensive CUDA file (extensive instructions)
|
||||
produce_extensive_cuda_file_prompt = """
|
||||
You are an AI agent tasked with producing a final, extensive .cu or .cuh file
|
||||
containing all functionality needed to run the code originally derived from
|
||||
PyTorch. This final file should encapsulate:
|
||||
- The core CUDA kernels.
|
||||
- The initialization and teardown logic for GPU resources.
|
||||
- Optimized computations and any specialized routines (e.g., custom backward
|
||||
passes, advanced math operations).
|
||||
- Comprehensive yet concise in-code documentation.
|
||||
|
||||
Detailed Instructions:
|
||||
1. Merge all relevant kernels, utility functions, and structures into a cohesive
|
||||
single file or a well-structured set of files.
|
||||
2. Ensure you apply all previous optimizations and cleanup efforts, reflecting
|
||||
them in this combined final output.
|
||||
3. Use robust function prototypes for any utility routines, paying attention to
|
||||
scope and avoiding unnecessary global variables.
|
||||
4. Keep the code easily navigable with clear sections:
|
||||
- Data structures and utility definitions.
|
||||
- Kernel definitions.
|
||||
- Host functions for kernel launches.
|
||||
- Initialization and cleanup logic.
|
||||
- Document each section thoroughly for ease of reference.
|
||||
5. Ensure it compiles error-free under standard NVCC compilation.
|
||||
6. Provide only the CUDA code, including all documentation within comments:
|
||||
- No additional external explanation should be outside these comments.
|
||||
- The code must be very long and extensive, containing all necessary functionality.
|
||||
- Output only the complete code file with no additional text.
|
||||
|
||||
Goal:
|
||||
- A single or modular set of .cu/.cuh files that stand as the definitive version
|
||||
of the CUDA codebase, balancing performance, reliability, and maintainability.
|
||||
- Ensure the output is a complete, extensive codebase with all functionality included.
|
||||
"""
|
||||
|
||||
|
||||
# Now create one agent for each prompt, similar to the example with the Financial-Analysis-Agent.
|
||||
translate_agent = Agent(
|
||||
agent_name="Translate-PyTorch-To-CUDA-Agent",
|
||||
system_prompt=translate_pytorch_to_cuda_prompt,
|
||||
model_name="openai/o1",
|
||||
max_loops=1,
|
||||
max_tokens=10000,
|
||||
output_type="str",
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
super_fast_agent = Agent(
|
||||
agent_name="Make-CUDA-Code-Super-Fast-Agent",
|
||||
system_prompt=make_cuda_super_reliable_and_fast_prompt,
|
||||
model_name="openai/o1",
|
||||
max_loops=1,
|
||||
max_tokens=10000,
|
||||
output_type="str",
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
cleanup_agent = Agent(
|
||||
agent_name="Cleanup-and-Document-CUDA-Agent",
|
||||
system_prompt=cleanup_and_document_cuda_code_prompt,
|
||||
model_name="openai/o1",
|
||||
max_loops=1,
|
||||
max_tokens=10000,
|
||||
output_type="str",
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
final_cuda_agent = Agent(
|
||||
agent_name="Produce-Final-Extensive-CUDA-File-Agent",
|
||||
system_prompt=produce_extensive_cuda_file_prompt,
|
||||
model_name="openai/o1",
|
||||
max_loops=1,
|
||||
max_tokens=10000,
|
||||
output_type="str",
|
||||
temperature=0,
|
||||
)
|
||||
|
||||
|
||||
class CudaSwarm:
|
||||
def __init__(
|
||||
self,
|
||||
name: str = "CudaSwarm",
|
||||
description: str = "A swarm of agents that convert PyTorch code into CUDA code",
|
||||
agents: list[Agent] = [
|
||||
translate_agent,
|
||||
super_fast_agent,
|
||||
cleanup_agent,
|
||||
final_cuda_agent,
|
||||
],
|
||||
max_loops: int = 1,
|
||||
file_path: str = "cuda_code.cu",
|
||||
):
|
||||
self.name = name
|
||||
self.description = description
|
||||
self.agents = agents
|
||||
self.max_loops = max_loops
|
||||
self.file_path = file_path
|
||||
|
||||
def write_file(self, content: str = "") -> None:
|
||||
"""
|
||||
Creates a new CUDA file or overwrites an existing one with the specified content.
|
||||
|
||||
Args:
|
||||
content (str): The content to write into the file. Defaults to an empty string.
|
||||
"""
|
||||
# Ensure the file has a .cu extension
|
||||
if not self.file_path.endswith(".cu"):
|
||||
self.file_path += ".cu"
|
||||
|
||||
# Create the directory if it doesn't exist
|
||||
directory = os.path.dirname(self.file_path)
|
||||
if directory and not os.path.exists(directory):
|
||||
os.makedirs(directory)
|
||||
|
||||
# Use write mode to overwrite the file
|
||||
mode = "w"
|
||||
|
||||
try:
|
||||
# Write content to the file
|
||||
with open(self.file_path, mode) as file:
|
||||
file.write(content)
|
||||
print(f"Content successfully written to {self.file_path}")
|
||||
except IOError as e:
|
||||
print(f"An error occurred while writing to the file: {e}")
|
||||
|
||||
def run(self, task: str):
|
||||
"""
|
||||
Runs the swarm of agents to complete the task.
|
||||
"""
|
||||
first_iteration = self.agents[0].run(task)
|
||||
first_iteration = extract_code_blocks(first_iteration)
|
||||
self.write_file(first_iteration)
|
||||
|
||||
# second_iteration = self.agents[1].run(task)
|
||||
# second_iteration = extract_code_blocks(second_iteration)
|
||||
# self.write_file(second_iteration)
|
||||
|
||||
# third_iteration = self.agents[2].run(task)
|
||||
# third_iteration = extract_code_blocks(third_iteration)
|
||||
# self.write_file(third_iteration)
|
||||
|
||||
# final_iteration = self.agents[3].run(task)
|
||||
# final_iteration = extract_code_blocks(final_iteration)
|
||||
# self.write_file(final_iteration)
|
||||
|
||||
return first_iteration
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
swarm = CudaSwarm(file_path="cuda_code.cu")
|
||||
swarm.run(
|
||||
"Create the cuda code for a highly optimized liquid continous learner ssm model"
|
||||
)
|
||||
@ -0,0 +1,49 @@
|
||||
from swarms import Agent
|
||||
from swarms.prompts.finance_agent_sys_prompt import (
|
||||
FINANCIAL_AGENT_SYS_PROMPT,
|
||||
)
|
||||
|
||||
# Initialize the equity analyst agents
|
||||
equity_analyst_1 = Agent(
|
||||
agent_name="Equity-Analyst-1",
|
||||
agent_description="Equity research analyst focused on fundamental analysis",
|
||||
system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
|
||||
max_loops=1,
|
||||
model_name="gpt-4o",
|
||||
dynamic_temperature_enabled=True,
|
||||
user_name="swarms_corp",
|
||||
retry_attempts=3,
|
||||
context_length=8192,
|
||||
return_step_meta=False,
|
||||
output_type="str",
|
||||
auto_generate_prompt=False,
|
||||
max_tokens=4000,
|
||||
saved_state_path="equity_analyst_1.json",
|
||||
interactive=False,
|
||||
roles="analyst"
|
||||
)
|
||||
|
||||
equity_analyst_2 = Agent(
|
||||
agent_name="Equity-Analyst-2",
|
||||
agent_description="Equity research analyst focused on technical analysis",
|
||||
system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
|
||||
max_loops=1,
|
||||
model_name="gpt-4o",
|
||||
dynamic_temperature_enabled=True,
|
||||
user_name="swarms_corp",
|
||||
retry_attempts=3,
|
||||
context_length=8192,
|
||||
return_step_meta=False,
|
||||
output_type="str",
|
||||
auto_generate_prompt=False,
|
||||
max_tokens=4000,
|
||||
saved_state_path="equity_analyst_2.json",
|
||||
interactive=False,
|
||||
roles="analyst"
|
||||
)
|
||||
|
||||
# Run analysis with both analysts
|
||||
equity_analyst_1.talk_to(
|
||||
equity_analyst_2,
|
||||
"Analyze high growth tech stocks focusing on fundamentals like revenue growth, margins, and market position. Create a detailed analysis table in markdown."
|
||||
)
|
||||
@ -0,0 +1,506 @@
|
||||
from typing import Callable, Any
|
||||
import psutil
|
||||
import asyncio
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
from rich.panel import Panel
|
||||
from rich.layout import Layout
|
||||
from rich.tree import Tree
|
||||
from rich.progress import (
|
||||
Progress,
|
||||
TimeElapsedColumn,
|
||||
TextColumn,
|
||||
SpinnerColumn,
|
||||
)
|
||||
from rich.live import Live
|
||||
from rich.text import Text
|
||||
from typing import Dict, Optional, List
|
||||
from dataclasses import dataclass
|
||||
from queue import Queue
|
||||
from datetime import datetime
|
||||
|
||||
try:
|
||||
import pynvml
|
||||
|
||||
pynvml.nvmlInit()
|
||||
GPU_ENABLED = True
|
||||
except ImportError:
|
||||
GPU_ENABLED = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class SwarmMetadata:
|
||||
name: str
|
||||
description: str
|
||||
version: str
|
||||
type: str # hierarchical, parallel, sequential
|
||||
created_at: datetime
|
||||
author: str
|
||||
tags: List[str]
|
||||
primary_objective: str
|
||||
secondary_objectives: List[str]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Agent:
|
||||
name: str
|
||||
role: str
|
||||
description: str
|
||||
agent_type: str # e.g., "LLM", "Neural", "Rule-based"
|
||||
capabilities: List[str]
|
||||
parameters: Dict[str, any]
|
||||
metadata: Dict[str, str]
|
||||
children: List["Agent"] = None
|
||||
parent: Optional["Agent"] = None
|
||||
output_stream: Queue = None
|
||||
|
||||
def __post_init__(self):
|
||||
self.children = self.children or []
|
||||
self.output_stream = Queue()
|
||||
|
||||
@property
|
||||
def hierarchy_level(self) -> int:
|
||||
level = 0
|
||||
current = self
|
||||
while current.parent:
|
||||
level += 1
|
||||
current = current.parent
|
||||
return level
|
||||
|
||||
|
||||
class SwarmVisualizationRich:
|
||||
def __init__(
|
||||
self,
|
||||
swarm_metadata: SwarmMetadata,
|
||||
root_agent: Agent,
|
||||
update_resources: bool = True,
|
||||
refresh_rate: float = 0.1,
|
||||
):
|
||||
self.swarm_metadata = swarm_metadata
|
||||
self.root_agent = root_agent
|
||||
self.update_resources = update_resources
|
||||
self.refresh_rate = refresh_rate
|
||||
self.console = Console()
|
||||
self.live = None
|
||||
self.output_history = {}
|
||||
|
||||
# System monitoring
|
||||
self.cores_available = 0
|
||||
self.memory_usage = "N/A"
|
||||
self.gpu_power = "N/A"
|
||||
self.start_time = datetime.now()
|
||||
|
||||
if self.update_resources:
|
||||
self._update_resource_stats()
|
||||
|
||||
def _format_uptime(self) -> str:
|
||||
"""Formats the swarm's uptime."""
|
||||
delta = datetime.now() - self.start_time
|
||||
hours, remainder = divmod(delta.seconds, 3600)
|
||||
minutes, seconds = divmod(remainder, 60)
|
||||
return f"{hours:02d}:{minutes:02d}:{seconds:02d}"
|
||||
|
||||
def _build_agent_tree(
|
||||
self, agent: Agent, tree: Optional[Tree] = None
|
||||
) -> Tree:
|
||||
"""Builds a detailed tree visualization of the agent hierarchy."""
|
||||
agent_info = [
|
||||
f"[bold cyan]{agent.name}[/bold cyan]",
|
||||
f"[yellow]Role:[/yellow] {agent.role}",
|
||||
f"[green]Type:[/green] {agent.agent_type}",
|
||||
f"[blue]Level:[/blue] {agent.hierarchy_level}",
|
||||
f"[magenta]Capabilities:[/magenta] {', '.join(agent.capabilities)}",
|
||||
]
|
||||
|
||||
# Add any custom metadata
|
||||
for key, value in agent.metadata.items():
|
||||
agent_info.append(f"[white]{key}:[/white] {value}")
|
||||
|
||||
# Parameters summary
|
||||
param_summary = ", ".join(
|
||||
f"{k}: {v}" for k, v in agent.parameters.items()
|
||||
)
|
||||
agent_info.append(
|
||||
f"[white]Parameters:[/white] {param_summary}"
|
||||
)
|
||||
|
||||
node_text = "\n".join(agent_info)
|
||||
|
||||
if tree is None:
|
||||
tree = Tree(node_text)
|
||||
else:
|
||||
branch = tree.add(node_text)
|
||||
tree = branch
|
||||
|
||||
for child in agent.children:
|
||||
self._build_agent_tree(child, tree)
|
||||
|
||||
return tree
|
||||
|
||||
def _count_agents(self, agent: Agent) -> int:
|
||||
"""Recursively counts total number of agents in the swarm."""
|
||||
count = 1 # Count current agent
|
||||
for child in agent.children or []:
|
||||
count += self._count_agents(child)
|
||||
return count
|
||||
|
||||
def _create_unified_info_panel(self) -> Panel:
|
||||
"""Creates a unified panel showing both swarm metadata and architecture."""
|
||||
# Create the main container
|
||||
info_layout = Layout()
|
||||
info_layout.split_column(
|
||||
Layout(name="metadata", size=13),
|
||||
Layout(name="architecture"),
|
||||
)
|
||||
|
||||
# Calculate total agents
|
||||
total_agents = self._count_agents(self.root_agent)
|
||||
|
||||
# Metadata section
|
||||
metadata_table = Table.grid(padding=1, expand=True)
|
||||
metadata_table.add_column("Label", style="bold cyan")
|
||||
metadata_table.add_column("Value", style="white")
|
||||
|
||||
# System resources
|
||||
if self.update_resources:
|
||||
self._update_resource_stats()
|
||||
|
||||
# Add description with proper wrapping
|
||||
description_text = Text(
|
||||
self.swarm_metadata.description, style="italic"
|
||||
)
|
||||
description_text.wrap(self.console, width=60, overflow="fold")
|
||||
|
||||
metadata_table.add_row("Swarm Name", self.swarm_metadata.name)
|
||||
metadata_table.add_row("Description", description_text)
|
||||
metadata_table.add_row("Version", self.swarm_metadata.version)
|
||||
metadata_table.add_row("Total Agents", str(total_agents))
|
||||
metadata_table.add_row("Author", self.swarm_metadata.author)
|
||||
metadata_table.add_row(
|
||||
"System",
|
||||
f"CPU: {self.cores_available} cores | Memory: {self.memory_usage}",
|
||||
)
|
||||
metadata_table.add_row(
|
||||
"Primary Objective", self.swarm_metadata.primary_objective
|
||||
)
|
||||
|
||||
info_layout["metadata"].update(metadata_table)
|
||||
|
||||
info_layout["metadata"].update(metadata_table)
|
||||
|
||||
# Architecture section with tree visualization
|
||||
architecture_tree = self._build_agent_tree(self.root_agent)
|
||||
info_layout["architecture"].update(architecture_tree)
|
||||
|
||||
return Panel(
|
||||
info_layout,
|
||||
title="[bold]Swarm Information & Architecture[/bold]",
|
||||
)
|
||||
|
||||
def _create_outputs_panel(self) -> Panel:
|
||||
"""Creates a panel that displays stacked message history for all agents."""
|
||||
# Create a container for all messages across all agents
|
||||
all_messages = []
|
||||
|
||||
def collect_agent_messages(agent: Agent):
|
||||
"""Recursively collect messages from all agents."""
|
||||
messages = self.output_history.get(agent.name, [])
|
||||
for msg in messages:
|
||||
all_messages.append(
|
||||
{
|
||||
"agent": agent.name,
|
||||
"time": msg["time"],
|
||||
"content": msg["content"],
|
||||
"style": msg["style"],
|
||||
}
|
||||
)
|
||||
for child in agent.children:
|
||||
collect_agent_messages(child)
|
||||
|
||||
# Collect all messages
|
||||
collect_agent_messages(self.root_agent)
|
||||
|
||||
# Sort messages by timestamp
|
||||
all_messages.sort(key=lambda x: x["time"])
|
||||
|
||||
# Create the stacked message display
|
||||
Layout()
|
||||
messages_container = []
|
||||
|
||||
for msg in all_messages:
|
||||
# Create a panel for each message
|
||||
message_text = Text()
|
||||
message_text.append(f"[{msg['time']}] ", style="dim")
|
||||
message_text.append(
|
||||
f"{msg['agent']}: ", style="bold cyan"
|
||||
)
|
||||
message_text.append(msg["content"], style=msg["style"])
|
||||
|
||||
messages_container.append(message_text)
|
||||
|
||||
# Join all messages with line breaks
|
||||
if messages_container:
|
||||
final_text = Text("\n").join(messages_container)
|
||||
else:
|
||||
final_text = Text("No messages yet...", style="dim")
|
||||
|
||||
# Create scrollable panel for all messages
|
||||
return Panel(
|
||||
final_text,
|
||||
title="[bold]Agent Communication Log[/bold]",
|
||||
border_style="green",
|
||||
padding=(1, 2),
|
||||
)
|
||||
|
||||
def _update_resource_stats(self):
|
||||
"""Updates system resource statistics."""
|
||||
self.cores_available = psutil.cpu_count(logical=True)
|
||||
mem_info = psutil.virtual_memory()
|
||||
total_gb = mem_info.total / (1024**3)
|
||||
used_gb = mem_info.used / (1024**3)
|
||||
self.memory_usage = f"{used_gb:.1f}GB / {total_gb:.1f}GB ({mem_info.percent}%)"
|
||||
|
||||
if GPU_ENABLED:
|
||||
try:
|
||||
device_count = pynvml.nvmlDeviceGetCount()
|
||||
gpu_info = []
|
||||
for i in range(device_count):
|
||||
handle = pynvml.nvmlDeviceGetHandleByIndex(i)
|
||||
name = pynvml.nvmlDeviceGetName(handle).decode()
|
||||
mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
||||
usage = (mem.used / mem.total) * 100
|
||||
gpu_info.append(f"{name}: {usage:.1f}%")
|
||||
self.gpu_power = " | ".join(gpu_info)
|
||||
except Exception as e:
|
||||
self.gpu_power = f"GPU Error: {str(e)}"
|
||||
else:
|
||||
self.gpu_power = "No GPU detected"
|
||||
|
||||
async def stream_output(
|
||||
self,
|
||||
agent: Agent,
|
||||
text: str,
|
||||
title: Optional[str] = None,
|
||||
style: str = "bold cyan",
|
||||
delay: float = 0.05,
|
||||
by_word: bool = False,
|
||||
):
|
||||
"""
|
||||
Streams output for a specific agent with sophisticated token-by-token animation.
|
||||
|
||||
Args:
|
||||
agent (Agent): The agent whose output is being streamed
|
||||
text (str): The text to stream
|
||||
title (Optional[str]): Custom title for the output panel
|
||||
style (str): Style for the output text
|
||||
delay (float): Delay between tokens
|
||||
by_word (bool): If True, streams word by word instead of character by character
|
||||
"""
|
||||
display_text = Text(style=style)
|
||||
current_output = ""
|
||||
|
||||
# Split into words or characters
|
||||
tokens = text.split() if by_word else text
|
||||
|
||||
# Create a panel for this agent's output
|
||||
title = title or f"{agent.name} Output"
|
||||
|
||||
for token in tokens:
|
||||
# Add appropriate spacing
|
||||
token_with_space = token + (" " if by_word else "")
|
||||
current_output += token_with_space
|
||||
display_text.append(token_with_space)
|
||||
|
||||
# Initialize history list if it doesn't exist
|
||||
if agent.name not in self.output_history:
|
||||
self.output_history[agent.name] = []
|
||||
|
||||
# Store the complete message when finished
|
||||
if token == tokens[-1]:
|
||||
timestamp = datetime.now().strftime("%H:%M:%S")
|
||||
self.output_history[agent.name].append(
|
||||
{
|
||||
"time": timestamp,
|
||||
"content": current_output,
|
||||
"style": style,
|
||||
}
|
||||
)
|
||||
|
||||
# Update live display if active
|
||||
if self.live:
|
||||
self.live.update(self._create_layout())
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
async def print_progress(
|
||||
self,
|
||||
description: str,
|
||||
task_fn: Callable,
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""
|
||||
Displays a progress spinner while executing a task.
|
||||
|
||||
Args:
|
||||
description (str): Task description
|
||||
task_fn (Callable): Function to execute
|
||||
*args (Any): Arguments for task_fn
|
||||
**kwargs (Any): Keyword arguments for task_fn
|
||||
|
||||
Returns:
|
||||
Any: Result of task_fn
|
||||
"""
|
||||
progress = Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
TimeElapsedColumn(),
|
||||
)
|
||||
|
||||
try:
|
||||
with progress:
|
||||
task = progress.add_task(description, total=None)
|
||||
result = await task_fn(*args, **kwargs)
|
||||
progress.update(task, completed=True)
|
||||
return result
|
||||
except Exception as e:
|
||||
progress.stop()
|
||||
raise e
|
||||
|
||||
def _create_layout(self) -> Layout:
|
||||
"""Creates the main visualization layout."""
|
||||
layout = Layout()
|
||||
layout.split_row(
|
||||
Layout(name="info", ratio=2),
|
||||
Layout(name="outputs", ratio=3),
|
||||
)
|
||||
|
||||
layout["info"].update(self._create_unified_info_panel())
|
||||
layout["outputs"].update(self._create_outputs_panel())
|
||||
|
||||
return layout
|
||||
|
||||
async def start(self):
|
||||
"""Starts the visualization with live updates."""
|
||||
with Live(
|
||||
self._create_layout(),
|
||||
refresh_per_second=int(1 / self.refresh_rate),
|
||||
) as self.live:
|
||||
while True:
|
||||
|
||||
def process_agent_streams(agent: Agent):
|
||||
while not agent.output_stream.empty():
|
||||
new_output = agent.output_stream.get()
|
||||
asyncio.create_task(
|
||||
self.stream_output(agent, new_output)
|
||||
)
|
||||
for child in agent.children:
|
||||
process_agent_streams(child)
|
||||
|
||||
process_agent_streams(self.root_agent)
|
||||
await asyncio.sleep(self.refresh_rate)
|
||||
|
||||
|
||||
# Example usage
|
||||
if __name__ == "__main__":
|
||||
# Create swarm metadata
|
||||
swarm_metadata = SwarmMetadata(
|
||||
name="Financial Advisory Swarm",
|
||||
description="Intelligent swarm for financial analysis and advisory",
|
||||
version="1.0.0",
|
||||
type="hierarchical",
|
||||
created_at=datetime.now(),
|
||||
author="AI Research Team",
|
||||
# tags=["finance", "analysis", "advisory"],
|
||||
primary_objective="Provide comprehensive financial analysis and recommendations",
|
||||
secondary_objectives=[
|
||||
"Monitor market trends",
|
||||
"Analyze competitor behavior",
|
||||
"Generate investment strategies",
|
||||
],
|
||||
)
|
||||
|
||||
# Create agent hierarchy with detailed parameters
|
||||
analyst = Agent(
|
||||
name="Financial Analyst",
|
||||
role="Analysis",
|
||||
description="Analyzes financial data and market trends",
|
||||
agent_type="LLM",
|
||||
capabilities=[
|
||||
"data analysis",
|
||||
"trend detection",
|
||||
"risk assessment",
|
||||
],
|
||||
parameters={"model": "gpt-4", "temperature": 0.7},
|
||||
metadata={
|
||||
"specialty": "Market Analysis",
|
||||
"confidence_threshold": "0.85",
|
||||
},
|
||||
)
|
||||
|
||||
researcher = Agent(
|
||||
name="Market Researcher",
|
||||
role="Research",
|
||||
description="Conducts market research and competitor analysis",
|
||||
agent_type="Neural",
|
||||
capabilities=[
|
||||
"competitor analysis",
|
||||
"market sentiment",
|
||||
"trend forecasting",
|
||||
],
|
||||
parameters={"batch_size": 32, "learning_rate": 0.001},
|
||||
metadata={
|
||||
"data_sources": "Bloomberg, Reuters",
|
||||
"update_frequency": "1h",
|
||||
},
|
||||
)
|
||||
|
||||
advisor = Agent(
|
||||
name="Investment Advisor",
|
||||
role="Advisory",
|
||||
description="Provides investment recommendations",
|
||||
agent_type="Hybrid",
|
||||
capabilities=[
|
||||
"portfolio optimization",
|
||||
"risk management",
|
||||
"strategy generation",
|
||||
],
|
||||
parameters={
|
||||
"risk_tolerance": "moderate",
|
||||
"time_horizon": "long",
|
||||
},
|
||||
metadata={
|
||||
"certification": "CFA Level 3",
|
||||
"specialization": "Equity",
|
||||
},
|
||||
children=[analyst, researcher],
|
||||
)
|
||||
|
||||
# Create visualization
|
||||
viz = SwarmVisualizationRich(
|
||||
swarm_metadata=swarm_metadata,
|
||||
root_agent=advisor,
|
||||
refresh_rate=0.1,
|
||||
)
|
||||
|
||||
# Example of streaming output simulation
|
||||
async def simulate_outputs():
|
||||
await viz.stream_output(
|
||||
advisor,
|
||||
"Analyzing market conditions...\nGenerating investment advice...",
|
||||
)
|
||||
await viz.stream_output(
|
||||
analyst,
|
||||
"Processing financial data...\nIdentifying trends...",
|
||||
)
|
||||
await viz.stream_output(
|
||||
researcher,
|
||||
"Researching competitor movements...\nAnalyzing market share...",
|
||||
)
|
||||
|
||||
# Run the visualization
|
||||
async def main():
|
||||
viz_task = asyncio.create_task(viz.start())
|
||||
await simulate_outputs()
|
||||
await viz_task
|
||||
|
||||
asyncio.run(main())
|
||||
Loading…
Reference in new issue