Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use https://registry.ollama.ai/v2 API to check model staleness #190

Merged
merged 1 commit into from
Jan 7, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
115 changes: 1 addition & 114 deletions package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

3 changes: 1 addition & 2 deletions package.json
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,6 @@
},
"dependencies": {
"@redhat-developer/vscode-redhat-telemetry": "^0.9.1",
"node-html-parser": "^6.1.13",
"systeminformation": "^5.23.24"
}
}
}
56 changes: 29 additions & 27 deletions src/ollama/ollamaLibrary.ts
Original file line number Diff line number Diff line change
@@ -1,46 +1,42 @@
import { parse } from 'node-html-parser';
import crypto from 'crypto';
import { ModelInfo } from '../commons/modelInfo';
import { formatSize } from '../commons/textUtils';

const cache = new Map<string, ModelInfo | undefined>();//TODO limit caching lifespan

const INFO_DELIMITER = ' · ';

// This is fugly, extremely brittle, but we have no other choice because The ollama library doesn't seem to expose an API we can query.
export async function getRemoteModelInfo(modelId: string): Promise<ModelInfo | undefined> {
// Check if the result is already cached
if (cache.has(modelId)) {
return cache.get(modelId);
}
const start = Date.now();

const url = `https://ollama.com/library/${modelId}`;
const [modelName, tag] = modelId.split(":");
const url = `https://registry.ollama.ai/v2/library/${modelName}/manifests/${tag}`;
try {
const response = await fetch(url, { signal: AbortSignal.timeout(3000) });

if (!response.ok) {
throw new Error(`Failed to fetch the model page: ${response.statusText}`);
}
const html = await response.text();
const root = parse(html);
const fileExplorer = root.querySelector('#file-explorer');
const itemsCenter = fileExplorer?.querySelector('.items-center');
const lastParagraphElement = itemsCenter?.querySelectorAll('p')?.pop();

if (lastParagraphElement) {
const lastParagraph = lastParagraphElement.text.trim();
if (lastParagraph.includes(INFO_DELIMITER)) {
const [digest, size] = lastParagraph.split(INFO_DELIMITER).map(item => item.trim());
const data: ModelInfo = {
id: modelId,
size,
digest
};
// Cache the successful result
cache.set(modelId, data);
console.log('Model info:', data);
return data;
}
}

// First, read the response body as an ArrayBuffer to compute the digest
const buffer = await response.arrayBuffer();
const digest = getDigest(buffer);

// Then, decode the ArrayBuffer into a string and parse it as JSON
const text = new TextDecoder().decode(buffer);
const manifest = JSON.parse(text) as { layers: { size: number }[] };
const modelSize = manifest.layers.reduce((sum, layer) => sum + layer.size, 0);

const data: ModelInfo = {
id: modelId,
size: formatSize(modelSize),
digest
};
// Cache the successful result
cache.set(modelId, data);
console.log('Model info:', data);
return data;
} catch (error) {
console.error(`Error fetching or parsing model info: ${error}`);
} finally {
Expand All @@ -51,4 +47,10 @@ export async function getRemoteModelInfo(modelId: string): Promise<ModelInfo | u
// Cache the failure
cache.set(modelId, undefined);
return undefined;
}

function getDigest(buffer: ArrayBuffer): string {
const hash = crypto.createHash('sha256');
hash.update(Buffer.from(buffer));
return hash.digest('hex');
}
Loading