knowledge_agent/llm.py

"""
llm.py — Groq-powered RAG answer generation for Knowledge Agent
Takes a question + retrieved context chunks → returns a grounded answer.
"""

import os
from typing import List, Dict, Optional
from groq import Groq

# ── Config ────────────────────────────────────────────────────────────────────
GROQ_API_KEY  = os.getenv("GROQ_API_KEY", "")
GROQ_MODEL    = os.getenv("GROQ_MODEL", "llama-3.3-70b-versatile")
MAX_TOKENS    = int(os.getenv("LLM_MAX_TOKENS", "1500"))
TEMPERATURE   = float(os.getenv("LLM_TEMPERATURE", "0.2"))

# How many context chunks to pass to the LLM
TOP_K_CHUNKS  = int(os.getenv("TOP_K_CHUNKS", "5"))

# Minimum similarity score to include a chunk (0–1, cosine)
MIN_SCORE     = float(os.getenv("MIN_CHUNK_SCORE", "0.30"))

# ── System prompt ─────────────────────────────────────────────────────────────
SYSTEM_PROMPT = """You are a precise, helpful personal knowledge assistant.
You answer questions STRICTLY based on the document excerpts provided below.

Rules:
- Ground every claim in the provided context. Cite source filenames inline like [source.pdf].
- If the context doesn't contain enough information, say so clearly — do NOT hallucinate.
- Be concise and structured. Use bullet points or numbered lists when helpful.
- If asked for a summary, provide a well-organised paragraph-style answer.
- When multiple documents cover the same topic, synthesise them coherently.
"""


def _build_context_block(chunks: List[Dict]) -> str:
    """Format retrieved chunks into a readable context block for the prompt."""
    filtered = [c for c in chunks if c.get("score", 0) >= MIN_SCORE]
    if not filtered:
        return "No relevant context found in your documents."

    lines = []
    for i, c in enumerate(filtered, 1):
        source = c.get("source", "unknown")
        score  = c.get("score", 0)
        text   = c.get("text", "").strip()
        lines.append(f"[Excerpt {i} | Source: {source} | Relevance: {score:.2f}]\n{text}")

    return "\n\n---\n\n".join(lines)


def _build_user_message(question: str, context: str) -> str:
    return f"""Here are relevant excerpts from your personal documents:

{context}

---

Question: {question}

Answer (cite sources inline):"""


# ── Main function ─────────────────────────────────────────────────────────────

def generate_answer(
    question: str,
    chunks: List[Dict],
    conversation_history: Optional[List[Dict]] = None,
) -> Dict:
    """
    Generate a grounded answer using Groq + retrieved chunks.

    Args:
        question: The user's question.
        chunks:   List of dicts from KnowledgeIndexer.query()
                  Each has keys: text, source, score.
        conversation_history: Optional list of prior {role, content} turns
                              for multi-turn chat support.

    Returns:
        {
            "answer": str,
            "sources": [str, ...],    # unique source filenames cited
            "chunks_used": int,
            "model": str,
        }
    """
    if not GROQ_API_KEY:
        raise EnvironmentError(
            "GROQ_API_KEY is not set. Add it to your .env file."
        )

    client = Groq(api_key=GROQ_API_KEY)

    # Build context from retrieved chunks
    context = _build_context_block(chunks)
    user_message = _build_user_message(question, context)

    # Build message list (supports conversational follow-ups)
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    if conversation_history:
        messages.extend(conversation_history)
    messages.append({"role": "user", "content": user_message})

    # ── Call Groq ──────────────────────────────────────────────────────────────
    response = client.chat.completions.create(
        model=GROQ_MODEL,
        messages=messages,
        temperature=TEMPERATURE,
        max_completion_tokens=MAX_TOKENS,
    )

    answer = response.choices[0].message.content.strip()

    # Extract unique source filenames from chunks that met the score threshold
    sources = list({
        c["source"]
        for c in chunks
        if c.get("score", 0) >= MIN_SCORE
    })

    return {
        "answer":      answer,
        "sources":     sources,
        "chunks_used": len([c for c in chunks if c.get("score", 0) >= MIN_SCORE]),
        "model":       GROQ_MODEL,
    }


# ── Streaming variant (used by FastAPI for real-time web UI) ──────────────────

def stream_answer(question: str, chunks: List[Dict]):
    """
    Generator that yields answer tokens one by one for SSE streaming.
    Usage: for token in stream_answer(q, chunks): ...
    """
    if not GROQ_API_KEY:
        yield "ERROR: GROQ_API_KEY not set."
        return

    client = Groq(api_key=GROQ_API_KEY)
    context = _build_context_block(chunks)
    user_message = _build_user_message(question, context)

    messages = [
        {"role": "system",  "content": SYSTEM_PROMPT},
        {"role": "user",    "content": user_message},
    ]

    stream = client.chat.completions.create(
        model=GROQ_MODEL,
        messages=messages,
        temperature=TEMPERATURE,
        max_completion_tokens=MAX_TOKENS,
        stream=True,
    )

    for chunk in stream:
        if chunk.choices[0].delta.content:
            yield chunk.choices[0].delta.content