mirror of
https://github.com/usestrix/strix.git
synced 2026-09-11 22:51:18 +00:00
Improve tool efficiency: connection pooling, parallel execution, dedup pre-filter, memory compression
- Add string-similarity pre-filter to vulnerability deduplication to limit LLM comparisons to the top 10 most similar reports instead of all reports - Replace per-request httpx.AsyncClient with persistent connection pool per sandbox, eliminating repeated TCP/TLS handshake overhead - Execute independent tools concurrently via asyncio.gather while keeping state-modifying tools sequential - Lower memory compression threshold from 100K to 60K tokens and cache token counts to avoid redundant litellm.token_counter calls - Double compression chunk size from 10 to 20 messages to halve LLM calls - Replace asyncio.sleep(0.5) polling with event-based wake signaling in agent state for immediate response to state changes https://claude.ai/code/session_012JYGtxVh4zRbzXKarNmb11
This commit is contained in:
parent
5a76fab4ae
commit
e014890de6
5 changed files with 186 additions and 31 deletions
|
|
@ -275,7 +275,7 @@ class BaseAgent(metaclass=AgentMeta):
|
|||
|
||||
return
|
||||
|
||||
await asyncio.sleep(0.5)
|
||||
await self.state.wait_for_wake(timeout=0.5)
|
||||
|
||||
async def _enter_waiting_state(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import asyncio
|
||||
import uuid
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
|
@ -10,6 +11,8 @@ def _generate_agent_id() -> str:
|
|||
|
||||
|
||||
class AgentState(BaseModel):
|
||||
model_config = {"arbitrary_types_allowed": True}
|
||||
|
||||
agent_id: str = Field(default_factory=_generate_agent_id)
|
||||
agent_name: str = "Strix Agent"
|
||||
parent_id: str | None = None
|
||||
|
|
@ -39,6 +42,9 @@ class AgentState(BaseModel):
|
|||
|
||||
errors: list[str] = Field(default_factory=list)
|
||||
|
||||
# Event for signaling state changes (excluded from serialization)
|
||||
_wake_event: asyncio.Event = Field(default_factory=asyncio.Event, exclude=True)
|
||||
|
||||
def increment_iteration(self) -> None:
|
||||
self.iteration += 1
|
||||
self.last_updated = datetime.now(UTC).isoformat()
|
||||
|
|
@ -49,6 +55,8 @@ class AgentState(BaseModel):
|
|||
message["thinking_blocks"] = thinking_blocks
|
||||
self.messages.append(message)
|
||||
self.last_updated = datetime.now(UTC).isoformat()
|
||||
if self.waiting_for_input:
|
||||
self._wake_event.set()
|
||||
|
||||
def add_action(self, action: dict[str, Any]) -> None:
|
||||
self.actions_taken.append(
|
||||
|
|
@ -106,6 +114,20 @@ class AgentState(BaseModel):
|
|||
if new_task:
|
||||
self.task = new_task
|
||||
self.last_updated = datetime.now(UTC).isoformat()
|
||||
self._wake_event.set()
|
||||
|
||||
def signal_wake(self) -> None:
|
||||
"""Signal the agent to wake up from waiting."""
|
||||
self._wake_event.set()
|
||||
|
||||
async def wait_for_wake(self, timeout: float = 0.5) -> bool:
|
||||
"""Wait for a wake signal with timeout. Returns True if signaled, False on timeout."""
|
||||
try:
|
||||
await asyncio.wait_for(self._wake_event.wait(), timeout=timeout)
|
||||
self._wake_event.clear()
|
||||
return True
|
||||
except TimeoutError:
|
||||
return False
|
||||
|
||||
def has_reached_max_iterations(self) -> bool:
|
||||
return self.iteration >= self.max_iterations
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import json
|
||||
import logging
|
||||
import re
|
||||
from difflib import SequenceMatcher
|
||||
from typing import Any
|
||||
|
||||
import litellm
|
||||
|
|
@ -10,6 +11,8 @@ from strix.config import Config
|
|||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MAX_COMPARISON_CANDIDATES = 10
|
||||
|
||||
DEDUPE_SYSTEM_PROMPT = """You are an expert vulnerability report deduplication judge.
|
||||
Your task is to determine if a candidate vulnerability report describes the SAME vulnerability
|
||||
as any existing report.
|
||||
|
|
@ -138,6 +141,59 @@ def _parse_dedupe_response(content: str) -> dict[str, Any]:
|
|||
}
|
||||
|
||||
|
||||
def _compute_similarity(report_a: dict[str, Any], report_b: dict[str, Any]) -> float:
|
||||
"""Compute lightweight string similarity between two reports for pre-filtering."""
|
||||
fields = ["title", "endpoint", "method", "target", "description"]
|
||||
|
||||
total_score = 0.0
|
||||
weights_sum = 0.0
|
||||
|
||||
# Weighted fields: title and endpoint matter most for duplicate detection
|
||||
field_weights = {
|
||||
"title": 3.0,
|
||||
"endpoint": 3.0,
|
||||
"method": 1.5,
|
||||
"target": 2.0,
|
||||
"description": 1.0,
|
||||
}
|
||||
|
||||
for field in fields:
|
||||
val_a = str(report_a.get(field, "")).lower().strip()
|
||||
val_b = str(report_b.get(field, "")).lower().strip()
|
||||
weight = field_weights.get(field, 1.0)
|
||||
|
||||
if not val_a or not val_b:
|
||||
continue
|
||||
|
||||
ratio = SequenceMatcher(None, val_a, val_b).ratio()
|
||||
total_score += ratio * weight
|
||||
weights_sum += weight
|
||||
|
||||
return total_score / weights_sum if weights_sum > 0 else 0.0
|
||||
|
||||
|
||||
def _prefilter_candidates(
|
||||
candidate: dict[str, Any],
|
||||
existing_reports: list[dict[str, Any]],
|
||||
max_candidates: int = MAX_COMPARISON_CANDIDATES,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Pre-filter existing reports using string similarity to reduce LLM calls.
|
||||
|
||||
Only the top-N most similar reports are sent to the LLM for detailed comparison,
|
||||
avoiding sending hundreds of reports in a single prompt.
|
||||
"""
|
||||
if len(existing_reports) <= max_candidates:
|
||||
return existing_reports
|
||||
|
||||
scored = []
|
||||
for report in existing_reports:
|
||||
score = _compute_similarity(candidate, report)
|
||||
scored.append((score, report))
|
||||
|
||||
scored.sort(key=lambda x: x[0], reverse=True)
|
||||
return [report for _, report in scored[:max_candidates]]
|
||||
|
||||
|
||||
def check_duplicate(
|
||||
candidate: dict[str, Any], existing_reports: list[dict[str, Any]]
|
||||
) -> dict[str, Any]:
|
||||
|
|
@ -150,8 +206,11 @@ def check_duplicate(
|
|||
}
|
||||
|
||||
try:
|
||||
# Pre-filter to only compare against the most similar existing reports
|
||||
filtered_reports = _prefilter_candidates(candidate, existing_reports)
|
||||
|
||||
candidate_cleaned = _prepare_report_for_comparison(candidate)
|
||||
existing_cleaned = [_prepare_report_for_comparison(r) for r in existing_reports]
|
||||
existing_cleaned = [_prepare_report_for_comparison(r) for r in filtered_reports]
|
||||
|
||||
comparison_data = {"candidate": candidate_cleaned, "existing_reports": existing_cleaned}
|
||||
|
||||
|
|
|
|||
|
|
@ -9,8 +9,9 @@ from strix.config import Config
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
MAX_TOTAL_TOKENS = 100_000
|
||||
MAX_TOTAL_TOKENS = 60_000
|
||||
MIN_RECENT_MESSAGES = 15
|
||||
COMPRESSION_CHUNK_SIZE = 20
|
||||
|
||||
SUMMARY_PROMPT_TEMPLATE = """You are an agent performing context
|
||||
condensation for a security agent. Your job is to compress scan data while preserving
|
||||
|
|
@ -43,13 +44,22 @@ Provide a technically precise summary that preserves all operational security co
|
|||
keeping the summary concise and to the point."""
|
||||
|
||||
|
||||
_token_cache: dict[int, int] = {}
|
||||
|
||||
|
||||
def _count_tokens(text: str, model: str) -> int:
|
||||
cache_key = hash(text)
|
||||
if cache_key in _token_cache:
|
||||
return _token_cache[cache_key]
|
||||
|
||||
try:
|
||||
count = litellm.token_counter(model=model, text=text)
|
||||
return int(count)
|
||||
count = int(litellm.token_counter(model=model, text=text))
|
||||
except Exception:
|
||||
logger.exception("Failed to count tokens")
|
||||
return len(text) // 4 # Rough estimate
|
||||
count = len(text) // 4 # Rough estimate
|
||||
|
||||
_token_cache[cache_key] = count
|
||||
return count
|
||||
|
||||
|
||||
def _get_message_tokens(msg: dict[str, Any], model: str) -> int:
|
||||
|
|
@ -215,7 +225,7 @@ class MemoryCompressor:
|
|||
return messages
|
||||
|
||||
compressed = []
|
||||
chunk_size = 10
|
||||
chunk_size = COMPRESSION_CHUNK_SIZE
|
||||
for i in range(0, len(old_msgs), chunk_size):
|
||||
chunk = old_msgs[i : i + chunk_size]
|
||||
summary = _summarize_messages(chunk, model_name, self.timeout)
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import asyncio
|
||||
import inspect
|
||||
import os
|
||||
from typing import Any
|
||||
|
|
@ -25,6 +26,31 @@ _SERVER_TIMEOUT = float(Config.get("strix_sandbox_execution_timeout") or "120")
|
|||
SANDBOX_EXECUTION_TIMEOUT = _SERVER_TIMEOUT + 30
|
||||
SANDBOX_CONNECT_TIMEOUT = float(Config.get("strix_sandbox_connect_timeout") or "10")
|
||||
|
||||
# Connection pool: reuse HTTP clients per sandbox instead of creating one per call
|
||||
_sandbox_clients: dict[str, httpx.AsyncClient] = {}
|
||||
|
||||
|
||||
def _get_sandbox_client(sandbox_id: str) -> httpx.AsyncClient:
|
||||
"""Get or create a persistent HTTP client for a sandbox, enabling connection reuse."""
|
||||
if sandbox_id not in _sandbox_clients:
|
||||
timeout = httpx.Timeout(
|
||||
timeout=SANDBOX_EXECUTION_TIMEOUT,
|
||||
connect=SANDBOX_CONNECT_TIMEOUT,
|
||||
)
|
||||
_sandbox_clients[sandbox_id] = httpx.AsyncClient(
|
||||
trust_env=False,
|
||||
timeout=timeout,
|
||||
limits=httpx.Limits(max_connections=10, max_keepalive_connections=5),
|
||||
)
|
||||
return _sandbox_clients[sandbox_id]
|
||||
|
||||
|
||||
async def close_sandbox_client(sandbox_id: str) -> None:
|
||||
"""Close and remove the HTTP client for a sandbox when it's torn down."""
|
||||
client = _sandbox_clients.pop(sandbox_id, None)
|
||||
if client:
|
||||
await client.aclose()
|
||||
|
||||
|
||||
async def execute_tool(tool_name: str, agent_state: Any | None = None, **kwargs: Any) -> Any:
|
||||
execute_in_sandbox = should_execute_in_sandbox(tool_name)
|
||||
|
|
@ -71,31 +97,27 @@ async def _execute_tool_in_sandbox(tool_name: str, agent_state: Any, **kwargs: A
|
|||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
timeout = httpx.Timeout(
|
||||
timeout=SANDBOX_EXECUTION_TIMEOUT,
|
||||
connect=SANDBOX_CONNECT_TIMEOUT,
|
||||
)
|
||||
client = _get_sandbox_client(agent_state.sandbox_id)
|
||||
|
||||
async with httpx.AsyncClient(trust_env=False) as client:
|
||||
try:
|
||||
response = await client.post(
|
||||
request_url, json=request_data, headers=headers, timeout=timeout
|
||||
)
|
||||
response.raise_for_status()
|
||||
response_data = response.json()
|
||||
if response_data.get("error"):
|
||||
posthog.error("tool_execution_error", f"{tool_name}: {response_data['error']}")
|
||||
raise RuntimeError(f"Sandbox execution error: {response_data['error']}")
|
||||
return response_data.get("result")
|
||||
except httpx.HTTPStatusError as e:
|
||||
posthog.error("tool_http_error", f"{tool_name}: HTTP {e.response.status_code}")
|
||||
if e.response.status_code == 401:
|
||||
raise RuntimeError("Authentication failed: Invalid or missing sandbox token") from e
|
||||
raise RuntimeError(f"HTTP error calling tool server: {e.response.status_code}") from e
|
||||
except httpx.RequestError as e:
|
||||
error_type = type(e).__name__
|
||||
posthog.error("tool_request_error", f"{tool_name}: {error_type}")
|
||||
raise RuntimeError(f"Request error calling tool server: {error_type}") from e
|
||||
try:
|
||||
response = await client.post(
|
||||
request_url, json=request_data, headers=headers
|
||||
)
|
||||
response.raise_for_status()
|
||||
response_data = response.json()
|
||||
if response_data.get("error"):
|
||||
posthog.error("tool_execution_error", f"{tool_name}: {response_data['error']}")
|
||||
raise RuntimeError(f"Sandbox execution error: {response_data['error']}")
|
||||
return response_data.get("result")
|
||||
except httpx.HTTPStatusError as e:
|
||||
posthog.error("tool_http_error", f"{tool_name}: HTTP {e.response.status_code}")
|
||||
if e.response.status_code == 401:
|
||||
raise RuntimeError("Authentication failed: Invalid or missing sandbox token") from e
|
||||
raise RuntimeError(f"HTTP error calling tool server: {e.response.status_code}") from e
|
||||
except httpx.RequestError as e:
|
||||
error_type = type(e).__name__
|
||||
posthog.error("tool_request_error", f"{tool_name}: {error_type}")
|
||||
raise RuntimeError(f"Request error calling tool server: {error_type}") from e
|
||||
|
||||
|
||||
async def _execute_tool_locally(tool_name: str, agent_state: Any | None, **kwargs: Any) -> Any:
|
||||
|
|
@ -310,6 +332,13 @@ def _get_tracer_and_agent_id(agent_state: Any | None) -> tuple[Any | None, str]:
|
|||
return tracer, agent_id
|
||||
|
||||
|
||||
# Tools that modify shared state and must run sequentially
|
||||
_SEQUENTIAL_TOOLS = frozenset({
|
||||
"finish_scan", "agent_finish", "delegate_task", "send_message",
|
||||
"wait_for_message", "create_agent",
|
||||
})
|
||||
|
||||
|
||||
async def process_tool_invocations(
|
||||
tool_invocations: list[dict[str, Any]],
|
||||
conversation_history: list[dict[str, Any]],
|
||||
|
|
@ -321,7 +350,42 @@ async def process_tool_invocations(
|
|||
|
||||
tracer, agent_id = _get_tracer_and_agent_id(agent_state)
|
||||
|
||||
# Partition into parallelizable and sequential tools
|
||||
parallel_batch: list[dict[str, Any]] = []
|
||||
sequential_queue: list[dict[str, Any]] = []
|
||||
|
||||
for tool_inv in tool_invocations:
|
||||
tool_name = tool_inv.get("toolName", "unknown")
|
||||
if tool_name in _SEQUENTIAL_TOOLS:
|
||||
sequential_queue.append(tool_inv)
|
||||
else:
|
||||
parallel_batch.append(tool_inv)
|
||||
|
||||
# Execute parallelizable tools concurrently
|
||||
if parallel_batch:
|
||||
tasks = [
|
||||
_execute_single_tool(tool_inv, agent_state, tracer, agent_id)
|
||||
for tool_inv in parallel_batch
|
||||
]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
for i, result in enumerate(results):
|
||||
if isinstance(result, Exception):
|
||||
tool_name = parallel_batch[i].get("toolName", "unknown")
|
||||
error_xml = (
|
||||
f"<tool_result>\n<tool_name>{tool_name}</tool_name>\n"
|
||||
f"<result>Error executing {tool_name}: {result!s}</result>\n</tool_result>"
|
||||
)
|
||||
observation_parts.append(error_xml)
|
||||
else:
|
||||
observation_xml, images, tool_should_finish = result
|
||||
observation_parts.append(observation_xml)
|
||||
all_images.extend(images)
|
||||
if tool_should_finish:
|
||||
should_agent_finish = True
|
||||
|
||||
# Execute sequential tools one at a time (order matters)
|
||||
for tool_inv in sequential_queue:
|
||||
observation_xml, images, tool_should_finish = await _execute_single_tool(
|
||||
tool_inv, agent_state, tracer, agent_id
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue