289 lines
8.9 KiB
Go
289 lines
8.9 KiB
Go
package main
|
|
|
|
import (
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/sirupsen/logrus"
|
|
)
|
|
|
|
// Controller struct to hold dependencies
|
|
// Add more dependencies as needed
|
|
// e.g. templates, services, etc.
|
|
type Controller struct {
|
|
repo ChatRepositoryAPI
|
|
llm LLMClientAPI
|
|
visitDB *VisitDB
|
|
template TemplateExecutor
|
|
uiDBEditTemplate TemplateExecutor
|
|
uiAdminChatsTemplate TemplateExecutor
|
|
uiAdminLoginTemplate TemplateExecutor
|
|
}
|
|
|
|
// NewController creates a new Controller instance
|
|
func NewController(repo ChatRepositoryAPI, llm LLMClientAPI, visitDB *VisitDB, template, uiDBEditTemplate, uiAdminChatsTemplate, uiAdminLoginTemplate TemplateExecutor) *Controller {
|
|
return &Controller{
|
|
repo: repo,
|
|
llm: llm,
|
|
visitDB: visitDB,
|
|
template: template,
|
|
uiDBEditTemplate: uiDBEditTemplate,
|
|
uiAdminChatsTemplate: uiAdminChatsTemplate,
|
|
uiAdminLoginTemplate: uiAdminLoginTemplate,
|
|
}
|
|
}
|
|
|
|
// Handler for root UI
|
|
func (ctrl *Controller) RootUI(c *gin.Context) {
|
|
c.Status(http.StatusOK)
|
|
if err := ctrl.template.Execute(c.Writer, nil); err != nil {
|
|
logrus.Errorf("Failed to execute ui.html template: %v", err)
|
|
}
|
|
}
|
|
|
|
// Handler for health check
|
|
func (ctrl *Controller) Health(c *gin.Context) {
|
|
c.Status(http.StatusOK)
|
|
c.JSON(http.StatusOK, gin.H{"status": "ok"})
|
|
}
|
|
|
|
// Handler for chat
|
|
func (ctrl *Controller) HandleChat(c *gin.Context) {
|
|
// Use your existing chatService logic here
|
|
// You may want to move NewChatService to a service file for full MVC
|
|
chatService := NewChatService(ctrl.llm, ctrl.visitDB, ctrl.repo)
|
|
chatService.HandleChat(c)
|
|
}
|
|
|
|
// Handler for admin UI
|
|
func (ctrl *Controller) AdminUI(c *gin.Context) {
|
|
c.Status(http.StatusOK)
|
|
if err := ctrl.uiDBEditTemplate.Execute(c.Writer, nil); err != nil {
|
|
logrus.Errorf("Failed to execute ui_dbedit.html template: %v", err)
|
|
}
|
|
}
|
|
|
|
// Handler for download db.yaml
|
|
func (ctrl *Controller) DownloadDB(c *gin.Context) {
|
|
c.File("config/db.yaml")
|
|
}
|
|
|
|
// Handler for saving knowledgeModel (snapshot)
|
|
func (ctrl *Controller) SaveKnowledgeModel(c *gin.Context) {
|
|
if ctrl.repo == nil {
|
|
c.JSON(http.StatusInternalServerError, 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(http.StatusBadRequest, gin.H{"error": "invalid request"})
|
|
return
|
|
}
|
|
if err := ctrl.repo.SaveKnowledgeModel(c.Request.Context(), req.Text); err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to save snapshot"})
|
|
return
|
|
}
|
|
c.JSON(http.StatusOK, gin.H{"status": "ok"})
|
|
}
|
|
|
|
// Handler for listing chat interactions
|
|
func (ctrl *Controller) ListChatInteractions(c *gin.Context) {
|
|
if ctrl.repo == nil {
|
|
c.JSON(http.StatusOK, 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 := ctrl.repo.ListChatInteractions(c.Request.Context(), limit, offset)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to list interactions"})
|
|
return
|
|
}
|
|
c.JSON(http.StatusOK, gin.H{"items": items, "pagination": gin.H{"limit": limit, "offset": offset, "count": len(items)}})
|
|
}
|
|
|
|
// Handler for listing LLM events
|
|
func (ctrl *Controller) ListLLMEvents(c *gin.Context) {
|
|
corr := c.Query("correlation_id")
|
|
if corr == "" {
|
|
if strings.Contains(c.GetHeader("Accept"), "application/json") {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "missing correlation_id"})
|
|
return
|
|
}
|
|
c.Status(http.StatusOK)
|
|
if err := ctrl.uiAdminChatsTemplate.Execute(c.Writer, nil); err != nil {
|
|
logrus.Errorf("Failed to execute ui_admin_chats.html template: %v", err)
|
|
}
|
|
return
|
|
}
|
|
if ctrl.repo == nil {
|
|
c.JSON(http.StatusOK, 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 := ctrl.repo.ListLLMRawEvents(c.Request.Context(), corr, limit, offset)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to list events"})
|
|
return
|
|
}
|
|
c.JSON(http.StatusOK, gin.H{"items": events, "pagination": gin.H{"limit": limit, "offset": offset, "count": len(events)}})
|
|
}
|
|
|
|
// Handler for admin chats UI
|
|
func (ctrl *Controller) AdminChatsUI(c *gin.Context) {
|
|
c.Status(http.StatusOK)
|
|
if err := ctrl.uiAdminChatsTemplate.Execute(c.Writer, nil); err != nil {
|
|
logrus.Errorf("Failed to execute ui_admin_chats.html template: %v", err)
|
|
}
|
|
}
|
|
|
|
// Handler for listing knowledgeModel entries
|
|
func (ctrl *Controller) ListKnowledgeModels(c *gin.Context) {
|
|
if ctrl.repo == nil {
|
|
c.JSON(http.StatusOK, gin.H{"items": []interface{}{}})
|
|
return
|
|
}
|
|
models, err := ctrl.repo.ListKnowledgeModels(c.Request.Context(), 100, 0)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to list snapshots"})
|
|
return
|
|
}
|
|
c.JSON(http.StatusOK, gin.H{"items": models})
|
|
}
|
|
|
|
// Handler for getting a specific knowledgeModel entry
|
|
func (ctrl *Controller) GetKnowledgeModel(c *gin.Context) {
|
|
if ctrl.repo == nil {
|
|
c.JSON(http.StatusNotFound, gin.H{"error": "repository not configured"})
|
|
return
|
|
}
|
|
idStr := c.Param("id")
|
|
id, err := strconv.ParseInt(idStr, 10, 64)
|
|
if err != nil {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid id"})
|
|
return
|
|
}
|
|
text, err := ctrl.repo.GetKnowledgeModelText(c.Request.Context(), id)
|
|
if err != nil {
|
|
c.JSON(http.StatusNotFound, gin.H{"error": "snapshot not found"})
|
|
return
|
|
}
|
|
c.Data(http.StatusOK, "text/yaml; charset=utf-8", []byte(text))
|
|
}
|
|
|
|
// AdminLoginPage serves the admin login page
|
|
func (ctrl *Controller) AdminLoginPage(c *gin.Context) {
|
|
c.Status(http.StatusOK)
|
|
if err := ctrl.uiAdminLoginTemplate.Execute(c.Writer, gin.H{"Error": c.Query("error")}); err != nil {
|
|
logrus.Errorf("Failed to execute ui_admin_login.html template: %v", err)
|
|
}
|
|
}
|
|
|
|
// AdminLogin handles admin login POST
|
|
func (ctrl *Controller) AdminLogin(c *gin.Context) {
|
|
username := c.PostForm("username")
|
|
password := c.PostForm("password")
|
|
user, err := ctrl.repo.GetUserByUsername(c.Request.Context(), username)
|
|
if err != nil || !CheckPasswordHash(password, user.PasswordHash) {
|
|
c.Redirect(http.StatusFound, "/admin/login?error=Invalid+credentials")
|
|
return
|
|
}
|
|
token, err := GenerateJWT(user.Username)
|
|
if err != nil {
|
|
c.Redirect(http.StatusFound, "/admin/login?error=Internal+error")
|
|
return
|
|
}
|
|
c.SetCookie("admin_jwt", token, 86400, "/", "", false, true)
|
|
c.Redirect(http.StatusFound, "/admin")
|
|
}
|
|
|
|
// AdminLogout logs out the admin
|
|
func (ctrl *Controller) AdminLogout(c *gin.Context) {
|
|
c.SetCookie("admin_jwt", "", -1, "/", "", false, true)
|
|
c.Redirect(http.StatusFound, "/admin/login")
|
|
}
|
|
|
|
// AdminAuthMiddleware checks for a valid admin JWT cookie
|
|
func AdminAuthMiddleware() gin.HandlerFunc {
|
|
return func(c *gin.Context) {
|
|
if c.Request.URL.Path == "/admin/login" || c.Request.URL.Path == "/admin/logout" {
|
|
c.Next()
|
|
return
|
|
}
|
|
jwtToken, err := c.Cookie("admin_jwt")
|
|
if err != nil {
|
|
c.Redirect(http.StatusFound, "/admin/login")
|
|
c.Abort()
|
|
return
|
|
}
|
|
username, err := ValidateJWT(jwtToken)
|
|
if err != nil || username == "" {
|
|
c.Redirect(http.StatusFound, "/admin/login")
|
|
c.Abort()
|
|
return
|
|
}
|
|
c.Set("admin_username", username)
|
|
c.Next()
|
|
}
|
|
}
|
|
|
|
// CreateInitialAdmin allows creation of the first admin user if none exist
|
|
func (ctrl *Controller) CreateInitialAdmin(c *gin.Context) {
|
|
// Only allow if no users exist
|
|
count, err := ctrl.repo.CountUsers(c.Request.Context())
|
|
if err != nil || count > 0 {
|
|
c.JSON(http.StatusForbidden, gin.H{"error": "Admin user already exists or error"})
|
|
return
|
|
}
|
|
username := c.PostForm("username")
|
|
password := c.PostForm("password")
|
|
if username == "" || password == "" {
|
|
c.JSON(http.StatusBadRequest, gin.H{"error": "Username and password required"})
|
|
return
|
|
}
|
|
hash, err := HashPassword(password)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to hash password"})
|
|
return
|
|
}
|
|
err = ctrl.repo.CreateUser(c.Request.Context(), username, hash)
|
|
if err != nil {
|
|
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to create user"})
|
|
return
|
|
}
|
|
c.JSON(http.StatusOK, gin.H{"status": "Admin user created"})
|
|
}
|
|
|
|
// TemplateExecutor is an interface for template execution
|
|
// This allows for easier testing and decoupling
|
|
// You can implement this interface for your templates
|
|
|
|
type TemplateExecutor interface {
|
|
Execute(wr http.ResponseWriter, data interface{}) error
|
|
}
|