Skip to content

Commit 13210c1

Browse files
GrigoriyPAGrigoriyPA
authored
Merge a2c7388 into 9c0c0f6 
2 parents 9c0c0f6 + a2c7388 commit 13210c1

27 files changed

+1507
-6
lines changed

‎ydb/core/base/validation.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
#pragma once
2+
3+
#ifndef NDEBUG
4+
#define Y_DEBUG_VERIFY Y_ABORT_UNLESS
5+
#else
6+
#define Y_DEBUG_VERIFY Y_VERIFY
7+
#endif

‎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: 6 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,10 @@ 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/common
66+
ydb/public/lib/ydb_cli/commands/ydb_ai/models
67+
ydb/public/lib/ydb_cli/commands/ydb_ai/tools
6368
ydb/public/lib/ydb_cli/commands/ydb_discovery
6469
ydb/public/lib/ydb_cli/common
6570
ydb/public/lib/ydb_cli/dump
@@ -95,5 +100,6 @@ RECURSE(
95100
sdk_core_access
96101
topic_workload
97102
transfer_workload
103+
ydb_ai
98104
ydb_discovery
99105
)

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

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
#include "ydb_ai.h"
2+
3+
#include <ydb/public/lib/ydb_cli/commands/ydb_ai/common/json_utils.h>
4+
#include <ydb/public/lib/ydb_cli/commands/ydb_ai/line_reader.h>
5+
#include <ydb/public/lib/ydb_cli/commands/ydb_ai/models/model_anthropic.h>
6+
#include <ydb/public/lib/ydb_cli/commands/ydb_ai/models/model_openai.h>
7+
#include <ydb/public/lib/ydb_cli/commands/ydb_ai/tools/exec_query_tool.h>
8+
#include <ydb/public/lib/ydb_cli/commands/ydb_ai/tools/list_directory_tool.h>
9+
10+
#include <util/string/strip.h>
11+
#include <util/system/env.h>
12+
13+
/*
14+
15+
FEATURES-TODO:
16+
17+
- Streamable model response printing
18+
- Streamable results printing
19+
- Adjusting errors, progress and response printing
20+
- Approving before tool use
21+
- Integration into common interactive mode
22+
- Think about helps
23+
- Think about robust
24+
- Provide system promt
25+
- Somehow render markdown
26+
27+
*/
28+
29+
namespace NYdb::NConsoleClient {
30+
31+
namespace {
32+
33+
void PrintExitMessage() {
34+
Cout << "\nBye" << Endl;
35+
}
36+
37+
} // anonymous namespace
38+
39+
TCommandAi::TCommandAi()
40+
: TBase("ai", {}, "AI-TODO: KIKIMR-24198 -- description")
41+
{}
42+
43+
void TCommandAi::Config(TConfig& config) {
44+
TBase::Config(config);
45+
config.Opts->SetTitle("AI-TODO: KIKIMR-24198 -- title");
46+
config.Opts->SetFreeArgsNum(0);
47+
}
48+
49+
int TCommandAi::Run(TConfig& config) {
50+
Cout << "AI-TODO: KIKIMR-24198 -- welcome message" << Endl;
51+
52+
// AI-TODO: KIKIMR-24202 - robust file creation
53+
NAi::TLineReader lineReader("ydb-ai> ", (TFsPath(HomeDir) / ".ydb-ai/history").GetPath());
54+
55+
// DeepSeek
56+
// const auto model = NAi::CreateOpenAiModel({
57+
// .BaseUrl = "https://api.eliza.yandex.net/raw/internal/deepseek", // AI-TODO: KIKIMR-24214 -- configure it
58+
// .ModelId = "deepseek-0324", // AI-TODO: KIKIMR-24214 -- configure it
59+
// .ApiKey = GetEnv("MODEL_TOKEN"), // AI-TODO: KIKIMR-24214 -- configure it
60+
// }, config);
61+
62+
// Claude 3.5 haiku
63+
// const auto model = NAi::CreateAnthropicModel({
64+
// .BaseUrl = "https://api.eliza.yandex.net/anthropic", // AI-TODO: KIKIMR-24214 -- configure it
65+
// .ModelId = "claude-3-5-haiku-20241022",
66+
// .ApiKey = GetEnv("MODEL_TOKEN"), // AI-TODO: KIKIMR-24214 -- configure it
67+
// }, config);
68+
69+
// YandexGPT Pro
70+
const auto model = NAi::CreateOpenAiModel({
71+
.BaseUrl = "https://api.eliza.yandex.net/internal/zeliboba/32b_aligned_quantized_202506/generative", // AI-TODO: KIKIMR-24214 -- configure it
72+
.ApiKey = GetEnv("MODEL_TOKEN"), // AI-TODO: KIKIMR-24214 -- configure it
73+
}, config);
74+
75+
std::unordered_map<TString, NAi::ITool::TPtr> tools = {
76+
{"execute_query", NAi::CreateExecQueryTool(config)},
77+
{"list_directory", NAi::CreateListDirectoryTool(config)},
78+
};
79+
for (const auto& [name, tool] : tools) {
80+
model->RegisterTool(name, tool->GetParametersSchema(), tool->GetDescription());
81+
}
82+
83+
// AI-TODO: there is strange highlighting of brackets
84+
std::vector<NAi::IModel::TMessage> messages;
85+
while (const auto& maybeLine = lineReader.ReadLine()) {
86+
const auto& input = *maybeLine;
87+
if (input.empty()) {
88+
continue;
89+
}
90+
91+
if (IsIn({"quit", "exit"}, to_lower(input))) {
92+
PrintExitMessage();
93+
return EXIT_SUCCESS;
94+
}
95+
96+
// AI-TODO: limit interaction number
97+
messages.emplace_back(NAi::IModel::TUserMessage{.Text = input});
98+
while (!messages.empty()) {
99+
// AI-TODO: progress visualization
100+
const auto output = model->HandleMessages(messages);
101+
messages.clear();
102+
103+
if (!output.Text && output.ToolCalls.empty()) {
104+
// AI-TODO: proper answer format
105+
Cout << "Model answer is empty(" << Endl;
106+
break;
107+
}
108+
109+
if (output.Text) {
110+
// AI-TODO: proper answer format
111+
Cout << "Model answer:\n" << output.Text << Endl;
112+
}
113+
114+
for (const auto& toolCall : output.ToolCalls) {
115+
const auto it = tools.find(toolCall.Name);
116+
if (it == tools.end()) {
117+
// AI-TODO: proper wrong tool handling
118+
Cout << "Unsupported tool: " << toolCall.Name << Endl;
119+
return EXIT_FAILURE;
120+
}
121+
122+
// AI-TODO: proper tool call printing
123+
Cout << "Calling tool: " << toolCall.Name << " with params:\n" << NAi::FormatJsonValue(toolCall.Parameters) << Endl;
124+
125+
// AI-TODO: add approving
126+
const auto& result = it->second->Execute(toolCall.Parameters);
127+
if (!result.IsSuccess) {
128+
// AI-TODO: proper error handling
129+
Cout << result.Text << Endl;
130+
}
131+
// AI-TODO: show progress
132+
133+
messages.push_back(NAi::IModel::TToolResponse{
134+
.Text = result.Text,
135+
.ToolCallId = toolCall.Id,
136+
.IsSuccess = result.IsSuccess,
137+
});
138+
}
139+
}
140+
}
141+
142+
PrintExitMessage();
143+
return EXIT_SUCCESS;
144+
}
145+
146+
} // 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

0 commit comments

Comments
 (0)