mega commit
This commit is contained in:
Lennart J. Kurzweg (Nx2)
322
libs/ollama_functions.py
Normal file
322
libs/ollama_functions.py
Normal file
@@ -0,0 +1,322 @@
|
||||
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"
|
||||
Reference in New Issue
Block a user