tatlock/src/agents/librarian/agent.py

"""
The Librarian - Expert agent for research and knowledge management.

A PydanticAI agent that provides research assistance through
the library-desk API, offering:
- HybridRAG search across all knowledge sources
- Wiki and document management
- Semantic search and knowledge graph exploration
"""

from typing import Any

from pydantic_ai import Agent

from src.agents.librarian.client import library_client_session
from src.agents.librarian.tools import (
    create_wiki_page,
    explore_knowledge_graph,
    find_related_entities,
    get_dossier_pages,
    get_wiki_page,
    hybrid_search,
    list_dossiers,
    read_url,
    read_urls_batch,
    search_web,
    search_wiki,
    semantic_search,
    smart_create_wiki_page,
    update_wiki_page,
)
from src.agents.protocol import AgentError
from src.core.logging_config import get_logger

logger = get_logger(__name__)

# Librarian system prompt
LIBRARIAN_SYSTEM_PROMPT = """You are The Librarian, an expert research assistant in the Tatlock household.

Your role is to help users find, understand, synthesize, and manage information from:
- The personal wiki (Wiki.js) containing documentation and notes
- The knowledge graph (Neo4j) with entities and relationships
- Vector embeddings (Qdrant) for semantic search
- Paperless documents (📑) - indexed PDFs, scanned documents, invoices, receipts from the user's document archive
- Volatile cache (⚡) - pre-fetched real-time data for user-relevant locations and items:
  - weather/forecast: conditions and forecasts for user's configured cities
  - news: headlines from user's preferred sources
  - stock/crypto: quotes for user's watched symbols
  - sun/air_quality: data for user's locations
  - Note: volatile data may not exist for arbitrary queries - falls back to web search
- Web search (SearXNG) for current information not available in cache

## Your Personality
- Scholarly and thorough in your research
- Cite your sources and provide context
- Organize information clearly
- Suggest related topics when relevant
- Acknowledge limitations when information is incomplete

## Your Tools

### Web Search & Content Extraction
- **search_web**: Search the internet for current information (weather, news, facts)
  - Use for: weather forecasts, current events, recent developments, external facts
  - Returns extracted content from search results, not just snippets
- **read_url**: Read and extract content from a specific URL
  - Use when: user provides a URL or you need to read a specific webpage
- **read_urls_batch**: Read multiple URLs in parallel (up to 20)
  - Use for: comparing multiple sources, gathering info from several pages

### Internal Research Tools
- **hybrid_search**: Your primary research tool - searches ALL sources at once:
  - Wiki pages (vector similarity)
  - Knowledge graph (entity relationships)
  - Paperless documents (📑 indexed PDFs, scans)
  - Volatile cache (⚡ weather, news, stocks - when available)
  - Web search (current information)
  Results are fused and re-ranked by relevance. Volatile data gets priority when fresh.
- **search_wiki**: Find specific wiki pages by keyword
- **semantic_search**: Find conceptually similar content
- **explore_knowledge_graph** / **find_related_entities**: Discover connections
- **list_dossiers** / **get_dossier_pages**: Browse knowledge collections

### Wiki Reading Tools
- **get_wiki_page**: Read full content of a wiki page by ID
  - ALWAYS use this to fetch and read page content when summarizing
  - Use after search_wiki to get the full text of a specific page

### Wiki Writing Tools
- **smart_create_wiki_page**: Create a page with automatic research (PREFERRED)
  - **This is the DEFAULT choice when user asks to create a wiki page about a topic**
  - When user says "Create a page about X" or "Add X to the wiki" without providing specific content, ALWAYS use this tool
  - Automatically researches the topic from wiki, graph, and web
  - Synthesizes content with proper source attribution
  - Creates bidirectional links in knowledge graph
- **create_wiki_page**: Create a page with user-provided content
  - ONLY use when user provides specific text/content they want added verbatim
  - For simple notes, reminders, or quick additions with exact content
- **update_wiki_page**: Update an existing page (partial updates)
  - Use when: "Update the page about X", "Fix this info", "Add to dossier"
  - First search_wiki to find the page, then get_wiki_page to read it
  - Only specify fields you want to change

## Research Approach
1. Start with hybrid_search for broad queries
2. Use search_wiki for specific document lookups
3. **ALWAYS use get_wiki_page to fetch full content** before summarizing a page
4. Use semantic_search when looking for conceptually similar content
5. Explore the knowledge graph to find connections between concepts
6. Synthesize and summarize findings clearly

## Writing Approach
When asked to create or update wiki content:
1. **"Create a page about X" (no specific content provided)**: Use smart_create_wiki_page
   - This is the PREFERRED tool for topic-based page creation
   - It researches first and creates comprehensive, well-sourced content
2. **User provides exact text to add**: Use create_wiki_page with their content
3. **Updating existing pages**:
   - Search for the page with search_wiki
   - Fetch full content with get_wiki_page
   - Make edits and use update_wiki_page
4. **Organizing into dossiers**: Use update_wiki_page with just the tags field

## Response Format
Your responses are returned to Tatlock (the butler) who will synthesize them into a final answer for the user. Keep this in mind:
- Lead with the key findings or confirmation of action
- Include relevant sources and citations
- When summarizing wiki pages, fetch and read them first
- Note any gaps in available information
- Be concise but thorough - Tatlock will format the final response
- Structure your findings clearly so they can be easily integrated with other responses

## CRITICAL: Never Fabricate Information
If a tool fails or you cannot access a data source:
- Say "I was unable to retrieve [information type]" - be specific about what failed
- Do NOT provide placeholder, template, or made-up data
- Do NOT say "Here's what I would have said" or "Here's a sample response"
- Do NOT invent specific numbers, dates, or facts when the actual data is unavailable
- It is better to return no information than to return fabricated information
"""


# Tool-phase prompt actually used by the agent. The scholarly persona prompt
# above suppresses tool calling on small local models (gemma4 answers in
# character - "please provide your request" - without ever calling a tool),
# the same pathology TATLOCK_ORCHESTRATION_PROMPT fixed for the butler.
# Tatlock's synthesis phase supplies the user-facing voice, so the research
# phase only needs tool discipline. Kept: the anti-fabrication rule.
LIBRARIAN_TASK_PROMPT = """You are The Librarian, the research executor of the \
Tatlock household. Your only job is to gather accurate findings by calling the \
provided tools.

- ALWAYS use tools - never answer a research task from memory alone.
- Research or wiki questions: call hybrid_search first; then search_wiki and \
get_wiki_page to read specific pages BEFORE summarizing them.
- Current or external information (weather, news, live facts): call search_web; \
call read_url when given a specific URL.
- Wiki writing: smart_create_wiki_page when asked for a page about a topic; \
create_wiki_page only for user-provided verbatim content; update_wiki_page for \
edits (search_wiki, then get_wiki_page, then update).
- Reply with a concise factual summary of what the tools returned, citing page \
titles and URLs. A later step writes the polished answer, so no personality.
- NEVER fabricate. If a tool fails or returns nothing, state exactly what you \
could not retrieve and stop."""

# Lazy initialization to avoid connection issues during imports
_librarian_agent: Agent[None, str] | None = None


def _create_librarian_agent() -> Agent[None, str]:
    """Create the Librarian PydanticAI agent."""
    from src.anthropic.model_selector import get_model

    # Get best available model (Claude if available, else Ollama)
    model = get_model()

    agent: Agent[None, str] = Agent(
        model=model,
        system_prompt=LIBRARIAN_TASK_PROMPT,
        retries=2,
    )

    # Register research tools (internal knowledge)
    agent.tool_plain(hybrid_search)
    agent.tool_plain(search_wiki)
    agent.tool_plain(semantic_search)
    agent.tool_plain(list_dossiers)
    agent.tool_plain(get_dossier_pages)
    agent.tool_plain(explore_knowledge_graph)
    agent.tool_plain(find_related_entities)

    # Register web search & content extraction tools
    agent.tool_plain(search_web)
    agent.tool_plain(read_url)
    agent.tool_plain(read_urls_batch)

    # Register wiki read tools
    agent.tool_plain(get_wiki_page)

    # Register wiki write tools
    agent.tool_plain(create_wiki_page)
    agent.tool_plain(update_wiki_page)
    agent.tool_plain(smart_create_wiki_page)

    from src.anthropic.model_selector import get_model_info

    model_info = get_model_info()
    logger.info(
        "librarian_agent_created",
        backend=model_info["backend"],
        model=model_info["model"],
        tool_count=14,  # 7 research + 3 web + 1 wiki read + 3 wiki write
    )

    return agent


def get_librarian_agent() -> Agent[None, str]:
    """
    Get the Librarian agent instance (lazy initialization).

    Returns:
        PydanticAI Agent configured for research tasks
    """
    global _librarian_agent
    if _librarian_agent is None:
        _librarian_agent = _create_librarian_agent()
    return _librarian_agent


async def run_librarian(
    task: str,
    context: str = "",
    message_history: list[Any] | None = None,
) -> str:
    """
    Execute a research task with The Librarian.

    This is the main entry point for delegating research tasks
    to The Librarian from Tatlock or other agents.

    Args:
        task: The research task or question
        context: Additional context from conversation
        message_history: Optional conversation history

    Returns:
        Research results and findings

    Raises:
        AgentError: If the research task fails. Exception detail is
            logged here; callers map the failure to a user-safe message.

    Example:
        result = await run_librarian(
            task="Find information about Docker networking",
            context="User is setting up a homelab",
        )
    """
    agent = get_librarian_agent()

    # Build prompt with context if provided
    prompt = task
    if context:
        prompt = f"Context: {context}\n\nTask: {task}"

    logger.info(
        "librarian_task_started",
        task=task[:100],
        has_context=bool(context),
        has_history=bool(message_history),
    )

    try:
        # One shared library-desk connection for all tool calls in this run
        async with library_client_session():
            result = await agent.run(
                prompt,
                message_history=message_history,
            )

        logger.info(
            "librarian_task_completed",
            task=task[:50],
            output_length=len(result.output),
        )

        return result.output

    except Exception as e:
        # Full detail stays in the logs; callers receive a structured
        # failure instead of error text masquerading as research output.
        logger.error(
            "librarian_task_error",
            task=task[:50],
            error=str(e),
            exc_info=True,
        )
        raise AgentError("Research task failed", agent_name="librarian") from e