Course ES
← back to chapter

KB agent

An agent whose one tool is `search` over a small product KB. It drives its own retrieval — deciding what to search, reading results, searching again if needed — and answers only from what it found.

Showcase — KB agent

An agent whose one tool is search over a small product KB. It drives its own retrieval — deciding what to search, reading results, searching again if needed — and answers only from what it found.

Run

bash bootstrap-secrets.sh              # reads ../../../../.env, writes secrets/
docker compose up --build              # default: PROVIDER=openai

Open http://localhost:3000. Gemini: PROVIDER=gemini docker compose up --build.

What's where

  • backend/ai_openai.py / backend/ai_gemini.py — the search tool + agent loop.
  • frontend/app/page.tsx — question box, grounded answer.

Stop

docker compose down

Run locally

Download the project as a ZIP and run it with Docker. Brings up a FastAPI backend + Next.js frontend on localhost:3000.

Download kb-agent.zip

unzip kb-agent.zip
cd kb-agent
bash bootstrap-secrets.sh   # one-time: pulls API keys into ./secrets
docker compose up --build   # default provider: openai
# or:  PROVIDER=gemini docker compose up --build

Type some input, pick a provider, and run the same code shown in Source against the live API. Sign-in required.


  

The same modules the Run button hits. The whole project (frontend, Dockerfile, compose) is in the ZIP under README.

backend/ai_openai.py

"""Showcase 2 (OpenAI): a knowledge-base agent.

An agent whose one tool is `search` over a small company KB. It decides what to
look up, reads the results, and answers from them — sometimes searching more than
once to piece an answer together. Tools plus grounding: the agent version of RAG,
where the model drives retrieval instead of a fixed pipeline.
"""
import json

from openai import OpenAI

_client = OpenAI()

_MODEL = "gpt-5.4-nano"

_KB = [
    {"topic": "pricing", "text": "Pro is $20/user/month billed annually, or $24 monthly. Free tier allows 2 users."},
    {"topic": "support", "text": "Pro includes 24/5 chat support; Enterprise adds a dedicated success manager."},
    {"topic": "security", "text": "Data is encrypted at rest (AES-256) and in transit (TLS 1.3). SOC 2 Type II certified."},
    {"topic": "integrations", "text": "Native integrations: Slack, GitHub, Jira, Google Drive. A REST API covers the rest."},
    {"topic": "limits", "text": "Free tier: 100 API calls/day. Pro: 10,000/day. Enterprise: negotiable."},
    {"topic": "onboarding", "text": "Enterprise plans include a guided onboarding and data migration from most competitors."},
]


def search(query: str) -> dict:
    words = [w for w in query.lower().split() if len(w) > 2]
    hits = [e for e in _KB if any(w in (e["text"] + " " + e["topic"]).lower() for w in words)]
    return {"results": hits[:4], "note": "" if hits else "no matches — try different keywords"}


_IMPL = {"search": search}
_TOOLS = [{"type": "function", "name": "search",
           "description": "Search the company knowledge base by keywords. Returns matching entries.",
           "parameters": {"type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"]}}]

_GOAL = ("You are a support agent for a SaaS product. Use `search` to find facts before "
         "answering — search more than once if needed. Answer ONLY from search results; if "
         "nothing relevant is found, say you don't have that information.")


def run(question: str) -> str:
    q = question.strip()
    if not q:
        return "Ask a question about the product (pricing, support, security, integrations, limits)."
    input_list = [{"role": "user", "content": q}]
    response = None
    for _ in range(6):
        response = _client.responses.create(model=_MODEL, instructions=_GOAL, input=input_list, tools=_TOOLS)
        calls = [i for i in response.output if i.type == "function_call"]
        if not calls:
            break
        input_list += response.output
        for c in calls:
            result = _IMPL[c.name](**json.loads(c.arguments))
            input_list.append({"type": "function_call_output", "call_id": c.call_id, "output": json.dumps(result)})
    return response.output_text if response else ""

backend/ai_gemini.py

"""Showcase 2 (Gemini): a knowledge-base agent.

Same `search`-tool agent, Gemini's tool-calling.
"""
import os

from google import genai
from google.genai import types

_client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

_MODEL = "gemini-3.1-flash-lite"

_KB = [
    {"topic": "pricing", "text": "Pro is $20/user/month billed annually, or $24 monthly. Free tier allows 2 users."},
    {"topic": "support", "text": "Pro includes 24/5 chat support; Enterprise adds a dedicated success manager."},
    {"topic": "security", "text": "Data is encrypted at rest (AES-256) and in transit (TLS 1.3). SOC 2 Type II certified."},
    {"topic": "integrations", "text": "Native integrations: Slack, GitHub, Jira, Google Drive. A REST API covers the rest."},
    {"topic": "limits", "text": "Free tier: 100 API calls/day. Pro: 10,000/day. Enterprise: negotiable."},
    {"topic": "onboarding", "text": "Enterprise plans include a guided onboarding and data migration from most competitors."},
]


def search(query: str) -> dict:
    words = [w for w in query.lower().split() if len(w) > 2]
    hits = [e for e in _KB if any(w in (e["text"] + " " + e["topic"]).lower() for w in words)]
    return {"results": hits[:4], "note": "" if hits else "no matches — try different keywords"}


_IMPL = {"search": search}

_CONFIG = types.GenerateContentConfig(
    system_instruction=("You are a support agent for a SaaS product. Use `search` to find facts "
                        "before answering — search more than once if needed. Answer ONLY from search "
                        "results; if nothing relevant is found, say you don't have that information."),
    tools=[types.Tool(function_declarations=[
        types.FunctionDeclaration(name="search", description="Search the company knowledge base by keywords.",
                                  parameters=types.Schema(type=types.Type.OBJECT,
                                      properties={"query": types.Schema(type=types.Type.STRING)}, required=["query"]))])],
)


def run(question: str) -> str:
    q = question.strip()
    if not q:
        return "Ask a question about the product (pricing, support, security, integrations, limits)."
    contents = [types.Content(role="user", parts=[types.Part(text=q)])]
    response = None
    for _ in range(6):
        response = _client.models.generate_content(model=_MODEL, contents=contents, config=_CONFIG)
        if not response.function_calls:
            break
        contents.append(response.candidates[0].content)
        parts = [types.Part.from_function_response(name=fc.name, response=_IMPL[fc.name](**dict(fc.args)))
                 for fc in response.function_calls]
        contents.append(types.Content(role="user", parts=parts))
    return (response.text or "") if response else ""

Project files

  • .gitignore
  • README.es.md
  • README.md
  • backend/Dockerfile
  • backend/ai_gemini.py
  • backend/ai_openai.py
  • backend/main.py
  • backend/requirements.txt
  • bootstrap-secrets.sh
  • docker-compose.yml
  • frontend/Dockerfile
  • frontend/app/layout.tsx
  • frontend/app/page.tsx
  • frontend/next.config.ts
  • frontend/package.json
  • frontend/tsconfig.json