mirror of
https://github.com/Sosokker/plain-rag.git
synced 2025-12-18 22:44:03 +01:00
23 lines
512 B
Python
23 lines
512 B
Python
from typing import Protocol, TypedDict
|
|
|
|
import numpy as np
|
|
|
|
|
|
class SearchResult(TypedDict):
|
|
"""Type definition for search results."""
|
|
|
|
content: str
|
|
source: str
|
|
|
|
|
|
class EmbeddingModel(Protocol):
|
|
def embed_documents(self, texts: list[str]) -> list[np.ndarray]: ...
|
|
|
|
def embed_query(self, text: str) -> np.ndarray: ...
|
|
|
|
|
|
class VectorDB(Protocol):
|
|
def upsert_documents(self, documents: list[dict]) -> None: ...
|
|
|
|
def search(self, vector: np.ndarray, top_k: int) -> list[SearchResult]: ...
|