はじめに
Open WebUIとVertex AI Searchの連携は、生成AIの活用を大きく広げる可能性を秘めています。しかし、現時点では公式なパイプラインが公開されていないため、自力でパイプラインを作成する必要があります。本記事では、Open WebUIとVertex AI Searchを連携させるためのカスタムパイプラインの作成方法を解説します。
Open WebUIとは?
Open WebUI は、大規模言語モデル(LLM) をローカル環境で手軽に利用するための、Web ベースのインターフェース です。
※詳しくはこちらを参照下さい。
Open WebUI
Gitソース
導入手順
※Open WebUIのソースをGitからclone済み
① Open WebUIのコンテナに必要なモジュールと環境変数をセットする
docker-compose.yaml
services: ollama: volumes: - ollama:/root/.ollama container_name: ollama pull_policy: always tty: true restart: unless-stopped image: ollama/ollama:${OLLAMA_DOCKER_TAG-latest} open-webui: build: context: . args: OLLAMA_BASE_URL: '/ollama' dockerfile: Dockerfile image: ghcr.io/open-webui/open-webui:${WEBUI_DOCKER_TAG-main} container_name: open-webui-test volumes: - open-webui:/app/backend/data - ${GOOGLE_APPLICATION_CREDENTIALS}:/tmp/keys/key.json:ro depends_on: - ollama ports: - ${OPEN_WEBUI_PORT-3000}:8080 environment: - 'OLLAMA_BASE_URL=http://ollama:11434' - 'WEBUI_SECRET_KEY=' - 'GCLOUD_PROJECT=${GCLOUD_PROJECT}' - 'GOOGLE_APPLICATION_CREDENTIALS=/tmp/keys/key.json' extra_hosts: - host.docker.internal:host-gateway restart: unless-stopped volumes: ollama: {} open-webui: {}
.env
GCLOUD_PROJECT={使用しているプロジェクト} GOOGLE_API_KEY={APIキー} GOOGLE_APPLICATION_CREDENTIALS=/{認証情報が置いてあるpath}/application_default_credentials.json
GOOGLE_API_KEY取得手順
Gemini Developer API
cretcredentials取得手順
※参照 アプリケーションのデフォルト認証情報を設定する
backend/requirements.txt
langchain-google-community==1.0.7 langchain-google-vertexai==1.0.10 langchain-google-community[vertexaisearch] ※追加
② Pipeをセットする
※今回はVertex AI SearchにLang Chainを使って接続
Vertex AI Search接続用のPipe
conectVertexAISearch.py
import os from pydantic import BaseModel, Field from typing import List, Union, Iterator from langchain_google_community import VertexAISearchRetriever from langchain_core.prompts import ChatPromptTemplate from textwrap import dedent from langchain_google_vertexai import ChatVertexAI from langchain_core.runnables import RunnablePassthrough from langchain_core.output_parsers import StrOutputParser import google.generativeai as genai # Set DEBUG to True to enable detailed logging DEBUG = True PROJECT_ID = os.getenv("GCLOUD_PROJECT") # Set to your Project ID class Pipe: class Valves(BaseModel): GOOGLE_API_KEY: str = Field(default="") STORE_ID: str = Field(default="") def __init__(self): self.id = "self-mode" self.type = "api connect" self.name = "Vertex AI Search connect:" self.valves = self.Valves( **{ "GOOGLE_API_KEY": os.getenv("GOOGLE_API_KEY", ""), } ) def get_vertex_aisearch_google_models(self): if not self.valves.GOOGLE_API_KEY: return [ { "id": "error", "name": "GOOGLE_API_KEY is not set. Please update the API Key in the valves.", } ] try: genai.configure(api_key=self.valves.GOOGLE_API_KEY) models = genai.list_models() valid_model_names = ["gemini-1.5-pro", "gemini-1.0-pro-001", "gemini-1.0-pro", "gemini-1.0-pro-002", "gemini-1.5-flash-001"] return [ { "id": model.name[7:], # remove the "models/" part "name": model.display_name, } for model in models if "generateContent" in model.supported_generation_methods if model.name[7:] in valid_model_names ] except Exception as e: if DEBUG: print(f"Error fetching Google models: {e}") return [ {"id": "error", "name": f"Could not fetch models from Google: {str(e)}"} ] def pipes(self) -> List[dict]: return self.get_vertex_aisearch_google_models() def pipe(self, body: dict) -> Union[str, Iterator[str]]: if not self.valves.STORE_ID: return f"Error: STORE_ID is not set. Please update the STORE_ID in the valves." try: retriever = self.get_retriever(self.valves.STORE_ID) system_prompt: str = """ {任意のシステムプロンプト} """ premise = dedent(system_prompt).strip() prompt = ChatPromptTemplate.from_template( dedent( """ Premise: {premise} Context: {context} Question: {question} """ ).strip(), partial_variables={"premise": premise}, ) model_id = self.setModel(body) if not model_id.startswith("gemini-"): return f"Error: Invalid model name format: {model_id}" llm = ChatVertexAI(model_name=model_id, temperature=0) def format_docs(docs): return "\n\n".join(doc.page_content for doc in docs) chain = ( {"context": retriever | format_docs, "question": RunnablePassthrough()} | prompt | llm | StrOutputParser() ) return chain.stream(self.setQuery(body)) except Exception as e: if DEBUG: print(f"Error in pipe method: {e}") return f"Error: {e}" @staticmethod def setQuery(body): messages = body["messages"] history = "\n".join([msg["content"] for msg in messages[:-1]]) current_query = messages[-1]["content"] return f"History:\n{history}\n\nCurrent Query:\n{current_query}" @staticmethod def setModel(body): model_id = body.get("model", "gemini-1.5-pro") if "." in model_id: model_id = model_id.split(".", 1)[-1] return model_id @staticmethod def get_retriever(store_id: str) -> VertexAISearchRetriever: return VertexAISearchRetriever( project_id=PROJECT_ID, data_store_id=store_id, location_id='global', max_documents=5, )
※ValvesのGOOGLE_API_KEY/STORE_ID入力必須
おまけ
Open WebUIからVertex AI Searchを搭載したバックエンドにAPIで接続するPipe
connect_api.py
import requests import os from pydantic import BaseModel, Field from typing import List, Union, Iterator import google.generativeai as genai # Set DEBUG to True to enable detailed logging DEBUG = True class Pipe: class Valves(BaseModel): CONNECT_API_URL: str = Field(default="") GOOGLE_API_KEY: str = Field(default="") STORE_ID: str = Field(default="") PROJECT_ID: str = Field(default="") def __init__(self): self.id = "self-mode" self.type = "api connect" self.name = "API connect: " self.valves = self.Valves( **{ "GOOGLE_API_KEY": os.getenv("GOOGLE_API_KEY", ""), } ) def get_vertex_aisearch_google_models(self): if not self.valves.GOOGLE_API_KEY: return [ { "id": "error", "name": "GOOGLE_API_KEY is not set. Please update the API Key in the valves.", } ] try: genai.configure(api_key=self.valves.GOOGLE_API_KEY) models = genai.list_models() valid_model_names = ["gemini-1.5-pro", "gemini-1.0-pro-001", "gemini-1.0-pro", "gemini-1.0-pro-002", "gemini-1.5-flash-001"] return [ { "id": model.name[7:], # remove the "models/" part "name": model.display_name, } for model in models if "generateContent" in model.supported_generation_methods if model.name[7:] in valid_model_names ] except Exception as e: if DEBUG: print(f"Error fetching Google models: {e}") return [ {"id": "error", "name": f"Could not fetch models from Google: {str(e)}"} ] def pipes(self) -> List[dict]: return self.get_vertex_aisearch_google_models() def pipe(self, body: dict) -> Union[str, Iterator[str]]: for key in ['CONNECT_API_URL', 'STORE_ID', 'PROJECT_ID', 'GOOGLE_API_KEY']: if not getattr(self.valves, key): return f"Error: {key} is not set. Please update the {key} in the valves." try: r = requests.post( self.valves.CONNECT_API_URL, json={"body": body, "valves": { 'PROJECT_ID': self.valves.PROJECT_ID, 'STORE_ID': self.valves.STORE_ID, }}, stream=True ) if r.status_code == 200: return (line.decode('utf-8') + '\n' for line in r.iter_lines(chunk_size=10)) else: return "Failed. status code: " + str(r.status_code) except Exception as e: if DEBUG: print(f"Error in pipe method: {e}") return f"Error: {e}"</pre>
※バックエンドの仕様によって、APIへ渡すパラメーターは適宜変更して下さい。
※バックエンドで処理する形にするとOpen WebUI側にモジュールのインストールが不要となるので、
Open WebUI側に手を入れたくない方はこちらがおすすめです。
まとめ
Open WebUIからVertex AI Searchを組み合わせることで、よりマネージドで高度な機能を備えたRAGシステムの構築が可能かと思います。今回ご紹介したPipeを軸として、要件に合わせて処理を追加・変更していくことが可能なのでより柔軟で幅の広いシステム構築が実現していけると思います。