Spaces:
Sleeping
Sleeping
| from collections.abc import Sequence | |
| import pytest | |
| from techdocs_graphrag.models import Chunk, GraphExtraction, GraphFact, RetrievedChunk | |
| class FakeLanguageModel: | |
| def __init__(self, completion: str = "réponse") -> None: | |
| self.completion = completion | |
| self.completions: list[tuple[str, str]] = [] | |
| def complete(self, system_prompt: str, user_prompt: str, *, json_mode: bool = False) -> str: | |
| self.completions.append((system_prompt, user_prompt)) | |
| return self.completion | |
| def embed(self, texts: Sequence[str]) -> list[list[float]]: | |
| return [[float(len(text))] for text in texts] | |
| class FakeKnowledgeGraph: | |
| def __init__( | |
| self, | |
| chunks: list[RetrievedChunk] | None = None, | |
| lexical_matches: list[str] | None = None, | |
| facts: list[GraphFact] | None = None, | |
| ) -> None: | |
| self.chunks = chunks or [] | |
| self.lexical_matches = lexical_matches or [] | |
| self.facts = facts or [] | |
| self.expanded_with: list[str] | None = None | |
| self.extractions: list[tuple[Chunk, GraphExtraction]] = [] | |
| self.indexes_awaited = False | |
| def ensure_schema(self) -> None: | |
| pass | |
| def await_indexes(self) -> None: | |
| self.indexes_awaited = True | |
| def upsert_chunks(self, chunks: Sequence[Chunk], embeddings: Sequence[list[float]]) -> None: | |
| pass | |
| def upsert_extraction(self, chunk: Chunk, extraction: GraphExtraction) -> None: | |
| self.extractions.append((chunk, extraction)) | |
| def vector_search(self, embedding: list[float], top_k: int) -> list[RetrievedChunk]: | |
| return self.chunks[:top_k] | |
| def search_entities(self, text: str, limit: int) -> list[str]: | |
| return self.lexical_matches[:limit] | |
| def expand(self, seed_keys: Sequence[str], max_hops: int, limit: int) -> list[GraphFact]: | |
| self.expanded_with = list(seed_keys) | |
| return self.facts[:limit] | |
| def retrieved_chunk(document: str, entity_keys: list[str]) -> RetrievedChunk: | |
| return RetrievedChunk( | |
| id=f"{document}-0", | |
| document=document, | |
| position=0, | |
| text=f"contenu de {document}", | |
| score=0.9, | |
| entity_keys=entity_keys, | |
| ) | |
| def fake_llm() -> FakeLanguageModel: | |
| return FakeLanguageModel() | |