knowledge-graph-agent/backend/internal/llm/client.go

package llm

import (
	"context"
	"fmt"
	"knowledge-graph-backend/internal/config"
	"log"

	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/llms/openai"
)

type Client struct {
	model  llms.Model
	config config.LLMConfig
}

func NewClient(cfg config.LLMConfig) (*Client, error) {
	model, err := openai.New(
		openai.WithToken(cfg.APIKey),
		openai.WithBaseURL(cfg.BaseURL),
		openai.WithModel(cfg.Model),
	)

	if err != nil {
		return nil, fmt.Errorf("failed to create llm client: %w", err)
	}

	return &Client{config: cfg, model: model}, nil
}

func (c *Client) Generate(ctx context.Context, messages []llms.MessageContent, opts ...llms.CallOption) (*llms.ContentResponse, error) {
	resp, err := c.model.GenerateContent(ctx, messages, opts...)
	if err != nil {
		log.Printf("[ERROR] GenerateContent failed: %v", err)
		return nil, fmt.Errorf("generate content error: %w", err)
	}
	if resp == nil || len(resp.Choices) == 0 {
		log.Printf("[WARN] Response is empty or has no choices")
		return nil, fmt.Errorf("no content choices returned")
	}
	return resp, nil
}

func (c *Client) Chat(ctx context.Context, messages []llms.MessageContent) (string, error) {
	resp, err := c.Generate(ctx, messages)
	if err != nil {
		return "", err
	}
	log.Printf("[INFO] Chat completed successfully, choice length: %d", len(resp.Choices[0].Content))
	return resp.Choices[0].Content, nil
}

func (c *Client) ChatWithSystemPrompt(ctx context.Context, systemPrompt string, messages []llms.MessageContent) (string, error) {
	allMessages := NewMessageBuilder().
		WithSystemPrompt(systemPrompt).
		WithMessages(messages).
		Build()

	resp, err := c.Generate(ctx, allMessages)
	if err != nil {
		return "", err
	}
	log.Printf("[INFO] Chat with system prompt completed successfully, choice length: %d", len(resp.Choices[0].Content))
	return resp.Choices[0].Content, nil
}