Building Enterprise Agent Systems: Core Component Design and Optimization

from typing import Protocol, Dict
from jinja2 import Template

class PromptTemplate(Protocol):
    def render(self, **kwargs) -> str:
        pass

class JinjaPromptTemplate:
    def __init__(self, template_string: str):
        self.template = Template(template_string)

    def render(self, **kwargs) -> str:
        return self.template.render(**kwargs)

class PromptLibrary:
    def __init__(self):
        self.templates: Dict[str, PromptTemplate] = {}

    def register_template(self, name: str, template: PromptTemplate):
        self.templates[name] = template

    def get_template(self, name: str) -> PromptTemplate:
        return self.templates[name]

class PromptVersion:
    def __init__(self, version: str, template: str, metadata: dict):
        self.version = version
        self.template = template
        self.metadata = metadata
        self.test_cases = []

    def add_test_case(self, inputs: dict, expected_output: str):
        self.test_cases.append((inputs, expected_output))

    def validate(self) -> bool:
        template = JinjaPromptTemplate(self.template)
        for inputs, expected in self.test_cases:
            result = template.render(**inputs)
            if not self._validate_output(result, expected):
                return False
        return True

from typing import Any, List
from datetime import datetime

class MemoryEntry:
    def __init__(self, content: Any, importance: float):
        self.content = content
        self.importance = importance
        self.timestamp = datetime.now()
        self.access_count = 0

class MemoryLayer:
    def __init__(self, capacity: int):
        self.capacity = capacity
        self.memories: List[MemoryEntry] = []

    def add(self, entry: MemoryEntry):
        if len(self.memories) >= self.capacity:
            self._evict()
        self.memories.append(entry)

    def _evict(self):
        # Implement memory eviction strategy
        self.memories.sort(key=lambda x: x.importance * x.access_count)
        self.memories.pop(0)

class HierarchicalMemory:
    def __init__(self):
        self.working_memory = MemoryLayer(capacity=5)
        self.short_term = MemoryLayer(capacity=50)
        self.long_term = MemoryLayer(capacity=1000)

    def store(self, content: Any, importance: float):
        entry = MemoryEntry(content, importance)

        if importance > 0.8:
            self.working_memory.add(entry)
        elif importance > 0.5:
            self.short_term.add(entry)
        else:
            self.long_term.add(entry)

from typing import List, Tuple
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

class MemoryIndex:
    def __init__(self, embedding_model):
        self.embedding_model = embedding_model
        self.embeddings = []
        self.memories = []

    def add(self, memory: MemoryEntry):
        embedding = self.embedding_model.embed(memory.content)
        self.embeddings.append(embedding)
        self.memories.append(memory)

    def search(self, query: str, k: int = 5) -> List[Tuple[MemoryEntry, float]]:
        query_embedding = self.embedding_model.embed(query)
        similarities = cosine_similarity(
            [query_embedding], 
            self.embeddings
        )[0]

        top_k_indices = np.argsort(similarities)[-k:]

        return [
            (self.memories[i], similarities[i]) 
            for i in top_k_indices
        ]

from typing import List, Optional
from dataclasses import dataclass
import uuid

@dataclass
class ThoughtNode:
    content: str
    confidence: float
    supporting_evidence: List[str]

class ReasoningChain:
    def __init__(self):
        self.chain_id = str(uuid.uuid4())
        self.nodes: List[ThoughtNode] = []
        self.metadata = {}

    def add_thought(self, thought: ThoughtNode):
        self.nodes.append(thought)

    def get_path(self) -> List[str]:
        return [node.content for node in self.nodes]

    def get_confidence(self) -> float:
        if not self.nodes:
            return 0.0
        return sum(n.confidence for n in self.nodes) / len(self.nodes)

import logging
from opentelemetry import trace
from prometheus_client import Histogram

reasoning_time = Histogram(
    'reasoning_chain_duration_seconds',
    'Time spent in reasoning chain'
)

class ChainMonitor:
    def __init__(self):
        self.tracer = trace.get_tracer(__name__)

    def monitor_chain(self, chain: ReasoningChain):
        with self.tracer.start_as_current_span("reasoning_chain") as span:
            span.set_attribute("chain_id", chain.chain_id)

            with reasoning_time.time():
                for node in chain.nodes:
                    with self.tracer.start_span("thought") as thought_span:
                        thought_span.set_attribute(
                            "confidence", 
                            node.confidence
                        )
                        logging.info(
                            f"Thought: {node.content} "
                            f"(confidence: {node.confidence})"
                        )

from abc import ABC, abstractmethod
from typing import Generic, TypeVar

T = TypeVar('T')

class Component(ABC, Generic[T]):
    @abstractmethod
    def process(self, input_data: T) -> T:
        pass

class Pipeline:
    def __init__(self):
        self.components: List[Component] = []

    def add_component(self, component: Component):
        self.components.append(component)

    def process(self, input_data: Any) -> Any:
        result = input_data
        for component in self.components:
            result = component.process(result)
        return result

class ComponentRegistry:
    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance.components = {}
        return cls._instance

    def register(self, name: str, component: Component):
        self.components[name] = component

    def get(self, name: str) -> Optional[Component]:
        return self.components.get(name)

    def create_pipeline(self, component_names: List[str]) -> Pipeline:
        pipeline = Pipeline()
        for name in component_names:
            component = self.get(name)
            if component:
                pipeline.add_component(component)
        return pipeline

from dataclasses import dataclass
from typing import Dict
import time

@dataclass
class PerformanceMetrics:
    latency: float
    memory_usage: float
    token_count: int
    success_rate: float

class PerformanceMonitor:
    def __init__(self):
        self.metrics: Dict[str, List[PerformanceMetrics]] = {}

    def record_operation(
        self,
        operation_name: str,
        metrics: PerformanceMetrics
    ):
        if operation_name not in self.metrics:
            self.metrics[operation_name] = []
        self.metrics[operation_name].append(metrics)

    def get_average_metrics(
        self,
        operation_name: str
    ) -> Optional[PerformanceMetrics]:
        if operation_name not in self.metrics:
            return None

        metrics_list = self.metrics[operation_name]
        return PerformanceMetrics(
            latency=sum(m.latency for m in metrics_list) / len(metrics_list),
            memory_usage=sum(m.memory_usage for m in metrics_list) / len(metrics_list),
            token_count=sum(m.token_count for m in metrics_list) / len(metrics_list),
            success_rate=sum(m.success_rate for m in metrics_list) / len(metrics_list)
        )

class PerformanceOptimizer:
    def __init__(self, monitor: PerformanceMonitor):
        self.monitor = monitor
        self.thresholds = {
            'latency': 1.0,  # seconds
            'memory_usage': 512,  # MB
            'token_count': 1000,
            'success_rate': 0.95
        }

    def analyze_performance(self, operation_name: str) -> List[str]:
        metrics = self.monitor.get_average_metrics(operation_name)
        if not metrics:
            return []

        recommendations = []

        if metrics.latency > self.thresholds['latency']:
            recommendations.append(
                "Consider implementing caching or parallel processing"
            )

        if metrics.memory_usage > self.thresholds['memory_usage']:
            recommendations.append(
                "Optimize memory usage through batch processing"
            )

        if metrics.token_count > self.thresholds['token_count']:
            recommendations.append(
                "Implement prompt optimization to reduce token usage"
            )

        if metrics.success_rate < self.thresholds['success_rate']:
            recommendations.append(
                "Review error handling and implement retry mechanisms"
            )

        return recommendations