chatbot/chat.py

"""
KALRO Maize Research Chatbot — Claude API + tiered retrieval.

Usage:
    from chatbot.chat import Chatbot
    bot = Chatbot()
    response = bot.ask("What water harvesting technologies work best in semi-arid Kenya?")
    print(response.answer)
    print(response.citations)
"""

from dataclasses import dataclass, field
from pathlib import Path

import anthropic
from dotenv import load_dotenv

load_dotenv(Path(__file__).parent.parent / ".env")

from chatbot.retriever import Retriever, Chunk

MODEL = "claude-sonnet-4-6"
MAX_CONTEXT_CHUNKS = 8          # max chunks to include in the prompt
MAX_HISTORY_TURNS = 6           # conversation turns to keep in memory

SYSTEM_PROMPT = """You are a research assistant for KALRO (Kenya Agricultural and Livestock Research Organization), \
specialising in maize production in Kenya. You answer questions based on a curated knowledge base of research papers, \
field trials, and training manuals.

## How to use the provided context

You will receive context chunks labelled [WIKI] or [PDF].
- [WIKI] chunks come from structured, reviewed wiki summaries — treat these as your primary source.
- [PDF] chunks come directly from raw research documents — use them for detail or verbatim data \
  not captured in the wiki.

## Answering rules

1. Base your answer on the provided context. Do not invent facts or statistics.
2. At the end of your answer, list the sources you used under a "**Sources**" heading, \
   using the citation labels provided with each chunk. One bullet per source.
3. If the context is insufficient to answer fully, say so clearly and state what is and is not covered.
4. Keep answers concise and practical — the audience is agricultural researchers and extension officers.
5. Use plain English. Spell out abbreviations on first use.
"""


@dataclass
class ChatResponse:
    answer: str
    citations: list[str]
    pdf_fallback_used: bool
    chunks_used: list[Chunk] = field(default_factory=list)


def _build_context_block(chunks: list[Chunk]) -> str:
    parts = []
    for i, chunk in enumerate(chunks, 1):
        layer_tag = "[WIKI]" if chunk.layer == "wiki" else "[PDF]"
        citation = chunk.citation()
        parts.append(f"{layer_tag} [{i}] {citation}\n{chunk.text}")
    return "\n\n---\n\n".join(parts)


class Chatbot:
    def __init__(self, retriever: Retriever | None = None):
        self._retriever = retriever or Retriever()
        self._client = anthropic.Anthropic()   # reads ANTHROPIC_API_KEY from env
        self._history: list[dict] = []

    def ask(self, question: str) -> ChatResponse:
        all_chunks, pdf_fallback = self._retriever.retrieve(question)
        top_chunks = all_chunks[:MAX_CONTEXT_CHUNKS]

        context_block = _build_context_block(top_chunks)
        user_message = f"""## Context\n\n{context_block}\n\n---\n\n## Question\n\n{question}"""

        messages = self._history[-MAX_HISTORY_TURNS * 2:] + [
            {"role": "user", "content": user_message}
        ]

        response = self._client.messages.create(
            model=MODEL,
            max_tokens=1024,
            system=SYSTEM_PROMPT,
            messages=messages,
        )

        answer = response.content[0].text

        self._history.append({"role": "user", "content": user_message})
        self._history.append({"role": "assistant", "content": answer})

        citations = [c.citation() for c in top_chunks]

        return ChatResponse(
            answer=answer,
            citations=citations,
            pdf_fallback_used=pdf_fallback,
            chunks_used=top_chunks,
        )

    def reset(self):
        self._history.clear()