From 109e31841eb3443be0a7ee9187eb224c534ef704 Mon Sep 17 00:00:00 2001 From: Collier King Date: Tue, 7 Jul 2026 15:14:34 -0500 Subject: [PATCH] add AI Search admin client --- CHANGELOG.md | 10 + README.md | 1 + docs/ai_search.ipynb | 414 +++++++ libs/langchain-cloudflare/README.md | 45 +- .../examples/workers/pyproject.toml | 2 +- .../examples/workers/src/entry.py | 67 + .../examples/workers/wrangler.jsonc | 7 + .../langchain_cloudflare/__init__.py | 2 + .../langchain_cloudflare/_errors.py | 8 + .../langchain_cloudflare/ai_search.py | 1093 +++++++++++++++++ libs/langchain-cloudflare/pyproject.toml | 2 +- libs/langchain-cloudflare/tests/conftest.py | 3 +- .../test_ai_search_client.py | 99 ++ .../tests/unit_tests/test_ai_search.py | 329 +++++ .../worker_tests/test_worker_integration.py | 58 + 15 files changed, 2136 insertions(+), 4 deletions(-) create mode 100644 docs/ai_search.ipynb create mode 100644 libs/langchain-cloudflare/langchain_cloudflare/ai_search.py create mode 100644 libs/langchain-cloudflare/tests/integration_tests/test_ai_search_client.py create mode 100644 libs/langchain-cloudflare/tests/unit_tests/test_ai_search.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 2988f98..78a3c99 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## langchain-cloudflare +### [0.3.7] + +#### Added + +- **`CloudflareAISearchClient`**: Adds an AI Search administration client for creating, listing, reading, updating, and deleting AI Search instances; reading instance stats; uploading, listing, reading, and deleting built-in-storage items; and running raw search and chat-completions requests. Supports REST API usage and async Python Worker `ai_search_namespaces` bindings for namespace-level instance management. +- **AI Search admin integration coverage**: Adds live REST lifecycle coverage that creates a temporary AI Search instance, uploads and searches a markdown item, and deletes the instance; adds Worker namespace-binding coverage for create/delete and namespace search. +- **AI Search notebook example**: Adds `docs/ai_search.ipynb` with a walkthrough for credentials, instance creation, item upload, raw search, retriever usage, and cleanup. + +--- + ### [0.3.6] #### Added diff --git a/README.md b/README.md index 28cde05..f230453 100644 --- a/README.md +++ b/README.md @@ -19,6 +19,7 @@ It contains the following packages. - [CloudflareWorkersAIEmbeddings](https://python.langchain.com/docs/integrations/text_embedding/cloudflare_workersai/) - [CloudflareWorkersAIReranker](https://developers.cloudflare.com/workers-ai/) — document reranking on Workers AI - [CloudflareVectorize](https://python.langchain.com/docs/integrations/vectorstores/cloudflare_vectorize/) +- [CloudflareAISearchClient](libs/langchain-cloudflare/README.md#ai-search-administration) — AI Search instance and item lifecycle - [CloudflareAISearchRetriever](libs/langchain-cloudflare/README.md#retrievers) — Cloudflare [AI Search](https://developers.cloudflare.com/ai-search/) (fka AutoRAG) ### LangGraph diff --git a/docs/ai_search.ipynb b/docs/ai_search.ipynb new file mode 100644 index 0000000..02edc34 --- /dev/null +++ b/docs/ai_search.ipynb @@ -0,0 +1,414 @@ +{ + "cells": [ + { + "cell_type": "raw", + "metadata": { + "collapsed": false + }, + "source": [ + "---\n", + "sidebar_label: CloudflareAISearch\n", + "---" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Cloudflare AI Search\n", + "\n", + "This notebook covers how to get started with Cloudflare AI Search instance administration and retrieval." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Setup\n", + "\n", + "This Python package wraps Cloudflare's REST API. To interact with AI Search, provide an API token with the appropriate privileges.\n", + "\n", + "You can create and manage API tokens here:\n", + "\n", + "https://dash.cloudflare.com/YOUR-ACCT-NUMBER/api-tokens\n", + "\n", + "### Credentials\n", + "\n", + "For this notebook, use a token with **AI Search:Edit** and **AI Search:Run** permissions.\n", + "\n", + "You can use `CF_AI_SEARCH_API_TOKEN` for AI Search-specific access or `CF_API_TOKEN` if you have a broader token. The examples also read `CF_ACCOUNT_ID` and optionally `CF_AI_SEARCH_NAMESPACE` from the environment." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from dotenv import load_dotenv\n", + "import os\n", + "\n", + "load_dotenv(\".env\")\n", + "\n", + "# Declare as variables for example's sake (you'll see why below)\n", + "cf_acct_id = os.getenv(\"CF_ACCOUNT_ID\")\n", + "\n", + "# AI Search token with AI Search:Edit and AI Search:Run permissions\n", + "cf_ai_search_token = os.getenv(\"CF_AI_SEARCH_API_TOKEN\")\n", + "\n", + "# OR, a single broader Cloudflare token\n", + "api_token = os.getenv(\"CF_API_TOKEN\")\n", + "\n", + "# AI Search instances live in a namespace. The default namespace is usually enough.\n", + "ai_search_namespace = os.getenv(\"CF_AI_SEARCH_NAMESPACE\", \"default\")\n", + "\n", + "token = cf_ai_search_token or api_token" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Initialization" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import uuid\n", + "import warnings\n", + "\n", + "from langchain_cloudflare import (\n", + " CloudflareAISearchClient,\n", + " CloudflareAISearchRetriever,\n", + ")\n", + "\n", + "warnings.filterwarnings(\"ignore\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# name your AI Search instance\n", + "ai_search_instance_name = f\"test-langchain-ai-search-{uuid.uuid4().hex[:8]}\"\n", + "\n", + "# use a unique phrase so we can search for exactly the content we upload\n", + "query_marker = f\"langchain-ai-search-notebook-{uuid.uuid4().hex}\"" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### CloudflareAISearchClient Class\n", + "\n", + "Now we can create the `CloudflareAISearchClient` instance. Here we passed:\n", + "\n", + "* The account ID\n", + "* An AI Search token or broader Cloudflare API token\n", + "* The AI Search namespace" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "client = CloudflareAISearchClient(\n", + " account_id=cf_acct_id,\n", + " api_token=token,\n", + " namespace=ai_search_namespace,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Cleanup\n", + "\n", + "Before we get started, let's delete any `test-langchain-ai-search-*` instances left over from earlier notebook runs." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "arr_instances = client.list_instances(search=\"test-langchain-ai-search\")\n", + "arr_instances = [\n", + " x for x in arr_instances if x.get(\"id\", \"\").startswith(\"test-langchain-ai-search-\")\n", + "]\n", + "\n", + "for instance in arr_instances:\n", + " client.delete_instance(instance.get(\"id\"), missing_ok=True)\n", + "\n", + "print(f\"Deleted {len(arr_instances)} old notebook instances\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Manage AI Search\n", + "\n", + "### Creating an Instance\n", + "\n", + "Let's start by creating a temporary AI Search instance. New AI Search instances include built-in storage, so we can upload files directly to the instance." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "instance = client.create_instance(ai_search_instance_name)\n", + "print(instance)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Listing Instances\n", + "\n", + "Now, we can list AI Search instances in the namespace on our account." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "instances = client.list_instances(search=ai_search_instance_name)\n", + "[x.get(\"id\") for x in instances]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Instance Info and Stats\n", + "\n", + "We can retrieve instance configuration and indexing stats." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "client.get_instance(ai_search_instance_name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "client.stats(ai_search_instance_name)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Uploading Items\n", + "\n", + "AI Search can index files uploaded to built-in storage. This example uploads one small Markdown document and waits for indexing to complete." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "content = \"\\n\".join(\n", + " [\n", + " \"# LangChain Cloudflare AI Search notebook\",\n", + " \"\",\n", + " f\"{query_marker} validates AI Search notebook retrieval.\",\n", + " \"Cloudflare AI Search indexes uploaded files for natural language search.\",\n", + " ]\n", + ")\n", + "\n", + "item = client.upload_item(\n", + " \"notebook-guide.md\",\n", + " content,\n", + " instance_name=ai_search_instance_name,\n", + " content_type=\"text/markdown\",\n", + " metadata={\"source\": \"notebook\"},\n", + " wait_for_completion=True,\n", + ")\n", + "\n", + "if item.get(\"status\") != \"completed\":\n", + " item = client.wait_for_item(item[\"id\"], instance_name=ai_search_instance_name)\n", + "\n", + "item" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Listing Items\n", + "\n", + "Now, we can inspect uploaded items for this instance." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "items = client.list_items(ai_search_instance_name)\n", + "[(x.get(\"id\"), x.get(\"key\"), x.get(\"status\")) for x in items]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Query AI Search\n", + "\n", + "We can run a raw AI Search query directly through the admin client. The result contains chunks from the indexed files." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "results = client.search(\n", + " query_marker,\n", + " instance_name=ai_search_instance_name,\n", + " ai_search_options={\n", + " \"retrieval\": {\"max_num_results\": 3},\n", + " \"query_rewrite\": {\"enabled\": False},\n", + " \"reranking\": {\"enabled\": False},\n", + " },\n", + ")\n", + "\n", + "[(chunk.get(\"text\"), chunk.get(\"score\")) for chunk in results.get(\"chunks\", [])]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Query by Turning into Retriever\n", + "\n", + "You can also use the AI Search instance as a LangChain retriever for easier usage in chains and agents." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "retriever = CloudflareAISearchRetriever(\n", + " account_id=cf_acct_id,\n", + " api_token=token,\n", + " namespace=ai_search_namespace,\n", + " instance_name=ai_search_instance_name,\n", + " k=3,\n", + " rewrite_query=False,\n", + " reranking=False,\n", + ")\n", + "\n", + "docs = retriever.invoke(query_marker)\n", + "[(doc.page_content, doc.metadata) for doc in docs]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Async Examples\n", + "\n", + "The admin client also exposes async methods for REST usage and Python Worker bindings. In a Python Worker, pass an `ai_search_namespaces` binding instead of REST credentials:\n", + "\n", + "```python\n", + "client = CloudflareAISearchClient(binding=env.AI_SEARCH)\n", + "instance = await client.acreate_instance(\"tenant-a\")\n", + "await client.adelete_instance(instance[\"id\"])\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Cleanup\n", + "\n", + "Let's finish by deleting the item and AI Search instance we created in this notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "if item.get(\"id\"):\n", + " client.delete_item(\n", + " item[\"id\"],\n", + " instance_name=ai_search_instance_name,\n", + " missing_ok=True,\n", + " )\n", + "\n", + "client.delete_instance(ai_search_instance_name, missing_ok=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## API Reference\n", + "\n", + "For more information, see:\n", + "\n", + "* [Cloudflare AI Search docs](https://developers.cloudflare.com/ai-search/)\n", + "* [AI Search REST API](https://developers.cloudflare.com/api/resources/ai_search/)\n", + "* [AI Search Workers binding](https://developers.cloudflare.com/ai-search/api/instances/workers-binding/)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "cc-langchain-G_cWTCcf-py3.11", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/libs/langchain-cloudflare/README.md b/libs/langchain-cloudflare/README.md index a8d4532..8cd3330 100644 --- a/libs/langchain-cloudflare/README.md +++ b/libs/langchain-cloudflare/README.md @@ -19,7 +19,7 @@ AND OR (if using separately scoped tokens) - `CF_AI_API_TOKEN` (CloudflareWorkersAI and CloudflareWorkersAIEmbeddings) -- `CF_AI_SEARCH_API_TOKEN` (CloudflareAISearchRetriever) +- `CF_AI_SEARCH_API_TOKEN` (CloudflareAISearchClient and CloudflareAISearchRetriever) - `CF_VECTORIZE_API_TOKEN` (CloudflareVectorize) - `CF_D1_API_TOKEN` (CloudflareVectorize) - `CF_D1_DATABASE_ID` (CloudflareVectorize) @@ -85,6 +85,49 @@ vst = CloudflareVectorize( vst.create_index(index_name="my-cool-vectorstore") ``` +## AI Search Administration + +`CloudflareAISearchClient` manages Cloudflare [AI Search](https://developers.cloudflare.com/ai-search/) instances and uploaded items. Use it for provisioning and lifecycle work; use `CloudflareAISearchRetriever` when you want a LangChain retriever for an existing instance. + +```python +from langchain_cloudflare import CloudflareAISearchClient + +client = CloudflareAISearchClient() + +instance = client.create_instance("my-instance") +item = client.upload_item( + "guide.md", + "AI Search indexes uploaded content for retrieval.", + instance_name=instance["id"], + content_type="text/markdown", + wait_for_completion=True, +) + +results = client.search( + "How does AI Search handle uploaded content?", + instance_name=instance["id"], +) + +client.delete_item(item["id"], instance_name=instance["id"], missing_ok=True) +client.delete_instance(instance["id"], missing_ok=True) +``` + +For REST usage, set: + +- `CF_ACCOUNT_ID` +- `CF_AI_SEARCH_API_TOKEN` — an `AI Search:Edit` and `AI Search:Run` token (falls back to `CF_API_TOKEN`) +- `CF_AI_SEARCH_NAMESPACE` — optional, defaults to `default` + +Inside a Python Worker, pass an `ai_search_namespaces` binding for instance lifecycle operations: + +```python +client = CloudflareAISearchClient(binding=env.AI_SEARCH) +instance = await client.acreate_instance("tenant-a") +await client.adelete_instance(instance["id"]) +``` + +An instance-specific `ai_search` binding can still call instance methods such as `asearch()`, `astats()`, and item operations, but it cannot create, list, or delete instances. + ## Retrievers `CloudflareAISearchRetriever` exposes Cloudflare [AI Search](https://developers.cloudflare.com/ai-search/) (the managed retrieval / RAG service, fka AutoRAG) as a LangChain retriever. diff --git a/libs/langchain-cloudflare/examples/workers/pyproject.toml b/libs/langchain-cloudflare/examples/workers/pyproject.toml index 2c48b7e..dd55076 100644 --- a/libs/langchain-cloudflare/examples/workers/pyproject.toml +++ b/libs/langchain-cloudflare/examples/workers/pyproject.toml @@ -6,7 +6,7 @@ readme = "README.md" requires-python = ">=3.12" dependencies = [ "webtypy>=0.1.7", - "langchain-cloudflare>=0.3.5,<0.3.6", + "langchain-cloudflare>=0.3.5,<0.4.0", # Keep pywrangler sync on the latest Pyodide-compatible 0.3.x core. # scripts/setup_pyodide_deps.sh replaces this with LangChain 1.x wheels. "langchain-core>=0.3.81,<1.0.0", diff --git a/libs/langchain-cloudflare/examples/workers/src/entry.py b/libs/langchain-cloudflare/examples/workers/src/entry.py index 4dd14e5..c640bfe 100644 --- a/libs/langchain-cloudflare/examples/workers/src/entry.py +++ b/libs/langchain-cloudflare/examples/workers/src/entry.py @@ -24,6 +24,7 @@ from langchain_cloudflare import ( ChatCloudflareWorkersAI, + CloudflareAISearchClient, CloudflareAISearchRetriever, CloudflareVectorize, ) @@ -125,6 +126,8 @@ async def fetch(self, request, env): # AI Search endpoint elif path == "ai-search": return await self.handle_ai_search(request) + elif path == "ai-search-admin": + return await self.handle_ai_search_admin(request) # Reranker endpoint elif path == "rerank": return await self.handle_rerank(request) @@ -164,6 +167,7 @@ async def handle_index(self): # Check binding availability vectorize_available = hasattr(self.env, "VECTORIZE") ai_search_available = hasattr(self.env, "AI_SEARCH") + ai_search_admin_available = hasattr(self.env, "AI_SEARCH_ADMIN") d1_available = hasattr(self.env, "D1") return Response.json( @@ -172,6 +176,7 @@ async def handle_index(self): "create_agent_available": CREATE_AGENT_AVAILABLE, "vectorize_available": vectorize_available, "ai_search_available": ai_search_available, + "ai_search_admin_available": ai_search_admin_available, "d1_available": d1_available, "supported_models": SUPPORTED_MODELS, "default_model": DEFAULT_MODEL, @@ -193,6 +198,7 @@ async def handle_index(self): "/vectorize-delete": "Delete documents from Vectorize", "/vectorize-info": "Get Vectorize index info", "/ai-search": "Search AI Search via binding", + "/ai-search-admin": "Manage AI Search via namespace binding", "/rerank": "Rerank documents by query relevance", "/ai-gateway-test": "Test AI Gateway with bindings", "/d1-health": "D1 database health check", @@ -1104,6 +1110,67 @@ async def handle_ai_search(self, request): } ) + async def handle_ai_search_admin(self, request): + """Handle AI Search administration through a namespace binding.""" + data = await request.json() + action = data.get("action", "create-delete") + + if not hasattr(self.env, "AI_SEARCH_ADMIN"): + return Response.json( + {"error": "AI_SEARCH_ADMIN binding not configured"}, + status=400, + ) + + client = CloudflareAISearchClient(binding=self.env.AI_SEARCH_ADMIN) + + if action == "create-delete": + instance_name = data.get("instance_name", "") + if not instance_name: + return Response.json({"error": "instance_name is required"}, status=400) + + try: + created = await client.acreate_instance(instance_name) + listed = await client.alist_instances(search=instance_name) + info = await client.aget_instance(instance_name) + stats = await client.astats(instance_name) + return Response.json( + { + "created_id": created.get("id"), + "listed": any( + item.get("id") == instance_name for item in listed + ), + "info_id": info.get("id"), + "stats_keys": list(stats.keys()), + } + ) + finally: + await client.adelete_instance(instance_name, missing_ok=True) + + if action == "search": + instance_name = data.get("instance_name", "") + query = data.get("query", "") + if not instance_name: + return Response.json({"error": "instance_name is required"}, status=400) + if not query: + return Response.json({"error": "query is required"}, status=400) + + result = await client.asearch( + query, + instance_name=instance_name, + ai_search_options=data.get("ai_search_options"), + ) + chunks = result.get("chunks") or [] + return Response.json( + { + "instance_name": instance_name, + "query": query, + "count": len(chunks), + "chunks": chunks, + } + ) + + return Response.json({"error": f"Unknown action: {action}"}, status=400) + # MARK: - Reranker Handler async def handle_rerank(self, request): diff --git a/libs/langchain-cloudflare/examples/workers/wrangler.jsonc b/libs/langchain-cloudflare/examples/workers/wrangler.jsonc index f0eff81..0b8c2ea 100644 --- a/libs/langchain-cloudflare/examples/workers/wrangler.jsonc +++ b/libs/langchain-cloudflare/examples/workers/wrangler.jsonc @@ -21,6 +21,13 @@ "remote": true } ], + "ai_search_namespaces": [ + { + "binding": "AI_SEARCH_ADMIN", + "namespace": "default", + "remote": true + } + ], "d1_databases": [ { "binding": "D1", diff --git a/libs/langchain-cloudflare/langchain_cloudflare/__init__.py b/libs/langchain-cloudflare/langchain_cloudflare/__init__.py index e003a80..0e3149d 100644 --- a/libs/langchain-cloudflare/langchain_cloudflare/__init__.py +++ b/libs/langchain-cloudflare/langchain_cloudflare/__init__.py @@ -1,6 +1,7 @@ # MARK: - Imports from importlib import metadata +from langchain_cloudflare.ai_search import CloudflareAISearchClient from langchain_cloudflare.bindings import ( # AI Search binding utilities convert_aisearch_response, @@ -33,6 +34,7 @@ # MARK: - Public API __all__ = [ "ChatCloudflareWorkersAI", + "CloudflareAISearchClient", "CloudflareAISearchRetriever", "CloudflareVectorize", "CloudflareWorkersAIEmbeddings", diff --git a/libs/langchain-cloudflare/langchain_cloudflare/_errors.py b/libs/langchain-cloudflare/langchain_cloudflare/_errors.py index c4eae73..ba487d0 100644 --- a/libs/langchain-cloudflare/langchain_cloudflare/_errors.py +++ b/libs/langchain-cloudflare/langchain_cloudflare/_errors.py @@ -49,6 +49,14 @@ class TokenErrors(StrEnum): "'binding' parameter (a dedicated ai_search binding) instead." ) + INSUFFICIENT_AI_SEARCH_ADMIN_TOKENS = ( + "A Cloudflare AI Search API token (with AI Search:Edit and " + "AI Search:Run permissions) must be provided either through the " + "api_token parameter or the CF_AI_SEARCH_API_TOKEN (or CF_API_TOKEN) " + "environment variable. Alternatively, when running in a Python Worker, " + "you can pass an ai_search_namespaces binding." + ) + NO_AI_SEARCH_INSTANCE = ( "An AI Search instance_name must be provided through the instance_name " "parameter or the CF_AI_SEARCH_INSTANCE_NAME environment variable when " diff --git a/libs/langchain-cloudflare/langchain_cloudflare/ai_search.py b/libs/langchain-cloudflare/langchain_cloudflare/ai_search.py new file mode 100644 index 0000000..fd7dbe2 --- /dev/null +++ b/libs/langchain-cloudflare/langchain_cloudflare/ai_search.py @@ -0,0 +1,1093 @@ +"""Cloudflare AI Search administration client.""" + +# MARK: - Imports +from __future__ import annotations + +import inspect +import json +import os +import time +from typing import Any, Dict, List, Optional + +import requests +from langchain_core.utils import from_env +from pydantic import SecretStr + +from ._errors import TokenErrors + +# MARK: - Constants +DEFAULT_BASE_URL = "https://api.cloudflare.com/client/v4" +DEFAULT_NAMESPACE = "default" +DEFAULT_WAIT_SECONDS = 3 +DEFAULT_TIMEOUT_SECONDS = 120 +FAILED_ITEM_STATUSES = {"error"} +DONE_ITEM_STATUSES = {"completed"} + + +# MARK: - Helpers +def _drop_none(data: Dict[str, Any]) -> Dict[str, Any]: + """Return a copy of ``data`` with all ``None``-valued keys removed.""" + return {key: value for key, value in data.items() if value is not None} + + +def _to_py(value: Any) -> Any: + """Convert Pyodide JS proxies into plain Python objects when possible.""" + if hasattr(value, "to_py"): + value = value.to_py() + + if isinstance(value, dict): + return {key: _to_py(item) for key, item in value.items()} + if isinstance(value, list): + return [_to_py(item) for item in value] + return value + + +async def _maybe_await(value: Any) -> Any: + """Await binding return values that are awaitable.""" + if inspect.isawaitable(value): + return await value + return value + + +def _check_api_success(data: Any) -> None: + """Raise when a Cloudflare API envelope reports ``success: false``.""" + if isinstance(data, dict) and data.get("success") is False: + errors = data.get("errors") or data + raise RuntimeError(f"AI Search API request failed: {errors}") + + +def _extract_result(data: Any) -> Any: + """Unwrap a Cloudflare API envelope if one is present.""" + data = _to_py(data) + _check_api_success(data) + if isinstance(data, dict) and "result" in data: + return data["result"] + return data + + +def _as_dict(data: Any) -> Dict[str, Any]: + """Return an unwrapped API response as a dict.""" + result = _extract_result(data) + if result is None: + return {} + if isinstance(result, dict): + return result + return {"value": result} + + +def _as_list(data: Any) -> List[Dict[str, Any]]: + """Return an unwrapped API response as a list of dicts.""" + result = _extract_result(data) + if isinstance(result, list): + return [item for item in result if isinstance(item, dict)] + return [] + + +def _json_or_empty(response: Any) -> Dict[str, Any]: + """Return a JSON response body, or ``{}`` for an empty response.""" + content = getattr(response, "content", b"") + if not content: + return {} + try: + data = response.json() + except ValueError: + return {} + return data if isinstance(data, dict) else {"result": data} + + +# MARK: - CloudflareAISearchClient +class CloudflareAISearchClient: + """Client for managing Cloudflare AI Search instances and uploaded items. + + Use this class for provisioning and administration: create or delete + instances, upload built-in-storage items, inspect stats, and run raw search + requests. For LangChain retrieval chains, use ``CloudflareAISearchRetriever``. + + The REST path uses an API token with ``AI Search:Edit`` and ``AI Search:Run``. + In Python Workers, pass an ``ai_search_namespaces`` binding for namespace + administration. Instance-specific ``ai_search`` bindings can call instance + methods such as ``stats`` and ``items`` but cannot list, create, or delete + instances. + """ + + # MARK: - Init + def __init__( + self, + account_id: Optional[str] = None, + api_token: Optional[str] = None, + instance_name: Optional[str] = None, + namespace: Optional[str] = None, + base_url: str = DEFAULT_BASE_URL, + binding: Any = None, + ) -> None: + """Initialize the AI Search client.""" + self.base_url = base_url.rstrip("/") + self.binding = binding + self.namespace = ( + namespace + if namespace is not None + else from_env("CF_AI_SEARCH_NAMESPACE", default=DEFAULT_NAMESPACE)() + ) + self.instance_name = ( + instance_name + if instance_name is not None + else from_env("CF_AI_SEARCH_INSTANCE_NAME", default="")() + ) + + if account_id is None: + account_id = from_env("CF_ACCOUNT_ID", default="")() + self.account_id = account_id + + token = ( + api_token + or os.environ.get("CF_AI_SEARCH_API_TOKEN") + or os.environ.get("TEST_CF_API_TOKEN") + or os.environ.get("CF_API_TOKEN") + or os.environ.get("CLOUDFLARE_API_TOKEN") + ) + self.api_token = SecretStr(token) if token else None + + if self.binding is not None: + self.headers: Dict[str, str] = {} + return + + if not self.account_id: + raise ValueError(TokenErrors.NO_ACCOUNT_ID_SET) + if not self.api_token or not self.api_token.get_secret_value(): + raise ValueError(TokenErrors.INSUFFICIENT_AI_SEARCH_ADMIN_TOKENS) + + self.headers = { + "Authorization": f"Bearer {self.api_token.get_secret_value()}", + } + + # MARK: - URL Builders + @property + def _ai_search_base_url(self) -> str: + """Return the account-scoped AI Search base URL.""" + return f"{self.base_url}/accounts/{self.account_id}/ai-search" + + def _resolve_namespace(self, namespace: Optional[str] = None) -> str: + """Resolve the namespace for an operation.""" + resolved = namespace if namespace is not None else self.namespace + return resolved or DEFAULT_NAMESPACE + + def _resolve_instance_name(self, instance_name: Optional[str] = None) -> str: + """Resolve the instance name for an operation.""" + resolved = instance_name or self.instance_name + if not resolved: + raise ValueError(TokenErrors.NO_AI_SEARCH_INSTANCE) + return resolved + + def _instances_url(self, namespace: Optional[str] = None) -> str: + """Build the REST URL for instance collection operations.""" + resolved_namespace = self._resolve_namespace(namespace) + return f"{self._ai_search_base_url}/namespaces/{resolved_namespace}/instances" + + def _instance_url( + self, + instance_name: Optional[str] = None, + namespace: Optional[str] = None, + ) -> str: + """Build the REST URL for an instance.""" + resolved_instance = self._resolve_instance_name(instance_name) + return f"{self._instances_url(namespace)}/{resolved_instance}" + + def _items_url( + self, + instance_name: Optional[str] = None, + namespace: Optional[str] = None, + ) -> str: + """Build the REST URL for an instance's items.""" + return f"{self._instance_url(instance_name, namespace)}/items" + + # MARK: - Request Helpers + def _require_rest(self) -> None: + """Raise if a synchronous REST method is called with a binding client.""" + if self.binding is not None: + raise NotImplementedError( + "AI Search bindings are async-only. Use the corresponding " + "`a...` method when passing a Worker binding." + ) + + def _request_json( + self, + method: str, + url: str, + request_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """Issue a REST request and return the JSON API body.""" + self._require_rest() + request_kwargs = dict(request_kwargs or {}) + headers = dict(self.headers) + headers.update(request_kwargs.pop("headers", {})) + response = requests.request( + method, + url, + headers=headers, + **request_kwargs, + **kwargs, + ) + response.raise_for_status() + data = _json_or_empty(response) + _check_api_success(data) + return data + + async def _arequest_json( + self, + method: str, + url: str, + request_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """Issue an async REST request and return the JSON API body.""" + self._require_rest() + import httpx + + request_kwargs = dict(request_kwargs or {}) + headers = dict(self.headers) + headers.update(request_kwargs.pop("headers", {})) + async with httpx.AsyncClient() as client: + response = await client.request( + method, + url, + headers=headers, + **request_kwargs, + **kwargs, + ) + response.raise_for_status() + data = _json_or_empty(response) + _check_api_success(data) + return data + + def _binding_instance(self, instance_name: Optional[str] = None) -> Any: + """Return an AI Search instance handle from a Worker binding.""" + if self.binding is None: + raise ValueError("An AI Search binding is required") + + resolved = instance_name or self.instance_name + if hasattr(self.binding, "get"): + return self.binding.get(self._resolve_instance_name(resolved)) + + if resolved and self.instance_name and resolved != self.instance_name: + raise ValueError( + "Instance-specific AI Search bindings cannot switch instances. " + "Pass an ai_search_namespaces binding for multi-instance access." + ) + + return self.binding + + async def _abinding_instance(self, instance_name: Optional[str] = None) -> Any: + """Asynchronously return an AI Search instance handle from a binding.""" + if self.binding is None: + raise ValueError("An AI Search binding is required") + + resolved = instance_name or self.instance_name + if hasattr(self.binding, "get"): + return await _maybe_await( + self.binding.get(self._resolve_instance_name(resolved)) + ) + + if resolved and self.instance_name and resolved != self.instance_name: + raise ValueError( + "Instance-specific AI Search bindings cannot switch instances. " + "Pass an ai_search_namespaces binding for multi-instance access." + ) + + return self.binding + + def _require_namespace_binding(self, method_name: str) -> Any: + """Return a namespace-binding method or raise a clear error.""" + method = getattr(self.binding, method_name, None) + if method is None: + raise NotImplementedError( + "This operation requires an ai_search_namespaces binding. " + "Instance-specific ai_search bindings do not support " + f"`{method_name}`." + ) + return method + + def _build_instance_body( + self, + instance_name: Optional[str], + instance_config: Optional[Dict[str, Any]], + extra_config: Dict[str, Any], + ) -> Dict[str, Any]: + """Build an instance create/update payload.""" + body = dict(instance_config or {}) + body.update(extra_config) + if instance_name is not None: + body["id"] = instance_name + if "id" not in body: + body["id"] = self._resolve_instance_name(None) + return body + + # MARK: - Instances + def list_instances( + self, + *, + namespace: Optional[str] = None, + page: Optional[int] = None, + per_page: Optional[int] = None, + search: Optional[str] = None, + order_by: Optional[str] = None, + order_by_direction: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> List[Dict[str, Any]]: + """List AI Search instances.""" + params = _drop_none( + { + "page": page, + "per_page": per_page, + "search": search, + "order_by": order_by, + "order_by_direction": order_by_direction, + } + ) + data = self._request_json( + "GET", + self._instances_url(namespace), + request_kwargs=request_kwargs, + params=params or None, + ) + return _as_list(data) + + async def alist_instances( + self, + *, + namespace: Optional[str] = None, + page: Optional[int] = None, + per_page: Optional[int] = None, + search: Optional[str] = None, + order_by: Optional[str] = None, + order_by_direction: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> List[Dict[str, Any]]: + """Asynchronously list AI Search instances.""" + if self.binding is not None: + list_method = self._require_namespace_binding("list") + params = _drop_none( + { + "page": page, + "per_page": per_page, + "search": search, + "order_by": order_by, + "order_by_direction": order_by_direction, + } + ) + from .bindings import convert_payload_for_binding + + response = await _maybe_await( + list_method(convert_payload_for_binding(params)) + ) + return _as_list(response) + + params = _drop_none( + { + "page": page, + "per_page": per_page, + "search": search, + "order_by": order_by, + "order_by_direction": order_by_direction, + } + ) + data = await self._arequest_json( + "GET", + self._instances_url(namespace), + request_kwargs=request_kwargs, + params=params or None, + ) + return _as_list(data) + + def create_instance( + self, + instance_name: Optional[str] = None, + *, + namespace: Optional[str] = None, + instance_config: Optional[Dict[str, Any]] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """Create an AI Search instance.""" + body = self._build_instance_body(instance_name, instance_config, kwargs) + data = self._request_json( + "POST", + self._instances_url(namespace), + request_kwargs=request_kwargs, + json=body, + ) + return _as_dict(data) + + async def acreate_instance( + self, + instance_name: Optional[str] = None, + *, + namespace: Optional[str] = None, + instance_config: Optional[Dict[str, Any]] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """Asynchronously create an AI Search instance.""" + body = self._build_instance_body(instance_name, instance_config, kwargs) + if self.binding is not None: + create_method = self._require_namespace_binding("create") + from .bindings import convert_payload_for_binding + + instance = await _maybe_await( + create_method(convert_payload_for_binding(body)) + ) + if hasattr(instance, "info"): + return _as_dict(await _maybe_await(instance.info())) + return _as_dict(instance) + + data = await self._arequest_json( + "POST", + self._instances_url(namespace), + request_kwargs=request_kwargs, + json=body, + ) + return _as_dict(data) + + def get_instance( + self, + instance_name: Optional[str] = None, + *, + namespace: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Read an AI Search instance.""" + data = self._request_json( + "GET", + self._instance_url(instance_name, namespace), + request_kwargs=request_kwargs, + ) + return _as_dict(data) + + async def aget_instance( + self, + instance_name: Optional[str] = None, + *, + namespace: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Asynchronously read an AI Search instance.""" + if self.binding is not None: + instance = await self._abinding_instance(instance_name) + return _as_dict(await _maybe_await(instance.info())) + + data = await self._arequest_json( + "GET", + self._instance_url(instance_name, namespace), + request_kwargs=request_kwargs, + ) + return _as_dict(data) + + def update_instance( + self, + instance_name: Optional[str] = None, + *, + namespace: Optional[str] = None, + instance_config: Optional[Dict[str, Any]] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """Update an AI Search instance.""" + body = dict(instance_config or {}) + body.update(kwargs) + data = self._request_json( + "PUT", + self._instance_url(instance_name, namespace), + request_kwargs=request_kwargs, + json=body, + ) + return _as_dict(data) + + async def aupdate_instance( + self, + instance_name: Optional[str] = None, + *, + namespace: Optional[str] = None, + instance_config: Optional[Dict[str, Any]] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """Asynchronously update an AI Search instance.""" + body = dict(instance_config or {}) + body.update(kwargs) + if self.binding is not None: + instance = await self._abinding_instance(instance_name) + from .bindings import convert_payload_for_binding + + response = await _maybe_await( + instance.update(convert_payload_for_binding(body)) + ) + return _as_dict(response) + + data = await self._arequest_json( + "PUT", + self._instance_url(instance_name, namespace), + request_kwargs=request_kwargs, + json=body, + ) + return _as_dict(data) + + def delete_instance( + self, + instance_name: Optional[str] = None, + *, + namespace: Optional[str] = None, + missing_ok: bool = False, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Delete an AI Search instance.""" + try: + data = self._request_json( + "DELETE", + self._instance_url(instance_name, namespace), + request_kwargs=request_kwargs, + ) + except requests.HTTPError as exc: + if ( + missing_ok + and exc.response is not None + and exc.response.status_code == 404 + ): + return {} + raise + return _as_dict(data) + + async def adelete_instance( + self, + instance_name: Optional[str] = None, + *, + namespace: Optional[str] = None, + missing_ok: bool = False, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Asynchronously delete an AI Search instance.""" + if self.binding is not None: + delete_method = self._require_namespace_binding("delete") + try: + await _maybe_await( + delete_method(self._resolve_instance_name(instance_name)) + ) + except Exception: + if missing_ok: + return {} + raise + return {} + + import httpx + + try: + data = await self._arequest_json( + "DELETE", + self._instance_url(instance_name, namespace), + request_kwargs=request_kwargs, + ) + except httpx.HTTPStatusError as exc: + if missing_ok and exc.response.status_code == 404: + return {} + raise + return _as_dict(data) + + def stats( + self, + instance_name: Optional[str] = None, + *, + namespace: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Return AI Search instance indexing stats.""" + data = self._request_json( + "GET", + f"{self._instance_url(instance_name, namespace)}/stats", + request_kwargs=request_kwargs, + ) + return _as_dict(data) + + async def astats( + self, + instance_name: Optional[str] = None, + *, + namespace: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Asynchronously return AI Search instance indexing stats.""" + if self.binding is not None: + instance = await self._abinding_instance(instance_name) + return _as_dict(await _maybe_await(instance.stats())) + + data = await self._arequest_json( + "GET", + f"{self._instance_url(instance_name, namespace)}/stats", + request_kwargs=request_kwargs, + ) + return _as_dict(data) + + # MARK: - Items + def _upload_data( + self, + *, + metadata: Optional[Dict[str, Any]] = None, + wait_for_completion: Optional[bool] = None, + ) -> Dict[str, str]: + """Build multipart form fields for item upload.""" + data: Dict[str, str] = {} + if metadata is not None: + data["metadata"] = json.dumps(metadata) + if wait_for_completion is not None: + data["wait_for_completion"] = str(wait_for_completion).lower() + return data + + def upload_item( + self, + filename: str, + content: Any, + *, + content_type: str = "text/plain", + metadata: Optional[Dict[str, Any]] = None, + wait_for_completion: Optional[bool] = None, + instance_name: Optional[str] = None, + namespace: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Upload a file to an AI Search instance's built-in storage.""" + upload_content = ( + content.encode("utf-8") if isinstance(content, str) else content + ) + file_value = (filename, upload_content, content_type) + data = self._request_json( + "POST", + self._items_url(instance_name, namespace), + request_kwargs=request_kwargs, + files={"file": file_value}, + data=self._upload_data( + metadata=metadata, + wait_for_completion=wait_for_completion, + ) + or None, + ) + return _as_dict(data) + + async def aupload_item( + self, + filename: str, + content: Any, + *, + content_type: str = "text/plain", + metadata: Optional[Dict[str, Any]] = None, + wait_for_completion: Optional[bool] = None, + poll_interval_ms: Optional[int] = None, + timeout_ms: Optional[int] = None, + instance_name: Optional[str] = None, + namespace: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Asynchronously upload a file to an AI Search instance.""" + if self.binding is not None: + instance = await self._abinding_instance(instance_name) + options = _drop_none( + { + "metadata": metadata, + "pollIntervalMs": poll_interval_ms, + "timeoutMs": timeout_ms, + } + ) + from .bindings import convert_payload_for_binding + + if wait_for_completion: + upload_method = instance.items.uploadAndPoll + else: + upload_method = instance.items.upload + + if options: + response = await _maybe_await( + upload_method( + filename, content, convert_payload_for_binding(options) + ) + ) + else: + response = await _maybe_await(upload_method(filename, content)) + return _as_dict(response) + + upload_content = ( + content.encode("utf-8") if isinstance(content, str) else content + ) + file_value = (filename, upload_content, content_type) + data = await self._arequest_json( + "POST", + self._items_url(instance_name, namespace), + request_kwargs=request_kwargs, + files={"file": file_value}, + data=self._upload_data( + metadata=metadata, + wait_for_completion=wait_for_completion, + ) + or None, + ) + return _as_dict(data) + + def list_items( + self, + instance_name: Optional[str] = None, + *, + namespace: Optional[str] = None, + page: Optional[int] = None, + per_page: Optional[int] = None, + status: Optional[str] = None, + sort_by: Optional[str] = None, + search: Optional[str] = None, + source: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> List[Dict[str, Any]]: + """List uploaded items in an AI Search instance.""" + params = _drop_none( + { + "page": page, + "per_page": per_page, + "status": status, + "sort_by": sort_by, + "search": search, + "source": source, + } + ) + data = self._request_json( + "GET", + self._items_url(instance_name, namespace), + request_kwargs=request_kwargs, + params=params or None, + ) + return _as_list(data) + + async def alist_items( + self, + instance_name: Optional[str] = None, + *, + namespace: Optional[str] = None, + page: Optional[int] = None, + per_page: Optional[int] = None, + status: Optional[str] = None, + sort_by: Optional[str] = None, + search: Optional[str] = None, + source: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> List[Dict[str, Any]]: + """Asynchronously list uploaded items in an AI Search instance.""" + if self.binding is not None: + instance = await self._abinding_instance(instance_name) + params = _drop_none( + { + "page": page, + "per_page": per_page, + "status": status, + "sort_by": sort_by, + "search": search, + "source": source, + } + ) + from .bindings import convert_payload_for_binding + + response = await _maybe_await( + instance.items.list(convert_payload_for_binding(params)) + ) + return _as_list(response) + + params = _drop_none( + { + "page": page, + "per_page": per_page, + "status": status, + "sort_by": sort_by, + "search": search, + "source": source, + } + ) + data = await self._arequest_json( + "GET", + self._items_url(instance_name, namespace), + request_kwargs=request_kwargs, + params=params or None, + ) + return _as_list(data) + + def get_item( + self, + item_id: str, + *, + instance_name: Optional[str] = None, + namespace: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Read item status and metadata.""" + data = self._request_json( + "GET", + f"{self._items_url(instance_name, namespace)}/{item_id}", + request_kwargs=request_kwargs, + ) + return _as_dict(data) + + async def aget_item( + self, + item_id: str, + *, + instance_name: Optional[str] = None, + namespace: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Asynchronously read item status and metadata.""" + if self.binding is not None: + instance = await self._abinding_instance(instance_name) + item = await _maybe_await(instance.items.get(item_id)) + return _as_dict(await _maybe_await(item.info())) + + data = await self._arequest_json( + "GET", + f"{self._items_url(instance_name, namespace)}/{item_id}", + request_kwargs=request_kwargs, + ) + return _as_dict(data) + + def delete_item( + self, + item_id: str, + *, + instance_name: Optional[str] = None, + namespace: Optional[str] = None, + missing_ok: bool = False, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Delete an uploaded item and its indexed chunks.""" + try: + data = self._request_json( + "DELETE", + f"{self._items_url(instance_name, namespace)}/{item_id}", + request_kwargs=request_kwargs, + ) + except requests.HTTPError as exc: + if ( + missing_ok + and exc.response is not None + and exc.response.status_code == 404 + ): + return {} + raise + return _as_dict(data) + + async def adelete_item( + self, + item_id: str, + *, + instance_name: Optional[str] = None, + namespace: Optional[str] = None, + missing_ok: bool = False, + request_kwargs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Asynchronously delete an uploaded item and its indexed chunks.""" + if self.binding is not None: + instance = await self._abinding_instance(instance_name) + try: + await _maybe_await(instance.items.delete(item_id)) + except Exception: + if missing_ok: + return {} + raise + return {} + + import httpx + + try: + data = await self._arequest_json( + "DELETE", + f"{self._items_url(instance_name, namespace)}/{item_id}", + request_kwargs=request_kwargs, + ) + except httpx.HTTPStatusError as exc: + if missing_ok and exc.response.status_code == 404: + return {} + raise + return _as_dict(data) + + def wait_for_item( + self, + item_id: str, + *, + instance_name: Optional[str] = None, + namespace: Optional[str] = None, + timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS, + poll_interval_seconds: int = DEFAULT_WAIT_SECONDS, + ) -> Dict[str, Any]: + """Wait until an uploaded item has finished indexing.""" + deadline = time.time() + timeout_seconds + last_item: Dict[str, Any] = {} + while time.time() < deadline: + last_item = self.get_item( + item_id, + instance_name=instance_name, + namespace=namespace, + ) + status = last_item.get("status") + if status in DONE_ITEM_STATUSES: + return last_item + if status in FAILED_ITEM_STATUSES: + raise RuntimeError(f"AI Search item indexing failed: {last_item}") + time.sleep(poll_interval_seconds) + + raise TimeoutError( + f"AI Search item {item_id!r} did not finish indexing within " + f"{timeout_seconds} seconds. Last item: {last_item}" + ) + + async def await_for_item( + self, + item_id: str, + *, + instance_name: Optional[str] = None, + namespace: Optional[str] = None, + timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS, + poll_interval_seconds: int = DEFAULT_WAIT_SECONDS, + ) -> Dict[str, Any]: + """Asynchronously wait until an uploaded item has finished indexing.""" + import asyncio + + deadline = time.time() + timeout_seconds + last_item: Dict[str, Any] = {} + while time.time() < deadline: + last_item = await self.aget_item( + item_id, + instance_name=instance_name, + namespace=namespace, + ) + status = last_item.get("status") + if status in DONE_ITEM_STATUSES: + return last_item + if status in FAILED_ITEM_STATUSES: + raise RuntimeError(f"AI Search item indexing failed: {last_item}") + await asyncio.sleep(poll_interval_seconds) + + raise TimeoutError( + f"AI Search item {item_id!r} did not finish indexing within " + f"{timeout_seconds} seconds. Last item: {last_item}" + ) + + # MARK: - Query + def _build_search_body( + self, + query: Optional[str], + messages: Optional[List[Dict[str, str]]], + ai_search_options: Optional[Dict[str, Any]], + extra_body: Dict[str, Any], + ) -> Dict[str, Any]: + """Build an AI Search search request body.""" + if query and messages: + raise ValueError("Provide either query or messages, not both") + + body = dict(extra_body) + if query is not None: + body["query"] = query + elif messages is not None: + body["messages"] = messages + + if ai_search_options is not None: + body["ai_search_options"] = ai_search_options + + return body + + def search( + self, + query: Optional[str] = None, + *, + messages: Optional[List[Dict[str, str]]] = None, + ai_search_options: Optional[Dict[str, Any]] = None, + instance_name: Optional[str] = None, + namespace: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """Run a raw AI Search query against an instance.""" + data = self._request_json( + "POST", + f"{self._instance_url(instance_name, namespace)}/search", + request_kwargs=request_kwargs, + json=self._build_search_body(query, messages, ai_search_options, kwargs), + ) + return _as_dict(data) + + async def asearch( + self, + query: Optional[str] = None, + *, + messages: Optional[List[Dict[str, str]]] = None, + ai_search_options: Optional[Dict[str, Any]] = None, + instance_name: Optional[str] = None, + namespace: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """Asynchronously run a raw AI Search query against an instance.""" + body = self._build_search_body(query, messages, ai_search_options, kwargs) + if self.binding is not None: + instance = await self._abinding_instance(instance_name) + from .bindings import convert_aisearch_response, convert_payload_for_binding + + response = await _maybe_await( + instance.search(convert_payload_for_binding(body)) + ) + return _as_dict(convert_aisearch_response(response)) + + data = await self._arequest_json( + "POST", + f"{self._instance_url(instance_name, namespace)}/search", + request_kwargs=request_kwargs, + json=body, + ) + return _as_dict(data) + + def chat_completions( + self, + messages: List[Dict[str, str]], + *, + instance_name: Optional[str] = None, + namespace: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """Run an AI Search chat-completions request.""" + body = dict(kwargs) + body["messages"] = messages + data = self._request_json( + "POST", + f"{self._instance_url(instance_name, namespace)}/chat/completions", + request_kwargs=request_kwargs, + json=body, + ) + return _as_dict(data) + + async def achat_completions( + self, + messages: List[Dict[str, str]], + *, + instance_name: Optional[str] = None, + namespace: Optional[str] = None, + request_kwargs: Optional[Dict[str, Any]] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """Asynchronously run an AI Search chat-completions request.""" + body = dict(kwargs) + body["messages"] = messages + if self.binding is not None: + instance = await self._abinding_instance(instance_name) + from .bindings import convert_payload_for_binding + + response = await _maybe_await( + instance.chatCompletions(convert_payload_for_binding(body)) + ) + return _as_dict(response) + + data = await self._arequest_json( + "POST", + f"{self._instance_url(instance_name, namespace)}/chat/completions", + request_kwargs=request_kwargs, + json=body, + ) + return _as_dict(data) diff --git a/libs/langchain-cloudflare/pyproject.toml b/libs/langchain-cloudflare/pyproject.toml index 92560d7..2ccfa81 100644 --- a/libs/langchain-cloudflare/pyproject.toml +++ b/libs/langchain-cloudflare/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "langchain-cloudflare" -version = "0.3.6" +version = "0.3.7" description = "Langchain Integrations for Cloudflare's WorkersAI and Vectorize" readme = "README.md" license = "MIT" diff --git a/libs/langchain-cloudflare/tests/conftest.py b/libs/langchain-cloudflare/tests/conftest.py index 316d8e3..603405a 100644 --- a/libs/langchain-cloudflare/tests/conftest.py +++ b/libs/langchain-cloudflare/tests/conftest.py @@ -476,6 +476,7 @@ def _upload_ai_search_fixture_documents( "POST", f"{instance_url}/items", files=files, + data={"wait_for_completion": "true"}, ) result = data.get("result") or {} item_id = result.get("id") @@ -513,7 +514,7 @@ def _search_ai_search_fixture( def _wait_for_ai_search_fixture( session: requests.Session, instance_url: str, - timeout_seconds: int = 120, + timeout_seconds: int = 240, ) -> None: """Wait until the fixture documents are indexed and searchable.""" deadline = time.time() + timeout_seconds diff --git a/libs/langchain-cloudflare/tests/integration_tests/test_ai_search_client.py b/libs/langchain-cloudflare/tests/integration_tests/test_ai_search_client.py new file mode 100644 index 0000000..0a7d086 --- /dev/null +++ b/libs/langchain-cloudflare/tests/integration_tests/test_ai_search_client.py @@ -0,0 +1,99 @@ +"""Integration tests for CloudflareAISearchClient. + +These tests create a temporary AI Search instance, upload one fixture document, +query it, and delete the instance during teardown. + +Required environment variables: + - CF_ACCOUNT_ID + - CF_AI_SEARCH_API_TOKEN (or TEST_CF_API_TOKEN or CF_API_TOKEN) +""" + +import os +import uuid + +import pytest + +from langchain_cloudflare.ai_search import CloudflareAISearchClient + +_ACCOUNT_ID = os.environ.get("CF_ACCOUNT_ID") or os.environ.get("CLOUDFLARE_ACCOUNT_ID") +_API_TOKEN = ( + os.environ.get("CF_AI_SEARCH_API_TOKEN") + or os.environ.get("TEST_CF_API_TOKEN") + or os.environ.get("CF_API_TOKEN") + or os.environ.get("CLOUDFLARE_API_TOKEN") +) +_NAMESPACE = os.environ.get("CF_AI_SEARCH_NAMESPACE", "default") +_HAS_CREDS = bool(_ACCOUNT_ID and _API_TOKEN) + +pytestmark = pytest.mark.skipif( + not _HAS_CREDS, + reason="AI Search credentials not configured", +) + + +def test_ai_search_client_instance_item_lifecycle() -> None: + """Create an instance, upload/search an item, and delete the instance.""" + instance_name = f"langchain-cloudflare-client-{uuid.uuid4().hex[:8]}" + query = f"langchaincloudflareclientfixture {uuid.uuid4().hex}" + item_id = "" + + client = CloudflareAISearchClient( + account_id=_ACCOUNT_ID, + api_token=_API_TOKEN, + namespace=_NAMESPACE, + ) + + try: + created = client.create_instance(instance_name) + assert created["id"] == instance_name + + listed = client.list_instances(search=instance_name) + assert any(instance.get("id") == instance_name for instance in listed) + + stats = client.stats(instance_name) + assert isinstance(stats, dict) + + item = client.upload_item( + f"{instance_name}.md", + "\n".join( + [ + "# LangChain Cloudflare AI Search client fixture", + "", + f"{query} validates temporary instance lifecycle tests.", + ] + ), + content_type="text/markdown", + metadata={"suite": "langchain-cloudflare"}, + wait_for_completion=True, + instance_name=instance_name, + ) + item_id = item["id"] + if item.get("status") != "completed": + item = client.wait_for_item(item_id, instance_name=instance_name) + assert item["status"] == "completed" + + items = client.list_items(instance_name) + assert any(uploaded.get("id") == item_id for uploaded in items) + + result = client.search( + query, + instance_name=instance_name, + ai_search_options={ + "retrieval": { + "max_num_results": 1, + }, + "query_rewrite": {"enabled": False}, + "reranking": {"enabled": False}, + }, + ) + chunks = result.get("chunks") or [] + assert chunks + assert query in chunks[0].get("text", "") + finally: + if item_id: + client.delete_item( + item_id, + instance_name=instance_name, + missing_ok=True, + ) + client.delete_instance(instance_name, missing_ok=True) diff --git a/libs/langchain-cloudflare/tests/unit_tests/test_ai_search.py b/libs/langchain-cloudflare/tests/unit_tests/test_ai_search.py new file mode 100644 index 0000000..cba699c --- /dev/null +++ b/libs/langchain-cloudflare/tests/unit_tests/test_ai_search.py @@ -0,0 +1,329 @@ +"""Unit tests for CloudflareAISearchClient (offline: no HTTP is issued).""" + +import re +from typing import Any + +import pytest +import requests + +from langchain_cloudflare._errors import TokenErrors +from langchain_cloudflare.ai_search import CloudflareAISearchClient + + +class FakeResponse: + """Small requests/httpx response stub.""" + + def __init__(self, data: Any, status_code: int = 200): + self._data = data + self.status_code = status_code + self.content = b"{}" if data is not None else b"" + + def json(self) -> Any: + return self._data + + def raise_for_status(self) -> None: + if self.status_code >= 400: + raise requests.HTTPError(response=self) + + +class RequestRecorder: + """Capture requests.request calls and return a Cloudflare API envelope.""" + + def __init__(self, result: Any): + self.result = result + self.calls: list[dict[str, Any]] = [] + + def __call__(self, method: str, url: str, **kwargs: Any) -> FakeResponse: + self.calls.append({"method": method, "url": url, **kwargs}) + return FakeResponse({"success": True, "result": self.result}) + + +def _make_client(**overrides: Any) -> CloudflareAISearchClient: + """Construct a client with valid dummy REST credentials.""" + params = { + "account_id": "abc123", + "api_token": "valid-token", + "instance_name": "test-instance", + } + params.update(overrides) + return CloudflareAISearchClient(**params) + + +# MARK: - REST Client Tests +class TestRESTClient: + """Test REST credential handling, URLs, and request bodies.""" + + def test_missing_account_id_raises(self) -> None: + """Missing account_id should raise ValueError.""" + with pytest.raises(ValueError, match="account ID"): + CloudflareAISearchClient(account_id="", api_token="tok") + + def test_missing_token_raises(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Missing token should raise the AI Search admin-token error.""" + for key in ( + "CF_AI_SEARCH_API_TOKEN", + "TEST_CF_API_TOKEN", + "CF_API_TOKEN", + "CLOUDFLARE_API_TOKEN", + ): + monkeypatch.delenv(key, raising=False) + + with pytest.raises( + ValueError, + match=re.escape(str(TokenErrors.INSUFFICIENT_AI_SEARCH_ADMIN_TOKENS)), + ): + CloudflareAISearchClient(account_id="abc123", api_token="") + + def test_create_instance_default_namespace( + self, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Default namespace uses the namespace-scoped collection endpoint.""" + recorder = RequestRecorder({"id": "docs"}) + monkeypatch.setattr(requests, "request", recorder) + + result = _make_client().create_instance( + "docs", + type="web-crawler", + source="developers.cloudflare.com", + ) + + assert result == {"id": "docs"} + assert recorder.calls[0]["method"] == "POST" + assert recorder.calls[0]["url"] == ( + "https://api.cloudflare.com/client/v4/accounts/abc123/" + "ai-search/namespaces/default/instances" + ) + assert recorder.calls[0]["json"] == { + "id": "docs", + "type": "web-crawler", + "source": "developers.cloudflare.com", + } + assert recorder.calls[0]["headers"] == {"Authorization": "Bearer valid-token"} + + def test_get_instance_namespace_url(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Non-default namespace uses the namespace-scoped endpoint.""" + recorder = RequestRecorder({"id": "docs"}) + monkeypatch.setattr(requests, "request", recorder) + + result = _make_client(namespace="tenant-a").get_instance("docs") + + assert result == {"id": "docs"} + assert recorder.calls[0]["url"] == ( + "https://api.cloudflare.com/client/v4/accounts/abc123/" + "ai-search/namespaces/tenant-a/instances/docs" + ) + + def test_list_instances_params(self, monkeypatch: pytest.MonkeyPatch) -> None: + """List instances sends pagination and search query params.""" + recorder = RequestRecorder([{"id": "docs"}]) + monkeypatch.setattr(requests, "request", recorder) + + result = _make_client().list_instances(page=2, per_page=10, search="docs") + + assert result == [{"id": "docs"}] + assert recorder.calls[0]["method"] == "GET" + assert recorder.calls[0]["params"] == { + "page": 2, + "per_page": 10, + "search": "docs", + } + + def test_upload_item_body(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Item upload uses multipart form fields and leaves Content-Type unset.""" + recorder = RequestRecorder({"id": "item-1", "key": "docs.md"}) + monkeypatch.setattr(requests, "request", recorder) + + result = _make_client().upload_item( + "docs.md", + "# Docs", + content_type="text/markdown", + metadata={"category": "docs"}, + wait_for_completion=True, + ) + + assert result["id"] == "item-1" + call = recorder.calls[0] + assert call["url"].endswith( + "/ai-search/namespaces/default/instances/test-instance/items" + ) + assert call["files"]["file"] == ("docs.md", b"# Docs", "text/markdown") + assert call["data"] == { + "metadata": '{"category": "docs"}', + "wait_for_completion": "true", + } + assert "Content-Type" not in call["headers"] + + def test_search_body(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Search builds a raw query body with AI Search options.""" + recorder = RequestRecorder({"chunks": []}) + monkeypatch.setattr(requests, "request", recorder) + + _make_client().search( + "hello", + ai_search_options={"retrieval": {"max_num_results": 3}}, + ) + + assert recorder.calls[0]["json"] == { + "query": "hello", + "ai_search_options": {"retrieval": {"max_num_results": 3}}, + } + + def test_search_rejects_query_and_messages(self) -> None: + """Search should not accept both query shapes at once.""" + with pytest.raises(ValueError, match="either query or messages"): + _make_client().search("hello", messages=[{"role": "user", "content": "hi"}]) + + +# MARK: - Binding Test Fakes +class FakeItemHandle: + """Fake item handle returned by items.get().""" + + def __init__(self, item_id: str): + self.item_id = item_id + + async def info(self) -> dict[str, Any]: + return {"id": self.item_id, "status": "completed"} + + +class FakeItems: + """Fake AI Search items binding.""" + + def __init__(self) -> None: + self.deleted: list[str] = [] + self.uploads: list[tuple[Any, ...]] = [] + + async def upload(self, *args: Any) -> dict[str, Any]: + self.uploads.append(args) + return {"id": "item-1", "key": args[0]} + + async def uploadAndPoll(self, *args: Any) -> dict[str, Any]: + self.uploads.append(args) + return {"id": "item-1", "key": args[0], "status": "completed"} + + async def list(self, params: dict[str, Any]) -> dict[str, Any]: + return {"result": [{"id": "item-1", "status": params.get("status")}]} + + def get(self, item_id: str) -> FakeItemHandle: + return FakeItemHandle(item_id) + + async def delete(self, item_id: str) -> None: + self.deleted.append(item_id) + + +class FakeInstance: + """Fake AI Search instance binding.""" + + def __init__(self, instance_id: str): + self.instance_id = instance_id + self.items = FakeItems() + self.updated: dict[str, Any] = {} + + async def info(self) -> dict[str, Any]: + return {"id": self.instance_id, "status": "active"} + + async def stats(self) -> dict[str, Any]: + return {"completed": 1, "queued": 0} + + async def update(self, payload: dict[str, Any]) -> dict[str, Any]: + self.updated = payload + return {"id": self.instance_id, **payload} + + async def search(self, payload: dict[str, Any]) -> dict[str, Any]: + return {"chunks": [{"text": payload.get("query", "")}]} + + async def chatCompletions(self, payload: dict[str, Any]) -> dict[str, Any]: + return {"messages": payload["messages"]} + + +class FakeNamespaceBinding: + """Fake ai_search_namespaces binding.""" + + def __init__(self) -> None: + self.instances: dict[str, FakeInstance] = {} + self.deleted: list[str] = [] + self.list_params: dict[str, Any] = {} + + def get(self, instance_name: str) -> FakeInstance: + return self.instances.setdefault(instance_name, FakeInstance(instance_name)) + + async def create(self, payload: dict[str, Any]) -> FakeInstance: + instance = FakeInstance(payload["id"]) + self.instances[payload["id"]] = instance + return instance + + async def list(self, params: dict[str, Any]) -> dict[str, Any]: + self.list_params = params + return {"result": [{"id": instance_id} for instance_id in self.instances]} + + async def delete(self, instance_name: str) -> None: + self.deleted.append(instance_name) + self.instances.pop(instance_name, None) + + +# MARK: - Binding Tests +class TestBindingClient: + """Test async Worker binding behavior.""" + + async def test_namespace_binding_lifecycle(self) -> None: + """Namespace binding supports create/list/get/delete.""" + binding = FakeNamespaceBinding() + client = CloudflareAISearchClient(binding=binding) + + created = await client.acreate_instance("docs", type="r2", source="bucket") + listed = await client.alist_instances(search="docs") + info = await client.aget_instance("docs") + await client.adelete_instance("docs") + + assert created == {"id": "docs", "status": "active"} + assert listed == [{"id": "docs"}] + assert binding.list_params == {"search": "docs"} + assert info == {"id": "docs", "status": "active"} + assert binding.deleted == ["docs"] + + async def test_binding_item_methods(self) -> None: + """Binding client can upload, list, get, and delete items.""" + binding = FakeNamespaceBinding() + client = CloudflareAISearchClient(binding=binding) + await client.acreate_instance("docs") + + uploaded = await client.aupload_item( + "docs.md", + "# Docs", + instance_name="docs", + wait_for_completion=True, + metadata={"category": "docs"}, + ) + items = await client.alist_items("docs", status="completed") + item = await client.aget_item("item-1", instance_name="docs") + await client.adelete_item("item-1", instance_name="docs") + + assert uploaded["status"] == "completed" + assert items == [{"id": "item-1", "status": "completed"}] + assert item == {"id": "item-1", "status": "completed"} + assert binding.get("docs").items.deleted == ["item-1"] + + async def test_binding_search_and_chat(self) -> None: + """Binding client can run query and chat methods through an instance.""" + binding = FakeNamespaceBinding() + client = CloudflareAISearchClient(binding=binding) + await client.acreate_instance("docs") + + search = await client.asearch("hello", instance_name="docs") + chat = await client.achat_completions( + [{"role": "user", "content": "hello"}], + instance_name="docs", + ) + + assert search == {"chunks": [{"text": "hello"}]} + assert chat == {"messages": [{"role": "user", "content": "hello"}]} + + async def test_instance_binding_rejects_namespace_operations(self) -> None: + """Instance-specific bindings cannot create/list/delete instances.""" + client = CloudflareAISearchClient( + binding=FakeInstance("docs"), + instance_name="docs", + ) + + with pytest.raises(NotImplementedError, match="ai_search_namespaces"): + await client.acreate_instance("other") diff --git a/libs/langchain-cloudflare/tests/worker_tests/test_worker_integration.py b/libs/langchain-cloudflare/tests/worker_tests/test_worker_integration.py index 1c7f058..1aaff1d 100644 --- a/libs/langchain-cloudflare/tests/worker_tests/test_worker_integration.py +++ b/libs/langchain-cloudflare/tests/worker_tests/test_worker_integration.py @@ -635,6 +635,64 @@ def test_ai_search_missing_query(self, dev_server): data = response.json() assert "error" in data + def test_ai_search_admin_create_delete(self, dev_server): + """POST /ai-search-admin should manage a temporary instance.""" + port = dev_server + instance_name = f"langchain-cloudflare-worker-{uuid.uuid4().hex[:8]}" + response = requests.post( + f"http://localhost:{port}/ai-search-admin", + json={ + "action": "create-delete", + "instance_name": instance_name, + }, + headers={"Content-Type": "application/json"}, + ) + + if response.status_code in (400, 500): + data = response.json() + if "AI_SEARCH_ADMIN binding not configured" in data.get("error", ""): + pytest.skip("AI Search namespace binding not configured") + + assert response.status_code == 200, response.text + data = response.json() + assert data["created_id"] == instance_name + assert data["listed"] is True + assert data["info_id"] == instance_name + assert isinstance(data["stats_keys"], list) + + def test_ai_search_admin_namespace_search(self, ai_search_test_data, dev_server): + """POST /ai-search-admin should search via namespace binding.""" + port = dev_server + response = requests.post( + f"http://localhost:{port}/ai-search-admin", + json={ + "action": "search", + "instance_name": ai_search_test_data["instance_name"], + "query": ai_search_test_data["query"], + "ai_search_options": { + "retrieval": { + "max_num_results": 3, + "retrieval_type": "hybrid", + }, + "query_rewrite": {"enabled": False}, + "reranking": {"enabled": False}, + }, + }, + headers={"Content-Type": "application/json"}, + ) + + if response.status_code in (400, 500): + data = response.json() + if "AI_SEARCH_ADMIN binding not configured" in data.get("error", ""): + pytest.skip("AI Search namespace binding not configured") + + assert response.status_code == 200, response.text + data = response.json() + assert data["instance_name"] == ai_search_test_data["instance_name"] + assert data["query"] == ai_search_test_data["query"] + assert data["count"] == 3 + assert len(data["chunks"]) == 3 + # MARK: - Error Handling Tests