ChatAgent

graph = StateGraph(ResearcherState)

graph.add_node('triage', nodes.triage_agent)
graph.add_node('sales', nodes.sales_agent)
graph.add_node('billing', nodes.billing_agent)
graph.add_node('support', nodes.support_agent)
graph.add_node('tools', nodes.tool_node)

graph.add_edge(START, 'triage')
graph.add_conditional_edges(
  'triage',
  routes.route_from_triage,
  {'sales': 'sales', 'billing': 'billing', 'support': 'support', 'end': END},
)
graph.add_conditional_edges(
  'billing', routes.should_continue, {'tools': 'tools', 'end': END}
)
# sales and support wire up identically
graph.add_conditional_edges(
  'tools',
  routes.route_from_tools,
  {'sales': 'sales', 'billing': 'billing', 'support': 'support'},
)
return graph.compile(checkpointer=saver)

class ResearcherRoutes:
  def route_from_triage(
    self, state: ResearcherState
  ) -> Literal['billing', 'sales', 'support', 'end']:
    agent = state['current_agent']
    routable_agents: list[ResearcherAgent] = ['billing', 'sales', 'support']
    if agent in routable_agents:
      return agent

    return 'end'

  def route_from_tools(
    self, state: ResearcherState
  ) -> Literal['sales', 'billing', 'support']:
    """Route back to whichever specialist invoked the tool call."""
    return cast(Literal['sales', 'billing', 'support'], state['current_agent'])

  def should_continue(self, state: ResearcherState) -> Literal['tools', 'end']:
    """Check if should continue to tools or end."""
    last_message = state['messages'][-1]

    if isinstance(last_message, AIMessage) and last_message.tool_calls:
      return 'tools'
    return 'end'

class ResearcherState(TypedDict):
  """State passed between nodes in the researcher graph."""

  messages: Annotated[list[BaseMessage], add_messages]
  current_agent: ResearcherAgent | None
  handoff_reason: str
  context_summary: str
  error: str | None
  retry_count: int
  model_used: str


class TriageUpdate(TypedDict):
  current_agent: ResearcherAgent
  handoff_reason: str
  context_summary: str
  messages: list[BaseMessage]


class SpecialistUpdate(TypedDict):
  messages: Annotated[list[BaseMessage], add_messages]
  current_agent: ResearcherAgent

type ResearcherAgent = Literal['sales', 'billing', 'support', 'end']


class HandoffDecision(BaseModel):
  handoff_to: ResearcherAgent = Field(description='Which agent to hand off to')
  reason: str = Field(description='Reason for handoff')
  context: str = Field(description='Key context to pass to the next agent')

def triage_agent(self, state: ResearcherState) -> TriageUpdate:
  """Initial triage to route the customer query to."""
  messages: list[BaseMessage] = [
    SystemMessage(content=TRIAGE_AGENT_PROMPT),
    *state['messages'],
  ]
  decision = cast(
    HandoffDecision,
    self.triage_llm.invoke(messages),
  )

  if decision.handoff_to == 'end':
    ...

  return {
    'context_summary': decision.context,
    'handoff_reason': decision.reason,
    'current_agent': decision.handoff_to,
    'messages': [
      AIMessage(
        f'[ TRIAGE ] Transferring to {decision.handoff_to}. {decision.reason}'
      )
    ],
  }

class ResearcherTools:
  def __init__(self, rag: Rag) -> None:
    self.rag = rag
    self.get_relevant_documents = tool(self._get_relevant_documents)

  def _get_relevant_documents(self, query: str) -> str:
    """Search for information relevant to the query."""
    return self.rag.ask(query)

@traceable(name='security_check_input')
def check_input(self, input: str) -> InputCheckResult:
  """Process the user input prompt through security checks."""
  notes: list[str] = []

  # 1. Check for suspicious input
  suspense = self.input_sanitiser.is_suspicious(input)
  if suspense.is_suspicious:
    notes.append(suspense.reason) if suspense.reason is not None else None
    return InputCheckResult(is_allowed=False, security_notes=notes, cleaned_text='')

  # 2. Sanitise the input
  cleaned = self.input_sanitiser.sanitise(input)

  # 3. Mask PII
  cleaned_masked = self.pii_detector.mask(cleaned)

  # 4. Go through security guard
  security_check_result = self.security_guard.security_check(cleaned_masked)
  if not security_check_result.safe:
    notes.append(security_check_result.reason)
    return InputCheckResult(is_allowed=False, security_notes=notes, cleaned_text='')

  # 5. Security checked
  return InputCheckResult(
    is_allowed=True, security_notes=notes, cleaned_text=cleaned_masked
  )

class InputSanitiser:
  """Sanitise user input before processing."""

  INJECTION_PATTERNS = [
    re.compile(
      r'ignore\s+(all\s+)?(previous|prior|above)\s+instructions', re.IGNORECASE
    ),
    re.compile(r'you\s+are\s+now\s+(a|an|in)\s', re.IGNORECASE),
    re.compile(r'reveal\s+(your|the)\s+(system\s+)?(prompt|instructions)', re.IGNORECASE),
    re.compile(r'</?(system|assistant|user)>', re.IGNORECASE),
    re.compile(r'\bDAN\s+mode\b', re.IGNORECASE),
    # …16 in total
  ]

  def is_suspicious(self, text: str) -> IsSuspiciousResult:
    """Check if input contains suspicious patterns"""
    for pattern in self.INJECTION_PATTERNS:
      if pattern.search(text):
        return IsSuspiciousResult(
          is_suspicious=True,
          reason=f'Suspicious pattern detected: {pattern.pattern}',
        )
    return IsSuspiciousResult(is_suspicious=False, reason=None)

async def get(self, query: str, thread_id: str) -> str | None:
  """Get cached response if the same query exists and has not expired"""
  query_hash = self._hash_query(query)

  with self._lock:
    entry = self.cache.get(query_hash)

    if entry is None:
      return None

    if entry['thread_id'] != thread_id:
      return None

    if self._is_expired(entry, int(time()), self.ttl_seconds):
      del self.cache[query_hash]
      return None

    entry['hits'] += 1
    self.cache.move_to_end(query_hash)
    return entry['response']

async def get(self, query: str, thread_id: str) -> str | None:
  """Find the one most semantically similar query
  WHERE thread_id equals the current thread
  AND timestamp is newer than the expiry cutoff"""

  expires_after = int(time()) - self.ttl_seconds

  results = await self.vectorstore.asimilarity_search_with_relevance_scores(
    query=query,
    k=1,
    efficient_filter={
      'bool': {
        'filter': [
          {'term': {'metadata.thread_id': thread_id}},
          {'range': {'metadata.timestamp': {'gte': expires_after}}},
        ]
      }
    },
  )

  if not results:
    return None

  doc, score = results[0]

  # Raw OpenSearch kNN score, not a cosine similarity: the score formula
  # depends on the index engine and space_type, so comparing it directly
  # avoids hardcoding a conversion that silently breaks if either changes.
  # Monotonic in cosine either way, so it is still a valid cutoff.
  if score < self.score_threshold:
    return None