Mistral Forward Deployed Engineer Interview Guide

1from __future__ import annotations
2
3import asyncio
4import json
5import math
6import re
7import uuid
8from dataclasses import dataclass
9from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple
10
11from fastapi import FastAPI, HTTPException
12from fastapi.responses import StreamingResponse
13from pydantic import BaseModel, Field, conint
14
15app = FastAPI(title="Mistral FDE RAG Demo")
16
17
18# -----------------------------
19# Data model
20# -----------------------------
21class RagAnswerRequest(BaseModel):
22    query: str = Field(..., min_length=1, max_length=2000)
23    documents: List[str] = Field(..., min_items=1, max_items=50)
24    k: conint(ge=1, le=10) = 4
25    max_context_chars: conint(ge=200, le=20000) = 6000
26
27
28@dataclass(frozen=True)
29class Chunk:
30    doc_id: int
31    chunk_id: int
32    text: str
33
34
35# -----------------------------
36# Deterministic chunking
37# -----------------------------
38_WS_RE = re.compile(r"\s+")
39
40
41def normalize_ws(text: str) -> str:
42    return _WS_RE.sub(" ", text).strip()
43
44
45def chunk_text(text: str, chunk_size: int = 500, overlap: int = 100) -> List[str]:
46    """Deterministic, character-based chunking with overlap.
47
48    This is intentionally simple and stable for interview purposes.
49    """
50    text = normalize_ws(text)
51    if not text:
52        return []
53
54    if overlap >= chunk_size:
55        raise ValueError("overlap must be smaller than chunk_size")
56
57    chunks: List[str] = []
58    start = 0
59    n = len(text)
60    while start < n:
61        end = min(n, start + chunk_size)
62        chunk = text[start:end]
63        chunks.append(chunk)
64        if end == n:
65            break
66        start = end - overlap
67    return chunks
68
69
70# -----------------------------
71# Tiny TF-IDF + cosine retrieval
72# -----------------------------
73_TOKEN_RE = re.compile(r"[A-Za-z0-9_]+")
74
75
76def tokenize(text: str) -> List[str]:
77    return [t.lower() for t in _TOKEN_RE.findall(text)]
78
79
80def build_tfidf_index(chunks: List[Chunk]) -> Tuple[Dict[str, float], List[Dict[str, float]]]:
81    """Returns (idf, tfidf_vectors_per_chunk)."""
82    token_lists = [tokenize(c.text) for c in chunks]
83    df: Dict[str, int] = {}
84    for toks in token_lists:
85        for tok in set(toks):
86            df[tok] = df.get(tok, 0) + 1
87
88    n_docs = len(chunks)
89    idf: Dict[str, float] = {}
90    for tok, dfi in df.items():
91        # Smoothed IDF: log((N + 1) / (df + 1)) + 1
92        idf[tok] = math.log((n_docs + 1.0) / (dfi + 1.0)) + 1.0
93
94    vectors: List[Dict[str, float]] = []
95    for toks in token_lists:
96        tf: Dict[str, int] = {}
97        for tok in toks:
98            tf[tok] = tf.get(tok, 0) + 1
99        # L2-normalized tf-idf
100        vec: Dict[str, float] = {}
101        norm2 = 0.0
102        for tok, cnt in tf.items():
103            val = float(cnt) * idf.get(tok, 0.0)
104            if val:
105                vec[tok] = val
106                norm2 += val * val
107        norm = math.sqrt(norm2) if norm2 > 0 else 1.0
108        for tok in list(vec.keys()):
109            vec[tok] /= norm
110        vectors.append(vec)
111
112    return idf, vectors
113
114
115def tfidf_vector(text: str, idf: Dict[str, float]) -> Dict[str, float]:
116    toks = tokenize(text)
117    tf: Dict[str, int] = {}
118    for tok in toks:
119        tf[tok] = tf.get(tok, 0) + 1
120    vec: Dict[str, float] = {}
121    norm2 = 0.0
122    for tok, cnt in tf.items():
123        if tok not in idf:
124            continue
125        val = float(cnt) * idf[tok]
126        vec[tok] = val
127        norm2 += val * val
128    norm = math.sqrt(norm2) if norm2 > 0 else 1.0
129    for tok in list(vec.keys()):
130        vec[tok] /= norm
131    return vec
132
133
134def cosine_sparse(a: Dict[str, float], b: Dict[str, float]) -> float:
135    if not a or not b:
136        return 0.0
137    # Iterate over smaller dict
138    if len(a) > len(b):
139        a, b = b, a
140    s = 0.0
141    for k, va in a.items():
142        vb = b.get(k)
143        if vb is not None:
144            s += va * vb
145    return s
146
147
148def top_k_chunks(query: str, chunks: List[Chunk], k: int) -> List[Tuple[Chunk, float]]:
149    idf, vectors = build_tfidf_index(chunks)
150    qv = tfidf_vector(query, idf)
151    scored: List[Tuple[int, float]] = []
152    for i, cv in enumerate(vectors):
153        scored.append((i, cosine_sparse(qv, cv)))
154    scored.sort(key=lambda x: x[1], reverse=True)
155    top = scored[:k]
156    return [(chunks[i], score) for i, score in top]
157
158
159# -----------------------------
160# Mistral streaming stub
161# -----------------------------
162async def mistral_stream_chat_completion(prompt: str) -> AsyncGenerator[str, None]:
163    """Stubbed token stream.
164
165    In production you would call Mistral's SDK with stream=True and yield deltas.
166    """
167    # Simulate tokens by splitting on whitespace
168    for tok in prompt.split():
169        await asyncio.sleep(0.005)
170        yield tok + " "
171
172
173# -----------------------------
174# SSE helpers
175# -----------------------------
176
177def sse_event(event: str, data: Dict[str, Any]) -> str:
178    payload = json.dumps(data, ensure_ascii=False)
179    # SSE format: event + data lines, then a blank line
180    return f"event: {event}\n" + f"data: {payload}\n\n"
181
182
183def build_prompt(query: str, retrieved: List[Tuple[Chunk, float]], max_context_chars: int) -> str:
184    context_parts: List[str] = []
185    used = 0
186    for chunk, score in retrieved:
187        snippet = chunk.text
188        block = f"[doc:{chunk.doc_id} chunk:{chunk.chunk_id} score:{score:.3f}] {snippet}"
189        if used + len(block) + 1 > max_context_chars:
190            break
191        context_parts.append(block)
192        used += len(block) + 1
193
194    context = "\n".join(context_parts)
195    return (
196        "You are a helpful assistant. Answer using only the context. "
197        "If insufficient, say you do not know. Cite sources as [doc:X chunk:Y].\n\n"
198        f"Context:\n{context}\n\n"
199        f"Question: {query}\nAnswer:"
200    )
201
202
203@app.post("/v1/rag/answer")
204async def rag_answer(req: RagAnswerRequest):
205    # Validate documents are not empty after normalization
206    normalized_docs = [normalize_ws(d) for d in req.documents]
207    if any(not d for d in normalized_docs):
208        raise HTTPException(status_code=400, detail="documents contains an empty string")
209
210    # Build chunks deterministically
211    chunks: List[Chunk] = []
212    for doc_id, doc in enumerate(normalized_docs):
213        parts = chunk_text(doc, chunk_size=500, overlap=100)
214        for chunk_id, part in enumerate(parts):
215            chunks.append(Chunk(doc_id=doc_id, chunk_id=chunk_id, text=part))
216
217    if not chunks:
218        raise HTTPException(status_code=400, detail="no chunks could be created")
219
220    retrieved = top_k_chunks(req.query, chunks, req.k)
221    prompt = build_prompt(req.query, retrieved, req.max_context_chars)
222
223    trace_id = str(uuid.uuid4())
224
225    async def event_stream() -> AsyncGenerator[bytes, None]:
226        # Start event
227        yield sse_event("start", {"trace_id": trace_id}).encode("utf-8")
228
229        # Stream tokens
230        async for token in mistral_stream_chat_completion(prompt):
231            yield sse_event("token", {"trace_id": trace_id, "text": token}).encode("utf-8")
232
233        # End event with retrieval metadata for debuggability
234        sources = [
235            {
236                "doc_id": c.doc_id,
237                "chunk_id": c.chunk_id,
238                "score": float(score),
239                "preview": c.text[:120],
240            }
241            for c, score in retrieved
242        ]
243        yield sse_event("end", {"trace_id": trace_id, "sources": sources}).encode("utf-8")
244
245    return StreamingResponse(event_stream(), media_type="text/event-stream")
246
247
248# For local testing:
249# uvicorn this_file:app --reload
250

1-- Parameters:
2--   :tenant_id           uuid
3--   :user_id             uuid
4--   :query_embedding     vector(1536)
5--   :match_count         int
6
7SELECT
8  c.id AS chunk_id,
9  c.document_id,
10  c.chunk_index,
11  c.content,
12  (c.embedding <-> :query_embedding) AS distance
13FROM rag_chunks AS c
14JOIN rag_documents AS d
15  ON d.id = c.document_id
16 AND d.tenant_id = :tenant_id
17 AND d.deleted_at IS NULL
18JOIN rag_document_acl AS a
19  ON a.document_id = d.id
20 AND a.user_id = :user_id
21 AND a.can_read = TRUE
22WHERE c.tenant_id = :tenant_id
23  AND c.deleted_at IS NULL
24ORDER BY c.embedding <-> :query_embedding
25LIMIT COALESCE(:match_count, 10);