262 lines
10 KiB
Python
262 lines
10 KiB
Python
"""Companion worker agents — Logistics, Caravan, and Scout.
|
||
|
||
Companions are the lowest tier — fast, specialized, single-purpose workers.
|
||
Each companion listens to its :class:`TaskMessage` queue, executes the
|
||
requested primitive against GABS, and emits a :class:`ResultMessage`.
|
||
|
||
Model: Qwen3:8b (or smaller) — sub-2-second response times.
|
||
Frequency: event-driven (triggered by vassal task messages).
|
||
|
||
Primitive vocabulary per companion:
|
||
Logistics: recruit_troop, buy_supplies, rest_party, sell_prisoners, upgrade_troops, build_project
|
||
Caravan: assess_prices, buy_goods, sell_goods, establish_caravan, abandon_route
|
||
Scout: track_lord, assess_garrison, map_patrol_routes, report_intel
|
||
|
||
Refs: #1097, #1099.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
from typing import Any
|
||
|
||
from bannerlord.gabs_client import GABSClient, GABSUnavailable
|
||
from bannerlord.models import ResultMessage, TaskMessage
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class BaseCompanion:
|
||
"""Shared companion lifecycle — polls task queue, executes primitives."""
|
||
|
||
name: str = "base_companion"
|
||
primitives: frozenset[str] = frozenset()
|
||
|
||
def __init__(
|
||
self,
|
||
gabs_client: GABSClient,
|
||
task_queue: asyncio.Queue[TaskMessage],
|
||
result_queue: asyncio.Queue[ResultMessage] | None = None,
|
||
) -> None:
|
||
self._gabs = gabs_client
|
||
self._task_queue = task_queue
|
||
self._result_queue = result_queue or asyncio.Queue()
|
||
self._running = False
|
||
|
||
@property
|
||
def result_queue(self) -> asyncio.Queue[ResultMessage]:
|
||
return self._result_queue
|
||
|
||
async def run(self) -> None:
|
||
"""Companion event loop — processes task messages."""
|
||
self._running = True
|
||
logger.info("%s started", self.name)
|
||
try:
|
||
while self._running:
|
||
try:
|
||
task = await asyncio.wait_for(self._task_queue.get(), timeout=1.0)
|
||
except TimeoutError:
|
||
continue
|
||
|
||
if task.to_agent != self.name:
|
||
# Not for us — put it back (another companion will handle it)
|
||
await self._task_queue.put(task)
|
||
await asyncio.sleep(0.05)
|
||
continue
|
||
|
||
result = await self._execute(task)
|
||
await self._result_queue.put(result)
|
||
self._task_queue.task_done()
|
||
|
||
except asyncio.CancelledError:
|
||
logger.info("%s cancelled", self.name)
|
||
raise
|
||
finally:
|
||
self._running = False
|
||
|
||
def stop(self) -> None:
|
||
self._running = False
|
||
|
||
async def _execute(self, task: TaskMessage) -> ResultMessage:
|
||
"""Dispatch *task.primitive* to its handler method."""
|
||
handler = getattr(self, f"_prim_{task.primitive}", None)
|
||
if handler is None:
|
||
logger.warning("%s: unknown primitive %r — skipping", self.name, task.primitive)
|
||
return ResultMessage(
|
||
from_agent=self.name,
|
||
to_agent=task.from_agent,
|
||
success=False,
|
||
outcome={"error": f"Unknown primitive: {task.primitive}"},
|
||
)
|
||
try:
|
||
outcome = await handler(task.args)
|
||
return ResultMessage(
|
||
from_agent=self.name,
|
||
to_agent=task.from_agent,
|
||
success=True,
|
||
outcome=outcome or {},
|
||
)
|
||
except GABSUnavailable as exc:
|
||
logger.warning("%s: GABS unavailable for %r: %s", self.name, task.primitive, exc)
|
||
return ResultMessage(
|
||
from_agent=self.name,
|
||
to_agent=task.from_agent,
|
||
success=False,
|
||
outcome={"error": str(exc)},
|
||
)
|
||
except Exception as exc: # noqa: BLE001
|
||
logger.warning("%s: %r failed: %s", self.name, task.primitive, exc)
|
||
return ResultMessage(
|
||
from_agent=self.name,
|
||
to_agent=task.from_agent,
|
||
success=False,
|
||
outcome={"error": str(exc)},
|
||
)
|
||
|
||
|
||
# ── Logistics Companion ───────────────────────────────────────────────────────
|
||
|
||
|
||
class LogisticsCompanion(BaseCompanion):
|
||
"""Party management — recruitment, supply, healing, troop upgrades.
|
||
|
||
Skill domain: Scouting / Steward / Medicine.
|
||
"""
|
||
|
||
name = "logistics_companion"
|
||
primitives = frozenset(
|
||
{
|
||
"recruit_troop",
|
||
"buy_supplies",
|
||
"rest_party",
|
||
"sell_prisoners",
|
||
"upgrade_troops",
|
||
"build_project",
|
||
}
|
||
)
|
||
|
||
async def _prim_recruit_troop(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
troop_type = args.get("troop_type", "infantry")
|
||
qty = int(args.get("quantity", 10))
|
||
result = await self._gabs.recruit_troops(troop_type, qty)
|
||
logger.info("Recruited %d %s", qty, troop_type)
|
||
return result or {"recruited": qty, "type": troop_type}
|
||
|
||
async def _prim_buy_supplies(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
qty = int(args.get("quantity", 50))
|
||
result = await self._gabs.call("party.buySupplies", {"quantity": qty})
|
||
logger.info("Bought %d food supplies", qty)
|
||
return result or {"purchased": qty}
|
||
|
||
async def _prim_rest_party(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
days = int(args.get("days", 3))
|
||
result = await self._gabs.call("party.rest", {"days": days})
|
||
logger.info("Resting party for %d days", days)
|
||
return result or {"rested_days": days}
|
||
|
||
async def _prim_sell_prisoners(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
location = args.get("location", "nearest_town")
|
||
result = await self._gabs.call("party.sellPrisoners", {"location": location})
|
||
logger.info("Selling prisoners at %s", location)
|
||
return result or {"sold_at": location}
|
||
|
||
async def _prim_upgrade_troops(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
result = await self._gabs.call("party.upgradeTroops", {})
|
||
logger.info("Upgraded available troops")
|
||
return result or {"upgraded": True}
|
||
|
||
async def _prim_build_project(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
settlement = args.get("settlement", "")
|
||
result = await self._gabs.call("settlement.buildProject", {"settlement": settlement})
|
||
logger.info("Building project in %s", settlement)
|
||
return result or {"settlement": settlement}
|
||
|
||
async def _prim_move_party(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
destination = args.get("destination", "")
|
||
result = await self._gabs.move_party(destination)
|
||
logger.info("Moving party to %s", destination)
|
||
return result or {"destination": destination}
|
||
|
||
|
||
# ── Caravan Companion ─────────────────────────────────────────────────────────
|
||
|
||
|
||
class CaravanCompanion(BaseCompanion):
|
||
"""Trade route management — price assessment, goods trading, caravan deployment.
|
||
|
||
Skill domain: Trade / Charm.
|
||
"""
|
||
|
||
name = "caravan_companion"
|
||
primitives = frozenset(
|
||
{"assess_prices", "buy_goods", "sell_goods", "establish_caravan", "abandon_route"}
|
||
)
|
||
|
||
async def _prim_assess_prices(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
town = args.get("town", "nearest")
|
||
result = await self._gabs.call("trade.assessPrices", {"town": town})
|
||
logger.info("Assessed prices at %s", town)
|
||
return result or {"town": town}
|
||
|
||
async def _prim_buy_goods(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
item = args.get("item", "grain")
|
||
qty = int(args.get("quantity", 10))
|
||
result = await self._gabs.call("trade.buyGoods", {"item": item, "quantity": qty})
|
||
logger.info("Buying %d × %s", qty, item)
|
||
return result or {"item": item, "quantity": qty}
|
||
|
||
async def _prim_sell_goods(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
item = args.get("item", "grain")
|
||
qty = int(args.get("quantity", 10))
|
||
result = await self._gabs.call("trade.sellGoods", {"item": item, "quantity": qty})
|
||
logger.info("Selling %d × %s", qty, item)
|
||
return result or {"item": item, "quantity": qty}
|
||
|
||
async def _prim_establish_caravan(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
town = args.get("town", "")
|
||
result = await self._gabs.call("trade.establishCaravan", {"town": town})
|
||
logger.info("Establishing caravan at %s", town)
|
||
return result or {"town": town}
|
||
|
||
async def _prim_abandon_route(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
result = await self._gabs.call("trade.abandonRoute", {})
|
||
logger.info("Caravan route abandoned — returning to main party")
|
||
return result or {"abandoned": True}
|
||
|
||
|
||
# ── Scout Companion ───────────────────────────────────────────────────────────
|
||
|
||
|
||
class ScoutCompanion(BaseCompanion):
|
||
"""Intelligence gathering — lord tracking, garrison assessment, patrol mapping.
|
||
|
||
Skill domain: Scouting / Roguery.
|
||
"""
|
||
|
||
name = "scout_companion"
|
||
primitives = frozenset({"track_lord", "assess_garrison", "map_patrol_routes", "report_intel"})
|
||
|
||
async def _prim_track_lord(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
lord_name = args.get("name", "")
|
||
result = await self._gabs.call("intelligence.trackLord", {"name": lord_name})
|
||
logger.info("Tracking lord: %s", lord_name)
|
||
return result or {"tracking": lord_name}
|
||
|
||
async def _prim_assess_garrison(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
settlement = args.get("settlement", "")
|
||
result = await self._gabs.call("intelligence.assessGarrison", {"settlement": settlement})
|
||
logger.info("Assessing garrison at %s", settlement)
|
||
return result or {"settlement": settlement}
|
||
|
||
async def _prim_map_patrol_routes(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
region = args.get("region", "")
|
||
result = await self._gabs.call("intelligence.mapPatrols", {"region": region})
|
||
logger.info("Mapping patrol routes in %s", region)
|
||
return result or {"region": region}
|
||
|
||
async def _prim_report_intel(self, args: dict[str, Any]) -> dict[str, Any]:
|
||
result = await self._gabs.call("intelligence.report", {})
|
||
logger.info("Scout intel report generated")
|
||
return result or {"reported": True}
|