XFE Git
XFE Studio Git
Git 首页 全局搜索
XFE 主站 文档 NuGet
公开
关注 0 Fork 0 Star 0
返回提交历史

XFEstudio/gpt4free

fix: update Ollama class to improve model retrieval and refactor base_url handling

1448f7ad
hlohaus <983577+hlohaus@users.noreply.github.com>
提交于

代码差异

2 个文件 +22 -37
Modified g4f/Provider/local/Ollama.py +16 -35
@@ -3,6 +3,7 @@ from __future__ import annotations
3 3 import json
4 4 import requests
5 5 import os
6 from typing import Optional
6 7
7 8 from ..template import OpenaiTemplate
8 9 from ...requests import StreamSession, raise_for_status
@@ -59,42 +60,22 @@ class Ollama(OpenaiTemplate):
59 60 messages: Messages,
60 61 api_key: str = None,
61 62 base_url: str = None,
62 proxy: str = None,
63 63 **kwargs
64 64 ) -> AsyncResult:
65 if base_url is None:
66 host = os.getenv("OLLAMA_HOST", "localhost")
67 port = os.getenv("OLLAMA_PORT", "11434")
68 base_url: str = f"http://{host}:{port}/v1"
65 if not cls.models:
66 cls.get_models(api_key=api_key, base_url=base_url)
69 67 if model in cls.local_models:
70 async with StreamSession(headers={"Authorization": f"Bearer {api_key}"}, proxy=proxy) as session:
71 async with session.post(f"{base_url.replace('/v1', '')}/api/chat", json={
72 "model": model,
73 "messages": messages,
74 }) as response:
75 await raise_for_status(response)
76 last_data = {}
77 async for chunk in response.iter_lines():
78 data = json.loads(chunk)
79 last_data = data
80 thinking = data.get("message", {}).get("thinking", "")
81 if thinking:
82 yield Reasoning(thinking)
83 content = data.get("message", {}).get("content", "")
84 if content:
85 yield content
86 yield Usage(
87 prompt_tokens=last_data.get("prompt_eval_count", 0),
88 completion_tokens=last_data.get("eval_count", 0),
89 total_tokens=last_data.get("prompt_eval_count", 0) + last_data.get("eval_count", 0),
90 )
68 if base_url is None:
69 host = os.getenv("OLLAMA_HOST", "localhost")
70 port = os.getenv("OLLAMA_PORT", "11434")
71 base_url: str = f"http://{host}:{port}/v1"
91 72 else:
92 async for chunk in super().create_async_generator(
93 model,
94 messages,
95 api_key=api_key,
96 base_url=cls.backup_url,
97 proxy=proxy,
98 **kwargs
99 ):
100 yield chunk
73 base_url = cls.backup_url
74 async for chunk in super().create_async_generator(
75 model,
76 messages,
77 api_key=api_key,
78 base_url=cls.backup_url,
79 **kwargs
80 ):
81 yield chunk
Modified g4f/client/stubs.py +6 -2
@@ -66,14 +66,18 @@ class ToolFunctionModel(BaseModel):
66 66 arguments: str
67 67
68 68 class ToolCallModel(BaseModel):
69 index: int
69 index: int = 0
70 70 id: str
71 71 type: str
72 72 function: ToolFunctionModel
73 73
74 74 @classmethod
75 def model_construct(cls, function=None, **kwargs):
75 def model_construct(cls, function=None, index=0, **kwargs):
76 # Ensure arguments is always a string
77 if function and "arguments" in function and not isinstance(function["arguments"], str):
78 function["arguments"] = str(function["arguments"])
76 79 return super().model_construct(
80 index=index,
77 81 **kwargs,
78 82 function=ToolFunctionModel.model_construct(**function),
79 83 )