wehub-resource-sync
52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
import itertools
|
|
|
|
from langchain.agents import create_agent
|
|
from langchain.tools import tool
|
|
from langchain_core.messages import AIMessageChunk, ToolCall
|
|
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
import mlflow
|
|
|
|
|
|
class FakeOpenAI(ChatOpenAI, extra="allow"):
|
|
# In normal LangChain tests, we use the fake OpenAI server to mock the OpenAI REST API.
|
|
# The fake server returns the input payload as it is. However, for agent tests, the
|
|
# response should be a specific format so that the agent can parse it correctly.
|
|
# Also, mocking with mock.patch does not work for testing model serving (as the server
|
|
# will run in a separate process).
|
|
# Therefore, we mock the OpenAI client in the model definition here.
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
# Using itertools.cycle to create an infinite iterator
|
|
self._responses = itertools.cycle([
|
|
AIMessageChunk(
|
|
content="",
|
|
tool_calls=[ToolCall(name="multiply", args={"a": 2, "b": 3}, id="123")],
|
|
),
|
|
AIMessageChunk(content="The result of 2 * 3 is 6."),
|
|
])
|
|
|
|
def _generate(self, *args, **kwargs):
|
|
return ChatResult(generations=[ChatGeneration(message=next(self._responses))])
|
|
|
|
def _stream(self, *args, **kwargs):
|
|
yield ChatGenerationChunk(message=next(self._responses))
|
|
|
|
|
|
@tool
|
|
def add(a: int, b: int) -> int:
|
|
"""Add two numbers."""
|
|
return a + b
|
|
|
|
|
|
@tool
|
|
def multiply(a: int, b: int) -> int:
|
|
"""Multiply two numbers."""
|
|
return a * b
|
|
|
|
|
|
llm = FakeOpenAI()
|
|
agent = create_agent(llm, [add, multiply], system_prompt="You are a helpful assistant")
|
|
mlflow.models.set_model(agent)
|