rdune71 commited on
Commit
fdf2492
·
1 Parent(s): 860bf55

Update OllamaProvider to use chat endpoint for consistency and fix README config

Browse files
Files changed (1) hide show
  1. 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
- url = f"{self.host}/api/generate"
 
46
  payload = {
47
  "model": self.model_name,
48
- "prompt": prompt,
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
- if 'response' in data:
73
- yield data['response']
 
74
  except:
75
  continue
76
  return stream_response()
77
  else:
78
- return response.json()["response"]
 
 
 
 
 
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)