File size: 9,148 Bytes
5583f46
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2f1e7ff
5583f46
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2f1e7ff
5583f46
 
2f1e7ff
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5583f46
2f1e7ff
 
 
 
 
 
 
 
 
5583f46
2f1e7ff
 
 
 
 
5583f46
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
"""Run through entrypoint.sh --qualify inside a local pre-production image."""

import json
import math
import os
from pathlib import Path
import sys
import time
from uuid import uuid4

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import httpx

from src.config import config
from src.constants.embeddings import EMBEDDING_STORES_SIZES
from src.lib.embeddings.client_ollama import OllamaEmbeddingsClient
from src.lib.embeddings.profile import EmbeddingProfileError
from src.lib.postgresql._provisioning import borrow, get_brain_pool
from src.lib.postgresql.vectors import PostgreSQLVectorStoreClient
from src.services.data.main import data_adapter


def check(condition, message):
    if not condition:
        raise RuntimeError(message)


def qualify():
    check(config.ollama.embedding_model == 'embeddinggemma-2:270m', 'Wrong embedding model')
    check(config.pipeline_mode == 'lightweight', 'Qualification requires lightweight ingestion')
    check(config.search_enabled and config.search_use_dense and config.search_use_bm25,
          'Qualification requires dense + BM25 fusion')
    check(set(EMBEDDING_STORES_SIZES.values()) == {512}, 'All stores must use 512 dimensions')
    token = config.system_pat
    api_url = os.getenv('QUALIFY_API_URL', 'http://127.0.0.1:7860')
    model_url = f'http://{config.ollama.host}:{config.ollama.port}'
    with httpx.Client(base_url=model_url, timeout=120, trust_env=False) as ollama:
        response = ollama.get('/api/tags')
        response.raise_for_status()
        tags = {item['name'] for item in response.json()['models']}
        check(tags == {'embeddinggemma-2:270m'}, f'Unexpected baked models: {sorted(tags)}')
        response = ollama.post('/api/embed', json={
            'model': config.ollama.embedding_model, 'input': 'task: search result | query: launch code',
            'truncate': False, 'options': {'num_ctx': 8192, 'num_thread': 2}})
        response.raise_for_status()
        raw = response.json()['embeddings'][0]
        check(len(raw) == 768 and all(math.isfinite(x) for x in raw), 'Invalid native embedding')
        response = ollama.post('/api/embed', json={
            'model': config.ollama.embedding_model,
            'input': 'task: search result | query: ' + 'hello ' * 9000,
            'truncate': False, 'options': {'num_ctx': 8192, 'num_thread': 2}})
        check(response.status_code == 400 and 'context' in response.json().get('error', '').lower(),
              'Ollama did not enforce the 8192-token input limit')
    client = OllamaEmbeddingsClient()
    try:
        for vector in (client.embed_text('The launch code is violet.'), client.embed_query('launch code')):
            check(len(vector) == 512 and all(math.isfinite(x) for x in vector), 'Invalid MRL vector')
            check(abs(math.hypot(*vector) - 1) < 1e-6, 'MRL vector is not normalized')
    finally:
        client.client.close()
    store = PostgreSQLVectorStoreClient()
    marker = 'violet-' + uuid4().hex
    text = f'The Orchid qualification launch code is {marker}.'
    workspaces = []
    issued_ids = []
    with httpx.Client(base_url=api_url, timeout=120, trust_env=False, follow_redirects=True) as api:
        try:
            for suffix in ('a', 'b'):
                slug = 'gemma-qa-' + uuid4().hex[:12] + '-' + suffix
                response = api.post('/system/workspaces', headers={'BrainPAT': token}, json={
                    'slug': slug, 'display_name': 'Gemma qualification ' + suffix})
                response.raise_for_status()
                brain_id = response.json()['brain_id']
                issued = api.post('/system/access-tokens', headers={'Authorization': 'Bearer ' + token}, json={
                    'brain_id': brain_id, 'label': 'embedding-qualification', 'permissions': ['memory:read', 'memory:write']})
                issued.raise_for_status()
                pat = issued.json()['token']
                issued_ids.append(issued.json()['id'])
                workspaces.append((slug, brain_id, pat))
            slug, brain_id, pat = workspaces[0]
            root = f'/brains/{slug}/api'
            headers = {'BrainPAT': pat}
            response = api.post(root + '/ingest/', headers=headers, json={
                'data': {'data_type': 'text', 'text_data': text}})
            response.raise_for_status()
            check(response.status_code == 202, 'Ingestion was not accepted')
            from src.services.task_contract import accepted_task_id
            task_id = accepted_task_id(response.status_code, response.json())
            deadline = time.monotonic() + 120
            while True:
                response = api.get(root + '/tasks/' + task_id, headers=headers)
                response.raise_for_status()
                state = response.json()
                check(state['status'] != 'failed', 'Worker failed: ' + str(state.get('error')))
                if state['status'] == 'completed':
                    break
                check(time.monotonic() < deadline, 'Worker did not complete before deadline')
                time.sleep(1)
            for name in EMBEDDING_STORES_SIZES:
                store._ensure_store(name, brain_id)
            with borrow(get_brain_pool(brain_id)) as connection:
                with connection.cursor() as cursor:
                    cursor.execute('SELECT vector_dims(embeddings), metadata FROM vectors_data')
                    rows = cursor.fetchall()
                    check(rows and all(row[0] == 512 for row in rows), 'Persisted vectors are not 512d')
                    for name in EMBEDDING_STORES_SIZES:
                        cursor.execute("SELECT atttypmod FROM pg_attribute WHERE attrelid = %s::regclass AND attname = 'embeddings'", ('vectors_' + name,))
                        check(cursor.fetchone()[0] == 512, 'Incorrect schema dimension: ' + name)
                    cursor.execute('SELECT profile FROM embedding_profiles')
                    profiles = [row[0] for row in cursor.fetchall()]
                    check(len(profiles) == 5 and all(p['model'] == 'embeddinggemma-2:270m' for p in profiles), 'Index profile not recorded')
            requests = [('search', {'query': 'Orchid qualification launch code', 'rerank': 'none'}),
                        ('context', {'text': 'Orchid qualification launch code', 'apply_fact_filter': False})]
            for endpoint, payload in requests:
                response = api.post(root + '/retrieve/' + endpoint, headers=headers, json=payload)
                response.raise_for_status()
                body = response.json()
                if endpoint == 'search':
                    check(any(marker in hit['snippet'] for hit in body['hits']), 'Search did not retrieve the unique fact')
                else:
                    check(marker in json.dumps(body), 'Context did not include the unique fact')
                other_slug, _, other_pat = workspaces[1]
                other = api.post(f'/brains/{other_slug}/api/retrieve/' + endpoint,
                                 headers={'BrainPAT': other_pat}, json=payload)
                other.raise_for_status()
                check(marker not in json.dumps(other.json()), 'Workspace data leaked')
                denied = api.post(f'/brains/{other_slug}/api/retrieve/' + endpoint,
                                  headers=headers, json=payload)
                check(denied.status_code == 403, 'Cross-workspace PAT was not denied')
            # Even a same-dimension model identity change must be rejected on a fresh client.
            previous = config.ollama.embedding_model
            try:
                config.ollama.embedding_model = 'qwen3-embedding:0.6b'
                try:
                    PostgreSQLVectorStoreClient()._ensure_store('data', brain_id)
                except EmbeddingProfileError:
                    pass
                else:
                    raise RuntimeError('A same-dimension model identity mismatch was accepted')
            finally:
                config.ollama.embedding_model = previous
        finally:
            for token_id in issued_ids:
                api.delete('/system/access-tokens/' + token_id, headers={'Authorization': 'Bearer ' + token}).raise_for_status()
            for slug, brain_id, _ in workspaces:
                api.delete('/system/brains/' + brain_id, headers={'Authorization': 'Bearer ' + token}).raise_for_status()
                api.patch('/system/workspaces/' + slug, json={'archived': True}, headers={'Authorization': 'Bearer ' + token}).raise_for_status()
    print(json.dumps({'result': 'passed', 'model': config.ollama.embedding_model,
                      'raw_dimension': 768, 'stored_dimension': 512, 'context_limit': 8192,
                      'schema_dimensions': dict(EMBEDDING_STORES_SIZES),
                      'worker': 'completed', 'search': 'unique fact retrieved', 'context': 'unique fact present',
                      'workspace_isolation': 'passed', 'model_mismatch': 'rejected',
                      'qualification_workspaces': [w[0] for w in workspaces]}))


if __name__ == '__main__':
    qualify()