feat: refactor agent initialization and add MemoryRepository and RagService

This commit is contained in:
2026-06-16 19:57:42 +08:00
parent 9dee4d1200
commit f51f7fc8f4
4 changed files with 123 additions and 99 deletions
+14 -99
View File
@@ -28,23 +28,15 @@ from __future__ import annotations
import time
from dataclasses import dataclass
import logfire
from pydantic_ai import Agent
from pydantic_ai.capabilities import MCP, ToolSearch, WebSearch
from pydantic_ai.models.openai import OpenAIResponsesModel
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai.tools import Tool
from pydantic_ai_harness import CodeMode
# Community packages, alphabetical:
from pydantic_ai_backends import LocalBackend
from pydantic_ai_shields import InputGuard, SecretRedaction, ToolGuard
from pydantic_deep import create_deep_agent
from pydantic_deep.deps import DeepAgentDeps
from subagents_pydantic_ai import SubAgentConfig
from backend.core.config import Settings
from backend.repositories.memory_repository import MemoryRepository
from backend.services.agent_factory import build_agent
from backend.services.rag_service import RagService
from pydantic_deep.deps import DeepAgentDeps
class AgentService:
"""Encapsulates agent construction, execution, and teardown.
@@ -55,8 +47,7 @@ class AgentService:
def __init__(self, settings: Settings) -> None:
self._settings = settings
self._agent: Agent[DeepAgentDeps, str] | None = None
self._model: OpenAIResponsesModel | None = None
self._agent: Agent[any, str] | None = None
# ── Lifecycle ──────────────────────────────────────────────────────────
@@ -65,58 +56,11 @@ class AgentService:
if self._agent is not None:
return
if self._settings.logfire_token:
try:
logfire.configure(token=self._settings.logfire_token)
logfire.instrument_pydantic_ai()
except Exception:
pass # Logfire is optional — don't crash if it fails
memory_repo = MemoryRepository()
rag_service = RagService()
self._model = OpenAIResponsesModel(
self._settings.llm_model,
provider=OpenAIProvider(
base_url=self._settings.llm_base_url,
api_key=self._settings.llm_api_key or "no-key-required",
),
)
self._agent = create_deep_agent(
model=self._model,
include_todo=True,
include_filesystem=True,
include_subagents=True,
include_skills=True,
include_plan=True,
include_execute=False,
include_memory=True,
memory_dir=_MEMORY_DIR,
web_search=False,
web_fetch=True,
thinking="xhigh",
context_manager=True,
context_manager_max_tokens=100_000,
cost_tracking=True,
cost_budget_usd=5.0,
stuck_loop_detection=True,
subagents=[
SubAgentConfig(
name="researcher",
description="Deep research on a topic",
instructions="You are a thorough research assistant.",
),
],
skill_directories=["./skills"],
tools=self._build_rag_tools(),
capabilities=[
CodeMode(),
ToolSearch(),
MCP("https://hn.caseyjhand.com/mcp", native=True),
WebSearch(local="duckduckgo"),
InputGuard(guard=lambda p: "ignore previous instructions" not in p.lower()),
ToolGuard(blocked=["rm"], require_approval=["write_file"]),
SecretRedaction(),
],
)
# The factory handles all complex initialization logic
self._agent = build_agent(self._settings, memory_repo, rag_service)
async def shutdown(self) -> None:
"""Release agent resources."""
@@ -128,7 +72,11 @@ class AgentService:
"""Send a prompt to the agent and return the text output with usage stats."""
self.initialize()
assert self._agent is not None
deps = DeepAgentDeps(backend=_backend)
# Dependency injection using the repository's backend
memory_repo = MemoryRepository()
deps = DeepAgentDeps(backend=memory_repo.backend)
t0 = time.monotonic()
result = await self._agent.run(prompt, deps=deps)
elapsed = time.monotonic() - t0
@@ -141,25 +89,6 @@ class AgentService:
elapsed_seconds=round(elapsed, 2),
)
# ── Internal helpers ───────────────────────────────────────────────────
@staticmethod
def _build_rag_tools() -> list[Tool]:
"""Build RAG search tools for document retrieval."""
from skills.rag_search.search import rag_search, rag_search_by_document, rag_list_collections
return [
Tool(rag_search, name="rag_search"),
Tool(rag_search_by_document, name="rag_search_by_document"),
Tool(rag_list_collections, name="rag_list_collections"),
]
# ---------------------------------------------------------------------------
# AskResult — returned by AgentService.ask() with output + token usage stats.
# ---------------------------------------------------------------------------
@dataclass
class AskResult:
"""Result of a single agent ``ask()`` call with token usage information."""
@@ -179,18 +108,6 @@ class AskResult:
elapsed_seconds: float = 0.0
"""Wall-clock time in seconds the agent took to respond."""
# ---------------------------------------------------------------------------
# Persistent backend — uses local filesystem so agent memory survives restarts.
# ---------------------------------------------------------------------------
import os as _os
_AGENT_ROOT = _os.path.dirname(_os.path.dirname(_os.path.abspath(__file__)))
_MEMORY_DIR = _os.path.join(_AGENT_ROOT, ".agent_memory")
_backend = LocalBackend(root_dir=_AGENT_ROOT)
# ---------------------------------------------------------------------------
# Singleton — lazy-initialised on the first call to ``ask()``.
# Keeps the existing module-level ``ask`` facade so downstream imports in
@@ -201,12 +118,10 @@ from backend.core.config import settings as _settings # noqa: E402
_service = AgentService(_settings)
def get_service() -> AgentService:
"""Return the singleton AgentService instance."""
return _service
async def ask(prompt: str, session_id: str | None = None) -> AskResult:
"""Send a prompt to the agent and return the text output with usage stats."""
return await _service.ask(prompt, session_id=session_id)
+19
View File
@@ -0,0 +1,19 @@
import os as _os
from pydantic_ai_backends import LocalBackend
class MemoryRepository:
"""Handles local filesystem persistence for agent memory."""
def __init__(self) -> None:
# In /backend/repositories/memory_repository.py,
# the project root is 3 levels up from this file.
self._agent_root = _os.path.dirname(_os.path.dirname(_os.path.dirname(_os.path.abspath(__file__))))
self._backend = LocalBackend(root_dir=self._agent_root)
@property
def backend(self) -> LocalBackend:
return self._backend
@property
def memory_dir(self) -> str:
return _os.path.join(self._agent_root, ".agent_memory")
+70
View File
@@ -0,0 +1,70 @@
from typing import List
import logfire
from pydantic_ai.models.openai import OpenAIResponsesModel
from pydantic_ai.providers.openai import OpenAIProvider
from pydantic_ai import Agent
from pydantic_deep import create_deep_agent
from subagents_pydantic_ai import SubAgentConfig
from pydantic_ai_shields import InputGuard, SecretRedaction, ToolGuard
from backend.core.config import Settings
from backend.repositories.memory_repository import MemoryRepository
from backend.services.rag_service import RagService
def build_agent(settings: Settings, memory_repository: MemoryRepository, rag_service: RagService) -> Agent:
"""
Factory method to create the PydanticDeep agent.
Moves initialization logic from AgentService into this factory.
"""
if settings.logfire_token:
try:
logfire.configure(token=settings.logfire_token)
logfire.instrument_pydantic_ai()
except Exception:
pass # Logfire is optional — don't crash if it fails
model = OpenAIResponsesModel(
settings.llm_model,
provider=OpenAIProvider(
base_url=settings.llm_base_url,
api_key=settings.llm_api_key or "no-key-required",
),
)
return create_deep_agent(
model=model,
include_todo=True,
include_filesystem=True,
include_subagents=True,
include_skills=True,
include_plan=True,
include_execute=False,
include_memory=True,
memory_dir=memory_repository.memory_dir,
web_search=False,
web_fetch=True,
thinking="xhigh",
context_manager=True,
context_manager_max_tokens=100_000,
cost_tracking=True,
cost_budget_usd=5.0,
stuck_loop_detection=True,
subagents=[
SubAgentConfig(
name="researcher",
description="Deep research on a topic",
instructions="You are a thorough research assistant.",
),
],
skill_directories=["./skills"],
tools=rag_service.get_rag_tools(),
capabilities=[
CodeMode(),
ToolSearch(),
MCP("https://hn.caseyjhand.com/mcp", native=True),
WebSearch(local="duckduckgo"),
InputGuard(guard=lambda p: "ignore previous instructions" not in p.lower()),
ToolGuard(blocked=["rm"], require_approval=["write_file"]),
SecretRedaction(),
],
)
+20
View File
@@ -0,0 +1,20 @@
from pydantic_ai.tools import Tool
from typing import List
# These are imported from the existing skills structure
try:
from skills.rag_search.search import rag_search, rag_search_by_document, rag_list_collections
except ImportError:
# Fallback if the package isn't installed but we want to maintain the code structure
pass
class RagService:
"""Handles RAG tool construction."""
def get_rag_tools(self) -> List[Tool]:
"""Builds and returns RAG search tools for document retrieval."""
return [
Tool(rag_search, name="rag_search"),
Tool(rag_search_by_document, name="rag_search_by_document"),
Tool(rag_list_collections, name="rag_list_collections"),
]