mirror of
https://github.com/FuzzingLabs/fuzzforge_ai.git
synced 2026-05-23 15:39:41 +02:00
11f294abdb
- 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
87 lines
2.8 KiB
Python
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()
|