python/src/agents/rag/tools.py

import json
import logging
from datetime import datetime

from pydantic_ai import RunContext

from src.agents.mcp_client import get_mcp_client

# Assuming RagDependencies will be imported by the agent that uses these tools.
# We will define it in rag_agent.py and import it here for type hinting.
from src.agents.rag_agent import RagDependencies

logger = logging.getLogger(__name__)


async def search_documents_tool(ctx: RunContext[RagDependencies], query: str, source_filter: str | None = None) -> str:
    """Search through documents using RAG query."""
    try:
        # Use source filter from context if not provided
        if source_filter is None:
            source_filter = ctx.deps.source_filter

        # Use MCP client to perform RAG query
        mcp_client = await get_mcp_client(agent_type="rag")
        result_json = await mcp_client.rag_search_knowledge_base(  # type: ignore
            query=query, source_id=source_filter, match_count=ctx.deps.match_count
        )

        # Parse the JSON response
        result = json.loads(result_json)

        if not result.get("success", False):
            return f"Search failed: {result.get('error', 'Unknown error')}"

        results = result.get("results", [])
        if not results:
            return "No results found for your query. Try using different search terms or removing filters."

        # Format results for display and capture citations
        formatted_results = []
        for i, res in enumerate(results, 1):
            similarity = res.get("similarity_score", res.get("similarity", 0))
            metadata = res.get("metadata", {})
            source = metadata.get("source", "Unknown")
            url = metadata.get("url", res.get("url", ""))
            content = res.get("content", "")

            # Capture physical citation
            ctx.deps.collected_citations.append(
                {
                    "index": i,
                    "source": source,
                    "url": url,
                    "similarity": similarity,
                    "snippet": content[:200] + "..." if len(content) > 200 else content,
                }
            )

            # Truncate content for LLM if too long
            llm_content = content[:500] + "..." if len(content) > 500 else content

            formatted_results.append(
                f"**Result {i}** (Relevance: {similarity:.2%})\nSource: {source}\nURL: {url}\nContent: {llm_content}\n"
            )

        return f"Found {len(results)} relevant results:\n\n" + "\n---\n".join(formatted_results)

    except Exception as e:
        logger.error(f"Error searching documents: {e}")
        return f"Error performing search: {str(e)}"


async def list_available_sources_tool(ctx: RunContext[RagDependencies]) -> str:
    """List all available sources that can be searched."""
    try:
        # Use MCP client to get available sources
        mcp_client = await get_mcp_client(agent_type="rag")
        result_json = await mcp_client.get_available_sources()

        # Parse the JSON response
        result = json.loads(result_json)

        if not result.get("success", False):
            return f"Failed to get sources: {result.get('error', 'Unknown error')}"

        sources = result.get("sources", [])
        if not sources:
            return "No sources are currently available. You may need to crawl some documentation first."

        source_list = []
        for source in sources:
            source_id = source.get("source_id", "Unknown")
            title = source.get("title", "Untitled")
            description = source.get("description", "")
            created = source.get("created_at", "")

            # Format the description if available
            desc_text = f" - {description}" if description else ""

            source_list.append(f"- **{source_id}**: {title}{desc_text} (added {created[:10]})")

        return f"Available sources ({len(sources)} total):\n" + "\n".join(source_list)

    except Exception as e:
        logger.error(f"Error listing sources: {e}")
        return f"Error retrieving sources: {str(e)}"


async def search_code_examples_tool(
    ctx: RunContext[RagDependencies], query: str, source_filter: str | None = None
) -> str:
    """Search for code examples related to the query."""
    try:
        # Use source filter from context if not provided
        if source_filter is None:
            source_filter = ctx.deps.source_filter

        # Use MCP client to search code examples
        mcp_client = await get_mcp_client(agent_type="rag")
        result_json = await mcp_client.call_tool(
            "rag_search_code_examples", query=query, source_id=source_filter, match_count=ctx.deps.match_count
        )

        # Parse the JSON response
        result = json.loads(result_json)

        if not result.get("success", False):
            return f"Code search failed: {result.get('error', 'Unknown error')}"

        examples = result.get("results", result.get("code_examples", []))
        if not examples:
            return "No code examples found for your query."

        formatted_examples = []
        for i, example in enumerate(examples, 1):
            similarity = example.get("similarity", 0)
            summary = example.get("summary", "No summary")
            code = example.get("code", example.get("code_block", ""))
            url = example.get("url", "")

            # Extract language from code block if available
            lang = "code"
            if code.startswith("```"):
                first_line = code.split("\n")[0]
                if len(first_line) > 3:
                    lang = first_line[3:].strip()

            formatted_examples.append(
                f"**Example {i}** (Relevance: {similarity:.2%})\n"
                f"Summary: {summary}\n"
                f"Source: {url}\n"
                f"```{lang}\n{code}\n```"
            )

        return f"Found {len(examples)} code examples:\n\n" + "\n---\n".join(formatted_examples)

    except Exception as e:
        logger.error(f"Error searching code examples: {e}")
        return f"Error searching code: {str(e)}"


async def refine_search_query_tool(ctx: RunContext[RagDependencies], original_query: str, context: str) -> str:
    """Refine a search query based on context to get better results."""
    try:
        # Simple query expansion based on context
        refined_parts = [original_query]

        # Add contextual keywords
        if "how" in original_query.lower():
            refined_parts.append("tutorial guide example")
        elif "what" in original_query.lower():
            refined_parts.append("definition explanation overview")
        elif "error" in original_query.lower() or "issue" in original_query.lower():
            refined_parts.append("troubleshooting solution fix")
        elif "api" in original_query.lower():
            refined_parts.append("endpoint method parameters response")

        # Add project-specific context if available
        if ctx.deps.project_id:
            refined_parts.append(f"project:{ctx.deps.project_id}")

        refined_query = " ".join(refined_parts)
        return f"Refined query: '{refined_query}' (original: '{original_query}')"

    except Exception as e:
        return f"Could not refine query: {str(e)}"


async def add_search_context_prompt(ctx: RunContext[RagDependencies]) -> str:
    source_info = f"Source Filter: {ctx.deps.source_filter}" if ctx.deps.source_filter else "No source filter"
    return f"""
**Current Search Context:**
- Project ID: {ctx.deps.project_id or "Global search"}
- {source_info}
- Max Results: {ctx.deps.match_count}
- Timestamp: {datetime.now().isoformat()}
"""