langchain-google-colab-ai

LangChain chat model and LLM for Google Colab's first-party google.colab.ai model proxy — free Gemini/Gemma access inside Colab with no API key, exposed through the standard LangChain Runnable interface (LCEL, .invoke/.stream/.batch, output parsers, retries, fallbacks) with automatic LangSmith tracing and real token usage.

⚠️ Limitations (read first)

The Colab model proxy (mp.kaggle.net) does not allow:

• Tool calling → bind_tools() raises NotImplementedError.
• Structured output (response_format) → with_structured_output() raises NotImplementedError.

Consequently tool-calling agents and deep agents are not supported. What does work: chat, multi-turn conversations with system prompts, streaming, token usage, RAG, and arbitrary LCEL chains.

The proxy is only reachable from an external Colab runtime (or anywhere the MODEL_PROXY_HOST / MODEL_PROXY_API_KEY credentials are available).

Install

pip install "git+https://github.com/langchain-ai/langchain-google-colab-ai.git"

uv add "git+https://github.com/langchain-ai/langchain-google-colab-ai.git"

Quickstart

from langchain_google_colab_ai import ChatGoogleColabAI, GoogleColabAI

# Chat model (recommended) — no API key needed inside Colab.
chat = ChatGoogleColabAI(model="google/gemini-2.5-flash")

print(chat.invoke("Explain the Colab model proxy in one sentence.").content)

# Streaming
for chunk in chat.stream("Write a haiku about notebooks."):
    print(chunk.content, end="", flush=True)

# Token usage is populated automatically
msg = chat.invoke("Hello!")
print(msg.usage_metadata)  # {'input_tokens': ..., 'output_tokens': ..., 'total_tokens': ...}

# String-in / string-out LLM (mirrors google.colab.ai.generate_text)
llm = GoogleColabAI()
print(llm.invoke("Say hi in French."))

Both classes are ordinary LangChain Runnables, so they compose in LCEL chains:

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

prompt = ChatPromptTemplate.from_messages(
    [("system", "You are a terse assistant."), ("human", "{question}")]
)
chain = prompt | ChatGoogleColabAI() | StrOutputParser()
print(chain.invoke({"question": "What is LCEL?"}))

Choosing a model

Pass any model the proxy exposes via model=, e.g. "google/gemini-2.5-flash". Available models vary by account tier; the default is "google/gemini-3.5-flash". If a model name is rejected, try one reported by google.colab.ai.list_models() inside your Colab runtime.

LangSmith tracing

No code changes required — both classes are standard Runnables. Set:

export LANGSMITH_TRACING=true
export LANGSMITH_API_KEY=...   # your LangSmith key

Traces include real token counts (input/output/total). Dollar-cost columns may be blank, because LangSmith's pricing tables don't know the google/gemini-* models served through this proxy.

Best-effort structured output (not guaranteed)

Since response_format and tool calling are blocked, structured output is only available as a best-effort prompt-and-parse recipe:

from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate

parser = JsonOutputParser()
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", 'Respond ONLY with JSON matching: {{"city": str, "population": int}}'),
        ("human", "{query}"),
    ]
)
chain = prompt | ChatGoogleColabAI() | parser
chain.invoke({"query": "Tell me about Tokyo."})  # may raise if output is not JSON

This is not guaranteed to produce valid JSON. For reliable structured output or tool calling, use a provider that supports it.

Credentials

Resolved in this order:

1. Explicit base_url / api_key constructor arguments.
2. MODEL_PROXY_HOST and MODEL_PROXY_API_KEY environment variables.
3. Colab notebook secret store (google.colab.userdata → MODEL_PROXY_API_KEY).

The bearer token is wrapped in pydantic.SecretStr and never logged or printed. Before the token is transmitted, the host is validated (HTTPS + allowlisted domain) to guard against SSRF. To point at a non-allowlisted host deliberately, set LANGCHAIN_COLAB_ALLOW_INSECURE_HOST=1.

Development

uv sync --all-groups
uv run lefthook install     # git hooks: ruff on commit, ty + pytest on push

uv run ruff check .
uv run ruff format --check .
uv run ty check src
uv run pytest tests/unit_tests   # 100% coverage gate

Integration tests (tests/integration_tests/) only run inside Colab with live proxy credentials; they are skipped everywhere else.

License

MIT — © LangChain, Inc. See LICENSE.