diff --git a/cmd/pievr/main.go b/cmd/pievr/main.go index b429e2c..fa9806d 100644 --- a/cmd/pievr/main.go +++ b/cmd/pievr/main.go @@ -4,10 +4,12 @@ import ( "context" "fmt" "log" + "os" + "os/signal" - "github.com/google/generative-ai-go/genai" - "google.golang.org/api/option" + "pievr/internal/bot" "pievr/internal/config" + "pievr/internal/gemini" ) func main() { @@ -15,22 +17,25 @@ func main() { cfg, err := config.Load("config.json") if err != nil { - log.Fatalf("Failed to load config: %v", err) + log.Fatalf("Failed to load config: %v (make sure config.json exists, see config.json.example)", err) } - ctx := context.Background() - opts := []option.ClientOption{option.WithAPIKey(cfg.GeminiApiKey)} - if cfg.GeminiBaseURL != "" { - log.Printf("Using custom Gemini BaseURL: %s", cfg.GeminiBaseURL) - // Custom base URL can be passed via transport options if needed in future, - // for now we log it. - } + ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt) + defer cancel() - client, err := genai.NewClient(ctx, opts...) + // Initialize Gemini client + geminiCli, err := gemini.NewClient(ctx, cfg.GeminiApiKey, cfg.GeminiBaseURL) if err != nil { - log.Fatalf("Failed to create generative AI client: %v", err) + log.Fatalf("Failed to create Gemini client: %v", err) } - defer client.Close() + defer geminiCli.Close() - fmt.Println("Configuration and Gemini client loaded successfully.") + // Initialize Telegram bot + telegramBot, err := bot.New(cfg.TelegramToken, geminiCli) + if err != nil { + log.Fatalf("Failed to create Telegram bot: %v", err) + } + + // Start bot (blocking) + telegramBot.Start(ctx) } diff --git a/go.mod b/go.mod index 345bb99..68e0a3d 100644 --- a/go.mod +++ b/go.mod @@ -13,6 +13,7 @@ require ( github.com/felixge/httpsnoop v1.0.4 // indirect github.com/go-logr/logr v1.4.3 // indirect github.com/go-logr/stdr v1.2.2 // indirect + github.com/go-telegram/bot v1.23.0 // indirect github.com/google/generative-ai-go v0.20.1 // indirect github.com/google/s2a-go v0.1.9 // indirect github.com/google/uuid v1.6.0 // indirect diff --git a/go.sum b/go.sum index 1a888fa..66b85ff 100644 --- a/go.sum +++ b/go.sum @@ -19,6 +19,8 @@ github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= +github.com/go-telegram/bot v1.23.0 h1:CKKQq115G/GUGBG8uuWl5uXbiBHyVjZBp/qqOLWZjJk= +github.com/go-telegram/bot v1.23.0/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM= github.com/google/generative-ai-go v0.20.1 h1:6dEIujpgN2V0PgLhr6c/M1ynRdc7ARtiIDPFzj45uNQ= github.com/google/generative-ai-go v0.20.1/go.mod h1:TjOnZJmZKzarWbjUJgy+r3Ee7HGBRVLhOIgupnwR4Bg= github.com/google/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0= diff --git a/internal/bot/bot.go b/internal/bot/bot.go new file mode 100644 index 0000000..bcae4db --- /dev/null +++ b/internal/bot/bot.go @@ -0,0 +1,65 @@ +package bot + +import ( + "context" + "fmt" + "log" + + "github.com/go-telegram/bot" + "github.com/go-telegram/bot/models" + "pievr/internal/gemini" +) + +type Bot struct { + b *bot.Bot + geminiCli *gemini.Client +} + +func New(token string, geminiCli *gemini.Client) (*Bot, error) { + myBot := &Bot{ + geminiCli: geminiCli, + } + + opts := []bot.Option{ + bot.WithDefaultHandler(myBot.handler), + } + + b, err := bot.New(token, opts...) + if err != nil { + return nil, fmt.Errorf("failed to create telegram bot: %w", err) + } + + myBot.b = b + return myBot, nil +} + +func (s *Bot) Start(ctx context.Context) { + log.Println("Telegram bot is starting...") + s.b.Start(ctx) +} + +func (s *Bot) handler(ctx context.Context, b *bot.Bot, update *models.Update) { + if update.Message == nil || update.Message.Text == "" { + return + } + + userMsg := update.Message.Text + chatID := update.Message.Chat.ID + + log.Printf("Received message from %d: %s", chatID, userMsg) + + // Send prompt to Gemini + reply, err := s.geminiCli.Generate(ctx, userMsg) + if err != nil { + log.Printf("Failed to generate response from Gemini: %v", err) + reply = "Извините, произошла ошибка при обращении к LLM." + } + + _, err = b.SendMessage(ctx, &bot.SendMessageParams{ + ChatID: chatID, + Text: reply, + }) + if err != nil { + log.Printf("Failed to send telegram message: %v", err) + } +} diff --git a/internal/gemini/client.go b/internal/gemini/client.go new file mode 100644 index 0000000..dd0e96b --- /dev/null +++ b/internal/gemini/client.go @@ -0,0 +1,90 @@ +package gemini + +import ( + "context" + "fmt" + "net/http" + "net/url" + + "github.com/google/generative-ai-go/genai" + "google.golang.org/api/option" +) + +type Client struct { + client *genai.Client + model string +} + +func NewClient(ctx context.Context, apiKey, baseURL string) (*Client, error) { + var opts []option.ClientOption + if apiKey != "" { + opts = append(opts, option.WithAPIKey(apiKey)) + } + + if baseURL != "" { + parsedURL, err := url.Parse(baseURL) + if err != nil { + return nil, fmt.Errorf("invalid base url: %w", err) + } + // Custom HTTP client endpoint override for Gemini API proxy/baseURL + customClient := &http.Client{ + Transport: &baseURITransport{ + BaseURL: parsedURL, + Transport: http.DefaultTransport, + }, + } + opts = append(opts, option.WithHTTPClient(customClient)) + } + + c, err := genai.NewClient(ctx, opts...) + if err != nil { + return nil, err + } + + return &Client{ + client: c, + model: "gemini-2.5-flash", // Default model for now + }, nil +} + +func (c *Client) Generate(ctx context.Context, prompt string) (string, error) { + model := c.client.GenerativeModel(c.model) + resp, err := model.GenerateContent(ctx, genai.Text(prompt)) + if err != nil { + return "", err + } + + if len(resp.Candidates) == 0 || len(resp.Candidates[0].Content.Parts) == 0 { + return "", fmt.Errorf("empty response from Gemini") + } + + part := resp.Candidates[0].Content.Parts[0] + if textPart, ok := part.(genai.Text); ok { + return string(textPart), nil + } + + return fmt.Sprintf("%v", part), nil +} + +func (c *Client) Close() { + c.client.Close() +} + +type baseURITransport struct { + BaseURL *url.URL + Transport http.RoundTripper +} + +func (t *baseURITransport) RoundTrip(req *http.Request) (*http.Response, error) { + // Rewrite request URL to use custom BaseURL while preserving path & query + req.URL.Scheme = t.BaseURL.Scheme + req.URL.Host = t.BaseURL.Host + if t.BaseURL.Path != "" { + req.URL.Path = t.BaseURL.Path + req.URL.Path + } + tr := t.Transport + if tr == nil { + tr = http.DefaultTransport + } + return tr.RoundTrip(req) +}