返回提交历史
Modified
g4f/Provider/local/Ollama.py
+53
-0
XFEstudio/gpt4free
feat: Add get_quota method to Ollama provider
c2da72c3
代码差异
1 个文件
+53
-0
@@ -1,5 +1,6 @@
1
1
from __future__ import annotations
2
2
3
import re
3
4
import json
4
5
import requests
5
6
import os
@@ -8,9 +9,11 @@ from typing import Optional
8
9
from ..template import OpenaiTemplate
9
10
from ...requests import StreamSession, raise_for_status
10
11
from ...providers.response import Usage, Reasoning
12
from ...cookies import get_cookies
11
13
from ...tools.run_tools import AuthManager
12
14
from ...typing import AsyncResult, Messages
13
15
from ...config import AppConfig
16
from ... import debug
14
17
15
18
class Ollama(OpenaiTemplate):
16
19
label = "Ollama 🦙"
@@ -26,6 +29,56 @@ class Ollama(OpenaiTemplate):
26
29
"gpt-oss-20b": "gpt-oss:20b"
27
30
}
28
31
32
@classmethod
33
async def get_quota(cls, api_key: Optional[str] = None) -> Optional[dict]:
34
cookies = {}
35
if api_key:
36
cookies = {"__Secure-session": api_key}
37
else:
38
api_key = AuthManager.load_api_key(cls)
39
if api_key:
40
cookies = {"__Secure-session": api_key}
41
else:
42
cookies = get_cookies("ollama.com", raise_requirements_error=False)
43
if not cookies:
44
return None
45
try:
46
async with StreamSession() as session:
47
async with session.get(
48
"https://ollama.com/settings",
49
cookies=cookies,
50
headers={"User-Agent": "Mozilla/5.0"},
51
) as response:
52
await raise_for_status(response)
53
html = await response.text()
54
if "sign in" in html.lower() or "/signin" in html.lower():
55
debug.error("Ollama session cookie is invalid or expired")
56
return None
57
quota = {}
58
for label in ("Session usage", "Hourly usage", "Weekly usage"):
59
idx = html.find(label)
60
if idx == -1:
61
continue
62
section = html[idx:idx + 500]
63
pct_match = re.search(r'(\d+(?:\.\d+)?)%\s*used', section)
64
if not pct_match:
65
width_match = re.search(r'width:\s*(\d+(?:\.\d+)?)%', section)
66
if width_match:
67
pct_match = width_match
68
pct = float(pct_match.group(1)) if pct_match else None
69
reset_match = re.search(r'data-time=["\']([^"\']+)["\']', section)
70
reset_time = reset_match.group(1) if reset_match else None
71
key = label.lower().replace(" ", "_")
72
quota[key] = {
73
"used_percent": pct,
74
"reset_time": reset_time,
75
}
76
if quota:
77
return quota
78
except Exception as e:
79
debug.error(f"Failed to get Ollama quota:", e)
80
return None
81
29
82
@classmethod
30
83
def get_models(cls, api_key: str = None, base_url: str = None, **kwargs):
31
84
if not cls.models: