mirror of
https://github.com/furyhawk/agent_alpha.git
synced 2026-07-21 02:25:34 +00:00
feat: refactor agent initialization and add MemoryRepository and RagService
This commit is contained in:
+14
-99
@@ -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)
|
||||
|
||||
@@ -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")
|
||||
@@ -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(),
|
||||
],
|
||||
)
|
||||
@@ -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"),
|
||||
]
|
||||
Reference in New Issue
Block a user