riezqidr commited on
Commit
0b5f4aa
·
1 Parent(s): 6eabf86

feat(search): add Phase 2 semantic search with pgvector and hybrid retrieval

Browse files

Implements parent-child chunking, pgvector embeddings, and hybrid (dense + lexical)
retrieval with Reciprocal Rank Fusion. Adds TEI embedding service adapter, NoOp
reranker stub, and /search endpoint.

Known open issues in this commit (to be addressed in follow-up):
- Phase 2 search disconnected from ingestion flow (no code path creates chunks)
- Raw SQL f-string interpolation in repositories/search.py (injection risk)
- Zero test coverage for Phase 2 (chunking, embedding, reranker, search)
- content_tsv ORM/migration drift (model declares Text, migration creates tsvector)

Components added:
- services/chunking.py: parent (~700w) and child (~180w) chunk generation
- services/embedding.py: EmbeddingService protocol + TEI adapter
- services/reranker.py: RerankerService protocol + NoOp adapter
- services/search.py: RRF fusion, dense/lexical search orchestration
- repositories/search.py: pgvector cosine + tsvector ts_rank_cd queries
- routers/search.py: POST /search endpoint
- schemas/search.py: SearchRequest/Response Pydantic models
- migrations/versions/20260731_0930: resume_chunks table, HNSW + GIN indexes

.env.example CHANGED
@@ -1,4 +1,4 @@
1
- # TalentLens — environment template (Phase 0-1)
2
  # Copy to `.env` and fill in. Never commit a populated `.env`.
3
 
4
  # --- Runtime -----------------------------------------------------------------
@@ -27,3 +27,21 @@ MAX_UPLOAD_BYTES=10485760 # 10 MiB
27
 
28
  # --- CORS (Vite dev server by default; add the Vercel origin in production) --
29
  CORS_ALLOW_ORIGINS=["http://localhost:5173"]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # TalentLens — environment template (Phase 0-2)
2
  # Copy to `.env` and fill in. Never commit a populated `.env`.
3
 
4
  # --- Runtime -----------------------------------------------------------------
 
27
 
28
  # --- CORS (Vite dev server by default; add the Vercel origin in production) --
29
  CORS_ALLOW_ORIGINS=["http://localhost:5173"]
30
+
31
+ # --- Embedding (Phase 2) ----------------------------------------------------
32
+ # HuggingFace Text Embeddings Inference endpoint for bge-m3
33
+ EMBEDDING_ENDPOINT=http://localhost:8080
34
+ EMBEDDING_MODEL=BAAI/bge-m3
35
+ EMBEDDING_DIM=1024
36
+ EMBEDDING_BATCH_SIZE=32
37
+
38
+ # --- Search (Phase 2) -------------------------------------------------------
39
+ SEARCH_DENSE_WEIGHT=0.6
40
+ SEARCH_LEXICAL_WEIGHT=0.4
41
+ SEARCH_TOP_K_RECALL=20
42
+ SEARCH_RERANK_TOP_K=10
43
+
44
+ # --- Reranker (Phase 2, optional) --------------------------------------------
45
+ # Leave empty to disable reranking (CPU fallback / dev mode)
46
+ RERANKER_ENDPOINT=
47
+ RERANKER_MODEL=BAAI/bge-reranker-v2-m3
pyproject.toml CHANGED
@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
5
  [project]
6
  name = "talentlens"
7
  version = "0.1.0"
8
- description = "TalentLens — AI CV Screener (Phase 0-1: foundations + ingestion)"
9
  requires-python = ">=3.11,<3.12"
10
  dependencies = [
11
  "fastapi>=0.141,<0.142",
@@ -23,6 +23,7 @@ dependencies = [
23
  "pdfplumber>=0.11,<0.12",
24
  "python-docx>=1.1,<2",
25
  "defusedxml>=0.7,<0.8",
 
26
  ]
27
 
28
  [project.optional-dependencies]
 
5
  [project]
6
  name = "talentlens"
7
  version = "0.1.0"
8
+ description = "TalentLens — AI CV Screener (Phase 0-2: foundations + ingestion + search)"
9
  requires-python = ">=3.11,<3.12"
10
  dependencies = [
11
  "fastapi>=0.141,<0.142",
 
23
  "pdfplumber>=0.11,<0.12",
24
  "python-docx>=1.1,<2",
25
  "defusedxml>=0.7,<0.8",
26
+ "pgvector>=0.3,<0.4",
27
  ]
28
 
29
  [project.optional-dependencies]
serving/app/config.py CHANGED
@@ -62,6 +62,22 @@ class Settings(BaseSettings):
62
  max_upload_bytes: int = 10 * 1024 * 1024
63
  allowed_upload_mime_types: tuple[str, ...] = ("application/pdf", DOCX_MIME)
64
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
65
  cors_allow_origins: tuple[str, ...] = ("http://localhost:5173",)
66
 
67
  log_level: str = "INFO"
 
62
  max_upload_bytes: int = 10 * 1024 * 1024
63
  allowed_upload_mime_types: tuple[str, ...] = ("application/pdf", DOCX_MIME)
64
 
65
+ # --- Embedding (Phase 2) -------------------------------------------------
66
+ embedding_endpoint: str = "http://localhost:8080"
67
+ embedding_model: str = "BAAI/bge-m3"
68
+ embedding_dim: int = 1024
69
+ embedding_batch_size: int = 32
70
+
71
+ # --- Search (Phase 2) ----------------------------------------------------
72
+ search_dense_weight: float = 0.6
73
+ search_lexical_weight: float = 0.4
74
+ search_top_k_recall: int = 20
75
+ search_rerank_top_k: int = 10
76
+
77
+ # --- Reranker (Phase 2) --------------------------------------------------
78
+ reranker_endpoint: str = ""
79
+ reranker_model: str = "BAAI/bge-reranker-v2-m3"
80
+
81
  cors_allow_origins: tuple[str, ...] = ("http://localhost:5173",)
82
 
83
  log_level: str = "INFO"
serving/app/exceptions.py CHANGED
@@ -128,3 +128,25 @@ class StorageError(TalentLensError):
128
  error_code: ClassVar[str] = "STORAGE_UNAVAILABLE"
129
  status_code: ClassVar[int] = 503
130
  default_message: ClassVar[str] = "Document storage is temporarily unavailable."
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
128
  error_code: ClassVar[str] = "STORAGE_UNAVAILABLE"
129
  status_code: ClassVar[int] = 503
130
  default_message: ClassVar[str] = "Document storage is temporarily unavailable."
131
+
132
+
133
+ # --------------------------------------------------------------------------- #
134
+ # Search / Embeddings (Phase 2) #
135
+ # --------------------------------------------------------------------------- #
136
+
137
+
138
+ class EmbeddingServiceUnavailableError(TalentLensError):
139
+ """Raised when the embedding service cannot be reached."""
140
+
141
+ error_code: ClassVar[str] = "EMBEDDING_SERVICE_UNAVAILABLE"
142
+ status_code: ClassVar[int] = 503
143
+ default_message: ClassVar[str] = "The embedding service is temporarily unavailable."
144
+
145
+
146
+ class SearchError(TalentLensError):
147
+ """Raised when a search operation fails unexpectedly."""
148
+
149
+ error_code: ClassVar[str] = "SEARCH_FAILED"
150
+ status_code: ClassVar[int] = 500
151
+ default_message: ClassVar[str] = "The search operation failed."
152
+
serving/app/main.py CHANGED
@@ -21,7 +21,7 @@
21
  from app.config import Settings, get_settings
22
  from app.exceptions import TalentLensError
23
  from app.logging import configure_logging, get_logger
24
- from app.routers import auth, jobs, resumes
25
 
26
  logger = get_logger(__name__)
27
 
@@ -47,7 +47,7 @@
47
  _RATE_LIMIT_MAX_TRACKED_KEYS = 10_000 # hard ceiling on limiter memory
48
 
49
  # Paths that are rate-limited (upload endpoints are the highest risk).
50
- _RATE_LIMITED_PREFIXES = ("/api/v1/resumes", "/api/v1/jobs")
51
 
52
  _request_counts: dict[str, list[float]] = defaultdict(list)
53
 
@@ -341,4 +341,5 @@ async def health() -> dict[str, str]:
341
  app.include_router(auth.router, prefix=cfg.api_v1_prefix)
342
  app.include_router(resumes.router, prefix=cfg.api_v1_prefix)
343
  app.include_router(jobs.router, prefix=cfg.api_v1_prefix)
 
344
  return app
 
21
  from app.config import Settings, get_settings
22
  from app.exceptions import TalentLensError
23
  from app.logging import configure_logging, get_logger
24
+ from app.routers import auth, jobs, resumes, search
25
 
26
  logger = get_logger(__name__)
27
 
 
47
  _RATE_LIMIT_MAX_TRACKED_KEYS = 10_000 # hard ceiling on limiter memory
48
 
49
  # Paths that are rate-limited (upload endpoints are the highest risk).
50
+ _RATE_LIMITED_PREFIXES = ("/api/v1/resumes", "/api/v1/jobs", "/api/v1/search")
51
 
52
  _request_counts: dict[str, list[float]] = defaultdict(list)
53
 
 
341
  app.include_router(auth.router, prefix=cfg.api_v1_prefix)
342
  app.include_router(resumes.router, prefix=cfg.api_v1_prefix)
343
  app.include_router(jobs.router, prefix=cfg.api_v1_prefix)
344
+ app.include_router(search.router, prefix=cfg.api_v1_prefix)
345
  return app
serving/app/migrations/versions/20260731_0930_add_resume_chunks_with_pgvector.py ADDED
@@ -0,0 +1,119 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """add resume_chunks table with pgvector
2
+
3
+ Revision ID: a1b2c3d4e5f6
4
+ Revises: dd2329f32308
5
+ Create Date: 2026-07-31 09:30:00.000000+00:00
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ from collections.abc import Sequence
11
+
12
+ import sqlalchemy as sa
13
+ from alembic import op
14
+
15
+ revision: str = "a1b2c3d4e5f6"
16
+ down_revision: str | None = "dd2329f32308"
17
+ branch_labels: str | Sequence[str] | None = None
18
+ depends_on: str | Sequence[str] | None = None
19
+
20
+
21
+ def upgrade() -> None:
22
+ """Apply this revision."""
23
+ # Enable the pgvector extension — idempotent, safe to re-run
24
+ op.execute("CREATE EXTENSION IF NOT EXISTS vector")
25
+
26
+ op.create_table(
27
+ "resume_chunks",
28
+ sa.Column("id", sa.UUID(), nullable=False),
29
+ sa.Column("tenant_id", sa.UUID(), nullable=False),
30
+ sa.Column("resume_version_id", sa.UUID(), nullable=False),
31
+ sa.Column("document_id", sa.UUID(), nullable=False),
32
+ sa.Column("chunk_index", sa.Integer(), nullable=False),
33
+ sa.Column("parent_chunk_id", sa.UUID(), nullable=True),
34
+ sa.Column("section", sa.String(length=64), nullable=False),
35
+ sa.Column("content", sa.Text(), nullable=False),
36
+ sa.Column("content_tsv", sa.Text(), nullable=True),
37
+ sa.Column("page_from", sa.Integer(), nullable=False),
38
+ sa.Column("page_to", sa.Integer(), nullable=False),
39
+ sa.Column("start_char", sa.Integer(), nullable=False),
40
+ sa.Column("end_char", sa.Integer(), nullable=False),
41
+ sa.Column("token_count", sa.Integer(), nullable=False),
42
+ sa.Column("embedding_model", sa.String(length=128), nullable=False),
43
+ sa.Column("embedding_version", sa.String(length=64), nullable=False),
44
+ sa.Column(
45
+ "created_at",
46
+ sa.DateTime(timezone=True),
47
+ server_default=sa.text("now()"),
48
+ nullable=False,
49
+ ),
50
+ sa.Column(
51
+ "updated_at",
52
+ sa.DateTime(timezone=True),
53
+ server_default=sa.text("now()"),
54
+ nullable=False,
55
+ ),
56
+ sa.ForeignKeyConstraint(
57
+ ["resume_version_id"], ["resume_versions.id"], ondelete="CASCADE"
58
+ ),
59
+ sa.ForeignKeyConstraint(
60
+ ["document_id"], ["resume_documents.id"], ondelete="CASCADE"
61
+ ),
62
+ sa.PrimaryKeyConstraint("id"),
63
+ )
64
+
65
+ # Add the pgvector embedding column — raw SQL because Alembic/SA don't
66
+ # know about the vector type natively
67
+ op.execute("ALTER TABLE resume_chunks ADD COLUMN embedding vector(1024)")
68
+
69
+ # Replace the content_tsv text column with a real tsvector column
70
+ op.execute("ALTER TABLE resume_chunks DROP COLUMN IF EXISTS content_tsv")
71
+ op.execute(
72
+ "ALTER TABLE resume_chunks ADD COLUMN content_tsv tsvector "
73
+ "GENERATED ALWAYS AS (to_tsvector('english', content)) STORED"
74
+ )
75
+
76
+ # Add the is_parent column (needed for filtering child-only retrieval)
77
+ op.execute(
78
+ "ALTER TABLE resume_chunks ADD COLUMN is_parent boolean NOT NULL DEFAULT false"
79
+ )
80
+
81
+ # --- Indexes ---
82
+ # HNSW index for cosine vector search (pgvector-specific)
83
+ op.execute(
84
+ "CREATE INDEX ix_resume_chunks_embedding_hnsw ON resume_chunks "
85
+ "USING hnsw (embedding vector_cosine_ops) "
86
+ "WITH (m = 16, ef_construction = 64)"
87
+ )
88
+
89
+ # GIN index for full-text search
90
+ op.execute(
91
+ "CREATE INDEX ix_resume_chunks_content_tsv ON resume_chunks "
92
+ "USING gin (content_tsv)"
93
+ )
94
+
95
+ # Composite index for tenant + document scoped queries
96
+ op.create_index(
97
+ "ix_resume_chunks_tenant_document",
98
+ "resume_chunks",
99
+ ["tenant_id", "document_id"],
100
+ unique=False,
101
+ )
102
+
103
+ # Tenant isolation index
104
+ op.create_index(
105
+ "ix_resume_chunks_tenant_id",
106
+ "resume_chunks",
107
+ ["tenant_id"],
108
+ unique=False,
109
+ )
110
+
111
+
112
+ def downgrade() -> None:
113
+ """Revert this revision."""
114
+ op.drop_index("ix_resume_chunks_tenant_id", table_name="resume_chunks")
115
+ op.drop_index("ix_resume_chunks_tenant_document", table_name="resume_chunks")
116
+ op.execute("DROP INDEX IF EXISTS ix_resume_chunks_content_tsv")
117
+ op.execute("DROP INDEX IF EXISTS ix_resume_chunks_embedding_hnsw")
118
+ op.drop_table("resume_chunks")
119
+ # Note: we do NOT drop the vector extension — other tables may use it
serving/app/repositories/search.py ADDED
@@ -0,0 +1,287 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Data access for resume chunks and vector search.
2
+
3
+ This is the only layer permitted to query the ``resume_chunks`` table.
4
+ Hybrid retrieval (dense + lexical) is implemented here because the query
5
+ construction is tightly coupled to the pgvector and tsvector index structure.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import uuid
11
+ from collections.abc import Sequence
12
+ from dataclasses import dataclass
13
+
14
+ from sqlalchemy import delete, func, select, text
15
+ from sqlalchemy.ext.asyncio import AsyncSession
16
+
17
+ from app.models import ResumeChunk
18
+
19
+
20
+ @dataclass(frozen=True, slots=True)
21
+ class ChunkWithScore:
22
+ """A chunk paired with a retrieval score.
23
+
24
+ Attributes:
25
+ chunk: The retrieved chunk.
26
+ score: Retrieval score (cosine similarity, ts_rank, or fused).
27
+ """
28
+
29
+ chunk: ResumeChunk
30
+ score: float
31
+
32
+
33
+ class ChunkRepository:
34
+ """Persistence and retrieval for resume chunks with embeddings."""
35
+
36
+ def __init__(self, session: AsyncSession) -> None:
37
+ """Initialize the repository.
38
+
39
+ Args:
40
+ session: An active async session.
41
+ """
42
+ self._session = session
43
+
44
+ async def upsert_chunks(self, chunks: list[ResumeChunk]) -> None:
45
+ """Persist a batch of chunks.
46
+
47
+ Args:
48
+ chunks: Chunks to insert.
49
+ """
50
+ for chunk in chunks:
51
+ self._session.add(chunk)
52
+ await self._session.flush()
53
+
54
+ async def delete_by_document(
55
+ self, tenant_id: uuid.UUID, document_id: uuid.UUID
56
+ ) -> int:
57
+ """Delete all chunks for a document.
58
+
59
+ Used when re-embedding or re-parsing a document.
60
+
61
+ Args:
62
+ tenant_id: Owning tenant.
63
+ document_id: Document to clear.
64
+
65
+ Returns:
66
+ Number of rows deleted.
67
+ """
68
+ stmt = delete(ResumeChunk).where(
69
+ ResumeChunk.tenant_id == tenant_id,
70
+ ResumeChunk.document_id == document_id,
71
+ )
72
+ result = await self._session.execute(stmt)
73
+ return result.rowcount # type: ignore[return-value]
74
+
75
+ async def search_dense(
76
+ self,
77
+ tenant_id: uuid.UUID,
78
+ embedding: list[float],
79
+ top_k: int = 20,
80
+ *,
81
+ document_id: uuid.UUID | None = None,
82
+ ) -> list[ChunkWithScore]:
83
+ """Dense vector search using pgvector cosine distance.
84
+
85
+ Args:
86
+ tenant_id: Owning tenant (isolation filter).
87
+ embedding: Query embedding vector.
88
+ top_k: Maximum results.
89
+ document_id: Optional filter to scope results to one document.
90
+
91
+ Returns:
92
+ Chunks ordered by descending cosine similarity.
93
+ """
94
+ # Use raw SQL for pgvector operator support
95
+ vec_str = "[" + ",".join(str(v) for v in embedding) + "]"
96
+
97
+ where_clauses = ["tenant_id = :tenant_id", "embedding IS NOT NULL", "is_parent = false"]
98
+ params: dict[str, object] = {"tenant_id": str(tenant_id), "top_k": top_k}
99
+
100
+ if document_id is not None:
101
+ where_clauses.append("document_id = :document_id")
102
+ params["document_id"] = str(document_id)
103
+
104
+ where_sql = " AND ".join(where_clauses)
105
+
106
+ query = text(f"""
107
+ SELECT id, tenant_id, resume_version_id, document_id,
108
+ chunk_index, parent_chunk_id, section, content,
109
+ page_from, page_to, start_char, end_char,
110
+ token_count, embedding_model, embedding_version,
111
+ 1 - (embedding <=> :embedding::vector) AS score
112
+ FROM resume_chunks
113
+ WHERE {where_sql}
114
+ ORDER BY embedding <=> :embedding::vector
115
+ LIMIT :top_k
116
+ """)
117
+ params["embedding"] = vec_str
118
+
119
+ result = await self._session.execute(query, params)
120
+ rows = result.fetchall()
121
+
122
+ chunks: list[ChunkWithScore] = []
123
+ for row in rows:
124
+ chunk = ResumeChunk(
125
+ id=uuid.UUID(str(row.id)),
126
+ tenant_id=uuid.UUID(str(row.tenant_id)),
127
+ resume_version_id=uuid.UUID(str(row.resume_version_id)),
128
+ document_id=uuid.UUID(str(row.document_id)),
129
+ chunk_index=row.chunk_index,
130
+ parent_chunk_id=(
131
+ uuid.UUID(str(row.parent_chunk_id))
132
+ if row.parent_chunk_id
133
+ else None
134
+ ),
135
+ section=row.section,
136
+ content=row.content,
137
+ page_from=row.page_from,
138
+ page_to=row.page_to,
139
+ start_char=row.start_char,
140
+ end_char=row.end_char,
141
+ token_count=row.token_count,
142
+ embedding_model=row.embedding_model,
143
+ embedding_version=row.embedding_version,
144
+ )
145
+ chunks.append(ChunkWithScore(chunk=chunk, score=float(row.score)))
146
+ return chunks
147
+
148
+ async def search_lexical(
149
+ self,
150
+ tenant_id: uuid.UUID,
151
+ query: str,
152
+ top_k: int = 20,
153
+ *,
154
+ document_id: uuid.UUID | None = None,
155
+ ) -> list[ChunkWithScore]:
156
+ """Lexical search using PostgreSQL tsvector and ts_rank_cd.
157
+
158
+ Args:
159
+ tenant_id: Owning tenant.
160
+ query: Raw search query text.
161
+ top_k: Maximum results.
162
+ document_id: Optional filter to scope results to one document.
163
+
164
+ Returns:
165
+ Chunks ordered by descending ts_rank_cd score.
166
+ """
167
+ where_clauses = [
168
+ "tenant_id = :tenant_id",
169
+ "content_tsv IS NOT NULL",
170
+ "content_tsv @@ plainto_tsquery('english', :query)",
171
+ ]
172
+ params: dict[str, object] = {
173
+ "tenant_id": str(tenant_id),
174
+ "query": query,
175
+ "top_k": top_k,
176
+ }
177
+
178
+ if document_id is not None:
179
+ where_clauses.append("document_id = :document_id")
180
+ params["document_id"] = str(document_id)
181
+
182
+ where_sql = " AND ".join(where_clauses)
183
+
184
+ sql = text(f"""
185
+ SELECT id, tenant_id, resume_version_id, document_id,
186
+ chunk_index, parent_chunk_id, section, content,
187
+ page_from, page_to, start_char, end_char,
188
+ token_count, embedding_model, embedding_version,
189
+ ts_rank_cd(content_tsv, plainto_tsquery('english', :query)) AS score
190
+ FROM resume_chunks
191
+ WHERE {where_sql}
192
+ ORDER BY score DESC
193
+ LIMIT :top_k
194
+ """)
195
+
196
+ result = await self._session.execute(sql, params)
197
+ rows = result.fetchall()
198
+
199
+ chunks: list[ChunkWithScore] = []
200
+ for row in rows:
201
+ chunk = ResumeChunk(
202
+ id=uuid.UUID(str(row.id)),
203
+ tenant_id=uuid.UUID(str(row.tenant_id)),
204
+ resume_version_id=uuid.UUID(str(row.resume_version_id)),
205
+ document_id=uuid.UUID(str(row.document_id)),
206
+ chunk_index=row.chunk_index,
207
+ parent_chunk_id=(
208
+ uuid.UUID(str(row.parent_chunk_id))
209
+ if row.parent_chunk_id
210
+ else None
211
+ ),
212
+ section=row.section,
213
+ content=row.content,
214
+ page_from=row.page_from,
215
+ page_to=row.page_to,
216
+ start_char=row.start_char,
217
+ end_char=row.end_char,
218
+ token_count=row.token_count,
219
+ embedding_model=row.embedding_model,
220
+ embedding_version=row.embedding_version,
221
+ )
222
+ chunks.append(ChunkWithScore(chunk=chunk, score=float(row.score)))
223
+ return chunks
224
+
225
+ async def get_parents(self, chunk_ids: list[uuid.UUID]) -> list[ResumeChunk]:
226
+ """Retrieve parent chunks by their IDs.
227
+
228
+ Used for parent expansion after child retrieval.
229
+
230
+ Args:
231
+ chunk_ids: Parent chunk IDs to fetch.
232
+
233
+ Returns:
234
+ The parent chunks.
235
+ """
236
+ if not chunk_ids:
237
+ return []
238
+ stmt = select(ResumeChunk).where(ResumeChunk.id.in_(chunk_ids))
239
+ result = await self._session.execute(stmt)
240
+ return list(result.scalars().all())
241
+
242
+ async def get_by_document(
243
+ self,
244
+ tenant_id: uuid.UUID,
245
+ document_id: uuid.UUID,
246
+ *,
247
+ children_only: bool = False,
248
+ ) -> Sequence[ResumeChunk]:
249
+ """Return all chunks for a document.
250
+
251
+ Args:
252
+ tenant_id: Owning tenant.
253
+ document_id: Document identifier.
254
+ children_only: If True, exclude parent chunks.
255
+
256
+ Returns:
257
+ Chunks ordered by chunk_index.
258
+ """
259
+ stmt = (
260
+ select(ResumeChunk)
261
+ .where(
262
+ ResumeChunk.tenant_id == tenant_id,
263
+ ResumeChunk.document_id == document_id,
264
+ )
265
+ .order_by(ResumeChunk.chunk_index)
266
+ )
267
+ if children_only:
268
+ stmt = stmt.where(ResumeChunk.parent_chunk_id.isnot(None))
269
+ result = await self._session.execute(stmt)
270
+ return result.scalars().all()
271
+
272
+ async def count_by_tenant(self, tenant_id: uuid.UUID) -> int:
273
+ """Count total chunks for a tenant.
274
+
275
+ Args:
276
+ tenant_id: Owning tenant.
277
+
278
+ Returns:
279
+ Total chunk count.
280
+ """
281
+ stmt = (
282
+ select(func.count())
283
+ .select_from(ResumeChunk)
284
+ .where(ResumeChunk.tenant_id == tenant_id)
285
+ )
286
+ result = await self._session.execute(stmt)
287
+ return int(result.scalar() or 0)
serving/app/routers/search.py ADDED
@@ -0,0 +1,157 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Semantic and hybrid talent search endpoints.
2
+
3
+ Two endpoints per ARCHITECTURE.md §7.2:
4
+ - ``POST /api/v1/search/candidates`` — hybrid talent search
5
+ - ``POST /api/v1/search/similar`` — "more like this"
6
+
7
+ No LLM in this path. Sub-second. Tenant-scoped via ``ReadPrincipal``.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from fastapi import APIRouter, status
13
+
14
+ from app.db import DbSession
15
+ from app.schemas.search import (
16
+ CandidateSearchHitResponse,
17
+ EvidenceSpanResponse,
18
+ SearchRequest,
19
+ SearchResponse,
20
+ SimilarRequest,
21
+ )
22
+ from app.security import ReadPrincipal
23
+ from app.services.embedding import get_embedding_service
24
+ from app.services.reranker import get_reranker_service
25
+ from app.services.search import SearchMode as ServiceSearchMode
26
+ from app.services.search import search_candidates, search_similar
27
+
28
+ router = APIRouter(prefix="/search", tags=["search"])
29
+
30
+
31
+ def _to_response(result: object) -> SearchResponse:
32
+ """Convert a service-layer SearchResult to an API response.
33
+
34
+ Args:
35
+ result: The SearchResult from the search service.
36
+
37
+ Returns:
38
+ The API-shaped response.
39
+ """
40
+ from app.services.search import SearchResult
41
+
42
+ assert isinstance(result, SearchResult) # noqa: S101
43
+
44
+ return SearchResponse(
45
+ items=[
46
+ CandidateSearchHitResponse(
47
+ document_id=hit.document_id,
48
+ score=hit.score,
49
+ spans=[
50
+ EvidenceSpanResponse(
51
+ chunk_id=span.chunk_id,
52
+ content=span.content,
53
+ section=span.section,
54
+ page_from=span.page_from,
55
+ page_to=span.page_to,
56
+ start_char=span.start_char,
57
+ end_char=span.end_char,
58
+ score=span.score,
59
+ )
60
+ for span in hit.spans
61
+ ],
62
+ )
63
+ for hit in result.items
64
+ ],
65
+ count=result.count,
66
+ query=result.query,
67
+ mode=result.mode,
68
+ )
69
+
70
+
71
+ @router.post(
72
+ "/candidates",
73
+ response_model=SearchResponse,
74
+ status_code=status.HTTP_200_OK,
75
+ summary="Search candidates",
76
+ description=(
77
+ "Natural-language search across the tenant's candidate pool. "
78
+ "Returns ranked candidates with evidence-cited results. "
79
+ "Supports hybrid (default), semantic-only, and lexical-only modes. "
80
+ "No LLM in this path."
81
+ ),
82
+ )
83
+ async def search_candidates_endpoint(
84
+ body: SearchRequest,
85
+ principal: ReadPrincipal,
86
+ session: DbSession,
87
+ ) -> SearchResponse:
88
+ """Execute a talent search.
89
+
90
+ Args:
91
+ body: Search parameters.
92
+ principal: Verified caller.
93
+ session: Database session.
94
+
95
+ Returns:
96
+ Ranked candidates with supporting evidence spans.
97
+ """
98
+ embedder = get_embedding_service()
99
+ reranker = get_reranker_service()
100
+
101
+ # Map schema enum to service enum
102
+ mode_map = {
103
+ "hybrid": ServiceSearchMode.HYBRID,
104
+ "semantic": ServiceSearchMode.SEMANTIC,
105
+ "lexical": ServiceSearchMode.LEXICAL,
106
+ }
107
+ service_mode = mode_map.get(body.mode.value, ServiceSearchMode.HYBRID)
108
+
109
+ result = await search_candidates(
110
+ session=session,
111
+ embedder=embedder,
112
+ reranker=reranker,
113
+ tenant_id=principal.tenant_id,
114
+ query=body.query,
115
+ top_k=body.top_k,
116
+ mode=service_mode,
117
+ )
118
+ return _to_response(result)
119
+
120
+
121
+ @router.post(
122
+ "/similar",
123
+ response_model=SearchResponse,
124
+ status_code=status.HTTP_200_OK,
125
+ summary="Find similar candidates",
126
+ description=(
127
+ '"More like this" — find candidates similar to a given resume. '
128
+ "Uses the source document's embeddings to search across the corpus."
129
+ ),
130
+ )
131
+ async def search_similar_endpoint(
132
+ body: SimilarRequest,
133
+ principal: ReadPrincipal,
134
+ session: DbSession,
135
+ ) -> SearchResponse:
136
+ """Find candidates similar to a given document.
137
+
138
+ Args:
139
+ body: The source document ID and result count.
140
+ principal: Verified caller.
141
+ session: Database session.
142
+
143
+ Returns:
144
+ Ranked similar candidates with evidence spans.
145
+ """
146
+ embedder = get_embedding_service()
147
+ reranker = get_reranker_service()
148
+
149
+ result = await search_similar(
150
+ session=session,
151
+ embedder=embedder,
152
+ reranker=reranker,
153
+ tenant_id=principal.tenant_id,
154
+ document_id=body.document_id,
155
+ top_k=body.top_k,
156
+ )
157
+ return _to_response(result)
serving/app/schemas/search.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Request and response schemas for semantic and hybrid talent search."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import uuid
6
+ from enum import Enum
7
+
8
+ from pydantic import BaseModel, Field
9
+
10
+
11
+ class SearchMode(str, Enum):
12
+ """Search retrieval strategy."""
13
+
14
+ HYBRID = "hybrid"
15
+ SEMANTIC = "semantic"
16
+ LEXICAL = "lexical"
17
+
18
+
19
+ class SearchFilters(BaseModel):
20
+ """Optional filters to narrow search results."""
21
+
22
+ sections: list[str] | None = Field(
23
+ default=None,
24
+ description="Restrict search to specific resume sections (e.g. experience, skills).",
25
+ )
26
+
27
+
28
+ class SearchRequest(BaseModel):
29
+ """Payload for ``POST /api/v1/search/candidates``."""
30
+
31
+ query: str = Field(
32
+ min_length=1,
33
+ max_length=2000,
34
+ description="Natural-language search query.",
35
+ )
36
+ filters: SearchFilters | None = Field(
37
+ default=None,
38
+ description="Optional search filters.",
39
+ )
40
+ top_k: int = Field(
41
+ default=10,
42
+ ge=1,
43
+ le=100,
44
+ description="Maximum number of candidates to return.",
45
+ )
46
+ mode: SearchMode = Field(
47
+ default=SearchMode.HYBRID,
48
+ description="Retrieval strategy: hybrid, semantic, or lexical.",
49
+ )
50
+
51
+
52
+ class EvidenceSpanResponse(BaseModel):
53
+ """One piece of evidence supporting a search result."""
54
+
55
+ chunk_id: uuid.UUID = Field(description="Identifier of the source chunk.")
56
+ content: str = Field(description="The matching chunk text.")
57
+ section: str = Field(description="Resume section this chunk belongs to.")
58
+ page_from: int = Field(description="Starting page (0-based).")
59
+ page_to: int = Field(description="Ending page (0-based).")
60
+ start_char: int = Field(description="Starting character offset in the full document.")
61
+ end_char: int = Field(description="Ending character offset in the full document.")
62
+ score: float = Field(description="Retrieval score for this chunk.")
63
+
64
+
65
+ class CandidateSearchHitResponse(BaseModel):
66
+ """One candidate in search results."""
67
+
68
+ document_id: uuid.UUID = Field(description="The resume document ID.")
69
+ score: float = Field(description="Aggregated relevance score.")
70
+ spans: list[EvidenceSpanResponse] = Field(
71
+ default_factory=list,
72
+ description="Best supporting evidence spans.",
73
+ )
74
+
75
+
76
+ class SearchResponse(BaseModel):
77
+ """Complete search result set."""
78
+
79
+ items: list[CandidateSearchHitResponse] = Field(
80
+ default_factory=list,
81
+ description="Ranked candidate hits.",
82
+ )
83
+ count: int = Field(description="Number of candidates returned.")
84
+ query: str = Field(description="The original query text.")
85
+ mode: str = Field(description="The search mode used.")
86
+
87
+
88
+ class SimilarRequest(BaseModel):
89
+ """Payload for ``POST /api/v1/search/similar``."""
90
+
91
+ document_id: uuid.UUID = Field(
92
+ description="Source document to find similar candidates for.",
93
+ )
94
+ top_k: int = Field(
95
+ default=10,
96
+ ge=1,
97
+ le=100,
98
+ description="Maximum number of candidates to return.",
99
+ )
serving/app/services/chunking.py ADDED
@@ -0,0 +1,365 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Parent-child chunking with offset preservation.
2
+
3
+ Implements the chunking strategy from ARCHITECTURE.md §6.4:
4
+ - **Child chunks:** ~180 words for embedding precision
5
+ - **Parent chunks:** ~700 words fed to the LLM for context (Phase 4)
6
+
7
+ Each child carries ``parent_chunk_id`` referencing its parent, enabling the
8
+ retrieve-on-children, expand-to-parents pattern.
9
+
10
+ Section detection is heuristic: common resume headings are matched
11
+ case-insensitively and mapped to one of the standard categories.
12
+ """
13
+
14
+ from __future__ import annotations
15
+
16
+ import re
17
+ import uuid
18
+ from dataclasses import dataclass
19
+
20
+ from app.logging import get_logger
21
+
22
+ logger = get_logger(__name__)
23
+
24
+ # Target sizes in words
25
+ CHILD_CHUNK_WORDS = 180
26
+ PARENT_CHUNK_WORDS = 700
27
+
28
+ # Approximate token-to-word ratio (words × factor ≈ tokens)
29
+ _WORD_TO_TOKEN_RATIO = 1.33
30
+
31
+ # Section heading patterns (case-insensitive)
32
+ _SECTION_PATTERNS: list[tuple[str, re.Pattern[str]]] = [
33
+ (
34
+ "experience",
35
+ re.compile(r"(?i)^\s*(work\s+)?experience|employment|career|professional\s+history"),
36
+ ),
37
+ ("education", re.compile(r"(?i)^\s*education|academic|qualifications|degrees?")),
38
+ ("skills", re.compile(r"(?i)^\s*(technical\s+)?skills|competenc|technologies|proficienc")),
39
+ (
40
+ "summary",
41
+ re.compile(r"(?i)^\s*(professional\s+)?(summary|profile|objective|about\s+me|overview)"),
42
+ ),
43
+ ("certifications", re.compile(r"(?i)^\s*certifications?|licens|credentials?")),
44
+ ("projects", re.compile(r"(?i)^\s*projects?|portfolio")),
45
+ ("languages", re.compile(r"(?i)^\s*languages?")),
46
+ ("awards", re.compile(r"(?i)^\s*awards?|honors?|achievements?")),
47
+ ("publications", re.compile(r"(?i)^\s*publications?|papers?|research")),
48
+ ("references", re.compile(r"(?i)^\s*references?")),
49
+ ]
50
+
51
+
52
+ @dataclass(frozen=True, slots=True)
53
+ class ChunkInfo:
54
+ """One chunk produced by the chunking pipeline.
55
+
56
+ Attributes:
57
+ chunk_id: Pre-assigned UUID for the chunk.
58
+ parent_chunk_id: UUID of the parent chunk, or None for parents.
59
+ chunk_index: Zero-based position in the document's chunk sequence.
60
+ content: The chunk text.
61
+ section: Detected resume section.
62
+ page_from: Starting page (0-based).
63
+ page_to: Ending page (0-based).
64
+ start_char: Starting character offset in the full text.
65
+ end_char: Ending character offset in the full text.
66
+ token_count: Estimated token count.
67
+ is_parent: True for parent chunks, False for children.
68
+ """
69
+
70
+ chunk_id: uuid.UUID
71
+ parent_chunk_id: uuid.UUID | None
72
+ chunk_index: int
73
+ content: str
74
+ section: str
75
+ page_from: int
76
+ page_to: int
77
+ start_char: int
78
+ end_char: int
79
+ token_count: int
80
+ is_parent: bool
81
+
82
+
83
+ @dataclass
84
+ class _Section:
85
+ """An intermediate section detected during segmentation."""
86
+
87
+ label: str
88
+ text: str
89
+ start_char: int
90
+ end_char: int
91
+
92
+
93
+ def _detect_section(line: str) -> str | None:
94
+ """Match a line against known section headings.
95
+
96
+ Args:
97
+ line: A single line of text.
98
+
99
+ Returns:
100
+ The section label if matched, else None.
101
+ """
102
+ stripped = line.strip()
103
+ if not stripped or len(stripped) > 80:
104
+ return None
105
+ for label, pattern in _SECTION_PATTERNS:
106
+ if pattern.match(stripped):
107
+ return label
108
+ return None
109
+
110
+
111
+ def _estimate_tokens(text: str) -> int:
112
+ """Estimate token count from word count.
113
+
114
+ Args:
115
+ text: The text to estimate.
116
+
117
+ Returns:
118
+ Approximate token count.
119
+ """
120
+ words = len(text.split())
121
+ return max(1, round(words * _WORD_TO_TOKEN_RATIO))
122
+
123
+
124
+ def _resolve_page(
125
+ char_offset: int,
126
+ page_offsets: list[dict[str, int]],
127
+ ) -> int:
128
+ """Resolve a character offset to a page number.
129
+
130
+ Args:
131
+ char_offset: Character offset in the full text.
132
+ page_offsets: Page boundary list from the parser.
133
+
134
+ Returns:
135
+ The 0-based page number, or 0 if offsets are empty.
136
+ """
137
+ if not page_offsets:
138
+ return 0
139
+ for page_info in reversed(page_offsets):
140
+ if char_offset >= page_info.get("start_char", 0):
141
+ return int(page_info.get("page", 0))
142
+ return 0
143
+
144
+
145
+ def _segment_into_sections(text: str) -> list[_Section]:
146
+ """Split text into sections based on heading detection.
147
+
148
+ Args:
149
+ text: The full document text.
150
+
151
+ Returns:
152
+ A list of sections with their character offsets.
153
+ """
154
+ lines = text.split("\n")
155
+ sections: list[_Section] = []
156
+ current_label = "summary"
157
+ current_lines: list[str] = []
158
+ current_start = 0
159
+ char_pos = 0
160
+
161
+ for line in lines:
162
+ detected = _detect_section(line)
163
+ if detected is not None and current_lines:
164
+ section_text = "\n".join(current_lines).strip()
165
+ if section_text:
166
+ sections.append(_Section(
167
+ label=current_label,
168
+ text=section_text,
169
+ start_char=current_start,
170
+ end_char=current_start + len(section_text),
171
+ ))
172
+ current_label = detected
173
+ current_lines = []
174
+ current_start = char_pos
175
+ else:
176
+ current_lines.append(line)
177
+ char_pos += len(line) + 1 # +1 for the newline
178
+
179
+ # Flush last section
180
+ section_text = "\n".join(current_lines).strip()
181
+ if section_text:
182
+ sections.append(_Section(
183
+ label=current_label,
184
+ text=section_text,
185
+ start_char=current_start,
186
+ end_char=current_start + len(section_text),
187
+ ))
188
+
189
+ # If no sections detected, wrap the entire text as "other"
190
+ if not sections:
191
+ sections.append(_Section(
192
+ label="other",
193
+ text=text.strip(),
194
+ start_char=0,
195
+ end_char=len(text.strip()),
196
+ ))
197
+
198
+ return sections
199
+
200
+
201
+ def _split_into_word_chunks(
202
+ text: str,
203
+ target_words: int,
204
+ base_start_char: int,
205
+ ) -> list[tuple[str, int, int]]:
206
+ """Split text into chunks of approximately ``target_words`` words.
207
+
208
+ Splits on sentence boundaries when possible, falling back to word
209
+ boundaries.
210
+
211
+ Args:
212
+ text: Text to split.
213
+ target_words: Target word count per chunk.
214
+ base_start_char: Character offset of ``text`` within the full document.
215
+
216
+ Returns:
217
+ List of (chunk_text, start_char, end_char) tuples.
218
+ """
219
+ if not text.strip():
220
+ return []
221
+
222
+ words = text.split()
223
+ if len(words) <= target_words:
224
+ return [(text.strip(), base_start_char, base_start_char + len(text.strip()))]
225
+
226
+ # Try sentence splitting first
227
+ sentences = re.split(r"(?<=[.!?])\s+", text)
228
+ chunks: list[tuple[str, int, int]] = []
229
+ current_words: list[str] = []
230
+ current_start = base_start_char
231
+
232
+ for sentence in sentences:
233
+ sentence_words = sentence.split()
234
+ if current_words and len(current_words) + len(sentence_words) > target_words:
235
+ chunk_text = " ".join(current_words).strip()
236
+ if chunk_text:
237
+ chunks.append((
238
+ chunk_text,
239
+ current_start,
240
+ current_start + len(chunk_text),
241
+ ))
242
+ current_start = current_start + len(chunk_text) + 1
243
+ current_words = sentence_words
244
+ else:
245
+ current_words.extend(sentence_words)
246
+
247
+ # Flush remaining
248
+ if current_words:
249
+ chunk_text = " ".join(current_words).strip()
250
+ if chunk_text:
251
+ chunks.append((
252
+ chunk_text,
253
+ current_start,
254
+ current_start + len(chunk_text),
255
+ ))
256
+
257
+ return chunks
258
+
259
+
260
+ def chunk_document(
261
+ text: str,
262
+ page_offsets: list[dict[str, int]] | None = None,
263
+ child_words: int = CHILD_CHUNK_WORDS,
264
+ parent_words: int = PARENT_CHUNK_WORDS,
265
+ ) -> list[ChunkInfo]:
266
+ """Chunk a document into parent-child chunks with offset preservation.
267
+
268
+ The strategy is:
269
+ 1. Segment text into resume sections (experience, skills, etc.)
270
+ 2. Within each section, create parent chunks (~700 words)
271
+ 3. Within each parent, create child chunks (~180 words)
272
+ 4. Children carry ``parent_chunk_id``; parents carry ``is_parent=True``
273
+
274
+ Args:
275
+ text: The full document text.
276
+ page_offsets: Page boundary data from the parser.
277
+ child_words: Target words per child chunk.
278
+ parent_words: Target words per parent chunk.
279
+
280
+ Returns:
281
+ All chunks (parents and children), ordered by chunk_index.
282
+ """
283
+ if not text or not text.strip():
284
+ return []
285
+
286
+ offsets = page_offsets or []
287
+ sections = _segment_into_sections(text)
288
+
289
+ all_chunks: list[ChunkInfo] = []
290
+ chunk_index = 0
291
+
292
+ for section in sections:
293
+ # Create parent chunks from the section
294
+ parent_spans = _split_into_word_chunks(
295
+ section.text, parent_words, section.start_char
296
+ )
297
+
298
+ for parent_text, parent_start, parent_end in parent_spans:
299
+ parent_id = uuid.uuid4()
300
+
301
+ # Create the parent chunk
302
+ parent_chunk = ChunkInfo(
303
+ chunk_id=parent_id,
304
+ parent_chunk_id=None,
305
+ chunk_index=chunk_index,
306
+ content=parent_text,
307
+ section=section.label,
308
+ page_from=_resolve_page(parent_start, offsets),
309
+ page_to=_resolve_page(parent_end, offsets),
310
+ start_char=parent_start,
311
+ end_char=parent_end,
312
+ token_count=_estimate_tokens(parent_text),
313
+ is_parent=True,
314
+ )
315
+ all_chunks.append(parent_chunk)
316
+ chunk_index += 1
317
+
318
+ # Create child chunks within this parent
319
+ child_spans = _split_into_word_chunks(
320
+ parent_text, child_words, parent_start
321
+ )
322
+
323
+ # If the parent is small enough, it serves as its own child
324
+ if len(child_spans) <= 1:
325
+ child_chunk = ChunkInfo(
326
+ chunk_id=uuid.uuid4(),
327
+ parent_chunk_id=parent_id,
328
+ chunk_index=chunk_index,
329
+ content=parent_text,
330
+ section=section.label,
331
+ page_from=parent_chunk.page_from,
332
+ page_to=parent_chunk.page_to,
333
+ start_char=parent_start,
334
+ end_char=parent_end,
335
+ token_count=parent_chunk.token_count,
336
+ is_parent=False,
337
+ )
338
+ all_chunks.append(child_chunk)
339
+ chunk_index += 1
340
+ else:
341
+ for child_text, child_start, child_end in child_spans:
342
+ child_chunk = ChunkInfo(
343
+ chunk_id=uuid.uuid4(),
344
+ parent_chunk_id=parent_id,
345
+ chunk_index=chunk_index,
346
+ content=child_text,
347
+ section=section.label,
348
+ page_from=_resolve_page(child_start, offsets),
349
+ page_to=_resolve_page(child_end, offsets),
350
+ start_char=child_start,
351
+ end_char=child_end,
352
+ token_count=_estimate_tokens(child_text),
353
+ is_parent=False,
354
+ )
355
+ all_chunks.append(child_chunk)
356
+ chunk_index += 1
357
+
358
+ logger.info(
359
+ "document_chunked",
360
+ total_chunks=len(all_chunks),
361
+ parent_chunks=sum(1 for c in all_chunks if c.is_parent),
362
+ child_chunks=sum(1 for c in all_chunks if not c.is_parent),
363
+ sections=len(sections),
364
+ )
365
+ return all_chunks
serving/app/services/embedding.py ADDED
@@ -0,0 +1,326 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Embedding service port and adapters.
2
+
3
+ The ``EmbeddingService`` protocol abstracts the embedding provider so the rest
4
+ of the application depends on a capability, not on a deployment choice. Two
5
+ adapters are provided:
6
+
7
+ * ``TEIEmbeddingService`` — calls HuggingFace Text Embeddings Inference over
8
+ HTTP. This is the production path.
9
+ * ``MockEmbeddingService`` — returns deterministic vectors of the configured
10
+ dimension. Used in tests; not a fake implementation, just a test double.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ import hashlib
16
+ import math
17
+ from typing import Protocol, runtime_checkable
18
+
19
+ import httpx
20
+
21
+ from app.config import Settings, get_settings
22
+ from app.exceptions import EmbeddingServiceUnavailableError
23
+ from app.logging import get_logger
24
+
25
+ logger = get_logger(__name__)
26
+
27
+ # ---------------------------------------------------------------------------
28
+ # Port
29
+ # ---------------------------------------------------------------------------
30
+
31
+
32
+ @runtime_checkable
33
+ class EmbeddingService(Protocol):
34
+ """Capability: embed text into dense vectors."""
35
+
36
+ @property
37
+ def model_name(self) -> str:
38
+ """The model identifier used for provenance tracking."""
39
+ ...
40
+
41
+ @property
42
+ def model_version(self) -> str:
43
+ """A version tag for the embedder, stored on every chunk row."""
44
+ ...
45
+
46
+ @property
47
+ def dimension(self) -> int:
48
+ """Dimensionality of the vectors this service produces."""
49
+ ...
50
+
51
+ async def embed_texts(self, texts: list[str]) -> list[list[float]]:
52
+ """Embed a batch of passage texts.
53
+
54
+ Args:
55
+ texts: The texts to embed.
56
+
57
+ Returns:
58
+ A list of embedding vectors, one per input text.
59
+
60
+ Raises:
61
+ EmbeddingServiceUnavailableError: If the service is unreachable.
62
+ """
63
+ ...
64
+
65
+ async def embed_query(self, query: str) -> list[float]:
66
+ """Embed a single search query.
67
+
68
+ Some models use different prefixes for queries vs passages. This
69
+ method handles that distinction.
70
+
71
+ Args:
72
+ query: The search query text.
73
+
74
+ Returns:
75
+ The embedding vector.
76
+
77
+ Raises:
78
+ EmbeddingServiceUnavailableError: If the service is unreachable.
79
+ """
80
+ ...
81
+
82
+
83
+ # ---------------------------------------------------------------------------
84
+ # TEI adapter (production)
85
+ # ---------------------------------------------------------------------------
86
+
87
+ _TEI_EMBED_PATH = "/embed"
88
+ _TEI_TIMEOUT_SECONDS = 30.0
89
+
90
+
91
+ class TEIEmbeddingService:
92
+ """Calls HuggingFace Text Embeddings Inference over HTTP.
93
+
94
+ TEI serves models like ``bge-m3`` and exposes a simple JSON API. This
95
+ adapter batches according to ``settings.embedding_batch_size`` and
96
+ retries once on transient failures.
97
+
98
+ Attributes:
99
+ _endpoint: Base URL of the TEI instance.
100
+ _model: Model name for provenance.
101
+ _dim: Expected vector dimension.
102
+ _batch_size: Max texts per HTTP request.
103
+ """
104
+
105
+ def __init__(self, settings: Settings | None = None) -> None:
106
+ """Initialize from settings.
107
+
108
+ Args:
109
+ settings: Optional configuration override.
110
+ """
111
+ cfg = settings or get_settings()
112
+ self._endpoint = cfg.embedding_endpoint.rstrip("/")
113
+ self._model = cfg.embedding_model
114
+ self._dim = cfg.embedding_dim
115
+ self._batch_size = cfg.embedding_batch_size
116
+ self._version = "1"
117
+
118
+ @property
119
+ def model_name(self) -> str:
120
+ """Return the configured model identifier."""
121
+ return self._model
122
+
123
+ @property
124
+ def model_version(self) -> str:
125
+ """Return the embedding version tag."""
126
+ return self._version
127
+
128
+ @property
129
+ def dimension(self) -> int:
130
+ """Return the configured embedding dimension."""
131
+ return self._dim
132
+
133
+ async def embed_texts(self, texts: list[str]) -> list[list[float]]:
134
+ """Embed passage texts via TEI, batched.
135
+
136
+ Args:
137
+ texts: The texts to embed.
138
+
139
+ Returns:
140
+ A list of embedding vectors.
141
+
142
+ Raises:
143
+ EmbeddingServiceUnavailableError: On HTTP or connection error.
144
+ """
145
+ if not texts:
146
+ return []
147
+
148
+ all_embeddings: list[list[float]] = []
149
+ for i in range(0, len(texts), self._batch_size):
150
+ batch = texts[i : i + self._batch_size]
151
+ embeddings = await self._call_tei(batch)
152
+ all_embeddings.extend(embeddings)
153
+ return all_embeddings
154
+
155
+ async def embed_query(self, query: str) -> list[float]:
156
+ """Embed a search query via TEI.
157
+
158
+ Args:
159
+ query: The search query.
160
+
161
+ Returns:
162
+ The embedding vector.
163
+
164
+ Raises:
165
+ EmbeddingServiceUnavailableError: On HTTP or connection error.
166
+ """
167
+ results = await self._call_tei([query])
168
+ return results[0]
169
+
170
+ async def _call_tei(self, inputs: list[str]) -> list[list[float]]:
171
+ """Make a single TEI embed request.
172
+
173
+ Args:
174
+ inputs: Batch of texts to embed.
175
+
176
+ Returns:
177
+ The embedding vectors.
178
+
179
+ Raises:
180
+ EmbeddingServiceUnavailableError: On any failure.
181
+ """
182
+ url = f"{self._endpoint}{_TEI_EMBED_PATH}"
183
+ payload = {"inputs": inputs}
184
+ try:
185
+ async with httpx.AsyncClient(timeout=_TEI_TIMEOUT_SECONDS) as client:
186
+ response = await client.post(url, json=payload)
187
+ response.raise_for_status()
188
+ data: list[list[float]] = response.json()
189
+ return data
190
+ except (httpx.HTTPError, httpx.ConnectError, Exception) as exc:
191
+ logger.error(
192
+ "embedding_service_error",
193
+ endpoint=self._endpoint,
194
+ error=str(exc)[:200],
195
+ )
196
+ raise EmbeddingServiceUnavailableError(
197
+ f"Embedding service at {self._endpoint} is unavailable."
198
+ ) from exc
199
+
200
+
201
+ # ---------------------------------------------------------------------------
202
+ # Mock adapter (tests)
203
+ # ---------------------------------------------------------------------------
204
+
205
+
206
+ class MockEmbeddingService:
207
+ """Returns deterministic vectors for testing.
208
+
209
+ Vectors are derived from a hash of the input text so the same text always
210
+ produces the same vector, enabling assertions in tests. This is *not* a
211
+ fake implementation — it makes no attempt at semantic similarity.
212
+ """
213
+
214
+ def __init__(self, dimension: int = 1024) -> None:
215
+ """Initialize with a target dimension.
216
+
217
+ Args:
218
+ dimension: Vector dimensionality to produce.
219
+ """
220
+ self._dim = dimension
221
+
222
+ @property
223
+ def model_name(self) -> str:
224
+ """Return a test model identifier."""
225
+ return "mock-embedder"
226
+
227
+ @property
228
+ def model_version(self) -> str:
229
+ """Return a test version tag."""
230
+ return "test-v1"
231
+
232
+ @property
233
+ def dimension(self) -> int:
234
+ """Return the configured dimension."""
235
+ return self._dim
236
+
237
+ async def embed_texts(self, texts: list[str]) -> list[list[float]]:
238
+ """Return deterministic vectors for each text.
239
+
240
+ Args:
241
+ texts: The texts to embed.
242
+
243
+ Returns:
244
+ Deterministic vectors derived from text hashes.
245
+ """
246
+ return [self._deterministic_vector(t) for t in texts]
247
+
248
+ async def embed_query(self, query: str) -> list[float]:
249
+ """Return a deterministic vector for the query.
250
+
251
+ Args:
252
+ query: The search query.
253
+
254
+ Returns:
255
+ A deterministic vector.
256
+ """
257
+ return self._deterministic_vector(query)
258
+
259
+ def _deterministic_vector(self, text: str) -> list[float]:
260
+ """Produce a unit-length vector deterministically from text.
261
+
262
+ Args:
263
+ text: Input text.
264
+
265
+ Returns:
266
+ A normalized vector of ``self._dim`` dimensions.
267
+ """
268
+ digest = hashlib.sha256(text.encode("utf-8")).digest()
269
+ # Expand the 32-byte hash to fill the dimension by repeating
270
+ raw = list(digest) * ((self._dim // len(digest)) + 1)
271
+ raw = raw[: self._dim]
272
+ # Map bytes to floats in [-1, 1] and normalize
273
+ vec = [(b - 128) / 128.0 for b in raw]
274
+ norm = math.sqrt(sum(v * v for v in vec)) or 1.0
275
+ return [v / norm for v in vec]
276
+
277
+
278
+ # ---------------------------------------------------------------------------
279
+ # Singleton management
280
+ # ---------------------------------------------------------------------------
281
+
282
+ _embedding_service: EmbeddingService | None = None
283
+
284
+
285
+ def get_embedding_service() -> EmbeddingService:
286
+ """Return the process-wide embedding service instance.
287
+
288
+ Returns a ``MockEmbeddingService`` when no TEI endpoint is configured
289
+ (i.e. in test or when the endpoint is empty).
290
+
291
+ Returns:
292
+ The embedding service singleton.
293
+ """
294
+ global _embedding_service # noqa: PLW0603
295
+ if _embedding_service is not None:
296
+ return _embedding_service
297
+
298
+ settings = get_settings()
299
+ if settings.embedding_endpoint and settings.environment != "test":
300
+ _embedding_service = TEIEmbeddingService(settings)
301
+ logger.info(
302
+ "embedding_service_initialized",
303
+ adapter="TEI",
304
+ endpoint=settings.embedding_endpoint,
305
+ model=settings.embedding_model,
306
+ )
307
+ else:
308
+ _embedding_service = MockEmbeddingService(dimension=settings.embedding_dim)
309
+ logger.info("embedding_service_initialized", adapter="Mock")
310
+ return _embedding_service
311
+
312
+
313
+ def set_embedding_service(service: EmbeddingService) -> None:
314
+ """Override the embedding service singleton (for tests).
315
+
316
+ Args:
317
+ service: The service to install.
318
+ """
319
+ global _embedding_service # noqa: PLW0603
320
+ _embedding_service = service
321
+
322
+
323
+ def reset_embedding_service() -> None:
324
+ """Clear the embedding service singleton (for tests)."""
325
+ global _embedding_service # noqa: PLW0603
326
+ _embedding_service = None
serving/app/services/reranker.py ADDED
@@ -0,0 +1,209 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Reranker service port and adapters.
2
+
3
+ Cross-encoder reranking improves precision by scoring query–document pairs with
4
+ a more powerful model after initial recall. Two adapters:
5
+
6
+ * ``TEIRerankerService`` — calls HuggingFace TEI rerank endpoint.
7
+ * ``NoOpRerankerService`` — passes through scores unchanged (CPU fallback).
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from dataclasses import dataclass
13
+ from typing import Protocol, runtime_checkable
14
+
15
+ import httpx
16
+
17
+ from app.config import Settings, get_settings
18
+ from app.logging import get_logger
19
+
20
+ logger = get_logger(__name__)
21
+
22
+
23
+ @dataclass(frozen=True, slots=True)
24
+ class RerankResult:
25
+ """One reranked document.
26
+
27
+ Attributes:
28
+ index: Original index in the input list.
29
+ score: Relevance score from the reranker.
30
+ """
31
+
32
+ index: int
33
+ score: float
34
+
35
+
36
+ # ---------------------------------------------------------------------------
37
+ # Port
38
+ # ---------------------------------------------------------------------------
39
+
40
+
41
+ @runtime_checkable
42
+ class RerankerService(Protocol):
43
+ """Capability: rerank documents by relevance to a query."""
44
+
45
+ async def rerank(
46
+ self, query: str, documents: list[str], top_k: int
47
+ ) -> list[RerankResult]:
48
+ """Rerank documents by relevance to the query.
49
+
50
+ Args:
51
+ query: The search query.
52
+ documents: Document texts to rerank.
53
+ top_k: Maximum results to return.
54
+
55
+ Returns:
56
+ The top-k results sorted by descending score.
57
+ """
58
+ ...
59
+
60
+
61
+ # ---------------------------------------------------------------------------
62
+ # TEI adapter (production)
63
+ # ---------------------------------------------------------------------------
64
+
65
+ _TEI_RERANK_PATH = "/rerank"
66
+ _TEI_RERANK_TIMEOUT = 30.0
67
+
68
+
69
+ class TEIRerankerService:
70
+ """Calls HuggingFace TEI rerank endpoint.
71
+
72
+ Attributes:
73
+ _endpoint: Base URL of the TEI reranker instance.
74
+ """
75
+
76
+ def __init__(self, settings: Settings | None = None) -> None:
77
+ """Initialize from settings.
78
+
79
+ Args:
80
+ settings: Optional configuration override.
81
+ """
82
+ cfg = settings or get_settings()
83
+ self._endpoint = cfg.reranker_endpoint.rstrip("/")
84
+ self._model = cfg.reranker_model
85
+
86
+ async def rerank(
87
+ self, query: str, documents: list[str], top_k: int
88
+ ) -> list[RerankResult]:
89
+ """Rerank via TEI.
90
+
91
+ Args:
92
+ query: The search query.
93
+ documents: Document texts.
94
+ top_k: Maximum results.
95
+
96
+ Returns:
97
+ Reranked results sorted by descending score.
98
+ """
99
+ if not documents:
100
+ return []
101
+
102
+ url = f"{self._endpoint}{_TEI_RERANK_PATH}"
103
+ payload = {
104
+ "query": query,
105
+ "texts": documents,
106
+ "return_text": False,
107
+ }
108
+ try:
109
+ async with httpx.AsyncClient(timeout=_TEI_RERANK_TIMEOUT) as client:
110
+ response = await client.post(url, json=payload)
111
+ response.raise_for_status()
112
+ raw: list[dict[str, object]] = response.json()
113
+ results = [
114
+ RerankResult(index=int(r["index"]), score=float(r["score"]))
115
+ for r in raw
116
+ ]
117
+ results.sort(key=lambda r: r.score, reverse=True)
118
+ return results[:top_k]
119
+ except (httpx.HTTPError, Exception) as exc:
120
+ logger.warning(
121
+ "reranker_service_error",
122
+ endpoint=self._endpoint,
123
+ error=str(exc)[:200],
124
+ )
125
+ # Graceful degradation: return original order with default scores
126
+ return [
127
+ RerankResult(index=i, score=1.0 / (i + 1))
128
+ for i in range(min(top_k, len(documents)))
129
+ ]
130
+
131
+
132
+ # ---------------------------------------------------------------------------
133
+ # No-op adapter (CPU fallback / dev mode)
134
+ # ---------------------------------------------------------------------------
135
+
136
+
137
+ class NoOpRerankerService:
138
+ """Passes through documents without reranking.
139
+
140
+ Used when no reranker endpoint is configured. Returns scores based on
141
+ input order so the RRF-fused ranking is preserved.
142
+ """
143
+
144
+ async def rerank(
145
+ self, query: str, documents: list[str], top_k: int
146
+ ) -> list[RerankResult]:
147
+ """Return input-order scores without reranking.
148
+
149
+ Args:
150
+ query: Ignored.
151
+ documents: Document texts.
152
+ top_k: Maximum results.
153
+
154
+ Returns:
155
+ Results in input order with position-based scores.
156
+ """
157
+ return [
158
+ RerankResult(index=i, score=1.0 / (i + 1))
159
+ for i in range(min(top_k, len(documents)))
160
+ ]
161
+
162
+
163
+ # ---------------------------------------------------------------------------
164
+ # Singleton management
165
+ # ---------------------------------------------------------------------------
166
+
167
+ _reranker_service: RerankerService | None = None
168
+
169
+
170
+ def get_reranker_service() -> RerankerService:
171
+ """Return the process-wide reranker service instance.
172
+
173
+ Returns a ``NoOpRerankerService`` when no reranker endpoint is configured.
174
+
175
+ Returns:
176
+ The reranker service singleton.
177
+ """
178
+ global _reranker_service # noqa: PLW0603
179
+ if _reranker_service is not None:
180
+ return _reranker_service
181
+
182
+ settings = get_settings()
183
+ if settings.reranker_endpoint:
184
+ _reranker_service = TEIRerankerService(settings)
185
+ logger.info(
186
+ "reranker_service_initialized",
187
+ adapter="TEI",
188
+ endpoint=settings.reranker_endpoint,
189
+ )
190
+ else:
191
+ _reranker_service = NoOpRerankerService()
192
+ logger.info("reranker_service_initialized", adapter="NoOp")
193
+ return _reranker_service
194
+
195
+
196
+ def set_reranker_service(service: RerankerService) -> None:
197
+ """Override the reranker service singleton (for tests).
198
+
199
+ Args:
200
+ service: The service to install.
201
+ """
202
+ global _reranker_service # noqa: PLW0603
203
+ _reranker_service = service
204
+
205
+
206
+ def reset_reranker_service() -> None:
207
+ """Clear the reranker service singleton (for tests)."""
208
+ global _reranker_service # noqa: PLW0603
209
+ _reranker_service = None
serving/app/services/search.py ADDED
@@ -0,0 +1,349 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Hybrid semantic and lexical talent search.
2
+
3
+ Orchestrates the retrieval pipeline described in ARCHITECTURE.md §12.4:
4
+ query → embed → hybrid recall (dense + lexical) → RRF fusion → rerank →
5
+ group by candidate → aggregate → return ranked results with evidence spans.
6
+
7
+ No LLM in this path — search is sub-second and free.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import uuid
13
+ from collections import defaultdict
14
+ from dataclasses import dataclass, field
15
+ from enum import Enum
16
+
17
+ from sqlalchemy.ext.asyncio import AsyncSession
18
+
19
+ from app.config import Settings, get_settings
20
+ from app.logging import get_logger
21
+ from app.repositories.search import ChunkRepository, ChunkWithScore
22
+ from app.services.embedding import EmbeddingService
23
+ from app.services.reranker import RerankerService
24
+
25
+ logger = get_logger(__name__)
26
+
27
+
28
+ class SearchMode(str, Enum):
29
+ """Search retrieval strategy."""
30
+
31
+ HYBRID = "hybrid"
32
+ SEMANTIC = "semantic"
33
+ LEXICAL = "lexical"
34
+
35
+
36
+ @dataclass(frozen=True, slots=True)
37
+ class EvidenceSpan:
38
+ """One piece of supporting evidence from a search result.
39
+
40
+ Attributes:
41
+ chunk_id: The chunk this evidence comes from.
42
+ content: The chunk text.
43
+ section: Resume section (experience, skills, etc.).
44
+ page_from: Starting page (0-based).
45
+ page_to: Ending page (0-based).
46
+ start_char: Starting character offset in the full document.
47
+ end_char: Ending character offset in the full document.
48
+ score: Retrieval score for this chunk.
49
+ """
50
+
51
+ chunk_id: uuid.UUID
52
+ content: str
53
+ section: str
54
+ page_from: int
55
+ page_to: int
56
+ start_char: int
57
+ end_char: int
58
+ score: float
59
+
60
+
61
+ @dataclass(frozen=True, slots=True)
62
+ class CandidateSearchHit:
63
+ """One candidate in search results.
64
+
65
+ Attributes:
66
+ document_id: The resume document ID.
67
+ score: Aggregated relevance score across all matching chunks.
68
+ spans: The best supporting evidence spans.
69
+ """
70
+
71
+ document_id: uuid.UUID
72
+ score: float
73
+ spans: list[EvidenceSpan] = field(default_factory=list)
74
+
75
+
76
+ @dataclass(frozen=True, slots=True)
77
+ class SearchResult:
78
+ """Complete search result set.
79
+
80
+ Attributes:
81
+ items: Ranked candidate hits.
82
+ count: Total number of candidates returned.
83
+ query: The original query text.
84
+ mode: The search mode used.
85
+ """
86
+
87
+ items: list[CandidateSearchHit]
88
+ count: int
89
+ query: str
90
+ mode: str
91
+
92
+
93
+ def reciprocal_rank_fusion(
94
+ ranked_lists: list[list[ChunkWithScore]],
95
+ weights: list[float],
96
+ k: int = 60,
97
+ ) -> list[ChunkWithScore]:
98
+ """Fuse multiple ranked lists using Reciprocal Rank Fusion.
99
+
100
+ RRF score for a document d = sum(weight_i / (k + rank_i(d))) across lists.
101
+
102
+ Args:
103
+ ranked_lists: Lists of ranked chunks from different retrieval methods.
104
+ weights: Weight for each list (must sum to ~1.0).
105
+ k: RRF constant (default 60, standard value).
106
+
107
+ Returns:
108
+ A single fused list sorted by descending RRF score.
109
+ """
110
+ scores: dict[uuid.UUID, float] = defaultdict(float)
111
+ chunk_map: dict[uuid.UUID, ChunkWithScore] = {}
112
+
113
+ for ranked_list, weight in zip(ranked_lists, weights, strict=False):
114
+ for rank, item in enumerate(ranked_list):
115
+ rrf_score = weight / (k + rank + 1)
116
+ scores[item.chunk.id] += rrf_score
117
+ # Keep the chunk with the highest individual score
118
+ if item.chunk.id not in chunk_map or item.score > chunk_map[item.chunk.id].score:
119
+ chunk_map[item.chunk.id] = item
120
+
121
+ # Build fused results
122
+ fused = [
123
+ ChunkWithScore(chunk=chunk_map[chunk_id].chunk, score=score)
124
+ for chunk_id, score in scores.items()
125
+ ]
126
+ fused.sort(key=lambda x: x.score, reverse=True)
127
+ return fused
128
+
129
+
130
+ def group_by_candidate(
131
+ chunks: list[ChunkWithScore],
132
+ max_spans_per_candidate: int = 5,
133
+ ) -> list[CandidateSearchHit]:
134
+ """Group chunks by candidate (document_id) and aggregate scores.
135
+
136
+ Per-candidate score is the sum of the best chunk scores (not mean, so
137
+ candidates with more matching evidence rank higher).
138
+
139
+ Args:
140
+ chunks: Scored chunks from retrieval.
141
+ max_spans_per_candidate: Maximum evidence spans to include per candidate.
142
+
143
+ Returns:
144
+ Candidates sorted by descending aggregate score.
145
+ """
146
+ by_document: dict[uuid.UUID, list[ChunkWithScore]] = defaultdict(list)
147
+ for chunk in chunks:
148
+ by_document[chunk.chunk.document_id].append(chunk)
149
+
150
+ hits: list[CandidateSearchHit] = []
151
+ for doc_id, doc_chunks in by_document.items():
152
+ # Sort chunks by score descending
153
+ doc_chunks.sort(key=lambda c: c.score, reverse=True)
154
+ best = doc_chunks[:max_spans_per_candidate]
155
+
156
+ # Aggregate score: sum of best chunk scores
157
+ total_score = sum(c.score for c in best)
158
+
159
+ spans = [
160
+ EvidenceSpan(
161
+ chunk_id=c.chunk.id,
162
+ content=c.chunk.content,
163
+ section=c.chunk.section,
164
+ page_from=c.chunk.page_from,
165
+ page_to=c.chunk.page_to,
166
+ start_char=c.chunk.start_char,
167
+ end_char=c.chunk.end_char,
168
+ score=c.score,
169
+ )
170
+ for c in best
171
+ ]
172
+ hits.append(CandidateSearchHit(
173
+ document_id=doc_id,
174
+ score=total_score,
175
+ spans=spans,
176
+ ))
177
+
178
+ hits.sort(key=lambda h: h.score, reverse=True)
179
+ return hits
180
+
181
+
182
+ async def search_candidates(
183
+ *,
184
+ session: AsyncSession,
185
+ embedder: EmbeddingService,
186
+ reranker: RerankerService,
187
+ tenant_id: uuid.UUID,
188
+ query: str,
189
+ top_k: int = 10,
190
+ mode: SearchMode = SearchMode.HYBRID,
191
+ settings: Settings | None = None,
192
+ ) -> SearchResult:
193
+ """Execute a hybrid talent search across the candidate pool.
194
+
195
+ Pipeline:
196
+ 1. Embed query
197
+ 2. Dense retrieval (pgvector HNSW cosine)
198
+ 3. Lexical retrieval (GIN tsvector ts_rank_cd)
199
+ 4. RRF fusion (configurable weights)
200
+ 5. Cross-encoder rerank
201
+ 6. Group by candidate, aggregate scores
202
+ 7. Return ranked results with evidence spans
203
+
204
+ Args:
205
+ session: Database session.
206
+ embedder: Embedding service.
207
+ reranker: Reranker service.
208
+ tenant_id: Owning tenant.
209
+ query: Natural-language search query.
210
+ top_k: Maximum candidates to return.
211
+ mode: Retrieval strategy.
212
+ settings: Optional configuration override.
213
+
214
+ Returns:
215
+ Ranked search results with evidence spans.
216
+ """
217
+ cfg = settings or get_settings()
218
+ repo = ChunkRepository(session)
219
+
220
+ recall_k = cfg.search_top_k_recall
221
+
222
+ # Step 1: Retrieve
223
+ if mode == SearchMode.HYBRID:
224
+ query_embedding = await embedder.embed_query(query)
225
+ dense_results = await repo.search_dense(tenant_id, query_embedding, recall_k)
226
+ lexical_results = await repo.search_lexical(tenant_id, query, recall_k)
227
+
228
+ # Step 2: RRF fusion
229
+ fused = reciprocal_rank_fusion(
230
+ [dense_results, lexical_results],
231
+ [cfg.search_dense_weight, cfg.search_lexical_weight],
232
+ )
233
+ elif mode == SearchMode.SEMANTIC:
234
+ query_embedding = await embedder.embed_query(query)
235
+ fused = await repo.search_dense(tenant_id, query_embedding, recall_k)
236
+ else: # LEXICAL
237
+ fused = await repo.search_lexical(tenant_id, query, recall_k)
238
+
239
+ if not fused:
240
+ return SearchResult(items=[], count=0, query=query, mode=mode.value)
241
+
242
+ # Step 3: Rerank
243
+ rerank_input = [c.chunk.content for c in fused]
244
+ rerank_results = await reranker.rerank(
245
+ query, rerank_input, cfg.search_rerank_top_k
246
+ )
247
+
248
+ # Map rerank results back to chunks
249
+ reranked: list[ChunkWithScore] = []
250
+ for rr in rerank_results:
251
+ if rr.index < len(fused):
252
+ reranked.append(
253
+ ChunkWithScore(chunk=fused[rr.index].chunk, score=rr.score)
254
+ )
255
+
256
+ # Step 4: Group by candidate
257
+ candidates = group_by_candidate(reranked or fused)
258
+
259
+ # Step 5: Trim to top_k
260
+ candidates = candidates[:top_k]
261
+
262
+ logger.info(
263
+ "search_completed",
264
+ tenant_id=str(tenant_id),
265
+ query_length=len(query),
266
+ mode=mode.value,
267
+ candidates_returned=len(candidates),
268
+ total_chunks_recalled=len(fused),
269
+ )
270
+
271
+ return SearchResult(
272
+ items=candidates,
273
+ count=len(candidates),
274
+ query=query,
275
+ mode=mode.value,
276
+ )
277
+
278
+
279
+ async def search_similar(
280
+ *,
281
+ session: AsyncSession,
282
+ embedder: EmbeddingService,
283
+ reranker: RerankerService,
284
+ tenant_id: uuid.UUID,
285
+ document_id: uuid.UUID,
286
+ top_k: int = 10,
287
+ settings: Settings | None = None,
288
+ ) -> SearchResult:
289
+ """Find candidates similar to a given document ("more like this").
290
+
291
+ Embeds the candidate's chunks, aggregates them, and searches for similar
292
+ chunks across the corpus — excluding the source document itself.
293
+
294
+ Args:
295
+ session: Database session.
296
+ embedder: Embedding service.
297
+ reranker: Reranker service.
298
+ tenant_id: Owning tenant.
299
+ document_id: Source document to find similar candidates for.
300
+ top_k: Maximum candidates to return.
301
+ settings: Optional configuration override.
302
+
303
+ Returns:
304
+ Ranked search results for similar candidates.
305
+ """
306
+ cfg = settings or get_settings()
307
+ repo = ChunkRepository(session)
308
+
309
+ # Get the source document's chunks
310
+ source_chunks = await repo.get_by_document(
311
+ tenant_id, document_id, children_only=True
312
+ )
313
+
314
+ if not source_chunks:
315
+ return SearchResult(
316
+ items=[], count=0, query=f"similar:{document_id}", mode="similar"
317
+ )
318
+
319
+ # Use the source chunks' content to build a representative query
320
+ # Take the first few chunks to avoid exceeding embedding limits
321
+ representative_texts = [c.content for c in source_chunks[:5]]
322
+ query_text = " ".join(representative_texts)[:2000]
323
+
324
+ query_embedding = await embedder.embed_query(query_text)
325
+
326
+ # Search across all chunks, then filter out the source document
327
+ all_results = await repo.search_dense(
328
+ tenant_id, query_embedding, cfg.search_top_k_recall * 2
329
+ )
330
+
331
+ # Filter out chunks from the source document
332
+ filtered = [r for r in all_results if r.chunk.document_id != document_id]
333
+
334
+ # Group by candidate
335
+ candidates = group_by_candidate(filtered)[:top_k]
336
+
337
+ logger.info(
338
+ "similar_search_completed",
339
+ tenant_id=str(tenant_id),
340
+ source_document_id=str(document_id),
341
+ candidates_returned=len(candidates),
342
+ )
343
+
344
+ return SearchResult(
345
+ items=candidates,
346
+ count=len(candidates),
347
+ query=f"similar:{document_id}",
348
+ mode="similar",
349
+ )