diff --git a/openrag/core/chunking/__init__.py b/openrag/core/chunking/__init__.py index e69de29bb..f2caca418 100644 --- a/openrag/core/chunking/__init__.py +++ b/openrag/core/chunking/__init__.py @@ -0,0 +1,6 @@ +"""ChunkingStrategy ABC + registry.""" + +from openrag.core.chunking.chunking_strategy import ChunkingStrategy +from openrag.core.chunking.registry import chunking_registry + +__all__ = ["ChunkingStrategy", "chunking_registry"] diff --git a/openrag/core/chunking/chunking_strategy.py b/openrag/core/chunking/chunking_strategy.py new file mode 100644 index 000000000..2f2dbf93d --- /dev/null +++ b/openrag/core/chunking/chunking_strategy.py @@ -0,0 +1,17 @@ +"""Abstract chunking strategy interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from openrag.core.models.chunk import Chunk +from openrag.core.models.document import ProcessedDocument + + +class ChunkingStrategy(ABC): + """Base class for all chunking strategies.""" + + @abstractmethod + def chunk(self, document: ProcessedDocument, partition: str = "default") -> list[Chunk]: + """Split a processed document into chunks.""" + ... diff --git a/openrag/core/chunking/registry.py b/openrag/core/chunking/registry.py new file mode 100644 index 000000000..80ca0986c --- /dev/null +++ b/openrag/core/chunking/registry.py @@ -0,0 +1,6 @@ +"""Chunking strategy registry.""" + +from openrag.core.chunking.chunking_strategy import ChunkingStrategy +from openrag.core.utils.registry import Registry + +chunking_registry: Registry[ChunkingStrategy] = Registry("chunking") diff --git a/openrag/core/embeddings/__init__.py b/openrag/core/embeddings/__init__.py index e69de29bb..bb83819db 100644 --- a/openrag/core/embeddings/__init__.py +++ b/openrag/core/embeddings/__init__.py @@ -0,0 +1,6 @@ +"""Embedder ABC + registry.""" + +from openrag.core.embeddings.embedder import Embedder +from openrag.core.embeddings.registry import embedder_registry + +__all__ = ["Embedder", "embedder_registry"] diff --git a/openrag/core/embeddings/embedder.py b/openrag/core/embeddings/embedder.py new file mode 100644 index 000000000..1ac34aa05 --- /dev/null +++ b/openrag/core/embeddings/embedder.py @@ -0,0 +1,25 @@ +"""Abstract embedder interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class Embedder(ABC): + """Base class for all embedding providers.""" + + @abstractmethod + async def embed(self, texts: list[str]) -> list[list[float]]: + """Embed a batch of texts, returning vectors.""" + ... + + @abstractmethod + async def embed_single(self, text: str) -> list[float]: + """Embed a single text.""" + ... + + @property + @abstractmethod + def dimension(self) -> int: + """Return the embedding dimension.""" + ... diff --git a/openrag/core/embeddings/registry.py b/openrag/core/embeddings/registry.py new file mode 100644 index 000000000..a79e08d7c --- /dev/null +++ b/openrag/core/embeddings/registry.py @@ -0,0 +1,6 @@ +"""Embedder registry.""" + +from openrag.core.embeddings.embedder import Embedder +from openrag.core.utils.registry import Registry + +embedder_registry: Registry[Embedder] = Registry("embedder") diff --git a/openrag/core/indexing/parsers/__init__.py b/openrag/core/indexing/parsers/__init__.py index e69de29bb..16e23bfb4 100644 --- a/openrag/core/indexing/parsers/__init__.py +++ b/openrag/core/indexing/parsers/__init__.py @@ -0,0 +1,6 @@ +"""DocumentParser ABC + registry.""" + +from openrag.core.indexing.parsers.document_parser import DocumentParser +from openrag.core.indexing.parsers.registry import parser_registry + +__all__ = ["DocumentParser", "parser_registry"] diff --git a/openrag/core/indexing/parsers/document_parser.py b/openrag/core/indexing/parsers/document_parser.py new file mode 100644 index 000000000..aedd0b383 --- /dev/null +++ b/openrag/core/indexing/parsers/document_parser.py @@ -0,0 +1,21 @@ +"""Abstract document parser interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from openrag.core.models.document import Document, ProcessedDocument + + +class DocumentParser(ABC): + """Base class for all document parsers (PDF, text, HTML, image, audio, etc.).""" + + @abstractmethod + async def parse(self, document: Document) -> ProcessedDocument: + """Parse a document into text blocks and images.""" + ... + + @abstractmethod + def supported_types(self) -> list[str]: + """Return list of DocumentType values this parser handles.""" + ... diff --git a/openrag/core/indexing/parsers/registry.py b/openrag/core/indexing/parsers/registry.py new file mode 100644 index 000000000..2975143b5 --- /dev/null +++ b/openrag/core/indexing/parsers/registry.py @@ -0,0 +1,6 @@ +"""Document parser registry.""" + +from openrag.core.indexing.parsers.document_parser import DocumentParser +from openrag.core.utils.registry import Registry + +parser_registry: Registry[DocumentParser] = Registry("parser") diff --git a/openrag/core/llm/__init__.py b/openrag/core/llm/__init__.py index e69de29bb..0a37f2307 100644 --- a/openrag/core/llm/__init__.py +++ b/openrag/core/llm/__init__.py @@ -0,0 +1,6 @@ +"""LLM ABC + registry.""" + +from openrag.core.llm.llm import LLM +from openrag.core.llm.registry import llm_registry + +__all__ = ["LLM", "llm_registry"] diff --git a/openrag/core/llm/llm.py b/openrag/core/llm/llm.py new file mode 100644 index 000000000..7e24d85b2 --- /dev/null +++ b/openrag/core/llm/llm.py @@ -0,0 +1,57 @@ +"""Abstract LLM interface.""" + +from __future__ import annotations + +import json +from abc import ABC, abstractmethod +from collections.abc import AsyncIterator + +from openrag.core.utils.exceptions import LLMParsingError + + +class LLM(ABC): + """Base class for all LLM providers.""" + + @abstractmethod + async def generate(self, prompt: str, **kwargs) -> str: + """Generate a completion for a prompt.""" + ... + + @abstractmethod + async def chat(self, messages: list[dict[str, str]], **kwargs) -> str: + """Chat completion with message list.""" + ... + + async def generate_json(self, prompt: str, **kwargs) -> dict: + """Generate a JSON response. Default: parse generate() output. + + Raises LLMParsingError if the LLM output is not valid JSON. + """ + response = await self.generate(prompt, **kwargs) + text = response.strip() + try: + return json.loads(text) + except json.JSONDecodeError as exc: + raise LLMParsingError( + raw_response=text, + parse_error=str(exc), + ) from exc + + async def chat_with_tools( + self, + messages: list[dict[str, str]], + tools: list[dict], + tool_choice: str | dict = "required", + **kwargs, + ) -> dict: + """Chat completion with function calling. + + Only supported by backends with function calling (e.g., vLLM). + Default raises NotImplementedError. + """ + raise NotImplementedError(f"{type(self).__name__} does not support tool calling") + + async def stream_chat(self, messages: list[dict[str, str]], **kwargs) -> AsyncIterator[str]: + """Stream chat completion. Default falls back to non-streaming.""" + result = await self.chat(messages, **kwargs) + yield result diff --git a/openrag/core/llm/registry.py b/openrag/core/llm/registry.py new file mode 100644 index 000000000..537faac4c --- /dev/null +++ b/openrag/core/llm/registry.py @@ -0,0 +1,6 @@ +"""LLM registry.""" + +from openrag.core.llm.llm import LLM +from openrag.core.utils.registry import Registry + +llm_registry: Registry[LLM] = Registry("llm") diff --git a/openrag/core/ports/__init__.py b/openrag/core/ports/__init__.py index e69de29bb..76344fc8b 100644 --- a/openrag/core/ports/__init__.py +++ b/openrag/core/ports/__init__.py @@ -0,0 +1,33 @@ +"""Port interfaces — repository ABCs + CatalogStore aggregate root.""" + +from openrag.core.ports.audit_log_repo import AuditLogRepository +from openrag.core.ports.catalog_store import CatalogStore +from openrag.core.ports.chunk_repo import ChunkRepository +from openrag.core.ports.conversation_repo import ConversationRepository +from openrag.core.ports.document_repo import DocumentRepository +from openrag.core.ports.entity_repo import EntityRepository +from openrag.core.ports.idempotency_repo import IdempotencyRepository +from openrag.core.ports.job_repo import JobRepository +from openrag.core.ports.model_endpoint_repo import ModelEndpointRepository +from openrag.core.ports.partition_repo import PartitionRepository +from openrag.core.ports.preset_repo import PresetRepository +from openrag.core.ports.prompt_repo import PromptRepository +from openrag.core.ports.topic_tag_repo import TopicTagRepository +from openrag.core.ports.user_repo import UserRepository + +__all__ = [ + "AuditLogRepository", + "CatalogStore", + "ChunkRepository", + "ConversationRepository", + "DocumentRepository", + "EntityRepository", + "IdempotencyRepository", + "JobRepository", + "ModelEndpointRepository", + "PartitionRepository", + "PresetRepository", + "PromptRepository", + "TopicTagRepository", + "UserRepository", +] diff --git a/openrag/core/ports/audit_log_repo.py b/openrag/core/ports/audit_log_repo.py new file mode 100644 index 000000000..2318bc4d1 --- /dev/null +++ b/openrag/core/ports/audit_log_repo.py @@ -0,0 +1,24 @@ +"""Audit log repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + + +class AuditLogRepository(ABC): + """Append-only audit trail.""" + + @abstractmethod + async def insert( + self, + user_id: int | None, + action: str, + resource_type: str, + resource_id: str | None = None, + details_json: dict | None = None, + request_id: str | None = None, + ) -> None: ... + + @abstractmethod + async def query(self, filters: dict[str, Any], offset: int = 0, limit: int = 50) -> list[dict]: ... diff --git a/openrag/core/ports/catalog_store.py b/openrag/core/ports/catalog_store.py new file mode 100644 index 000000000..5f275f6ae --- /dev/null +++ b/openrag/core/ports/catalog_store.py @@ -0,0 +1,88 @@ +"""CatalogStore — aggregate root composing all repository ports. + +Concrete implementations (e.g. PostgresStore) own the connection pool +and compose per-entity repository instances. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from openrag.core.ports.audit_log_repo import AuditLogRepository +from openrag.core.ports.chunk_repo import ChunkRepository +from openrag.core.ports.conversation_repo import ConversationRepository +from openrag.core.ports.document_repo import DocumentRepository +from openrag.core.ports.entity_repo import EntityRepository +from openrag.core.ports.idempotency_repo import IdempotencyRepository +from openrag.core.ports.job_repo import JobRepository +from openrag.core.ports.model_endpoint_repo import ModelEndpointRepository +from openrag.core.ports.partition_repo import PartitionRepository +from openrag.core.ports.preset_repo import PresetRepository +from openrag.core.ports.prompt_repo import PromptRepository +from openrag.core.ports.topic_tag_repo import TopicTagRepository +from openrag.core.ports.user_repo import UserRepository + + +class CatalogStore(ABC): + """Abstract interface for the relational catalog backing store. + + Concrete implementations (e.g. PostgresStore) live in the services layer. + """ + + @abstractmethod + async def initialize(self) -> None: ... + + @abstractmethod + async def shutdown(self) -> None: ... + + @property + @abstractmethod + def document_repo(self) -> DocumentRepository: ... + + @property + @abstractmethod + def job_repo(self) -> JobRepository: ... + + @property + @abstractmethod + def user_repo(self) -> UserRepository: ... + + @property + @abstractmethod + def prompt_repo(self) -> PromptRepository: ... + + @property + @abstractmethod + def partition_repo(self) -> PartitionRepository: ... + + @property + @abstractmethod + def model_endpoint_repo(self) -> ModelEndpointRepository: ... + + @property + @abstractmethod + def preset_repo(self) -> PresetRepository: ... + + @property + @abstractmethod + def chunk_repo(self) -> ChunkRepository: ... + + @property + @abstractmethod + def entity_repo(self) -> EntityRepository: ... + + @property + @abstractmethod + def topic_tag_repo(self) -> TopicTagRepository: ... + + @property + @abstractmethod + def conversation_repo(self) -> ConversationRepository: ... + + @property + @abstractmethod + def audit_log_repo(self) -> AuditLogRepository: ... + + @property + @abstractmethod + def idempotency_repo(self) -> IdempotencyRepository: ... diff --git a/openrag/core/ports/chunk_repo.py b/openrag/core/ports/chunk_repo.py new file mode 100644 index 000000000..e4b986d8e --- /dev/null +++ b/openrag/core/ports/chunk_repo.py @@ -0,0 +1,39 @@ +"""Chunk repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class ChunkRepository(ABC): + """Bulk CRUD operations for text chunks.""" + + @abstractmethod + async def bulk_insert(self, chunks: list[dict]) -> int: + """Insert multiple chunks. Returns count of inserted rows.""" + ... + + @abstractmethod + async def get_by_ids(self, chunk_ids: list[str]) -> list[dict]: + """Batch fetch chunks by IDs.""" + ... + + @abstractmethod + async def get_by_document_id(self, document_id: str) -> list[dict]: + """Fetch all chunks for a document, ordered by chunk_index.""" + ... + + @abstractmethod + async def delete_by_document_id(self, document_id: str) -> int: + """Delete all chunks belonging to a document. Returns count.""" + ... + + @abstractmethod + async def delete_by_partition(self, partition: str) -> int: + """Delete all chunks in a partition. Returns count.""" + ... + + @abstractmethod + async def bm25_search(self, query_text: str, partition: str, top_k: int = 20) -> list[dict]: + """Full-text search using tsvector column with ts_rank scoring.""" + ... diff --git a/openrag/core/ports/conversation_repo.py b/openrag/core/ports/conversation_repo.py new file mode 100644 index 000000000..6b1938f11 --- /dev/null +++ b/openrag/core/ports/conversation_repo.py @@ -0,0 +1,29 @@ +"""Conversation repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from openrag.core.models.conversation import Conversation, Message + + +class ConversationRepository(ABC): + """CRUD operations for conversations and messages.""" + + @abstractmethod + async def create_conversation(self, conversation: Conversation) -> Conversation: ... + + @abstractmethod + async def get_conversation(self, conversation_id: str) -> Conversation | None: ... + + @abstractmethod + async def list_conversations(self, user_id: int, partition: str | None = None) -> list[Conversation]: ... + + @abstractmethod + async def delete_conversation(self, conversation_id: str) -> bool: ... + + @abstractmethod + async def add_message(self, message: Message) -> Message: ... + + @abstractmethod + async def list_messages(self, conversation_id: str) -> list[Message]: ... diff --git a/openrag/core/ports/document_repo.py b/openrag/core/ports/document_repo.py new file mode 100644 index 000000000..289feed6f --- /dev/null +++ b/openrag/core/ports/document_repo.py @@ -0,0 +1,42 @@ +"""Document repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +from openrag.core.models.catalog import DocumentRecord + + +class DocumentRepository(ABC): + """CRUD operations for documents.""" + + @abstractmethod + async def create_document(self, doc: DocumentRecord) -> DocumentRecord: ... + + @abstractmethod + async def get_document(self, document_id: str) -> DocumentRecord | None: ... + + @abstractmethod + async def list_documents( + self, + partition: str | list[str] | None = None, + status: str | None = None, + offset: int = 0, + limit: int = 50, + ) -> list[DocumentRecord]: ... + + @abstractmethod + async def update_document(self, document_id: str, **fields: Any) -> DocumentRecord | None: ... + + @abstractmethod + async def delete_document(self, document_id: str) -> bool: ... + + @abstractmethod + async def delete_documents_by_partition(self, partition: str) -> int: ... + + @abstractmethod + async def count_documents(self, partition: str | list[str] | None = None, status: str | None = None) -> int: ... + + @abstractmethod + async def file_exists_in_partition(self, file_id: str, partition: str) -> bool: ... diff --git a/openrag/core/ports/entity_repo.py b/openrag/core/ports/entity_repo.py new file mode 100644 index 000000000..a8520484f --- /dev/null +++ b/openrag/core/ports/entity_repo.py @@ -0,0 +1,21 @@ +"""Entity repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class EntityRepository(ABC): + """CRUD operations for extracted entities.""" + + @abstractmethod + async def upsert(self, partition: str, entity_type: str, canonical_name: str, aliases: list[str]) -> str: ... + + @abstractmethod + async def search(self, partition: str, query: str, top_k: int = 10) -> list[dict]: ... + + @abstractmethod + async def get_by_document(self, document_id: str) -> list[dict]: ... + + @abstractmethod + async def delete_by_document(self, document_id: str) -> int: ... diff --git a/openrag/core/ports/idempotency_repo.py b/openrag/core/ports/idempotency_repo.py new file mode 100644 index 000000000..8a5f65dcd --- /dev/null +++ b/openrag/core/ports/idempotency_repo.py @@ -0,0 +1,15 @@ +"""Idempotency key repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class IdempotencyRepository(ABC): + """Cache for request idempotency keys.""" + + @abstractmethod + async def get_by_hash(self, key_hash: str) -> dict | None: ... + + @abstractmethod + async def store(self, key_hash: str, http_method: str, status_code: int, response_body: bytes) -> None: ... diff --git a/openrag/core/ports/job_repo.py b/openrag/core/ports/job_repo.py new file mode 100644 index 000000000..5e20b6830 --- /dev/null +++ b/openrag/core/ports/job_repo.py @@ -0,0 +1,24 @@ +"""Job repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +from openrag.core.models.catalog import IndexationJob + + +class JobRepository(ABC): + """CRUD operations for indexation jobs.""" + + @abstractmethod + async def create_job(self, job: IndexationJob) -> IndexationJob: ... + + @abstractmethod + async def get_job(self, job_id: str) -> IndexationJob | None: ... + + @abstractmethod + async def list_jobs(self, status: str | None = None, offset: int = 0, limit: int = 50) -> list[IndexationJob]: ... + + @abstractmethod + async def update_job(self, job_id: str, **fields: Any) -> IndexationJob | None: ... diff --git a/openrag/core/ports/model_endpoint_repo.py b/openrag/core/ports/model_endpoint_repo.py new file mode 100644 index 000000000..e907a0ce0 --- /dev/null +++ b/openrag/core/ports/model_endpoint_repo.py @@ -0,0 +1,21 @@ +"""Model endpoint repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class ModelEndpointRepository(ABC): + """CRUD operations for model endpoint configurations.""" + + @abstractmethod + async def get(self, name: str, model_type: str) -> dict | None: ... + + @abstractmethod + async def list_all(self, model_type: str | None = None) -> list[dict]: ... + + @abstractmethod + async def upsert(self, name: str, model_type: str, config: dict) -> dict: ... + + @abstractmethod + async def delete(self, name: str, model_type: str) -> bool: ... diff --git a/openrag/core/ports/partition_repo.py b/openrag/core/ports/partition_repo.py new file mode 100644 index 000000000..00898ce76 --- /dev/null +++ b/openrag/core/ports/partition_repo.py @@ -0,0 +1,24 @@ +"""Partition repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class PartitionRepository(ABC): + """CRUD operations for partitions.""" + + @abstractmethod + async def create_partition(self, name: str, user_id: int | None = None) -> dict: ... + + @abstractmethod + async def get_partition(self, name: str) -> dict | None: ... + + @abstractmethod + async def list_partitions(self) -> list[dict]: ... + + @abstractmethod + async def delete_partition(self, name: str) -> bool: ... + + @abstractmethod + async def partition_exists(self, name: str) -> bool: ... diff --git a/openrag/core/ports/preset_repo.py b/openrag/core/ports/preset_repo.py new file mode 100644 index 000000000..363b654fc --- /dev/null +++ b/openrag/core/ports/preset_repo.py @@ -0,0 +1,21 @@ +"""Preset repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class PresetRepository(ABC): + """CRUD operations for pipeline presets.""" + + @abstractmethod + async def get(self, name: str, preset_type: str) -> dict | None: ... + + @abstractmethod + async def list_all(self, preset_type: str | None = None) -> list[dict]: ... + + @abstractmethod + async def upsert(self, name: str, preset_type: str, config: dict) -> dict: ... + + @abstractmethod + async def delete(self, name: str, preset_type: str) -> bool: ... diff --git a/openrag/core/ports/prompt_repo.py b/openrag/core/ports/prompt_repo.py new file mode 100644 index 000000000..43336f2f5 --- /dev/null +++ b/openrag/core/ports/prompt_repo.py @@ -0,0 +1,32 @@ +"""Prompt repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +from openrag.core.models.prompt import Prompt + + +class PromptRepository(ABC): + """CRUD operations for prompt templates.""" + + @abstractmethod + async def create_prompt(self, prompt: Prompt) -> Prompt: ... + + @abstractmethod + async def get_prompt(self, prompt_id: str) -> Prompt | None: ... + + @abstractmethod + async def get_by_type(self, prompt_type: str) -> list[Prompt]: ... + + @abstractmethod + async def get_active(self, prompt_type: str) -> Prompt | None: ... + + @abstractmethod + async def list_prompts(self) -> list[Prompt]: ... + + @abstractmethod + async def update_prompt(self, prompt_id: str, content: str) -> Prompt | None: ... + + @abstractmethod + async def delete_prompt(self, prompt_id: str) -> bool: ... diff --git a/openrag/core/ports/topic_tag_repo.py b/openrag/core/ports/topic_tag_repo.py new file mode 100644 index 000000000..99287f46c --- /dev/null +++ b/openrag/core/ports/topic_tag_repo.py @@ -0,0 +1,21 @@ +"""Topic tag repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class TopicTagRepository(ABC): + """CRUD operations for document topic tags.""" + + @abstractmethod + async def bulk_insert(self, tags: list[dict]) -> int: ... + + @abstractmethod + async def get_by_document(self, document_id: str) -> list[dict]: ... + + @abstractmethod + async def delete_by_document(self, document_id: str) -> int: ... + + @abstractmethod + async def search(self, partition: str, tag: str, top_k: int = 10) -> list[dict]: ... diff --git a/openrag/core/ports/user_repo.py b/openrag/core/ports/user_repo.py new file mode 100644 index 000000000..0e7982a23 --- /dev/null +++ b/openrag/core/ports/user_repo.py @@ -0,0 +1,50 @@ +"""User repository interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +from openrag.core.models.user import User, UserPartition + + +class UserRepository(ABC): + """CRUD operations for users and partition memberships.""" + + @abstractmethod + async def create_user(self, user: User) -> User: ... + + @abstractmethod + async def get_user(self, user_id: int) -> User | None: ... + + @abstractmethod + async def get_user_by_token(self, token_hash: str) -> User | None: ... + + @abstractmethod + async def get_user_by_external_id(self, external_id: str) -> User | None: ... + + @abstractmethod + async def list_users(self, offset: int = 0, limit: int = 50) -> list[User]: ... + + @abstractmethod + async def update_user(self, user_id: int, **fields: Any) -> User | None: ... + + @abstractmethod + async def delete_user(self, user_id: int) -> bool: ... + + # ── Partition memberships ───────────────────────────────────────── + + @abstractmethod + async def assign_partition(self, assignment: UserPartition) -> UserPartition: ... + + @abstractmethod + async def remove_partition(self, user_id: int, partition: str) -> bool: ... + + @abstractmethod + async def list_user_partitions(self, user_id: int) -> list[UserPartition]: ... + + @abstractmethod + async def list_partition_users(self, partition: str) -> list[UserPartition]: ... + + @abstractmethod + async def update_partition_role(self, user_id: int, partition: str, role: str) -> bool: ... diff --git a/openrag/core/rerankers/__init__.py b/openrag/core/rerankers/__init__.py index e69de29bb..f27129c2a 100644 --- a/openrag/core/rerankers/__init__.py +++ b/openrag/core/rerankers/__init__.py @@ -0,0 +1,6 @@ +"""Reranker ABC + registry.""" + +from openrag.core.rerankers.registry import reranker_registry +from openrag.core.rerankers.reranker import Reranker + +__all__ = ["Reranker", "reranker_registry"] diff --git a/openrag/core/rerankers/registry.py b/openrag/core/rerankers/registry.py new file mode 100644 index 000000000..103feef02 --- /dev/null +++ b/openrag/core/rerankers/registry.py @@ -0,0 +1,6 @@ +"""Reranker registry.""" + +from openrag.core.rerankers.reranker import Reranker +from openrag.core.utils.registry import Registry + +reranker_registry: Registry[Reranker] = Registry("reranker") diff --git a/openrag/core/rerankers/reranker.py b/openrag/core/rerankers/reranker.py new file mode 100644 index 000000000..9f2ea60da --- /dev/null +++ b/openrag/core/rerankers/reranker.py @@ -0,0 +1,17 @@ +"""Abstract reranker interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class Reranker(ABC): + """Base class for all reranking providers.""" + + @abstractmethod + async def rerank(self, query: str, documents: list[str], top_k: int | None = None) -> list[tuple[int, float]]: + """Rerank documents for a query. + + Returns list of (original_index, score) sorted by relevance. + """ + ... diff --git a/openrag/core/vector_stores/__init__.py b/openrag/core/vector_stores/__init__.py index e69de29bb..501aa1d4d 100644 --- a/openrag/core/vector_stores/__init__.py +++ b/openrag/core/vector_stores/__init__.py @@ -0,0 +1,5 @@ +"""VectorStore ABC.""" + +from openrag.core.vector_stores.vector_store import VectorStore + +__all__ = ["VectorStore"] diff --git a/openrag/core/vector_stores/vector_store.py b/openrag/core/vector_stores/vector_store.py new file mode 100644 index 000000000..221d5c0c3 --- /dev/null +++ b/openrag/core/vector_stores/vector_store.py @@ -0,0 +1,63 @@ +"""Abstract vector store interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Any + +from openrag.core.models.chunk import Chunk + + +class VectorStore(ABC): + """Base class for vector database backends.""" + + @abstractmethod + async def upsert(self, chunks: list[Chunk], collection: str = "default") -> int: + """Insert or update chunks. Returns count of upserted items.""" + ... + + @abstractmethod + async def search( + self, + embedding: list[float], + top_k: int = 10, + collection: str = "default", + filters: dict[str, Any] | None = None, + ) -> list[dict[str, Any]]: + """Search by embedding vector. Returns raw results.""" + ... + + @abstractmethod + async def delete(self, ids: list[str], collection: str = "default") -> int: + """Delete chunks by ID. Returns count of deleted items.""" + ... + + @abstractmethod + async def ensure_collection(self, name: str, dimension: int, **kwargs: Any) -> None: + """Create collection if it doesn't exist.""" + ... + + @abstractmethod + async def drop_collection(self, name: str) -> None: + """Drop a collection entirely.""" + ... + + @abstractmethod + async def collection_exists(self, name: str) -> bool: + """Check if collection exists.""" + ... + + @abstractmethod + async def query_ids_by_filter(self, collection: str, filters: dict[str, Any]) -> list[str]: + """Return chunk IDs matching the given filter expression.""" + ... + + @abstractmethod + async def query_chunks_by_filter( + self, + collection: str, + filters: dict[str, Any], + output_fields: list[str] | None = None, + ) -> list[dict[str, Any]]: + """Return full chunk data matching the given filter expression.""" + ... diff --git a/openrag/core/vlm/__init__.py b/openrag/core/vlm/__init__.py index e69de29bb..b562bec80 100644 --- a/openrag/core/vlm/__init__.py +++ b/openrag/core/vlm/__init__.py @@ -0,0 +1,6 @@ +"""VLM ABC + registry.""" + +from openrag.core.vlm.registry import vlm_registry +from openrag.core.vlm.vlm import VLM + +__all__ = ["VLM", "vlm_registry"] diff --git a/openrag/core/vlm/registry.py b/openrag/core/vlm/registry.py new file mode 100644 index 000000000..91d57dce9 --- /dev/null +++ b/openrag/core/vlm/registry.py @@ -0,0 +1,6 @@ +"""VLM registry.""" + +from openrag.core.utils.registry import Registry +from openrag.core.vlm.vlm import VLM + +vlm_registry: Registry[VLM] = Registry("vlm") diff --git a/openrag/core/vlm/vlm.py b/openrag/core/vlm/vlm.py new file mode 100644 index 000000000..2da69b6d7 --- /dev/null +++ b/openrag/core/vlm/vlm.py @@ -0,0 +1,19 @@ +"""Abstract VLM (Vision-Language Model) interface.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + + +class VLM(ABC): + """Base class for all vision-language model providers.""" + + @abstractmethod + async def caption_image(self, image_bytes: bytes, prompt: str | None = None) -> str: + """Generate a caption/description for an image.""" + ... + + @abstractmethod + async def caption_images_batch(self, images: list[bytes], prompt: str | None = None) -> list[str]: + """Batch caption multiple images.""" + ...