diff --git a/tests/trace/builtin_objects/backend_models.ipynb b/tests/trace/builtin_objects/backend_models.ipynb new file mode 100644 index 00000000000..ba835ae269d --- /dev/null +++ b/tests/trace/builtin_objects/backend_models.ipynb @@ -0,0 +1,251 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "os.environ[\"WF_TRACE_SERVER_URL\"] = \"http://127.0.01:6345\"" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "import weave" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Logged in as Weights & Biases user: timssweeney.\n", + "View Weave data at https://wandb.ai/timssweeney/remote_model_demo/weave\n" + ] + } + ], + "source": [ + "client = weave.init(\"remote_model_demo_2\")" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "🍩 https://wandb.ai/timssweeney/remote_model_demo/r/call/0193b497-27e9-7f31-8f26-a24d76254d32\n" + ] + }, + { + "data": { + "text/plain": [ + "{'name': 'Fred', 'age': 30}" + ] + }, + "execution_count": 9, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Demonstrates creating a model in python\n", + "\n", + "from weave.builtin_objects.models.CompletionModel import LiteLLMCompletionModel\n", + "\n", + "model = LiteLLMCompletionModel(\n", + " model=\"gpt-4o\",\n", + " messages_template=[\n", + " {\"role\": \"system\", \"content\": \"Please extract the name and age from the following text\"},\n", + " {\"role\": \"user\", \"content\": \"{user_input}\"},\n", + " ],\n", + " response_format={\n", + " \"type\": \"json_schema\",\n", + " \"json_schema\": {\n", + " \"name\": \"Person\",\n", + " \"schema\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"age\": {\"type\": \"integer\"},\n", + " \"name\": {\"type\": \"string\"},\n", + " },\n", + " }\n", + " }\n", + " },\n", + ")\n", + "\n", + "res = model.predict(user_input=\"Hello, my name is Fred and I am 30 years old.\")\n", + "\n", + "print(model.ref.uri())\n", + "\n", + "print(res)" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "CallMethodRes(call_id='0193b497-3c5b-7121-9a88-f14798eac933', output={'name': 'Charles', 'age': 40})" + ] + }, + "execution_count": 10, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Demonstrates calling a model created in python\n", + "\n", + "from weave.trace_server.trace_server_interface import CallMethodReq\n", + "\n", + "call_res = client.server.call_method(CallMethodReq(\n", + " project_id=client._project_id(),\n", + " object_ref=model.ref.uri(),\n", + " method_name=\"predict\",\n", + " args={\n", + " \"user_input\": \"Hello, my name is Charles and I am 40 years old.\"\n", + " }\n", + "))\n", + "\n", + "call_res" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "ObjCreateRes(digest='GSK0zJwYfB7Scv2StQWsxGtoxuvuGxTSEOFzKSrmkAg')" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Demonstrates creating a model in the UI - notice the digest match\n", + "\n", + "from weave.trace_server.trace_server_interface import ObjCreateReq\n", + "\n", + "obj_res = client.server.obj_create(ObjCreateReq.model_validate({'obj': {\n", + " \"project_id\": client._project_id(),\n", + " \"object_id\": \"LiteLLMCompletionModel\",\n", + " \"val\": {\n", + " \"model\": \"gpt-4o\",\n", + " \"messages_template\": [{\"role\": \"user\", \"content\": \"{user_input}\"}],\n", + " \"response_format\": {\n", + " \"type\": \"json_schema\",\n", + " \"json_schema\": {\n", + " \"name\": \"Person\",\n", + " \"schema\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"age\": {\"type\": \"integer\"},\n", + " \"name\": {\"type\": \"string\"},\n", + " },\n", + " },\n", + " },\n", + " },\n", + " },\n", + " \"set_leaf_object_class\": \"LiteLLMCompletionModel\",\n", + "}}))\n", + "\n", + "obj_res" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "🍩 https://wandb.ai/timssweeney/remote_model_demo/r/call/0193b498-a593-7391-b059-b32909b46f93\n" + ] + }, + { + "data": { + "text/plain": [ + "{'name': 'Fred', 'age': 30}" + ] + }, + "execution_count": 15, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Demonstrates fetching a model in python that was created in the UI\n", + "\n", + "from weave.trace.refs import ObjectRef\n", + "\n", + "gotten_model = weave.ref(ObjectRef(\n", + " entity=client.entity,\n", + " project=client.project,\n", + " name=\"LiteLLMCompletionModel\",\n", + " _digest=obj_res.digest\n", + ").uri()).get()\n", + "\n", + "res = gotten_model.predict(user_input=\"Hello, my name is Fred and I am 30 years old.\")\n", + "res" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "wandb-weave", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.8" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/weave/trace_server_bindings/remote_http_trace_server.py b/weave/trace_server_bindings/remote_http_trace_server.py index b749af60d24..2e7b9054a97 100644 --- a/weave/trace_server_bindings/remote_http_trace_server.py +++ b/weave/trace_server_bindings/remote_http_trace_server.py @@ -568,6 +568,11 @@ def completions_create( tsi.CompletionsCreateReq, tsi.CompletionsCreateRes, ) + + def call_method(self, req: tsi.CallMethodReq) -> tsi.CallMethodRes: + return self._generic_request( + "/execute/method", req, tsi.CallMethodReq, tsi.CallMethodRes + ) __docspec__ = [