Spaces:
Sleeping
Sleeping
File size: 2,279 Bytes
da21448 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | 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,
)
@pytest.fixture
def fake_llm() -> FakeLanguageModel:
return FakeLanguageModel()
|