Files
DocsGPT/docsgpt/agents/agentic_agent.py
T
Alex 6b9b193337 feat(agents): let agents walk a source's graph when they can search it
A one-shot graph ranking diffuses over the whole neighbourhood; a question
whose answer sits two hops away is better served by following the edges. The
graph_search tool gives an agent search_entities, get_relationships and
read_entity_pages over its graph sources. On a multi-hop corpus where the
bridging entity is never named, answers went from 0/10 with classic vector
retrieval to 10/10 with the tool, end to end through /stream.

The tool has no setting of its own. It is offered exactly where the agent can
already search: agentic and research agents always, a classic agent only for
sources exposed as a search tool. A graph source left at prefetch is used for
ranking only.

Tests now pin GRAPHRAG_ENABLED to its shipped default, as CI has: with a dev
.env enabling it, every agent test's graph check read the developer's real
database and left a pool to it behind.
2026-09-19 14:07:41 +01:00

61 lines
2.0 KiB
Python

import logging
from typing import Dict, Generator, Optional
from docsgpt.agents.base import BaseAgent
from docsgpt.agents.tools.graph_search import add_graph_search_tool
from docsgpt.agents.tools.internal_search import add_internal_search_tool
from docsgpt.agents.tools.wiki import add_wiki_tool
from docsgpt.logging import LogContext
logger = logging.getLogger(__name__)
class AgenticAgent(BaseAgent):
"""Agent where the LLM controls retrieval via tools.
Unlike ClassicAgent which pre-fetches docs into the prompt,
AgenticAgent gives the LLM an internal_search tool so it can
decide when, what, and whether to search.
"""
def __init__(
self,
retriever_config: Optional[Dict] = None,
wiki_config: Optional[Dict] = None,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.retriever_config = retriever_config or {}
self.wiki_config = wiki_config or {}
def _gen_inner(
self, query: str, log_context: LogContext
) -> Generator[Dict, None, None]:
tools_dict = self.tool_executor.get_tools()
add_internal_search_tool(tools_dict, self.retriever_config)
add_graph_search_tool(tools_dict, self.retriever_config)
if self.wiki_config:
add_wiki_tool(tools_dict, self.wiki_config)
self._prepare_tools(tools_dict)
# 4. Build messages (prompt has NO pre-fetched docs)
messages = self._build_messages(self.prompt, query)
# 5. Call LLM — the handler manages the tool loop
llm_response = self._llm_gen(messages, log_context)
yield from self._handle_response(
llm_response, tools_dict, messages, log_context
)
# 6. Collect sources from internal search tool results
self._collect_internal_sources()
yield {"sources": self.retrieved_docs}
yield {"tool_calls": self._get_truncated_tool_calls()}
log_context.stacks.append(
{"component": "agent", "data": {"tool_calls": self.tool_calls.copy()}}
)