Create chatbot_gemini.py

chatbot_gemini.py

@@ -0,0 +1,218 @@
+# NOTE: this is the Hugging Face Space copy of the Google AI Studio build. It is
+# the same shape as ../chatbot_ollama.py - only the provider changed. Compare the
+# two side by side: the observability code is identical, which is the point.
+import asyncio  # lets Python do jobs that involve waiting, like calls over the internet
+import json  # reads and writes data in the JSON text format
+import os  # lets us read settings and secrets from the computer
+import re  # finds and removes patterns inside a piece of text
+import time  # used to measure how long something takes
+import urllib.error  # holds the kinds of errors a web request can give back
+import urllib.request  # lets us send a request to a website by hand
+import uuid  # makes random one-of-a-kind id numbers
+from pathlib import Path  # an easy way to point at files and folders
+from typing import Annotated  # lets us attach a short description to a value
+
+from agent_framework import Agent, Message  # the chatbot itself, and one chat message
+from agent_framework.observability import enable_instrumentation  # switches on the recording of what the bot does
+from agent_framework.openai import OpenAIChatCompletionClient  # the part that talks to the model provider
+from dotenv import load_dotenv  # reads secrets out of a .env file
+from langfuse import Langfuse, propagate_attributes  # Langfuse keeps the recordings; the second one labels them
+from pydantic import Field  # used to describe what a piece of information means
+
+load_dotenv(Path(__file__).parent / ".env")  # load the secrets from the .env file sitting next to this file
+MODEL = os.getenv("GOOGLE_MODEL", "gemma-4-31b-it")  # which model to use, with a sensible default
+API_KEY = os.environ["GOOGLE_API_KEY"]  # the secret password for Google AI Studio; the program stops here if it was never set
+
+# Google speaks its own protocol, but it also offers an OpenAI-shaped door into the
+# same models. We use that door, so the client below is the very same class the
+# Ollama build uses - only the address and the key are different.
+BASE_URL = os.getenv("GOOGLE_BASE_URL", "https://generativelanguage.googleapis.com/v1beta/openai")
+
+langfuse = Langfuse()  # open the connection to Langfuse, our recording tool
+enable_instrumentation(enable_sensitive_data=True)  # start recording, including the real message text
+
+
+# Asks a Google endpoint a question and hands back the parsed JSON answer.
+def google_get(path: str) -> dict:  # start of the small helper that calls Google for us
+    """GET one path under BASE_URL, signed with the API key."""
+    request = urllib.request.Request(  # build the request, because Google needs the key in a header
+        f"{BASE_URL}/{path}", headers={"Authorization": f"Bearer {API_KEY}"}
+    )
+    with urllib.request.urlopen(request, timeout=30) as response:  # send it, giving up after 30 seconds
+        return json.load(response)  # turn the JSON text we got back into Python data
+
+
+# Explains an API-key problem in plain words, then stops the program.
+def stop_with_key_help(status: int, detail: str) -> None:  # start of the give-up-and-explain helper
+    """Google answers 400 for a bad key and 403 for a key that is not allowed here."""
+    env_file = Path(__file__).parent / ".env"  # work out where the .env file lives, to mention it below
+    raise SystemExit(  # stop right now and print something the reader can act on
+        f"\n[google] Google AI Studio would not accept GOOGLE_API_KEY "
+        f"(HTTP {status}).\n"
+        f"         The key ends '...{API_KEY[-6:]}' ({len(API_KEY)} chars).\n"
+        f"         Google said: {detail.strip()[:160]}\n\n"
+        f"         Get a working key at https://aistudio.google.com/apikey\n"
+        f"         and put it in {env_file}\n"
+        f"         On a Hugging Face Space there is no .env - add it under\n"
+        f"         Settings > Variables and secrets instead.\n"
+    ) from None
+
+
+try:  # try to read the model list, and catch the trouble if it goes wrong
+    listing = google_get("models")  # ask Google which models this key may use
+except urllib.error.HTTPError as error:  # this runs only if Google answered with an error
+    stop_with_key_help(error.code, error.read().decode(errors="replace"))  # explain the key problem and stop
+
+# Google returns names like "models/gemma-4-31b-it"; keep only the part after the slash.
+AVAILABLE = [m["id"].split("/")[-1] for m in listing["data"]]  # the plain name of every model on offer
+
+if MODEL not in AVAILABLE:  # if the model we asked for is not on that list
+    raise SystemExit(  # stop the program right now and print a helpful message instead
+        f"[google] '{MODEL}' is not available to this key. Available:\n  "
+        + "\n  ".join(sorted(AVAILABLE))
+    )
+
+
+# Sends one tiny message at startup to check that the key really works for chatting.
+def check_key() -> None:  # start of the key check; it hands nothing back
+    """Send the cheapest possible chat request just to see if the key is accepted."""
+    request = urllib.request.Request(  # build a very small chat request by hand
+        f"{BASE_URL}/chat/completions",
+        data=json.dumps({"model": MODEL, "max_tokens": 1,
+                         "messages": [{"role": "user", "content": "hi"}]}).encode(),
+        headers={"Authorization": f"Bearer {API_KEY}",
+                 "Content-Type": "application/json"},
+    )
+    try:  # try the next line, and catch the trouble if it goes wrong
+        urllib.request.urlopen(request, timeout=30).read()  # send it and wait up to 30 seconds for a reply
+    except urllib.error.HTTPError as error:  # this runs only if Google answered with an error
+        if error.code not in (400, 401, 403):  # if it is not a key problem and not a permission problem
+            raise  # pass the error on, because we have no better advice to offer
+        stop_with_key_help(error.code, error.read().decode(errors="replace"))  # explain and stop
+
+
+check_key()  # run that check now, before anything else gets started
+print(f"[google] {MODEL} via AI Studio - key accepted, {len(AVAILABLE)} models offered")  # tell the person all is well
+
+ACCOUNTS = {"SB-9001": 84_215.50, "SB-9002": 12_430.00, "SB-9003": 3_46_890.25}  # a pretend bank: account number and money in it
+
+
+# The tool the model can call: gives back the balance for one account number.
+def check_balance(  # start of the balance tool that the chatbot is allowed to use
+    account_id: Annotated[str, Field(description="Account id, e.g. SB-9001")],
+) -> str:
+    """Look up the balance of a Meridian Bank account."""
+    balance = ACCOUNTS.get(account_id.upper())  # look the account up, ignoring small or capital letters
+    return f"{account_id}: Rs {balance:,.2f}" if balance else f"No account {account_id}."  # give back the money, or say there is no such account
+
+
+agent = Agent(  # build the chatbot: who it talks to, its rules, its name, and the tools it may use
+    OpenAIChatCompletionClient(model=MODEL, api_key=API_KEY, base_url=BASE_URL),
+    "You are Meridian Bank's assistant. Branches open Mon-Fri 9:30-16:30, "
+    "Sat 9:30-13:30. Savings pays 2.75%, fixed deposits 7.25%. Use the tool for "
+    "any balance - never guess one. Be brief.",
+    name="MeridianAssist",
+    tools=[check_balance],
+)
+
+# Gemma thinks out loud before answering, and wraps that thinking in <thought> tags.
+# Google will not let us switch it off for this model ("Thinking budget is not
+# supported for this model"), so we cut the tags out before showing the reply.
+THOUGHTS = re.compile(r"<thought>.*?</thought>", re.DOTALL)  # a pattern that matches one whole thinking block
+UNCLOSED = re.compile(r"<thought>.*", re.DOTALL)  # a thinking block that was opened and never closed
+
+
+# Removes Gemma's thinking blocks, leaving just the answer meant for the reader.
+def strip_thoughts(text: str) -> str:  # start of the tidy-up helper for Gemma's replies
+    """Drop Gemma's thinking and leave only the answer.
+
+    The tags are not always balanced. A reply can end with a stray '</thought>'
+    that never had an opener, and a reply that was cut short can open a block it
+    never closes - so all three cases are handled, in this order.
+    """
+    text = THOUGHTS.sub("", text)  # first take out every properly closed thinking block
+    text = UNCLOSED.sub("", text)  # then drop an opener with no closer, and everything after it
+    return text.replace("</thought>", "").strip()  # finally remove any orphan closing tag, and trim
+
+
+# Turns whatever Gradio hands us as a message into a plain string.
+def plain_text(content: object) -> str:  # start of the tidy-up helper; it always hands back plain text
+    """Flatten one Gradio history entry's content down to a string."""
+    if isinstance(content, str):  # if it is already plain text
+        return content  # hand it straight back, untouched
+    if isinstance(content, list):  # if it is a list of several pieces
+        return "".join(plain_text(part) for part in content)  # tidy up each piece and glue them together
+    if isinstance(content, dict):  # if it is a labelled box of values
+        return str(content.get("text") or content.get("content") or "")  # pull the text out of it, or use nothing
+    return str(content)  # anything else: simply turn it into text
+
+
+# Starts a new conversation: a fresh id and counters set back to zero.
+def new_session() -> dict:  # start of the new-conversation helper
+    """Fresh, empty running totals for one conversation."""
+    return {"id": f"chat-{uuid.uuid4().hex[:12]}", "turns": 0,  # hand back a brand new id plus counters starting at zero
+            "tokens_in": 0, "tokens_out": 0, "tokens_think": 0, "seconds": 0.0}
+
+
+# Builds the Langfuse link that shows every turn of one conversation together.
+def session_url(trace_id: str, session_id: str) -> str:  # start of the link builder
+    """Link to the session view in Langfuse."""
+    trace_url = langfuse.get_trace_url(trace_id=trace_id) or ""  # ask Langfuse for the link to this one turn
+    return f"{trace_url.rsplit('/traces/', 1)[0]}/sessions/{session_id}"  # change the end of that link so it points at the whole chat
+
+
+# Runs one chat turn: asks the model, traces it, returns the reply and the token report.
+def ask(message: str, history: list, session: dict | None = None) -> tuple[str, str]:  # start of one chat turn
+    session = new_session() if session is None else session  # begin a fresh conversation if none was handed in
+    past = [Message(m["role"], [plain_text(m["content"])]) for m in history]  # put the older messages into the shape the bot expects
+    started = time.perf_counter()  # note the time now, so we can measure how long this takes
+
+    with propagate_attributes(session_id=session["id"]), \
+            langfuse.start_as_current_observation(name="chat turn", as_type="agent") as span:
+        trace_id = span.trace_id  # remember this recording's id so we can link to it later
+        try:  # try to answer, and catch the trouble if it goes wrong
+            result = asyncio.run(agent.run([*past, Message("user", [message])]))  # send the old messages plus the new one, and wait for the answer
+        except Exception as error:  # this runs only if something went wrong along the way
+            span.update(level="ERROR", status_message=str(error))  # mark this recording as failed and save the reason
+            langfuse.flush()  # push the recording over to Langfuse straight away
+            return (f"That turn failed: {type(error).__name__}",  # hand back a short apology plus a link to the failed recording
+                    f"**Error** · {error}\n\n"
+                    f"[See the failed trace in Langfuse]"
+                    f"({langfuse.get_trace_url(trace_id=trace_id)})")
+
+    elapsed = time.perf_counter() - started  # work out how many seconds the whole turn took
+    langfuse.flush()  # push the recording over to Langfuse now, rather than whenever it feels like it
+
+    used = result.usage_details or {}  # the token counts, or an empty box if none were reported
+    tokens_in = used.get("input_token_count") or 0  # how many tokens we sent to the model
+    tokens_out = used.get("output_token_count") or 0  # how many tokens of answer the model sent back
+    total = used.get("total_token_count") or 0  # every token the model charged us for, thinking included
+    # Google bills the thinking but does not report it as output, so the two never add
+    # up. The difference is the thinking, and naming it is the whole lesson here.
+    tokens_think = max(total - tokens_in - tokens_out, 0)  # whatever is left over is thinking, never below zero
+
+    session["turns"] += 1  # count this turn in the conversation's running total
+    session["tokens_in"] += tokens_in  # add the tokens we sent to the running total
+    session["tokens_out"] += tokens_out  # add the tokens we received to the running total
+    session["tokens_think"] += tokens_think  # add the hidden thinking tokens to the running total
+    session["seconds"] += elapsed  # add this turn's seconds to the running total
+
+    report = (f"**This turn** · **{tokens_in:,}** tokens in · **{tokens_out:,}** out · "  # build the little summary shown beside the chat
+              f"**{tokens_think:,}** thinking · **{elapsed:.2f}s**\n\n"
+              f"**All {session['turns']} turn(s)** · **{session['tokens_in']:,}** in · "
+              f"**{session['tokens_out']:,}** out · "
+              f"**{session['tokens_think']:,}** thinking · "
+              f"**{session['tokens_in'] + session['tokens_out'] + session['tokens_think']:,}** total · "
+              f"**{session['seconds']:.2f}s**\n\n"
+              f"[This turn's trace]({langfuse.get_trace_url(trace_id=trace_id)}) · "
+              f"[The whole conversation]({session_url(trace_id, session['id'])})")
+    return strip_thoughts(result.text), report  # hand back the bot's answer, thinking removed, and that summary
+
+
+if __name__ == "__main__":  # only run the little test below when this file is started directly
+    session, history = new_session(), []  # begin a fresh conversation with nothing said yet
+    for question in ["What's the balance on SB-9001?", "And SB-9003?"]:  # ask these two questions, one after the other
+        answer, summary = ask(question, history, session)  # send the question and collect the answer and summary
+        history += [{"role": "user", "content": question},  # remember what we asked and what the bot replied
+                    {"role": "assistant", "content": answer}]
+        print(f"\n> {question}\n{answer}\n\n{summary}")  # show the question, the answer and the summary on screen