Building a Retrieval Tool an Agent Can Call

from __future__ import annotations
 
from typing import Any
 
# embed_query / vector_search are adapters you implement for OpenAI+Pinecone,
# Chroma, pgvector, etc. Keep secrets on the server only.
 
def embed_query(text: str) -> list[float]:
    raise NotImplementedError  # call your embedding API
 
def vector_search(
    vector: list[float],
    *,
    top_k: int,
    filters: dict[str, Any],
) -> list[dict[str, Any]]:
    raise NotImplementedError  # return [{id, text, score, metadata}, ...]
 
def search_knowledge(
    query: str,
    top_k: int = 4,
    corpus: str | None = "policies",
) -> dict[str, Any]:
    """
    Search internal documentation and policies by semantic similarity.
 
    Use for refund rules, support hours, plan limits, and product behavior.
    Do not use for live order status, payments, or other customers' data.
    """
    if not query or not query.strip():
        return {"ok": False, "error_type": "bad_request", "message": "query required"}
 
    top_k = max(1, min(int(top_k), 10))
    filters: dict[str, Any] = {}
    if corpus:
        filters["corpus"] = corpus
    # Inject tenant from auth context in real systems - never trust model-supplied tenant ids alone.
 
    try:
        vec = embed_query(query.strip())
        hits = vector_search(vec, top_k=top_k, filters=filters)
    except Exception as exc:  # narrow in production
        return {
            "ok": False,
            "error_type": "upstream",
            "retryable": True,
            "message": "retrieval failed",
            "detail": str(exc)[:200],
        }
 
    return {
        "ok": True,
        "query": query.strip(),
        "results": [
            {
                "id": h["id"],
                "text": h["text"][:2000],
                "score": h.get("score"),
                "metadata": h.get("metadata") or {},
            }
            for h in hits
        ],
    }

retrieval_tool_schema = {
    "type": "function",
    "function": {
        "name": "search_knowledge",
        "description": (
            "Semantic search over internal policies and product docs. "
            "Use when the user asks about refunds, support hours, or plan rules. "
            "Not for live account balances or order tracking."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Full natural-language question to embed and search.",
                },
                "top_k": {
                    "type": "integer",
                    "minimum": 1,
                    "maximum": 10,
                    "default": 4,
                },
                "corpus": {
                    "type": "string",
                    "enum": ["policies", "product", "runbooks"],
                    "description": "Which document partition to search.",
                },
            },
            "required": ["query"],
            "additionalProperties": False,
        },
    },
}
 
TOOL_IMPL = {"search_knowledge": search_knowledge}
 
def dispatch_tool(name: str, args: dict[str, Any]) -> dict[str, Any]:
    fn = TOOL_IMPL.get(name)
    if not fn:
        return {"ok": False, "error_type": "unknown_tool", "message": name}
    return fn(**args)

# Pseudocode - use your framework's native tool loop when possible
messages = [{"role": "user", "content": "Can I get a refund after 12 days?"}]
# 1) model returns tool_call name=search_knowledge args={query: "..."}
# 2) observation = dispatch_tool(name, args)
# 3) append tool result; model writes final answer citing result ids