From 41ea5750714921ffe1880b335e67f3b910cb18fc Mon Sep 17 00:00:00 2001 From: Sunrise Date: Tue, 30 Apr 2024 03:21:18 +0800 Subject: [PATCH] Use `#connect_timeout=(Time::Span)` --- src/mockgpt.cr | 86 ++++++------- src/mockgpt/controllers.cr | 214 ++++++++++++++++----------------- src/mockgpt/cors.cr | 42 +++---- src/mockgpt/routes.cr | 32 ++--- src/mockgpt/struct/request.cr | 22 ++-- src/mockgpt/struct/response.cr | 110 ++++++++--------- 6 files changed, 253 insertions(+), 253 deletions(-) diff --git a/src/mockgpt.cr b/src/mockgpt.cr index 880358d..818a44d 100644 --- a/src/mockgpt.cr +++ b/src/mockgpt.cr @@ -1,43 +1,43 @@ -require "grip" -require "uri" -require "json" -require "colorize" -require "option_parser" -require "./mockgpt/*" -require "./mockgpt/struct/*" - -module Mocker - class_property ip : String = "localhost" - class_property port : Int32 = 3000 - class_property model : String = "codellama:13b" - class_property gpt : String = "gpt-4" -end - -homePath = "#{Path.home}/mocker.json" -exePath = "#{File.dirname Process.executable_path.not_nil!}/mocker.json" -confPath = File.exists?(exePath) ? exePath : homePath -if File.file?(confPath) - mocker = JSON.parse(File.read(confPath)) - Mocker.ip = mocker["ip"].as_s if mocker["ip"]? - Mocker.port = mocker["port"].as_i if mocker["port"]? - Mocker.model = mocker["model"].as_s if mocker["model"]? -end - -OptionParser.parse do |parser| - parser.banner = "Usage: mockgpt [arguments]" - parser.on("-b HOST", "--binding HOST", "Bind to the specified IP") { |_host| Mocker.ip = _host } - parser.on("-p PORT", "--port PORT", "Run on the specified port") { |_port| Mocker.port = _port.to_i } - parser.on("-m MODEL", "--mocker MODEL", "Employ the specified model") { |_model| Mocker.model = _model } - parser.on("-h", "--help", "Show this help") do - puts parser - exit - end - parser.invalid_option do |flag| - STDERR.puts "ERROR: #{flag} is not a valid option." - STDERR.puts parser - exit(1) - end -end - -mockgpt = Application.new(Mocker.ip, Mocker.port) -mockgpt.run +require "grip" +require "uri" +require "json" +require "colorize" +require "option_parser" +require "./mockgpt/*" +require "./mockgpt/struct/*" + +module Mocker + class_property ip : String = "localhost" + class_property port : Int32 = 3000 + class_property model : String = "codellama:13b" + class_property gpt : String = "gpt-4" +end + +homePath = "#{Path.home}/mocker.json" +exePath = "#{File.dirname Process.executable_path.not_nil!}/mocker.json" +confPath = File.exists?(exePath) ? exePath : homePath +if File.file?(confPath) + mocker = JSON.parse(File.read(confPath)) + Mocker.ip = mocker["ip"].as_s if mocker["ip"]? + Mocker.port = mocker["port"].as_i if mocker["port"]? + Mocker.model = mocker["model"].as_s if mocker["model"]? +end + +OptionParser.parse do |parser| + parser.banner = "Usage: mockgpt [arguments]" + parser.on("-b HOST", "--binding HOST", "Bind to the specified IP") { |_host| Mocker.ip = _host } + parser.on("-p PORT", "--port PORT", "Run on the specified port") { |_port| Mocker.port = _port.to_i } + parser.on("-m MODEL", "--mocker MODEL", "Employ the specified model") { |_model| Mocker.model = _model } + parser.on("-h", "--help", "Show this help") do + puts parser + exit + end + parser.invalid_option do |flag| + STDERR.puts "ERROR: #{flag} is not a valid option." + STDERR.puts parser + exit(1) + end +end + +mockgpt = Application.new(Mocker.ip, Mocker.port) +mockgpt.run diff --git a/src/mockgpt/controllers.cr b/src/mockgpt/controllers.cr index 42a442f..3e682da 100644 --- a/src/mockgpt/controllers.cr +++ b/src/mockgpt/controllers.cr @@ -1,107 +1,107 @@ -class MockGPT < Grip::Controllers::Http - def connect(context : Context) - context - .put_status(HTTP::Status::OK) - .halt - end - - def ollama(context : Context) - params = context.fetch_json_params - params["model"] = Mocker.model - presetUri = URI.parse "http://#{ENV["OLLAMA_HOST"]}" - url = URI.new scheme: presetUri.scheme || "http", host: presetUri.host || "localhost", port: presetUri.port || 11434 - agent = HTTP::Client.new url - agent.connect_timeout = 5 - print " POST ".colorize.bright.on_blue, "| #{url}/api/chat | " - currentTime = Time.local.to_unix - if params["stream"]? - print "stream | " - begin - agent.post(path: "/api/chat", body: params.to_json) do |response| - # context.send_resp response.body_io.try &.gets_to_end - - leadChunk = { - "choices" => [ - { - "delta" => { - "role" => "assistant", - "content" => "", - }, - "index" => 0, - "finish_reason" => nil, - }, - ], - "created" => currentTime, - "id" => "chatcmpl", - "object" => "chat.completion.chunk", - "model" => Mocker.gpt, - } - context - .send_resp("data: #{leadChunk.to_json}\n\n") - .response.flush - - # response.body_io.each_line.skip(1).each do |line| - response.body_io.each_line do |line| - chunk = JSON.parse line - - transformedChunk = { - "choices" => [ - { - "delta" => { - "content" => chunk["message"]["content"].as_s, - }, - "index" => 0, - "finish_reason" => chunk["done"].as_bool ? "stop" : nil, - }, - ], - "created" => currentTime, - "id" => "chatcmpl", - "object" => "chat.completion.chunk", - "model" => Mocker.gpt, - } - context - .send_resp("data: #{transformedChunk.to_json}\n\n") - .response.flush - end - puts " Done ".colorize.bright.on_green - context.send_resp "data: [DONE]" - end - rescue - puts " Fail ".colorize.bright.on_red - context - .put_status(HTTP::Status::BAD_REQUEST) - .halt - end - else - print "normal | " - begin - plainRequest = PlainRequest.from_json(params.to_json) - plainResponse = agent.post(path: "/api/chat", body: plainRequest.to_json) - respJson = JSON.parse plainResponse.body - gptResponse = PlainResponse.new( - created: currentTime, - choices: [ - Choice.new( - message: Message.new( - role: respJson["message"]["role"].as_s, - content: respJson["message"]["content"].as_s, - ), - ), - ], - usage: Usage.new( - prompt_tokens: respJson["prompt_eval_count"].as_i, - completion_tokens: respJson["eval_count"].as_i, - total_tokens: respJson["prompt_eval_count"].as_i + respJson["eval_count"].as_i, - ), - ) - puts " Done ".colorize.bright.on_green - context.send_resp gptResponse.to_json - rescue - puts " Fail ".colorize.bright.on_red - context - .put_status(HTTP::Status::BAD_REQUEST) - .halt - end - end - end -end +class MockGPT < Grip::Controllers::Http + def connect(context : Context) + context + .put_status(HTTP::Status::OK) + .halt + end + + def ollama(context : Context) + params = context.fetch_json_params + params["model"] = Mocker.model + presetUri = URI.parse "http://#{ENV["OLLAMA_HOST"]}" + url = URI.new scheme: presetUri.scheme || "http", host: presetUri.host || "localhost", port: presetUri.port || 11434 + agent = HTTP::Client.new url + agent.connect_timeout = 5.seconds + print " POST ".colorize.bright.on_blue, "| #{url}/api/chat | " + currentTime = Time.local.to_unix + if params["stream"]? + print "stream | " + begin + agent.post(path: "/api/chat", body: params.to_json) do |response| + # context.send_resp response.body_io.try &.gets_to_end + + leadChunk = { + "choices" => [ + { + "delta" => { + "role" => "assistant", + "content" => "", + }, + "index" => 0, + "finish_reason" => nil, + }, + ], + "created" => currentTime, + "id" => "chatcmpl", + "object" => "chat.completion.chunk", + "model" => Mocker.gpt, + } + context + .send_resp("data: #{leadChunk.to_json}\n\n") + .response.flush + + # response.body_io.each_line.skip(1).each do |line| + response.body_io.each_line do |line| + chunk = JSON.parse line + + transformedChunk = { + "choices" => [ + { + "delta" => { + "content" => chunk["message"]["content"].as_s, + }, + "index" => 0, + "finish_reason" => chunk["done"].as_bool ? "stop" : nil, + }, + ], + "created" => currentTime, + "id" => "chatcmpl", + "object" => "chat.completion.chunk", + "model" => Mocker.gpt, + } + context + .send_resp("data: #{transformedChunk.to_json}\n\n") + .response.flush + end + puts " Done ".colorize.bright.on_green + context.send_resp "data: [DONE]" + end + rescue + puts " Fail ".colorize.bright.on_red + context + .put_status(HTTP::Status::BAD_REQUEST) + .halt + end + else + print "normal | " + begin + plainRequest = PlainRequest.from_json(params.to_json) + plainResponse = agent.post(path: "/api/chat", body: plainRequest.to_json) + respJson = JSON.parse plainResponse.body + gptResponse = PlainResponse.new( + created: currentTime, + choices: [ + Choice.new( + message: Message.new( + role: respJson["message"]["role"].as_s, + content: respJson["message"]["content"].as_s, + ), + ), + ], + usage: Usage.new( + prompt_tokens: respJson["prompt_eval_count"].as_i, + completion_tokens: respJson["eval_count"].as_i, + total_tokens: respJson["prompt_eval_count"].as_i + respJson["eval_count"].as_i, + ), + ) + puts " Done ".colorize.bright.on_green + context.send_resp gptResponse.to_json + rescue + puts " Fail ".colorize.bright.on_red + context + .put_status(HTTP::Status::BAD_REQUEST) + .halt + end + end + end +end diff --git a/src/mockgpt/cors.cr b/src/mockgpt/cors.cr index 4c899ec..08b5b46 100644 --- a/src/mockgpt/cors.cr +++ b/src/mockgpt/cors.cr @@ -1,21 +1,21 @@ -class CrossOriginResourceSharing - include HTTP::Handler - - def call(context : HTTP::Server::Context) - context.response.headers.add "Server", "Grip/v2" - context.response.headers.add "Access-Control-Allow-Origin", "*" - context.response.headers.add "Access-Control-Allow-Headers", "*" - context.response.headers.add "Access-Control-Allow-Credentials", "true" - context.response.headers.add "Content-Type", "text/event-stream; charset=utf-8" - context.response.headers.add "Cache-Control", "no-cache" - context.response.headers.add "X-Accel-Buffering", "no" - - unless context.request.method.in? ["POST", "OPTIONS"] - context.response.headers.add "Access-Control-Allow-Methods", "POST" - - return context.put_status(HTTP::Status::METHOD_NOT_ALLOWED) - end - - call_next(context) - end -end +class CrossOriginResourceSharing + include HTTP::Handler + + def call(context : HTTP::Server::Context) + context.response.headers.add "Server", "Grip/v2" + context.response.headers.add "Access-Control-Allow-Origin", "*" + context.response.headers.add "Access-Control-Allow-Headers", "*" + context.response.headers.add "Access-Control-Allow-Credentials", "true" + context.response.headers.add "Content-Type", "text/event-stream; charset=utf-8" + context.response.headers.add "Cache-Control", "no-cache" + context.response.headers.add "X-Accel-Buffering", "no" + + unless context.request.method.in? ["POST", "OPTIONS"] + context.response.headers.add "Access-Control-Allow-Methods", "POST" + + return context.put_status(HTTP::Status::METHOD_NOT_ALLOWED) + end + + call_next(context) + end +end diff --git a/src/mockgpt/routes.cr b/src/mockgpt/routes.cr index 92c7f5c..6f7ee95 100644 --- a/src/mockgpt/routes.cr +++ b/src/mockgpt/routes.cr @@ -1,16 +1,16 @@ -class Application < Grip::Application - def initialize(@host : String, @port : Int32) - super(environment: "production", serve_static: false) - - router.insert(1, CrossOriginResourceSharing.new) - - post "/v1/chat/completions", MockGPT, as: :ollama - options "/v1/chat/completions", MockGPT, as: :connect - end - - getter host : String - getter port : Int32 - getter reuse_port : Bool = true - getter fallthrough : Bool = true - getter directory_listing : Bool = false -end +class Application < Grip::Application + def initialize(@host : String, @port : Int32) + super(environment: "production", serve_static: false) + + router.insert(1, CrossOriginResourceSharing.new) + + post "/v1/chat/completions", MockGPT, as: :ollama + options "/v1/chat/completions", MockGPT, as: :connect + end + + getter host : String + getter port : Int32 + getter reuse_port : Bool = true + getter fallthrough : Bool = true + getter directory_listing : Bool = false +end diff --git a/src/mockgpt/struct/request.cr b/src/mockgpt/struct/request.cr index 3f5d77e..97e97d3 100644 --- a/src/mockgpt/struct/request.cr +++ b/src/mockgpt/struct/request.cr @@ -1,11 +1,11 @@ -record Dialog, role : String, content : String do - include JSON::Serializable -end - -struct PlainRequest - include JSON::Serializable - - property model : String = Mocker.model - property stream : Bool = false - property messages : Array(Dialog) -end +record Dialog, role : String, content : String do + include JSON::Serializable +end + +struct PlainRequest + include JSON::Serializable + + property model : String = Mocker.model + property stream : Bool = false + property messages : Array(Dialog) +end diff --git a/src/mockgpt/struct/response.cr b/src/mockgpt/struct/response.cr index 46c837a..0a4bd47 100644 --- a/src/mockgpt/struct/response.cr +++ b/src/mockgpt/struct/response.cr @@ -1,55 +1,55 @@ -struct PlainResponse - include JSON::Serializable - - def initialize(@created, @choices, @usage) - end - - @id : String = "chatcmpl" - - @object : String = "chat.completion" - - @model : String = Mocker.gpt - - property created : Int64 - - property choices : Array(Choice) - - property usage : Usage -end - -struct Choice - include JSON::Serializable - - def initialize(@message) - end - - property message : Message - - @index : Int32 = 0 - - @finish_reason : String = "stop" -end - -struct Message - include JSON::Serializable - - def initialize(@role, @content) - end - - property role : String - - property content : String -end - -struct Usage - include JSON::Serializable - - def initialize(@prompt_tokens, @completion_tokens, @total_tokens) - end - - property prompt_tokens : Int32 - - property completion_tokens : Int32 - - property total_tokens : Int32 -end +struct PlainResponse + include JSON::Serializable + + def initialize(@created, @choices, @usage) + end + + @id : String = "chatcmpl" + + @object : String = "chat.completion" + + @model : String = Mocker.gpt + + property created : Int64 + + property choices : Array(Choice) + + property usage : Usage +end + +struct Choice + include JSON::Serializable + + def initialize(@message) + end + + property message : Message + + @index : Int32 = 0 + + @finish_reason : String = "stop" +end + +struct Message + include JSON::Serializable + + def initialize(@role, @content) + end + + property role : String + + property content : String +end + +struct Usage + include JSON::Serializable + + def initialize(@prompt_tokens, @completion_tokens, @total_tokens) + end + + property prompt_tokens : Int32 + + property completion_tokens : Int32 + + property total_tokens : Int32 +end