cmd/agent-server/main.go

// Package main provides the entry point for the Agent Orchestrator service.
// This is the primary service that coordinates query processing and orchestrates
// interactions between the Retriever and Generator services.
package main

import (
	"context"
	"fmt"
	"net"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/AmaniQuery/amaniquery/internal/agent"
	"github.com/AmaniQuery/amaniquery/internal/cache"
	"github.com/AmaniQuery/amaniquery/internal/generator"
	"github.com/AmaniQuery/amaniquery/internal/generator/llm"
	"github.com/AmaniQuery/amaniquery/internal/retriever"
	"github.com/AmaniQuery/amaniquery/internal/retriever/keyword"
	"github.com/AmaniQuery/amaniquery/internal/retriever/vector"
	"github.com/AmaniQuery/amaniquery/internal/router"
	"github.com/AmaniQuery/amaniquery/pkg/config"
	"github.com/AmaniQuery/amaniquery/pkg/observability"

	"go.uber.org/zap"
	"google.golang.org/grpc"
	"google.golang.org/grpc/health"
	"google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/reflection"
)

func main() {
	// Load configuration
	cfg, err := config.Load()
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to load configuration: %v\n", err)
		os.Exit(1)
	}

	// Initialize logger
	logger, err := observability.NewLogger(cfg.Observability.LogLevel, cfg.Observability.LogFormat)
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to initialize logger: %v\n", err)
		os.Exit(1)
	}
	defer logger.Sync()

	logger.Info("starting AmaniQuery agent server",
		zap.String("version", cfg.Version),
		zap.String("environment", cfg.Environment),
		zap.Strings("llm_providers", cfg.LLM.GetConfiguredProviders()),
	)

	// Initialize observability (tracing + metrics)
	tracingShutdown, err := observability.InitProvider(observability.Config{
		ServiceName:     "amaniquery-agent",
		ServiceVersion:  cfg.Version,
		TracingEnabled:  cfg.Observability.TracingEnabled,
		TracingEndpoint: cfg.Observability.TracingEndpoint,
		MetricsEnabled:  cfg.Observability.MetricsEnabled,
		MetricsPort:     cfg.Observability.MetricsPort,
	})
	if err != nil {
		logger.Warn("failed to initialize tracing, continuing without", zap.Error(err))
	} else {
		defer tracingShutdown(context.Background())
	}

	// Start metrics server
	if cfg.Observability.MetricsEnabled {
		metricsServer := observability.StartMetricsServer(cfg.Observability.MetricsPort)
		defer metricsServer.Shutdown(context.Background())
		logger.Info("metrics server started", zap.Int("port", cfg.Observability.MetricsPort))
	}

	// Build dependencies
	deps, err := buildDependencies(cfg, logger)
	if err != nil {
		logger.Fatal("failed to build dependencies", zap.Error(err))
	}
	defer deps.Close()

	// Create gRPC server with interceptors
	grpcServer := grpc.NewServer(
		grpc.ChainUnaryInterceptor(
			observability.UnaryServerInterceptor(),
			agent.LoggingInterceptor(logger),
			agent.RecoveryInterceptor(),
		),
		grpc.ChainStreamInterceptor(
			observability.StreamServerInterceptor(),
		),
	)

	// Register Agent service
	agentServer := agent.NewServer(deps, logger)
	agent.RegisterAgentServiceServer(grpcServer, agentServer)

	// Register health service
	healthServer := health.NewServer()
	grpc_health_v1.RegisterHealthServer(grpcServer, healthServer)
	healthServer.SetServingStatus("", grpc_health_v1.HealthCheckResponse_SERVING)

	// Enable reflection for grpcurl
	reflection.Register(grpcServer)

	// Start gRPC server
	addr := fmt.Sprintf(":%d", cfg.Server.GRPCPort)
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		logger.Fatal("failed to listen", zap.String("addr", addr), zap.Error(err))
	}

	logger.Info("gRPC server starting",
		zap.String("addr", addr),
		zap.Int("http_port", cfg.Server.HTTPPort),
	)

	// Graceful shutdown handling
	errChan := make(chan error, 1)
	go func() {
		errChan <- grpcServer.Serve(listener)
	}()

	// Wait for shutdown signal
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)

	select {
	case err := <-errChan:
		logger.Fatal("server error", zap.Error(err))
	case sig := <-quit:
		logger.Info("shutting down", zap.String("signal", sig.String()))
	}

	// Graceful shutdown with timeout
	ctx, cancel := context.WithTimeout(context.Background(), cfg.Server.GracefulTimeout)
	defer cancel()

	healthServer.SetServingStatus("", grpc_health_v1.HealthCheckResponse_NOT_SERVING)
	grpcServer.GracefulStop()

	logger.Info("server stopped gracefully")
	_ = ctx // used for cleanup operations
}

// buildDependencies initializes all service dependencies
func buildDependencies(cfg *config.Config, logger *zap.Logger) (*agent.Dependencies, error) {
	deps := &agent.Dependencies{}

	// Initialize vector store client (Qdrant)
	logger.Info("connecting to vector store",
		zap.String("type", cfg.VectorStore.Type),
		zap.String("host", cfg.VectorStore.Host),
		zap.Int("port", cfg.VectorStore.Port),
	)
	vectorClient, err := vector.NewQdrantClient(vector.Config{
		Host:       cfg.VectorStore.Host,
		Port:       cfg.VectorStore.Port,
		APIKey:     cfg.VectorStore.APIKey,
		Collection: cfg.VectorStore.Collection,
		Dimension:  cfg.VectorStore.Dimension,
		Distance:   cfg.VectorStore.Distance,
	})
	if err != nil {
		logger.Warn("failed to connect to vector store, continuing without", zap.Error(err))
	} else {
		deps.VectorStore = vectorClient
		logger.Info("connected to vector store")
	}

	// Initialize cache
	logger.Info("connecting to cache", zap.String("url", cfg.Cache.RedisURL))
	cacheClient, err := cache.New(cache.Config{
		RedisURL:   cfg.Cache.RedisURL,
		LocalSize:  cfg.Cache.LocalSize,
		TTL:        cfg.Cache.TTL,
		MaxRetries: cfg.Cache.MaxRetries,
		PoolSize:   cfg.Cache.PoolSize,
	})
	if err != nil {
		logger.Warn("failed to initialize cache, continuing without", zap.Error(err))
	} else {
		deps.Cache = cacheClient
		logger.Info("cache initialized")
	}

	// Initialize embedding client
	logger.Info("initializing embedding client",
		zap.String("provider", cfg.Embedding.Provider),
		zap.String("model", cfg.Embedding.Model),
	)
	embeddingClient := generator.NewOpenAIEmbeddingClient(generator.EmbeddingConfig{
		Provider:   cfg.Embedding.Provider,
		APIKey:     cfg.Embedding.APIKey,
		Model:      cfg.Embedding.Model,
		Dimension:  cfg.Embedding.Dimension,
		BatchSize:  cfg.Embedding.BatchSize,
		Timeout:    30 * time.Second,
		MaxRetries: 3,
	})
	deps.EmbeddingClient = embeddingClient
	logger.Info("embedding client initialized")

	// Initialize multi-provider LLM client with fallback
	// Fallback order: Gemini → Moonshot → Ollama → OpenAI → Anthropic
	logger.Info("initializing LLM client with fallback",
		zap.Strings("providers", cfg.LLM.GetConfiguredProviders()),
	)
	llmClient := llm.NewFallbackClient(llm.Config{
		GeminiAPIKey:    cfg.LLM.GeminiAPIKey,
		MoonshotAPIKey:  cfg.LLM.MoonshotAPIKey,
		OllamaBaseURL:   cfg.LLM.OllamaBaseURL,
		OpenAIAPIKey:    cfg.LLM.OpenAIAPIKey,
		AnthropicAPIKey: cfg.LLM.AnthropicAPIKey,
		DefaultModel:    cfg.LLM.DefaultModel,
		MaxTokens:       cfg.LLM.MaxTokens,
		Temperature:     cfg.LLM.Temperature,
		Timeout:         cfg.LLM.Timeout,
		MaxRetries:      cfg.LLM.MaxRetries,
		EnableFallback:  cfg.LLM.EnableFallback,
		Logger:          logger,
	})
	deps.LLMClient = llmClient
	logger.Info("LLM client initialized",
		zap.Int("provider_count", len(llmClient.GetAvailableProviders())),
		zap.Strings("available_providers", toStringSlice(llmClient.GetAvailableProviders())),
	)

	// Initialize keyword search engine
	logger.Info("initializing keyword search engine")
	keywordEngine, err := keyword.NewBleveEngine(keyword.Config{
		InMemory: true, // Use in-memory for development
	})
	if err != nil {
		logger.Warn("failed to initialize keyword engine, continuing without", zap.Error(err))
	} else {
		deps.KeywordEngine = keywordEngine
		logger.Info("keyword search engine initialized")
	}

	// Initialize query router
	queryRouter := router.NewRouter(router.DefaultRouterConfig())
	deps.Router = queryRouter
	logger.Info("query router initialized")

	// Initialize hybrid retriever
	if deps.VectorStore != nil || deps.KeywordEngine != nil {
		hybridRetriever := retriever.NewHybridRetriever(
			deps.VectorStore,
			deps.KeywordEngine,
			deps.EmbeddingClient,
			retriever.DefaultConfig(),
		)
		deps.Retriever = hybridRetriever
		logger.Info("hybrid retriever initialized")
	}

	return deps, nil
}

// toStringSlice converts Provider slice to string slice for logging
func toStringSlice(providers []llm.Provider) []string {
	result := make([]string, len(providers))
	for i, p := range providers {
		result[i] = string(p)
	}
	return result
}