refactor: add agent refactoring

Author: pmutua

agent_manager/__init__.py

@@ -0,0 +1 @@
+from . import agent

agent_manager/agent.py

@@ -0,0 +1,81 @@
+import uuid
+import re
+import asyncio
+from dotenv import load_dotenv
+
+from google.adk.runners import Runner
+from google.adk.agents import Agent
+from google.adk.models.lite_llm import LiteLlm
+from google.adk.sessions import DatabaseSessionService
+from google.genai import types
+
+from .sub_agents.detector.agent import detector_agent
+from .sub_agents.planner.agent import planner_agent
+from .sub_agents.fixer.agent import fixer_agent
+from .sub_agents.notifier.agent import notifier_agent
+from .sub_agents.action_recommender.agent import action_recommender_agent
+from .config import MODEL, APP_NAME
+
+# --- Constants ---
+AZURE_OPENAI_MODEL_PATTERN = r"\bazure\S*"
+GEMINI_MODEL_PATTERN = r"\bgemini\S*"
+USER_ID = "chaospilot_user"
+SESSION_ID = str(uuid.uuid4())
+DB_URL = "sqlite:///./agent_Manager_data.db"
+
+# --- Model Selection ---
+if re.findall(AZURE_OPENAI_MODEL_PATTERN, MODEL):
+    model_instance = LiteLlm(model="azure/gpt-4o")
+elif re.findall(GEMINI_MODEL_PATTERN, MODEL):
+    model_instance = MODEL  # MODEL is assumed to be a valid Gemini string
+else:
+    raise ValueError(f"Unsupported model format: {MODEL}")
+
+# --- Root Agent Definition ---
+root_agent = Agent(
+    name="agent_manager",
+    model=model_instance,
+    description="Orchestration layer for the chaos engineering system.",
+    sub_agents=[
+        detector_agent,
+        planner_agent,
+        action_recommender_agent,
+        fixer_agent,
+        notifier_agent,
+    ],
+)
+
+# --- Session Service ---
+session_service = DatabaseSessionService(db_url=DB_URL)
+
+
+# --- Main Runner ---
+async def run():
+    session = await session_service.create_session(
+        app_name=APP_NAME,
+        user_id=USER_ID,
+        session_id=SESSION_ID,
+        state={"initialized_at_startup": True},
+    )
+
+    runner = Runner(
+        agent=root_agent, app_name=APP_NAME, session_service=session_service
+    )
+
+    new_message = types.Content(
+        role="user",
+        parts=[
+            types.Part(
+                text="Use the detector agent to: Analyze the latest logs for anomalies"
+            )
+        ],
+    )
+
+    async for event in runner.run_async(
+        user_id=USER_ID, session_id=SESSION_ID, new_message=new_message
+    ):
+        if event.is_final_response():
+            print("Final Response:", event.content.parts[0].text)
+
+    print("=========FINAL SESSION STATE==============")
+

agent_manager/config.py

@@ -0,0 +1,38 @@
+# agent.py
+
+import os
+import re
+from dotenv import load_dotenv
+from toolbox_core import ToolboxSyncClient
+
+# Load .env file
+load_dotenv()
+
+# --- App Configuration ---
+APP_NAME = os.getenv("APP_NAME", "agent_manager")
+VERSION = os.getenv("VERSION", "0.1.0")
+MODEL = os.getenv("MODEL")
+TOOLBOX_URL = os.getenv("TOOLBOX_URL", "http:localhost:5000")
+
+# --- Model Patterns ---
+GEMINI_MODEL_PATTERN = r"\bgemini-[\w-]+"
+AZURE_OPENAI_MODEL_PATTERN = r"\bazure\w*"
+
+# --- Keys/Secrets (optional until model matches) ---
+AZURE_API_KEY = None
+AZURE_API_BASE = None
+AZURE_API_VERSION = None
+GOOGLE_API_KEY = None
+
+toolbox = None
+
+toolbox = ToolboxSyncClient(TOOLBOX_URL)
+
+# --- Load API Keys based on model pattern ---
+if re.findall(GEMINI_MODEL_PATTERN, MODEL):
+    GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+
+elif re.findall(AZURE_OPENAI_MODEL_PATTERN, MODEL):
+    AZURE_API_KEY = os.getenv("AZURE_API_KEY")
+    AZURE_API_BASE = os.getenv("AZURE_API_BASE")
+    AZURE_API_VERSION = os.getenv("AZURE_API_VERSION")

agent_manager/models.py

@@ -0,0 +1,50 @@
+from sqlalchemy import Column, Integer, String, DateTime, Text, JSON
+from sqlalchemy.ext.declarative import declarative_base
+import datetime
+
+Base = declarative_base()
+
+
+class Feedback(Base):
+    __tablename__ = "feedback"
+    id = Column(Integer, primary_key=True)
+    agent_name = Column(String)
+    output_id = Column(String)
+    user_id = Column(String)
+    timestamp = Column(DateTime, default=datetime.datetime.utcnow)
+    feedback_type = Column(String)  # e.g., 'clarity', 'effectiveness'
+    feedback_text = Column(Text)
+    effectiveness_score = Column(Integer)  # 1-5 scale
+    context_json = Column(JSON)
+    tools_used = Column(JSON)  # List of tool names
+
+
+class AgentRegistry(Base):
+    __tablename__ = "agent_registry"
+    id = Column(Integer, primary_key=True)
+    name = Column(String, unique=True)
+    description = Column(Text)
+    default_toolset = Column(JSON)
+    capabilities = Column(JSON)
+    created_at = Column(DateTime, default=datetime.datetime.utcnow)
+
+
+class AuditLog(Base):
+    __tablename__ = "audit_log"
+    id = Column(Integer, primary_key=True)
+    event_type = Column(String)  # e.g., 'agent_invocation', 'tool_call', 'feedback'
+    agent_name = Column(String)
+    tool_name = Column(String, nullable=True)
+    user_id = Column(String, nullable=True)
+    details = Column(JSON)
+    timestamp = Column(DateTime, default=datetime.datetime.utcnow)
+
+
+class IncidentHistory(Base):
+    __tablename__ = "incident_history"
+    id = Column(Integer, primary_key=True)
+    incident_id = Column(String, unique=True)
+    summary_json = Column(JSON)
+    resolution_json = Column(JSON)
+    agent_outputs_json = Column(JSON)
+    created_at = Column(DateTime, default=datetime.datetime.utcnow)