Skip to content

Commit 7ab303d

Browse files
GrigoriyPAGrigoriyPA
authored
Merge 5593ac6 into b8bc549 
2 parents b8bc549 + 5593ac6 commit 7ab303d

File tree

16 files changed

+745
-6
lines changed

16 files changed

+745
-6
lines changed

‎ydb/library/yql/providers/common/http_gateway/yql_http_gateway.h‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,21 +2,21 @@
22

33
#include "yql_http_header.h"
44

5-
#include <yql/essentials/providers/common/proto/gateways_config.pb.h>
6-
75
#include <yql/essentials/public/issue/yql_issue.h>
6+
7+
#include <contrib/libs/curl/include/curl/curl.h>
8+
89
#include <library/cpp/containers/stack_vector/stack_vec.h>
910
#include <library/cpp/monlib/dynamic_counters/counters.h>
1011
#include <library/cpp/retry/retry_policy.h>
1112

12-
#include <contrib/libs/curl/include/curl/curl.h>
13-
1413
#include <atomic>
15-
#include <variant>
1614
#include <functional>
1715

1816
namespace NYql {
1917

18+
class THttpGatewayConfig;
19+
2020
class IHTTPGateway {
2121
public:
2222
using TPtr = std::shared_ptr<IHTTPGateway>;
@@ -142,4 +142,4 @@ class IHTTPGateway {
142142
const TString& awsSigV4 = {});
143143
};
144144

145-
}
145+
} // namespace NYql

‎ydb/public/lib/ydb_cli/commands/ya.make‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ SRCS(
1111
topic_write_scenario.cpp
1212
topic_readwrite_scenario.cpp
1313
ydb_admin.cpp
14+
ydb_ai.cpp
1415
ydb_benchmark.cpp
1516
ydb_bridge.cpp
1617
ydb_cluster.cpp
@@ -60,6 +61,9 @@ PEERDIR(
6061
ydb/public/lib/ydb_cli/commands/sdk_core_access
6162
ydb/public/lib/ydb_cli/commands/topic_workload
6263
ydb/public/lib/ydb_cli/commands/transfer_workload
64+
ydb/public/lib/ydb_cli/commands/ydb_ai
65+
ydb/public/lib/ydb_cli/commands/ydb_ai/models
66+
ydb/public/lib/ydb_cli/commands/ydb_ai/tools
6367
ydb/public/lib/ydb_cli/commands/ydb_discovery
6468
ydb/public/lib/ydb_cli/common
6569
ydb/public/lib/ydb_cli/dump
@@ -95,5 +99,6 @@ RECURSE(
9599
sdk_core_access
96100
topic_workload
97101
transfer_workload
102+
ydb_ai
98103
ydb_discovery
99104
)

‎ydb/public/lib/ydb_cli/commands/ydb_ai.cpp‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
#include "ydb_ai.h"
2+
3+
#include <ydb/public/lib/ydb_cli/commands/ydb_ai/line_reader.h>
4+
#include <ydb/public/lib/ydb_cli/commands/ydb_ai/models/model_openai.h>
5+
#include <ydb/public/lib/ydb_cli/commands/ydb_ai/tools/exec_query_tool.h>
6+
7+
#include <util/string/strip.h>
8+
#include <util/system/env.h>
9+
10+
namespace NYdb::NConsoleClient {
11+
12+
namespace {
13+
14+
void PrintExitMessage() {
15+
Cout << "\nBye" << Endl;
16+
}
17+
18+
} // anonymous namespace
19+
20+
TCommandAi::TCommandAi()
21+
: TBase("ai", {}, "AI-TODO: KIKIMR-24198 -- description")
22+
{}
23+
24+
void TCommandAi::Config(TConfig& config) {
25+
TBase::Config(config);
26+
config.Opts->SetTitle("AI-TODO: KIKIMR-24198 -- title");
27+
config.Opts->SetFreeArgsNum(0);
28+
}
29+
30+
int TCommandAi::Run(TConfig& config) {
31+
Y_UNUSED(config);
32+
33+
Cout << "AI-TODO: KIKIMR-24198 -- welcome message" << Endl;
34+
35+
// AI-TODO: KIKIMR-24202 - robust file creation
36+
NAi::TLineReader lineReader("ydb-ai> ", (TFsPath(HomeDir) / ".ydb-ai/history").GetPath());
37+
38+
// DeepSeek
39+
// const auto model = NAi::CreateOpenAiModel({
40+
// .BaseUrl = "https://api.eliza.yandex.net/raw/internal/deepseek/v1", // AI-TODO: KIKIMR-24214 -- configure it
41+
// .ModelId = "deepseek-0324", // AI-TODO: KIKIMR-24214 -- configure it
42+
// .ApiKey = GetEnv("MODEL_TOKEN"), // AI-TODO: KIKIMR-24214 -- configure it
43+
// });
44+
45+
// YandexGPT Pro
46+
const auto model = NAi::CreateOpenAiModel({
47+
.BaseUrl = "https://api.eliza.yandex.net/internal/zeliboba/32b_aligned_quantized_202506/generative/v1", // AI-TODO: KIKIMR-24214 -- configure it
48+
.ApiKey = GetEnv("MODEL_TOKEN"), // AI-TODO: KIKIMR-24214 -- configure it
49+
});
50+
51+
const auto sqlTool = NAi::CreateExecQueryTool(config); // AI-TODO: more generic tools registration
52+
model->RegisterTool(sqlTool->GetName(), sqlTool->GetParametersSchema(), sqlTool->GetDescription());
53+
54+
// AI-TODO: there is strange highlighting of brackets
55+
while (const auto& maybeLine = lineReader.ReadLine()) {
56+
TString input = *maybeLine;
57+
if (input.empty()) {
58+
continue;
59+
}
60+
61+
if (IsIn({"quit", "exit"}, to_lower(input))) {
62+
PrintExitMessage();
63+
return EXIT_SUCCESS;
64+
}
65+
66+
// AI-TODO: limit interaction number
67+
std::optional<TString> toolCallId;
68+
while (input) {
69+
// AI-TODO: progress visualization
70+
auto output = model->HandleMessage(input);
71+
Y_ENSURE(output.Text || output.ToolCall);
72+
73+
if (output.Text) {
74+
// AI-TODO: proper answer format
75+
Cout << "Model answer:\n" << *output.Text << Endl;
76+
}
77+
78+
if (!output.ToolCall) {
79+
break;
80+
}
81+
82+
const auto& toolCall = *output.ToolCall;
83+
if (toolCall.Name != sqlTool->GetName()) {
84+
// AI-TODO: proper wrong tool handling
85+
Cout << "Unsupported tool: " << toolCall.Name << Endl;
86+
return EXIT_FAILURE;
87+
}
88+
89+
// AI-TODO: ask permission before call and show progress
90+
toolCallId = toolCall.Id;
91+
input = sqlTool->Execute(toolCall.Parameters);
92+
}
93+
}
94+
95+
PrintExitMessage();
96+
return EXIT_SUCCESS;
97+
}
98+
99+
} // namespace NYdb::NConsoleClient

‎ydb/public/lib/ydb_cli/commands/ydb_ai.h‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
#pragma once
2+
3+
#include "ydb_command.h"
4+
5+
namespace NYdb::NConsoleClient {
6+
7+
class TCommandAi final : public TYdbCommand {
8+
using TBase = TYdbCommand;
9+
10+
public:
11+
TCommandAi();
12+
13+
void Config(TConfig& config) final;
14+
15+
int Run(TConfig& config) final;
16+
};
17+
18+
} // namespace NYdb::NConsoleClient

‎ydb/public/lib/ydb_cli/commands/ydb_ai/line_reader.cpp‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
#include "line_reader.h"
2+
3+
#include <util/string/strip.h>
4+
5+
namespace NYdb::NConsoleClient::NAi {
6+
7+
TLineReader::TLineReader(TString prompt, TString historyFilePath)
8+
: Prompt(std::move(prompt))
9+
, HistoryFilePath(historyFilePath)
10+
, HistoryFileLock(historyFilePath)
11+
{
12+
Rx.install_window_change_handler();
13+
14+
Rx.bind_key(replxx::Replxx::KEY::control('N'), [this](char32_t code) {
15+
return Rx.invoke(replxx::Replxx::ACTION::HISTORY_NEXT, code);
16+
});
17+
Rx.bind_key(replxx::Replxx::KEY::control('P'), [this](char32_t code) {
18+
return Rx.invoke(replxx::Replxx::ACTION::HISTORY_PREVIOUS, code);
19+
});
20+
Rx.bind_key(replxx::Replxx::KEY::control('D'), [](char32_t) {
21+
return replxx::Replxx::ACTION_RESULT::BAIL;
22+
});
23+
Rx.bind_key(replxx::Replxx::KEY::control('J'), [this](char32_t code) {
24+
return Rx.invoke(replxx::Replxx::ACTION::COMMIT_LINE, code);
25+
});
26+
27+
Rx.enable_bracketed_paste();
28+
Rx.set_unique_history(true);
29+
30+
if (const auto guard = TryLockHistory(); guard && !Rx.history_load(HistoryFilePath)) {
31+
Rx.print("Loading history failed: %s\n", strerror(errno));
32+
}
33+
}
34+
35+
std::optional<TString> TLineReader::ReadLine() {
36+
do {
37+
if (const char* input = Rx.input(Prompt.c_str())) {
38+
auto result = Strip(input);
39+
AddToHistory(result);
40+
return std::move(result);
41+
}
42+
} while (errno == EAGAIN);
43+
44+
return std::nullopt;
45+
}
46+
47+
TTryGuard<TFileLock> TLineReader::TryLockHistory() {
48+
// AI-TODO: KIKIMR-24202 - robust file creation and handling
49+
TTryGuard guard(HistoryFileLock);
50+
51+
if (!guard) {
52+
Rx.print("Lock of history file failed: %s\n", strerror(errno));
53+
}
54+
55+
return guard;
56+
}
57+
58+
void TLineReader::AddToHistory(const TString& line) {
59+
Rx.history_add(line);
60+
61+
if (const auto guard = TryLockHistory(); guard && !Rx.history_save(HistoryFilePath)) {
62+
Rx.print("Save history failed: %s\n", strerror(errno));
63+
}
64+
}
65+
66+
} // namespace NYdb::NConsoleClient::NAi

‎ydb/public/lib/ydb_cli/commands/ydb_ai/line_reader.h‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
#pragma once
2+
3+
#include <contrib/restricted/patched/replxx/include/replxx.hxx>
4+
5+
#include <util/generic/string.h>
6+
#include <util/system/file_lock.h>
7+
8+
#include <optional>
9+
10+
namespace NYdb::NConsoleClient::NAi {
11+
12+
class TLineReader {
13+
public:
14+
TLineReader(TString prompt, TString historyFilePath);
15+
16+
std::optional<TString> ReadLine();
17+
18+
private:
19+
TTryGuard<TFileLock> TryLockHistory();
20+
21+
void AddToHistory(const TString& line);
22+
23+
private:
24+
TString Prompt;
25+
TString HistoryFilePath;
26+
TFileLock HistoryFileLock;
27+
replxx::Replxx Rx;
28+
};
29+
30+
} // namespace NYdb::NConsoleClient::NAi

‎ydb/public/lib/ydb_cli/commands/ydb_ai/models/model_interface.h‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
#pragma once
2+
3+
#include <library/cpp/json/writer/json_value.h>
4+
5+
#include <util/generic/string.h>
6+
7+
#include <memory>
8+
9+
namespace NYdb::NConsoleClient::NAi {
10+
11+
class IModel {
12+
public:
13+
using TPtr = std::shared_ptr<IModel>;
14+
15+
virtual ~IModel() = default;
16+
17+
struct TResponse {
18+
struct TToolCall {
19+
TString Id;
20+
TString Name;
21+
NJson::TJsonValue Parameters;
22+
};
23+
24+
std::optional<TString> Text;
25+
std::optional<TToolCall> ToolCall;
26+
};
27+
28+
virtual TResponse HandleMessage(const TString& input, std::optional<TString> toolCallId = std::nullopt) = 0;
29+
30+
virtual void RegisterTool(const TString& name, const NJson::TJsonValue& parametersSchema, const TString& description) = 0;
31+
};
32+
33+
} // namespace NYdb::NConsoleClient::NAi

0 commit comments

Comments
 (0)