Building a Smarter FAQ Bot (with Gemini, RAG, and Structured Output)

chromadb import Documents, EmbeddingFunction, Embeddings from google.api_core
import retry from google import genai from google.genai import types
is_retriable = lambda e: (isinstance(e, genai.errors.APIError) and e.code in
{429, 503}) class GeminiEmbeddingFunction(EmbeddingFunction): document_mode =
True # Toggle between indexing docs and embedding queries
@retry.Retry(predicate=is_retriable) def __call__(self, input_texts:
Documents) -> Embeddings: task = "retrieval_document" if self.document_mode
else "retrieval_query" print(f"Embedding {'documents' if self.document_mode
else 'query'} ({len(input_texts)})...") try: # Assuming 'client' is
initialized Google GenAI client response = client.models.embed_content(
model="models/text-embedding-004", contents=input_texts,
config=types.EmbedContentConfig(task_type=task), # Specify task type ) return
[e.values for e in response.embeddings] except Exception as e: print(f"Error
during embedding: {e}") return [[] for _ in input_texts]
</div>

# --- 5. Setup ChromaDB Vector Store ---
import chromadb
import time

print("Setting up ChromaDB...")
DB_NAME = "googlecar_faq_db"
embed_fn = GeminiEmbeddingFunction()
chroma_client = chromadb.Client()  # In-memory client

try:
    db = chroma_client.get_or_create_collection(name=DB_NAME, embedding_function=embed_fn)
    print(f"Collection '{DB_NAME}' ready. Current count: {db.count()}")
    # Assuming 'documents' and 'doc_ids' are defined earlier
    if db.count() < len(documents):
        print(f"Adding/Updating documents in '{DB_NAME}'...")
        embed_fn.document_mode = True  # Set mode for indexing
        db.upsert(documents=documents, ids=doc_ids)  # Use upsert for safety
        time.sleep(2)  # Allow indexing to settle
        print(f"Documents added/updated. New count: {db.count()}")
    else:
        print("Documents already seem to be indexed.")
except Exception as e:
    print(f"Error setting up ChromaDB collection: {e}")
    raise SystemExit("ChromaDB setup failed. Exiting.")

# --- 6. Define Retrieval Function ---
def retrieve_documents(query: str, n_results: int = 1) -> list[str]:
    print(f"\nRetrieving documents for query: '{query}'")
    embed_fn.document_mode = False  # Switch to query mode
    try:
        results = db.query(query_texts=[query], n_results=n_results)
        if results and results.get("documents"):
            retrieved_docs = results["documents"][0]
            print(f"Retrieved {len(retrieved_docs)} documents.")
            return retrieved_docs
        else:
            print("No documents retrieved.")
            return []
    except Exception as e:
        print(f"Error querying ChromaDB: {e}")
        return []

# --- 7. Define Structured Output Schema ---
from typing_extensions import Literal
from pydantic import BaseModel

class AnswerWithConfidence(BaseModel):
    answer: str
    confidence: Literal["High", "Medium", "Low"]

# --- 8. Define Augmented Generation Function ---
def generate_structured_answer(query: str, context_docs: list[str]) -> dict | None:
    if not context_docs:
        print("No context provided, cannot generate answer.")
        return {
            "answer": "I couldn't find relevant information in the provided documents to answer this question.",
            "confidence": "Low",
        }

    context = "\n---\n".join(context_docs)

    prompt = f\"\"\"You are an AI assistant answering questions about a Google car based ONLY on the provided documents.
Context Documents:
---
{context}
---
Question: {query}
Based *only* on the information in the context documents above, answer the question.
Also, assess your confidence in the answer based *only* on the provided text:
- "High" if the answer is directly and clearly stated in the documents.
- "Medium" if the answer can be inferred but isn't explicitly stated.
- "Low" if the documents don't seem to contain the answer or are ambiguous.
Return your response ONLY as a JSON object with the keys "answer" and "confidence". Example format:
{
  "answer": "Your answer here.",
  "confidence": "High/Medium/Low"
}
\"\"\"
    try:
        generation_config = types.GenerateContentConfig(
            temperature=0.2,
            response_mime_type="application/json",  # Request JSON
            response_schema=AnswerWithConfidence,  # Provide the schema
        )
        # Assuming 'client' is initialized Google GenAI client
        response = client.models.generate_content(
            model="gemini-2.0-flash",
            contents=prompt,
            generation_config=generation_config,  # Pass the config object
        )

        # Safe access to parsed output
        if (
            response.candidates
            and response.candidates[0].content
            and response.candidates[0].content.parts
        ):
            parsed_output = response.candidates[0].content.parts[0].function_call
            # Fallback check if .parsed is used
            if not parsed_output and hasattr(
                response.candidates[0].content.parts[0], "parsed"
            ):
                parsed_output = response.candidates[0].content.parts[0].parsed

            if isinstance(parsed_output, dict) and "answer" in parsed_output and "confidence" in parsed_output:
                print("Generated Answer:", parsed_output)
                return parsed_output
            else:
                print("Warning: Could not extract valid JSON from response.")
                print("Raw response part:", response.candidates[0].content.parts[0])
                # Attempt to parse the text part if it exists and looks like JSON
                try:
                    import json

                    text_part = response.candidates[0].content.parts[0].text
                    if text_part and text_part.strip().startswith("{") and text_part.strip().endswith("}"):
                        parsed_json = json.loads(text_part)
                        if isinstance(parsed_json, dict) and "answer" in parsed_json and "confidence" in parsed_json:
                            print("Recovered JSON from text part:", parsed_json)
                            return parsed_json
                except Exception as json_e:
                    print(f"Could not parse text part as JSON: {json_e}")

        print("Error: Could not generate/parse structured response correctly.")
        return {"answer": "Error: Could not generate or parse the structured response from the AI.", "confidence": "Low"}

    except Exception as e:
        print(f"Error during content generation call: {e}")
        return {"answer": f"Error during generation API call: {e}", "confidence": "Low"}