323 lines
12 KiB
Python
323 lines
12 KiB
Python
import json
|
|
import uuid
|
|
from typing import (
|
|
Any,
|
|
Callable,
|
|
Dict,
|
|
List,
|
|
Optional,
|
|
Sequence,
|
|
Type,
|
|
TypeVar,
|
|
Union,
|
|
Tuple,
|
|
)
|
|
from types import NoneType
|
|
|
|
from langchain_ollama.chat_models import ChatOllama
|
|
from langchain_core.callbacks import CallbackManagerForLLMRun
|
|
from langchain_core.language_models import LanguageModelInput
|
|
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage, BaseMessage, ToolCall
|
|
from langchain_core.outputs import ChatGeneration, ChatResult
|
|
from langchain_core.prompts import SystemMessagePromptTemplate
|
|
from langchain_core.pydantic_v1 import BaseModel
|
|
from langchain_core.runnables import Runnable
|
|
from langchain_core.tools import BaseTool, Tool
|
|
from langchain_core.utils.pydantic import is_basemodel_instance, is_basemodel_subclass
|
|
|
|
from libs.functions import nxhash
|
|
|
|
DEFAULT_SYTEM_PROMPT = """You have access to the following tools:
|
|
|
|
{tools}
|
|
|
|
You must always select one of the above tools and respond with only a JSON object matching the following schema:
|
|
|
|
{{
|
|
"tool": <name of the selected tool>,
|
|
"tool_input": <parameters for the selected tool, matching the tool's JSON schema>
|
|
}}
|
|
"""
|
|
|
|
|
|
DEFAULT_SYTEM_PROMPT_WITH_HISTORY = """{system_msg}
|
|
|
|
You continue a chat history either conversationally or with a tool call.
|
|
|
|
You have access to the following tools:
|
|
|
|
{tools}
|
|
|
|
You must either select one of the above tools and respond with only a JSON object matching the following schema:
|
|
|
|
{{
|
|
"tool": <name of the selected tool>,
|
|
"tool_input": <parameters for the selected tool, matching the tool's JSON schema>
|
|
}}
|
|
|
|
or answer conversationally normally.
|
|
|
|
The conversation before consisted of the following messages:
|
|
|
|
{history}
|
|
|
|
Now you must answer accordingly either conversationally or with another tool call.
|
|
|
|
For conversational answers: Answer as if it was a continuous conversation. The Human only sees the conversational responses, and not anything about the tools. Do not mention the tools or the process of using them.
|
|
"""
|
|
CONVERSATIONAL_RESPONSE_TOOL = {
|
|
"name": "__conversational_response",
|
|
"description": (
|
|
"Respond conversationally if no other tools should be called for a given query."
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"response": {
|
|
"type": "string",
|
|
"description": "Conversational response to the user.",
|
|
},
|
|
},
|
|
"required": ["response"],
|
|
},
|
|
}
|
|
|
|
|
|
_BM = TypeVar("_BM", bound=BaseModel)
|
|
_DictOrPydantic = Union[Dict, _BM]
|
|
|
|
|
|
def _is_pydantic_class(obj: Any) -> bool:
|
|
return isinstance(obj, type) and (
|
|
is_basemodel_subclass(obj) or BaseModel in obj.__bases__
|
|
)
|
|
|
|
|
|
def convert_to_ollama_tool(tool: Any) -> Dict:
|
|
"""Convert a tool to an Ollama tool."""
|
|
description = None
|
|
if _is_pydantic_class(tool):
|
|
schema = tool.construct().schema()
|
|
name = schema["title"]
|
|
elif isinstance(tool, BaseTool):
|
|
schema = tool.tool_call_schema.schema()
|
|
name = tool.get_name()
|
|
description = tool.description
|
|
elif is_basemodel_instance(tool):
|
|
schema = tool.get_input_schema().schema()
|
|
name = tool.get_name()
|
|
description = tool.description
|
|
elif isinstance(tool, dict) and "name" in tool and "parameters" in tool:
|
|
return tool.copy()
|
|
else:
|
|
raise ValueError(
|
|
f"""Cannot convert {tool} to an Ollama tool.
|
|
{tool} needs to be a Pydantic class, model, or a dict."""
|
|
)
|
|
definition = {"name": name, "parameters": schema}
|
|
if description:
|
|
definition["description"] = description
|
|
|
|
return definition
|
|
|
|
def parse_response(message: BaseMessage) -> str:
|
|
"""Extract `function_call` from `AIMessage`."""
|
|
if isinstance(message, AIMessage):
|
|
kwargs = message.additional_kwargs
|
|
tool_calls = message.tool_calls
|
|
if len(tool_calls) > 0:
|
|
tool_call = tool_calls[-1]
|
|
args = tool_call.get("args")
|
|
return json.dumps(args)
|
|
elif "function_call" in kwargs:
|
|
if "arguments" in kwargs["function_call"]:
|
|
return kwargs["function_call"]["arguments"]
|
|
raise ValueError(f"`arguments` missing from `function_call` within AIMessage: {message}")
|
|
else:
|
|
raise ValueError("`tool_calls` missing from AIMessage: {message}")
|
|
raise ValueError(f"`message` is not an instance of `AIMessage`: {message}")
|
|
|
|
|
|
|
|
class OllamaFunctions(ChatOllama):
|
|
"""Function chat model that uses Ollama API."""
|
|
|
|
tool_system_prompt_template: str = DEFAULT_SYTEM_PROMPT
|
|
tool_system_prompt_template_with_history: str = DEFAULT_SYTEM_PROMPT_WITH_HISTORY
|
|
|
|
def __init__(self, **kwargs: Any) -> None:
|
|
super().__init__(**kwargs)
|
|
|
|
def bind_tools(
|
|
self,
|
|
tools: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool]],
|
|
**kwargs: Any,
|
|
) -> Runnable[LanguageModelInput, BaseMessage]:
|
|
return self.bind(functions=tools, **kwargs)
|
|
|
|
|
|
def _generate(self, messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any) -> ChatResult:
|
|
|
|
def _get_system_msg_and_formatted_history(self, messages: list) -> Tuple[str, str]:
|
|
def _format_tools_for_history(tool_calls: list[ToolCall]) -> str:
|
|
call_list = []
|
|
for c in tool_calls:
|
|
call_list.append({
|
|
"id": nxhash(c['id'])[-4:],
|
|
"tool": c['name'],
|
|
"args": c['args']
|
|
})
|
|
if len(call_list) == 1:
|
|
return json.dumps(obj=call_list[0], ensure_ascii=False, indent=2)
|
|
else:
|
|
return json.dumps(obj=call_list, ensure_ascii=False, indent=2)
|
|
formated_history = ""
|
|
system_msg = ""
|
|
for m in messages:
|
|
|
|
if formated_history != "":
|
|
formated_history += "\n\n"
|
|
|
|
if isinstance(m, SystemMessage):
|
|
system_msg += str(m.content)
|
|
elif isinstance(m, HumanMessage):
|
|
formated_history += "The Human said:\n" + str(m.content)
|
|
elif isinstance(m, AIMessage) and m.tool_calls:
|
|
formated_history += "So you called the tool" + (":\n" if len(m.tool_calls) == 1 else "s:\n") + _format_tools_for_history(m.tool_calls)
|
|
elif isinstance(m, ToolMessage):
|
|
formated_history += "To which the tool (" + nxhash(m.tool_call_id)[-4:] + ") replied with:\n" + str(m.content)
|
|
elif isinstance(m, AIMessage) and not m.tool_calls:
|
|
formated_history += "You said:\n" + str(m.content)
|
|
else:
|
|
raise TypeError("OllamaFunctions only supports SystemMessage HumanMessage ToolMessage AIMessage but got " + str(type(m)))
|
|
|
|
return system_msg, formated_history
|
|
|
|
|
|
def _get_parsed_chat_result(self, chat_result_str: str) -> Union[dict, str]:
|
|
try:
|
|
parsed_chat_result = json.loads(chat_result_str)
|
|
except json.JSONDecodeError:
|
|
parsed_chat_result = chat_result_str
|
|
return parsed_chat_result
|
|
|
|
def _get_called_tool(self, d: dict, functions_list: list[dict]) -> dict|NoneType:
|
|
if not parsed_chat_result:
|
|
called_tool_name = None
|
|
elif "tool" in parsed_chat_result:
|
|
called_tool_name = d["tool"] # per spec
|
|
elif "name" in d:
|
|
called_tool_name = d["name"] # Phi3 often does this
|
|
elif "tool_name" in d:
|
|
called_tool_name = d["tool_name"] # Phi3 often does this
|
|
elif "action" in d:
|
|
called_tool_name = d["action"] # Phi3 does this
|
|
else:
|
|
return None
|
|
|
|
try:
|
|
called_tool = [tool for tool in functions_list if tool['name'] == called_tool_name][0]
|
|
except IndexError:
|
|
return None # when a tool is called, but the tool doesnt exist
|
|
|
|
return called_tool
|
|
|
|
def _extract_conversaional_response(self, d: dict) -> str:
|
|
if ("tool_input" in d and "response" in d["tool_input"]):
|
|
response = d["tool_input"]["response"]
|
|
elif ("input" in d and "response" in d["input"]):
|
|
response = d["input"]["response"]
|
|
elif ("args" in d and "response" in d["args"]):
|
|
response = d["args"]["response"]
|
|
elif "response" in d:
|
|
response = d["response"]
|
|
elif "input" in d:
|
|
response = d["input"]
|
|
elif "args" in d:
|
|
response = d["args"]
|
|
elif "tool_input" in d:
|
|
response = d["tool_input"]
|
|
else:
|
|
raise ValueError(f"Failed to parse a response from {self.model} output: {chat_result}")
|
|
|
|
try:
|
|
assert isinstance(response, str)
|
|
except AssertionError:
|
|
raise ValueError(f"Failed to parse a response from {self.model} output: {chat_result}")
|
|
|
|
return response
|
|
|
|
def _extract_tool_args(self, d: dict) -> dict:
|
|
if "tool_input" in parsed_chat_result:
|
|
called_tool_args = d["tool_input"] # per spec
|
|
elif "input" in d:
|
|
called_tool_args = d["input"] # Phi3 often does this
|
|
elif "args" in d:
|
|
called_tool_args = d["args"]
|
|
else:
|
|
called_tool_args = {}
|
|
return called_tool_args
|
|
|
|
# prepare generation
|
|
functions_list = [convert_to_ollama_tool(fn) for fn in kwargs.get("functions", [])]
|
|
functions_list.append(CONVERSATIONAL_RESPONSE_TOOL)
|
|
functions_str = json.dumps(functions_list, indent=2)
|
|
|
|
# prepare generation with history
|
|
if True in [ isinstance(m, ToolMessage) for m in messages ]:
|
|
system_msg, formated_history = _get_system_msg_and_formatted_history(self, messages=messages)
|
|
|
|
system_message_prompt_template = SystemMessagePromptTemplate.from_template(self.tool_system_prompt_template_with_history)
|
|
system_message = system_message_prompt_template.format(
|
|
tools=functions_str,
|
|
history=formated_history,
|
|
system_msg=system_msg
|
|
)
|
|
final_messages = [ system_message ]
|
|
|
|
# prepare generation without history
|
|
else:
|
|
system_message_prompt_template = SystemMessagePromptTemplate.from_template(self.tool_system_prompt_template)
|
|
system_message = system_message_prompt_template.format(
|
|
tools=functions_str
|
|
)
|
|
final_messages = [ system_message ] + messages
|
|
|
|
# genrerate chat result
|
|
response_message = super()._generate(final_messages, stop=stop, run_manager=run_manager, **kwargs)
|
|
chat_result = response_message.generations[0].text
|
|
|
|
# chekc for validity
|
|
if not isinstance(chat_result, str):
|
|
raise ValueError("OllamaFunctions does not support non-string output.")
|
|
|
|
# make str to dict
|
|
parsed_chat_result = _get_parsed_chat_result(self, chat_result_str=chat_result)
|
|
# if model failed to return vailid json, just retrun the whole thing
|
|
if isinstance(parsed_chat_result, str):
|
|
return ChatResult(generations=[ChatGeneration(message=AIMessage(content=parsed_chat_result))])
|
|
|
|
|
|
# get the called tool from the dict
|
|
called_tool = _get_called_tool(self, d=parsed_chat_result, functions_list=functions_list)
|
|
|
|
if not called_tool:
|
|
response_msg = AIMessage(content=_extract_conversaional_response(self, d=parsed_chat_result))
|
|
elif called_tool == CONVERSATIONAL_RESPONSE_TOOL:
|
|
response_msg = AIMessage(content=_extract_conversaional_response(self, d=parsed_chat_result))
|
|
else:
|
|
response_msg = AIMessage(
|
|
content="",
|
|
tool_calls=[ToolCall(
|
|
name=called_tool['name'],
|
|
args=_extract_tool_args(self, d=parsed_chat_result),
|
|
id=f"call_{str(uuid.uuid4()).replace('-', '')}",
|
|
)],
|
|
)
|
|
|
|
return ChatResult(generations=[ChatGeneration(message=response_msg)])
|
|
|
|
@property
|
|
def _llm_type(self) -> str:
|
|
return "ollama_functions"
|