Files
fuzzforge_ai/ai/agents/task_agent/litellm_agent/state.py
T
Songbird 11f294abdb Add AI module with A2A wrapper and task agent
- Disable FuzzForge MCP connection (no Prefect backend)
- Add a2a_wrapper module for programmatic A2A agent tasks
- Add task_agent (LiteLLM A2A agent) on port 10900
- Create volumes/env/ for centralized Docker config
- Update docker-compose.yml with task-agent service
- Remove workflow_automation_skill from agent card
2025-10-14 13:05:35 +02:00

87 lines
2.8 KiB
Python

"""Session state utilities for the LiteLLM hot-swap agent."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Mapping, MutableMapping, Optional
from .config import (
DEFAULT_MODEL,
DEFAULT_PROVIDER,
STATE_MODEL_KEY,
STATE_PROMPT_KEY,
STATE_PROVIDER_KEY,
)
@dataclass(slots=True)
class HotSwapState:
"""Lightweight view of the hot-swap session state."""
model: str = DEFAULT_MODEL
provider: Optional[str] = None
prompt: Optional[str] = None
@classmethod
def from_mapping(cls, mapping: Optional[Mapping[str, Any]]) -> "HotSwapState":
if not mapping:
return cls()
raw_model = mapping.get(STATE_MODEL_KEY, DEFAULT_MODEL)
raw_provider = mapping.get(STATE_PROVIDER_KEY)
raw_prompt = mapping.get(STATE_PROMPT_KEY)
model = raw_model.strip() if isinstance(raw_model, str) else DEFAULT_MODEL
provider = raw_provider.strip() if isinstance(raw_provider, str) else None
if not provider and DEFAULT_PROVIDER:
provider = DEFAULT_PROVIDER.strip() or None
prompt = raw_prompt.strip() if isinstance(raw_prompt, str) else None
return cls(
model=model or DEFAULT_MODEL,
provider=provider or None,
prompt=prompt or None,
)
def persist(self, store: MutableMapping[str, object]) -> None:
store[STATE_MODEL_KEY] = self.model
if self.provider:
store[STATE_PROVIDER_KEY] = self.provider
else:
store[STATE_PROVIDER_KEY] = None
store[STATE_PROMPT_KEY] = self.prompt
def describe(self) -> str:
prompt_value = self.prompt if self.prompt else "(default prompt)"
provider_value = self.provider if self.provider else "(default provider)"
return (
"📊 Current Configuration\n"
"━━━━━━━━━━━━━━━━━━━━━━\n"
f"Model: {self.model}\n"
f"Provider: {provider_value}\n"
f"System Prompt: {prompt_value}\n"
"━━━━━━━━━━━━━━━━━━━━━━"
)
def instantiate_llm(self):
"""Create a LiteLlm instance for the current state."""
from google.adk.models.lite_llm import LiteLlm # Lazy import to avoid cycle
kwargs = {"model": self.model}
if self.provider:
kwargs["custom_llm_provider"] = self.provider
return LiteLlm(**kwargs)
@property
def display_model(self) -> str:
if self.provider:
return f"{self.provider}/{self.model}"
return self.model
def apply_state_to_agent(invocation_context, state: HotSwapState) -> None:
"""Update the provided agent with a LiteLLM instance matching state."""
agent = invocation_context.agent
agent.model = state.instantiate_llm()