Spaces:
Running
Running
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()
|