Spaces:
Running
Running
| from typing import List, AsyncIterator | |
| import os | |
| from autogen_agentchat.agents import AssistantAgent | |
| from autogen_agentchat.conditions import TextMentionTermination | |
| from autogen_agentchat.messages import MultiModalMessage | |
| from autogen_agentchat.teams import RoundRobinGroupChat | |
| from autogen_core import Image as AGImage | |
| from autogen_core.tools import FunctionTool | |
| from autogen_ext.models.openai import OpenAIChatCompletionClient | |
| from dotenv import load_dotenv, find_dotenv | |
| from PIL import Image | |
| import yaml | |
| from tools import Tools | |
| class MinervaTeam: | |
| """ | |
| Minerva Team tasked with Scam Protection using an Agentic Workflow | |
| """ | |
| def __init__(self, config_path: str = "config/agents.yaml"): | |
| """ | |
| Initialize Minerva with agents and tools | |
| """ | |
| self.load_environment() | |
| self.model = self.initialize_model() | |
| self.config = self.load_config(config_path) | |
| self.tools = Tools() | |
| self.agents = self.create_agents() | |
| self.team = self.create_team() | |
| def load_environment(self): | |
| """Load environment variables""" | |
| load_dotenv(find_dotenv()) | |
| def load_config(self, config_path: str) -> dict: | |
| """Load agent configurations from YAML file""" | |
| with open(config_path, 'r') as file: | |
| return yaml.safe_load(file) | |
| def initialize_model(self) -> OpenAIChatCompletionClient: | |
| """Initialize OpenAI model""" | |
| return OpenAIChatCompletionClient( | |
| model="gpt-4o", | |
| api_key=os.getenv("OPENAI_API_KEY") | |
| ) | |
| def create_agents(self) -> List[AssistantAgent]: | |
| """Create all required agents with their specialized roles and tools""" | |
| ocr_tool = FunctionTool( | |
| self.tools.ocr, | |
| description="Extracts text from an image path" | |
| ) | |
| url_checker_tool = FunctionTool( | |
| self.tools.is_url_safe, | |
| description="Checks if a URL is safe" | |
| ) | |
| agents = [] | |
| agents.append(AssistantAgent( | |
| name="OCR_Specialist", | |
| description="Extracts text from an image", | |
| system_message=self.config['ocr_agent']['assignment'], | |
| model_client=self.model, | |
| #tools=[ocr_tool] # Default OCR to GPT-4o vision capabilities. Uncomment to OCR with tool calling (requires pytesseract) | |
| )) | |
| agents.append(AssistantAgent( | |
| name="Link_Checker", | |
| description="Checks if a Link is safe", | |
| system_message=self.config['link_checker_agent']['assignment'], | |
| model_client=self.model, | |
| tools=[url_checker_tool] | |
| )) | |
| agents.append(AssistantAgent( | |
| name="Content_Analyst", | |
| description="Analyzes the text for scam patterns", | |
| system_message=self.config['content_agent']['assignment'], | |
| model_client=self.model | |
| )) | |
| agents.append(AssistantAgent( | |
| name="Decision_Maker", | |
| description="Synthesizes the analyses and make final determination", | |
| system_message=self.config['decision_agent']['assignment'], | |
| model_client=self.model | |
| )) | |
| agents.append(AssistantAgent( | |
| name="Summary_Agent", | |
| description="Generate a summary of the final determination", | |
| system_message=self.config['summary_agent']['assignment'], | |
| model_client=self.model | |
| )) | |
| agents.append(AssistantAgent( | |
| name="Language_Translation_Agent", | |
| description="Translate the summary to the user language", | |
| system_message=self.config['language_translation_agent']['assignment'], | |
| model_client=self.model | |
| )) | |
| return agents | |
| def create_team(self) -> RoundRobinGroupChat: | |
| """Create a team of agents that work together in Round Robin fashion""" | |
| termination = TextMentionTermination("NO_TEXT_FOUND") | |
| return RoundRobinGroupChat( | |
| self.agents, | |
| max_turns=6, | |
| termination_condition=termination | |
| ) | |
| async def reset(self): | |
| """Reset team state""" | |
| await self.team.reset() | |
| async def analyze(self, image: Image) -> AsyncIterator: | |
| """ | |
| Analyze an image for potential scams. | |
| """ | |
| img = AGImage(image) | |
| mm_message = MultiModalMessage(content=[img], source="User") | |
| return self.team.run_stream(task=mm_message) |