letta-nightly 0.1.7.dev20240924104148__py3-none-any.whl

letta/local_llm/chat_completion_proxy.py

@@ -0,0 +1,279 @@
+"""Key idea: create drop-in replacement for agent's ChatCompletion call that runs on an OpenLLM backend"""
+
+import uuid
+
+import requests
+
+from letta.constants import CLI_WARNING_PREFIX
+from letta.errors import LocalLLMConnectionError, LocalLLMError
+from letta.local_llm.constants import DEFAULT_WRAPPER
+from letta.local_llm.function_parser import patch_function
+from letta.local_llm.grammars.gbnf_grammar_generator import (
+    create_dynamic_model_from_function,
+    generate_gbnf_grammar_and_documentation,
+)
+from letta.local_llm.groq.api import get_groq_completion
+from letta.local_llm.koboldcpp.api import get_koboldcpp_completion
+from letta.local_llm.llamacpp.api import get_llamacpp_completion
+from letta.local_llm.llm_chat_completion_wrappers import simple_summary_wrapper
+from letta.local_llm.lmstudio.api import get_lmstudio_completion
+from letta.local_llm.ollama.api import get_ollama_completion
+from letta.local_llm.utils import count_tokens, get_available_wrappers
+from letta.local_llm.vllm.api import get_vllm_completion
+from letta.local_llm.webui.api import get_webui_completion
+from letta.local_llm.webui.legacy_api import (
+    get_webui_completion as get_webui_completion_legacy,
+)
+from letta.prompts.gpt_summarize import SYSTEM as SUMMARIZE_SYSTEM_MESSAGE
+from letta.schemas.openai.chat_completion_response import (
+    ChatCompletionResponse,
+    Choice,
+    Message,
+    ToolCall,
+    UsageStatistics,
+)
+from letta.utils import get_tool_call_id, get_utc_time, json_dumps
+
+has_shown_warning = False
+grammar_supported_backends = ["koboldcpp", "llamacpp", "webui", "webui-legacy"]
+
+
+def get_chat_completion(
+    model,
+    # no model required (except for Ollama), since the model is fixed to whatever you set in your own backend
+    messages,
+    functions=None,
+    functions_python=None,
+    function_call="auto",
+    context_window=None,
+    user=None,
+    # required
+    wrapper=None,
+    endpoint=None,
+    endpoint_type=None,
+    # optional cleanup
+    function_correction=True,
+    # extra hints to allow for additional prompt formatting hacks
+    # TODO this could alternatively be supported via passing function_call="send_message" into the wrapper
+    first_message=False,
+    # optional auth headers
+    auth_type=None,
+    auth_key=None,
+) -> ChatCompletionResponse:
+    from letta.utils import printd
+
+    assert context_window is not None, "Local LLM calls need the context length to be explicitly set"
+    assert endpoint is not None, "Local LLM calls need the endpoint (eg http://localendpoint:1234) to be explicitly set"
+    assert endpoint_type is not None, "Local LLM calls need the endpoint type (eg webui) to be explicitly set"
+    global has_shown_warning
+    grammar = None
+
+    # TODO: eventually just process Message object
+    if not isinstance(messages[0], dict):
+        messages = [m.to_openai_dict() for m in messages]
+
+    if function_call is not None and function_call != "auto":
+        raise ValueError(f"function_call == {function_call} not supported (auto or None only)")
+
+    available_wrappers = get_available_wrappers()
+    documentation = None
+
+    # Special case for if the call we're making is coming from the summarizer
+    if messages[0]["role"] == "system" and messages[0]["content"].strip() == SUMMARIZE_SYSTEM_MESSAGE.strip():
+        llm_wrapper = simple_summary_wrapper.SimpleSummaryWrapper()
+
+    # Select a default prompt formatter
+    elif wrapper is None:
+        # Warn the user that we're using the fallback
+        if not has_shown_warning:
+            print(
+                f"{CLI_WARNING_PREFIX}no wrapper specified for local LLM, using the default wrapper (you can remove this warning by specifying the wrapper with --model-wrapper)"
+            )
+            has_shown_warning = True
+
+        llm_wrapper = DEFAULT_WRAPPER()
+
+    # User provided an incorrect prompt formatter
+    elif wrapper not in available_wrappers:
+        raise ValueError(f"Could not find requested wrapper '{wrapper} in available wrappers list:\n{', '.join(available_wrappers)}")
+
+    # User provided a correct prompt formatter
+    else:
+        llm_wrapper = available_wrappers[wrapper]
+
+    # If the wrapper uses grammar, generate the grammar using the grammar generating function
+    # TODO move this to a flag
+    if wrapper is not None and "grammar" in wrapper:
+        # When using grammars, we don't want to do any extras output tricks like appending a response prefix
+        setattr(llm_wrapper, "assistant_prefix_extra_first_message", "")
+        setattr(llm_wrapper, "assistant_prefix_extra", "")
+
+        # TODO find a better way to do this than string matching (eg an attribute)
+        if "noforce" in wrapper:
+            # "noforce" means that the prompt formatter expects inner thoughts as a top-level parameter
+            # this is closer to the OpenAI style since it allows for messages w/o any function calls
+            # however, with bad LLMs it makes it easier for the LLM to "forget" to call any of the functions
+            grammar, documentation = generate_grammar_and_documentation(
+                functions_python=functions_python,
+                add_inner_thoughts_top_level=True,
+                add_inner_thoughts_param_level=False,
+                allow_only_inner_thoughts=True,
+            )
+        else:
+            # otherwise, the other prompt formatters will insert inner thoughts as a function call parameter (by default)
+            # this means that every response from the LLM will be required to call a function
+            grammar, documentation = generate_grammar_and_documentation(
+                functions_python=functions_python,
+                add_inner_thoughts_top_level=False,
+                add_inner_thoughts_param_level=True,
+                allow_only_inner_thoughts=False,
+            )
+        printd(grammar)
+
+    if grammar is not None and endpoint_type not in grammar_supported_backends:
+        print(
+            f"{CLI_WARNING_PREFIX}grammars are currently not supported when using {endpoint_type} as the Letta local LLM backend (supported: {', '.join(grammar_supported_backends)})"
+        )
+        grammar = None
+
+    # First step: turn the message sequence into a prompt that the model expects
+    try:
+        # if hasattr(llm_wrapper, "supports_first_message"):
+        if hasattr(llm_wrapper, "supports_first_message") and llm_wrapper.supports_first_message:
+            prompt = llm_wrapper.chat_completion_to_prompt(
+                messages=messages, functions=functions, first_message=first_message, function_documentation=documentation
+            )
+        else:
+            prompt = llm_wrapper.chat_completion_to_prompt(messages=messages, functions=functions, function_documentation=documentation)
+
+        printd(prompt)
+    except Exception as e:
+        print(e)
+        raise LocalLLMError(
+            f"Failed to convert ChatCompletion messages into prompt string with wrapper {str(llm_wrapper)} - error: {str(e)}"
+        )
+
+    try:
+        if endpoint_type == "webui":
+            result, usage = get_webui_completion(endpoint, auth_type, auth_key, prompt, context_window, grammar=grammar)
+        elif endpoint_type == "webui-legacy":
+            result, usage = get_webui_completion_legacy(endpoint, auth_type, auth_key, prompt, context_window, grammar=grammar)
+        elif endpoint_type == "lmstudio":
+            result, usage = get_lmstudio_completion(endpoint, auth_type, auth_key, prompt, context_window, api="completions")
+        elif endpoint_type == "lmstudio-legacy":
+            result, usage = get_lmstudio_completion(endpoint, auth_type, auth_key, prompt, context_window, api="chat")
+        elif endpoint_type == "llamacpp":
+            result, usage = get_llamacpp_completion(endpoint, auth_type, auth_key, prompt, context_window, grammar=grammar)
+        elif endpoint_type == "koboldcpp":
+            result, usage = get_koboldcpp_completion(endpoint, auth_type, auth_key, prompt, context_window, grammar=grammar)
+        elif endpoint_type == "ollama":
+            result, usage = get_ollama_completion(endpoint, auth_type, auth_key, model, prompt, context_window)
+        elif endpoint_type == "vllm":
+            result, usage = get_vllm_completion(endpoint, auth_type, auth_key, model, prompt, context_window, user)
+        elif endpoint_type == "groq":
+            result, usage = get_groq_completion(endpoint, auth_type, auth_key, model, prompt, context_window)
+        else:
+            raise LocalLLMError(
+                f"Invalid endpoint type {endpoint_type}, please set variable depending on your backend (webui, lmstudio, llamacpp, koboldcpp)"
+            )
+    except requests.exceptions.ConnectionError as e:
+        raise LocalLLMConnectionError(f"Unable to connect to endpoint {endpoint}")
+
+    if result is None or result == "":
+        raise LocalLLMError(f"Got back an empty response string from {endpoint}")
+    printd(f"Raw LLM output:\n====\n{result}\n====")
+
+    try:
+        if hasattr(llm_wrapper, "supports_first_message") and llm_wrapper.supports_first_message:
+            chat_completion_result = llm_wrapper.output_to_chat_completion_response(result, first_message=first_message)
+        else:
+            chat_completion_result = llm_wrapper.output_to_chat_completion_response(result)
+        printd(json_dumps(chat_completion_result, indent=2))
+    except Exception as e:
+        raise LocalLLMError(f"Failed to parse JSON from local LLM response - error: {str(e)}")
+
+    # Run through some manual function correction (optional)
+    if function_correction:
+        chat_completion_result = patch_function(message_history=messages, new_message=chat_completion_result)
+
+    # Fill in potential missing usage information (used for tracking token use)
+    if not ("prompt_tokens" in usage and "completion_tokens" in usage and "total_tokens" in usage):
+        raise LocalLLMError(f"usage dict in response was missing fields ({usage})")
+
+    if usage["prompt_tokens"] is None:
+        printd(f"usage dict was missing prompt_tokens, computing on-the-fly...")
+        usage["prompt_tokens"] = count_tokens(prompt)
+
+    # NOTE: we should compute on-the-fly anyways since we might have to correct for errors during JSON parsing
+    usage["completion_tokens"] = count_tokens(json_dumps(chat_completion_result))
+    """
+    if usage["completion_tokens"] is None:
+        printd(f"usage dict was missing completion_tokens, computing on-the-fly...")
+        # chat_completion_result is dict with 'role' and 'content'
+        # token counter wants a string
+        usage["completion_tokens"] = count_tokens(json_dumps(chat_completion_result))
+    """
+
+    # NOTE: this is the token count that matters most
+    if usage["total_tokens"] is None:
+        printd(f"usage dict was missing total_tokens, computing on-the-fly...")
+        usage["total_tokens"] = usage["prompt_tokens"] + usage["completion_tokens"]
+
+    # unpack with response.choices[0].message.content
+    response = ChatCompletionResponse(
+        id=str(uuid.uuid4()),  # TODO something better?
+        choices=[
+            Choice(
+                finish_reason="stop",
+                index=0,
+                message=Message(
+                    role=chat_completion_result["role"],
+                    content=chat_completion_result["content"],
+                    tool_calls=(
+                        [ToolCall(id=get_tool_call_id(), type="function", function=chat_completion_result["function_call"])]
+                        if "function_call" in chat_completion_result
+                        else []
+                    ),
+                ),
+            )
+        ],
+        created=get_utc_time(),
+        model=model,
+        # "This fingerprint represents the backend configuration that the model runs with."
+        # system_fingerprint=user if user is not None else "null",
+        system_fingerprint=None,
+        object="chat.completion",
+        usage=UsageStatistics(**usage),
+    )
+    printd(response)
+    return response
+
+
+def generate_grammar_and_documentation(
+    functions_python: dict,
+    add_inner_thoughts_top_level: bool,
+    add_inner_thoughts_param_level: bool,
+    allow_only_inner_thoughts: bool,
+):
+    from letta.utils import printd
+
+    assert not (
+        add_inner_thoughts_top_level and add_inner_thoughts_param_level
+    ), "Can only place inner thoughts in one location in the grammar generator"
+
+    grammar_function_models = []
+    # create_dynamic_model_from_function will add inner thoughts to the function parameters if add_inner_thoughts is True.
+    # generate_gbnf_grammar_and_documentation will add inner thoughts to the outer object of the function parameters if add_inner_thoughts is True.
+    for key, func in functions_python.items():
+        grammar_function_models.append(create_dynamic_model_from_function(func, add_inner_thoughts=add_inner_thoughts_param_level))
+    grammar, documentation = generate_gbnf_grammar_and_documentation(
+        grammar_function_models,
+        outer_object_name="function",
+        outer_object_content="params",
+        model_prefix="function",
+        fields_prefix="params",
+        add_inner_thoughts=add_inner_thoughts_top_level,
+        allow_only_inner_thoughts=allow_only_inner_thoughts,
+    )
+    printd(grammar)
+    return grammar, documentation

letta/local_llm/constants.py

@@ -0,0 +1,31 @@
+# import letta.local_llm.llm_chat_completion_wrappers.airoboros as airoboros
+from letta.local_llm.llm_chat_completion_wrappers.chatml import (
+    ChatMLInnerMonologueWrapper,
+)
+
+DEFAULT_ENDPOINTS = {
+    # Local
+    "koboldcpp": "http://localhost:5001",
+    "llamacpp": "http://localhost:8080",
+    "lmstudio": "http://localhost:1234",
+    "lmstudio-legacy": "http://localhost:1234",
+    "ollama": "http://localhost:11434",
+    "webui-legacy": "http://localhost:5000",
+    "webui": "http://localhost:5000",
+    "vllm": "http://localhost:8000",
+    # APIs
+    "openai": "https://api.openai.com",
+    "anthropic": "https://api.anthropic.com",
+    "groq": "https://api.groq.com/openai",
+}
+
+DEFAULT_OLLAMA_MODEL = "dolphin2.2-mistral:7b-q6_K"
+
+# DEFAULT_WRAPPER = airoboros.Airoboros21InnerMonologueWrapper
+# DEFAULT_WRAPPER_NAME = "airoboros-l2-70b-2.1"
+
+DEFAULT_WRAPPER = ChatMLInnerMonologueWrapper
+DEFAULT_WRAPPER_NAME = "chatml"
+
+INNER_THOUGHTS_KWARG = "inner_thoughts"
+INNER_THOUGHTS_KWARG_DESCRIPTION = "Deep inner monologue private to you only."

letta/local_llm/function_parser.py

@@ -0,0 +1,68 @@
+import copy
+import json
+
+from letta.utils import json_dumps, json_loads
+
+NO_HEARTBEAT_FUNCS = ["send_message", "pause_heartbeats"]
+
+
+def insert_heartbeat(message):
+    # message_copy = message.copy()
+    message_copy = copy.deepcopy(message)
+
+    if message_copy.get("function_call"):
+        # function_name = message.get("function_call").get("name")
+        params = message_copy.get("function_call").get("arguments")
+        params = json_loads(params)
+        params["request_heartbeat"] = True
+        message_copy["function_call"]["arguments"] = json_dumps(params)
+
+    elif message_copy.get("tool_call"):
+        # function_name = message.get("tool_calls")[0].get("function").get("name")
+        params = message_copy.get("tool_calls")[0].get("function").get("arguments")
+        params = json_loads(params)
+        params["request_heartbeat"] = True
+        message_copy["tools_calls"][0]["function"]["arguments"] = json_dumps(params)
+
+    return message_copy
+
+
+def heartbeat_correction(message_history, new_message):
+    """Add heartbeats where we think the agent forgot to add them themselves
+
+    If the last message in the stack is a user message and the new message is an assistant func call, fix the heartbeat
+
+    See: https://github.com/cpacker/Letta/issues/601
+    """
+    if len(message_history) < 1:
+        return None
+
+    last_message_was_user = False
+    if message_history[-1]["role"] == "user":
+        try:
+            content = json_loads(message_history[-1]["content"])
+        except json.JSONDecodeError:
+            return None
+        # Check if it's a user message or system message
+        if content["type"] == "user_message":
+            last_message_was_user = True
+
+    new_message_is_heartbeat_function = False
+    if new_message["role"] == "assistant":
+        if new_message.get("function_call") or new_message.get("tool_calls"):
+            if new_message.get("function_call"):
+                function_name = new_message.get("function_call").get("name")
+            elif new_message.get("tool_calls"):
+                function_name = new_message.get("tool_calls")[0].get("function").get("name")
+            if function_name not in NO_HEARTBEAT_FUNCS:
+                new_message_is_heartbeat_function = True
+
+    if last_message_was_user and new_message_is_heartbeat_function:
+        return insert_heartbeat(new_message)
+    else:
+        return None
+
+
+def patch_function(message_history, new_message):
+    corrected_output = heartbeat_correction(message_history=message_history, new_message=new_message)
+    return corrected_output if corrected_output is not None else new_message