diff --git a/Plans.md b/Plans.md index 25e3116..5ea113b 100644 --- a/Plans.md +++ b/Plans.md @@ -58,8 +58,9 @@ | 79 | **都市マップ WebSocket リアルタイム更新** | `skills/city_map_realtime.py` | 4541 | v0.82 | | 80 | **TraceCompiler — LCSスキルマイニング → 決定論的WF** | `skills/trace_compiler.py` | 4613 | v0.83 | | 81 | **レコメンドエンジン統合** | `skills/recommender.py` | 4705 | v0.84 | +| 82 | **EmbeddingGemma RAG エンジン** | `skills/gemma_rag.py` | 4796 | v0.85 | -> **累計テスト数**: 4705 PASS (Sprint 81: +92) — Sprint 82 候補検討中 +> **累計テスト数**: 4796 PASS (Sprint 82: +91) — Sprint 83 候補検討中 --- @@ -459,6 +460,39 @@ --- +## Sprint 82 詳細 (完了) + +### Sprint 82: EmbeddingGemma RAG エンジン — v0.85.0 +> Google EmbeddingGemma (text-embedding-004) を利用した Retrieval-Augmented Generation パイプライン。 +> awesome-gemma の知見をベースに Mock / 本番モード両対応で実装。 + +| task-id | 説明 | 状態 | +|---------|------|------| +| 82.1 | `GemmaEmbeddingModel` / `ChunkingStrategy` / `RAGStatus` (Enum 層) | cc:完了 | +| 82.2 | `EmbeddingConfig` / `RAGDocument` / `RAGChunk` / `ChunkingConfig` (データモデル) | cc:完了 | +| 82.3 | `RetrievalResult` / `RAGAnswer` / `IndexStats` (結果モデル) | cc:完了 | +| 82.4 | `GemmaEmbeddingProvider` — Mock (deterministic hash) + HTTP (Google AI Studio REST) | cc:完了 | +| 82.5 | `DocumentChunker` — FIXED_SIZE / SENTENCE / PARAGRAPH 3戦略 + min_chars フィルタ | cc:完了 | +| 82.6 | `_ChunkVectorStore` — in-memory コサイン類似度検索 + doc_filter 対応 | cc:完了 | +| 82.7 | `RAGIndexer` — chunk → embed → store (上書き/削除/一括対応) | cc:完了 | +| 82.8 | `RAGRetriever` — query embed → similarity search → RetrievalResult | cc:完了 | +| 82.9 | `MockLLMGenerator` / `GemmaLLMGenerator` — 生成層 (FunctionGemma 統合ポイント) | cc:完了 | +| 82.10 | `RAGPipeline` — create_mock / create_production ファクトリ + query / search / index 統合 | cc:完了 | +| 82.11 | `serve/api.py` — `/v1/rag/*` 6 エンドポイント (index/batch/delete/search/query/status) | cc:完了 | +| 82.T | `tests/test_sprint82.py` — 91 PASS (累計 4796) | cc:完了 | + +### コンポーネント対応表 +| EmbeddingGemma 概念 | gemma_rag.py | 役割 | +|--------------------|--------------|------| +| EmbeddingGemma API | `GemmaEmbeddingProvider._http_embed` | text-embedding-004 REST 呼び出し | +| Mock embedding | `GemmaEmbeddingProvider._mock_embed` | MD5 ハッシュ → deterministic ベクター | +| Chunk | `RAGChunk` | チャンキング後断片 | +| Vector Store | `_ChunkVectorStore` | in-memory cosine 検索 (FAISS 代替) | +| RAG Query | `RAGPipeline.query` | 検索 + 生成統合 | +| Gemma 生成 | `GemmaLLMGenerator` | FunctionGemma / Gemma 4 E2B 統合ポイント | + +--- + ## Sprint 78 詳細 (完了) ### Sprint 78: 3D 都市マップビジュアライゼーション — v0.81.0 diff --git a/open_mythos/skills/gemma_rag.py b/open_mythos/skills/gemma_rag.py new file mode 100644 index 0000000..dc55801 --- /dev/null +++ b/open_mythos/skills/gemma_rag.py @@ -0,0 +1,728 @@ +""" +Sprint 82 — EmbeddingGemma RAG エンジン + +Google EmbeddingGemma (text-embedding-004) を利用した +Retrieval-Augmented Generation パイプラインの OpenMythos 移植。 + +ref: awesome-gemma / Google AI Studio Embedding API + https://ai.google.dev/gemini-api/docs/embeddings + +オブジェクト: + GemmaEmbeddingModel : 利用モデル識別 enum + EmbeddingConfig : 埋め込みプロバイダー設定 + RAGDocument : インデックス対象ドキュメント + RAGChunk : チャンキング後の断片 + RetrievalResult : 検索結果 (chunk + score) + RAGAnswer : 最終 RAG 回答 + GemmaEmbeddingProvider: EmbeddingGemma 呼び出し (Mock / HTTP) + DocumentChunker : 固定長 + 文境界チャンキング + RAGIndexer : chunk → embed → VectorStore(FAISS in-memory) + RAGRetriever : query embed → similarity search → RetrievalResult + RAGPipeline : インデックス + 検索 + 生成の統合ファサード + +使用例:: + pipeline = RAGPipeline.create_mock() + pipeline.index_document(RAGDocument(id="d1", title="製品説明", content="...")) + answer = pipeline.query("この製品の主な特徴は？") + print(answer.text) + print(answer.sources) # [{"chunk_id": ..., "score": ...}] +""" +from __future__ import annotations + +import hashlib +import json +import math +import re +import time +import urllib.error +import urllib.request +import uuid +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# 型定義 / Enum +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class GemmaEmbeddingModel(str, Enum): + """Google AI Studio embedding モデル""" + TEXT_EMBEDDING_004 = "text-embedding-004" # EmbeddingGemma 最新 + EMBEDDING_001 = "embedding-001" # 旧世代 + MOCK = "mock" # テスト用決定論的埋め込み + + +class ChunkingStrategy(str, Enum): + FIXED_SIZE = "fixed_size" # 固定文字数 + SENTENCE = "sentence" # 文境界 + PARAGRAPH = "paragraph" # 段落境界 + + +class RAGStatus(str, Enum): + IDLE = "idle" + INDEXING = "indexing" + READY = "ready" + ERROR = "error" + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# データモデル +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +@dataclass +class EmbeddingConfig: + """埋め込みプロバイダー設定""" + model: GemmaEmbeddingModel = GemmaEmbeddingModel.TEXT_EMBEDDING_004 + api_key: str = "" + timeout: int = 30 # 秒 + dim: int = 768 # text-embedding-004 の次元数 + task_type: str = "RETRIEVAL_DOCUMENT" + # RETRIEVAL_DOCUMENT | RETRIEVAL_QUERY | SEMANTIC_SIMILARITY | CLASSIFICATION + + +@dataclass +class RAGDocument: + """インデックス対象ドキュメント""" + id: str + title: str + content: str + metadata: Dict[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not self.id: + self.id = str(uuid.uuid4()) + + +@dataclass +class RAGChunk: + """チャンキング後の断片""" + id: str + doc_id: str + doc_title: str + text: str + chunk_index: int + metadata: Dict[str, Any] = field(default_factory=dict) + + @property + def full_text(self) -> str: + """タイトルを先頭に付与した検索用テキスト""" + return f"{self.doc_title}\n{self.text}" + + +@dataclass +class ChunkingConfig: + """チャンキング設定""" + strategy: ChunkingStrategy = ChunkingStrategy.SENTENCE + chunk_size: int = 512 # 文字数 (FIXED_SIZE / SENTENCE 上限) + overlap: int = 64 # オーバーラップ文字数 + min_chars: int = 20 # この未満は無視 + + +@dataclass +class RetrievalResult: + """検索結果""" + chunk: RAGChunk + score: float # cosine similarity (0.0〜1.0) + rank: int = 0 + + def to_dict(self) -> Dict[str, Any]: + return { + "chunk_id": self.chunk.id, + "doc_id": self.chunk.doc_id, + "doc_title": self.chunk.doc_title, + "text": self.chunk.text, + "score": round(self.score, 4), + "rank": self.rank, + } + + +@dataclass +class RAGAnswer: + """RAG 最終回答""" + query: str + text: str + sources: List[Dict[str, Any]] = field(default_factory=list) + latency_ms: float = 0.0 + model_used: str = "" + + @property + def success(self) -> bool: + return bool(self.text) + + def to_dict(self) -> Dict[str, Any]: + return { + "query": self.query, + "answer": self.text, + "sources": self.sources, + "latency_ms": round(self.latency_ms, 1), + "model_used": self.model_used, + } + + +@dataclass +class IndexStats: + """インデックス統計""" + doc_count: int = 0 + chunk_count: int = 0 + status: RAGStatus = RAGStatus.IDLE + last_indexed: Optional[str] = None # ISO 8601 + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# GemmaEmbeddingProvider +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class GemmaEmbeddingProvider: + """ + EmbeddingGemma (text-embedding-004) の埋め込み生成プロバイダー。 + + mock モード時は deterministic ハッシュベース埋め込みを返す（テスト用）。 + 本番モード時は Google AI Studio REST API を呼ぶ。 + """ + + _ENDPOINT = ( + "https://generativelanguage.googleapis.com/v1beta/models/" + "{model}:embedContent?key={api_key}" + ) + + def __init__(self, config: EmbeddingConfig) -> None: + self.config = config + self._is_mock = (config.model == GemmaEmbeddingModel.MOCK or not config.api_key) + + # ── 公開 API ────────────────────────────────────────────────── + + def embed(self, text: str, task_type: Optional[str] = None) -> List[float]: + """テキスト 1 件を埋め込みベクターに変換する。""" + t = task_type or self.config.task_type + if self._is_mock: + return self._mock_embed(text, self.config.dim) + return self._http_embed(text, t) + + def embed_batch( + self, + texts: List[str], + task_type: Optional[str] = None, + ) -> List[List[float]]: + """複数テキストを一括で埋め込む。""" + return [self.embed(t, task_type) for t in texts] + + @property + def dim(self) -> int: + return self.config.dim + + # ── Mock 実装 ────────────────────────────────────────────────── + + @staticmethod + def _mock_embed(text: str, dim: int) -> List[float]: + """ + テスト用 deterministic 埋め込み。 + 同じテキストは常に同じベクターを返し、コサイン類似度が + 意味的近さをある程度反映するよう単語ハッシュで構成する。 + """ + vec = [0.0] * dim + words = re.findall(r'\w+', text.lower()) + for word in words: + h = int(hashlib.md5(word.encode()).hexdigest(), 16) + idx = h % dim + vec[idx] += 1.0 + # L2 正規化 + norm = math.sqrt(sum(v * v for v in vec)) or 1.0 + return [v / norm for v in vec] + + # ── HTTP 実装 ───────────────────────────────────────────────── + + def _http_embed(self, text: str, task_type: str) -> List[float]: + model_name = self.config.model.value + url = self._ENDPOINT.format( + model=model_name, + api_key=self.config.api_key, + ) + body = json.dumps({ + "model": f"models/{model_name}", + "content": {"parts": [{"text": text}]}, + "taskType": task_type, + }).encode() + req = urllib.request.Request( + url, + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=self.config.timeout) as resp: + data = json.loads(resp.read()) + return data["embedding"]["values"] + except urllib.error.HTTPError as e: + raise RuntimeError(f"EmbeddingGemma HTTP {e.code}: {e.read().decode()}") from e + except (KeyError, json.JSONDecodeError) as e: + raise RuntimeError(f"EmbeddingGemma レスポンス解析失敗: {e}") from e + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# DocumentChunker +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class DocumentChunker: + """ + ドキュメントを検索しやすい断片に分割する。 + + strategy: + FIXED_SIZE — chunk_size 文字ごとに分割 (overlap あり) + SENTENCE — 文末 (。.!?) で分割、上限 chunk_size 以内に収める + PARAGRAPH — 空行で段落分割 + """ + + def __init__(self, config: Optional[ChunkingConfig] = None) -> None: + self.config = config or ChunkingConfig() + + def chunk(self, doc: RAGDocument) -> List[RAGChunk]: + strategy = self.config.strategy + if strategy == ChunkingStrategy.FIXED_SIZE: + texts = self._fixed_size(doc.content) + elif strategy == ChunkingStrategy.SENTENCE: + texts = self._sentence(doc.content) + else: + texts = self._paragraph(doc.content) + + chunks: List[RAGChunk] = [] + for i, text in enumerate(texts): + if len(text) < self.config.min_chars: + continue + chunk_id = f"{doc.id}__c{i}" + chunks.append(RAGChunk( + id=chunk_id, + doc_id=doc.id, + doc_title=doc.title, + text=text.strip(), + chunk_index=i, + metadata={**doc.metadata, "doc_id": doc.id}, + )) + return chunks + + # ── チャンキング戦略 ───────────────────────────────────────── + + def _fixed_size(self, text: str) -> List[str]: + size = self.config.chunk_size + overlap = self.config.overlap + step = max(1, size - overlap) + return [text[i:i + size] for i in range(0, len(text), step) if text[i:i + size]] + + def _sentence(self, text: str) -> List[str]: + """文末記号で分割し、chunk_size を超えたら強制分割。""" + # 文末記号: 。.!?！？（全角・半角両対応） + raw_sents = re.split(r'(?<=[。.!?！？])\s*', text) + chunks: List[str] = [] + current = "" + for sent in raw_sents: + if not sent: + continue + if len(current) + len(sent) <= self.config.chunk_size: + current += sent + else: + if current: + chunks.append(current) + # sent 自体が chunk_size を超えるなら固定分割 + if len(sent) > self.config.chunk_size: + chunks.extend(self._fixed_size(sent)) + current = "" + else: + current = sent + if current: + chunks.append(current) + return chunks + + def _paragraph(self, text: str) -> List[str]: + """空行で段落分割。""" + paras = re.split(r'\n\s*\n', text) + return [p.strip() for p in paras if p.strip()] + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# In-Memory Vector Store (FAISS 代替) +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +def _cosine_sim(a: List[float], b: List[float]) -> float: + dot = sum(x * y for x, y in zip(a, b)) + na = math.sqrt(sum(x * x for x in a)) or 1e-9 + nb = math.sqrt(sum(x * x for x in b)) or 1e-9 + return dot / (na * nb) + + +class _ChunkVectorStore: + """RAG チャンク専用の in-memory ベクターストア。""" + + def __init__(self) -> None: + self._store: Dict[str, Tuple[RAGChunk, List[float]]] = {} + + def upsert(self, chunk: RAGChunk, vector: List[float]) -> None: + self._store[chunk.id] = (chunk, vector) + + def delete_by_doc(self, doc_id: str) -> int: + keys = [k for k, (c, _) in self._store.items() if c.doc_id == doc_id] + for k in keys: + del self._store[k] + return len(keys) + + def search( + self, + query_vec: List[float], + top_k: int = 5, + doc_filter: Optional[List[str]] = None, + ) -> List[Tuple[RAGChunk, float]]: + results: List[Tuple[RAGChunk, float]] = [] + for chunk, vec in self._store.values(): + if doc_filter and chunk.doc_id not in doc_filter: + continue + score = _cosine_sim(query_vec, vec) + results.append((chunk, score)) + results.sort(key=lambda x: x[1], reverse=True) + return results[:top_k] + + @property + def chunk_count(self) -> int: + return len(self._store) + + @property + def doc_ids(self) -> List[str]: + return list({c.doc_id for c, _ in self._store.values()}) + + def clear(self) -> None: + self._store.clear() + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# RAGIndexer +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class RAGIndexer: + """ + ドキュメントをチャンク化 → 埋め込み → ベクターストアへ登録する。 + """ + + def __init__( + self, + embedder: GemmaEmbeddingProvider, + chunker: Optional[DocumentChunker] = None, + store: Optional[_ChunkVectorStore] = None, + ) -> None: + self.embedder = embedder + self.chunker = chunker or DocumentChunker() + self.store = store or _ChunkVectorStore() + self._doc_registry: Dict[str, RAGDocument] = {} + + def index(self, doc: RAGDocument) -> List[RAGChunk]: + """1 ドキュメントをインデックスする。既存の同 ID は上書き。""" + # 旧チャンクを削除 + if doc.id in self._doc_registry: + self.store.delete_by_doc(doc.id) + + chunks = self.chunker.chunk(doc) + texts = [c.full_text for c in chunks] + vectors = self.embedder.embed_batch(texts, task_type="RETRIEVAL_DOCUMENT") + + for chunk, vec in zip(chunks, vectors): + self.store.upsert(chunk, vec) + + self._doc_registry[doc.id] = doc + return chunks + + def index_batch(self, docs: List[RAGDocument]) -> Dict[str, int]: + """複数ドキュメントを一括インデックス。{doc_id: chunk_count} を返す。""" + result: Dict[str, int] = {} + for doc in docs: + chunks = self.index(doc) + result[doc.id] = len(chunks) + return result + + def delete(self, doc_id: str) -> int: + """ドキュメントをインデックスから削除。削除チャンク数を返す。""" + removed = self.store.delete_by_doc(doc_id) + self._doc_registry.pop(doc_id, None) + return removed + + @property + def stats(self) -> IndexStats: + return IndexStats( + doc_count=len(self._doc_registry), + chunk_count=self.store.chunk_count, + status=RAGStatus.READY if self._doc_registry else RAGStatus.IDLE, + ) + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# RAGRetriever +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class RAGRetriever: + """ + クエリを埋め込み → ベクターストア検索 → RetrievalResult リストを返す。 + """ + + def __init__( + self, + embedder: GemmaEmbeddingProvider, + store: _ChunkVectorStore, + ) -> None: + self.embedder = embedder + self.store = store + + def retrieve( + self, + query: str, + top_k: int = 5, + doc_filter: Optional[List[str]] = None, + score_threshold: float = 0.0, + ) -> List[RetrievalResult]: + """ + クエリに最も関連するチャンクを返す。 + + Args: + query: 検索クエリ文字列 + top_k: 返す上位件数 + doc_filter: 特定 doc_id のみ検索（None = 全件） + score_threshold: この未満のスコアは除外 + """ + q_vec = self.embedder.embed(query, task_type="RETRIEVAL_QUERY") + hits = self.store.search(q_vec, top_k=top_k, doc_filter=doc_filter) + results: List[RetrievalResult] = [] + for rank, (chunk, score) in enumerate(hits, start=1): + if score < score_threshold: + continue + results.append(RetrievalResult(chunk=chunk, score=score, rank=rank)) + return results + + def retrieve_texts( + self, + query: str, + top_k: int = 5, + ) -> List[str]: + """チャンクのテキストだけをリストで返す（LLM プロンプト組み立て用）。""" + results = self.retrieve(query, top_k=top_k) + return [r.chunk.text for r in results] + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# GenerationMixin (LLM 生成の薄いラッパー) +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class MockLLMGenerator: + """テスト用 LLM — 取得チャンクを結合して返す。""" + + def generate(self, query: str, context_chunks: List[str]) -> str: + if not context_chunks: + return f"「{query}」に関する情報が見つかりませんでした。" + joined = "\n---\n".join(context_chunks[:3]) + return f"[RAG 回答]\n参照情報:\n{joined}\n\n質問: {query}" + + @property + def model_name(self) -> str: + return "mock-llm" + + +class GemmaLLMGenerator: + """ + Google AI Studio Gemma モデルを使った生成。 + (FunctionGemma / Gemma 4 等との統合ポイント) + """ + + _ENDPOINT = ( + "https://generativelanguage.googleapis.com/v1beta/models/" + "{model}:generateContent?key={api_key}" + ) + + def __init__( + self, + api_key: str, + model: str = "gemma-3-27b-it", + timeout: int = 60, + ) -> None: + self.api_key = api_key + self.model = model + self.timeout = timeout + + def generate(self, query: str, context_chunks: List[str]) -> str: + context = "\n\n".join(context_chunks) + prompt = ( + f"以下のコンテキスト情報を参照して、質問に日本語で答えてください。\n\n" + f"コンテキスト:\n{context}\n\n" + f"質問: {query}\n\n回答:" + ) + body = json.dumps({ + "contents": [{"parts": [{"text": prompt}]}], + "generationConfig": {"temperature": 0.3, "maxOutputTokens": 512}, + }).encode() + url = self._ENDPOINT.format(model=self.model, api_key=self.api_key) + req = urllib.request.Request( + url, + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=self.timeout) as resp: + data = json.loads(resp.read()) + return data["candidates"][0]["content"]["parts"][0]["text"] + except (urllib.error.HTTPError, KeyError, json.JSONDecodeError) as e: + raise RuntimeError(f"GemmaLLM 生成失敗: {e}") from e + + @property + def model_name(self) -> str: + return self.model + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# RAGPipeline (統合ファサード) +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class RAGPipeline: + """ + インデックス・検索・生成を統合したメインファサード。 + + 使用例: + # Mock モード（テスト） + pipeline = RAGPipeline.create_mock() + pipeline.index_document(RAGDocument(id="d1", title="製品Q&A", content="...")) + answer = pipeline.query("返品ポリシーは？") + print(answer.text) + + # 本番モード（API キーあり） + pipeline = RAGPipeline.create_production( + embedding_api_key="...", + generation_api_key="...", + ) + """ + + def __init__( + self, + indexer: RAGIndexer, + retriever: RAGRetriever, + generator: Any, # MockLLMGenerator | GemmaLLMGenerator + top_k: int = 5, + ) -> None: + self.indexer = indexer + self.retriever = retriever + self.generator = generator + self.top_k = top_k + + # ── ファクトリ ──────────────────────────────────────────────── + + @classmethod + def create_mock( + cls, + chunking_config: Optional[ChunkingConfig] = None, + top_k: int = 5, + ) -> "RAGPipeline": + """テスト用 Mock パイプライン（API キー不要）。""" + emb_cfg = EmbeddingConfig(model=GemmaEmbeddingModel.MOCK) + embedder = GemmaEmbeddingProvider(emb_cfg) + store = _ChunkVectorStore() + chunker = DocumentChunker(chunking_config) + indexer = RAGIndexer(embedder, chunker, store) + retriever = RAGRetriever(embedder, store) + generator = MockLLMGenerator() + return cls(indexer, retriever, generator, top_k) + + @classmethod + def create_production( + cls, + embedding_api_key: str, + generation_api_key: str = "", + embedding_model: GemmaEmbeddingModel = GemmaEmbeddingModel.TEXT_EMBEDDING_004, + generation_model: str = "gemma-3-27b-it", + chunking_config: Optional[ChunkingConfig] = None, + top_k: int = 5, + ) -> "RAGPipeline": + """本番用パイプライン（Google AI Studio API キー必要）。""" + emb_cfg = EmbeddingConfig(model=embedding_model, api_key=embedding_api_key) + embedder = GemmaEmbeddingProvider(emb_cfg) + store = _ChunkVectorStore() + chunker = DocumentChunker(chunking_config) + indexer = RAGIndexer(embedder, chunker, store) + retriever = RAGRetriever(embedder, store) + if generation_api_key: + generator = GemmaLLMGenerator(generation_api_key, model=generation_model) + else: + generator = MockLLMGenerator() + return cls(indexer, retriever, generator, top_k) + + # ── インデックス操作 ───────────────────────────────────────── + + def index_document(self, doc: RAGDocument) -> int: + """1 ドキュメントをインデックス。登録チャンク数を返す。""" + chunks = self.indexer.index(doc) + return len(chunks) + + def index_batch(self, docs: List[RAGDocument]) -> Dict[str, int]: + return self.indexer.index_batch(docs) + + def delete_document(self, doc_id: str) -> int: + return self.indexer.delete(doc_id) + + # ── 検索 ───────────────────────────────────────────────────── + + def search( + self, + query: str, + top_k: Optional[int] = None, + doc_filter: Optional[List[str]] = None, + score_threshold: float = 0.0, + ) -> List[RetrievalResult]: + """検索のみ（生成なし）。""" + return self.retriever.retrieve( + query, + top_k=top_k or self.top_k, + doc_filter=doc_filter, + score_threshold=score_threshold, + ) + + # ── RAG 生成 ───────────────────────────────────────────────── + + def query( + self, + question: str, + top_k: Optional[int] = None, + doc_filter: Optional[List[str]] = None, + score_threshold: float = 0.0, + ) -> RAGAnswer: + """ + 検索 + 生成の統合クエリ。 + + 1. question を embedding してチャンク検索 + 2. 上位 top_k チャンクをコンテキストとして LLM へ渡す + 3. RAGAnswer を返す + """ + t0 = time.perf_counter() + results = self.search( + question, + top_k=top_k or self.top_k, + doc_filter=doc_filter, + score_threshold=score_threshold, + ) + context_chunks = [r.chunk.text for r in results] + answer_text = self.generator.generate(question, context_chunks) + latency_ms = (time.perf_counter() - t0) * 1000 + + return RAGAnswer( + query=question, + text=answer_text, + sources=[r.to_dict() for r in results], + latency_ms=latency_ms, + model_used=self.generator.model_name, + ) + + # ── ステータス ──────────────────────────────────────────────── + + @property + def stats(self) -> IndexStats: + return self.indexer.stats + + def to_status_dict(self) -> Dict[str, Any]: + stats = self.stats + return { + "status": stats.status.value, + "doc_count": stats.doc_count, + "chunk_count": stats.chunk_count, + "embedder": self.indexer.embedder.config.model.value, + "generator": self.generator.model_name, + } diff --git a/serve/api.py b/serve/api.py index 97b9dd5..3d0399d 100644 --- a/serve/api.py +++ b/serve/api.py @@ -9771,3 +9771,110 @@ def rec_evaluate(body: _RecEvalBody): @app.get("/v1/rec/status", tags=["recommender"], summary="パイプライン状態 — Sprint 81H") def rec_status(): return _rec_pipeline.status() + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Sprint 82 — EmbeddingGemma RAG エンドポイント +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +from open_mythos.skills.gemma_rag import ( + RAGPipeline as _RAGPipeline, + RAGDocument as _RAGDocument, + ChunkingConfig as _ChunkingConfig, + ChunkingStrategy as _ChunkingStrategy, +) + +# シングルトン（サーバー起動時に Mock で初期化、API キーがあれば本番モードへ） +_GEMMA_API_KEY = os.getenv("GEMMA_API_KEY", "") +if _GEMMA_API_KEY: + _rag_pipeline: _RAGPipeline = _RAGPipeline.create_production( + embedding_api_key=_GEMMA_API_KEY, + ) +else: + _rag_pipeline: _RAGPipeline = _RAGPipeline.create_mock() + + +# ── Pydantic モデル ─────────────────────────────────────────────── + +class _RAGIndexBody(BaseModel): + id: str = Field(..., description="ドキュメント ID") + title: str = Field(..., description="ドキュメントタイトル") + content: str = Field(..., description="本文テキスト") + metadata: dict = Field(default_factory=dict) + + +class _RAGBatchIndexBody(BaseModel): + documents: List[_RAGIndexBody] + + +class _RAGQueryBody(BaseModel): + question: str = Field(..., description="質問文") + top_k: int = Field(5, ge=1, le=20) + doc_filter: Optional[List[str]] = Field(None, description="対象 doc_id リスト") + score_threshold: float = Field(0.0, ge=0.0, le=1.0) + + +class _RAGSearchBody(BaseModel): + query: str + top_k: int = Field(5, ge=1, le=20) + doc_filter: Optional[List[str]] = None + score_threshold: float = Field(0.0, ge=0.0, le=1.0) + + +# ── エンドポイント ──────────────────────────────────────────────── + +@app.post("/v1/rag/index", tags=["rag"], summary="ドキュメントインデックス登録 — Sprint 82") +def rag_index(body: _RAGIndexBody): + doc = _RAGDocument( + id=body.id, title=body.title, + content=body.content, metadata=body.metadata, + ) + count = _rag_pipeline.index_document(doc) + return {"doc_id": body.id, "chunks_indexed": count} + + +@app.post("/v1/rag/index/batch", tags=["rag"], summary="ドキュメント一括インデックス — Sprint 82") +def rag_index_batch(body: _RAGBatchIndexBody): + docs = [ + _RAGDocument(id=d.id, title=d.title, content=d.content, metadata=d.metadata) + for d in body.documents + ] + result = _rag_pipeline.index_batch(docs) + return {"indexed": result, "total_docs": len(result)} + + +@app.delete("/v1/rag/index/{doc_id}", tags=["rag"], summary="ドキュメント削除 — Sprint 82") +def rag_delete(doc_id: str): + removed = _rag_pipeline.delete_document(doc_id) + return {"doc_id": doc_id, "chunks_removed": removed} + + +@app.post("/v1/rag/search", tags=["rag"], summary="セマンティック検索（生成なし）— Sprint 82") +def rag_search(body: _RAGSearchBody): + results = _rag_pipeline.search( + body.query, + top_k=body.top_k, + doc_filter=body.doc_filter, + score_threshold=body.score_threshold, + ) + return { + "query": body.query, + "results": [r.to_dict() for r in results], + "count": len(results), + } + + +@app.post("/v1/rag/query", tags=["rag"], summary="RAG クエリ（検索 + 生成）— Sprint 82") +def rag_query(body: _RAGQueryBody): + answer = _rag_pipeline.query( + body.question, + top_k=body.top_k, + doc_filter=body.doc_filter, + score_threshold=body.score_threshold, + ) + return answer.to_dict() + + +@app.get("/v1/rag/status", tags=["rag"], summary="RAG パイプライン状態 — Sprint 82") +def rag_status(): + return _rag_pipeline.to_status_dict() diff --git a/tests/test_sprint82.py b/tests/test_sprint82.py new file mode 100644 index 0000000..47d5267 --- /dev/null +++ b/tests/test_sprint82.py @@ -0,0 +1,746 @@ +""" +tests/test_sprint82.py — Sprint 82: EmbeddingGemma RAG エンジン テストスイート + +対象: open_mythos/skills/gemma_rag.py +目標: 80+ PASS +""" +from __future__ import annotations + +import math +import re +import pytest +from open_mythos.skills.gemma_rag import ( + # Enums + GemmaEmbeddingModel, ChunkingStrategy, RAGStatus, + # Models + EmbeddingConfig, RAGDocument, RAGChunk, ChunkingConfig, + RetrievalResult, RAGAnswer, IndexStats, + # Components + GemmaEmbeddingProvider, DocumentChunker, + _ChunkVectorStore, _cosine_sim, + RAGIndexer, RAGRetriever, + MockLLMGenerator, GemmaLLMGenerator, + RAGPipeline, +) + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# Fixtures +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +@pytest.fixture +def mock_embedder(): + cfg = EmbeddingConfig(model=GemmaEmbeddingModel.MOCK) + return GemmaEmbeddingProvider(cfg) + + +@pytest.fixture +def basic_doc(): + return RAGDocument( + id="doc1", + title="テスト製品説明", + content=( + "この製品は高品質な素材で作られています。" + "耐久性に優れ、長期間使用できます。" + "また、環境に配慮した製造プロセスを採用しています。" + "返品は購入後30日以内に限ります。" + ), + ) + + +@pytest.fixture +def pipeline(): + return RAGPipeline.create_mock() + + +@pytest.fixture +def indexed_pipeline(basic_doc): + p = RAGPipeline.create_mock() + p.index_document(basic_doc) + return p + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# 1. Enum & Config +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class TestEnumsAndConfig: + def test_embedding_model_enum_values(self): + assert GemmaEmbeddingModel.TEXT_EMBEDDING_004.value == "text-embedding-004" + assert GemmaEmbeddingModel.EMBEDDING_001.value == "embedding-001" + assert GemmaEmbeddingModel.MOCK.value == "mock" + + def test_chunking_strategy_enum(self): + assert ChunkingStrategy.FIXED_SIZE.value == "fixed_size" + assert ChunkingStrategy.SENTENCE.value == "sentence" + assert ChunkingStrategy.PARAGRAPH.value == "paragraph" + + def test_rag_status_enum(self): + assert RAGStatus.IDLE.value == "idle" + assert RAGStatus.READY.value == "ready" + + def test_embedding_config_defaults(self): + cfg = EmbeddingConfig() + assert cfg.model == GemmaEmbeddingModel.TEXT_EMBEDDING_004 + assert cfg.dim == 768 + assert cfg.task_type == "RETRIEVAL_DOCUMENT" + assert cfg.timeout == 30 + + def test_chunking_config_defaults(self): + cfg = ChunkingConfig() + assert cfg.strategy == ChunkingStrategy.SENTENCE + assert cfg.chunk_size == 512 + assert cfg.overlap == 64 + assert cfg.min_chars == 20 + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# 2. Data Models +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class TestDataModels: + def test_rag_document_defaults(self): + doc = RAGDocument(id="d1", title="タイトル", content="コンテンツ") + assert doc.id == "d1" + assert doc.metadata == {} + + def test_rag_document_auto_id(self): + doc = RAGDocument(id="", title="T", content="C") + assert len(doc.id) > 0 # __post_init__ で UUID 生成 + + def test_rag_chunk_full_text(self): + chunk = RAGChunk( + id="c1", doc_id="d1", doc_title="製品説明", + text="高品質な素材です。", chunk_index=0, + ) + assert "製品説明" in chunk.full_text + assert "高品質な素材です。" in chunk.full_text + + def test_retrieval_result_to_dict(self): + chunk = RAGChunk( + id="c1", doc_id="d1", doc_title="T", text="本文", chunk_index=0, + ) + result = RetrievalResult(chunk=chunk, score=0.85, rank=1) + d = result.to_dict() + assert d["chunk_id"] == "c1" + assert d["score"] == 0.85 + assert d["rank"] == 1 + assert d["doc_title"] == "T" + + def test_rag_answer_properties(self): + answer = RAGAnswer(query="質問", text="回答") + assert answer.success is True + d = answer.to_dict() + assert d["query"] == "質問" + assert d["answer"] == "回答" + + def test_rag_answer_empty_text(self): + answer = RAGAnswer(query="質問", text="") + assert answer.success is False + + def test_index_stats_defaults(self): + stats = IndexStats() + assert stats.doc_count == 0 + assert stats.chunk_count == 0 + assert stats.status == RAGStatus.IDLE + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# 3. GemmaEmbeddingProvider (Mock) +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class TestGemmaEmbeddingProviderMock: + def test_mock_embed_returns_correct_dim(self, mock_embedder): + vec = mock_embedder.embed("テストテキスト") + assert len(vec) == 768 + + def test_mock_embed_is_normalized(self, mock_embedder): + vec = mock_embedder.embed("hello world") + norm = math.sqrt(sum(v * v for v in vec)) + assert abs(norm - 1.0) < 1e-6 + + def test_mock_embed_is_deterministic(self, mock_embedder): + text = "同じテキストは同じ埋め込み" + v1 = mock_embedder.embed(text) + v2 = mock_embedder.embed(text) + assert v1 == v2 + + def test_mock_embed_different_texts_different_vectors(self, mock_embedder): + v1 = mock_embedder.embed("りんご") + v2 = mock_embedder.embed("自動車") + assert v1 != v2 + + def test_mock_embed_similar_texts_higher_cosine(self, mock_embedder): + v_dog1 = mock_embedder.embed("犬 ペット 動物") + v_dog2 = mock_embedder.embed("犬 かわいい ペット") + v_car = mock_embedder.embed("自動車 エンジン 速度") + sim_dog = _cosine_sim(v_dog1, v_dog2) + sim_diff = _cosine_sim(v_dog1, v_car) + assert sim_dog > sim_diff + + def test_mock_embed_batch(self, mock_embedder): + texts = ["テキスト1", "テキスト2", "テキスト3"] + vecs = mock_embedder.embed_batch(texts) + assert len(vecs) == 3 + for vec in vecs: + assert len(vec) == 768 + + def test_mock_embed_empty_text(self, mock_embedder): + vec = mock_embedder.embed("") + assert len(vec) == 768 + # 全ゼロになるはずはない（正規化後は 1/768 ≈ 0） + # ゼロ除算は 1e-9 で保護されているので全ゼロも返る可能性あり + assert isinstance(vec, list) + + def test_provider_dim_property(self, mock_embedder): + assert mock_embedder.dim == 768 + + def test_is_mock_flag_true_without_api_key(self): + cfg = EmbeddingConfig( + model=GemmaEmbeddingModel.TEXT_EMBEDDING_004, + api_key="", + ) + provider = GemmaEmbeddingProvider(cfg) + assert provider._is_mock is True + + def test_is_mock_flag_false_with_api_key(self): + cfg = EmbeddingConfig( + model=GemmaEmbeddingModel.TEXT_EMBEDDING_004, + api_key="dummy_key", + ) + provider = GemmaEmbeddingProvider(cfg) + assert provider._is_mock is False + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# 4. DocumentChunker +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class TestDocumentChunker: + def test_sentence_chunker_basic(self, basic_doc): + chunker = DocumentChunker() + chunks = chunker.chunk(basic_doc) + assert len(chunks) >= 1 + for c in chunks: + assert c.doc_id == "doc1" + assert c.doc_title == "テスト製品説明" + assert len(c.text) >= 20 + + def test_chunk_ids_are_unique(self, basic_doc): + chunker = DocumentChunker() + chunks = chunker.chunk(basic_doc) + ids = [c.id for c in chunks] + assert len(ids) == len(set(ids)) + + def test_chunk_index_sequential(self, basic_doc): + chunker = DocumentChunker(ChunkingConfig(strategy=ChunkingStrategy.FIXED_SIZE, chunk_size=50)) + chunks = chunker.chunk(basic_doc) + for i, c in enumerate(chunks): + assert c.chunk_index == i + + def test_fixed_size_chunker(self): + doc = RAGDocument(id="d1", title="T", content="A" * 300) + cfg = ChunkingConfig(strategy=ChunkingStrategy.FIXED_SIZE, chunk_size=100, overlap=0) + chunker = DocumentChunker(cfg) + chunks = chunker.chunk(doc) + assert len(chunks) == 3 + + def test_fixed_size_with_overlap(self): + doc = RAGDocument(id="d1", title="T", content="A" * 200) + cfg = ChunkingConfig(strategy=ChunkingStrategy.FIXED_SIZE, chunk_size=100, overlap=50) + chunker = DocumentChunker(cfg) + chunks = chunker.chunk(doc) + # オーバーラップあり → 通常より多くのチャンク + assert len(chunks) > 2 + + def test_paragraph_chunker(self): + content = ( + "段落1の内容です。詳細な説明が含まれます。\n\n" + "段落2の内容です。さらに詳しい情報があります。\n\n" + "段落3の内容です。最後の段落になります。" + ) + doc = RAGDocument(id="d1", title="T", content=content) + cfg = ChunkingConfig(strategy=ChunkingStrategy.PARAGRAPH) + chunker = DocumentChunker(cfg) + chunks = chunker.chunk(doc) + assert len(chunks) == 3 + + def test_min_chars_filter(self): + content = "短い。\n\nこれは十分な長さのコンテンツです。詳細な説明があります。" + doc = RAGDocument(id="d1", title="T", content=content) + cfg = ChunkingConfig(strategy=ChunkingStrategy.PARAGRAPH, min_chars=20) + chunker = DocumentChunker(cfg) + chunks = chunker.chunk(doc) + # 「短い。」は min_chars=20 未満なので除外 + for c in chunks: + assert len(c.text) >= 20 + + def test_empty_content_returns_no_chunks(self): + doc = RAGDocument(id="d1", title="T", content="") + chunker = DocumentChunker() + chunks = chunker.chunk(doc) + assert len(chunks) == 0 + + def test_full_text_includes_title(self, basic_doc): + chunker = DocumentChunker() + chunks = chunker.chunk(basic_doc) + for c in chunks: + assert basic_doc.title in c.full_text + + def test_metadata_propagated(self): + doc = RAGDocument( + id="d1", title="T", content="コンテンツです。" * 5, + metadata={"source": "wiki", "lang": "ja"}, + ) + chunker = DocumentChunker() + chunks = chunker.chunk(doc) + for c in chunks: + assert c.metadata.get("source") == "wiki" + assert c.metadata.get("lang") == "ja" + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# 5. _ChunkVectorStore +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class TestChunkVectorStore: + def _make_chunk(self, cid: str, doc_id: str) -> RAGChunk: + return RAGChunk(id=cid, doc_id=doc_id, doc_title="T", text="テスト", chunk_index=0) + + def test_upsert_and_count(self): + store = _ChunkVectorStore() + chunk = self._make_chunk("c1", "d1") + store.upsert(chunk, [0.1, 0.2, 0.3]) + assert store.chunk_count == 1 + + def test_upsert_overwrite(self): + store = _ChunkVectorStore() + chunk = self._make_chunk("c1", "d1") + store.upsert(chunk, [0.1, 0.2]) + store.upsert(chunk, [0.3, 0.4]) + assert store.chunk_count == 1 + + def test_delete_by_doc(self): + store = _ChunkVectorStore() + for i in range(3): + store.upsert(self._make_chunk(f"c{i}", "d1"), [float(i), 0.0]) + store.upsert(self._make_chunk("cx", "d2"), [1.0, 0.0]) + removed = store.delete_by_doc("d1") + assert removed == 3 + assert store.chunk_count == 1 + + def test_search_returns_top_k(self): + store = _ChunkVectorStore() + for i in range(10): + chunk = self._make_chunk(f"c{i}", "d1") + store.upsert(chunk, [float(i), 0.0]) + results = store.search([1.0, 0.0], top_k=3) + assert len(results) == 3 + + def test_search_sorted_by_score(self): + store = _ChunkVectorStore() + store.upsert(self._make_chunk("c1", "d1"), [1.0, 0.0]) + store.upsert(self._make_chunk("c2", "d1"), [0.0, 1.0]) + store.upsert(self._make_chunk("c3", "d1"), [0.7, 0.7]) + results = store.search([1.0, 0.0], top_k=3) + scores = [s for _, s in results] + assert scores == sorted(scores, reverse=True) + + def test_doc_filter(self): + store = _ChunkVectorStore() + store.upsert(self._make_chunk("c1", "d1"), [1.0, 0.0]) + store.upsert(self._make_chunk("c2", "d2"), [0.9, 0.1]) + results = store.search([1.0, 0.0], top_k=5, doc_filter=["d1"]) + assert all(c.doc_id == "d1" for c, _ in results) + + def test_doc_ids_property(self): + store = _ChunkVectorStore() + store.upsert(self._make_chunk("c1", "d1"), [1.0, 0.0]) + store.upsert(self._make_chunk("c2", "d2"), [0.5, 0.5]) + assert set(store.doc_ids) == {"d1", "d2"} + + def test_clear(self): + store = _ChunkVectorStore() + store.upsert(self._make_chunk("c1", "d1"), [1.0]) + store.clear() + assert store.chunk_count == 0 + + def test_cosine_sim_perfect(self): + a = [1.0, 0.0, 0.0] + assert abs(_cosine_sim(a, a) - 1.0) < 1e-6 + + def test_cosine_sim_orthogonal(self): + a = [1.0, 0.0] + b = [0.0, 1.0] + assert abs(_cosine_sim(a, b)) < 1e-6 + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# 6. RAGIndexer +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class TestRAGIndexer: + def _make_indexer(self): + cfg = EmbeddingConfig(model=GemmaEmbeddingModel.MOCK) + embedder = GemmaEmbeddingProvider(cfg) + return RAGIndexer(embedder) + + def test_index_returns_chunks(self, basic_doc): + indexer = self._make_indexer() + chunks = indexer.index(basic_doc) + assert len(chunks) >= 1 + + def test_stats_after_index(self, basic_doc): + indexer = self._make_indexer() + indexer.index(basic_doc) + stats = indexer.stats + assert stats.doc_count == 1 + assert stats.chunk_count >= 1 + assert stats.status == RAGStatus.READY + + def test_stats_before_index(self): + indexer = self._make_indexer() + stats = indexer.stats + assert stats.doc_count == 0 + assert stats.status == RAGStatus.IDLE + + def test_index_overwrites_existing(self, basic_doc): + indexer = self._make_indexer() + indexer.index(basic_doc) + count_before = indexer.store.chunk_count + indexer.index(basic_doc) + count_after = indexer.store.chunk_count + # 上書きなので同じ件数のはず + assert count_after == count_before + + def test_index_multiple_docs(self): + indexer = self._make_indexer() + docs = [ + RAGDocument(id=f"d{i}", title=f"Doc{i}", content=f"内容です。" * 5) + for i in range(3) + ] + for doc in docs: + indexer.index(doc) + assert indexer.stats.doc_count == 3 + + def test_index_batch(self): + indexer = self._make_indexer() + docs = [RAGDocument(id=f"d{i}", title=f"T{i}", content="コンテンツ。" * 5) for i in range(5)] + result = indexer.index_batch(docs) + assert len(result) == 5 + assert all(v >= 1 for v in result.values()) + + def test_delete_doc(self, basic_doc): + indexer = self._make_indexer() + indexer.index(basic_doc) + removed = indexer.delete("doc1") + assert removed >= 1 + assert indexer.stats.doc_count == 0 + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# 7. RAGRetriever +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class TestRAGRetriever: + def _make_retriever_with_indexer(self): + cfg = EmbeddingConfig(model=GemmaEmbeddingModel.MOCK) + embedder = GemmaEmbeddingProvider(cfg) + store = _ChunkVectorStore() + indexer = RAGIndexer(embedder, store=store) + retriever = RAGRetriever(embedder, store) + return indexer, retriever + + def test_retrieve_returns_results(self, basic_doc): + indexer, retriever = self._make_retriever_with_indexer() + indexer.index(basic_doc) + results = retriever.retrieve("素材") + assert len(results) >= 1 + + def test_retrieve_result_type(self, basic_doc): + indexer, retriever = self._make_retriever_with_indexer() + indexer.index(basic_doc) + results = retriever.retrieve("製品") + assert all(isinstance(r, RetrievalResult) for r in results) + + def test_retrieve_ranks_assigned(self, basic_doc): + indexer, retriever = self._make_retriever_with_indexer() + indexer.index(basic_doc) + results = retriever.retrieve("製品", top_k=3) + ranks = [r.rank for r in results] + assert ranks == list(range(1, len(results) + 1)) + + def test_retrieve_score_threshold(self, basic_doc): + indexer, retriever = self._make_retriever_with_indexer() + indexer.index(basic_doc) + results = retriever.retrieve("製品", score_threshold=0.99) + # 極めて高い閾値では結果が少なくなる + all_scores = [r.score for r in results] + assert all(s >= 0.99 for s in all_scores) + + def test_retrieve_top_k_limit(self, basic_doc): + indexer, retriever = self._make_retriever_with_indexer() + indexer.index(basic_doc) + results = retriever.retrieve("製品", top_k=1) + assert len(results) <= 1 + + def test_retrieve_texts(self, basic_doc): + indexer, retriever = self._make_retriever_with_indexer() + indexer.index(basic_doc) + texts = retriever.retrieve_texts("製品", top_k=3) + assert isinstance(texts, list) + assert all(isinstance(t, str) for t in texts) + + def test_retrieve_empty_store(self): + _, retriever = self._make_retriever_with_indexer() + results = retriever.retrieve("クエリ") + assert results == [] + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# 8. Generators +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class TestGenerators: + def test_mock_generator_with_context(self): + gen = MockLLMGenerator() + text = gen.generate("質問", ["チャンク1", "チャンク2"]) + assert len(text) > 0 + + def test_mock_generator_no_context(self): + gen = MockLLMGenerator() + text = gen.generate("質問", []) + assert "見つかりませんでした" in text + + def test_mock_generator_model_name(self): + gen = MockLLMGenerator() + assert gen.model_name == "mock-llm" + + def test_gemma_llm_generator_model_name(self): + gen = GemmaLLMGenerator(api_key="dummy", model="gemma-3-27b-it") + assert gen.model_name == "gemma-3-27b-it" + + def test_gemma_llm_generator_default_model(self): + gen = GemmaLLMGenerator(api_key="dummy") + assert "gemma" in gen.model_name.lower() + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# 9. RAGPipeline — 統合テスト +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class TestRAGPipeline: + def test_create_mock(self): + p = RAGPipeline.create_mock() + assert isinstance(p, RAGPipeline) + + def test_create_production_no_gen_key(self): + p = RAGPipeline.create_production( + embedding_api_key="dummy_emb", + generation_api_key="", + ) + assert isinstance(p, RAGPipeline) + # api_key がないので embedder は mock 扱いになる + # (api_key が dummy_emb なので _is_mock=False だが httpエラーは実際の呼び出し時) + + def test_index_document(self, pipeline, basic_doc): + count = pipeline.index_document(basic_doc) + assert count >= 1 + + def test_index_batch(self, pipeline): + docs = [ + RAGDocument(id=f"d{i}", title=f"Doc{i}", content="内容です。" * 5) + for i in range(4) + ] + result = pipeline.index_batch(docs) + assert len(result) == 4 + + def test_stats_idle_initially(self, pipeline): + stats = pipeline.stats + assert stats.status == RAGStatus.IDLE + + def test_stats_ready_after_index(self, indexed_pipeline): + stats = indexed_pipeline.stats + assert stats.status == RAGStatus.READY + assert stats.doc_count == 1 + assert stats.chunk_count >= 1 + + def test_search_returns_results(self, indexed_pipeline): + results = indexed_pipeline.search("製品の素材") + assert isinstance(results, list) + assert all(isinstance(r, RetrievalResult) for r in results) + + def test_search_empty_pipeline(self, pipeline): + results = pipeline.search("質問") + assert results == [] + + def test_query_returns_rag_answer(self, indexed_pipeline): + answer = indexed_pipeline.query("この製品はどんな素材ですか？") + assert isinstance(answer, RAGAnswer) + assert answer.success is True + assert len(answer.sources) >= 0 + + def test_query_answer_has_query(self, indexed_pipeline): + q = "返品期限は？" + answer = indexed_pipeline.query(q) + assert answer.query == q + + def test_query_latency_tracked(self, indexed_pipeline): + answer = indexed_pipeline.query("素材") + assert answer.latency_ms >= 0 + + def test_query_model_used(self, indexed_pipeline): + answer = indexed_pipeline.query("素材") + assert answer.model_used == "mock-llm" + + def test_query_sources_dict_format(self, indexed_pipeline): + answer = indexed_pipeline.query("製品") + for src in answer.sources: + assert "chunk_id" in src + assert "score" in src + assert "rank" in src + + def test_delete_document(self, pipeline, basic_doc): + pipeline.index_document(basic_doc) + removed = pipeline.delete_document("doc1") + assert removed >= 1 + assert pipeline.stats.doc_count == 0 + + def test_re_index_after_delete(self, pipeline, basic_doc): + pipeline.index_document(basic_doc) + pipeline.delete_document("doc1") + count = pipeline.index_document(basic_doc) + assert count >= 1 + assert pipeline.stats.doc_count == 1 + + def test_multiple_docs_query(self, pipeline): + docs = [ + RAGDocument(id="d1", title="猫", content="猫はかわいい動物です。肉食で夜行性の生き物。"), + RAGDocument(id="d2", title="車", content="自動車はガソリンまたは電気で動く乗り物。"), + RAGDocument(id="d3", title="料理", content="料理は食材を加熱や調味料で味付けする技術。"), + ] + for doc in docs: + pipeline.index_document(doc) + answer = pipeline.query("猫について教えてください") + assert answer.success is True + + def test_to_status_dict(self, indexed_pipeline): + d = indexed_pipeline.to_status_dict() + assert "status" in d + assert "doc_count" in d + assert "chunk_count" in d + assert "embedder" in d + assert "generator" in d + + def test_status_dict_values(self, indexed_pipeline): + d = indexed_pipeline.to_status_dict() + assert d["status"] == RAGStatus.READY.value + assert d["doc_count"] == 1 + assert d["embedder"] == GemmaEmbeddingModel.MOCK.value + + def test_top_k_limit_in_query(self, pipeline): + for i in range(5): + doc = RAGDocument( + id=f"d{i}", title=f"文書{i}", + content=f"これは文書{i}の内容です。詳細な説明が含まれます。" + ) + pipeline.index_document(doc) + answer = pipeline.query("内容", top_k=2) + assert len(answer.sources) <= 2 + + def test_doc_filter_in_search(self, pipeline): + docs = [ + RAGDocument(id="d1", title="A", content="アルファの情報です。" * 3), + RAGDocument(id="d2", title="B", content="ベータの情報です。" * 3), + ] + for doc in docs: + pipeline.index_document(doc) + results = pipeline.search("情報", doc_filter=["d1"]) + assert all(r.chunk.doc_id == "d1" for r in results) + + def test_score_threshold_in_query(self, indexed_pipeline): + answer = indexed_pipeline.query("全く関係ない話", score_threshold=0.999) + # 非常に高い閾値 → ソースなし or 少数のみ + for src in answer.sources: + assert src["score"] >= 0.999 + + def test_rag_answer_to_dict_keys(self, indexed_pipeline): + answer = indexed_pipeline.query("製品") + d = answer.to_dict() + assert set(d.keys()) == {"query", "answer", "sources", "latency_ms", "model_used"} + + +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# 10. エッジケース +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +class TestEdgeCases: + def test_very_long_document(self, pipeline): + """長文ドキュメント（10000字）でもクラッシュしない。""" + content = "これは長い文書です。詳細な内容が続きます。" * 200 + doc = RAGDocument(id="long", title="長文", content=content) + count = pipeline.index_document(doc) + assert count >= 1 + + def test_unicode_content(self, pipeline): + """Unicode / emoji / 記号を含んでもクラッシュしない。""" + doc = RAGDocument( + id="u1", title="Unicode テスト", + content="🎉 祝！ 特殊文字 ①②③ αβγ → ← ↑ \n\n次の段落。", + ) + answer = pipeline.index_document(doc) + assert answer >= 0 + + def test_single_word_content(self, pipeline): + """非常に短いコンテンツ。""" + doc = RAGDocument(id="short", title="T", content="x" * 25) + pipeline.index_document(doc) + results = pipeline.search("x") + # チャンク化できるかどうか依存、少なくともクラッシュしない + assert isinstance(results, list) + + def test_query_after_all_docs_deleted(self, indexed_pipeline): + """全削除後のクエリ。""" + indexed_pipeline.delete_document("doc1") + answer = indexed_pipeline.query("製品") + assert isinstance(answer, RAGAnswer) + + def test_duplicate_doc_id_overwrites(self, pipeline): + """同 ID のドキュメントを二度インデックスすると上書き。""" + doc = RAGDocument(id="dup", title="初版", content="最初の内容です。詳細情報。" * 3) + pipeline.index_document(doc) + count_before = pipeline.stats.chunk_count + + doc2 = RAGDocument(id="dup", title="改版", content="更新された内容です。" * 3) + pipeline.index_document(doc2) + count_after = pipeline.stats.chunk_count + + assert pipeline.stats.doc_count == 1 # 2 ではなく 1 + # チャンク数は置き換え後の件数 + assert count_after >= 1 + + def test_chunking_config_sentence_strategy(self): + """文境界チャンキングで文末記号が正しく分割される。""" + content = ( + "第一文は長めの文章です。詳細な内容が含まれています。" + "第二文も同様に長い文章となっています。詳しい説明があります。" + ) + doc = RAGDocument(id="s1", title="T", content=content) + cfg = ChunkingConfig(strategy=ChunkingStrategy.SENTENCE, chunk_size=30) + chunker = DocumentChunker(cfg) + chunks = chunker.chunk(doc) + # chunk_size=30 で分割されるので複数チャンクになるはず + assert len(chunks) >= 1 + + def test_mock_embed_dimension_custom(self): + """カスタム dim で埋め込みが正しいサイズを返す。""" + cfg = EmbeddingConfig(model=GemmaEmbeddingModel.MOCK, dim=128) + provider = GemmaEmbeddingProvider(cfg) + vec = provider.embed("テスト") + assert len(vec) == 128 + + def test_retrieval_result_score_range(self, indexed_pipeline): + """検索スコアは 0〜1 の範囲。""" + results = indexed_pipeline.search("製品") + for r in results: + assert 0.0 <= r.score <= 1.0 + 1e-6 # 浮動小数誤差許容