Large-Scale Agent Architecture: Complete Guide to Building Scalable Multi-Agent Systems with AutoGen, Kubernetes, and Vector Databases Author: Daniel Kliewer Date: 2025-03-25 Tags: Multi-Agent Systems, AutoGen, Kubernetes, Vector Databases, Scalable Architecture, Distributed Computing, AI Agents, System Design, Enterprise AI, Microservices Description: An in-depth systems engineering guide to designing and implementing scalable multi-agent frameworks using AutoGen, Kubernetes, Kafka, vector databases, and local LLMs for enterprise-grade AI applications. ---![Image](/images/ComfyUI_00194_.png) # Building Large-Scale AI Agents: A Deep-Dive Guide for Experienced Engineers ## 1. Introduction ### Why AI Agents Are Revolutionizing Industries In today's high-velocity enterprise environments, the paradigm has shifted from monolithic AI models to orchestrated, purpose-built AI agents working in concert. These agent-based systems represent a fundamental evolution in how we architect intelligent applications, enabling autonomous decision-making and task execution at unprecedented scale. Financial institutions like JP Morgan Chase have deployed agent networks for algorithmic trading that dynamically respond to market conditions, executing complex strategies across multiple asset classes while maintaining regulatory compliance. Healthcare providers including Mayo Clinic have implemented diagnostic agent ecosystems that collaborate across specialties, analyzing patient data and providing treatment recommendations with 97% concordance with specialist physicians. The key differentiator between traditional AI systems and modern agent architectures lies in their ability to decompose complex problems into specialized sub-tasks, maintain persistent state across interactions, and intelligently route information through distributed processing pipelines—all while scaling horizontally across compute resources. ``` "AI agents represent a shift from passive inference to active computation. Where traditional models wait for queries, agents proactively identify problems and orchestrate solutions across organizational boundaries." — Andrej Karpathy, Former Director of AI at Tesla ``` ### Choosing the Right Tech Stack for Your AI Agent System Building enterprise-grade AI agent systems requires careful consideration of your infrastructure components, with each layer of the stack influencing performance, scalability, and operational complexity: | Layer | Key Technologies | Selection Criteria | |-------|-----------------|-------------------| | Orchestration | Kubernetes, Nomad, ECS | Deployment density, autoscaling capabilities, service mesh integration | | Compute Framework | Ray, Dask, Spark | Parallelization model, scheduling overhead, fault tolerance | | Agent Framework | AutoGen, LangChain, CrewAI | Agent cooperation models, reasoning capabilities, tool integration | | Vector Storage | ChromaDB, Pinecone, Weaviate, Snowflake | Query latency, indexing performance, embedding model compatibility | | Message Bus | Kafka, RabbitMQ, Pulsar | Throughput requirements, ordering guarantees, retention policies | | API Layer | FastAPI, Django, Flask | Request handling, async support, middleware ecosystem | | Monitoring | Prometheus, Grafana, Datadog | Observability coverage, alerting capabilities, performance impact | Your selection should be driven by specific workload characteristics, scaling requirements, and existing infrastructure investments. For real-time processing with strict latency requirements, a Ray + FastAPI + Kafka combination offers exceptional performance. For batch-oriented enterprise workflows with strong governance requirements, an Airflow + AutoGen + Snowflake stack provides robust auditability and integration with data warehousing. ### How This Guide Can Help You Build a Scalable AI Agent Framework This guide approaches AI agent architecture through the lens of production engineering, focusing on the challenges that emerge at scale: - **Stateful Agent Coordination**: How to maintain context across distributed agent clusters while preventing state explosion - **Intelligent Workload Distribution**: Techniques for dynamic task routing among specialized agents - **Knowledge Management**: Strategies for efficient retrieval and updates to agent knowledge bases - **Observability and Debugging**: Tracing causal chains of reasoning across multi-agent systems - **Performance Optimization**: Reducing token usage, latency, and compute costs in large deployments Rather than theoretical concepts, we'll examine concrete implementations with battle-tested infrastructure components. You'll learn how companies like Stripe have reduced their manual review workload by 85% using agent networks for fraud detection, and how Netflix has implemented content recommendation agents that reduce churn by dynamically personalizing user experiences. By the end of this guide, you'll be equipped to architect, implement, and scale AI agent systems that deliver measurable business impact—whether you're building customer-facing applications or internal automation tools. ## 2. Understanding the Core Technologies ### What is AutoGen? A Breakdown of Multi-Agent Systems AutoGen represents a paradigm shift in AI agent orchestration, offering a framework for building systems where multiple specialized agents collaborate to solve complex tasks. Developed by Microsoft Research, AutoGen moves beyond simple prompt engineering to enable sophisticated multi-agent conversations with memory, tool use, and dynamic conversation control. At its core, AutoGen defines a computational graph of conversational agents, each with distinct capabilities: ```python from autogen import AssistantAgent, UserProxyAgent, config_list_from_json # Load LLM configuration config_list = config_list_from_json("llm_config.json") # Define the system architecture with specialized agents assistant = AssistantAgent( name="CTO", llm_config={"config_list": config_list}, system_message="You are a CTO who makes executive technology decisions based on data." ) data_analyst = AssistantAgent( name="DataAnalyst", llm_config={"config_list": config_list}, system_message="You analyze data and provide insights to the CTO." ) engineer = AssistantAgent( name="Engineer", llm_config={"config_list": config_list}, system_message="You implement solutions proposed by the CTO." ) # User proxy agent with capabilities to execute code and retrieve data user_proxy = UserProxyAgent( name="DevOps", human_input_mode="NEVER", max_consecutive_auto_reply=10, code_execution_config={"work_dir": "workspace"}, system_message="You execute code and return results to other agents." ) # Initiate a group conversation with a specific task user_proxy.initiate_chat( assistant, message="Analyze our production logs to identify performance bottlenecks.", clear_history=True, groupchat_agents=[assistant, data_analyst, engineer] ) ``` What distinguishes AutoGen from simpler frameworks is its ability to handle: 1. **Conversational Memory**: Agents maintain context across multi-turn conversations 2. **Tool Usage**: Native integration with code execution and external APIs 3. **Dynamic Agent Selection**: Intelligent routing of tasks to specialized agents 4. **Hierarchical Planning**: Breaking complex tasks into subtasks with appropriate delegation In production environments, AutoGen's flexibility enables diverse agent architectures: - **Hierarchical Teams**: Manager agents delegate to specialist agents - **Competitive Evaluation**: Multiple agents generate solutions evaluated by a judge agent - **Consensus-Based**: Collaborative problem-solving with voting mechanisms Unlike other frameworks that primarily focus on prompt chaining, AutoGen is designed for true multi-agent systems where autonomous entities negotiate, collaborate, and resolve conflicts to achieve goals. ### Key Infrastructure Components: Kubernetes, Kafka, Airflow, and More Building scalable AI agent systems requires robust infrastructure components that can handle the unique demands of distributed agent workloads: #### Kubernetes for Agent Orchestration Kubernetes provides the foundation for deploying, scaling, and managing containerized AI agents. For production deployments, consider these Kubernetes patterns: ```yaml # Kubernetes manifest for a scalable AutoGen agent deployment apiVersion: apps/v1 kind: Deployment metadata: name: agent-deployment labels: app: ai-agent-system spec: replicas: 3 selector: matchLabels: app: ai-agent template: metadata: labels: app: ai-agent spec: containers: - name: agent-container image: your-registry/ai-agent:1.0.0 resources: requests: memory: "2Gi" cpu: "1" limits: memory: "4Gi" cpu: "2" env: - name: AGENT_ROLE value: "analyst" - name: REDIS_HOST value: "redis-service" - name: OPENAI_API_KEY valueFrom: secretKeyRef: name: openai-credentials key: api-key volumeMounts: - name: agent-config mountPath: /app/config volumes: - name: agent-config configMap: name: agent-config ``` Key considerations for Kubernetes deployments: - **Horizontal Pod Autoscaling**: Configure HPA based on CPU/memory metrics or custom metrics like queue depth - **Affinity/Anti-Affinity Rules**: Ensure agents that frequently communicate are co-located for reduced latency - **Resource Quotas**: Implement namespace quotas to prevent AI agent workloads from consuming all cluster resources - **Readiness Probes**: Properly configure readiness checks to ensure agents are fully initialized before receiving traffic #### Apache Kafka for Event-Driven Agent Communication Kafka serves as the nervous system for large-scale agent deployments, enabling asynchronous communication patterns essential for resilient systems: ```python # Producer code for agent event publishing from kafka import KafkaProducer import json producer = KafkaProducer( bootstrap_servers=['kafka-broker-1:9092', 'kafka-broker-2:9092'], value_serializer=lambda v: json.dumps(v).encode('utf-8'), acks='all', retries=3, linger_ms=5 # Batch messages for 5ms for better throughput ) # Agent publishing a task for another agent def publish_analysis_task(data, priority="high"): producer.send( topic='agent.tasks.analysis', key=data['request_id'].encode('utf-8'), # Ensures related messages go to same partition value={ 'task_type': 'analyze_document', 'priority': priority, 'timestamp': time.time(), 'payload': data, 'source_agent': 'document_processor' }, headers=[('priority', priority.encode('utf-8'))] ) producer.flush() ``` For high-throughput agent systems, implement these Kafka optimizations: - **Topic Partitioning Strategy**: Partition topics by agent task type for parallel processing - **Consumer Group Design**: Group consumers by agent role for workload distribution - **Compacted Topics**: Use compacted topics for agent state to maintain latest values - **Exactly-Once Semantics**: Enable transactions for critical agent workflows #### Airflow for Complex Agent Workflow Orchestration For sophisticated agent pipelines with dependencies and scheduling requirements, Apache Airflow provides enterprise-grade orchestration: ```python # Airflow DAG for a multi-agent financial analysis pipeline from airflow import DAG from airflow.operators.python import PythonOperator from datetime import datetime, timedelta default_args = { 'owner': 'ai_team', 'depends_on_past': False, 'start_date': datetime(2023, 1, 1), 'email_on_failure': True, 'retries': 1, 'retry_delay': timedelta(minutes=5), } def initialize_agents(**context): # Initialize agent system with necessary credentials and configuration from agent_framework import AgentCluster cluster = AgentCluster(config_path="/path/to/agent_config.json") return cluster.get_session_id() def run_data_gathering_agents(**context): # Retrieve session ID from previous task session_id = context['ti'].xcom_pull(task_ids='initialize_agents') # Activate data gathering agents to collect financial data from agent_tasks import DataGatheringTask task = DataGatheringTask(session_id=session_id) results = task.execute(sources=['bloomberg', 'reuters', 'sec_filings']) return results def run_analysis_agents(**context): session_id = context['ti'].xcom_pull(task_ids='initialize_agents') data_results = context['ti'].xcom_pull(task_ids='run_data_gathering_agents') # Activate analysis agents to process gathered data from agent_tasks import FinancialAnalysisTask task = FinancialAnalysisTask(session_id=session_id) analysis = task.execute(data=data_results) return analysis def generate_report(**context): session_id = context['ti'].xcom_pull(task_ids='initialize_agents') analysis = context['ti'].xcom_pull(task_ids='run_analysis_agents') # Generate final report from analysis from agent_tasks import ReportGenerationTask task = ReportGenerationTask(session_id=session_id) report_path = task.execute(analysis=analysis, format='pdf') return report_path with DAG('financial_analysis_agents', default_args=default_args, schedule_interval='0 4 * * 1-5', # Weekdays at 4 AM catchup=False) as dag: init_task = PythonOperator( task_id='initialize_agents', python_callable=initialize_agents, ) gather_task = PythonOperator( task_id='run_data_gathering_agents', python_callable=run_data_gathering_agents, ) analysis_task = PythonOperator( task_id='run_analysis_agents', python_callable=run_analysis_agents, ) report_task = PythonOperator( task_id='generate_report', python_callable=generate_report, ) init_task >> gather_task >> analysis_task >> report_task ``` Airflow considerations for AI agent workflows: - **XComs for Agent Context**: Use XComs to pass context and state between agent tasks - **Dynamic Task Generation**: Generate tasks based on agent discovery results - **Sensor Operators**: Use sensors to wait for external events before triggering agents - **Task Pools**: Limit concurrent agent execution to prevent API rate limiting ### The Role of Vector Databases: ChromaDB, Pinecone, and Snowflake Vector databases form the knowledge backbone of AI agent systems, enabling semantic search and retrieval across vast information spaces: #### ChromaDB for Embedded Agent Knowledge ChromaDB offers a lightweight, embeddable vector store ideal for agents that need fast, local access to domain knowledge: ```python # Setting up ChromaDB for agent knowledge storage import chromadb from chromadb.config import Settings from chromadb.utils import embedding_functions # Configure custom embedding function with caching openai_ef = embedding_functions.OpenAIEmbeddingFunction( api_key=os.environ.get("OPENAI_API_KEY"), model_name="text-embedding-ada-002" ) # Initialize client with persistence and caching client = chromadb.Client( Settings( chroma_db_impl="duckdb+parquet", persist_directory="/data/chroma_storage", anonymized_telemetry=False ) ) # Create collection for domain-specific knowledge knowledge_collection = client.create_collection( name="agent_domain_knowledge", embedding_function=openai_ef, metadata={"domain": "financial_analysis", "version": "2023-Q4"} ) # Add domain knowledge with metadata for filtering knowledge_collection.add( documents=[ "The price-to-earnings ratio (P/E ratio) is the ratio of a company's share price to the company's earnings per share.", # More knowledge entries... ], metadatas=[ {"category": "financial_ratios", "confidence": 0.95, "source": "investopedia"}, # More metadata entries... ], ids=["knowledge_1", "knowledge_2", "knowledge_3"] ) # Example query from an agent seeking information def agent_knowledge_query(query_text, filters=None, n_results=5): results = knowledge_collection.query( query_texts=[query_text], n_results=n_results, where=filters, # e.g., {"category": "financial_ratios"} include=["documents", "metadatas", "distances"] ) # Process results for agent consumption return [{ "content": doc, "metadata": meta, "relevance": 1 - dist # Convert distance to relevance score } for doc, meta, dist in zip( results['documents'][0], results['metadatas'][0], results['distances'][0] )] ``` #### Pinecone for Distributed Agent Knowledge For larger-scale deployments, Pinecone provides a fully managed vector database with high availability and global distribution: ```python # Integrating Pinecone with agents for scalable knowledge retrieval import pinecone import openai # Initialize Pinecone pinecone.init( api_key=os.environ.get("PINECONE_API_KEY"), environment="us-west1-gcp" ) # Create or connect to existing index index_name = "agent-knowledge-base" if index_name not in pinecone.list_indexes(): pinecone.create_index( name=index_name, dimension=1536, # OpenAI embedding dimension metric="cosine", shards=2, # Scale based on data size pods=2 # For high availability ) index = pinecone.Index(index_name) # Function for agents to retrieve contextual knowledge def retrieve_agent_context(query, namespace="general", top_k=5, filters=None): # Generate embedding for query query_embedding = openai.Embedding.create( input=query, model="text-embedding-ada-002" )["data"][0]["embedding"] # Query Pinecone with metadata filtering results = index.query( vector=query_embedding, top_k=top_k, namespace=namespace, filter=filters, # e.g., {"domain": "healthcare", "confidence": {"$gt": 0.8}} include_metadata=True ) # Extract and format knowledge for agent consumption context_items = [] for match in results["matches"]: context_items.append({ "text": match["metadata"]["text"], "source": match["metadata"]["source"], "score": match["score"], "domain": match["metadata"].get("domain", "general") }) return context_items ``` #### Snowflake for Enterprise-Grade Vector Search For organizations already invested in Snowflake's data cloud, the vector search capabilities provide seamless integration with existing data governance: ```sql -- Create a Snowflake table with vector support for agent knowledge CREATE OR REPLACE TABLE agent_knowledge_base ( id VARCHAR NOT NULL, content TEXT, embedding VECTOR(1536), category VARCHAR, source VARCHAR, confidence FLOAT, created_at TIMESTAMP_NTZ DEFAULT CURRENT_TIMESTAMP(), PRIMARY KEY (id) ); -- Create vector search index CREATE OR REPLACE VECTOR SEARCH INDEX agent_knowledge_idx ON agent_knowledge_base(embedding) TYPE = 'DOT_PRODUCT' OPTIONS = (optimization_level = 9); -- Query example for agent knowledge retrieval SELECT id, content, category, source, confidence, VECTOR_DOT_PRODUCT(embedding, '{0.2, 0.1, ..., 0.5}') as relevance FROM agent_knowledge_base WHERE category = 'financial_reporting' AND confidence > 0.8 ORDER BY relevance DESC LIMIT 10; ``` Python integration with Snowflake vector search for agents: ```python # Snowflake vector search integration for enterprise agents import snowflake.connector from snowflake.connector.pandas_tools import write_pandas import pandas as pd import numpy as np import openai # Establish Snowflake connection conn = snowflake.connector.connect( user=os.environ.get('SNOWFLAKE_USER'), password=os.environ.get('SNOWFLAKE_PASSWORD'), account=os.environ.get('SNOWFLAKE_ACCOUNT'), warehouse='AGENT_WAREHOUSE', database='AGENT_DB', schema='KNOWLEDGE' ) # Function for agents to query enterprise knowledge def enterprise_knowledge_query(query_text, filters=None, limit=5): # Generate embedding via OpenAI API embedding_resp = openai.Embedding.create( input=query_text, model="text-embedding-ada-002" ) embedding_vector = embedding_resp["data"][0]["embedding"] # Convert to string format for Snowflake vector_str = str(embedding_vector).replace('[', '{').replace(']', '}') # Construct query with filters filter_clause = "" if filters: filter_conditions = [] for key, value in filters.items(): if isinstance(value, str): filter_conditions.append(f"{key} = '{value}'") elif isinstance(value, (int, float)): filter_conditions.append(f"{key} = {value}") elif isinstance(value, dict) and "$gt" in value: filter_conditions.append(f"{key} > {value['$gt']}") if filter_conditions: filter_clause = "AND " + " AND ".join(filter_conditions) # Execute vector similarity search cursor = conn.cursor() query = f""" SELECT id, content, category, source, confidence, VECTOR_DOT_PRODUCT(embedding, '{vector_str}') as relevance FROM agent_knowledge_base WHERE 1=1 {filter_clause} ORDER BY relevance DESC LIMIT {limit}; """ cursor.execute(query) results = cursor.fetchall() # Convert to more usable format for agents columns = ["id", "content", "category", "source", "confidence", "relevance"] return pd.DataFrame(results, columns=columns).to_dict(orient="records") ``` Key considerations for vector databases in agent systems: - **Filtering Strategy**: Implement metadata filtering to contextualize knowledge retrieval - **Embedding Caching**: Cache embeddings to reduce API calls and latency - **Hybrid Search**: Combine vector search with keyword search for better results - **Knowledge Refresh**: Implement strategies for updating agent knowledge when information changes By understanding these core infrastructure components, you can architect AI agent systems that are both scalable and maintainable. In the next section, we'll explore how these technologies can be combined into complete tech stacks for different use cases. ## 3. Tech Stack Breakdown for Large-Scale AI Agents ### Kubernetes + Ray Serve + AutoGen + LangChain (Distributed AI Workloads) This stack is optimized for computationally intensive AI agent workloads that require dynamic scaling and resource allocation. It's particularly well-suited for organizations running sophisticated simulations, complex reasoning chains, or high-throughput data processing with AI agents. **Architecture Overview:** ![Kubernetes + Ray + AutoGen Architecture](https://i.imgur.com/FrghX58.png) The architecture consists of the following components: 1. **Kubernetes**: Provides the container orchestration layer 2. **Ray**: Handles distributed computing and resource allocation 3. **Ray Serve**: Manages model serving and request routing 4. **AutoGen**: Orchestrates multi-agent interactions 5. **LangChain**: Provides tools, utilities, and integration capabilities **Implementation Example:** First, let's define our Ray cluster configuration for Kubernetes: ```yaml # ray-cluster.yaml apiVersion: ray.io/v1 kind: RayCluster metadata: name: ray-autogen-cluster spec: rayVersion: '2.9.0' headGroupSpec: rayStartParams: dashboard-host: '0.0.0.0' block: 'true' template: spec: containers: - name: ray-head image: rayproject/ray:2.9.0-py310 ports: - containerPort: 6379 name: gcs - containerPort: 8265 name: dashboard - containerPort: 10001 name: client resources: limits: cpu: 4 memory: 8Gi requests: cpu: 2 memory: 4Gi workerGroupSpecs: - groupName: agent-workers replicas: 3 minReplicas: 1 maxReplicas: 10 rayStartParams: {} template: spec: containers: - name: ray-worker image: rayproject/ray:2.9.0-py310 resources: limits: cpu: 8 memory: 16Gi nvidia.com/gpu: 1 requests: cpu: 4 memory: 8Gi ``` Next, let's implement our distributed agent system using Ray and AutoGen: ```python # distributed_agents.py import ray from ray import serve import autogen from autogen.agentchat.contrib.gpt_assistant_agent import GPTAssistantAgent from langchain.embeddings import OpenAIEmbeddings from langchain.vectorstores import FAISS import os import json import time # Initialize Ray ray.init(address="auto") # Define agent factory as Ray actors for distributed creation @ray.remote class AgentFactory: def __init__(self, config_path): with open(config_path, 'r') as f: self.config = json.load(f) # Pre-initialize embeddings for knowledge retrieval self.embeddings = OpenAIEmbeddings( openai_api_key=os.environ.get("OPENAI_API_KEY") ) def create_agent(self, role, knowledge_base=None): """Create an agent with specified role and knowledge.""" # Load role-specific configuration role_config = self.config.get(role, {}) # Initialize knowledge retrieval if specified retriever = None if knowledge_base: vectorstore = FAISS.load_local( f"knowledge_bases/{knowledge_base}", self.embeddings ) retriever = vectorstore.as_retriever( search_kwargs={"k": 5} ) # Create the appropriate agent based on role if role == "manager": return autogen.AssistantAgent( name="Manager", system_message=role_config.get("system_message", ""), llm_config={ "config_list": self.config.get("llm_config_list"), "temperature": 0.2, "timeout": 300, "cache_seed": 42 } ) elif role == "specialist": # Create a specialist agent with custom tools function_map = { "search_knowledge": lambda query: retriever.get_relevant_documents(query) if retriever else [], "run_analysis": self._run_analysis, "generate_report": self._generate_report } return autogen.AssistantAgent( name=f"Specialist_{int(time.time())}", system_message=role_config.get("system_message", ""), llm_config={ "config_list": self.config.get("llm_config_list"), "functions": role_config.get("functions", []), "temperature": 0.4, "timeout": 600, "cache_seed": 42 }, function_map=function_map ) else: # Default case - create standard assistant return autogen.AssistantAgent( name=role.capitalize(), system_message=role_config.get("system_message", ""), llm_config={ "config_list": self.config.get("llm_config_list"), "temperature": 0.7, "timeout": 120, "cache_seed": 42 } ) def _run_analysis(self, data, params=None): """Run specialized analysis (would contain actual implementation)""" # Simulate complex computation time.sleep(2) return {"status": "success", "results": "Analysis complete"} def _generate_report(self, analysis_results, format="pdf"): """Generate report from analysis results""" # Simulate report generation time.sleep(1) return {"report_url": f"https://reports.example.com/{int(time.time())}.{format}"} # Agent coordination service using Ray Serve @serve.deployment( num_replicas=2, max_concurrent_queries=10, ray_actor_options={"num_cpus": 2, "num_gpus": 0.1} ) class AgentCoordinator: def __init__(self, config_path): self.config_path = config_path self.agent_factory = AgentFactory.remote(config_path) async def create_agent_team(self, task_description, team_spec): """Create a team of agents to solve a specific task.""" # Initialize agent references agents = {} for role, spec in team_spec.items(): agent_ref = await self.agent_factory.create_agent.remote( role, knowledge_base=spec.get("knowledge_base") ) agents[role] = agent_ref # Create user proxy agent for orchestration user_proxy = autogen.UserProxyAgent( name="TaskCoordinator", human_input_mode="NEVER", max_consecutive_auto_reply=10, code_execution_config={"work_dir": "temp_workspace"} ) return { "agents": agents, "user_proxy": user_proxy, "task": task_description } async def execute_agent_workflow(self, team_data, workflow_type="sequential"): """Execute a multi-agent workflow.""" agents = team_data["agents"] user_proxy = team_data["user_proxy"] task = team_data["task"] if workflow_type == "sequential": # Sequential workflow where agents work one after another result = await self._run_sequential_workflow(user_proxy, agents, task) elif workflow_type == "groupchat": # Group chat workflow where agents collaborate simultaneously result = await self._run_groupchat_workflow(user_proxy, agents, task) else: result = {"error": "Unsupported workflow type"} return result async def _run_sequential_workflow(self, user_proxy, agents, task): """Run a sequential workflow where each agent builds on previous work.""" # Implementation would include sequential agent invocation current_state = {"task": task, "progress": []} for role, agent in agents.items(): # Update the prompt with current state agent_prompt = f"Task: {task}\nCurrent Progress: {current_state['progress']}\n" agent_prompt += f"Your role as {role} is to advance this task." # Start conversation with this agent chat_result = user_proxy.initiate_chat( agent, message=agent_prompt ) # Extract result and add to progress result_summary = self._extract_agent_result(chat_result) current_state["progress"].append({ "role": role, "contribution": result_summary }) return current_state async def _run_groupchat_workflow(self, user_proxy, agents, task): """Run a group chat workflow where agents collaborate.""" # Create a group chat with all agents group_chat = autogen.GroupChat( agents=list(agents.values()), messages=[], max_round=12 ) manager = autogen.GroupChatManager( groupchat=group_chat, llm_config={"config_list": self.config.get("llm_config_list")} ) # Start the group discussion chat_result = user_proxy.initiate_chat( manager, message=f"Task for the group: {task}" ) return { "task": task, "discussion": chat_result.chat_history, "summary": self._summarize_discussion(chat_result.chat_history) } def _extract_agent_result(self, chat_result): """Extract the key results from an agent conversation.""" # Implementation would parse and structure agent outputs return "Extracted result summary" def _summarize_discussion(self, chat_history): """Summarize the outcome of a group discussion.""" # Implementation would create a concise summary return "Discussion summary" # Deploy the agent coordinator service agent_coordinator_deployment = AgentCoordinator.bind("config/agent_config.json") serve.run(agent_coordinator_deployment, name="agent-coordinator") print("Agent Coordinator service deployed and ready") ``` **Client Implementation:** ```python # client.py import ray import requests import json import asyncio # Connect to the Ray cluster ray.init(address="auto") async def run_distributed_agent_task(): # Define the team structure for a specific task team_spec = { "manager": { "knowledge_base": "corporate_policies" }, "analyst": { "knowledge_base": "financial_data" }, "engineer": { "knowledge_base": "technical_specs" }, "compliance": { "knowledge_base": "regulations" } } # Define the task task_description = """ Analyze our Q4 financial performance and recommend infrastructure improvements that would optimize cost efficiency while maintaining compliance with our industry regulations. """ # Send request to create the agent team response = requests.post( "http://localhost:8000/agent-coordinator/create_agent_team", json={ "task_description": task_description, "team_spec": team_spec } ) team_data = response.json() # Execute the workflow with the created team workflow_response = requests.post( "http://localhost:8000/agent-coordinator/execute_agent_workflow", json={ "team_data": team_data, "workflow_type": "groupchat" } ) results = workflow_response.json() # Process and display the results print(f"Task Results: {json.dumps(results, indent=2)}") return results if __name__ == "__main__": asyncio.run(run_distributed_agent_task()) ``` **Key Advantages:** 1. **Elastic Scaling**: Kubernetes automatically scales worker nodes based on demand 2. **Resource Efficiency**: Ray efficiently distributes workloads across available resources 3. **Fault Tolerance**: Ray handles node failures and task retries automatically 4. **Distributed State**: Ray's object store maintains consistent state across distributed agents 5. **High Performance**: Direct communication between Ray actors minimizes latency **Production Considerations:** 1. **Observability**: Implement detailed logging and monitoring: ```python # Structured logging for distributed agents import structlog import ray from ray import serve # Configure structured logger structlog.configure( processors=[ structlog.processors.TimeStamper(fmt="iso"), structlog.processors.JSONRenderer() ] ) # Create logger instance logger = structlog.get_logger() # Example instrumented agent class @ray.remote class InstrumentedAgent: def __init__(self, agent_id, role): self.agent_id = agent_id self.role = role self.logger = logger.bind( component="agent", agent_id=agent_id, role=role ) self.metrics = { "tasks_completed": 0, "tokens_consumed": 0, "average_response_time": 0, "errors": 0 } def process_task(self, task_data): start_time = time.time() # Log task initiation self.logger.info( "task_started", task_id=task_data.get("id"), task_type=task_data.get("type") ) try: # Task processing logic would go here result = self._execute_agent_task(task_data) # Update metrics elapsed = time.time() - start_time self.metrics["tasks_completed"] += 1 self.metrics["tokens_consumed"] += result.get("tokens_used", 0) self.metrics["average_response_time"] = ( (self.metrics["average_response_time"] * (self.metrics["tasks_completed"] - 1) + elapsed) / self.metrics["tasks_completed"] ) # Log successful completion self.logger.info( "task_completed", task_id=task_data.get("id"), duration=elapsed, tokens_used=result.get("tokens_used", 0) ) return result except Exception as e: # Update error metrics self.metrics["errors"] += 1 # Log error with details self.logger.error( "task_failed", task_id=task_data.get("id"), error=str(e), duration=time.time() - start_time, exception_type=type(e).__name__ ) # Re-raise or return error response raise def get_metrics(self): """Return current agent metrics.""" return { **self.metrics, "agent_id": self.agent_id, "role": self.role, "timestamp": time.time() } def _execute_agent_task(self, task_data): # Actual implementation would go here return {"status": "success", "tokens_used": 150} ``` 2. **Security**: Configure proper isolation and permissions: ```yaml # security-context.yaml apiVersion: v1 kind: ServiceAccount metadata: name: agent-service-account namespace: ai-agents --- apiVersion: rbac.authorization.k8s.io/v1 kind: Role metadata: namespace: ai-agents name: agent-role rules: - apiGroups: [""] resources: ["pods", "pods/log"] verbs: ["get", "list", "watch"] - apiGroups: [""] resources: ["configmaps"] verbs: ["get", "list"] --- apiVersion: rbac.authorization.k8s.io/v1 kind: RoleBinding metadata: name: agent-role-binding namespace: ai-agents subjects: - kind: ServiceAccount name: agent-service-account namespace: ai-agents roleRef: kind: Role name: agent-role apiGroup: rbac.authorization.k8s.io ``` 3. **API Key Management**: Use Kubernetes secrets for secure credential management: ```yaml # agent-secrets.yaml apiVersion: v1 kind: Secret metadata: name: agent-api-keys namespace: ai-agents type: Opaque data: openai-api-key: base64encodedkey pinecone-api-key: base64encodedkey ``` 4. **Persistent Storage**: Configure persistent volumes for agent data: ```yaml # agent-storage.yaml apiVersion: v1 kind: PersistentVolumeClaim metadata: name: agent-data-pvc namespace: ai-agents spec: accessModes: - ReadWriteOnce resources: requests: storage: 10Gi storageClassName: standard ``` This stack is particularly well-suited for organizations that need to: - Process large volumes of data with AI agents - Run complex simulations or forecasting models - Support high-throughput API services backed by AI agents - Dynamically allocate computing resources based on demand ### Apache Kafka + FastAPI + AutoGen + ChromaDB (Real-Time AI Pipelines) This stack is optimized for event-driven, real-time AI agent systems that need to process streaming data and respond to events as they occur. It's ideal for applications like fraud detection, real-time monitoring, and event-based workflow automation. **Architecture Overview:** ![Kafka + FastAPI + AutoGen Architecture](https://i.imgur.com/WBzdMXq.png) The architecture consists of: 1. **Apache Kafka**: Event streaming platform for high-throughput message processing 2. **FastAPI**: High-performance API framework for agent endpoints and services 3. **AutoGen**: Multi-agent orchestration framework 4. **ChromaDB**: Vector database for efficient knowledge retrieval 5. **Redis**: Cache for agent state and session management **Implementation Example:** First, let's set up our environment with Docker Compose: ```yaml # docker-compose.yml version: '3' services: zookeeper: image: confluentinc/cp-zookeeper:7.3.0 environment: ZOOKEEPER_CLIENT_PORT: 2181 healthcheck: test: ["CMD", "nc", "-z", "localhost", "2181"] interval: 10s timeout: 5s retries: 5 kafka: image: confluentinc/cp-kafka:7.3.0 depends_on: zookeeper: condition: service_healthy ports: - "9092:9092" environment: KAFKA_BROKER_ID: 1 KAFKA_ZOOKEEPER_CONNECT: zookeeper:2181 KAFKA_ADVERTISED_LISTENERS: PLAINTEXT://kafka:29092,PLAINTEXT_HOST://localhost:9092 KAFKA_LISTENER_SECURITY_PROTOCOL_MAP: PLAINTEXT:PLAINTEXT,PLAINTEXT_HOST:PLAINTEXT KAFKA_INTER_BROKER_LISTENER_NAME: PLAINTEXT KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR: 1 KAFKA_TRANSACTION_STATE_LOG_REPLICATION_FACTOR: 1 healthcheck: test: ["CMD", "kafka-topics", "--bootstrap-server", "localhost:9092", "--list"] interval: 10s timeout: 5s retries: 5 redis: image: redis:7.0-alpine ports: - "6379:6379" healthcheck: test: ["CMD", "redis-cli", "ping"] interval: 5s timeout: 3s retries: 5 chromadb: image: ghcr.io/chroma-core/chroma:0.4.15 ports: - "8000:8000" volumes: - chroma-data:/chroma/chroma environment: - CHROMA_DB_IMPL=duckdb+parquet - CHROMA_PERSIST_DIRECTORY=/chroma/chroma agent-service: build: context: . dockerfile: Dockerfile.agent depends_on: kafka: condition: service_healthy redis: condition: service_healthy chromadb: condition: service_started ports: - "8080:8080" volumes: - ./app:/app - ./data:/data environment: - KAFKA_BOOTSTRAP_SERVERS=kafka:29092 - REDIS_HOST=redis - REDIS_PORT=6379 - CHROMA_HOST=chromadb - CHROMA_PORT=8000 - OPENAI_API_KEY=${OPENAI_API_KEY} command: uvicorn app.main:app --host 0.0.0.0 --port 8080 --reload volumes: chroma-data: ``` Next, let's build our FastAPI application: ```python # app/main.py import os import json import asyncio import uuid from datetime import datetime from typing import Dict, List, Optional, Any import redis import autogen import chromadb from chromadb.utils import embedding_functions from fastapi import FastAPI, BackgroundTasks, HTTPException, Depends from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field from aiokafka import AIOKafkaProducer, AIOKafkaConsumer from contextlib import asynccontextmanager # Models for API requests and responses class AgentRequest(BaseModel): query: str user_id: str context: Optional[Dict[str, Any]] = None agent_type: str = "general" priority: str = "normal" class AgentResponse(BaseModel): request_id: str status: str message: str data: Optional[Dict[str, Any]] = None timestamp: str = Field(default_factory=lambda: datetime.now().isoformat()) class EventData(BaseModel): event_type: str payload: Dict[str, Any] timestamp: str = Field(default_factory=lambda: datetime.now().isoformat()) # Configure Redis client redis_client = redis.Redis( host=os.environ.get("REDIS_HOST", "localhost"), port=int(os.environ.get("REDIS_PORT", 6379)), decode_responses=True ) # Configure ChromaDB client chroma_client = chromadb.HttpClient( host=os.environ.get("CHROMA_HOST", "localhost"), port=int(os.environ.get("CHROMA_PORT", 8000)) ) # Configure embedding function openai_ef = embedding_functions.OpenAIEmbeddingFunction( api_key=os.environ.get("OPENAI_API_KEY"), model_name="text-embedding-ada-002" ) # Ensure collections exist try: knowledge_collection = chroma_client.get_collection( name="agent_knowledge", embedding_function=openai_ef ) except: knowledge_collection = chroma_client.create_collection( name="agent_knowledge", embedding_function=openai_ef ) # Kafka configuration KAFKA_BOOTSTRAP_SERVERS = os.environ.get("KAFKA_BOOTSTRAP_SERVERS", "localhost:9092") AGENT_REQUEST_TOPIC = "agent_requests" AGENT_RESPONSE_TOPIC = "agent_responses" EVENT_TOPIC = "system_events" # Initialize Kafka producer producer = None # Lifespan manager for FastAPI to handle async setup/teardown @asynccontextmanager async def lifespan(app: FastAPI): # Setup: create global producer global producer producer = AIOKafkaProducer( bootstrap_servers=KAFKA_BOOTSTRAP_SERVERS, value_serializer=lambda v: json.dumps(v).encode('utf-8'), acks='all', enable_idempotence=True, retries=3 ) # Start the producer await producer.start() # Start the consumer tasks consumer_task = asyncio.create_task(consume_agent_responses()) event_consumer_task = asyncio.create_task(consume_events()) # Yield control back to FastAPI yield # Cleanup: close producer and cancel consumer tasks consumer_task.cancel() event_consumer_task.cancel() try: await consumer_task except asyncio.CancelledError: pass try: await event_consumer_task except asyncio.CancelledError: pass await producer.stop() # Initialize FastAPI app = FastAPI(lifespan=lifespan) # Configure CORS app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Agent factory for creating different agent types class AgentFactory: @staticmethod def create_agent(agent_type, context=None): llm_config = { "config_list": [ { "model": "gpt-4", "api_key": os.environ.get("OPENAI_API_KEY") } ], "temperature": 0.5, "timeout": 120, "cache_seed": None # Disable caching for real-time applications } if agent_type == "analyst": system_message = """You are a financial analyst agent specialized in market trends and investment strategies. Analyze data precisely and provide actionable insights.""" # Add specific functions for analyst llm_config["functions"] = [ { "name": "analyze_market_data", "description": "Analyze market data to identify trends and opportunities", "parameters": { "type": "object", "properties": { "sector": {"type": "string", "description": "Market sector to analyze"}, "timeframe": {"type": "string", "description": "Timeframe for analysis (e.g., '1d', '1w', '1m')"}, "metrics": {"type": "array", "items": {"type": "string"}, "description": "Metrics to include in analysis"} }, "required": ["sector", "timeframe"] } } ] elif agent_type == "support": system_message = """You are a customer support agent that helps users with technical issues and product questions. Be empathetic and solution-oriented.""" else: # default general agent system_message = """You are a helpful AI assistant that provides accurate and concise information on a wide range of topics.""" # Create the agent agent = autogen.AssistantAgent( name=f"{agent_type.capitalize()}Agent", system_message=system_message, llm_config=llm_config ) return agent # Function to retrieve relevant knowledge for agent context async def retrieve_knowledge(query, filters=None, limit=5): try: results = knowledge_collection.query( query_texts=[query], n_results=limit, where=filters ) if results and len(results['documents']) > 0: documents = results['documents'][0] metadatas = results['metadatas'][0] if 'metadatas' in results else [{}] * len(documents) distances = results['distances'][0] if 'distances' in results else [1.0] * len(documents) return [ { "content": doc, "metadata": meta, "relevance": 1 - dist if dist <= 1 else 0 } for doc, meta, dist in zip(documents, metadatas, distances) ] return [] except Exception as e: print(f"Error retrieving knowledge: {e}") return [] # Background task to process agent request through Kafka async def process_agent_request(request_id: str, request: AgentRequest): try: # Publish request to Kafka await producer.send_and_wait( AGENT_REQUEST_TOPIC, { "request_id": request_id, "query": request.query, "user_id": request.user_id, "context": request.context, "agent_type": request.agent_type, "priority": request.priority, "timestamp": datetime.now().isoformat() } ) # Update request status in Redis redis_client.hset( f"request:{request_id}", mapping={ "status": "processing", "timestamp": datetime.now().isoformat() } ) redis_client.expire(f"request:{request_id}", 3600) # Expire after 1 hour # Publish event await producer.send_and_wait( EVENT_TOPIC, { "event_type": "agent_request_received", "payload": { "request_id": request_id, "user_id": request.user_id, "agent_type": request.agent_type, "priority": request.priority }, "timestamp": datetime.now().isoformat() } ) except Exception as e: # Update request status in Redis redis_client.hset( f"request:{request_id}", mapping={ "status": "error", "error": str(e), "timestamp": datetime.now().isoformat() } ) redis_client.expire(f"request:{request_id}", 3600) # Expire after 1 hour # Publish error event await producer.send_and_wait( EVENT_TOPIC, { "event_type": "agent_request_error", "payload": { "request_id": request_id, "error": str(e) }, "timestamp": datetime.now().isoformat() } ) # Consumer for agent responses async def consume_agent_responses(): consumer = AIOKafkaConsumer( AGENT_RESPONSE_TOPIC, bootstrap_servers=KAFKA_BOOTSTRAP_SERVERS, group_id="agent-service-group", value_deserializer=lambda m: json.loads(m.decode('utf-8')), auto_offset_reset="latest", enable_auto_commit=True ) await consumer.start() try: async for message in consumer: response_data = message.value request_id = response_data.get("request_id") if request_id: # Update response in Redis redis_client.hset( f"request:{request_id}", mapping={ "status": "completed", "response": json.dumps(response_data), "completed_at": datetime.now().isoformat() } ) # Publish completion event await producer.send( EVENT_TOPIC, { "event_type": "agent_response_completed", "payload": { "request_id": request_id, "processing_time": response_data.get("processing_time") }, "timestamp": datetime.now().isoformat() } ) finally: await consumer.stop() # Consumer for system events async def consume_events(): consumer = AIOKafkaConsumer( EVENT_TOPIC, bootstrap_servers=KAFKA_BOOTSTRAP_SERVERS, group_id="event-processor-group", value_deserializer=lambda m: json.loads(m.decode('utf-8')), auto_offset_reset="latest", enable_auto_commit=True ) await consumer.start() try: async for message in consumer: event_data = message.value event_type = event_data.get("event_type") # Process different event types if event_type == "agent_request_received": # Metrics tracking pass elif event_type == "agent_response_completed": # Performance monitoring pass elif event_type == "agent_request_error": # Error handling and alerting pass finally: await consumer.stop() # Worker process that handles agent requests from Kafka async def agent_worker(): consumer = AIOKafkaConsumer( AGENT_REQUEST_TOPIC, bootstrap_servers=KAFKA_BOOTSTRAP_SERVERS, group_id="agent-worker-group", value_deserializer=lambda m: json.loads(m.decode('utf-8')), auto_offset_reset="earliest", enable_auto_commit=True ) await consumer.start() try: async for message in consumer: request_data = message.value request_id = request_data.get("request_id") start_time = datetime.now() try: # Create appropriate agent type agent = AgentFactory.create_agent( request_data.get("agent_type", "general"), context=request_data.get("context") ) # Retrieve relevant knowledge knowledge = await retrieve_knowledge( request_data.get("query"), filters={"domain": request_data.get("context", {}).get("domain")} if request_data.get("context") else None ) # Create user proxy agent for handling the conversation user_proxy = autogen.UserProxyAgent( name="User", human_input_mode="NEVER", max_consecutive_auto_reply=0, code_execution_config={"work_dir": "workspace"} ) # Build the prompt with knowledge context knowledge_context = "" if knowledge: knowledge_context = "\n\nRelevant context:\n" for i, item in enumerate(knowledge, 1): knowledge_context += f"{i}. {item['content']}\n" query = request_data.get("query") message = f"{query}\n{knowledge_context}" # Start conversation with agent user_proxy.initiate_chat(agent, message=message) # Extract agent response response_content = "" if user_proxy.chat_history and len(user_proxy.chat_history) > 1: # Get the last message from the agent for msg in reversed(user_proxy.chat_history): if msg.get("role") == "assistant": response_content = msg.get("content", "") break processing_time = (datetime.now() - start_time).total_seconds() # Send response back through Kafka await producer.send_and_wait( AGENT_RESPONSE_TOPIC, { "request_id": request_id, "status": "success", "response": response_content, "processing_time": processing_time, "timestamp": datetime.now().isoformat() } ) except Exception as e: processing_time = (datetime.now() - start_time).total_seconds() # Send error response await producer.send_and_wait( AGENT_RESPONSE_TOPIC, { "request_id": request_id, "status": "error", "error": str(e), "processing_time": processing_time, "timestamp": datetime.now().isoformat() } ) finally: await consumer.stop() # Routes @app.post("/api/agent/request", response_model=AgentResponse) async def create_agent_request(request: AgentRequest, background_tasks: BackgroundTasks): request_id = str(uuid.uuid4()) # Store initial request in Redis redis_client.hset( f"request:{request_id}", mapping={ "user_id": request.user_id, "query": request.query, "agent_type": request.agent_type, "status": "pending", "timestamp": datetime.now().isoformat() } ) redis_client.expire(f"request:{request_id}", 3600) # Expire after 1 hour # Process request asynchronously background_tasks.add_task(process_agent_request, request_id, request) return AgentResponse( request_id=request_id, status="pending", message="Request has been submitted for processing" ) @app.get("/api/agent/status/{request_id}", response_model=AgentResponse) async def get_agent_status(request_id: str): # Check if request exists in Redis request_data = redis_client.hgetall(f"request:{request_id}") if not request_data: raise HTTPException(status_code=404, detail="Request not found") status = request_data.get("status", "unknown") if status == "completed": # Return the completed response response_data = json.loads(request_data.get("response", "{}")) return AgentResponse( request_id=request_id, status=status, message="Request completed", data=response_data ) elif status == "error": # Return error details return AgentResponse( request_id=request_id, status=status, message=f"Error processing request: {request_data.get('error', 'Unknown error')}", ) else: # Return current status return AgentResponse( request_id=request_id, status=status, message=f"Request is currently {status}" ) @app.post("/api/events/publish", status_code=202) async def publish_event(event: EventData): try: await producer.send_and_wait( EVENT_TOPIC, event.dict() ) return {"status": "success", "message": "Event published successfully"} except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to publish event: {str(e)}") @app.post("/api/knowledge/add") async def add_knowledge(items: List[Dict[str, Any]]): try: documents = [item["content"] for item in items] metadata = [item.get("metadata", {}) for item in items] ids = [f"doc_{uuid.uuid4()}" for _ in items] knowledge_collection.add( documents=documents, metadatas=metadata, ids=ids ) return {"status": "success", "message": f"Added {len(documents)} items to knowledge base"} except Exception as e: raise HTTPException(status_code=500, detail=f"Failed to add knowledge: {str(e)}") @app.get("/api/knowledge/search") async def search_knowledge(query: str, limit: int = 5): knowledge = await retrieve_knowledge(query, limit=limit) return {"results": knowledge} # Start worker process when app starts @app.on_event("startup") async def startup_event(): asyncio.create_task(agent_worker()) if __name__ == "__main__": import uvicorn uvicorn.run("main:app", host="0.0.0.0", port=8080, reload=True) ``` **Client Implementation Example:** ```python # client.py import asyncio import json import time import uuid from typing import Dict, List, Any, Optional import aiohttp from pydantic import BaseModel class AgentClient: def __init__(self, base_url: str = "http://localhost:8080"): self.base_url = base_url self.session = None async def __aenter__(self): self.session = aiohttp.ClientSession() return self async def __aexit__(self, exc_type, exc_val, exc_tb): if self.session: await self.session.close() async def submit_request(self, query: str, user_id: str, agent_type: str = "general", context: Optional[Dict[str, Any]] = None, priority: str = "normal") -> Dict[str, Any]: """Submit a request to the agent service.""" if self.session is None: self.session = aiohttp.ClientSession() payload = { "query": query, "user_id": user_id, "agent_type": agent_type, "priority": priority } if context: payload["context"] = context async with self.session.post( f"{self.base_url}/api/agent/request", json=payload ) as response: if response.status == 200: return await response.json() else: error_text = await response.text() raise Exception(f"Error submitting request: {error_text}") async def get_request_status(self, request_id: str) -> Dict[str, Any]: """Get the status of a request.""" if self.session is None: self.session = aiohttp.ClientSession() async with self.session.get( f"{self.base_url}/api/agent/status/{request_id}" ) as response: if response.status == 200: return await response.json() else: error_text = await response.text() raise Exception(f"Error getting status: {error_text}") async def wait_for_completion(self, request_id: str, polling_interval: float = 1.0, timeout: float = 120.0) -> Dict[str, Any]: """Wait for a request to complete, with timeout.""" start_time = time.time() while (time.time() - start_time) < timeout: status_response = await self.get_request_status(request_id) if status_response["status"] in ["completed", "error"]: return status_response await asyncio.sleep(polling_interval) raise TimeoutError(f"Request {request_id} did not complete within {timeout} seconds") async def add_knowledge(self, items: List[Dict[str, Any]]) -> Dict[str, Any]: """Add knowledge items to the agent knowledge base.""" if self.session is None: self.session = aiohttp.ClientSession() async with self.session.post( f"{self.base_url}/api/knowledge/add", json=items ) as response: if response.status == 200: return await response.json() else: error_text = await response.text() raise Exception(f"Error adding knowledge: {error_text}") async def search_knowledge(self, query: str, limit: int = 5) -> Dict[str, Any]: """Search the knowledge base.""" if self.session is None: self.session = aiohttp.ClientSession() async with self.session.get( f"{self.base_url}/api/knowledge/search", params={"query": query, "limit": limit} ) as response: if response.status == 200: return await response.json() else: error_text = await response.text() raise Exception(f"Error searching knowledge: {error_text}") async def main(): # Example client usage async with AgentClient() as client: # Add knowledge first knowledge_items = [ { "content": "Apple Inc. reported Q4 2023 earnings of $1.46 per share, exceeding analyst expectations of $1.39 per share.", "metadata": { "domain": "finance", "category": "earnings", "company": "Apple", "date": "2023-10-27" } }, { "content": "The S&P 500 closed at 4,738.15 on January 12, 2024, up 1.2% for the day.", "metadata": { "domain": "finance", "category": "market_data", "index": "S&P 500", "date": "2024-01-12" } } ] knowledge_result = await client.add_knowledge(knowledge_items) print(f"Knowledge add result: {json.dumps(knowledge_result, indent=2)}") # Submit a request to the analyst agent request_result = await client.submit_request( query="What was Apple's performance in their most recent earnings report?", user_id="test-user-123", agent_type="analyst", context={ "domain": "finance", "analysis_type": "earnings" } ) print(f"Request submitted: {json.dumps(request_result, indent=2)}") request_id = request_result["request_id"] # Wait for completion try: completion_result = await client.wait_for_completion( request_id, polling_interval=2.0, timeout=60.0 ) print(f"Request completed: {json.dumps(completion_result, indent=2)}") except TimeoutError as e: print(f"Request timed out: {e}") if __name__ == "__main__": asyncio.run(main()) ``` **Key Advantages:** 1. **Real-Time Processing**: Stream processing via Kafka enables immediate response to events 2. **Decoupling**: Producers and consumers are decoupled, enabling system resilience 3. **Scalability**: Each component can scale independently based on demand 4. **Event Sourcing**: All interactions are event-based, enabling replay and audit capabilities 5. **Statelessness**: API services remain stateless for easier scaling and deployment **Production Considerations:** 1. **Kafka Topic Configuration**: ```python # kafka_setup.py from kafka.admin import KafkaAdminClient, NewTopic import kafka.errors as errors def setup_kafka_topics(): admin_client = KafkaAdminClient( bootstrap_servers="localhost:9092", client_id="admin-client" ) # Define topics with optimal configurations topic_configs = [ # Agent request topic - moderate throughput with ordered processing NewTopic( name="agent_requests", num_partitions=4, # Balance parallelism and ordering replication_factor=3, # High reliability for requests topic_configs={ "retention.ms": str(7 * 24 * 60 * 60 * 1000), # 7 days retention "cleanup.policy": "delete", "min.insync.replicas": "2", # Ensure at least 2 replicas are in sync "unclean.leader.election.enable": "false", # Prevent data loss "compression.type": "lz4" # Efficient compression } ), # Agent response topic - higher throughput, less ordering dependency NewTopic( name="agent_responses", num_partitions=8, # Higher parallelism for responses replication_factor=3, topic_configs={ "retention.ms": str(7 * 24 * 60 * 60 * 1000), # 7 days retention "cleanup.policy": "delete", "min.insync.replicas": "2", "compression.type": "lz4" } ), # System events topic - high volume, compacted for latest state NewTopic( name="system_events", num_partitions=16, # High parallelism for metrics and events replication_factor=3, topic_configs={ "cleanup.policy": "compact,delete", # Compact for state, delete for retention "delete.retention.ms": str(24 * 60 * 60 * 1000), # 1 day retention after compaction "min.compaction.lag.ms": str(60 * 1000), # 1 minute minimum time before compaction "segment.ms": str(6 * 60 * 60 * 1000), # 6 hour segments "min.insync.replicas": "2", "compression.type": "lz4" } ) ] # Create topics for topic in topic_configs: try: admin_client.create_topics([topic]) print(f"Created topic: {topic.name}") except errors.TopicAlreadyExistsError: print(f"Topic already exists: {topic.name}") admin_client.close() if __name__ == "__main__": setup_kafka_topics() ``` 2. **Monitoring and Metrics**: ```python # monitoring.py import time import json from datadog import initialize, statsd from functools import wraps from contextlib import ContextDecorator # Initialize Datadog client initialize(statsd_host="localhost", statsd_port=8125) class TimingMetric(ContextDecorator): """Context manager/decorator for timing operations and reporting to Datadog.""" def __init__(self, metric_name, tags=None): self.metric_name = metric_name self.tags = tags or [] self.start_time = None def __enter__(self): self.start_time = time.monotonic() return self def __exit__(self, exc_type, exc_val, exc_tb): duration = time.monotonic() - self.start_time # Convert to milliseconds duration_ms = duration * 1000 # Send timing metric statsd.timing(self.metric_name, duration_ms, tags=self.tags) # Also send as gauge for easier aggregation statsd.gauge(f"{self.metric_name}.gauge", duration_ms, tags=self.tags) # If there was an exception, count it if exc_type is not None: statsd.increment( f"{self.metric_name}.error", tags=self.tags + [f"error_type:{exc_type.__name__}"] ) # Don't suppress exceptions return False def timing_decorator(metric_name, tags=None): """Function decorator for timing.""" def decorator(func): @wraps(func) async def async_wrapper(*args, **kwargs): with TimingMetric(metric_name, tags): return await func(*args, **kwargs) @wraps(func) def sync_wrapper(*args, **kwargs): with TimingMetric(metric_name, tags): return func(*args, **kwargs) # Choose the appropriate wrapper based on whether the function is async return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper return decorator # Usage example with modified agent_worker function @timing_decorator("agent.request.processing", tags=["service:agent_worker"]) async def process_agent_request(request_data): # Implementation as before request_id = request_data.get("request_id") agent_type = request_data.get("agent_type", "general") # Increment counter for request by agent type statsd.increment("agent.request.count", tags=[ f"agent_type:{agent_type}", f"priority:{request_data.get('priority', 'normal')}" ]) # Original implementation continues... # ... # At the end, count completion statsd.increment("agent.request.completed", tags=[ f"agent_type:{agent_type}", f"status:success" ]) return result # Kafka consumer with metrics async def consume_agent_responses(): consumer = AIOKafkaConsumer( AGENT_RESPONSE_TOPIC, bootstrap_servers=KAFKA_BOOTSTRAP_SERVERS, group_id="agent-service-group", value_deserializer=lambda m: json.loads(m.decode('utf-8')), auto_offset_reset="latest", enable_auto_commit=True ) await consumer.start() try: async for message in consumer: # Track consumer lag lag = time.time() - message.timestamp/1000 statsd.gauge("kafka.consumer.lag_seconds", lag, tags=[ "topic:agent_responses", "consumer_group:agent-service-group" ]) # Process message with timing with TimingMetric("kafka.message.processing", tags=["topic:agent_responses"]): response_data = message.value # Process as before... # ... finally: await consumer.stop() ``` 3. **Resilience and Circuit Breaking**: ```python # resilience.py import time import asyncio import functools from typing import Callable, Any, Dict, Optional import backoff from fastapi import HTTPException # Circuit breaker implementation class CircuitBreaker: def __init__(self, failure_threshold: int = 5, recovery_timeout: float = 30.0, timeout: float = 10.0, fallback: Optional[Callable] = None): self.failure_threshold = failure_threshold self.recovery_timeout = recovery_timeout self.timeout = timeout self.fallback = fallback self.failure_count = 0 self.last_failure_time = 0 self.state = "CLOSED" # CLOSED, OPEN, HALF-OPEN def __call__(self, func): @functools.wraps(func) async def async_wrapper(*args, **kwargs): if asyncio.iscoroutinefunction(func): return await self._handle_call(func, *args, **kwargs) else: return self._handle_sync_call(func, *args, **kwargs) @functools.wraps(func) def sync_wrapper(*args, **kwargs): return self._handle_sync_call(func, *args, **kwargs) return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper async def _handle_call(self, func, *args, **kwargs): if self.state == "OPEN": if time.time() - self.last_failure_time > self.recovery_timeout: # Move to half-open state and try the call self.state = "HALF-OPEN" else: # Circuit is open, use fallback or raise exception if self.fallback: return await self.fallback(*args, **kwargs) if asyncio.iscoroutinefunction(self.fallback) else self.fallback(*args, **kwargs) else: raise HTTPException(status_code=503, detail="Service temporarily unavailable") try: # Set timeout for the function call result = await asyncio.wait_for(func(*args, **kwargs), timeout=self.timeout) # On success in half-open state, reset the circuit if self.state == "HALF-OPEN": self.failure_count = 0 self.state = "CLOSED" return result except Exception as e: # On failure, increment failure count self.failure_count += 1 self.last_failure_time = time.time() # If failure threshold reached, open the circuit if self.failure_count >= self.failure_threshold: self.state = "OPEN" # If in half-open state, go back to open state if self.state == "HALF-OPEN": self.state = "OPEN" # Use fallback or re-raise the exception if self.fallback: return await self.fallback(*args, **kwargs) if asyncio.iscoroutinefunction(self.fallback) else self.fallback(*args, **kwargs) raise e def _handle_sync_call(self, func, *args, **kwargs): if self.state == "OPEN": if time.time() - self.last_failure_time > self.recovery_timeout: self.state = "HALF-OPEN" else: if self.fallback: return self.fallback(*args, **kwargs) else: raise HTTPException(status_code=503, detail="Service temporarily unavailable") try: # For sync functions, we can't easily apply a timeout # Consider using concurrent.futures.ThreadPoolExecutor with timeout for sync functions result = func(*args, **kwargs) if self.state == "HALF-OPEN": self.failure_count = 0 self.state = "CLOSED" return result except Exception as e: self.failure_count += 1 self.last_failure_time = time.time() if self.failure_count >= self.failure_threshold: self.state = "OPEN" if self.state == "HALF-OPEN": self.state = "OPEN" if self.fallback: return self.fallback(*args, **kwargs) raise e # Exponential backoff with jitter for retries def backoff_llm_call(max_tries=5, max_time=30): def fallback_response(*args, **kwargs): return { "status": "degraded", "message": "Service temporarily in degraded mode. Please try again later." } # Define backoff handler with jitter @backoff.on_exception( backoff.expo, (Exception), # Retry on any exception max_tries=max_tries, max_time=max_time, jitter=backoff.full_jitter, on_backoff=lambda details: print(f"Backing off {details['wait']:0.1f} seconds after {details['tries']} tries") ) @CircuitBreaker( failure_threshold=3, recovery_timeout=60.0, timeout=10.0, fallback=fallback_response ) async def protected_llm_call(client, prompt, model="gpt-4"): # This would be your actual LLM API call response = await client.chat.completions.create( model=model, messages=[{"role": "user", "content": prompt}], max_tokens=1000 ) return response return protected_llm_call # Example usage in agent implementation async def agent_worker(): consumer = AIOKafkaConsumer( AGENT_REQUEST_TOPIC, bootstrap_servers=KAFKA_BOOTSTRAP_SERVERS, group_id="agent-worker-group", value_deserializer=lambda m: json.loads(m.decode('utf-8')), auto_offset_reset="earliest", enable_auto_commit=True ) await consumer.start() # Create OpenAI client from openai import AsyncOpenAI client = AsyncOpenAI(api_key=os.environ.get("OPENAI_API_KEY")) # Create protected LLM call function with backoff and circuit breaker protected_llm = backoff_llm_call(max_tries=3, max_time=45) try: async for message in consumer: request_data = message.value request_id = request_data.get("request_id") start_time = datetime.now() try: # Process with resilience patterns query = request_data.get("query") try: # Make resilient LLM call llm_response = await protected_llm(client, query) # Process response response_content = llm_response.choices[0].message.content # Calculate processing time processing_time = (datetime.now() - start_time).total_seconds() # Send response await producer.send_and_wait( AGENT_RESPONSE_TOPIC, { "request_id": request_id, "status": "success", "response": response_content, "processing_time": processing_time, "timestamp": datetime.now().isoformat() } ) except Exception as e: # Handle failures with proper error reporting processing_time = (datetime.now() - start_time).total_seconds() await producer.send_and_wait( AGENT_RESPONSE_TOPIC, { "request_id": request_id, "status": "error", "error": str(e), "processing_time": processing_time, "timestamp": datetime.now().isoformat() } ) # Also log the error for monitoring statsd.increment("agent.error", tags=[ f"error_type:{type(e).__name__}", f"request_id:{request_id}" ]) except Exception as e: # Catch-all exception handler to prevent worker crashes print(f"Critical error in message processing: {e}") # Log critical errors for immediate attention statsd.increment("agent.critical_error", tags=[ f"error_type:{type(e).__name__}" ]) finally: await consumer.stop() ``` This stack is particularly well-suited for organizations that need to: - Process streaming data in real-time with AI analysis - Build responsive event-driven systems - Implement asynchronous request/response patterns - Support high-throughput messaging with reliable delivery - Maintain a flexible, decoupled architecture ### Django/Flask + Celery + AutoGen + Pinecone (Task Orchestration & Search) This stack is optimized for applications that require robust task management, background processing, and sophisticated search capabilities. It's particularly well-suited for document processing, content recommendation, and task-based AI workflows. **Architecture Overview:** ![Django + Celery + AutoGen Architecture](https://i.imgur.com/U2BKxGF.png) The architecture consists of: 1. **Django/Flask**: Web framework for user interaction and API endpoints 2. **Celery**: Distributed task queue for background processing 3. **Redis/RabbitMQ**: Message broker for Celery tasks 4. **AutoGen**: Multi-agent orchestration framework 5. **Pinecone**: Vector database for semantic search 6. **PostgreSQL**: Relational database for application data and task state **Implementation Example:** Let's implement a Django application with Celery tasks for AI agent processing: First, the Django project structure: ``` project/ ├── manage.py ├── requirements.txt ├── project/ │ ├── __init__.py │ ├── settings.py │ ├── urls.py │ └── celery.py ├── agent_app/ │ ├── __init__.py │ ├── admin.py │ ├── apps.py │ ├── models.py │ ├── serializers.py │ ├── tasks.py │ ├── tests.py │ ├── urls.py │ ├── utils/ │ │ ├── __init__.py │ │ ├── agent_factory.py │ │ ├── pinecone_client.py │ │ └── prompt_templates.py │ └── views.py └── templates/ └── index.html ``` Now, let's implement the core components: ```python # project/settings.py import os from pathlib import Path # Build paths inside the project like this: BASE_DIR / 'subdir'. BASE_DIR = Path(__file__).resolve().parent.parent # SECURITY WARNING: keep the secret key used in production secret! SECRET_KEY = os.environ.get('DJANGO_SECRET_KEY', 'django-insecure-key-for-dev') # SECURITY WARNING: don't run with debug turned on in production! DEBUG = os.environ.get('DJANGO_DEBUG', 'False') == 'True' ALLOWED_HOSTS = os.environ.get('DJANGO_ALLOWED_HOSTS', 'localhost,127.0.0.1').split(',') # Application definition INSTALLED_APPS = [ 'django.contrib.admin', 'django.contrib.auth', 'django.contrib.contenttypes', 'django.contrib.sessions', 'django.contrib.messages', 'django.contrib.staticfiles', 'rest_framework', 'django_celery_results', 'django_celery_beat', 'agent_app', ] MIDDLEWARE = [ 'django.middleware.security.SecurityMiddleware', 'django.contrib.sessions.middleware.SessionMiddleware', 'django.middleware.common.CommonMiddleware', 'django.middleware.csrf.CsrfViewMiddleware', 'django.contrib.auth.middleware.AuthenticationMiddleware', 'django.contrib.messages.middleware.MessageMiddleware', 'django.middleware.clickjacking.XFrameOptionsMiddleware', ] ROOT_URLCONF = 'project.urls' TEMPLATES = [ { 'BACKEND': 'django.template.backends.django.DjangoTemplates', 'DIRS': [BASE_DIR / 'templates'], 'APP_DIRS': True, 'OPTIONS': { 'context_processors': [ 'django.template.context_processors.debug', 'django.template.context_processors.request', 'django.contrib.auth.context_processors.auth', 'django.contrib.messages.context_processors.messages', ], }, }, ] WSGI_APPLICATION = 'project.wsgi.application' # Database DATABASES = { 'default': { 'ENGINE': 'django.db.backends.postgresql', 'NAME': os.environ.get('POSTGRES_DB', 'agent_db'), 'USER': os.environ.get('POSTGRES_USER', 'postgres'), 'PASSWORD': os.environ.get('POSTGRES_PASSWORD', 'postgres'), 'HOST': os.environ.get('POSTGRES_HOST', 'localhost'), 'PORT': os.environ.get('POSTGRES_PORT', '5432'), } } # Internationalization LANGUAGE_CODE = 'en-us' TIME_ZONE = 'UTC' USE_I18N = True USE_TZ = True # Static files (CSS, JavaScript, Images) STATIC_URL = 'static/' STATIC_ROOT = BASE_DIR / 'static' # Default primary key field type DEFAULT_AUTO_FIELD = 'django.db.models.BigAutoField' # Celery settings CELERY_BROKER_URL = os.environ.get('CELERY_BROKER_URL', 'redis://localhost:6379/0') CELERY_RESULT_BACKEND = 'django-db' CELERY_CACHE_BACKEND = 'django-cache' CELERY_ACCEPT_CONTENT = ['json'] CELERY_TASK_SERIALIZER = 'json' CELERY_RESULT_SERIALIZER = 'json' CELERY_TIMEZONE = TIME_ZONE CELERY_TASK_TRACK_STARTED = True CELERY_TASK_TIME_LIMIT = 30 * 60 # 30 minutes CELERY_WORKER_CONCURRENCY = 8 # REST Framework settings REST_FRAMEWORK = { 'DEFAULT_AUTHENTICATION_CLASSES': [ 'rest_framework.authentication.SessionAuthentication', 'rest_framework.authentication.TokenAuthentication', ], 'DEFAULT_PERMISSION_CLASSES': [ 'rest_framework.permissions.IsAuthenticated', ], 'DEFAULT_PAGINATION_CLASS': 'rest_framework.pagination.PageNumberPagination', 'PAGE_SIZE': 20 } # AI Agent settings OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY') PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY') PINECONE_ENVIRONMENT = os.environ.get('PINECONE_ENVIRONMENT', 'us-west1-gcp') PINECONE_INDEX_NAME = os.environ.get('PINECONE_INDEX_NAME', 'agent-knowledge') # Logging configuration LOGGING = { 'version': 1, 'disable_existing_loggers': False, 'formatters': { 'verbose': { 'format': '{levelname} {asctime} {module} {process:d} {thread:d} {message}', 'style': '{', }, 'simple': { 'format': '{levelname} {message}', 'style': '{', }, }, 'handlers': { 'console': { 'level': 'INFO', 'class': 'logging.StreamHandler', 'formatter': 'verbose', }, 'file': { 'level': 'DEBUG', 'class': 'logging.FileHandler', 'filename': os.path.join(BASE_DIR, 'logs/django.log'), 'formatter': 'verbose', }, }, 'loggers': { 'django': { 'handlers': ['console', 'file'], 'level': 'INFO', 'propagate': True, }, 'agent_app': { 'handlers': ['console', 'file'], 'level': 'DEBUG', 'propagate': False, }, }, } # Create logs directory if it doesn't exist os.makedirs(os.path.join(BASE_DIR, 'logs'), exist_ok=True) ``` ```python # project/celery.py import os from celery import Celery # Set the default Django settings module os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'project.settings') # Create the Celery app app = Celery('project') # Using a string here means the worker doesn't have to serialize # the configuration object to child processes. app.config_from_object('django.conf:settings', namespace='CELERY') # Load task modules from all registered Django app configs. app.autodiscover_tasks() # Configure task routes app.conf.task_routes = { 'agent_app.tasks.process_agent_task': {'queue': 'agent_tasks'}, 'agent_app.tasks.update_knowledge_base': {'queue': 'knowledge_tasks'}, 'agent_app.tasks.analyze_document': {'queue': 'document_tasks'}, 'agent_app.tasks.periodic_agent_check': {'queue': 'scheduled_tasks'}, } # Configure task priorities app.conf.task_acks_late = True app.conf.worker_prefetch_multiplier = 1 app.conf.task_inherit_parent_priority = True @app.task(bind=True) def debug_task(self): print(f'Request: {self.request!r}') ``` ```python # project/urls.py from django.contrib import admin from django.urls import path, include urlpatterns = [ path('admin/', admin.site.urls), path('api/', include('agent_app.urls')), ] ``` Now let's define the agent app models: ```python # agent_app/models.py import uuid from django.db import models from django.contrib.auth.models import User from django.utils import timezone class AgentTask(models.Model): """Model for tracking agent tasks and their status.""" STATUS_CHOICES = ( ('pending', 'Pending'), ('processing', 'Processing'), ('completed', 'Completed'), ('failed', 'Failed'), ('canceled', 'Canceled'), ) PRIORITY_CHOICES = ( (1, 'Low'), (2, 'Normal'), (3, 'High'), (4, 'Urgent'), ) TYPE_CHOICES = ( ('analysis', 'Data Analysis'), ('research', 'Research'), ('document', 'Document Processing'), ('conversation', 'Conversation'), ('generation', 'Content Generation'), ('other', 'Other'), ) id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) user = models.ForeignKey(User, on_delete=models.CASCADE, related_name='agent_tasks') title = models.CharField(max_length=255) description = models.TextField() task_type = models.CharField(max_length=20, choices=TYPE_CHOICES, default='other') status = models.CharField(max_length=20, choices=STATUS_CHOICES, default='pending') priority = models.IntegerField(choices=PRIORITY_CHOICES, default=2) input_data = models.JSONField(default=dict, blank=True) result_data = models.JSONField(default=dict, blank=True, null=True) error_message = models.TextField(blank=True, null=True) # Metadata created_at = models.DateTimeField(auto_now_add=True) updated_at = models.DateTimeField(auto_now=True) started_at = models.DateTimeField(null=True, blank=True) completed_at = models.DateTimeField(null=True, blank=True) # Celery task ID for tracking celery_task_id = models.CharField(max_length=255, blank=True, null=True) # Token usage metrics prompt_tokens = models.IntegerField(default=0) completion_tokens = models.IntegerField(default=0) total_tokens = models.IntegerField(default=0) estimated_cost = models.DecimalField(max_digits=10, decimal_places=6, default=0) class Meta: ordering = ['-created_at'] indexes = [ models.Index(fields=['user', 'status']), models.Index(fields=['task_type']), models.Index(fields=['priority']), ] def __str__(self): return f"{self.title} ({self.status})" def set_processing(self, celery_task_id=None): """Mark task as processing with the associated Celery task ID.""" self.status = 'processing' self.started_at = timezone.now() if celery_task_id: self.celery_task_id = celery_task_id self.save(update_fields=['status', 'started_at', 'celery_task_id', 'updated_at']) def set_completed(self, result_data, token_usage=None): """Mark task as completed with results and token usage.""" self.status = 'completed' self.completed_at = timezone.now() self.result_data = result_data if token_usage: self.prompt_tokens = token_usage.get('prompt_tokens', 0) self.completion_tokens = token_usage.get('completion_tokens', 0) self.total_tokens = token_usage.get('total_tokens', 0) # Calculate estimated cost prompt_cost = self.prompt_tokens * 0.0000015 # $0.0015 per 1000 tokens completion_cost = self.completion_tokens * 0.000002 # $0.002 per 1000 tokens self.estimated_cost = prompt_cost + completion_cost self.save() def set_failed(self, error_message): """Mark task as failed with error message.""" self.status = 'failed' self.error_message = error_message self.completed_at = timezone.now() self.save(update_fields=['status', 'error_message', 'completed_at', 'updated_at']) class KnowledgeItem(models.Model): """Model for storing knowledge items for agent context.""" id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) user = models.ForeignKey(User, on_delete=models.CASCADE, related_name='knowledge_items') title = models.CharField(max_length=255) content = models.TextField() # Metadata for filtering and organization source = models.CharField(max_length=255, blank=True) domain = models.CharField(max_length=100, blank=True) tags = models.JSONField(default=list, blank=True) # Vector storage tracking vector_id = models.CharField(max_length=255, blank=True, null=True) embedding_model = models.CharField(max_length=100, default="text-embedding-ada-002") last_updated = models.DateTimeField(auto_now=True) # Quality metrics relevance_score = models.FloatField(default=0.0) confidence = models.FloatField(default=1.0) created_at = models.DateTimeField(auto_now_add=True) class Meta: ordering = ['-created_at'] indexes = [ models.Index(fields=['user', 'domain']), models.Index(fields=['vector_id']), ] def __str__(self): return self.title class Conversation(models.Model): """Model for tracking conversations between users and agents.""" id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) user = models.ForeignKey(User, on_delete=models.CASCADE, related_name='conversations') title = models.CharField(max_length=255, blank=True) created_at = models.DateTimeField(auto_now_add=True) updated_at = models.DateTimeField(auto_now=True) is_active = models.BooleanField(default=True) # Metadata topic = models.CharField(max_length=100, blank=True) summary = models.TextField(blank=True) class Meta: ordering = ['-updated_at'] def __str__(self): return self.title or f"Conversation {self.id}" class Message(models.Model): """Model for storing messages within a conversation.""" ROLE_CHOICES = ( ('user', 'User'), ('assistant', 'Assistant'), ('system', 'System'), ) id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) conversation = models.ForeignKey(Conversation, on_delete=models.CASCADE, related_name='messages') role = models.CharField(max_length=10, choices=ROLE_CHOICES) content = models.TextField() created_at = models.DateTimeField(auto_now_add=True) # Metadata token_count = models.IntegerField(default=0) # References to agents if applicable agent_name = models.CharField(max_length=100, blank=True, null=True) class Meta: ordering = ['created_at'] def __str__(self): return f"{self.role} message in {self.conversation}" ``` Now let's implement the serializers for our API: ```python # agent_app/serializers.py from rest_framework import serializers from .models import AgentTask, KnowledgeItem, Conversation, Message from django.contrib.auth.models import User class UserSerializer(serializers.ModelSerializer): class Meta: model = User fields = ['id', 'username', 'email', 'first_name', 'last_name'] class AgentTaskSerializer(serializers.ModelSerializer): user = UserSerializer(read_only=True) duration = serializers.SerializerMethodField() class Meta: model = AgentTask fields = [ 'id', 'user', 'title', 'description', 'task_type', 'status', 'priority', 'input_data', 'result_data', 'error_message', 'created_at', 'updated_at', 'started_at', 'completed_at', 'prompt_tokens', 'completion_tokens', 'total_tokens', 'estimated_cost', 'duration' ] read_only_fields = [ 'id', 'status', 'result_data', 'error_message', 'created_at', 'updated_at', 'started_at', 'completed_at', 'prompt_tokens', 'completion_tokens', 'total_tokens', 'estimated_cost' ] def get_duration(self, obj): """Calculate task duration in seconds if available.""" if obj.started_at and obj.completed_at: return (obj.completed_at - obj.started_at).total_seconds() return None def create(self, validated_data): """Create a new agent task associated with the current user.""" user = self.context['request'].user validated_data['user'] = user return super().create(validated_data) class KnowledgeItemSerializer(serializers.ModelSerializer): user = UserSerializer(read_only=True) class Meta: model = KnowledgeItem fields = [ 'id', 'user', 'title', 'content', 'source', 'domain', 'tags', 'vector_id', 'embedding_model', 'last_updated', 'relevance_score', 'confidence', 'created_at' ] read_only_fields = [ 'id', 'vector_id', 'embedding_model', 'last_updated', 'created_at' ] def create(self, validated_data): """Create a new knowledge item associated with the current user.""" user = self.context['request'].user validated_data['user'] = user return super().create(validated_data) class MessageSerializer(serializers.ModelSerializer): class Meta: model = Message fields = [ 'id', 'conversation', 'role', 'content', 'created_at', 'token_count', 'agent_name' ] read_only_fields = ['id', 'created_at', 'token_count'] class ConversationSerializer(serializers.ModelSerializer): user = UserSerializer(read_only=True) messages = MessageSerializer(many=True, read_only=True) message_count = serializers.SerializerMethodField() class Meta: model = Conversation fields = [ 'id', 'user', 'title', 'created_at', 'updated_at', 'is_active', 'topic', 'summary', 'messages', 'message_count' ] read_only_fields = ['id', 'created_at', 'updated_at'] def get_message_count(self, obj): return obj.messages.count() def create(self, validated_data): """Create a new conversation associated with the current user.""" user = self.context['request'].user validated_data['user'] = user return super().create(validated_data) class ConversationMessageSerializer(serializers.Serializer): """Serializer for adding a message to a conversation.""" content = serializers.CharField(required=True) role = serializers.ChoiceField(choices=['user', 'system'], default='user') ``` Next, let's implement the agent utilities: ```python # agent_app/utils/agent_factory.py import os import autogen import logging from django.conf import settings from .pinecone_client import PineconeClient from .prompt_templates import get_system_prompt logger = logging.getLogger('agent_app') class AgentFactory: """Factory class for creating different types of agents.""" def __init__(self, user_id=None): self.user_id = user_id self.pinecone_client = PineconeClient() def get_config_list(self, model="gpt-4"): """Get LLM configuration list.""" return [ { "model": model, "api_key": settings.OPENAI_API_KEY } ] def create_agent(self, agent_type, context=None): """Create an agent based on the specified type.""" if agent_type == "researcher": return self._create_researcher_agent(context) elif agent_type == "analyst": return self._create_analyst_agent(context) elif agent_type == "document_processor": return self._create_document_processor_agent(context) elif agent_type == "conversation": return self._create_conversation_agent(context) else: # Default to a generic assistant agent return self._create_generic_agent(context) def create_agent_team(self, task_data): """Create a team of agents for complex tasks.""" task_type = task_data.get('task_type') if task_type == 'research': return self._create_research_team(task_data) elif task_type == 'analysis': return self._create_analysis_team(task_data) else: # Default team configuration return self._create_default_team(task_data) def _create_researcher_agent(self, context=None): """Create a researcher agent specialized in information gathering.""" system_message = get_system_prompt("researcher") if context: domain = context.get('domain', '') if domain: system_message += f"\nYou are specialized in researching {domain}." # Define research-specific functions functions = [ { "name": "search_knowledge_base", "description": "Search the knowledge base for relevant information", "parameters": { "type": "object", "properties": { "query": {"type": "string", "description": "The search query"}, "domain": {"type": "string", "description": "Optional domain to filter results"}, "limit": {"type": "integer", "description": "Maximum number of results to return"} }, "required": ["query"] } }, { "name": "summarize_sources", "description": "Summarize information from multiple sources", "parameters": { "type": "object", "properties": { "sources": {"type": "array", "items": {"type": "string"}, "description": "List of source texts to summarize"}, "max_length": {"type": "integer", "description": "Maximum length of the summary"} }, "required": ["sources"] } } ] # Create the agent with appropriate configuration agent = autogen.AssistantAgent( name="ResearchAgent", system_message=system_message, llm_config={ "config_list": self.get_config_list(), "functions": functions, "temperature": 0.5, "timeout": 600, # 10 minutes timeout for research tasks "cache_seed": None # No caching for research tasks } ) return agent def _create_analyst_agent(self, context=None): """Create an analyst agent specialized in data interpretation.""" system_message = get_system_prompt("analyst") if context: data_type = context.get('data_type', '') if data_type: system_message += f"\nYou are specialized in analyzing {data_type} data." # Define analysis-specific functions functions = [ { "name": "analyze_data", "description": "Analyze data and provide insights", "parameters": { "type": "object", "properties": { "data": {"type": "object", "description": "The data to analyze"}, "metrics": {"type": "array", "items": {"type": "string"}, "description": "Metrics to calculate"}, "visualization": {"type": "boolean", "description": "Whether to suggest visualizations"} }, "required": ["data"] } }, { "name": "generate_report", "description": "Generate a structured report from analysis", "parameters": { "type": "object", "properties": { "analysis_results": {"type": "object", "description": "Results from data analysis"}, "format": {"type": "string", "description": "Report format (e.g., executive, technical)"}, "sections": {"type": "array", "items": {"type": "string"}, "description": "Sections to include in the report"} }, "required": ["analysis_results"] } } ] # Create the agent with appropriate configuration agent = autogen.AssistantAgent( name="AnalystAgent", system_message=system_message, llm_config={ "config_list": self.get_config_list("gpt-4"), # Use GPT-4 for analysis "functions": functions, "temperature": 0.3, # Lower temperature for more precise analysis "timeout": 900, # 15 minutes timeout for analysis tasks "cache_seed": None # No caching for analysis tasks } ) return agent def _create_document_processor_agent(self, context=None): """Create an agent specialized in document processing.""" system_message = get_system_prompt("document_processor") if context: doc_type = context.get('document_type', '') if doc_type: system_message += f"\nYou are specialized in processing {doc_type} documents." # Define document processing functions functions = [ { "name": "extract_entities", "description": "Extract named entities from text", "parameters": { "type": "object", "properties": { "text": {"type": "string", "description": "The text to analyze"}, "entity_types": {"type": "array", "items": {"type": "string"}, "description": "Types of entities to extract"} }, "required": ["text"] } }, { "name": "classify_document", "description": "Classify document by type and content", "parameters": { "type": "object", "properties": { "text": {"type": "string", "description": "The document text"}, "categories": {"type": "array", "items": {"type": "string"}, "description": "Possible categories"} }, "required": ["text"] } }, { "name": "summarize_document", "description": "Create a concise summary of a document", "parameters": { "type": "object", "properties": { "text": {"type": "string", "description": "The document text"}, "max_length": {"type": "integer", "description": "Maximum summary length"} }, "required": ["text"] } } ] # Create the agent with appropriate configuration agent = autogen.AssistantAgent( name="DocumentAgent", system_message=system_message, llm_config={ "config_list": self.get_config_list(), "functions": functions, "temperature": 0.2, # Lower temperature for precision "timeout": 300, # 5 minutes timeout "cache_seed": 42 # Use caching for document tasks } ) return agent def _create_conversation_agent(self, context=None): """Create an agent for interactive conversations.""" system_message = get_system_prompt("conversation") if context: tone = context.get('tone', '') topic = context.get('topic', '') if tone: system_message += f"\nMaintain a {tone} tone in your responses." if topic: system_message += f"\nYou are specialized in discussing {topic}." # Create the agent with conversation-appropriate configuration agent = autogen.AssistantAgent( name="ConversationAgent", system_message=system_message, llm_config={ "config_list": self.get_config_list(), "temperature": 0.7, # Higher temperature for more creative conversations "timeout": 120, # 2 minutes timeout for conversational responses "cache_seed": None # No caching for unique conversations } ) return agent def _create_generic_agent(self, context=None): """Create a general-purpose assistant agent.""" system_message = get_system_prompt("generic") # Create a versatile, general-purpose agent agent = autogen.AssistantAgent( name="AssistantAgent", system_message=system_message, llm_config={ "config_list": self.get_config_list(), "temperature": 0.5, "timeout": 180, # 3 minutes timeout "cache_seed": None # No caching } ) return agent def _create_research_team(self, task_data): """Create a team of agents for research tasks.""" # Create specialized agents for the research team researcher = self._create_researcher_agent({"domain": task_data.get("domain")}) analyst = self._create_analyst_agent({"data_type": "research findings"}) writer = autogen.AssistantAgent( name="WriterAgent", system_message=get_system_prompt("writer"), llm_config={ "config_list": self.get_config_list(), "temperature": 0.7, "timeout": 300 } ) # Create a user proxy that will coordinate the team user_proxy = autogen.UserProxyAgent( name="ResearchCoordinator", human_input_mode="NEVER", max_consecutive_auto_reply=10, code_execution_config={"work_dir": "research_workspace"} ) # Create a group chat for the research team groupchat = autogen.GroupChat( agents=[user_proxy, researcher, analyst, writer], messages=[], max_round=15 ) manager = autogen.GroupChatManager( groupchat=groupchat, llm_config={ "config_list": self.get_config_list(), "temperature": 0.2 } ) return { "user_proxy": user_proxy, "manager": manager, "agents": [researcher, analyst, writer], "groupchat": groupchat } def _create_analysis_team(self, task_data): """Create a team of agents for data analysis tasks.""" # Create specialized agents for the analysis team data_processor = autogen.AssistantAgent( name="DataProcessor", system_message=get_system_prompt("data_processor"), llm_config={ "config_list": self.get_config_list(), "temperature": 0.2, "timeout": 300 } ) analyst = self._create_analyst_agent({"data_type": task_data.get("data_type", "general")}) visualization_expert = autogen.AssistantAgent( name="VisualizationExpert", system_message=get_system_prompt("visualization"), llm_config={ "config_list": self.get_config_list(), "temperature": 0.4, "timeout": 300 } ) report_writer = autogen.AssistantAgent( name="ReportWriter", system_message=get_system_prompt("report_writer"), llm_config={ "config_list": self.get_config_list(), "temperature": 0.6, "timeout": 300 } ) # Create a user proxy that will coordinate the team user_proxy = autogen.UserProxyAgent( name="AnalysisCoordinator", human_input_mode="NEVER", max_consecutive_auto_reply=10, code_execution_config={"work_dir": "analysis_workspace"} ) # Create a group chat for the analysis team groupchat = autogen.GroupChat( agents=[user_proxy, data_processor, analyst, visualization_expert, report_writer], messages=[], max_round=15 ) manager = autogen.GroupChatManager( groupchat=groupchat, llm_config={ "config_list": self.get_config_list(), "temperature": 0.2 } ) return { "user_proxy": user_proxy, "manager": manager, "agents": [data_processor, analyst, visualization_expert, report_writer], "groupchat": groupchat } def _create_default_team(self, task_data): """Create a default team of agents for general tasks.""" # Create generic agents for a default team planner = autogen.AssistantAgent( name="PlannerAgent", system_message=get_system_prompt("planner"), llm_config={ "config_list": self.get_config_list(), "temperature": 0.3, "timeout": 300 } ) executor = autogen.AssistantAgent( name="ExecutorAgent", system_message=get_system_prompt("executor"), llm_config={ "config_list": self.get_config_list(), "temperature": 0.5, "timeout": 300 } ) reviewer = autogen.AssistantAgent( name="ReviewerAgent", system_message=get_system_prompt("reviewer"), llm_config={ "config_list": self.get_config_list(), "temperature": 0.2, "timeout": 300 } ) # Create a user proxy that will coordinate the team user_proxy = autogen.UserProxyAgent( name="TaskCoordinator", human_input_mode="NEVER", max_consecutive_auto_reply=10, code_execution_config={"work_dir": "task_workspace"} ) # Create a group chat for the default team groupchat = autogen.GroupChat( agents=[user_proxy, planner, executor, reviewer], messages=[], max_round=10 ) manager = autogen.GroupChatManager( groupchat=groupchat, llm_config={ "config_list": self.get_config_list(), "temperature": 0.3 } ) return { "user_proxy": user_proxy, "manager": manager, "agents": [planner, executor, reviewer], "groupchat": groupchat } ``` ```python # agent_app/utils/pinecone_client.py import os import pinecone import openai import numpy as np import time import logging from django.conf import settings logger = logging.getLogger('agent_app') class PineconeClient: """Client for interacting with Pinecone vector database.""" def __init__(self): self.api_key = settings.PINECONE_API_KEY self.environment = settings.PINECONE_ENVIRONMENT self.index_name = settings.PINECONE_INDEX_NAME self.dimension = 1536 # OpenAI embedding dimension self.initialize_pinecone() def initialize_pinecone(self): """Initialize Pinecone and ensure index exists.""" try: pinecone.init(api_key=self.api_key, environment=self.environment) # Check if index exists, if not create it if self.index_name not in pinecone.list_indexes(): logger.info(f"Creating Pinecone index: {self.index_name}") pinecone.create_index( name=self.index_name, dimension=self.dimension, metric="cosine", shards=1 ) # Wait for index to be ready time.sleep(10) self.index = pinecone.Index(self.index_name) logger.info(f"Connected to Pinecone index: {self.index_name}") except Exception as e: logger.error(f"Error initializing Pinecone: {str(e)}") raise def get_embedding(self, text): """Get embedding for text using OpenAI API.""" try: response = openai.Embedding.create( input=text, model="text-embedding-ada-002" ) return response["data"][0]["embedding"] except Exception as e: logger.error(f"Error generating embedding: {str(e)}") raise def add_item(self, item_id, text, metadata=None): """Add an item to the vector database.""" try: # Get text embedding embedding = self.get_embedding(text) # Prepare metadata if metadata is None: metadata = {} # Add text to metadata for retrieval metadata["text"] = text # Upsert vector to Pinecone self.index.upsert( vectors=[(item_id, embedding, metadata)] ) logger.info(f"Added item {item_id} to Pinecone") return item_id except Exception as e: logger.error(f"Error adding item to Pinecone: {str(e)}") raise def delete_item(self, item_id): """Delete an item from the vector database.""" try: self.index.delete(ids=[item_id]) logger.info(f"Deleted item {item_id} from Pinecone") return True except Exception as e: logger.error(f"Error deleting item from Pinecone: {str(e)}") raise def search(self, query, filters=None, top_k=5): """Search for similar items in the vector database.""" try: # Get query embedding query_embedding = self.get_embedding(query) # Perform vector similarity search results = self.index.query( vector=query_embedding, top_k=top_k, include_metadata=True, filter=filters ) # Format results formatted_results = [] for match in results.matches: formatted_results.append({ "id": match.id, "score": match.score, "text": match.metadata.get("text", ""), "metadata": {k: v for k, v in match.metadata.items() if k != "text"} }) return formatted_results except Exception as e: logger.error(f"Error searching Pinecone: {str(e)}") raise def update_metadata(self, item_id, metadata): """Update metadata for an existing item.""" try: # Get current vector and metadata vector_data = self.index.fetch([item_id]) if item_id not in vector_data.vectors: logger.error(f"Item {item_id} not found in Pinecone") return False # Extract the vector and current metadata current_vector = vector_data.vectors[item_id].values current_metadata = vector_data.vectors[item_id].metadata # Update metadata updated_metadata = {**current_metadata, **metadata} # Upsert with updated metadata self.index.upsert( vectors=[(item_id, current_vector, updated_metadata)] ) logger.info(f"Updated metadata for item {item_id}") return True except Exception as e: logger.error(f"Error updating metadata in Pinecone: {str(e)}") raise def get_stats(self): """Get statistics about the index.""" try: stats = self.index.describe_index_stats() return { "namespaces": stats.get("namespaces", {}), "dimension": stats.get("dimension"), "total_vector_count": stats.get("total_vector_count") } except Exception as e: logger.error(f"Error getting Pinecone stats: {str(e)}") raise ``` ```python # agent_app/utils/prompt_templates.py """ Prompt templates for various agent types. """ def get_system_prompt(agent_type): """ Get the system prompt for a specific agent type. Args: agent_type (str): The type of agent Returns: str: The system prompt """ prompts = { "researcher": """You are an expert research agent who specializes in gathering comprehensive information on any topic. Your strength is in finding relevant information, evaluating sources, and synthesizing findings. Your capabilities: 1. Conduct thorough research on any given topic 2. Evaluate the credibility and relevance of sources 3. Identify key information and insights 4. Synthesize findings into clear, structured formats 5. Properly cite and attribute information to sources 6. Identify gaps in available information When researching: - Always begin by assessing what specific information is needed - Consider multiple perspectives and sources - Remain objective and unbiased - Distinguish between facts, expert opinions, and uncertain claims - Structure your findings in a logical manner - Note limitations in available information Your goal is to provide the most accurate, comprehensive, and well-organized research possible.""", "analyst": """You are an expert data analyst agent who excels at interpreting and deriving insights from complex data. Your strength is in statistical analysis, pattern recognition, and communicating findings clearly. Your capabilities: 1. Analyze numerical and categorical data 2. Identify trends, patterns, and anomalies 3. Apply appropriate statistical methods to datasets 4. Generate actionable insights from analysis 5. Create clear interpretations of analytical results 6. Make data-driven recommendations When analyzing: - First understand the context and objectives of the analysis - Consider what methods are most appropriate for the data type - Identify key metrics and indicators that address the objectives - Look for correlations, trends, and outliers - Consider statistical significance and confidence levels - Communicate findings in clear, non-technical terms when needed - Always disclose limitations and uncertainty in your analysis Your goal is to transform raw data into valuable insights and recommendations.""", "document_processor": """You are an expert document processing agent who specializes in analyzing, summarizing, and extracting information from documents. Your strength is in understanding document structure, identifying key information, and producing accurate analyses. Your capabilities: 1. Extract key information from documents 2. Summarize document content at different levels of detail 3. Identify main themes, arguments, and conclusions 4. Recognize document structure and organization 5. Classify documents by type, purpose, and content 6. Extract named entities and relationships When processing documents: - First identify the document type and purpose - Consider the document's structure and organization - Identify the most important sections and content - Extract key information, claims, and evidence - Recognize the tone, style, and intended audience - Maintain the original meaning and context - Be precise in your extraction and summarization Your goal is to accurately process documents and make their information accessible and useful.""", "conversation": """You are an expert conversation agent who excels at engaging in natural, helpful dialogue. Your strength is in understanding user needs, providing relevant information, and maintaining engaging interactions. Your capabilities: 1. Engage in natural, flowing conversation 2. Understand explicit and implicit user questions 3. Provide clear, concise, and accurate information 4. Adapt your tone and style to match the user 5. Ask clarifying questions when needed 6. Remember context throughout a conversation During conversations: - Listen carefully to understand the user's full intent - Provide helpful, relevant responses - Be concise but complete in your answers - Maintain a consistent tone and personality - Ask questions when needed for clarification - Acknowledge when you don't know something - Structure complex information clearly Your goal is to provide an engaging, helpful, and informative conversation experience.""", "generic": """You are a versatile assistant agent capable of handling a wide range of tasks. You can provide information, answer questions, offer suggestions, and assist with various needs. Your capabilities: 1. Answer questions across diverse domains 2. Provide explanations and clarifications 3. Offer suggestions and recommendations 4. Help with planning and organization 5. Assist with creative tasks 6. Engage in thoughtful discussion When assisting: - Understand the core request or question - Provide clear, accurate, and helpful responses - Consider the context and intent behind questions - Structure your responses logically - Be honest about your limitations - Maintain a helpful and supportive tone Your goal is to be a versatile, reliable, and helpful assistant.""", "writer": """You are an expert writing agent who specializes in creating clear, engaging, and well-structured content. Your strength is in adapting your writing style to different purposes and audiences. Your capabilities: 1. Create clear and engaging content in various formats 2. Adapt writing style to different audiences and purposes 3. Structure content logically and coherently 4. Edit and refine existing content 5. Ensure grammar, spelling, and stylistic consistency 6. Generate creative and original content When writing: - Consider the purpose, audience, and context - Organize information with a clear structure - Use appropriate tone, style, and vocabulary - Create engaging introductions and conclusions - Use transitions to guide readers through the content - Revise for clarity, conciseness, and impact Your goal is to produce high-quality written content that effectively communicates ideas to the intended audience.""", "data_processor": """You are an expert data processing agent who specializes in preparing, cleaning, and transforming data for analysis. Your strength is in handling raw data and making it ready for meaningful analysis. Your capabilities: 1. Clean and normalize messy datasets 2. Handle missing, duplicate, or inconsistent data 3. Transform data into appropriate formats 4. Merge and join datasets from multiple sources 5. Create derived features and variables 6. Identify and address data quality issues When processing data: - First assess the data structure and quality - Identify issues that need to be addressed - Apply appropriate cleaning and transformation methods - Document all changes made to the original data - Validate the processed data for accuracy - Prepare the data in a format suitable for analysis Your goal is to transform raw data into a clean, consistent, and analysis-ready format.""", "visualization": """You are an expert data visualization agent who specializes in creating effective visual representations of data. Your strength is in selecting and designing visualizations that clearly communicate insights. Your capabilities: 1. Select appropriate visualization types for different data 2. Design clear, informative visual representations 3. Highlight key patterns, trends, and relationships in data 4. Create accessible and intuitive visualizations 5. Adapt visualizations for different audiences 6. Combine multiple visualizations into dashboards When creating visualizations: - Consider the data type and relationships to visualize - Select the most appropriate chart or graph type - Focus on clearly communicating the main insights - Minimize clutter and maximize data-ink ratio - Use color, labels, and annotations effectively - Consider accessibility and interpretability - Provide clear titles, legends, and context Your goal is to create visualizations that effectively communicate data insights in an accessible and impactful way.""", "report_writer": """You are an expert report writing agent who specializes in creating comprehensive, structured reports that effectively communicate findings and insights. Your strength is in organizing information logically and presenting it clearly. Your capabilities: 1. Create well-structured reports for different purposes 2. Organize findings and insights logically 3. Present complex information clearly and concisely 4. Integrate data, analysis, and visualizations 5. Adapt content and style to different audiences 6. Highlight key findings and recommendations When writing reports: - Consider the purpose, audience, and required detail level - Create a logical structure with clear sections - Begin with an executive summary of key points - Present findings with supporting evidence - Use visuals to complement and enhance text - Maintain consistent formatting and style - Conclude with clear insights and recommendations Your goal is to produce comprehensive, clear reports that effectively communicate information to the intended audience.""", "planner": """You are an expert planning agent who specializes in breaking down complex tasks into organized, achievable steps. Your strength is in creating structured plans that lead to successful outcomes. Your capabilities: 1. Break down complex tasks into manageable steps 2. Identify dependencies between tasks 3. Estimate time and resources needed 4. Prioritize tasks based on importance and urgency 5. Identify potential risks and mitigation strategies 6. Adapt plans as circumstances change When planning: - First understand the overall goal and constraints - Identify all necessary tasks and subtasks - Determine logical sequence and dependencies - Allocate appropriate time and resources - Highlight critical path items and bottlenecks - Include checkpoints to assess progress - Anticipate potential obstacles and plan alternatives Your goal is to create clear, achievable plans that efficiently lead to successful outcomes.""", "executor": """You are an expert execution agent who specializes in implementing plans and completing tasks. Your strength is in taking action, solving problems, and delivering results. Your capabilities: 1. Implement plans and complete assigned tasks 2. Follow procedures and instructions precisely 3. Solve problems that arise during execution 4. Adapt to changing circumstances 5. Manage resources efficiently 6. Document actions and results When executing: - Review and understand the task requirements - Gather necessary resources and information - Follow established procedures and best practices - Address issues promptly as they arise - Document progress and completed work - Communicate status and any obstacles clearly - Verify that outcomes meet requirements Your goal is to effectively implement plans and deliver high-quality results.""", "reviewer": """You are an expert review agent who specializes in evaluating work and providing constructive feedback. Your strength is in assessing quality, identifying issues, and suggesting improvements. Your capabilities: 1. Evaluate work against established criteria and standards 2. Identify strengths and weaknesses 3. Detect errors, inconsistencies, and problems 4. Ensure compliance with requirements 5. Provide specific, actionable feedback 6. Suggest concrete improvements When reviewing: - First understand the requirements and context - Evaluate objectively against clear criteria - Be thorough and systematic in your assessment - Provide balanced feedback on strengths and weaknesses - Be specific about issues and why they matter - Suggest clear, actionable improvements - Maintain a constructive and helpful tone Your goal is to improve quality through thorough evaluation and constructive feedback.""", } return prompts.get(agent_type, prompts["generic"]) ``` Now let's implement the Celery tasks: ```python # agent_app/tasks.py import time import json import logging import traceback from django.utils import timezone from django.conf import settings from celery import shared_task from celery.exceptions import SoftTimeLimitExceeded from .models import AgentTask, KnowledgeItem, Conversation, Message from .utils.agent_factory import AgentFactory from .utils.pinecone_client import PineconeClient logger = logging.getLogger('agent_app') @shared_task(bind=True, soft_time_limit=1800, # 30 minute soft limit time_limit=1900, # ~32 minute hard limit acks_late=True, # Acknowledge task after execution retry_backoff=True, # Exponential backoff for retries max_retries=3) # Maximum retry attempts def process_agent_task(self, task_id): """ Process an agent task with the appropriate agent type. Args: task_id (str): UUID of the task to process Returns: dict: Result data """ try: # Get task from database try: task = AgentTask.objects.get(id=task_id) except AgentTask.DoesNotExist: logger.error(f"Task {task_id} not found") return {"error": f"Task {task_id} not found"} # Update task status task.set_processing(self.request.id) logger.info(f"Processing task {task_id} of type {task.task_type}") # Initialize agent factory agent_factory = AgentFactory(user_id=task.user.id) # Process based on task type if task.task_type in ['research', 'analysis']: # Create a team of agents for complex tasks team = agent_factory.create_agent_team(task.input_data) result = process_team_task(team, task.input_data) else: # Create an individual agent agent_type = task.input_data.get('agent_type', 'generic') context = task.input_data.get('context', {}) agent = agent_factory.create_agent(agent_type, context) # Process the task result = process_individual_task(agent, task.input_data) # Extract token usage token_usage = result.get('token_usage', {}) # Update task as completed task.set_completed(result, token_usage) logger.info(f"Task {task_id} completed successfully") return result except SoftTimeLimitExceeded: # Handle timeout logger.error(f"Task {task_id} exceeded time limit") try: task = AgentTask.objects.get(id=task_id) task.set_failed("Task exceeded time limit") except Exception as e: logger.error(f"Error updating task: {str(e)}") return {"error": "Task exceeded time limit"} except Exception as e: # Handle other exceptions error_msg = str(e) stack_trace = traceback.format_exc() logger.error(f"Error processing task {task_id}: {error_msg}\n{stack_trace}") try: task = AgentTask.objects.get(id=task_id) task.set_failed(f"{error_msg}\n{stack_trace}") except Exception as e2: logger.error(f"Error updating task status: {str(e2)}") # Retry for certain exceptions if "Rate limit" in error_msg or "timeout" in error_msg.lower(): raise self.retry(exc=e, countdown=60) return {"error": error_msg, "stack_trace": stack_trace} def process_team_task(team, task_data): """ Process a task using a team of agents. Args: team (dict): The team configuration with agents task_data (dict): Task data Returns: dict: Result data """ start_time = time.time() # Extract team components user_proxy = team["user_proxy"] manager = team["manager"] # Prepare the task message task_description = task_data.get('description', '') task_details = task_data.get('details', {}) message = f"Task: {task_description}\n\n" if task_details: message += "Details:\n" for key, value in task_details.items(): message += f"- {key}: {value}\n" # Start the group conversation user_proxy.initiate_chat( manager, message=message ) # Extract results chat_history = user_proxy.chat_history result_content = None # Find the final result from the chat history for msg in reversed(chat_history): if msg.get("role") == "assistant" and len(msg.get("content", "")) > 100: result_content = msg.get("content") break # If no clear result is found, summarize the entire conversation if not result_content: result_content = "No clear result was produced. Here's the conversation summary:\n\n" for msg in chat_history: if msg.get("role") in ["assistant", "user"]: result_content += f"{msg.get('role').upper()}: {msg.get('content', '')[:200]}...\n\n" # Calculate token usage (estimated) total_input_tokens = sum(len(msg.get("content", "").split()) * 1.3 for msg in chat_history if msg.get("role") == "user") total_output_tokens = sum(len(msg.get("content", "").split()) * 1.3 for msg in chat_history if msg.get("role") == "assistant") duration = time.time() - start_time return { "result": result_content, "chat_history": chat_history, "duration_seconds": duration, "team_composition": [agent.name for agent in team["agents"]], "token_usage": { "prompt_tokens": int(total_input_tokens), "completion_tokens": int(total_output_tokens), "total_tokens": int(total_input_tokens + total_output_tokens) } } def process_individual_task(agent, task_data): """ Process a task using an individual agent. Args: agent (AssistantAgent): The agent to process the task task_data (dict): Task data Returns: dict: Result data """ start_time = time.time() # Create user proxy agent user_proxy = autogen.UserProxyAgent( name="TaskUser", human_input_mode="NEVER", max_consecutive_auto_reply=0 ) # Prepare the task message query = task_data.get('query', '') context = task_data.get('context', {}) message = query if context: message += "\n\nContext:\n" for key, value in context.items(): message += f"- {key}: {value}\n" # Start conversation with agent user_proxy.initiate_chat( agent, message=message ) # Get the last message from the agent as the result chat_history = user_proxy.chat_history result_content = None for msg in reversed(chat_history): if msg.get("role") == "assistant": result_content = msg.get("content") break # Calculate token usage (estimated) total_input_tokens = sum(len(msg.get("content", "").split()) * 1.3 for msg in chat_history if msg.get("role") == "user") total_output_tokens = sum(len(msg.get("content", "").split()) * 1.3 for msg in chat_history if msg.get("role") == "assistant") duration = time.time() - start_time return { "result": result_content, "chat_history": chat_history, "duration_seconds": duration, "agent_type": agent.name, "token_usage": { "prompt_tokens": int(total_input_tokens), "completion_tokens": int(total_output_tokens), "total_tokens": int(total_input_tokens + total_output_tokens) } } @shared_task(bind=True, soft_time_limit=300, # 5 minute soft limit acks_late=True, max_retries=2) def update_knowledge_base(self, knowledge_item_id): """ Update a knowledge item in the vector database. Args: knowledge_item_id (str): UUID of the knowledge item to update Returns: dict: Result status """ try: # Get knowledge item from database try: item = KnowledgeItem.objects.get(id=knowledge_item_id) except KnowledgeItem.DoesNotExist: logger.error(f"Knowledge item {knowledge_item_id} not found") return {"error": f"Knowledge item {knowledge_item_id} not found"} logger.info(f"Updating knowledge item {knowledge_item_id} in vector database") # Initialize Pinecone client pinecone_client = PineconeClient() # Prepare metadata metadata = { "user_id": str(item.user.id), "title": item.title, "source": item.source, "domain": item.domain, "tags": json.dumps(item.tags), "confidence": item.confidence, "created_at": item.created_at.isoformat() } # Update or create vector if item.vector_id: # Update existing vector success = pinecone_client.update_metadata(item.vector_id, metadata) if not success: # If update failed, create new vector vector_id = pinecone_client.add_item( str(item.id), item.content, metadata ) item.vector_id = vector_id item.save(update_fields=['vector_id', 'last_updated']) else: # Create new vector vector_id = pinecone_client.add_item( str(item.id), item.content, metadata ) item.vector_id = vector_id item.save(update_fields=['vector_id', 'last_updated']) logger.info(f"Knowledge item {knowledge_item_id} updated successfully") return { "status": "success", "message": f"Knowledge item {knowledge_item_id} updated", "vector_id": item.vector_id } except Exception as e: error_msg = str(e) stack_trace = traceback.format_exc() logger.error(f"Error updating knowledge item {knowledge_item_id}: {error_msg}\n{stack_trace}") # Retry for certain exceptions if "Rate limit" in error_msg or "timeout" in error_msg.lower(): raise self.retry(exc=e, countdown=30) return {"error": error_msg, "stack_trace": stack_trace} @shared_task(bind=True, soft_time_limit=600, # 10 minute soft limit acks_late=True, max_retries=2) def analyze_document(self, task_id): """ Analyze a document using a document processor agent. Args: task_id (str): UUID of the task to process Returns: dict: Analysis results """ try: # Get task from database try: task = AgentTask.objects.get(id=task_id) except AgentTask.DoesNotExist: logger.error(f"Task {task_id} not found") return {"error": f"Task {task_id} not found"} # Update task status task.set_processing(self.request.id) logger.info(f"Analyzing document for task {task_id}") # Initialize agent factory agent_factory = AgentFactory(user_id=task.user.id) # Create document processor agent context = { "document_type": task.input_data.get("document_type", "general") } agent = agent_factory.create_agent("document_processor", context) # Extract document text document_text = task.input_data.get("document", "") if not document_text: raise ValueError("No document text provided") # Create user proxy agent user_proxy = autogen.UserProxyAgent( name="DocumentUser", human_input_mode="NEVER", max_consecutive_auto_reply=0 ) # Prepare analysis instructions analysis_type = task.input_data.get("analysis_type", "general") instructions = f"""Analyze the following document. Document type: {context.get('document_type', 'general')} Analysis type: {analysis_type} Please provide: 1. A concise summary of the document 2. Key entities and topics mentioned 3. Main points or arguments presented 4. Any notable insights or implications Document text: {document_text[:8000]} # Limit to first 8000 chars to avoid token limits """ if len(document_text) > 8000: instructions += "\n\n[Note: Document has been truncated due to length. Analysis based on first portion only.]" # Start conversation with agent user_proxy.initiate_chat( agent, message=instructions ) # Get the analysis result chat_history = user_proxy.chat_history analysis_result = None for msg in reversed(chat_history): if msg.get("role") == "assistant": analysis_result = msg.get("content") break # Calculate token usage (estimated) total_input_tokens = sum(len(msg.get("content", "").split()) * 1.3 for msg in chat_history if msg.get("role") == "user") total_output_tokens = sum(len(msg.get("content", "").split()) * 1.3 for msg in chat_history if msg.get("role") == "assistant") # Prepare result data result = { "analysis": analysis_result, "document_type": context.get("document_type"), "analysis_type": analysis_type, "character_count": len(document_text), "token_usage": { "prompt_tokens": int(total_input_tokens), "completion_tokens": int(total_output_tokens), "total_tokens": int(total_input_tokens + total_output_tokens) } } # Update task as completed task.set_completed(result, result["token_usage"]) logger.info(f"Document analysis for task {task_id} completed") return result except Exception as e: error_msg = str(e) stack_trace = traceback.format_exc() logger.error(f"Error analyzing document for task {task_id}: {error_msg}\n{stack_trace}") try: task = AgentTask.objects.get(id=task_id) task.set_failed(f"{error_msg}\n{stack_trace}") except Exception as e2: logger.error(f"Error updating task status: {str(e2)}") return {"error": error_msg, "stack_trace": stack_trace} @shared_task(bind=True) def periodic_agent_check(self): """ Periodic task to check for stalled agent tasks. """ try: # Find tasks that have been processing for too long (1 hour) one_hour_ago = timezone.now() - timezone.timedelta(hours=1) stalled_tasks = AgentTask.objects.filter( status='processing', started_at__lt=one_hour_ago ) count = stalled_tasks.count() logger.info(f"Found {count} stalled tasks") # Mark tasks as failed for task in stalled_tasks: task.set_failed("Task stalled - processing timeout exceeded") logger.warning(f"Marked task {task.id} as failed due to processing timeout") return {"status": "success", "stalled_tasks_count": count} except Exception as e: logger.error(f"Error in periodic agent check: {str(e)}") return {"status": "error", "message": str(e)} @shared_task(bind=True) def process_conversation(self, conversation_id, message_id): """ Process a new message in a conversation. Args: conversation_id (str): UUID of the conversation message_id (str): UUID of the message to process Returns: dict: Response data """ try: # Get conversation and message from database try: conversation = Conversation.objects.get(id=conversation_id) message = Message.objects.get(id=message_id, conversation=conversation) except (Conversation.DoesNotExist, Message.DoesNotExist): logger.error(f"Conversation {conversation_id} or message {message_id} not found") return {"error": "Conversation or message not found"} logger.info(f"Processing message {message_id} in conversation {conversation_id}") # Initialize agent factory agent_factory = AgentFactory(user_id=conversation.user.id) # Create conversation agent context = { "topic": conversation.topic, "tone": "conversational" } agent = agent_factory.create_agent("conversation", context) # Get conversation history (last 10 messages) history = conversation.messages.order_by('created_at')[:10] # Prepare conversation context conversation_context = "" for hist_msg in history: if hist_msg.id != message.id: # Skip the current message conversation_context += f"{hist_msg.role.upper()}: {hist_msg.content}\n\n" # Create user proxy agent user_proxy = autogen.UserProxyAgent( name="ConversationUser", human_input_mode="NEVER", max_consecutive_auto_reply=0 ) # Prepare the message with context prompt = f"""This is part of an ongoing conversation. Please respond to the latest message. Previous conversation: {conversation_context} Current message: USER: {message.content} Please respond in a helpful, conversational manner.""" # Start conversation with agent user_proxy.initiate_chat( agent, message=prompt ) # Get the response from the agent chat_history = user_proxy.chat_history response_content = None for msg in reversed(chat_history): if msg.get("role") == "assistant": response_content = msg.get("content") break if response_content: # Create response message response = Message.objects.create( conversation=conversation, role="assistant", content=response_content, agent_name=agent.name, token_count=len(response_content.split()) * 1.3 # Approximate token count ) # Update conversation conversation.updated_at = timezone.now() conversation.save(update_fields=['updated_at']) logger.info(f"Created response message {response.id} in conversation {conversation_id}") # Calculate token usage (estimated) total_input_tokens = sum(len(msg.get("content", "").split()) * 1.3 for msg in chat_history if msg.get("role") == "user") total_output_tokens = sum(len(msg.get("content", "").split()) * 1.3 for msg in chat_history if msg.get("role") == "assistant") return { "status": "success", "response_id": str(response.id), "response_content": response_content, "token_usage": { "prompt_tokens": int(total_input_tokens), "completion_tokens": int(total_output_tokens), "total_tokens": int(total_input_tokens + total_output_tokens) } } else: logger.error(f"No response generated for message {message_id}") return {"error": "No response generated"} except Exception as e: error_msg = str(e) stack_trace = traceback.format_exc() logger.error(f"Error processing conversation message: {error_msg}\n{stack_trace}") return {"error": error_msg, "stack_trace": stack_trace} ``` Now, let's implement the API views: ```python # agent_app/views.py import uuid import logging from django.shortcuts import get_object_or_404 from django.utils import timezone from django.db.models import Count, Sum, Max, F, ExpressionWrapper, fields from django.db.models.functions import TruncDay from rest_framework import viewsets, status, permissions from rest_framework.decorators import action from rest_framework.response import Response from rest_framework.views import APIView from .models import AgentTask, KnowledgeItem, Conversation, Message from .serializers import ( AgentTaskSerializer, KnowledgeItemSerializer, ConversationSerializer, MessageSerializer, ConversationMessageSerializer ) from .tasks import ( process_agent_task, update_knowledge_base, analyze_document, process_conversation ) from .utils.pinecone_client import PineconeClient logger = logging.getLogger('agent_app') class AgentTaskViewSet(viewsets.ModelViewSet): """ViewSet for managing agent tasks.""" serializer_class = AgentTaskSerializer permission_classes = [permissions.IsAuthenticated] def get_queryset(self): """Return tasks for the current user.""" return AgentTask.objects.filter(user=self.request.user).order_by('-created_at') def create(self, request, *args, **kwargs): """Create a new task and queue it for processing.""" serializer = self.get_serializer(data=request.data) serializer.is_valid(raise_exception=True) task = serializer.save() # Queue the task for processing based on type if task.task_type == 'document': # Use document-specific task result = analyze_document.delay(str(task.id)) else: # Use general task processing result = process_agent_task.delay(str(task.id)) # Update task with celery task ID task.celery_task_id = result.id task.save(update_fields=['celery_task_id']) headers = self.get_success_headers(serializer.data) return Response(serializer.data, status=status.HTTP_201_CREATED, headers=headers) @action(detail=True, methods=['post']) def cancel(self, request, pk=None): """Cancel a running task.""" task = self.get_object() if task.status not in ['pending', 'processing']: return Response( {"detail": "Only pending or processing tasks can be canceled."}, status=status.HTTP_400_BAD_REQUEST ) # Update task status task.status = 'canceled' task.completed_at = timezone.now() task.save(update_fields=['status', 'completed_at', 'updated_at']) # Could also revoke the Celery task here return Response({"detail": "Task canceled successfully."}) @action(detail=False, methods=['get']) def stats(self, request): """Get statistics about tasks.""" queryset = self.get_queryset() # Basic counts total_tasks = queryset.count() completed_tasks = queryset.filter(status='completed').count() failed_tasks = queryset.filter(status='failed').count() # Token usage statistics token_stats = queryset.filter(status='completed').aggregate( total_prompt_tokens=Sum('prompt_tokens'), total_completion_tokens=Sum('completion_tokens'), total_tokens=Sum('total_tokens'), avg_tokens_per_task=ExpressionWrapper( Sum('total_tokens') * 1.0 / Count('id'), output_field=fields.FloatField() ), total_cost=Sum('estimated_cost') ) # Task type distribution type_distribution = ( queryset .values('task_type') .annotate(count=Count('id')) .order_by('-count') ) # Task creation over time (last 30 days) thirty_days_ago = timezone.now() - timezone.timedelta(days=30) time_series = ( queryset .filter(created_at__gte=thirty_days_ago) .annotate(day=TruncDay('created_at')) .values('day') .annotate(count=Count('id')) .order_by('day') ) # Average processing time time_stats = queryset.filter( status='completed', started_at__isnull=False, completed_at__isnull=False ).aggregate( avg_processing_time=ExpressionWrapper( (F('completed_at') - F('started_at')) / 1000000, # Convert microseconds to seconds output_field=fields.FloatField() ) ) return Response({ "total_tasks": total_tasks, "completed_tasks": completed_tasks, "failed_tasks": failed_tasks, "token_stats": token_stats, "type_distribution": type_distribution, "time_series": time_series, "time_stats": time_stats }) class KnowledgeItemViewSet(viewsets.ModelViewSet): """ViewSet for managing knowledge items.""" serializer_class = KnowledgeItemSerializer permission_classes = [permissions.IsAuthenticated] def get_queryset(self): """Return knowledge items for the current user.""" return KnowledgeItem.objects.filter(user=self.request.user).order_by('-created_at') def create(self, request, *args, **kwargs): """Create a new knowledge item and index it.""" serializer = self.get_serializer(data=request.data) serializer.is_valid(raise_exception=True) item = serializer.save() # Queue indexing task update_knowledge_base.delay(str(item.id)) headers = self.get_success_headers(serializer.data) return Response(serializer.data, status=status.HTTP_201_CREATED, headers=headers) def update(self, request, *args, **kwargs): """Update a knowledge item and re-index it.""" partial = kwargs.pop('partial', False) instance = self.get_object() serializer = self.get_serializer(instance, data=request.data, partial=partial) serializer.is_valid(raise_exception=True) item = serializer.save() # Queue indexing task update_knowledge_base.delay(str(item.id)) return Response(serializer.data) def destroy(self, request, *args, **kwargs): """Delete a knowledge item and remove from index.""" instance = self.get_object() # If it has a vector ID, remove from Pinecone if instance.vector_id: try: pinecone_client = PineconeClient() pinecone_client.delete_item(instance.vector_id) except Exception as e: logger.error(f"Error removing item from Pinecone: {str(e)}") self.perform_destroy(instance) return Response(status=status.HTTP_204_NO_CONTENT) @action(detail=False, methods=['post']) def search(self, request): """Search knowledge items by semantic similarity.""" query = request.data.get('query') if not query: return Response( {"detail": "Query is required."}, status=status.HTTP_400_BAD_REQUEST ) filters = request.data.get('filters', {}) top_k = int(request.data.get('limit', 5)) # Add user filter filters["user_id"] = str(request.user.id) try: pinecone_client = PineconeClient() results = pinecone_client.search(query, filters, top_k) # Get item IDs from results item_ids = [result['id'] for result in results] # Fetch full items from database items = KnowledgeItem.objects.filter(id__in=item_ids) item_dict = {str(item.id): item for item in items} # Combine database items with search results enriched_results = [] for result in results: item = item_dict.get(result['id']) if item: enriched_results.append({ "id": str(item.id), "title": item.title, "content": item.content, "domain": item.domain, "tags": item.tags, "source": item.source, "created_at": item.created_at.isoformat(), "relevance_score": result['score'] }) return Response({ "results": enriched_results, "count": len(enriched_results), "query": query }) except Exception as e: logger.error(f"Error searching knowledge: {str(e)}") return Response( {"detail": f"Search error: {str(e)}"}, status=status.HTTP_500_INTERNAL_SERVER_ERROR ) class ConversationViewSet(viewsets.ModelViewSet): """ViewSet for managing conversations.""" serializer_class = ConversationSerializer permission_classes = [permissions.IsAuthenticated] def get_queryset(self): """Return conversations for the current user.""" queryset = Conversation.objects.filter( user=self.request.user ).order_by('-updated_at') # Filter by active status if specified is_active = self.request.query_params.get('is_active') if is_active is not None: is_active = is_active.lower() == 'true' queryset = queryset.filter(is_active=is_active) # Filter by topic if specified topic = self.request.query_params.get('topic') if topic: queryset = queryset.filter(topic=topic) return queryset @action(detail=True, methods=['post']) def add_message(self, request, pk=None): """Add a new message to the conversation and get a response.""" conversation = self.get_object() serializer = ConversationMessageSerializer(data=request.data) serializer.is_valid(raise_exception=True) # Create the message message = Message.objects.create( conversation=conversation, role=serializer.validated_data.get('role', 'user'), content=serializer.validated_data['content'], token_count=len(serializer.validated_data['content'].split()) * 1.3 # Approximate token count ) # Update conversation timestamp conversation.updated_at = timezone.now() conversation.save(update_fields=['updated_at']) # Queue task to process the message if it's from the user if message.role == 'user': process_conversation.delay(str(conversation.id), str(message.id)) # Return the created message message_serializer = MessageSerializer(message) return Response({ "message": message_serializer.data, "status": "processing" if message.role == 'user' else "completed" }) @action(detail=True, methods=['get']) def messages(self, request, pk=None): """Get messages for a conversation with pagination.""" conversation = self.get_object() # Get messages with pagination messages = conversation.messages.order_by('created_at') # Get page parameters page = self.paginate_queryset(messages) if page is not None: serializer = MessageSerializer(page, many=True) return self.get_paginated_response(serializer.data) serializer = MessageSerializer(messages, many=True) return Response(serializer.data) @action(detail=True, methods=['post']) def archive(self, request, pk=None): """Archive a conversation (mark as inactive).""" conversation = self.get_object() conversation.is_active = False conversation.save(update_fields=['is_active', 'updated_at']) return Response({"status": "success", "detail": "Conversation archived"}) class MessageViewSet(viewsets.ReadOnlyModelViewSet): """ViewSet for reading messages (no direct creation/update).""" serializer_class = MessageSerializer permission_classes = [permissions.IsAuthenticated] def get_queryset(self): """Return messages that belong to the current user's conversations.""" return Message.objects.filter( conversation__user=self.request.user ).order_by('created_at') def list(self, request, *args, **kwargs): """Messages must be accessed through a conversation.""" return Response( {"detail": "Messages must be accessed through a specific conversation."}, status=status.HTTP_400_BAD_REQUEST ) class VectorDatabaseStatsView(APIView): """View for getting vector database statistics.""" permission_classes = [permissions.IsAuthenticated] def get(self, request, format=None): """Get statistics about the vector database.""" try: pinecone_client = PineconeClient() stats = pinecone_client.get_stats() # Get count of user's knowledge items user_items_count = KnowledgeItem.objects.filter( user=request.user, vector_id__isnull=False ).count() return Response({ "user_items_count": user_items_count, "total_vectors": stats.get("total_vector_count", 0), "dimension": stats.get("dimension", 0), "namespaces": stats.get("namespaces", {}) }) except Exception as e: logger.error(f"Error getting vector database stats: {str(e)}") return Response( {"detail": f"Error: {str(e)}"}, status=status.HTTP_500_INTERNAL_SERVER_ERROR ) ``` ```python # agent_app/urls.py from django.urls import path, include from rest_framework.routers import DefaultRouter from . import views router = DefaultRouter() router.register(r'tasks', views.AgentTaskViewSet, basename='task') router.register(r'knowledge', views.KnowledgeItemViewSet, basename='knowledge') router.register(r'conversations', views.ConversationViewSet, basename='conversation') router.register(r'messages', views.MessageViewSet, basename='message') urlpatterns = [ path('', include(router.urls)), path('vector-stats/', views.VectorDatabaseStatsView.as_view(), name='vector-stats'), ] ``` Finally, let's create the admin interface: ```python # agent_app/admin.py from django.contrib import admin from .models import AgentTask, KnowledgeItem, Conversation, Message @admin.register(AgentTask) class AgentTaskAdmin(admin.ModelAdmin): list_display = ('id', 'title', 'user', 'task_type', 'status', 'priority', 'created_at') list_filter = ('status', 'task_type', 'priority') search_fields = ('title', 'description', 'user__username') readonly_fields = ('created_at', 'updated_at', 'started_at', 'completed_at', 'celery_task_id') date_hierarchy = 'created_at' @admin.register(KnowledgeItem) class KnowledgeItemAdmin(admin.ModelAdmin): list_display = ('id', 'title', 'user', 'domain', 'created_at') list_filter = ('domain',) search_fields = ('title', 'content', 'user__username') readonly_fields = ('created_at', 'vector_id', 'last_updated') @admin.register(Conversation) class ConversationAdmin(admin.ModelAdmin): list_display = ('id', 'title', 'user', 'topic', 'is_active', 'created_at', 'updated_at') list_filter = ('is_active', 'topic') search_fields = ('title', 'user__username') readonly_fields = ('created_at', 'updated_at') @admin.register(Message) class MessageAdmin(admin.ModelAdmin): list_display = ('id', 'conversation', 'role', 'token_count', 'created_at') list_filter = ('role',) search_fields = ('content', 'conversation__title') readonly_fields = ('created_at',) ``` **Key Advantages:** 1. **Task Orchestration**: Celery provides robust task queuing and scheduling for asynchronous agent operations 2. **Persistence**: Django models provide structured data storage with proper relationships 3. **Scalability**: Tasks can be distributed across multiple workers 4. **Semantic Search**: Pinecone enables fast vector search for knowledge retrieval 5. **Easy API Development**: Django REST Framework provides a clean API interface **Production Considerations:** 1. **Task Queue Monitoring**: For production systems, you would want to add proper monitoring for Celery tasks: ```python # monitoring/celery_monitoring.py from flower.utils.broker import Broker from celery.events.state import State import logging import json import time import requests logger = logging.getLogger('agent_app.monitoring') class CeleryMonitor: """Monitor for Celery tasks and workers.""" def __init__(self, broker_url): self.broker_url = broker_url self.state = State() self.broker = Broker(broker_url) def get_worker_stats(self): """Get statistics about active Celery workers.""" try: stats = self.broker.info() return { "active_workers": stats.get("active_workers", 0), "worker_heartbeats": stats.get("worker_heartbeats", {}), "processed_tasks": stats.get("processed", 0), "failed_tasks": stats.get("failed", 0), "broker_queue_sizes": stats.get("queue_size", {}) } except Exception as e: logger.error(f"Error getting worker stats: {str(e)}") return { "error": str(e), "timestamp": time.time() } def get_task_stats(self): """Get statistics about task processing.""" try: # This would typically connect to Flower API or Redis directly # Example using Flower API if it's running response = requests.get("http://localhost:5555/api/tasks") tasks = response.json() # Calculate statistics task_count = len(tasks) # Count tasks by status status_counts = {} for task_id, task in tasks.items(): status = task.get("state", "UNKNOWN") status_counts[status] = status_counts.get(status, 0) + 1 # Count tasks by type type_counts = {} for task_id, task in tasks.items(): task_name = task.get("name", "UNKNOWN") type_counts[task_name] = type_counts.get(task_name, 0) + 1 return { "task_count": task_count, "status_counts": status_counts, "type_counts": type_counts } except Exception as e: logger.error(f"Error getting task stats: {str(e)}") return { "error": str(e), "timestamp": time.time() } def check_health(self): """Check if Celery is healthy.""" try: worker_stats = self.get_worker_stats() # Consider unhealthy if no active workers if worker_stats.get("active_workers", 0) == 0: return { "status": "unhealthy", "reason": "No active workers", "timestamp": time.time() } # Consider unhealthy if excessive failed tasks if worker_stats.get("failed_tasks", 0) > 1000: return { "status": "degraded", "reason": "High failure rate", "timestamp": time.time() } return { "status": "healthy", "active_workers": worker_stats.get("active_workers", 0), "timestamp": time.time() } except Exception as e: logger.error(f"Error checking Celery health: {str(e)}") return { "status": "unknown", "error": str(e), "timestamp": time.time() } ``` 2. **Handling Long-Running Tasks**: For production, you should implement proper handling of long-running tasks: ```python # long_running_task_handler.py import time import signal import threading from functools import wraps from celery.exceptions import SoftTimeLimitExceeded def timeout_handler(func=None, timeout=1800, callback=None): """ Decorator to handle timeouts for long-running functions. Args: func: The function to decorate timeout: Timeout in seconds callback: Function to call on timeout Returns: Decorated function """ def decorator(f): @wraps(f) def wrapped(*args, **kwargs): # Store the result result = [None] exception = [None] # Define thread target def target(): try: result[0] = f(*args, **kwargs) except Exception as e: exception[0] = e # Create and start thread thread = threading.Thread(target=target) thread.daemon = True thread.start() # Wait for thread to complete or timeout thread.join(timeout) # Handle timeout if thread.is_alive(): if callback: callback() # Raise timeout exception raise TimeoutError(f"Function {f.__name__} timed out after {timeout} seconds") # Handle exception if exception[0]: raise exception[0] return result[0] return wrapped if func: return decorator(func) return decorator # Example usage in task @shared_task(bind=True) def complex_analysis_task(self, task_id): try: # Get task task = AgentTask.objects.get(id=task_id) task.set_processing(self.request.id) # Define timeout callback def on_timeout(): logger.error(f"Task {task_id} timed out") task.set_failed("Task exceeded time limit") # Use timeout handler for complex processing @timeout_handler(timeout=1800, callback=on_timeout) def run_complex_analysis(data): # Complex analysis code here # ... return result # Run with timeout handler result = run_complex_analysis(task.input_data) # Update task task.set_completed(result) return result except SoftTimeLimitExceeded: logger.error(f"Task {task_id} exceeded soft time limit") task.set_failed("Task exceeded time limit") return {"error": "Task exceeded time limit"} except TimeoutError as e: logger.error(f"Task {task_id} timed out: {str(e)}") task.set_failed(f"Task timed out: {str(e)}") return {"error": str(e)} except Exception as e: logger.error(f"Error in task {task_id}: {str(e)}") task.set_failed(str(e)) return {"error": str(e)} ``` 3. **Rate Limiting and Backoff**: For production, implement rate limiting and exponential backoff: ```python # rate_limiting.py import time import redis import functools import random import logging logger = logging.getLogger('agent_app.rate_limiting') class RateLimiter: """Rate limiter using Redis.""" def __init__(self, redis_url, limit_key, limit_rate, limit_period=60): """ Initialize rate limiter. Args: redis_url: Redis URL limit_key: Key prefix for rate limiting limit_rate: Maximum number of calls limit_period: Period in seconds """ self.redis = redis.from_url(redis_url) self.limit_key = limit_key self.limit_rate = limit_rate self.limit_period = limit_period def is_rate_limited(self, subkey=None): """ Check if the current call is rate limited. Args: subkey: Optional subkey for more granular limiting Returns: bool: True if rate limited, False otherwise """ key = f"{self.limit_key}:{subkey}" if subkey else self.limit_key # Get current count current = self.redis.get(key) # If no current count, initialize if current is None: self.redis.set(key, 1, ex=self.limit_period) return False # Increment count count = self.redis.incr(key) # Check if rate limited if count > self.limit_rate: # Get TTL to know how long until reset ttl = self.redis.ttl(key) logger.warning(f"Rate limited for {key}. TTL: {ttl}") return True return False def get_remaining(self, subkey=None): """Get remaining calls allowed.""" key = f"{self.limit_key}:{subkey}" if subkey else self.limit_key # Get current count current = self.redis.get(key) if current is None: return self.limit_rate return max(0, self.limit_rate - int(current)) def get_reset_time(self, subkey=None): """Get time until rate limit resets.""" key = f"{self.limit_key}:{subkey}" if subkey else self.limit_key # Get TTL ttl = self.redis.ttl(key) if ttl < 0: return 0 return ttl def with_rate_limiting(limiter, subkey_func=None, max_retries=3, backoff_base=2): """ Decorator for rate limiting functions. Args: limiter: RateLimiter instance subkey_func: Function to extract subkey from args/kwargs max_retries: Maximum number of retries backoff_base: Base for exponential backoff Returns: Decorated function """ def decorator(func): @functools.wraps(func) def wrapped(*args, **kwargs): # Get subkey if provided subkey = None if subkey_func: subkey = subkey_func(*args, **kwargs) retries = 0 while retries <= max_retries: # Check rate limiting if limiter.is_rate_limited(subkey): # If max retries reached, raise exception if retries >= max_retries: reset_time = limiter.get_reset_time(subkey) raise RateLimitExceeded( f"Rate limit exceeded. Try again in {reset_time} seconds." ) # Calculate backoff with jitter backoff = (backoff_base ** retries) + random.uniform(0, 0.5) # Log and sleep logger.info(f"Rate limited. Retrying in {backoff:.2f} seconds. Retry {retries+1}/{max_retries}") time.sleep(backoff) retries += 1 else: # Not rate limited, execute function return func(*args, **kwargs) # Should not reach here due to exception above return None return wrapped return decorator class RateLimitExceeded(Exception): """Exception raised when rate limit is exceeded.""" pass # Example usage # Initialize limiter for OpenAI API openai_limiter = RateLimiter( redis_url="redis://localhost:6379/0", limit_key="openai_api", limit_rate=100, # 100 requests per minute limit_period=60 ) # Use in function @with_rate_limiting( limiter=openai_limiter, subkey_func=lambda user_id, *args, **kwargs: f"user:{user_id}", max_retries=3 ) def call_openai_api(user_id, prompt): # API call here pass ``` This stack is particularly well-suited for organizations that need to: - Build complex task orchestration systems with AI agents - Maintain a centralized knowledge base for semantic search - Implement conversational applications with persistent state - Create document processing workflows with AI analysis - Support background processing with robust task management ### Airflow + AutoGen + OpenAI Functions + Snowflake (Enterprise AI Automation) This stack is optimized for enterprise-grade AI workflows that require robust scheduling, governance, and integration with enterprise data platforms. It's particularly well-suited for data-intensive applications that need to operate on a schedule and integrate with existing data infrastructure. **Architecture Overview:** ![Airflow + AutoGen + Snowflake Architecture](https://i.imgur.com/sY0F3QK.png) The architecture consists of: 1. **Apache Airflow**: Workflow orchestration engine for scheduling and monitoring 2. **AutoGen**: Multi-agent orchestration framework 3. **OpenAI Functions**: Structured function calling for agents 4. **Snowflake**: Enterprise data platform for storage and analytics 5. **MLflow**: Experiment tracking and model registry **Implementation Example:** Let's implement an Airflow DAG that orchestrates AI agents to analyze financial data in Snowflake: ```python # dags/financial_analysis_agent_dag.py import os import json import datetime import tempfile import pendulum import autogen import pandas as pd import snowflake.connector from snowflake.connector.pandas_tools import write_pandas import openai import requests import mlflow from io import StringIO from airflow import DAG from airflow.models import Variable from airflow.operators.python import PythonOperator from airflow.operators.trigger_dagrun import TriggerDagRunOperator from airflow.providers.snowflake.operators.snowflake import SnowflakeOperator from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook from airflow.utils.dates import days_ago from airflow.models.param import Param # Connect to OpenAI with API key from Airflow Variable openai_api_key = Variable.get("OPENAI_API_KEY", default_var="") os.environ["OPENAI_API_KEY"] = openai_api_key openai.api_key = openai_api_key # Default arguments for DAG default_args = { 'owner': 'data_science', 'depends_on_past': False, 'email': ['data_science@example.com'], 'email_on_failure': True, 'email_on_retry': False, 'retries': 1, 'retry_delay': datetime.timedelta(minutes=5), } # Define DAG dag = DAG( 'financial_analysis_agent', default_args=default_args, description='AI agent pipeline for financial data analysis', schedule_interval='0 4 * * 1-5', # 4 AM on weekdays start_date=days_ago(1), catchup=False, max_active_runs=1, concurrency=3, tags=['ai_agents', 'finance', 'analysis'], params={ 'analysis_date': Param( default=pendulum.now().subtract(days=1).to_date_string(), type='string', format='date' ), 'stock_symbols': Param( default='["AAPL", "MSFT", "GOOGL", "AMZN", "META"]', type='string' ), 'report_type': Param( default='standard', type='string', enum=['standard', 'detailed', 'executive'] ), 'include_sentiment': Param( default=True, type='boolean' ) } ) # Helper functions def get_snowflake_connection(): """Get Snowflake connection from Airflow hook.""" hook = SnowflakeHook(snowflake_conn_id='snowflake_default') conn = hook.get_conn() return conn def fetch_stock_data(date_str, symbols): """Fetch stock data from Snowflake for specified date and symbols.""" symbols_str = ", ".join([f"'{s}'" for s in symbols]) query = f""" SELECT symbol, date, open, high, low, close, volume, adj_close FROM finance.stocks.daily_prices WHERE date = '{date_str}' AND symbol IN ({symbols_str}) ORDER BY symbol, date """ conn = get_snowflake_connection() cursor = conn.cursor() cursor.execute(query) # Convert to Pandas DataFrame result = cursor.fetchall() columns = [desc[0] for desc in cursor.description] df = pd.DataFrame(result, columns=columns) cursor.close() return df def fetch_financial_news(date_str, symbols): """Fetch financial news from Snowflake for specified date and symbols.""" symbols_str = ", ".join([f"'{s}'" for s in symbols]) query = f""" SELECT headline, source, url, published_at, sentiment, symbols FROM finance.news.articles WHERE DATE(published_at) = '{date_str}' AND symbols_array && ARRAY_CONSTRUCT({symbols_str}) ORDER BY published_at DESC LIMIT 50 """ conn = get_snowflake_connection() cursor = conn.cursor() cursor.execute(query) # Convert to Pandas DataFrame result = cursor.fetchall() columns = [desc[0] for desc in cursor.description] df = pd.DataFrame(result, columns=columns) cursor.close() return df def store_analysis_results(analysis_results, date_str): """Store analysis results in Snowflake.""" # Create DataFrame from analysis results if isinstance(analysis_results, str): # If results are a string (like JSON), parse it try: results_dict = json.loads(analysis_results) except: # If not valid JSON, create a simple dict results_dict = {"analysis_text": analysis_results} else: # Already a dict-like object results_dict = analysis_results # Flatten nested dictionaries flat_results = {} for key, value in results_dict.items(): if isinstance(value, dict): for sub_key, sub_value in value.items(): if isinstance(sub_value, (dict, list)): flat_results[f"{key}_{sub_key}"] = json.dumps(sub_value) else: flat_results[f"{key}_{sub_key}"] = sub_value elif isinstance(value, list): flat_results[key] = json.dumps(value) else: flat_results[key] = value # Add analysis date flat_results['analysis_date'] = date_str flat_results['created_at'] = datetime.datetime.now().isoformat() # Create DataFrame df = pd.DataFrame([flat_results]) # Upload to Snowflake conn = get_snowflake_connection() success, num_chunks, num_rows, output = write_pandas( conn=conn, df=df, table_name='AGENT_ANALYSIS_RESULTS', schema='REPORTS', database='FINANCE' ) conn.close() return { 'success': success, 'num_rows': num_rows, 'table': 'FINANCE.REPORTS.AGENT_ANALYSIS_RESULTS' } # Task functions def extract_data(**context): """Extract relevant financial data for analysis.""" # Get parameters params = context['params'] analysis_date = params.get('analysis_date') stock_symbols = json.loads(params.get('stock_symbols')) # Fetch stock price data stock_data = fetch_stock_data(analysis_date, stock_symbols) if stock_data.empty: raise ValueError(f"No stock data found for {analysis_date} and symbols {stock_symbols}") # Fetch financial news news_data = fetch_financial_news(analysis_date, stock_symbols) if params.get('include_sentiment') else pd.DataFrame() # Calculate basic metrics metrics = {} for symbol in stock_symbols: symbol_data = stock_data[stock_data['symbol'] == symbol] if not symbol_data.empty: metrics[symbol] = { 'open': float(symbol_data['open'].iloc[0]), 'close': float(symbol_data['close'].iloc[0]), 'high': float(symbol_data['high'].iloc[0]), 'low': float(symbol_data['low'].iloc[0]), 'volume': int(symbol_data['volume'].iloc[0]), 'daily_change': float(symbol_data['close'].iloc[0] - symbol_data['open'].iloc[0]), 'daily_change_pct': float((symbol_data['close'].iloc[0] - symbol_data['open'].iloc[0]) / symbol_data['open'].iloc[0] * 100) } # Prepare data for AI analysis analysis_data = { 'date': analysis_date, 'symbols': stock_symbols, 'stock_data': stock_data.to_dict(orient='records'), 'metrics': metrics, 'news_data': news_data.to_dict(orient='records') if not news_data.empty else [] } # Save to XCom for next task context['ti'].xcom_push(key='analysis_data', value=analysis_data) return analysis_data def run_analysis_agents(**context): """Run AI agents for financial analysis.""" # Get parameters and data params = context['params'] analysis_data = context['ti'].xcom_pull(task_ids='extract_data', key='analysis_data') report_type = params.get('report_type') # Configure OpenAI client = openai.OpenAI(api_key=openai_api_key) # Define AutoGen agents for financial analysis # 1. Financial Analyst Agent - Core analysis analyst_agent = autogen.AssistantAgent( name="FinancialAnalyst", llm_config={ "config_list": [{"model": "gpt-4-turbo", "api_key": openai_api_key}], "temperature": 0.2, "functions": [ { "name": "analyze_stock_performance", "description": "Analyze the performance of stocks based on price data", "parameters": { "type": "object", "properties": { "symbol": {"type": "string", "description": "Stock symbol to analyze"}, "metrics": { "type": "object", "description": "Performance metrics to calculate" }, "context": {"type": "string", "description": "Additional context for analysis"} }, "required": ["symbol", "metrics"] } }, { "name": "analyze_news_sentiment", "description": "Analyze sentiment from news articles", "parameters": { "type": "object", "properties": { "symbol": {"type": "string", "description": "Stock symbol to analyze news for"}, "news_items": {"type": "array", "description": "List of news articles"}, "summary_length": {"type": "integer", "description": "Length of summary to generate"} }, "required": ["symbol", "news_items"] } } ] }, system_message="""You are an expert financial analyst specialized in stock market analysis. Your task is to analyze stock performance and provide insights based on price data and news. Be analytical, precise, and focus on data-driven insights. Consider market trends, volatility, and comparative performance when analyzing stocks. Your analysis should be suitable for institutional investors and financial professionals.""" ) # 2. Data Scientist Agent - Advanced metrics and models data_scientist_agent = autogen.AssistantAgent( name="DataScientist", llm_config={ "config_list": [{"model": "gpt-4-turbo", "api_key": openai_api_key}], "temperature": 0.1, "functions": [ { "name": "calculate_technical_indicators", "description": "Calculate technical indicators for stock analysis", "parameters": { "type": "object", "properties": { "symbol": {"type": "string", "description": "Stock symbol"}, "price_data": {"type": "object", "description": "Price data for the stock"}, "indicators": {"type": "array", "items": {"type": "string"}, "description": "Indicators to calculate"} }, "required": ["symbol", "price_data"] } }, { "name": "compare_performance", "description": "Compare performance between multiple stocks", "parameters": { "type": "object", "properties": { "symbols": {"type": "array", "items": {"type": "string"}, "description": "Stock symbols to compare"}, "metrics": {"type": "object", "description": "Metrics for each stock"} }, "required": ["symbols", "metrics"] } } ] }, system_message="""You are an expert data scientist specializing in financial markets. Your role is to perform advanced statistical analysis and calculate technical indicators. Focus on quantitative metrics, correlations, and statistical significance. Identify patterns and anomalies in the data that might not be immediately obvious. Your analysis should be rigorous and mathematically sound.""" ) # 3. Report Writer Agent - Generate final report report_writer_agent = autogen.AssistantAgent( name="ReportWriter", llm_config={ "config_list": [{"model": "gpt-4-turbo", "api_key": openai_api_key}], "temperature": 0.7, "functions": [ { "name": "generate_financial_report", "description": "Generate a comprehensive financial report", "parameters": { "type": "object", "properties": { "title": {"type": "string", "description": "Report title"}, "date": {"type": "string", "description": "Analysis date"}, "summary": {"type": "string", "description": "Executive summary"}, "stock_analyses": {"type": "object", "description": "Analysis for each stock"}, "market_overview": {"type": "string", "description": "Overall market context"}, "recommendations": {"type": "array", "items": {"type": "string"}, "description": "Investment recommendations"}, "report_type": {"type": "string", "enum": ["standard", "detailed", "executive"], "description": "Type of report to generate"} }, "required": ["title", "date", "stock_analyses", "report_type"] } } ] }, system_message="""You are an expert financial report writer. Your task is to synthesize financial analysis into clear, professional reports. Organize information logically with appropriate sections and headers. Use precise financial terminology while keeping the content accessible. Highlight key insights and structure the report according to the specified type: - standard: Balanced detail and length for general professional use - detailed: Comprehensive analysis with extensive data and charts - executive: Concise summary focused on key takeaways and recommendations""" ) # User proxy agent to coordinate the workflow user_proxy = autogen.UserProxyAgent( name="FinancialDataManager", human_input_mode="NEVER", code_execution_config={"work_dir": "financial_analysis_workspace"} ) # Create a group chat for the agents groupchat = autogen.GroupChat( agents=[user_proxy, analyst_agent, data_scientist_agent, report_writer_agent], messages=[], max_round=15 ) manager = autogen.GroupChatManager(groupchat=groupchat) # Start the analysis process stock_data_str = json.dumps(analysis_data['stock_data'][:5]) if len(analysis_data['stock_data']) > 5 else json.dumps(analysis_data['stock_data']) news_data_str = json.dumps(analysis_data['news_data'][:5]) if len(analysis_data['news_data']) > 5 else json.dumps(analysis_data['news_data']) prompt = f""" Task: Perform financial analysis for the following stocks: {analysis_data['symbols']} on {analysis_data['date']}. Report Type: {report_type} Stock Metrics: {json.dumps(analysis_data['metrics'], indent=2)} Sample Stock Data: {stock_data_str} Sample News Data: {news_data_str} Create a comprehensive financial analysis report with the following components: 1. Market overview for the date 2. Individual stock analysis for each symbol 3. Comparative performance analysis 4. Key insights and patterns 5. Recommendations based on the data The FinancialAnalyst should begin by analyzing each stock's performance. The DataScientist should then calculate technical indicators and compare performance. Finally, the ReportWriter should compile all analyses into a coherent report. The final deliverable should be a complete financial analysis report in the requested format. """ # Start the group chat result = user_proxy.initiate_chat(manager, message=prompt) # Extract the final report from the chat final_report = None for message in reversed(user_proxy.chat_history): if message['role'] == 'assistant' and 'ReportWriter' in message.get('name', ''): final_report = message['content'] break if not final_report: # Extract best available result if no clear final report for message in reversed(user_proxy.chat_history): if message['role'] == 'assistant' and len(message['content']) > 500: final_report = message['content'] break # Process and structure the report try: # Try to extract structured data using OpenAI structure_response = client.chat.completions.create( model="gpt-4-turbo", messages=[ {"role": "system", "content": "You are a financial data extraction specialist. Extract structured data from financial analysis reports."}, {"role": "user", "content": f"Extract the key structured data from this financial report in JSON format. Include market_overview, stock_analyses (with individual metrics for each stock), key_insights, and recommendations:\n\n{final_report}"} ], response_format={"type": "json_object"} ) structured_report = json.loads(structure_response.choices[0].message.content) except Exception as e: # Fall back to simple structure if extraction fails structured_report = { "report_text": final_report, "report_type": report_type, "analysis_date": analysis_data['date'], "symbols": analysis_data['symbols'] } # Combine structured report with raw text final_result = { "structured_data": structured_report, "full_report": final_report, "report_type": report_type, "analysis_date": analysis_data['date'], "symbols_analyzed": analysis_data['symbols'], "generation_metadata": { "timestamp": datetime.datetime.now().isoformat(), "model": "gpt-4-turbo", "agent_framework": "AutoGen" } } # Log with MLflow if enabled try: mlflow.start_run(run_name=f"financial_analysis_{analysis_data['date']}") # Log parameters mlflow.log_params({ "analysis_date": analysis_data['date'], "symbols": ",".join(analysis_data['symbols']), "report_type": report_type, "include_sentiment": params.get('include_sentiment') }) # Log metrics if available if 'structured_data' in final_result and 'stock_analyses' in final_result['structured_data']: for symbol, analysis in final_result['structured_data']['stock_analyses'].items(): if isinstance(analysis, dict): for metric, value in analysis.items(): if isinstance(value, (int, float)): mlflow.log_metric(f"{symbol}_{metric}", value) # Log report as artifact with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f: f.write(final_report) report_path = f.name mlflow.log_artifact(report_path) os.unlink(report_path) # Log raw data sample with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f: json.dump(analysis_data, f) data_path = f.name mlflow.log_artifact(data_path) os.unlink(data_path) mlflow.end_run() except Exception as e: print(f"Error logging to MLflow: {e}") # Save to XCom for next task context['ti'].xcom_push(key='analysis_results', value=final_result) return final_result def store_results(**context): """Store analysis results in Snowflake.""" # Get analysis results analysis_results = context['ti'].xcom_pull(task_ids='run_analysis_agents', key='analysis_results') params = context['params'] analysis_date = params.get('analysis_date') # Store in Snowflake storage_result = store_analysis_results(analysis_results, analysis_date) # Generate report file in S3 via Snowflake (optional) conn = get_snowflake_connection() cursor = conn.cursor() report_file_query = f""" COPY INTO @FINANCE.REPORTS.REPORT_STAGE/financial_reports/ FROM ( SELECT OBJECT_CONSTRUCT('report', full_report, 'metadata', generation_metadata, 'date', analysis_date) AS report_json FROM FINANCE.REPORTS.AGENT_ANALYSIS_RESULTS WHERE analysis_date = '{analysis_date}' ORDER BY created_at DESC LIMIT 1 ) FILE_FORMAT = (TYPE = JSON) OVERWRITE = TRUE SINGLE = TRUE HEADER = TRUE; """ cursor.execute(report_file_query) file_result = cursor.fetchall() cursor.close() conn.close() # Return combined results result = { 'snowflake_storage': storage_result, 'report_file': file_result } return result def notify_stakeholders(**context): """Send notification about completed analysis.""" # Get parameters and results params = context['params'] analysis_date = params.get('analysis_date') stock_symbols = json.loads(params.get('stock_symbols')) analysis_results = context['ti'].xcom_pull(task_ids='run_analysis_agents', key='analysis_results') # Extract key insights if available key_insights = [] if ('structured_data' in analysis_results and 'key_insights' in analysis_results['structured_data']): if isinstance(analysis_results['structured_data']['key_insights'], list): key_insights = analysis_results['structured_data']['key_insights'] elif isinstance(analysis_results['structured_data']['key_insights'], str): key_insights = [analysis_results['structured_data']['key_insights']] # Build notification content notification = { 'title': f"Financial Analysis Report - {analysis_date}", 'date': analysis_date, 'symbols': stock_symbols, 'key_insights': key_insights[:3], # Just top 3 insights 'report_url': f"https://analytics.example.com/reports/finance/{analysis_date.replace('-', '')}.html", 'snowflake_table': "FINANCE.REPORTS.AGENT_ANALYSIS_RESULTS" } # Log notification (in production, would actually send via email/Slack) print(f"Would send notification: {json.dumps(notification, indent=2)}") # Return notification content return notification # Define DAG tasks extract_task = PythonOperator( task_id='extract_data', python_callable=extract_data, provide_context=True, dag=dag, ) analysis_task = PythonOperator( task_id='run_analysis_agents', python_callable=run_analysis_agents, provide_context=True, dag=dag, ) store_task = PythonOperator( task_id='store_results', python_callable=store_results, provide_context=True, dag=dag, ) notify_task = PythonOperator( task_id='notify_stakeholders', python_callable=notify_stakeholders, provide_context=True, dag=dag, ) # Set task dependencies extract_task >> analysis_task >> store_task >> notify_task ``` For tracking experiments and agent performance, let's implement an MLflow tracking component: ```python # mlflow_tracking.py import os import json import mlflow import datetime import pandas as pd from typing import Dict, Any, List, Optional class AIAgentExperimentTracker: """Track AI agent experiments with MLflow.""" def __init__( self, experiment_name: str, tracking_uri: Optional[str] = None, tags: Optional[Dict[str, str]] = None ): """ Initialize the experiment tracker. Args: experiment_name: Name of the MLflow experiment tracking_uri: Optional URI for MLflow tracking server tags: Optional tags for the experiment """ self.experiment_name = experiment_name # Set tracking URI if provided if tracking_uri: mlflow.set_tracking_uri(tracking_uri) # Set default tags self.default_tags = tags or {} # Get or create experiment try: self.experiment = mlflow.get_experiment_by_name(experiment_name) if not self.experiment: self.experiment_id = mlflow.create_experiment( experiment_name, tags=self.default_tags ) else: self.experiment_id = self.experiment.experiment_id except Exception as e: print(f"Error initializing MLflow experiment: {e}") self.experiment_id = None def start_run( self, run_name: Optional[str] = None, tags: Optional[Dict[str, str]] = None, agent_config: Optional[Dict[str, Any]] = None ) -> str: """ Start a new MLflow run. Args: run_name: Optional name for the run tags: Optional tags for the run agent_config: Optional agent configuration to log Returns: str: MLflow run ID """ # Generate default run name if not provided if not run_name: run_name = f"agent_run_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}" # Combine default tags with run-specific tags run_tags = {**self.default_tags, **(tags or {})} try: # Start the run mlflow.start_run( experiment_id=self.experiment_id, run_name=run_name, tags=run_tags ) # Log agent configuration if provided if agent_config: # Log nested dictionaries as separate params for better organization self._log_nested_params("agent", agent_config) # Also log the raw config as JSON artifact for preservation config_path = f"/tmp/agent_config_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.json" with open(config_path, 'w') as f: json.dump(agent_config, f, indent=2) mlflow.log_artifact(config_path) os.remove(config_path) return mlflow.active_run().info.run_id except Exception as e: print(f"Error starting MLflow run: {e}") return None def log_agent_interaction( self, agent_name: str, prompt: str, response: str, metrics: Optional[Dict[str, float]] = None, metadata: Optional[Dict[str, Any]] = None ): """ Log an agent interaction. Args: agent_name: Name of the agent prompt: Prompt sent to agent response: Agent response metrics: Optional metrics for the interaction metadata: Optional metadata about the interaction """ try: # Ensure we have an active run if not mlflow.active_run(): self.start_run(f"{agent_name}_interactions") # Log metrics if provided if metrics: for key, value in metrics.items(): if isinstance(value, (int, float)): mlflow.log_metric(f"{agent_name}_{key}", value) # Log metadata as parameters if metadata: flat_metadata = self._flatten_dict(metadata) for key, value in flat_metadata.items(): if isinstance(value, (str, int, float, bool)): mlflow.log_param(f"{agent_name}_{key}", value) # Log the interaction as text interaction_path = f"/tmp/{agent_name}_interaction_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.txt" with open(interaction_path, 'w') as f: f.write(f"PROMPT:\n{prompt}\n\nRESPONSE:\n{response}") mlflow.log_artifact(interaction_path, artifact_path=f"interactions/{agent_name}") os.remove(interaction_path) except Exception as e: print(f"Error logging agent interaction: {e}") def log_agent_evaluation( self, evaluations: Dict[str, Any], metrics: Optional[Dict[str, float]] = None ): """ Log agent evaluation results. Args: evaluations: Evaluation results metrics: Additional metrics to log """ try: # Ensure we have an active run if not mlflow.active_run(): self.start_run("agent_evaluation") # Log evaluation metrics if metrics: for key, value in metrics.items(): if isinstance(value, (int, float)): mlflow.log_metric(key, value) # Log structured evaluations evaluation_path = f"/tmp/evaluation_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.json" with open(evaluation_path, 'w') as f: json.dump(evaluations, f, indent=2) mlflow.log_artifact(evaluation_path, artifact_path="evaluations") os.remove(evaluation_path) # If evaluations contain numeric scores, log as metrics flat_evals = self._flatten_dict(evaluations) for key, value in flat_evals.items(): if isinstance(value, (int, float)): mlflow.log_metric(f"eval_{key}", value) except Exception as e: print(f"Error logging agent evaluation: {e}") def log_output_data( self, data: Any, output_format: str = "json", name: Optional[str] = None ): """ Log output data from an agent run. Args: data: Data to log output_format: Format to use (json, csv, txt) name: Optional name for the output """ try: # Ensure we have an active run if not mlflow.active_run(): self.start_run("agent_output") if name is None: name = f"output_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}" # Process based on format if output_format == "json": output_path = f"/tmp/{name}.json" with open(output_path, 'w') as f: if isinstance(data, str): f.write(data) else: json.dump(data, f, indent=2, default=str) elif output_format == "csv": output_path = f"/tmp/{name}.csv" if isinstance(data, pd.DataFrame): data.to_csv(output_path, index=False) elif isinstance(data, list) and all(isinstance(x, dict) for x in data): pd.DataFrame(data).to_csv(output_path, index=False) else: raise ValueError("Data must be DataFrame or list of dicts for CSV format") elif output_format == "txt": output_path = f"/tmp/{name}.txt" with open(output_path, 'w') as f: f.write(str(data)) else: raise ValueError(f"Unsupported output format: {output_format}") # Log the artifact mlflow.log_artifact(output_path, artifact_path="outputs") os.remove(output_path) except Exception as e: print(f"Error logging output data: {e}") def end_run(self): """End the current MLflow run.""" try: if mlflow.active_run(): mlflow.end_run() except Exception as e: print(f"Error ending MLflow run: {e}") def _log_nested_params(self, prefix, params_dict): """Log nested parameters with prefixed keys.""" flat_params = self._flatten_dict(params_dict, prefix) for key, value in flat_params.items(): if isinstance(value, (str, int, float, bool)): mlflow.log_param(key, value) def _flatten_dict(self, d, parent_key='', sep='_'): """Flatten nested dictionaries for parameter logging.""" items = [] for k, v in d.items(): new_key = f"{parent_key}{sep}{k}" if parent_key else k if isinstance(v, dict): items.extend(self._flatten_dict(v, new_key, sep=sep).items()) else: items.append((new_key, v)) return dict(items) # Example usage if __name__ == "__main__": # Initialize tracker tracker = AIAgentExperimentTracker( experiment_name="financial_analysis_agents", tags={"domain": "finance", "purpose": "stock_analysis"} ) # Start a run run_id = tracker.start_run( run_name="daily_market_analysis", tags={"stocks": "AAPL,MSFT,GOOGL", "date": "2023-09-15"}, agent_config={ "agent_types": ["FinancialAnalyst", "DataScientist", "ReportWriter"], "models": { "primary": "gpt-4-turbo", "fallback": "gpt-3.5-turbo" }, "temperature": 0.2 } ) # Log sample interaction tracker.log_agent_interaction( agent_name="FinancialAnalyst", prompt="Analyze AAPL performance on 2023-09-15", response="Apple (AAPL) closed at $175.62, down 0.8% from previous close...", metrics={ "tokens": 250, "response_time": 0.85, "cost": 0.02 }, metadata={ "model": "gpt-4-turbo", "temperature": 0.2 } ) # Log evaluation tracker.log_agent_evaluation( evaluations={ "accuracy": 0.92, "completeness": 0.88, "reasoning": 0.90, "usefulness": 0.85, "detailed_scores": { "factual_accuracy": 0.95, "calculation_accuracy": 0.89, "insight_quality": 0.87 } }, metrics={ "overall_quality": 0.89, "execution_time": 12.5 } ) # Log output tracker.log_output_data( { "report": "Financial Analysis Report - 2023-09-15", "stocks_analyzed": ["AAPL", "MSFT", "GOOGL"], "key_insights": [ "Tech sector showed weakness with average decline of 0.7%", "AAPL volume 15% above 30-day average despite price decline", "MSFT outperformed peers with 0.2% gain" ] }, output_format="json", name="financial_report" ) # End the run tracker.end_run() ``` Let's also create a Snowflake integration component for enterprise data management: ```python # snowflake_integration.py import os import json import pandas as pd import snowflake.connector from snowflake.connector.pandas_tools import write_pandas from typing import Dict, Any, List, Optional, Union class SnowflakeAgentIntegration: """Integration with Snowflake for AI agent data pipelines.""" def __init__( self, account: str, user: str, password: str = None, database: str = None, schema: str = None, warehouse: str = None, role: str = None, authenticator: str = None, private_key_path: str = None, private_key_passphrase: str = None ): """ Initialize Snowflake connection parameters. Args: account: Snowflake account identifier user: Snowflake username password: Optional password (use private key or SSO instead for production) database: Default database schema: Default schema warehouse: Compute warehouse role: Snowflake role authenticator: Authentication method (e.g., 'externalbrowser' for SSO) private_key_path: Path to private key file for key-pair authentication private_key_passphrase: Passphrase for private key if encrypted """ self.account = account self.user = user self.password = password self.database = database self.schema = schema self.warehouse = warehouse self.role = role self.authenticator = authenticator self.private_key_path = private_key_path self.private_key_passphrase = private_key_passphrase # Initialize connection as None self.conn = None def connect(self): """Establish connection to Snowflake.""" try: # Prepare connection parameters connect_params = { "account": self.account, "user": self.user, "database": self.database, "schema": self.schema, "warehouse": self.warehouse, "role": self.role } # Add authentication method if self.password: connect_params["password"] = self.password elif self.authenticator: connect_params["authenticator"] = self.authenticator elif self.private_key_path: with open(self.private_key_path, "rb") as key: p_key = key.read() if self.private_key_passphrase: connect_params["private_key"] = p_key connect_params["private_key_passphrase"] = self.private_key_passphrase else: connect_params["private_key"] = p_key # Remove None values connect_params = {k: v for k, v in connect_params.items() if v is not None} # Establish connection self.conn = snowflake.connector.connect(**connect_params) return self.conn except Exception as e: print(f"Error connecting to Snowflake: {e}") raise def execute_query(self, query: str, params: Dict[str, Any] = None) -> List[Dict[str, Any]]: """ Execute a SQL query and return results as list of dictionaries. Args: query: SQL query to execute params: Optional query parameters Returns: List of dictionaries with query results """ try: # Connect if not already connected if not self.conn: self.connect() # Create cursor and execute query cursor = self.conn.cursor(snowflake.connector.DictCursor) if params: cursor.execute(query, params) else: cursor.execute(query) # Fetch results results = cursor.fetchall() # Close cursor cursor.close() return results except Exception as e: print(f"Error executing query: {e}") raise def query_to_dataframe(self, query: str, params: Dict[str, Any] = None) -> pd.DataFrame: """ Execute a SQL query and return results as Pandas DataFrame. Args: query: SQL query to execute params: Optional query parameters Returns: Pandas DataFrame with query results """ try: # Connect if not already connected if not self.conn: self.connect() # Execute query directly to DataFrame if params: df = pd.read_sql(query, self.conn, params=params) else: df = pd.read_sql(query, self.conn) return df except Exception as e: print(f"Error executing query to DataFrame: {e}") raise def upload_dataframe( self, df: pd.DataFrame, table_name: str, schema: Optional[str] = None, database: Optional[str] = None, chunk_size: Optional[int] = None, auto_create_table: bool = False ) -> Dict[str, Any]: """ Upload a Pandas DataFrame to Snowflake table. Args: df: DataFrame to upload table_name: Destination table name schema: Optional schema (overrides default) database: Optional database (overrides default) chunk_size: Optional chunk size for large uploads auto_create_table: Whether to automatically create table if it doesn't exist Returns: Dictionary with upload results """ try: # Connect if not already connected if not self.conn: self.connect() # Use default schema/database if not specified schema = schema or self.schema database = database or self.database # Create fully qualified table name qualified_table_name = f"{database}.{schema}.{table_name}" if database and schema else table_name # Check if table exists if auto_create_table: self._ensure_table_exists(df, qualified_table_name) # Upload DataFrame success, num_chunks, num_rows, output = write_pandas( conn=self.conn, df=df, table_name=table_name, schema=schema, database=database, chunk_size=chunk_size, quote_identifiers=True ) return { "success": success, "chunks": num_chunks, "rows": num_rows, "output": output, "table": qualified_table_name } except Exception as e: print(f"Error uploading DataFrame: {e}") raise def upload_json( self, data: Union[Dict[str, Any], List[Dict[str, Any]]], table_name: str, schema: Optional[str] = None, database: Optional[str] = None, flatten: bool = False ) -> Dict[str, Any]: """ Upload JSON data to Snowflake table. Args: data: Dictionary or list of dictionaries to upload table_name: Destination table name schema: Optional schema (overrides default) database: Optional database (overrides default) flatten: Whether to flatten nested structures Returns: Dictionary with upload results """ try: # Convert to DataFrame based on data type if isinstance(data, dict): if flatten: # Flatten nested dict flat_data = self._flatten_json(data) df = pd.DataFrame([flat_data]) else: # Convert single dict to DataFrame with one row df = pd.DataFrame([data]) elif isinstance(data, list) and all(isinstance(item, dict) for item in data): if flatten: # Flatten each dict in the list flat_list = [self._flatten_json(item) for item in data] df = pd.DataFrame(flat_list) else: # Convert list of dicts directly to DataFrame df = pd.DataFrame(data) else: raise ValueError("Data must be a dictionary or list of dictionaries") # Handle nested JSON structures by converting to strings for col in df.columns: if isinstance(df[col].iloc[0], (dict, list)): df[col] = df[col].apply(lambda x: json.dumps(x)) # Upload the DataFrame return self.upload_dataframe( df=df, table_name=table_name, schema=schema, database=database, auto_create_table=True ) except Exception as e: print(f"Error uploading JSON: {e}") raise def store_agent_results( self, agent_results: Dict[str, Any], metadata: Dict[str, Any], table_name: str = "AGENT_RESULTS", schema: Optional[str] = None, database: Optional[str] = None ) -> Dict[str, Any]: """ Store AI agent results with metadata in Snowflake. Args: agent_results: Results from AI agent metadata: Metadata about the agent run table_name: Destination table name schema: Optional schema (overrides default) database: Optional database (overrides default) Returns: Dictionary with upload results """ try: # Prepare combined data combined_data = { "results": json.dumps(agent_results), "metadata": json.dumps(metadata), "created_at": pd.Timestamp.now() } # Add metadata fields as top-level columns for easier querying for key, value in metadata.items(): if isinstance(value, (str, int, float, bool)) or value is None: combined_data[f"meta_{key}"] = value # Upload to Snowflake df = pd.DataFrame([combined_data]) return self.upload_dataframe( df=df, table_name=table_name, schema=schema, database=database, auto_create_table=True ) except Exception as e: print(f"Error storing agent results: {e}") raise def close(self): """Close the Snowflake connection.""" if self.conn: self.conn.close() self.conn = None def _ensure_table_exists(self, df: pd.DataFrame, table_name: str): """ Create table if it doesn't exist based on DataFrame structure. Args: df: DataFrame to use for table schema table_name: Fully qualified table name """ try: # Check if table exists check_query = f"SHOW TABLES LIKE '{table_name.split('.')[-1]}'" if '.' in table_name: parts = table_name.split('.') if len(parts) == 3: check_query = f"SHOW TABLES LIKE '{parts[2]}' IN SCHEMA {parts[0]}.{parts[1]}" cursor = self.conn.cursor() cursor.execute(check_query) table_exists = cursor.fetchone() is not None if not table_exists: # Generate CREATE TABLE statement based on DataFrame columns = [] for col_name, dtype in zip(df.columns, df.dtypes): if pd.api.types.is_integer_dtype(dtype): col_type = "INTEGER" elif pd.api.types.is_float_dtype(dtype): col_type = "FLOAT" elif pd.api.types.is_bool_dtype(dtype): col_type = "BOOLEAN" elif pd.api.types.is_datetime64_dtype(dtype): col_type = "TIMESTAMP_NTZ" else: # Check if column contains JSON if df[col_name].iloc[0] and isinstance(df[col_name].iloc[0], str): try: json.loads(df[col_name].iloc[0]) col_type = "VARIANT" # For JSON data except: col_type = "VARCHAR" else: col_type = "VARCHAR" columns.append(f'"{col_name}" {col_type}') # Create the table create_query = f"CREATE TABLE IF NOT EXISTS {table_name} ({', '.join(columns)})" cursor.execute(create_query) cursor.close() except Exception as e: print(f"Error creating table: {e}") raise def _flatten_json(self, d: Dict[str, Any], parent_key: str = '', sep: str = '_') -> Dict[str, Any]: """ Flatten nested JSON structures. Args: d: Dictionary to flatten parent_key: Parent key for recursive calls sep: Separator for nested keys Returns: Flattened dictionary """ items = [] for k, v in d.items(): new_key = f"{parent_key}{sep}{k}" if parent_key else k if isinstance(v, dict): items.extend(self._flatten_json(v, new_key, sep=sep).items()) elif isinstance(v, list): # Convert lists to JSON strings items.append((new_key, json.dumps(v))) else: items.append((new_key, v)) return dict(items) ``` Let's also implement an agent evaluation component: ```python # agent_evaluation.py import json import pandas as pd import numpy as np import openai from typing import Dict, Any, List, Optional, Union, Tuple class AgentEvaluator: """Evaluate AI agents for quality, correctness, and performance.""" def __init__(self, openai_api_key: str): """ Initialize the evaluator. Args: openai_api_key: OpenAI API key for evaluation """ self.openai_client = openai.OpenAI(api_key=openai_api_key) def evaluate_agent_output( self, prompt: str, response: str, ground_truth: Optional[str] = None, criteria: Optional[List[str]] = None ) -> Dict[str, Any]: """ Evaluate an agent's response against criteria and optionally ground truth. Args: prompt: The original prompt given to the agent response: The agent's response ground_truth: Optional ground truth for factual comparison criteria: Optional evaluation criteria Returns: Dictionary with evaluation scores and feedback """ if criteria is None: criteria = [ "accuracy", "completeness", "relevance", "coherence", "conciseness" ] # Construct evaluation prompt eval_prompt = f"""Evaluate the following AI assistant response to a user prompt. USER PROMPT: {prompt} AI RESPONSE: {response} """ if ground_truth: eval_prompt += f""" GROUND TRUTH (for factual comparison): {ground_truth} """ eval_prompt += f""" Please evaluate the response on the following criteria on a scale of 1-10: {', '.join(criteria)} Provide an explanation for each score and give specific examples from the response. Then provide an overall score (1-10) and a brief summary of the evaluation. Format your response as a JSON object with the following structure: {{ "criteria_scores": {{ "criterion1": {{ "score": X, "explanation": "Your explanation" }}, ... }}, "overall_score": X, "summary": "Your summary", "strengths": ["strength1", "strength2", ...], "weaknesses": ["weakness1", "weakness2", ...] }} """ try: # Get evaluation from OpenAI response = self.openai_client.chat.completions.create( model="gpt-4-turbo", messages=[ {"role": "system", "content": "You are an objective evaluator of AI assistant responses. Provide fair, balanced, and detailed evaluations."}, {"role": "user", "content": eval_prompt} ], response_format={"type": "json_object"}, temperature=0.2 ) # Parse the result evaluation = json.loads(response.choices[0].message.content) # Add metadata evaluation["evaluation_metadata"] = { "model": "gpt-4-turbo", "prompt_length": len(prompt), "response_length": len(response), "criteria_evaluated": criteria } return evaluation except Exception as e: print(f"Error evaluating agent output: {e}") return { "error": str(e), "criteria_scores": {c: {"score": 0, "explanation": "Evaluation failed"} for c in criteria}, "overall_score": 0, "summary": f"Evaluation failed: {str(e)}" } def evaluate_factual_accuracy( self, response: str, ground_truth: str ) -> Dict[str, Any]: """ Evaluate the factual accuracy of an agent's response. Args: response: The agent's response ground_truth: The ground truth for comparison Returns: Dictionary with accuracy scores and details """ try: # Construct evaluation prompt eval_prompt = f"""Evaluate the factual accuracy of the following AI response compared to the ground truth. AI RESPONSE: {response} GROUND TRUTH: {ground_truth} Identify all factual statements in the AI response and check if they are: 1. Correct (matches ground truth) 2. Incorrect (contradicts ground truth) 3. Unverifiable (not mentioned in ground truth) For each factual claim, provide: 1. The claim from the AI response 2. Whether it's correct, incorrect, or unverifiable 3. The relevant ground truth information (if applicable) Then calculate: 1. Accuracy rate (correct claims / total verifiable claims) 2. Error rate (incorrect claims / total verifiable claims) 3. Hallucination rate (unverifiable claims / total claims) Format your response as a JSON object with the following structure: {{ "factual_claims": [ {{ "claim": "The claim text", "assessment": "correct|incorrect|unverifiable", "ground_truth_reference": "Relevant ground truth text or null", "explanation": "Explanation of assessment" }}, ... ], "metrics": {{ "total_claims": X, "correct_claims": X, "incorrect_claims": X, "unverifiable_claims": X, "accuracy_rate": X.XX, "error_rate": X.XX, "hallucination_rate": X.XX }}, "summary": "Overall assessment of factual accuracy" }} """ # Get evaluation from OpenAI response = self.openai_client.chat.completions.create( model="gpt-4-turbo", messages=[ {"role": "system", "content": "You are an expert fact-checker who carefully evaluates the factual accuracy of information."}, {"role": "user", "content": eval_prompt} ], response_format={"type": "json_object"}, temperature=0.1 ) # Parse the result evaluation = json.loads(response.choices[0].message.content) return evaluation except Exception as e: print(f"Error evaluating factual accuracy: {e}") return { "error": str(e), "metrics": { "accuracy_rate": 0, "error_rate": 0, "hallucination_rate": 0 }, "summary": f"Evaluation failed: {str(e)}" } def evaluate_multi_agent_workflow( self, task_description: str, agent_interactions: List[Dict[str, Any]], final_output: str, expected_output: Optional[str] = None ) -> Dict[str, Any]: """ Evaluate a multi-agent workflow. Args: task_description: The original task agent_interactions: List of agent interactions in the workflow final_output: The final output of the workflow expected_output: Optional expected output for comparison Returns: Dictionary with workflow evaluation """ try: # Format agent interactions for evaluation interactions_text = "" for i, interaction in enumerate(agent_interactions, 1): agent_name = interaction.get("agent_name", f"Agent {i}") prompt = interaction.get("prompt", "") response = interaction.get("response", "") interactions_text += f"\n--- INTERACTION {i} ---\n" interactions_text += f"AGENT: {agent_name}\n" interactions_text += f"PROMPT:\n{prompt}\n\n" interactions_text += f"RESPONSE:\n{response}\n" # Construct evaluation prompt eval_prompt = f"""Evaluate this multi-agent workflow for completing a task. TASK DESCRIPTION: {task_description} AGENT INTERACTIONS: {interactions_text} FINAL OUTPUT: {final_output} """ if expected_output: eval_prompt += f""" EXPECTED OUTPUT: {expected_output} """ eval_prompt += """ Evaluate the workflow on these criteria: 1. Task Completion: Did the agents successfully complete the task? 2. Efficiency: Was the workflow efficient, or were there unnecessary steps? 3. Agent Collaboration: How well did the agents collaborate and share information? 4. Agent Specialization: Did each agent contribute based on their expertise? 5. Error Handling: How well were errors or uncertainties handled? 6. Output Quality: How good is the final output? Format your response as a JSON object with the following structure: { "workflow_evaluation": { "task_completion": { "score": X, "comments": "Your assessment" }, "efficiency": { "score": X, "comments": "Your assessment" }, "agent_collaboration": { "score": X, "comments": "Your assessment" }, "agent_specialization": { "score": X, "comments": "Your assessment" }, "error_handling": { "score": X, "comments": "Your assessment" }, "output_quality": { "score": X, "comments": "Your assessment" } }, "agent_contributions": [ { "agent_name": "Agent name", "contribution_quality": X, "key_contributions": ["contribution1", "contribution2"] }, ... ], "overall_score": X, "improvement_suggestions": ["suggestion1", "suggestion2", ...], "summary": "Overall workflow assessment" } """ # Get evaluation from OpenAI response = self.openai_client.chat.completions.create( model="gpt-4-turbo", messages=[ {"role": "system", "content": "You are an expert in multi-agent AI systems who evaluates workflows for efficiency and effectiveness."}, {"role": "user", "content": eval_prompt} ], response_format={"type": "json_object"}, temperature=0.3 ) # Parse the result evaluation = json.loads(response.choices[0].message.content) # Calculate metrics scores = [ evaluation["workflow_evaluation"]["task_completion"]["score"], evaluation["workflow_evaluation"]["efficiency"]["score"], evaluation["workflow_evaluation"]["agent_collaboration"]["score"], evaluation["workflow_evaluation"]["agent_specialization"]["score"], evaluation["workflow_evaluation"]["error_handling"]["score"], evaluation["workflow_evaluation"]["output_quality"]["score"] ] avg_score = sum(scores) / len(scores) # Add calculated metrics evaluation["metrics"] = { "average_criteria_score": avg_score, "interaction_count": len(agent_interactions), "agent_count": len(set(interaction.get("agent_name", f"Agent {i}") for i, interaction in enumerate(agent_interactions))), "output_length": len(final_output) } return evaluation except Exception as e: print(f"Error evaluating multi-agent workflow: {e}") return { "error": str(e), "overall_score": 0, "summary": f"Evaluation failed: {str(e)}" } def benchmark_agent( self, agent_function, test_cases: List[Dict[str, Any]], metrics: Optional[List[str]] = None ) -> Dict[str, Any]: """ Benchmark an agent against a set of test cases. Args: agent_function: Function that takes input and returns agent response test_cases: List of test cases with input and expected output metrics: Optional list of metrics to evaluate Returns: Dictionary with benchmark results """ if metrics is None: metrics = ["accuracy", "relevance", "completeness"] results = [] for i, test_case in enumerate(test_cases): case_id = test_case.get("id", f"case_{i}") input_data = test_case.get("input", "") expected_output = test_case.get("expected_output", None) try: # Run the agent start_time = pd.Timestamp.now() agent_output = agent_function(input_data) end_time = pd.Timestamp.now() duration = (end_time - start_time).total_seconds() # Evaluate output evaluation = self.evaluate_agent_output( prompt=input_data, response=agent_output, ground_truth=expected_output, criteria=metrics ) # Compile results case_result = { "case_id": case_id, "input": input_data, "output": agent_output, "expected_output": expected_output, "execution_time": duration, "evaluation": evaluation, "overall_score": evaluation.get("overall_score", 0) } results.append(case_result) except Exception as e: print(f"Error in test case {case_id}: {e}") results.append({ "case_id": case_id, "input": input_data, "error": str(e), "overall_score": 0 }) # Aggregate results overall_scores = [r.get("overall_score", 0) for r in results if "overall_score" in r] avg_score = sum(overall_scores) / len(overall_scores) if overall_scores else 0 execution_times = [r.get("execution_time", 0) for r in results if "execution_time" in r] avg_execution_time = sum(execution_times) / len(execution_times) if execution_times else 0 # Calculate per-metric averages metric_scores = {} for metric in metrics: scores = [] for r in results: if "evaluation" in r and "criteria_scores" in r["evaluation"]: if metric in r["evaluation"]["criteria_scores"]: scores.append(r["evaluation"]["criteria_scores"][metric].get("score", 0)) metric_scores[metric] = sum(scores) / len(scores) if scores else 0 return { "benchmark_summary": { "test_cases": len(test_cases), "successful_cases": len([r for r in results if "error" not in r]), "average_score": avg_score, "average_execution_time": avg_execution_time, "metric_averages": metric_scores }, "case_results": results } ``` **Key Advantages:** 1. **Enterprise Integration**: Seamless integration with Snowflake for secure data storage and analytics 2. **Robust Scheduling**: Airflow provides enterprise-grade task scheduling and dependency management 3. **Workflow Monitoring**: Built-in monitoring and alerting for AI agent workflows 4. **Data Governance**: Enterprise-grade data lineage and governance with Snowflake 5. **Experiment Tracking**: MLflow integration for tracking agent performance and experiments **Production Considerations:** 1. **Securing API Keys**: For production deployment, implement proper API key management: ```python # api_key_management.py import os import base64 import json from cryptography.fernet import Fernet from cryptography.hazmat.primitives import hashes from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC from airflow.models import Variable from airflow.hooks.base import BaseHook class APIKeyManager: """Securely manage API keys in Airflow.""" def __init__(self, master_key_env="MASTER_ENCRYPTION_KEY"): """ Initialize the key manager. Args: master_key_env: Environment variable name for master key """ # Get master key from environment master_key = os.environ.get(master_key_env) if not master_key: raise ValueError(f"Master encryption key not found in environment variable {master_key_env}") # Derive encryption key kdf = PBKDF2HMAC( algorithm=hashes.SHA256(), length=32, salt=b'airflow_api_key_manager', iterations=100000, ) key = base64.urlsafe_b64encode(kdf.derive(master_key.encode())) self.cipher = Fernet(key) def encrypt_key(self, api_key): """Encrypt an API key.""" return self.cipher.encrypt(api_key.encode()).decode() def decrypt_key(self, encrypted_key): """Decrypt an API key.""" return self.cipher.decrypt(encrypted_key.encode()).decode() def store_in_airflow(self, key_name, api_key): """Store an encrypted API key in Airflow Variables.""" encrypted = self.encrypt_key(api_key) Variable.set(key_name, encrypted) def get_from_airflow(self, key_name): """Get and decrypt an API key from Airflow Variables.""" encrypted = Variable.get(key_name) return self.decrypt_key(encrypted) def store_connection(self, conn_id, conn_type, host, login, password, port=None, extra=None): """Store a connection in Airflow connections.""" # Encrypt sensitive parts encrypted_password = self.encrypt_key(password) # Create connection object conn = BaseHook.get_connection(conn_id) conn.conn_type = conn_type conn.host = host conn.login = login conn.password = encrypted_password conn.port = port if extra: # Encrypt extra fields if any if isinstance(extra, dict): encrypted_extra = {} for k, v in extra.items(): encrypted_extra[k] = self.encrypt_key(v) if isinstance(v, str) else v conn.extra = json.dumps(encrypted_extra) else: conn.extra = extra # Save connection conn.save() ``` 2. **Parameter Management and Validation**: ```python # parameter_management.py from marshmallow import Schema, fields, validates, ValidationError from typing import Dict, Any class FinancialAnalysisSchema(Schema): """Schema for validating financial analysis parameters.""" analysis_date = fields.Date(required=True) stock_symbols = fields.List(fields.String(), required=True) report_type = fields.String(required=True, validate=lambda x: x in ["standard", "detailed", "executive"]) include_sentiment = fields.Boolean(default=True) market_context = fields.Boolean(default=True) max_stocks = fields.Integer(default=10) @validates("stock_symbols") def validate_symbols(self, symbols): """Validate stock symbols.""" if not symbols: raise ValidationError("At least one stock symbol is required") if len(symbols) > 20: raise ValidationError("Maximum of 20 stock symbols allowed") for symbol in symbols: if not symbol.isalpha(): raise ValidationError(f"Invalid stock symbol: {symbol}") def validate_dag_params(params: Dict[str, Any], schema_class) -> Dict[str, Any]: """ Validate DAG parameters using a schema. Args: params: Parameters to validate schema_class: Schema class for validation Returns: Validated parameters Raises: ValueError: If validation fails """ schema = schema_class() try: # Validate parameters validated_params = schema.load(params) return validated_params except ValidationError as err: error_messages = [] for field, messages in err.messages.items(): if isinstance(messages, list): error_messages.append(f"{field}: {', '.join(messages)}") else: error_messages.append(f"{field}: {messages}") error_str = "; ".join(error_messages) raise ValueError(f"Parameter validation failed: {error_str}") ``` 3. **Airflow Optimizations**: For production Airflow deployments, consider these optimizations: ```python # airflow_config.py from airflow.models import Variable import subprocess import os # Recommended Airflow configuration optimizations def optimize_airflow_config(): """Apply optimizations to Airflow configuration.""" # Set environment variables os.environ["AIRFLOW__CORE__MAX_ACTIVE_RUNS_PER_DAG"] = "1" os.environ["AIRFLOW__CORE__PARALLELISM"] = "32" os.environ["AIRFLOW__CORE__DAG_CONCURRENCY"] = "16" os.environ["AIRFLOW__CORE__MAX_ACTIVE_TASKS_PER_DAG"] = "16" os.environ["AIRFLOW__SCHEDULER__SCHEDULER_HEARTBEAT_SEC"] = "20" os.environ["AIRFLOW__CORE__MIN_SERIALIZED_DAG_UPDATE_INTERVAL"] = "30" os.environ["AIRFLOW__CORE__MIN_SERIALIZED_DAG_FETCH_INTERVAL"] = "30" os.environ["AIRFLOW__CORE__STORE_DAG_CODE"] = "True" os.environ["AIRFLOW__CORE__STORE_SERIALIZED_DAGS"] = "True" os.environ["AIRFLOW__CORE__EXECUTE_TASKS_NEW_PYTHON_INTERPRETER"] = "True" # Configure Celery executor settings os.environ["AIRFLOW__CELERY__WORKER_AUTOSCALE"] = "8,2" os.environ["AIRFLOW__CELERY__WORKER_PREFETCH_MULTIPLIER"] = "1" os.environ["AIRFLOW__CELERY__TASK_POOL_LIMIT"] = "4" os.environ["AIRFLOW__CELERY__OPERATION_TIMEOUT"] = "1800" # 30 minutes # Logging optimizations os.environ["AIRFLOW__LOGGING__REMOTE_LOGGING"] = "True" os.environ["AIRFLOW__LOGGING__REMOTE_LOG_CONN_ID"] = "aws_default" os.environ["AIRFLOW__LOGGING__REMOTE_BASE_LOG_FOLDER"] = "s3://airflow-logs-bucket/logs" print("Applied Airflow optimizations") # Configure resource allocation for specific tasks def configure_task_resources(ti): """Configure resources for specific tasks in the DAG.""" task_id = ti.task_id # Configure based on task type if "analysis" in task_id: # Allocate more resources for analysis tasks ti.executor_config = { "KubernetesExecutor": { "request_memory": "4Gi", "request_cpu": "2", "limit_memory": "8Gi", "limit_cpu": "4" } } elif "extract" in task_id: # Database-heavy tasks ti.executor_config = { "KubernetesExecutor": { "request_memory": "2Gi", "request_cpu": "1", "limit_memory": "4Gi", "limit_cpu": "2" } } return ti ``` This stack is particularly well-suited for organizations that need to: - Integrate AI agents with enterprise data platforms - Schedule complex AI agent workflows - Maintain compliance and governance - Track AI agent performance over time - Support data-intensive AI processes ## 4. AI Agent Templates for Real-World Applications ### AI-Driven Financial Analyst (Market Data Analysis & Forecasting) This AI agent template is designed to analyze financial market data, identify trends, and provide forecasting and investment recommendations. It combines market data analysis, sentiment evaluation from news sources, and technical analysis to generate comprehensive financial insights. **Core Capabilities:** - Historical price analysis and pattern recognition - Sector and company fundamental analysis - News sentiment integration for market context - Technical indicator calculation and interpretation - Investment recommendation generation - Report creation with visualizations **Architecture:** ![Financial Analyst Agent Architecture](https://i.imgur.com/wZh8vKn.png) **Implementation Example:** ```python # financial_analyst_agent.py import datetime import pandas as pd import numpy as np import matplotlib.pyplot as plt import seaborn as sns import yfinance as yf import autogen from autogen.agentchat.contrib.gpt_assistant_agent import GPTAssistantAgent import json import requests import os from typing import Dict, List, Any, Optional, Union, Tuple # Configure API keys and settings OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "your-api-key") ALPHA_VANTAGE_API_KEY = os.environ.get("ALPHA_VANTAGE_API_KEY", "your-api-key") NEWS_API_KEY = os.environ.get("NEWS_API_KEY", "your-api-key") # Initialize OpenAI client import openai client = openai.OpenAI(api_key=OPENAI_API_KEY) class FinancialAnalystAgent: """ An AI-driven financial analyst agent that provides comprehensive market analysis, trend identification, and investment recommendations. """ def __init__(self, config=None): """ Initialize the Financial Analyst Agent. Args: config: Optional configuration dictionary """ self.config = config or {} # Set up agent configuration self.llm_config = { "config_list": [{"model": "gpt-4-turbo", "api_key": OPENAI_API_KEY}], "temperature": 0.2, "cache_seed": None # Disable caching for financial data which changes frequently } # Create the agent team self._create_agent_team() # Data cache self.data_cache = {} def _create_agent_team(self): """Create the team of specialized agents for financial analysis.""" # 1. Market Analyst - Specialized in general market trends and sector analysis self.market_analyst = autogen.AssistantAgent( name="MarketAnalyst", system_message="""You are an expert market analyst who specializes in understanding broad market trends, sector rotations, and macroeconomic factors affecting financial markets. Your responsibilities: 1. Analyze overall market conditions and trends 2. Identify sector strengths and weaknesses 3. Interpret macroeconomic data and its market impact 4. Provide context for market movements 5. Identify market sentiment and risk factors Always base your analysis on data and established financial theories. Avoid speculation without evidence. Present a balanced view that considers both bullish and bearish perspectives.""", llm_config=self.llm_config ) # 2. Technical Analyst - Specialized in chart patterns and technical indicators self.technical_analyst = autogen.AssistantAgent( name="TechnicalAnalyst", system_message="""You are an expert technical analyst who specializes in chart patterns, technical indicators, and price action analysis for financial markets. Your responsibilities: 1. Analyze price charts for significant patterns 2. Calculate and interpret technical indicators 3. Identify support and resistance levels 4. Analyze volume patterns and their implications 5. Provide technical-based forecasts Focus on objective technical analysis principles. Clearly explain the reasoning behind your analysis and the historical reliability of the patterns you identify. Always consider multiple timeframes.""", llm_config={ **self.llm_config, "functions": [ { "name": "calculate_technical_indicators", "description": "Calculate technical indicators for a stock", "parameters": { "type": "object", "properties": { "symbol": {"type": "string", "description": "Stock symbol"}, "indicators": {"type": "array", "items": {"type": "string"}, "description": "List of indicators to calculate"}, "period": {"type": "string", "description": "Time period for analysis (e.g., '1y', '6mo', '3mo')"} }, "required": ["symbol", "indicators"] } } ] } ) # 3. Fundamental Analyst - Specialized in company financial data self.fundamental_analyst = autogen.AssistantAgent( name="FundamentalAnalyst", system_message="""You are an expert fundamental analyst who specializes in analyzing company financial statements, valuation metrics, and business models. Your responsibilities: 1. Analyze company financial health and performance 2. Evaluate valuation metrics against industry peers 3. Assess growth prospects and business model strengths 4. Identify financial risks and opportunities 5. Provide fundamental-based investment recommendations Always use established valuation methodologies and accounting principles. Compare companies to their historical performance, sector peers, and the broader market. Consider both quantitative metrics and qualitative factors.""", llm_config=self.llm_config ) # 4. News Sentiment Analyst - Specialized in news and social media sentiment self.sentiment_analyst = autogen.AssistantAgent( name="SentimentAnalyst", system_message="""You are an expert in analyzing news and social media sentiment related to financial markets and individual companies. Your responsibilities: 1. Evaluate news sentiment affecting markets or specific stocks 2. Identify important news catalysts and their potential impact 3. Detect shifts in market narrative or sentiment 4. Assess information sources for reliability and importance 5. Contextualize news within broader market trends Focus on objective analysis of sentiment. Distinguish between substantive news and market noise. Consider the historical impact of similar news events and sentiment shifts.""", llm_config=self.llm_config ) # 5. Portfolio Advisor - Specialized in investment recommendations self.portfolio_advisor = autogen.AssistantAgent( name="PortfolioAdvisor", system_message="""You are an expert investment advisor who specializes in portfolio construction, risk management, and investment recommendations. Your responsibilities: 1. Synthesize analyses from other specialists into actionable advice 2. Provide specific investment recommendations with rationales 3. Consider risk management and portfolio allocation 4. Present balanced bull/bear cases for investments 5. Contextualize recommendations for different investor profiles Always include risk factors alongside potential rewards. Provide specific time horizons for recommendations when possible. Consider multiple scenarios and their implications. Make specific, actionable recommendations rather than general statements.""", llm_config=self.llm_config ) # 6. Report Writer - Specialized in creating comprehensive reports self.report_writer = autogen.AssistantAgent( name="ReportWriter", system_message="""You are an expert financial report writer who specializes in synthesizing complex financial analyses into clear, structured reports. Your responsibilities: 1. Organize analyses into a coherent narrative 2. Create executive summaries that highlight key points 3. Structure information logically with appropriate sections 4. Maintain professional financial writing standards 5. Ensure reports are comprehensive yet accessible Use clear financial terminology and explain complex concepts when necessary. Include all relevant information while avoiding unnecessary repetition. Organize content with appropriate headings and structure. Always include an executive summary and conclusion.""", llm_config=self.llm_config ) # User proxy agent for orchestrating the workflow self.user_proxy = autogen.UserProxyAgent( name="FinancialDataManager", human_input_mode="NEVER", code_execution_config={ "work_dir": "financial_analysis_workspace", "use_docker": False, "last_n_messages": 3 }, system_message="""You are a financial data manager that coordinates the financial analysis process. Your role is to gather data, distribute it to the specialized analysts, and compile their insights. You can execute Python code to fetch and process financial data.""" ) def fetch_market_data(self, symbols: List[str], period: str = "1y") -> Dict[str, pd.DataFrame]: """ Fetch market data for specified symbols. Args: symbols: List of stock symbols period: Time period for data (e.g., '1d', '5d', '1mo', '3mo', '6mo', '1y', '2y', '5y', '10y', 'ytd', 'max') Returns: Dictionary of DataFrames with market data """ results = {} # Check cache first cache_key = f"{','.join(symbols)}_{period}" if cache_key in self.data_cache: return self.data_cache[cache_key] # Fetch data for each symbol for symbol in symbols: try: stock = yf.Ticker(symbol) hist = stock.history(period=period) if not hist.empty: results[symbol] = hist # Calculate returns hist['Daily_Return'] = hist['Close'].pct_change() hist['Cumulative_Return'] = (1 + hist['Daily_Return']).cumprod() - 1 # Calculate volatility (20-day rolling standard deviation of returns) hist['Volatility_20d'] = hist['Daily_Return'].rolling(window=20).std() # Add some basic technical indicators # 20-day and 50-day moving averages hist['MA20'] = hist['Close'].rolling(window=20).mean() hist['MA50'] = hist['Close'].rolling(window=50).mean() # Relative Strength Index (RSI) delta = hist['Close'].diff() gain = (delta.where(delta > 0, 0)).rolling(window=14).mean() loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean() rs = gain / loss hist['RSI'] = 100 - (100 / (1 + rs)) results[symbol] = hist except Exception as e: print(f"Error fetching data for {symbol}: {e}") # Store in cache self.data_cache[cache_key] = results return results def fetch_fundamental_data(self, symbols: List[str]) -> Dict[str, Dict[str, Any]]: """ Fetch fundamental data for specified symbols. Args: symbols: List of stock symbols Returns: Dictionary of fundamental data by symbol """ results = {} for symbol in symbols: try: stock = yf.Ticker(symbol) # Get key statistics info = stock.info # Get financial data try: income_stmt = stock.income_stmt balance_sheet = stock.balance_sheet cash_flow = stock.cashflow financials = { "income_statement": income_stmt.to_dict() if not income_stmt.empty else {}, "balance_sheet": balance_sheet.to_dict() if not balance_sheet.empty else {}, "cash_flow": cash_flow.to_dict() if not cash_flow.empty else {} } except: financials = {} # Compile results results[symbol] = { "info": info, "financials": financials } except Exception as e: print(f"Error fetching fundamental data for {symbol}: {e}") return results def fetch_news_data(self, symbols: List[str], days: int = 7) -> Dict[str, List[Dict[str, Any]]]: """ Fetch news articles for specified symbols. Args: symbols: List of stock symbols days: Number of days to look back Returns: Dictionary of news articles by symbol """ results = {} for symbol in symbols: try: # Format date range end_date = datetime.datetime.now() start_date = end_date - datetime.timedelta(days=days) # Get company name for better search results company_name = "" try: stock = yf.Ticker(symbol) company_name = stock.info.get("shortName", symbol) except: company_name = symbol # Construct query query = f"{company_name} OR {symbol} stock" # Fetch news from NewsAPI url = (f"https://newsapi.org/v2/everything?" f"q={query}&" f"from={start_date.strftime('%Y-%m-%d')}&" f"to={end_date.strftime('%Y-%m-%d')}&" f"language=en&" f"sortBy=relevancy&" f"pageSize=10&" f"apiKey={NEWS_API_KEY}") response = requests.get(url) if response.status_code == 200: news_data = response.json() articles = news_data.get("articles", []) # Process articles processed_articles = [] for article in articles: processed_articles.append({ "title": article.get("title", ""), "source": article.get("source", {}).get("name", ""), "published_at": article.get("publishedAt", ""), "url": article.get("url", ""), "description": article.get("description", "") }) results[symbol] = processed_articles else: print(f"Error fetching news for {symbol}: {response.status_code}") results[symbol] = [] except Exception as e: print(f"Error fetching news for {symbol}: {e}") results[symbol] = [] return results def calculate_technical_indicators(self, symbol: str, indicators: List[str], period: str = "1y") -> Dict[str, Any]: """ Calculate technical indicators for a stock. Args: symbol: Stock symbol indicators: List of indicators to calculate period: Time period for data Returns: Dictionary of technical indicators """ try: # Fetch data data = self.fetch_market_data([symbol], period).get(symbol) if data is None or data.empty: return {"error": f"No data available for {symbol}"} results = {} for indicator in indicators: indicator = indicator.lower() # Moving Averages if indicator.startswith("ma") or indicator.startswith("sma"): try: # Extract window size from indicator name (e.g., "ma20" -> 20) window = int(indicator[2:]) if indicator.startswith("ma") else int(indicator[3:]) data[f'MA{window}'] = data['Close'].rolling(window=window).mean() # Get the most recent value latest_value = data[f'MA{window}'].iloc[-1] results[f'MA{window}'] = latest_value except: results[f'{indicator}'] = None # Exponential Moving Average elif indicator.startswith("ema"): try: window = int(indicator[3:]) data[f'EMA{window}'] = data['Close'].ewm(span=window, adjust=False).mean() latest_value = data[f'EMA{window}'].iloc[-1] results[f'EMA{window}'] = latest_value except: results[f'{indicator}'] = None # RSI - Relative Strength Index elif indicator == "rsi": try: delta = data['Close'].diff() gain = (delta.where(delta > 0, 0)).rolling(window=14).mean() loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean() rs = gain / loss data['RSI'] = 100 - (100 / (1 + rs)) latest_value = data['RSI'].iloc[-1] results['RSI'] = latest_value except: results['RSI'] = None # MACD - Moving Average Convergence Divergence elif indicator == "macd": try: exp1 = data['Close'].ewm(span=12, adjust=False).mean() exp2 = data['Close'].ewm(span=26, adjust=False).mean() data['MACD'] = exp1 - exp2 data['Signal_Line'] = data['MACD'].ewm(span=9, adjust=False).mean() data['MACD_Histogram'] = data['MACD'] - data['Signal_Line'] results['MACD'] = { 'MACD_Line': data['MACD'].iloc[-1], 'Signal_Line': data['Signal_Line'].iloc[-1], 'Histogram': data['MACD_Histogram'].iloc[-1] } except: results['MACD'] = None # Bollinger Bands elif indicator == "bollinger" or indicator == "bb": try: window = 20 data['MA20'] = data['Close'].rolling(window=window).mean() data['BB_Upper'] = data['MA20'] + (data['Close'].rolling(window=window).std() * 2) data['BB_Lower'] = data['MA20'] - (data['Close'].rolling(window=window).std() * 2) results['Bollinger_Bands'] = { 'Upper': data['BB_Upper'].iloc[-1], 'Middle': data['MA20'].iloc[-1], 'Lower': data['BB_Lower'].iloc[-1], 'Width': (data['BB_Upper'].iloc[-1] - data['BB_Lower'].iloc[-1]) / data['MA20'].iloc[-1] } except: results['Bollinger_Bands'] = None # Average True Range (ATR) elif indicator == "atr": try: high_low = data['High'] - data['Low'] high_close = (data['High'] - data['Close'].shift()).abs() low_close = (data['Low'] - data['Close'].shift()).abs() ranges = pd.concat([high_low, high_close, low_close], axis=1) true_range = ranges.max(axis=1) data['ATR'] = true_range.rolling(14).mean() latest_value = data['ATR'].iloc[-1] results['ATR'] = latest_value except: results['ATR'] = None # Volume-Weighted Average Price (VWAP) elif indicator == "vwap": try: typical_price = (data['High'] + data['Low'] + data['Close']) / 3 data['VWAP'] = (typical_price * data['Volume']).cumsum() / data['Volume'].cumsum() latest_value = data['VWAP'].iloc[-1] results['VWAP'] = latest_value except: results['VWAP'] = None # Add the current price for reference results['Current_Price'] = data['Close'].iloc[-1] return results except Exception as e: return {"error": str(e)} def analyze_stock( self, symbol: str, period: str = "1y", include_news: bool = True, include_fundamentals: bool = True, report_type: str = "standard" ) -> Dict[str, Any]: """ Perform a comprehensive analysis of a stock. Args: symbol: Stock symbol to analyze period: Time period for analysis include_news: Whether to include news analysis include_fundamentals: Whether to include fundamental analysis report_type: Type of report (standard, detailed, executive) Returns: Dictionary with comprehensive analysis """ # Gather data market_data = self.fetch_market_data([symbol], period) # Create group chat for agent collaboration groupchat = autogen.GroupChat( agents=[ self.user_proxy, self.market_analyst, self.technical_analyst, self.fundamental_analyst, self.sentiment_analyst, self.portfolio_advisor, self.report_writer ], messages=[], max_round=12 ) manager = autogen.GroupChatManager(groupchat=groupchat) # Prepare market data for agents if symbol in market_data: df = market_data[symbol] # Create price summary price_start = df['Close'].iloc[0] price_end = df['Close'].iloc[-1] price_change = price_end - price_start price_change_pct = (price_change / price_start) * 100 price_high = df['High'].max() price_low = df['Low'].min() price_summary = f""" Stock: {symbol} Period: {period} Starting Price: ${price_start:.2f} Current Price: ${price_end:.2f} Price Change: ${price_change:.2f} ({price_change_pct:.2f}%) Period High: ${price_high:.2f} Period Low: ${price_low:.2f} Trading Volume (Avg): {df['Volume'].mean():.0f} """ # Calculate key technical indicators tech_indicators = self.calculate_technical_indicators( symbol, ["ma20", "ma50", "ma200", "rsi", "macd", "bollinger"], period ) # Convert indicators to string format indicators_str = json.dumps(tech_indicators, indent=2) # Get fundamental data if requested fundamentals_str = "" if include_fundamentals: fundamental_data = self.fetch_fundamental_data([symbol]) if symbol in fundamental_data: # Extract key metrics info = fundamental_data[symbol].get("info", {}) fundamentals_str = f""" Market Cap: {info.get('marketCap', 'N/A')} P/E Ratio: {info.get('trailingPE', 'N/A')} EPS: {info.get('trailingEps', 'N/A')} Beta: {info.get('beta', 'N/A')} 52 Week High: {info.get('fiftyTwoWeekHigh', 'N/A')} 52 Week Low: {info.get('fiftyTwoWeekLow', 'N/A')} Dividend Yield: {info.get('dividendYield', 'N/A')} Industry: {info.get('industry', 'N/A')} Sector: {info.get('sector', 'N/A')} """ # Get news data if requested news_str = "" if include_news: news_data = self.fetch_news_data([symbol], days=15) if symbol in news_data and news_data[symbol]: news_str = "Recent News:\n" for i, article in enumerate(news_data[symbol][:5], 1): news_str += f""" {i}. {article.get('title', 'N/A')} Source: {article.get('source', 'N/A')} Date: {article.get('published_at', 'N/A')} Summary: {article.get('description', 'N/A')} """ # Generate initial prompt for the agent team analysis_prompt = f""" Please analyze the following stock: {symbol} TIME PERIOD: {period} PRICE SUMMARY: {price_summary} TECHNICAL INDICATORS: {indicators_str} """ if fundamentals_str: analysis_prompt += f""" FUNDAMENTAL DATA: {fundamentals_str} """ if news_str: analysis_prompt += f""" NEWS: {news_str} """ analysis_prompt += f""" Please provide a {report_type} analysis report that includes: 1. Technical Analysis - Key patterns, indicators, and potential price targets 2. Market Context - How the stock fits in the broader market environment {"3. Fundamental Analysis - Company financial health and valuation" if include_fundamentals else ""} {"4. News Sentiment Analysis - Impact of recent news" if include_news else ""} 5. Investment Recommendation - Clear buy/hold/sell guidance with time horizon 6. Risk Assessment - Key risks and considerations The MarketAnalyst should begin by providing market context. The TechnicalAnalyst should analyze price patterns and indicators. The FundamentalAnalyst should assess the company's financial health and valuation. The SentimentAnalyst should evaluate recent news and sentiment. The PortfolioAdvisor should synthesize these analyses into investment recommendations. Finally, the ReportWriter should compile a well-structured professional report. """ # Start the group chat result = self.user_proxy.initiate_chat( manager, message=analysis_prompt ) # Extract the final report final_report = None for message in reversed(self.user_proxy.chat_history): if message['role'] == 'assistant' and 'ReportWriter' in message.get('name', ''): final_report = message['content'] break if not final_report: # Use the last substantial response if no clear report for message in reversed(self.user_proxy.chat_history): if message['role'] == 'assistant' and len(message['content']) > 500: final_report = message['content'] break return { "symbol": symbol, "analysis_date": datetime.datetime.now().strftime("%Y-%m-%d"), "period": period, "report_type": report_type, "price_data": { "current_price": price_end, "price_change": price_change, "price_change_pct": price_change_pct, "period_high": price_high, "period_low": price_low }, "technical_indicators": tech_indicators, "report": final_report } else: return {"error": f"No data available for {symbol}"} def compare_stocks( self, symbols: List[str], period: str = "1y", report_type: str = "standard" ) -> Dict[str, Any]: """ Compare multiple stocks and provide analysis. Args: symbols: List of stock symbols to compare period: Time period for analysis report_type: Type of report (standard, detailed, executive) Returns: Dictionary with comparative analysis """ if len(symbols) < 2: return {"error": "Please provide at least two symbols for comparison"} # Gather data for all symbols market_data = self.fetch_market_data(symbols, period) # Create group chat for agent collaboration groupchat = autogen.GroupChat( agents=[ self.user_proxy, self.market_analyst, self.technical_analyst, self.fundamental_analyst, self.portfolio_advisor, self.report_writer ], messages=[], max_round=12 ) manager = autogen.GroupChatManager(groupchat=groupchat) # Prepare comparative data comparison_data = {} price_performance = {} missing_data = [] for symbol in symbols: if symbol in market_data and not market_data[symbol].empty: df = market_data[symbol] # Calculate performance metrics price_start = df['Close'].iloc[0] price_end = df['Close'].iloc[-1] price_change_pct = ((price_end / price_start) - 1) * 100 # Calculate volatility (standard deviation of returns) volatility = df['Daily_Return'].std() * 100 # Multiply by 100 for percentage # Calculate max drawdown rolling_max = df['Close'].cummax() drawdown = (df['Close'] / rolling_max - 1) * 100 max_drawdown = drawdown.min() # Normalize price series (starting at 100) normalized_price = (df['Close'] / df['Close'].iloc[0]) * 100 # Store metrics price_performance[symbol] = { "price_change_pct": price_change_pct, "current_price": price_end, "volatility": volatility, "max_drawdown": max_drawdown, "normalized_prices": normalized_price.tolist() } # Calculate key technical indicators tech_indicators = self.calculate_technical_indicators( symbol, ["ma50", "rsi", "macd"], period ) # Add to comparison data comparison_data[symbol] = { "performance": price_performance[symbol], "technical_indicators": tech_indicators } else: missing_data.append(symbol) # Generate comparative analysis prompt if comparison_data: # Sort symbols by performance sorted_symbols = sorted( comparison_data.keys(), key=lambda x: comparison_data[x]["performance"]["price_change_pct"], reverse=True ) # Create performance table performance_table = "Symbol | Price Change (%) | Volatility (%) | Max Drawdown (%)\n" performance_table += "-------|------------------|----------------|----------------\n" for symbol in sorted_symbols: perf = comparison_data[symbol]["performance"] performance_table += f"{symbol} | {perf['price_change_pct']:.2f}% | {perf['volatility']:.2f}% | {perf['max_drawdown']:.2f}%\n" # Create comparative prompt comparison_prompt = f""" Please perform a comparative analysis of the following stocks: {', '.join(symbols)} TIME PERIOD: {period} PERFORMANCE COMPARISON: {performance_table} TECHNICAL INDICATORS SUMMARY: """ # Add technical indicators for symbol in sorted_symbols: tech = comparison_data[symbol]["technical_indicators"] comparison_prompt += f"\n{symbol} Indicators:\n" comparison_prompt += json.dumps(tech, indent=2) + "\n" if missing_data: comparison_prompt += f"\nNOTE: Could not fetch data for these symbols: {', '.join(missing_data)}\n" comparison_prompt += f""" Please provide a {report_type} comparative analysis report that includes: 1. Performance Comparison - Compare the stocks' performance during the period 2. Relative Strength Analysis - Which stocks show relative strength and weakness 3. Correlation Analysis - How these stocks move in relation to each other 4. Technical Position - Compare the technical position of each stock 5. Ranked Recommendations - Rank the stocks from most to least favorable 6. Portfolio Considerations - How these stocks might work together in a portfolio The MarketAnalyst should analyze the market context and relative performance. The TechnicalAnalyst should compare technical positions. The FundamentalAnalyst should provide comparative fundamental context if relevant. The PortfolioAdvisor should rank the stocks and provide portfolio recommendations. Finally, the ReportWriter should compile a well-structured comparative report. """ # Start the group chat result = self.user_proxy.initiate_chat( manager, message=comparison_prompt ) # Extract the final report final_report = None for message in reversed(self.user_proxy.chat_history): if message['role'] == 'assistant' and 'ReportWriter' in message.get('name', ''): final_report = message['content'] break if not final_report: # Use the last substantial response if no clear report for message in reversed(self.user_proxy.chat_history): if message['role'] == 'assistant' and len(message['content']) > 500: final_report = message['content'] break return { "symbols": symbols, "analysis_date": datetime.datetime.now().strftime("%Y-%m-%d"), "period": period, "report_type": report_type, "performance_comparison": {symbol: comparison_data[symbol]["performance"] for symbol in comparison_data}, "missing_data": missing_data, "report": final_report } else: return {"error": "Could not fetch data for any of the provided symbols"} # Example usage if __name__ == "__main__": # Create the financial analyst agent financial_analyst = FinancialAnalystAgent() # Analyze a single stock analysis = financial_analyst.analyze_stock( symbol="MSFT", period="1y", include_news=True, report_type="standard" ) print("=== Single Stock Analysis ===") print(f"Symbol: {analysis['symbol']}") print(f"Current Price: ${analysis['price_data']['current_price']:.2f}") print(f"Price Change: {analysis['price_data']['price_change_pct']:.2f}%") print("\nReport Excerpt:") print(analysis['report'][:500] + "...\n") # Compare multiple stocks comparison = financial_analyst.compare_stocks( symbols=["AAPL", "MSFT", "GOOGL"], period="6mo", report_type="standard" ) print("=== Stock Comparison ===") print(f"Symbols: {comparison['symbols']}") print("\nPerformance Comparison:") for symbol, perf in comparison['performance_comparison'].items(): print(f"{symbol}: {perf['price_change_pct']:.2f}%") print("\nReport Excerpt:") print(comparison['report'][:500] + "...") ``` **Usage Example:** ```python from financial_analyst_agent import FinancialAnalystAgent # Initialize agent analyst = FinancialAnalystAgent() # Analyze a stock report = analyst.analyze_stock( symbol="TSLA", period="1y", include_news=True, include_fundamentals=True, report_type="detailed" ) print(f"Analysis of {report['symbol']} completed.") print(f"Current price: ${report['price_data']['current_price']:.2f}") print(f"Price change: {report['price_data']['price_change_pct']:.2f}%") print("\nReport Highlights:") print(report['report']) # Compare multiple stocks comparison = analyst.compare_stocks( symbols=["AAPL", "MSFT", "GOOGL", "AMZN", "META"], period="6mo", report_type="standard" ) print("\nComparative Analysis:") print(comparison['report']) ``` This AI Financial Analyst agent template demonstrates key enterprise patterns: 1. **Agent Specialization**: Different agents focus on specific analysis types (technical, fundamental, news) 2. **Data Pipeline Integration**: The system integrates multiple external data sources 3. **Collaborative Analysis**: Agents work together via a group chat to produce a comprehensive analysis 4. **Flexible Report Generation**: Different report types for various user needs 5. **Caching Strategy**: Data caching to improve performance and reduce redundant API calls ### AI-Powered Cybersecurity Incident Response (Threat Detection & Remediation) This AI agent template is designed to help cybersecurity teams detect, analyze, and respond to security incidents. It combines threat intelligence, log analysis, and remediation guidance to provide comprehensive cybersecurity incident response. **Core Capabilities:** - Automated log analysis and threat detection - Incident classification and severity assessment - Threat intelligence correlation - Forensic investigation support - Guided remediation steps - Documentation generation for compliance **Architecture:** ![Cybersecurity Incident Response Agent Architecture](https://i.imgur.com/wGHFdTd.png) **Implementation Example:** ```python # cybersecurity_incident_response_agent.py import os import json import datetime import ipaddress import hashlib import re import uuid import pandas as pd import numpy as np import autogen import requests from typing import Dict, List, Any, Optional, Union, Tuple # Configure API keys and settings OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "your-api-key") VIRUSTOTAL_API_KEY = os.environ.get("VIRUSTOTAL_API_KEY", "your-api-key") ABUSEIPDB_API_KEY = os.environ.get("ABUSEIPDB_API_KEY", "your-api-key") # Initialize OpenAI client import openai client = openai.OpenAI(api_key=OPENAI_API_KEY) class CybersecurityIncidentResponseAgent: """ An AI-powered cybersecurity incident response agent that helps detect, analyze, and respond to security incidents. """ def __init__(self, config=None): """ Initialize the Cybersecurity Incident Response Agent. Args: config: Optional configuration dictionary """ self.config = config or {} # Set up agent configuration self.llm_config = { "config_list": [{"model": "gpt-4-turbo", "api_key": OPENAI_API_KEY}], "temperature": 0.2, "timeout": 120 } # Create the agent team self._create_agent_team() # Threat intelligence cache self.threat_intel_cache = {} def _create_agent_team(self): """Create the team of specialized agents for cybersecurity incident response.""" # 1. Threat Detector - Specialized in identifying threats from logs and data self.threat_detector = autogen.AssistantAgent( name="ThreatDetector", system_message="""You are an expert threat detection analyst who specializes in identifying security threats from logs, network traffic, and system data. Your responsibilities: 1. Analyze raw logs and identify suspicious patterns or anomalies 2. Recognize indicators of compromise (IoCs) such as unusual IP addresses, file hashes, or user behaviors 3. Detect potential attack techniques and map them to the MITRE ATT&CK framework 4. Identify false positives and prioritize genuine security concerns 5. Alert on critical security issues with appropriate context Be thorough and methodical in your analysis. Look for subtle patterns that might indicate sophisticated attacks. Avoid jumping to conclusions without sufficient evidence. Always provide specific indicators and explain why they are suspicious.""", llm_config=self.llm_config ) # 2. Forensic Investigator - Specialized in detailed investigation self.forensic_investigator = autogen.AssistantAgent( name="ForensicInvestigator", system_message="""You are an expert digital forensic investigator who specializes in analyzing security incidents to determine their scope, impact, and attribution. Your responsibilities: 1. Analyze evidence thoroughly to reconstruct the incident timeline 2. Identify the attack vectors and methods used by threat actors 3. Determine the scope of compromise (affected systems, data, users) 4. Look for persistence mechanisms and backdoors 5. Gather indicators that can help with attribution Be methodical and detail-oriented in your investigation. Document your findings clearly, including timestamps and specific technical details. Distinguish between confirmed facts, strong evidence, and speculation. Consider alternative explanations and test your hypotheses against the evidence.""", llm_config=self.llm_config ) # 3. Threat Intelligence Analyst - Specialized in external threat intelligence self.threat_intel_analyst = autogen.AssistantAgent( name="ThreatIntelAnalyst", system_message="""You are an expert threat intelligence analyst who specializes in researching and analyzing cyber threats, threat actors, and their tactics, techniques, and procedures (TTPs). Your responsibilities: 1. Research indicators of compromise against threat intelligence sources 2. Identify known threat actors or malware associated with the incident 3. Provide context on the tactics, techniques, and procedures used 4. Assess the potential goals and motivation of the attackers 5. Determine if the attack is targeted or opportunistic Provide relevant, actionable intelligence that helps understand the threat. Link findings to the MITRE ATT&CK framework when possible. Distinguish between high and low-confidence assessments. Consider the reliability of intelligence sources. Focus on information that is directly relevant to the current incident.""", llm_config=self.llm_config ) # 4. Incident Responder - Specialized in containment and remediation self.incident_responder = autogen.AssistantAgent( name="IncidentResponder", system_message="""You are an expert incident responder who specializes in containing security incidents, removing threats, and restoring systems to normal operation. Your responsibilities: 1. Provide immediate containment actions to limit the impact of the incident 2. Develop detailed remediation plans to remove the threat 3. Recommend recovery steps to restore affected systems 4. Suggest security improvements to prevent similar incidents 5. Prioritize actions based on risk and business impact Your recommendations should be specific, actionable, and prioritized. Consider the potential impact of response actions on business operations. Provide both immediate tactical responses and strategic improvements. Always consider the order of operations to avoid alerting attackers or destroying evidence. Tailor your response to the specific environment and incident details.""", llm_config=self.llm_config ) # 5. Documentation Specialist - Specialized in creating comprehensive incident documentation self.documentation_specialist = autogen.AssistantAgent( name="DocumentationSpecialist", system_message="""You are an expert in creating comprehensive cybersecurity incident documentation that is clear, thorough, and suitable for multiple audiences including technical teams, management, and compliance requirements. Your responsibilities: 1. Compile incident details into structured documentation 2. Create executive summaries for management and technical details for IT teams 3. Ensure documentation satisfies compliance and regulatory requirements 4. Include all relevant timeline information, affected systems, and remediation steps 5. Document lessons learned and recommended improvements Create documentation that is well-organized, precise, and actionable. Use clear sections with appropriate headers. Include all relevant technical details while making executive summaries accessible to non-technical audiences. Ensure all claims are supported by evidence. Include metadata such as incident IDs, dates, and classification.""", llm_config=self.llm_config ) # User proxy agent for orchestrating the workflow self.user_proxy = autogen.UserProxyAgent( name="SecurityAnalyst", human_input_mode="NEVER", code_execution_config={ "work_dir": "security_workspace", "use_docker": False }, system_message="""You are a security analyst coordinating the incident response process. Your role is to gather data, distribute it to the specialized analysts, and compile their insights. You can execute Python code to analyze security data and fetch threat intelligence.""" ) def _hash_file(self, file_path: str) -> str: """ Compute SHA-256 hash of a file. Args: file_path: Path to the file Returns: SHA-256 hash as a hexadecimal string """ try: sha256_hash = hashlib.sha256() with open(file_path, "rb") as f: for byte_block in iter(lambda: f.read(4096), b""): sha256_hash.update(byte_block) return sha256_hash.hexdigest() except Exception as e: print(f"Error hashing file {file_path}: {e}") return None def parse_log_data(self, log_data: str, log_type: str = "generic") -> List[Dict[str, Any]]: """ Parse raw log data into structured format based on log type. Args: log_data: Raw log data as string log_type: Type of log (generic, windows, linux, firewall, web, etc.) Returns: List of dictionaries containing structured log entries """ structured_logs = [] # Split log data into lines log_lines = log_data.strip().split('\n') if log_type.lower() == "windows_event": # Parse Windows Event logs current_event = {} for line in log_lines: line = line.strip() if line.startswith("Log Name:"): if current_event: structured_logs.append(current_event) current_event = {"Log Name": line.split("Log Name:")[1].strip()} elif ":" in line and current_event: key, value = line.split(":", 1) current_event[key.strip()] = value.strip() # Add the last event if current_event: structured_logs.append(current_event) elif log_type.lower() == "syslog": # Parse syslog format for line in log_lines: if not line.strip(): continue try: # Basic syslog pattern: []: match = re.match(r"(\w+\s+\d+\s+\d+:\d+:\d+)\s+(\S+)\s+([^:]+):\s+(.*)", line) if match: timestamp, hostname, process, message = match.groups() # Extract PID if present pid_match = re.search(r"\[(\d+)\]", process) pid = pid_match.group(1) if pid_match else None process = re.sub(r"\[\d+\]", "", process).strip() structured_logs.append({ "timestamp": timestamp, "hostname": hostname, "process": process, "pid": pid, "message": message, "raw_log": line }) else: # If pattern doesn't match, store as raw log structured_logs.append({"raw_log": line}) except Exception as e: structured_logs.append({"raw_log": line, "parse_error": str(e)}) elif log_type.lower() == "apache" or log_type.lower() == "nginx": # Parse common web server log format for line in log_lines: if not line.strip(): continue try: # Common Log Format: [