vetrag/main.go

package main

import (
	"context"
	"os"
	"strconv"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/sirupsen/logrus"
)

func main() {
	logrus.SetFormatter(&logrus.TextFormatter{FullTimestamp: true})
	logrus.SetLevel(logrus.InfoLevel)

	if err := loadConfig("config.yaml"); err != nil {
		logrus.Fatalf("Failed to load config.yaml: %v", err)
	}
	logrus.Infof("Loaded config: %+v", appConfig)
	visitDB := NewVisitDB()

	if err := loadUITemplate("ui.html"); err != nil {
		logrus.Fatalf("Failed to load ui.html: %v", err)
	}

	// Initialize PostgreSQL repository first
	dsn := buildDefaultDSN()
	logrus.Info("Connecting to PostgreSQL with DSN: ", dsn)
	repo, err := NewPGChatRepository(context.Background(), dsn)
	if err != nil {
		logrus.WithError(err).Warn("PostgreSQL repository disabled (connection failed)")
	} else if repo == nil {
		logrus.Info("PostgreSQL repository not configured (no DSN)")
	}
	// defer repo.Close() // optionally enable

	// Initialize LLM client
	llmClient := NewLLMClient(
		os.Getenv("OPENAI_API_KEY"),
		os.Getenv("OPENAI_BASE_URL"),
		os.Getenv("OPENAI_MODEL"),
		repo,
	)
	var llm LLMClientAPI = llmClient

	chatService := NewChatService(llm, &visitDB, repo)
	r := gin.Default()
	// Routes
	r.GET("/", func(c *gin.Context) {
		c.Status(200)
		if err := uiTemplate.Execute(c.Writer, nil); err != nil {
			logrus.Errorf("Failed to execute ui.html template: %v", err)
		}
	})
	r.GET("/health", func(c *gin.Context) {
		c.Status(200)
		c.JSON(200, gin.H{"status": "ok"})
	})
	r.POST("/chat", chatService.HandleChat)

	if err := loadDBEditTemplate("ui_dbedit.html"); err != nil {
		logrus.Fatalf("Failed to load ui_dbedit.html: %v", err)
	}
	if err := loadAdminChatsTemplate("ui_admin_chats.html"); err != nil {
		logrus.Fatalf("Failed to load ui_admin_chats.html: %v", err)
	}
	r.GET("/admin", func(c *gin.Context) {
		c.Status(200)
		if err := uiDBEditTemplate.Execute(c.Writer, nil); err != nil {
			logrus.Errorf("Failed to execute ui_dbedit.html template: %v", err)
		}
	})
	r.GET("/db.yaml", func(c *gin.Context) { c.File("db.yaml") })

	// Add knowledgeModel save endpoint (replaces snapshot)
	r.POST("/admin/chats/snapshot", func(c *gin.Context) {
		if repo == nil {
			c.JSON(500, gin.H{"error": "repository not configured"})
			return
		}
		var req struct {
			Text string `json:"text"`
		}
		if err := c.BindJSON(&req); err != nil || req.Text == "" {
			c.JSON(400, gin.H{"error": "invalid request"})
			return
		}
		if err := repo.SaveKnowledgeModel(c.Request.Context(), req.Text); err != nil {
			c.JSON(500, gin.H{"error": "failed to save snapshot"})
			return
		}
		c.JSON(200, gin.H{"status": "ok"})
	})

	// JSON: list chat interactions
	r.GET("/admin/chats", func(c *gin.Context) {
		if repo == nil {
			c.JSON(200, gin.H{"items": []ChatInteraction{}, "pagination": gin.H{"limit": 0, "offset": 0, "count": 0}, "warning": "repository not configured"})
			return
		}
		limit := 50
		if ls := c.Query("limit"); ls != "" {
			if v, err := strconv.Atoi(ls); err == nil {
				limit = v
			}
		}
		offset := 0
		if osf := c.Query("offset"); osf != "" {
			if v, err := strconv.Atoi(osf); err == nil {
				offset = v
			}
		}
		items, err := repo.ListChatInteractions(c.Request.Context(), limit, offset)
		if err != nil {
			c.JSON(500, gin.H{"error": "failed to list interactions"})
			return
		}
		c.JSON(200, gin.H{"items": items, "pagination": gin.H{"limit": limit, "offset": offset, "count": len(items)}})
	})
	// JSON: list raw LLM events for a correlation id OR serve UI when no correlation_id provided
	r.GET("/admin/chats/events", func(c *gin.Context) {
		corr := c.Query("correlation_id")
		if corr == "" {
			if strings.Contains(c.GetHeader("Accept"), "application/json") {
				c.JSON(400, gin.H{"error": "missing correlation_id"})
				return
			}
			c.Status(200)
			if err := uiAdminChatsTemplate.Execute(c.Writer, nil); err != nil {
				logrus.Errorf("Failed to execute ui_admin_chats.html template: %v", err)
			}
			return
		}
		if repo == nil { // repository not configured, return empty JSON set for events
			c.JSON(200, gin.H{"items": []RawLLMEvent{}, "pagination": gin.H{"limit": 0, "offset": 0, "count": 0}, "warning": "repository not configured"})
			return
		}
		limit := 100
		if ls := c.Query("limit"); ls != "" {
			if v, err := strconv.Atoi(ls); err == nil {
				limit = v
			}
		}
		offset := 0
		if osf := c.Query("offset"); osf != "" {
			if v, err := strconv.Atoi(osf); err == nil {
				offset = v
			}
		}
		events, err := repo.ListLLMRawEvents(c.Request.Context(), corr, limit, offset)
		if err != nil {
			c.JSON(500, gin.H{"error": "failed to list events"})
			return
		}
		c.JSON(200, gin.H{"items": events, "pagination": gin.H{"limit": limit, "offset": offset, "count": len(events)}})
	})
	// HTML UI for chats & events
	r.GET("/admin/chats/ui", func(c *gin.Context) {
		c.Status(200)
		if err := uiAdminChatsTemplate.Execute(c.Writer, nil); err != nil {
			logrus.Errorf("Failed to execute ui_admin_chats.html template: %v", err)
		}
	})

	// List knowledgeModel entries (replaces snapshots)
	r.GET("/admin/chats/snapshots", func(c *gin.Context) {
		if repo == nil {
			c.JSON(200, gin.H{"items": []interface{}{}})
			return
		}
		models, err := repo.ListKnowledgeModels(c.Request.Context(), 100, 0)
		if err != nil {
			c.JSON(500, gin.H{"error": "failed to list snapshots"})
			return
		}
		c.JSON(200, gin.H{"items": models})
	})

	// Get a specific knowledgeModel entry (replaces snapshot)
	r.GET("/admin/chats/snapshot/:id", func(c *gin.Context) {
		if repo == nil {
			c.JSON(404, gin.H{"error": "repository not configured"})
			return
		}
		idStr := c.Param("id")
		id, err := strconv.ParseInt(idStr, 10, 64)
		if err != nil {
			c.JSON(400, gin.H{"error": "invalid id"})
			return
		}
		text, err := repo.GetKnowledgeModelText(c.Request.Context(), id)
		if err != nil {
			c.JSON(404, gin.H{"error": "snapshot not found"})
			return
		}
		c.Data(200, "text/yaml; charset=utf-8", []byte(text))
	})

	r.Run(":8080")
}