yoga1290.rag 0.1.0__tar.gz

yoga1290_rag-0.1.0/PKG-INFO

@@ -0,0 +1,7 @@
+Metadata-Version: 2.4
+Name: yoga1290.rag
+Version: 0.1.0
+Summary: Document AI and RAG application
+Project-URL: Homepage, https://github.com/yoga1290/rag
+Project-URL: Issues, https://github.com/yoga1290/rag/issues
+Requires-Python: >=3.11

yoga1290_rag-0.1.0/pyproject.toml

@@ -0,0 +1,16 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "yoga1290.rag"
+version = "0.1.0"
+description = "Document AI and RAG application"
+requires-python = ">=3.11"
+
+[tool.setuptools.packages.find]
+where = ["yoga1290/rag"]
+
+[project.urls]
+Homepage = "https://github.com/yoga1290/rag"
+Issues = "https://github.com/yoga1290/rag/issues"

yoga1290_rag-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

yoga1290_rag-0.1.0/yoga1290/rag/api/main.py

@@ -0,0 +1,48 @@
+from fastapi import FastAPI, File, HTTPException, Query, UploadFile
+from pathlib import Path
+from uuid import uuid4
+
+app = FastAPI()
+
+UPLOAD_DIR = Path("uploads")
+UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
+
+
+@app.post("/embed")
+async def upload_file(
+    user_id: str = Query(...),
+    file: UploadFile = File(...),
+):
+    try:
+        # Preserve the original file extension
+        extension = Path(file.filename).suffix
+
+        # Generate a unique filename
+        unique_filename = f"{uuid4()}{extension}"
+
+        # Optionally organize files by user
+        user_dir = UPLOAD_DIR / user_id
+        user_dir.mkdir(parents=True, exist_ok=True)
+
+        destination = user_dir / unique_filename
+
+        with destination.open("wb") as buffer:
+            while chunk := await file.read(1024 * 1024):
+                buffer.write(chunk)
+
+        return {
+            "user_id": user_id,
+            "filename": unique_filename,
+            "original_filename": file.filename,
+            "content_type": file.content_type,
+            "path": str(destination),
+        }
+
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Failed to upload file: {e}",
+        )
+
+    finally:
+        await file.close()

yoga1290_rag-0.1.0/yoga1290/rag/application/pipelines/__init__.py

@@ -0,0 +1,7 @@
+from .ingestion_pipeline import IngestionPipeline
+from .retrieval_pipeline import RetrievalPipeline
+
+__all__ = [
+    "IngestionPipeline",
+    "RetrievalPipeline",
+]

yoga1290_rag-0.1.0/yoga1290/rag/application/pipelines/indexing_pipeline.py

@@ -0,0 +1,34 @@
+from yoga1290.rag.domain.ports.embedder import Embedder
+from yoga1290.rag.domain.ports.vector_store import VectorStore
+
+class IndexingPipeline:
+    def __init__(
+        self,
+        embedder: Embedder,
+        vector_store: VectorStore,
+    ) -> None:
+        self.embedder = embedder
+        self.vector_store = vector_store
+
+    def index(self, search_document):
+        chunks = list(search_document.chunks)
+        texts = [chunk.text for chunk in chunks]
+        embeddings = self.embedder.embed_batch(texts)
+
+        indexed_chunks = [
+            chunk.__class__(
+                id=chunk.id,
+                document_id=chunk.document_id,
+                text=chunk.text,
+                index=chunk.index,
+                embedding=embedding,
+                metadata=chunk.metadata,
+            )
+            for chunk, embedding in zip(
+                chunks,
+                embeddings,
+                strict=True,
+            )
+        ]
+
+        self.vector_store.upsert(indexed_chunks)

yoga1290_rag-0.1.0/yoga1290/rag/application/pipelines/ingestion_pipeline.py

@@ -0,0 +1,54 @@
+from yoga1290.rag.domain.models import (
+                            Document,
+                            ParsedDocument,
+                            SearchDocument)
+from yoga1290.rag.domain.ports import (
+                            DocumentParser,
+                            DocumentExtractor,
+                            DocumentClassifier,
+                            SearchPreparer,
+                            VectorStore,)
+from collections.abc import Sequence
+
+class IngestionPipeline:
+    def __init__(
+        self,
+        parser: DocumentParser,
+        extractor: DocumentExtractor,
+        classifier: DocumentClassifier,
+        search_preparer: SearchPreparer,
+        vectorstore: VectorStore,
+    ) -> None:
+        self.parser = parser
+        self.extractor = extractor
+        self.classifier = classifier
+        self.search_preparer = search_preparer
+        self.vectorstore = vectorstore
+    
+    def process(self,
+                documents: Sequence[Document],
+                # extraction_schema,
+                # classification_schema,
+        ):
+        search_documents = []
+        for document in documents:
+            parsed_document: ParsedDocument = (
+                                self.parser.parse(document))
+            search_document: SearchDocument = (
+                                self.search_preparer.prepare(parsed_document))
+            search_documents.append(search_document)
+        self.vectorstore.add_documents(search_documents)
+
+        return {
+            # "parsed": parsed,
+            # "extracted": extracted,
+            # "classification": classification,
+            "search_document": search_documents,
+        }
+
+
+        # extracted = self.extractor.extract(parsed, extraction_schema)
+        # classification = self.classifier.classify(
+        #     parsed,
+        #     classification_schema,
+        # )

yoga1290_rag-0.1.0/yoga1290/rag/application/pipelines/retrieval_pipeline.py

@@ -0,0 +1,84 @@
+from yoga1290.rag.domain.models.rag_result import (
+    RagResult,
+    Source,
+)
+from yoga1290.rag.domain.ports import (
+                                    LLM,
+                                    Retriever, )
+
+
+class RetrievalPipeline:
+
+    def __init__(
+        self,
+        llm: LLM,
+        retriever: Retriever,
+        top_k: int = 2,
+    ) -> None:
+        self._llm = llm
+        self._retriever = retriever
+        self._top_k = top_k
+
+    def run(self, question: str) -> RagResult:
+
+        import time
+
+        start = time.perf_counter()
+
+        chunks = self._retriever.retrieve(
+            question,
+            top_k=self._top_k,
+        )
+
+        context = self._build_context(chunks)
+
+        prompt = self._build_prompt(
+            question=question,
+            context=context,
+        )
+
+        answer = self._llm.generate(prompt)
+
+        duration = time.perf_counter() - start
+
+        sources = tuple(
+            Source(
+                document_id=chunk.document_id,
+                chunk_id=chunk.id,
+                metadata=chunk.metadata,
+            )
+            for chunk in chunks
+        )
+
+        return RagResult(
+            question=question,
+            answer=answer,
+            sources=sources,
+            duration_seconds=duration,
+        )
+
+    @staticmethod
+    def _build_context(chunks) -> str:
+
+        return "\n\n".join(
+            chunk.text
+            for chunk in chunks
+        )
+
+    @staticmethod
+    def _build_prompt(
+        question: str,
+        context: str,
+    ) -> str:
+
+        return f"""
+Use the following context to answer the question.
+
+Context:
+{context}
+
+Question:
+{question}
+
+Answer:
+""".strip()

yoga1290_rag-0.1.0/yoga1290/rag/application/services/document_service.py

@@ -0,0 +1,15 @@
+from yoga1290.rag.domain.ports.document_parser import DocumentParser
+from yoga1290.rag.domain.ports.document_source import DocumentSource
+
+class DocumentService:
+    def __init__(
+        self,
+        source: DocumentSource,
+        parser: DocumentParser,
+    ) -> None:
+        self.source = source
+        self.parser = parser
+
+    def process(self, document_id: str):
+        document = self.source.get(document_id)
+        return self.parser.parse(document)

yoga1290_rag-0.1.0/yoga1290/rag/application/services/search_service.py

@@ -0,0 +1,25 @@
+from yoga1290.rag.domain.ports.embedder import Embedder
+from yoga1290.rag.domain.ports.vector_store import VectorStore
+
+class SearchService:
+    def __init__(
+        self,
+        embedder: Embedder,
+        vector_store: VectorStore,
+    ) -> None:
+        self.embedder = embedder
+        self.vector_store = vector_store
+
+    def search(
+        self,
+        query: str,
+        *,
+        top_k: int = 10,
+        filters: dict[str, object] | None = None,
+    ):
+        embedding = self.embedder.embed(query)
+        return self.vector_store.search(
+            embedding,
+            top_k=top_k,
+            filters=filters,
+        )

yoga1290_rag-0.1.0/yoga1290/rag/config/settings.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+
+@dataclass(frozen=True)
+class Settings:
+    app_env: str = "local"
+    ai_provider: str = "local"
+    parser_provider: str = "docling"
+    extractor_provider: str = "qwen"
+    classifier_provider: str = "qwen"
+    search_preparer_provider: str = "local"
+    embedding_provider: str = "qwen"
+    vector_store_provider: str = "pgvector"

yoga1290_rag-0.1.0/yoga1290/rag/domain/models/__init__.py

@@ -0,0 +1,20 @@
+from .document import Document
+from .parsed_document import ParsedDocument, ParsedPage, ParsedTable
+from .extracted_data import ExtractedData
+from .classification import Classification, ClassificationLabel
+from .search_document import SearchDocument
+from .chunk import Chunk
+from .search_result import SearchResult
+
+__all__ = [
+    "Document",
+    "ParsedDocument",
+    "ParsedPage",
+    "ParsedTable",
+    "ExtractedData",
+    "Classification",
+    "ClassificationLabel",
+    "SearchDocument",
+    "Chunk",
+    "SearchResult",
+]

yoga1290_rag-0.1.0/yoga1290/rag/domain/models/chunk.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Mapping, Sequence
+
+@dataclass(frozen=True, slots=True)
+class Chunk:
+    id: str
+    document_id: str
+    text: str
+    index: int
+    embedding: Sequence[float] | None = None
+    metadata: Mapping[str, object] = field(default_factory=dict)

yoga1290_rag-0.1.0/yoga1290/rag/domain/models/classification.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Mapping, Sequence
+
+@dataclass(frozen=True, slots=True)
+class ClassificationLabel:
+    name: str
+    confidence: float | None = None
+    metadata: Mapping[str, object] = field(default_factory=dict)
+
+@dataclass(frozen=True, slots=True)
+class Classification:
+    document_id: str
+    labels: Sequence[ClassificationLabel]
+    provider: str | None = None
+
+    @property
+    def primary_label(self) -> ClassificationLabel | None:
+        return self.labels[0] if self.labels else None

yoga1290_rag-0.1.0/yoga1290/rag/domain/models/document.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import Mapping
+
+@dataclass(frozen=True, slots=True)
+class Document:
+    id: str | None = None
+    size_bytes: int | None = None
+    uri: str | None = None
+    content: bytes | None = None
+    filename: str = None
+    content_type: str = None
+    metadata: Mapping[str, object] = field(default_factory=dict)
+    created_at: datetime | None = None
+
+    @property
+    def suffix(self) -> str:
+        return Path(self.filename).suffix.lower()

yoga1290_rag-0.1.0/yoga1290/rag/domain/models/extracted_data.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Generic, TypeVar
+
+T = TypeVar("T")
+
+@dataclass(frozen=True, slots=True)
+class ExtractedData(Generic[T]):
+    document_id: str
+    data: T
+    confidence: float | None = None
+    provider: str | None = None
+    extracted_at: datetime | None = None

yoga1290_rag-0.1.0/yoga1290/rag/domain/models/parsed_document.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Mapping, Sequence
+
+@dataclass(frozen=True, slots=True)
+class ParsedPage:
+    page_number: int
+    text: str
+    metadata: Mapping[str, object] = field(default_factory=dict)
+
+@dataclass(frozen=True, slots=True)
+class ParsedTable:
+    table_id: str
+    page_number: int | None
+    headers: Sequence[str] = field(default_factory=tuple)
+    rows: Sequence[Sequence[object]] = field(default_factory=tuple)
+    metadata: Mapping[str, object] = field(default_factory=dict)
+
+@dataclass(frozen=True, slots=True)
+class ParsedDocument:
+    document_id: str
+    text: str
+    pages: Sequence[ParsedPage] = field(default_factory=tuple)
+    tables: Sequence[ParsedTable] = field(default_factory=tuple)
+    metadata: Mapping[str, object] = field(default_factory=dict)
+
+    @property
+    def page_count(self) -> int:
+        return len(self.pages)

yoga1290_rag-0.1.0/yoga1290/rag/domain/models/rag_result.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True, slots=True)
+class Source:
+    document_id: str
+    chunk_id: str
+    metadata: dict[str, Any]
+
+
+@dataclass(frozen=True, slots=True)
+class RagResult:
+    question: str
+    answer: str
+    sources: tuple[Source, ...]
+    duration_seconds: float

yoga1290_rag-0.1.0/yoga1290/rag/domain/models/search_document.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Mapping, Sequence
+
+from .chunk import Chunk
+
+@dataclass(frozen=True, slots=True)
+class SearchDocument:
+    document_id: str
+    chunks: Sequence[Chunk] = field(default_factory=tuple)
+    metadata: Mapping[str, object] = field(default_factory=dict)
+
+    @property
+    def chunk_count(self) -> int:
+        return len(self.chunks)

yoga1290_rag-0.1.0/yoga1290/rag/domain/models/search_result.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Mapping
+
+@dataclass(frozen=True, slots=True)
+class SearchResult:
+    chunk_id: str
+    document_id: str
+    text: str
+    score: float
+    metadata: Mapping[str, object] = field(default_factory=dict)

yoga1290_rag-0.1.0/yoga1290/rag/domain/ports/__init__.py

@@ -0,0 +1,21 @@
+from .document_classifier import DocumentClassifier
+from .document_extractor import DocumentExtractor
+from .document_parser import DocumentParser
+from .document_source import DocumentSource
+from .embedder import Embedder
+from .search_preparer import SearchPreparer
+from .vector_store import VectorStore
+from .llm import LLM
+from .retriever import Retriever
+
+__all__ = [
+    "DocumentSource",
+    "DocumentParser",
+    "DocumentExtractor",
+    "DocumentClassifier",
+    "SearchPreparer",
+    "Embedder",
+    "LLM",
+    "VectorStore",
+    "Retriever",
+]

yoga1290_rag-0.1.0/yoga1290/rag/domain/ports/document_classifier.py

@@ -0,0 +1,31 @@
+from yoga1290.rag.domain.models.classification import Classification
+from yoga1290.rag.domain.models.parsed_document import ParsedDocument
+from yoga1290.rag.domain.schemas.classification_schema import ClassificationSchema
+
+
+class DocumentClassifier():
+    """
+    Classifies a document according to a domain-defined classification schema.
+
+    Implementations may use:
+      - Databricks ai_classify
+      - Qwen
+      - BERT
+      - Other classifiers
+    """
+    def classify(
+        self,
+        document: ParsedDocument,
+        schema: ClassificationSchema,
+    ) -> Classification:
+        """
+        Classify a document.
+
+        Args:
+            document: Parsed document.
+            schema: Allowed categories and classification configuration.
+
+        Returns:
+            Classification result.
+        """
+        raise NotImplementedError

yoga1290_rag-0.1.0/yoga1290/rag/domain/ports/document_extractor.py

@@ -0,0 +1,17 @@
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+from yoga1290.rag.domain.models.extracted_data import ExtractedData
+from yoga1290.rag.domain.models.parsed_document import ParsedDocument
+from yoga1290.rag.domain.schemas.extraction_schema import ExtractionSchema
+
+T = TypeVar("T")
+
+class DocumentExtractor(ABC, Generic[T]):
+    @abstractmethod
+    def extract(
+        self,
+        document: ParsedDocument,
+        schema: ExtractionSchema[T],
+    ) -> ExtractedData[T]:
+        raise NotImplementedError

yoga1290_rag-0.1.0/yoga1290/rag/domain/ports/document_parser.py

@@ -0,0 +1,30 @@
+from yoga1290.rag.domain.models.document import Document
+from yoga1290.rag.domain.models.parsed_document import ParsedDocument
+
+class DocumentParser():
+    """
+    Converts a raw document into a normalized ParsedDocument.
+
+    Implementations may use:
+      - Databricks ai_parse_document
+      - Docling
+      - PyMuPDF
+      - OCR
+      - Other document parsing engines
+    """
+    def parse(self, document: Document) -> ParsedDocument:
+        """
+        Parse a document into a normalized representation.
+
+        Args:
+            document: Raw document metadata and content reference.
+
+        Returns:
+            ParsedDocument containing normalized text, pages,
+            tables, images, metadata, etc.
+
+        Raises:
+            DocumentParseError: If the document cannot be parsed.
+        """
+        raise NotImplementedError
+

yoga1290_rag-0.1.0/yoga1290/rag/domain/ports/document_source.py

@@ -0,0 +1,29 @@
+from collections.abc import Iterator
+from abc import ABC, abstractmethod
+from yoga1290.rag.domain.models.document import Document
+
+
+class DocumentSource:
+    """
+    Port for retrieving documents from an external source.
+
+    Implementations may read from:
+      - Local filesystem
+      - S3
+      - Azure Blob Storage
+      - Google Cloud Storage
+      - Databricks Volumes
+      - Database
+      - Message queue
+      - HTTP/API
+    """
+
+    @abstractmethod
+    def __iter__(self) -> Iterator[Document]:
+        """
+        Iterate over documents available from the source.
+
+        Returns:
+            Iterator of Document objects.
+        """
+        raise NotImplementedError