79 lines
3.8 KiB
Clojure
79 lines
3.8 KiB
Clojure
(ns tbd.agent.test-support
|
|
"A fake LangChain4j StreamingChatModel for tests: no real network, no LM Studio. `model-of`
|
|
builds one from a script — a function of the ChatRequest that calls the handler's callbacks
|
|
directly, exactly as a real provider's client would, so tbd.agent.core's pump code is exercised
|
|
for real. `text-model`/`tool-call-model`/`error-model`/`thinking-model` cover the common cases."
|
|
(:require [tbd.agent.core :as agent])
|
|
(:import (dev.langchain4j.agent.tool ToolExecutionRequest)
|
|
(dev.langchain4j.data.message AiMessage)
|
|
(dev.langchain4j.model.chat StreamingChatModel)
|
|
(dev.langchain4j.model.chat.response
|
|
ChatResponse ChatResponseMetadata$Builder CompleteToolCall PartialResponse
|
|
PartialResponseContext PartialThinking PartialThinkingContext StreamingHandle)
|
|
(dev.langchain4j.model.output FinishReason TokenUsage)))
|
|
|
|
(defn fake-handle
|
|
"A StreamingHandle whose `cancel` is observable through the atom `cancelled?`."
|
|
([] (fake-handle (atom false)))
|
|
([cancelled?]
|
|
(reify StreamingHandle
|
|
(cancel [_] (reset! cancelled? true))
|
|
(isCancelled [_] @cancelled?))))
|
|
|
|
(defn model-of
|
|
"A StreamingChatModel whose `doChat` calls `(script req handler)`."
|
|
^StreamingChatModel [script]
|
|
(reify StreamingChatModel
|
|
(doChat [_ req handler] (script req handler) nil)))
|
|
|
|
(defn- complete [handler ai-message finish-reason usage]
|
|
(.onCompleteResponse handler
|
|
(-> (ChatResponse/builder)
|
|
(.aiMessage ai-message)
|
|
(.metadata (cond-> (ChatResponseMetadata$Builder.)
|
|
finish-reason (.finishReason finish-reason)
|
|
usage (.tokenUsage (TokenUsage. (int (first usage)) (int (second usage))))
|
|
true .build))
|
|
.build)))
|
|
|
|
(defn text-model
|
|
"A model that streams `pieces` as :content, then finishes with their concatenation as the reply."
|
|
[& pieces]
|
|
(model-of (fn [_ handler]
|
|
(let [h (fake-handle)]
|
|
(doseq [p pieces] (.onPartialResponse handler (PartialResponse. ^String p) (PartialResponseContext. h)))
|
|
(complete handler (AiMessage/from (apply str pieces)) FinishReason/STOP [10 5])))))
|
|
|
|
(defn thinking-model
|
|
"A model that streams `thinking`, then `pieces` as content."
|
|
[thinking & pieces]
|
|
(model-of (fn [_ handler]
|
|
(let [h (fake-handle)]
|
|
(.onPartialThinking handler (PartialThinking. ^String thinking) (PartialThinkingContext. h))
|
|
(doseq [p pieces] (.onPartialResponse handler (PartialResponse. ^String p) (PartialResponseContext. h)))
|
|
(complete handler (AiMessage/from (apply str pieces)) FinishReason/STOP [10 5])))))
|
|
|
|
(defn tool-call-model
|
|
"A model that asks for one tool call: {:id :name :arguments}."
|
|
[{:keys [id name arguments]}]
|
|
(model-of (fn [_ handler]
|
|
(let [req (-> (ToolExecutionRequest/builder) (.id id) (.name name) (.arguments arguments) .build)]
|
|
(.onCompleteToolCall handler (CompleteToolCall. (int 0) req))
|
|
(complete handler (AiMessage/from "" [req]) FinishReason/TOOL_EXECUTION nil)))))
|
|
|
|
(defn error-model
|
|
"A model whose call always fails with `message`."
|
|
[message]
|
|
(model-of (fn [_ handler] (.onError handler (RuntimeException. ^String message)))))
|
|
|
|
(defn no-reply-model
|
|
"A model that ends the turn with nothing said and no calls (the LLM sent no reply)."
|
|
[]
|
|
(model-of (fn [_ handler] (complete handler (AiMessage/from "") FinishReason/STOP nil))))
|
|
|
|
(defn config
|
|
"A chat config (see tbd.agent.core/env-config) using `model` instead of a real LM Studio."
|
|
([model] (config model {}))
|
|
([model overrides]
|
|
(merge (agent/env-config {"LMSTUDIO_URL" "http://llm.test"}) {:chat-model model} overrides)))
|