parent
82a2d8954b
commit
570f8a2020
@ -0,0 +1,92 @@
|
||||
from typing import Dict, List, Optional
|
||||
from pydantic import BaseModel
|
||||
import numpy as np
|
||||
from datetime import datetime
|
||||
|
||||
class UpdateMetadata(BaseModel):
|
||||
"""Metadata for vector DB updates"""
|
||||
topic: str
|
||||
importance: float
|
||||
timestamp: datetime
|
||||
embedding: Optional[List[float]] = None
|
||||
affected_areas: List[str] = []
|
||||
|
||||
class AgentProfile(BaseModel):
|
||||
"""Profile for agent notification preferences"""
|
||||
agent_id: str
|
||||
expertise_areas: List[str]
|
||||
importance_threshold: float = 0.5
|
||||
current_task_context: Optional[str] = None
|
||||
embedding: Optional[List[float]] = None
|
||||
|
||||
class NotificationManager:
|
||||
"""Manages selective notifications for vector DB updates"""
|
||||
|
||||
def __init__(self):
|
||||
self.agent_profiles: Dict[str, AgentProfile] = {}
|
||||
|
||||
def register_agent(self, profile: AgentProfile):
|
||||
"""Register an agent's notification preferences"""
|
||||
self.agent_profiles[profile.agent_id] = profile
|
||||
|
||||
def unregister_agent(self, agent_id: str):
|
||||
"""Remove an agent's notification preferences"""
|
||||
if agent_id in self.agent_profiles:
|
||||
del self.agent_profiles[agent_id]
|
||||
|
||||
def calculate_relevance(
|
||||
self,
|
||||
update_metadata: UpdateMetadata,
|
||||
agent_profile: AgentProfile
|
||||
) -> float:
|
||||
"""Calculate relevance score between update and agent"""
|
||||
# Topic/expertise overlap score
|
||||
topic_score = len(
|
||||
set(agent_profile.expertise_areas) &
|
||||
set(update_metadata.affected_areas)
|
||||
) / max(
|
||||
len(agent_profile.expertise_areas),
|
||||
len(update_metadata.affected_areas)
|
||||
)
|
||||
|
||||
# Embedding similarity if available
|
||||
embedding_score = 0.0
|
||||
if update_metadata.embedding and agent_profile.embedding:
|
||||
embedding_score = np.dot(
|
||||
update_metadata.embedding,
|
||||
agent_profile.embedding
|
||||
)
|
||||
|
||||
# Combine scores (can be tuned)
|
||||
relevance = 0.7 * topic_score + 0.3 * embedding_score
|
||||
|
||||
return relevance
|
||||
|
||||
def should_notify_agent(
|
||||
self,
|
||||
update_metadata: UpdateMetadata,
|
||||
agent_profile: AgentProfile
|
||||
) -> bool:
|
||||
"""Determine if an agent should be notified of an update"""
|
||||
# Check importance threshold
|
||||
if update_metadata.importance < agent_profile.importance_threshold:
|
||||
return False
|
||||
|
||||
# Calculate relevance
|
||||
relevance = self.calculate_relevance(update_metadata, agent_profile)
|
||||
|
||||
# Notification threshold (can be tuned)
|
||||
return relevance > 0.5
|
||||
|
||||
def get_agents_to_notify(
|
||||
self,
|
||||
update_metadata: UpdateMetadata
|
||||
) -> List[str]:
|
||||
"""Get list of agent IDs that should be notified of an update"""
|
||||
agents_to_notify = []
|
||||
|
||||
for agent_id, profile in self.agent_profiles.items():
|
||||
if self.should_notify_agent(update_metadata, profile):
|
||||
agents_to_notify.append(agent_id)
|
||||
|
||||
return agents_to_notify
|
@ -0,0 +1,32 @@
|
||||
from typing import List
|
||||
from .notification_manager import NotificationManager, UpdateMetadata
|
||||
from .agent import Agent
|
||||
|
||||
class Orchestrator:
|
||||
def __init__(self):
|
||||
self.notification_manager = NotificationManager()
|
||||
self.agents: List[Agent] = []
|
||||
|
||||
def register_agent(self, agent: Agent):
|
||||
"""Register an agent with the orchestrator"""
|
||||
self.agents.append(agent)
|
||||
self.notification_manager.register_agent(agent.notification_profile)
|
||||
|
||||
def handle_vector_db_update(self, update_metadata: UpdateMetadata):
|
||||
"""Handle a vector DB update and notify relevant agents"""
|
||||
# Get list of agents to notify
|
||||
agents_to_notify = self.notification_manager.get_agents_to_notify(
|
||||
update_metadata
|
||||
)
|
||||
|
||||
# Notify relevant agents
|
||||
for agent in self.agents:
|
||||
if agent.agent_name in agents_to_notify:
|
||||
agent.handle_vector_db_update(update_metadata)
|
||||
|
||||
def update_agent_task_context(self, agent_name: str, task_context: str):
|
||||
"""Update an agent's current task context"""
|
||||
if agent_name in self.notification_manager.agent_profiles:
|
||||
self.notification_manager.agent_profiles[
|
||||
agent_name
|
||||
].current_task_context = task_context
|
Loading…
Reference in new issue