Add document chunk embeddings to workflow

backend/alembic/versions/585aa57edb27_add_chunk_embeddings.py

@@ -0,0 +1,30 @@
+"""add chunk embeddings
+
+Revision ID: 585aa57edb27
+Revises: 5b7572e0875a
+Create Date: 2026-08-10 22:06:30.806027
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+from pgvector.sqlalchemy import Vector
+
+revision: str = '585aa57edb27'
+down_revision: Union[str, None] = '5b7572e0875a'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column('document_chunks', sa.Column('embedding', Vector(384), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column('document_chunks', 'embedding')
+    # ### end Alembic commands ###

backend/app/models/document_chunk.py

@@ -8,7 +8,7 @@ from sqlalchemy.dialects.postgresql import UUID
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from app.database.database import Base
-
+from pgvector.sqlalchemy import Vector
 
 class DocumentChunk(Base):
     """
@@ -56,6 +56,11 @@ class DocumentChunk(Base):
         nullable=True,
     )
 
+    embedding: Mapped[list[float] | None] = mapped_column(
+        Vector(384),
+        nullable=True,
+    )
+
     created_at: Mapped[datetime] = mapped_column(
         DateTime(timezone=True),
         server_default=func.now(),

backend/app/services/embedding_service.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from sentence_transformers import SentenceTransformer
+
+
+class EmbeddingService:
+    """
+    Generates dense vector embeddings for document chunks.
+
+    Model output dimension:
+        384
+    """
+
+    MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+    DIMENSION = 384
+
+    def __init__(self):
+        self.model = SentenceTransformer(self.MODEL_NAME)
+
+    def embed(self, text: str) -> list[float]:
+        if not text or not text.strip():
+            return []
+
+        vector = self.model.encode(
+            text,
+            normalize_embeddings=True,
+        )
+
+        return vector.tolist()
+
+    def embed_many(self, texts: list[str]) -> list[list[float]]:
+        if not texts:
+            return []
+
+        vectors = self.model.encode(
+            texts,
+            normalize_embeddings=True,
+        )
+
+        return vectors.tolist()

backend/app/workflow/graph.py

@@ -16,7 +16,8 @@ from app.workflow.nodes.knowledge import knowledge
 from app.workflow.nodes.link import link
 from app.workflow.nodes.reconcile import reconcile
 from app.workflow.state import WorkflowState
-
+from app.services.embedding_service import EmbeddingService
+from app.workflow.nodes.embed import embed
 
 def build_workflow(db: Session):
     """
@@ -27,7 +28,7 @@ def build_workflow(db: Session):
     knowledge_repository = KnowledgeRepository(db)
     chunking_service = ChunkingService(db)
     knowledge_link_service = KnowledgeLinkService(db)
-
+    embedding_service = EmbeddingService()
     graph = StateGraph(WorkflowState)
 
     # ------------------------------------------------------------------
@@ -73,13 +74,22 @@ def build_workflow(db: Session):
 
     graph.add_node("complete", complete)
 
+    graph.add_node(
+        "embedding",
+        lambda state: embed(
+            state,
+            db,
+            embedding_service,
+        ),
+    )
     # ------------------------------------------------------------------
     # Workflow edges
     # ------------------------------------------------------------------
 
     graph.add_edge(START, "extract")
     graph.add_edge("extract", "chunk")
-    graph.add_edge("chunk", "classify")
+    graph.add_edge("chunk", "embedding")
+    graph.add_edge("embedding", "classify")
     graph.add_edge("classify", "knowledge")
     graph.add_edge("knowledge", "reconciliation")
     graph.add_edge("reconciliation", "link")

backend/app/workflow/nodes/embed.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from sqlalchemy.orm import Session
+
+from app.models.document_chunk import DocumentChunk
+from app.services.embedding_service import EmbeddingService
+from app.workflow.state import WorkflowState
+
+
+def embed(
+    state: WorkflowState,
+    db: Session,
+    embedding_service: EmbeddingService,
+) -> WorkflowState:
+    """
+    Generate and persist embeddings for the current document's chunks.
+    """
+
+    state.current_node = "EMBEDDING"
+
+    chunks = (
+        db.query(DocumentChunk)
+        .filter(
+            DocumentChunk.document_version_id
+            == state.document_version_id,
+            DocumentChunk.embedding.is_(None),
+        )
+        .order_by(DocumentChunk.chunk_index)
+        .all()
+    )
+
+    if not chunks:
+        state.metadata["embeddings_created"] = 0
+        return state
+
+    texts = [chunk.text for chunk in chunks]
+
+    vectors = embedding_service.embed_many(texts)
+
+    if len(vectors) != len(chunks):
+        raise RuntimeError(
+            "Embedding count does not match chunk count."
+        )
+
+    for chunk, vector in zip(chunks, vectors):
+        chunk.embedding = vector
+
+    db.commit()
+
+    state.metadata["embeddings_created"] = len(chunks)
+
+    return state

backend/requirements.txt

@@ -49,4 +49,7 @@ email-validator==2.2.0
 # Agentic workflow
 langgraph
 langchain
-langchain-groq
+langchain-groq
+
+
+grandalf==0.8