51 lines
1.7 KiB
Python
51 lines
1.7 KiB
Python
from itertools import tee
|
|
from typing import Optional, Dict, List
|
|
|
|
from .docker_module import DockerModule
|
|
from .llm_engine import LlmEngine
|
|
from .context_template import generate_context
|
|
from .util import get_valid_root_elements, split_response
|
|
|
|
class AgentCore:
|
|
"""
|
|
Core orchestration class for SIA that manages the interaction between
|
|
different modules and runs the main agent loop.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
system_prompt: str,
|
|
action_schema: str,
|
|
docker_module: DockerModule,
|
|
llm_engine: LlmEngine
|
|
):
|
|
"""
|
|
Initialize the AgentCore with required components.
|
|
|
|
Args:
|
|
system_prompt: System prompt to use for the LLM
|
|
action_schema_path: Path to the XML schema defining valid actions
|
|
docker_module: DockerModule instance
|
|
llm_engine: LLmEngine instance
|
|
"""
|
|
self.system_prompt = system_prompt
|
|
self.action_schema = action_schema
|
|
self.docker_module = docker_module
|
|
self.llm_engine = llm_engine
|
|
|
|
self.valid_elements = get_valid_root_elements(self.action_schema)
|
|
|
|
def run_iteration(self) -> None:
|
|
"""Run a single iteration of the main agent loop."""
|
|
containers = self.docker_module.get_all_container_statuses()
|
|
context = generate_context(containers)
|
|
tokens = self.llm_engine.infer(
|
|
self.system_prompt,
|
|
context
|
|
)
|
|
print_tokens, response_tokens = tee(tokens)
|
|
for token in print_tokens:
|
|
print(token, end="", flush=True)
|
|
response_tokens = ''.join(response_tokens)
|
|
result = split_response(response_tokens, self.valid_elements)
|
|
print(f"result: {result}") |