|
|
|
|
@ -22,6 +22,7 @@ from core.model_runtime.entities.message_entities import (
|
|
|
|
|
PromptMessageTool,
|
|
|
|
|
SystemPromptMessage,
|
|
|
|
|
TextPromptMessageContent,
|
|
|
|
|
ToolPromptMessage,
|
|
|
|
|
UserPromptMessage,
|
|
|
|
|
)
|
|
|
|
|
from core.model_runtime.entities.model_entities import (
|
|
|
|
|
@ -86,6 +87,7 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
|
|
|
|
|
credentials=credentials,
|
|
|
|
|
prompt_messages=prompt_messages,
|
|
|
|
|
model_parameters=model_parameters,
|
|
|
|
|
tools=tools,
|
|
|
|
|
stop=stop,
|
|
|
|
|
stream=stream,
|
|
|
|
|
user=user,
|
|
|
|
|
@ -153,6 +155,7 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
|
|
|
|
|
credentials: dict,
|
|
|
|
|
prompt_messages: list[PromptMessage],
|
|
|
|
|
model_parameters: dict,
|
|
|
|
|
tools: Optional[list[PromptMessageTool]] = None,
|
|
|
|
|
stop: Optional[list[str]] = None,
|
|
|
|
|
stream: bool = True,
|
|
|
|
|
user: Optional[str] = None,
|
|
|
|
|
@ -196,6 +199,8 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
|
|
|
|
|
if completion_type is LLMMode.CHAT:
|
|
|
|
|
endpoint_url = urljoin(endpoint_url, "api/chat")
|
|
|
|
|
data["messages"] = [self._convert_prompt_message_to_dict(m) for m in prompt_messages]
|
|
|
|
|
if tools:
|
|
|
|
|
data["tools"] = [self._convert_prompt_message_tool_to_dict(tool) for tool in tools]
|
|
|
|
|
else:
|
|
|
|
|
endpoint_url = urljoin(endpoint_url, "api/generate")
|
|
|
|
|
first_prompt_message = prompt_messages[0]
|
|
|
|
|
@ -232,7 +237,7 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
|
|
|
|
|
if stream:
|
|
|
|
|
return self._handle_generate_stream_response(model, credentials, completion_type, response, prompt_messages)
|
|
|
|
|
|
|
|
|
|
return self._handle_generate_response(model, credentials, completion_type, response, prompt_messages)
|
|
|
|
|
return self._handle_generate_response(model, credentials, completion_type, response, prompt_messages, tools)
|
|
|
|
|
|
|
|
|
|
def _handle_generate_response(
|
|
|
|
|
self,
|
|
|
|
|
@ -241,6 +246,7 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
|
|
|
|
|
completion_type: LLMMode,
|
|
|
|
|
response: requests.Response,
|
|
|
|
|
prompt_messages: list[PromptMessage],
|
|
|
|
|
tools: Optional[list[PromptMessageTool]],
|
|
|
|
|
) -> LLMResult:
|
|
|
|
|
"""
|
|
|
|
|
Handle llm completion response
|
|
|
|
|
@ -253,14 +259,16 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
|
|
|
|
|
:return: llm result
|
|
|
|
|
"""
|
|
|
|
|
response_json = response.json()
|
|
|
|
|
|
|
|
|
|
tool_calls = []
|
|
|
|
|
if completion_type is LLMMode.CHAT:
|
|
|
|
|
message = response_json.get("message", {})
|
|
|
|
|
response_content = message.get("content", "")
|
|
|
|
|
response_tool_calls = message.get("tool_calls", [])
|
|
|
|
|
tool_calls = [self._extract_response_tool_call(tool_call) for tool_call in response_tool_calls]
|
|
|
|
|
else:
|
|
|
|
|
response_content = response_json["response"]
|
|
|
|
|
|
|
|
|
|
assistant_message = AssistantPromptMessage(content=response_content)
|
|
|
|
|
assistant_message = AssistantPromptMessage(content=response_content, tool_calls=tool_calls)
|
|
|
|
|
|
|
|
|
|
if "prompt_eval_count" in response_json and "eval_count" in response_json:
|
|
|
|
|
# transform usage
|
|
|
|
|
@ -405,9 +413,28 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
|
|
|
|
|
|
|
|
|
|
chunk_index += 1
|
|
|
|
|
|
|
|
|
|
def _convert_prompt_message_tool_to_dict(self, tool: PromptMessageTool) -> dict:
|
|
|
|
|
"""
|
|
|
|
|
Convert PromptMessageTool to dict for Ollama API
|
|
|
|
|
|
|
|
|
|
:param tool: tool
|
|
|
|
|
:return: tool dict
|
|
|
|
|
"""
|
|
|
|
|
return {
|
|
|
|
|
"type": "function",
|
|
|
|
|
"function": {
|
|
|
|
|
"name": tool.name,
|
|
|
|
|
"description": tool.description,
|
|
|
|
|
"parameters": tool.parameters,
|
|
|
|
|
},
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
def _convert_prompt_message_to_dict(self, message: PromptMessage) -> dict:
|
|
|
|
|
"""
|
|
|
|
|
Convert PromptMessage to dict for Ollama API
|
|
|
|
|
|
|
|
|
|
:param message: prompt message
|
|
|
|
|
:return: message dict
|
|
|
|
|
"""
|
|
|
|
|
if isinstance(message, UserPromptMessage):
|
|
|
|
|
message = cast(UserPromptMessage, message)
|
|
|
|
|
@ -432,6 +459,9 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
|
|
|
|
|
elif isinstance(message, SystemPromptMessage):
|
|
|
|
|
message = cast(SystemPromptMessage, message)
|
|
|
|
|
message_dict = {"role": "system", "content": message.content}
|
|
|
|
|
elif isinstance(message, ToolPromptMessage):
|
|
|
|
|
message = cast(ToolPromptMessage, message)
|
|
|
|
|
message_dict = {"role": "tool", "content": message.content}
|
|
|
|
|
else:
|
|
|
|
|
raise ValueError(f"Got unknown type {message}")
|
|
|
|
|
|
|
|
|
|
@ -452,6 +482,29 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
|
|
|
|
|
|
|
|
|
|
return num_tokens
|
|
|
|
|
|
|
|
|
|
def _extract_response_tool_call(self, response_tool_call: dict) -> AssistantPromptMessage.ToolCall:
|
|
|
|
|
"""
|
|
|
|
|
Extract response tool call
|
|
|
|
|
"""
|
|
|
|
|
tool_call = None
|
|
|
|
|
if response_tool_call and "function" in response_tool_call:
|
|
|
|
|
# Convert arguments to JSON string if it's a dict
|
|
|
|
|
arguments = response_tool_call.get("function").get("arguments")
|
|
|
|
|
if isinstance(arguments, dict):
|
|
|
|
|
arguments = json.dumps(arguments)
|
|
|
|
|
|
|
|
|
|
function = AssistantPromptMessage.ToolCall.ToolCallFunction(
|
|
|
|
|
name=response_tool_call.get("function").get("name"),
|
|
|
|
|
arguments=arguments,
|
|
|
|
|
)
|
|
|
|
|
tool_call = AssistantPromptMessage.ToolCall(
|
|
|
|
|
id=response_tool_call.get("function").get("name"),
|
|
|
|
|
type="function",
|
|
|
|
|
function=function,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
return tool_call
|
|
|
|
|
|
|
|
|
|
def get_customizable_model_schema(self, model: str, credentials: dict) -> AIModelEntity:
|
|
|
|
|
"""
|
|
|
|
|
Get customizable model schema.
|
|
|
|
|
@ -461,10 +514,15 @@ class OllamaLargeLanguageModel(LargeLanguageModel):
|
|
|
|
|
|
|
|
|
|
:return: model schema
|
|
|
|
|
"""
|
|
|
|
|
extras = {}
|
|
|
|
|
extras = {
|
|
|
|
|
"features": [],
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
if "vision_support" in credentials and credentials["vision_support"] == "true":
|
|
|
|
|
extras["features"] = [ModelFeature.VISION]
|
|
|
|
|
extras["features"].append(ModelFeature.VISION)
|
|
|
|
|
if "function_call_support" in credentials and credentials["function_call_support"] == "true":
|
|
|
|
|
extras["features"].append(ModelFeature.TOOL_CALL)
|
|
|
|
|
extras["features"].append(ModelFeature.MULTI_TOOL_CALL)
|
|
|
|
|
|
|
|
|
|
entity = AIModelEntity(
|
|
|
|
|
model=model,
|
|
|
|
|
|