11日目-embeddingの実行-

Azure OpenAIとChromaDBを使用したベクトルデータベースアプリケーションの実装

はじめに

こんにちは！今回は、Azure OpenAIとChromaDBを使用したベクトルデータベースアプリケーションの実装において、特にAzure OpenAIの初期化エラー処理とドキュメントアップロード機能の実装について共有したいと思います。

環境設定

# requirements.txt
fastapi==0.68.0
python-dotenv==0.19.0
azure-openai==1.0.0
chromadb==0.4.0
streamlit==1.8.0

Azure OpenAIクライアントの初期化とエラーハンドリング

1. 環境変数の設定

# .env
AZURE_OPENAI_API_KEY=your_api_key
AZURE_OPENAI_ENDPOINT=your_endpoint
AZURE_OPENAI_API_VERSION=2023-05-15

2. クライアント初期化処理

from azure.openai import AzureOpenAI
from dotenv import load_dotenv
import os

def initialize_azure_openai():
    try:
        load_dotenv()
        required_env_vars = [
            "AZURE_OPENAI_API_KEY",
            "AZURE_OPENAI_ENDPOINT",
            "AZURE_OPENAI_API_VERSION"
        ]
        
        # 環境変数の存在チェック
        missing_vars = [var for var in required_env_vars if not os.getenv(var)]
        if missing_vars:
            raise ValueError(
                f"Azure OpenAIの設定が不完全です。以下の環境変数を確認してください: {', '.join(missing_vars)}"
            )
            
        client = AzureOpenAI(
            api_key=os.getenv("AZURE_OPENAI_API_KEY"),
            api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
            azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT")
        )
        print("Azure OpenAI client initialized successfully")
        return client
        
    except Exception as e:
        print(f"Azure OpenAI initialization error: {str(e)}")
        return None

3. エラーハンドリングの実装

def generate_embeddings(text: str, client: AzureOpenAI = None):
    if not client:
        raise ValueError("Azure OpenAIクライアントが初期化されていません。環境変数を確認してください。")
        
    try:
        response = client.embeddings.create(
            input=text,
            model="text-embedding-ada-002"
        )
        return response.data[0].embedding
        
    except Exception as e:
        print(f"Embedding generation error: {str(e)}")
        return None

ドキュメントアップロード機能の実装

1. FastAPIエンドポイント

from fastapi import FastAPI, HTTPException, UploadFile, File
from typing import List

app = FastAPI()

@app.post("/collections/{collection_name}/add")
async def add_documents(
    collection_name: str,
    documents: List[str],
    metadata: List[dict] = None
):
    try:
        # コレクションの取得または作成
        collection = get_or_create_collection(collection_name)
        if not collection:
            raise HTTPException(status_code=404, detail="コレクションの作成に失敗しました")
            
        # ドキュメントの埋め込みベクトル生成
        embeddings = []
        for doc in documents:
            embedding = generate_embeddings(doc, azure_client)
            if not embedding:
                raise HTTPException(
                    status_code=500,
                    detail="埋め込みベクトルの生成に失敗しました"
                )
            embeddings.append(embedding)
            
        # ChromaDBへの保存
        collection.add(
            documents=documents,
            embeddings=embeddings,
            metadatas=metadata if metadata else [{}] * len(documents)
        )
        
        return {"status": "success", "message": "ドキュメントが正常に追加されました"}
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

2. Streamlitフロントエンド

import streamlit as st
import requests
import json

def upload_documents():
    st.header("ドキュメントのアップロード")
    
    # コレクション名の入力
    collection_name = st.text_input("コレクション名")
    
    # ドキュメントのアップロード
    uploaded_files = st.file_uploader(
        "ドキュメントをアップロード",
        accept_multiple_files=True,
        type=["txt"]
    )
    
    if st.button("アップロード") and uploaded_files and collection_name:
        documents = []
        metadata = []
        
        for file in uploaded_files:
            content = file.read().decode()
            documents.append(content)
            metadata.append({"filename": file.name})
            
        try:
            response = requests.post(
                f"{API_URL}/collections/{collection_name}/add",
                json={
                    "documents": documents,
                    "metadata": metadata
                }
            )
            
            if response.status_code == 200:
                st.success("ドキュメントが正常にアップロードされました")
            else:
                st.error(f"エラーが発生しました: {response.json()['detail']}")
                
        except Exception as e:
            st.error(f"アップロード中にエラーが発生しました: {str(e)}")

単体テストの実装

import pytest
from unittest.mock import Mock, patch

def test_azure_openai_initialization():
    with patch.dict('os.environ', {
        'AZURE_OPENAI_API_KEY': 'test_key',
        'AZURE_OPENAI_ENDPOINT': 'test_endpoint',
        'AZURE_OPENAI_API_VERSION': 'test_version'
    }):
        client = initialize_azure_openai()
        assert client is not None

def test_embedding_generation():
    mock_client = Mock()
    mock_client.embeddings.create.return_value = Mock(
        data=[Mock(embedding=[0.1, 0.2, 0.3])]
    )
    
    embedding = generate_embeddings("test text", mock_client)
    assert embedding is not None
    assert len(embedding) > 0

@pytest.mark.asyncio
async def test_document_upload():
    with patch('app.generate_embeddings') as mock_generate:
        mock_generate.return_value = [0.1, 0.2, 0.3]
        
        response = await add_documents(
            "test_collection",
            ["test document"],
            [{"test": "metadata"}]
        )
        
        assert response["status"] == "success"

主な実装のポイント

1. Azure OpenAI初期化のエラーハンドリング

   • 環境変数の存在確認
   • クライアント初期化時の例外処理
   • エラーメッセージの日本語化

2. ドキュメントアップロード機能

   • ファイルのバリデーション
   • 非同期処理の実装
   • エラー時のユーザーフィードバック

3. テストの実装

   • モックを使用した外部依存の分離
   • 異常系のテスト
   • 非同期処理のテスト

まとめ

今回は、Azure OpenAIとChromaDBを使用したベクトルデータベースアプリケーションの実装について、特にエラーハンドリングとテストに焦点を当てて解説しました。

実装時の注意点：

1. 環境変数の適切な管理
2. エラーメッセージの明確化
3. 適切な例外処理の実装
4. テストケースの充実

参考リンク

• Azure OpenAI Documentation
• ChromaDB Documentation
• FastAPI Testing

#Python #AzureOpenAI #ChromaDB #FastAPI #Streamlit #Testing