Update OllamaProvider to use chat endpoint for consistency and fix README config
Browse files- core/llm.py +12 -5
core/llm.py
CHANGED
|
@@ -42,10 +42,11 @@ class OllamaProvider(LLMProvider):
|
|
| 42 |
|
| 43 |
def generate(self, prompt: str, max_tokens: int = 500, stream: bool = False) -> Union[str, Generator[str, None, None]]:
|
| 44 |
def _make_request():
|
| 45 |
-
|
|
|
|
| 46 |
payload = {
|
| 47 |
"model": self.model_name,
|
| 48 |
-
"
|
| 49 |
"stream": stream,
|
| 50 |
"options": {
|
| 51 |
"num_predict": max_tokens
|
|
@@ -69,13 +70,19 @@ class OllamaProvider(LLMProvider):
|
|
| 69 |
if line:
|
| 70 |
try:
|
| 71 |
data = json.loads(line.decode('utf-8'))
|
| 72 |
-
|
| 73 |
-
|
|
|
|
| 74 |
except:
|
| 75 |
continue
|
| 76 |
return stream_response()
|
| 77 |
else:
|
| 78 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 79 |
|
| 80 |
# Fixed: Moved return outside the _make_request function
|
| 81 |
return self._retry_request(_make_request)
|
|
|
|
| 42 |
|
| 43 |
def generate(self, prompt: str, max_tokens: int = 500, stream: bool = False) -> Union[str, Generator[str, None, None]]:
|
| 44 |
def _make_request():
|
| 45 |
+
# Use the chat endpoint instead of generate for better compatibility
|
| 46 |
+
url = f"{self.host}/api/chat"
|
| 47 |
payload = {
|
| 48 |
"model": self.model_name,
|
| 49 |
+
"messages": [{"role": "user", "content": prompt}],
|
| 50 |
"stream": stream,
|
| 51 |
"options": {
|
| 52 |
"num_predict": max_tokens
|
|
|
|
| 70 |
if line:
|
| 71 |
try:
|
| 72 |
data = json.loads(line.decode('utf-8'))
|
| 73 |
+
# Handle chat endpoint response format
|
| 74 |
+
if 'message' in data and 'content' in data['message']:
|
| 75 |
+
yield data['message']['content']
|
| 76 |
except:
|
| 77 |
continue
|
| 78 |
return stream_response()
|
| 79 |
else:
|
| 80 |
+
# Handle chat endpoint response format
|
| 81 |
+
data = response.json()
|
| 82 |
+
if 'message' in data and 'content' in data['message']:
|
| 83 |
+
return data['message']['content']
|
| 84 |
+
else:
|
| 85 |
+
raise Exception("Unexpected response format from Ollama")
|
| 86 |
|
| 87 |
# Fixed: Moved return outside the _make_request function
|
| 88 |
return self._retry_request(_make_request)
|