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": , "tool_input": }} """ 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": , "tool_input": }} 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"