mirror of
https://github.com/Dvorinka/Trackeep.git
synced 2026-06-04 12:32:58 +00:00
first test
This commit is contained in:
Tomas Dvorak
@@ -0,0 +1,360 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/trackeep/backend/models"
|
||||
"github.com/trackeep/backend/services"
|
||||
)
|
||||
|
||||
// MistralConfig holds configuration for Mistral AI
|
||||
type MistralConfig struct {
|
||||
APIKey string
|
||||
BaseURL string
|
||||
Model string
|
||||
MaxTokens int
|
||||
Temperature float64
|
||||
}
|
||||
|
||||
// ChatRequest represents a chat message request
|
||||
type ChatRequest struct {
|
||||
Message string `json:"message" binding:"required"`
|
||||
SessionID *string `json:"session_id,omitempty"`
|
||||
Context map[string]bool `json:"context,omitempty"` // what data to include
|
||||
Provider string `json:"provider,omitempty"` // "mistral", "longcat", "grok", "deepseek", "ollama", "openrouter"
|
||||
ModelType string `json:"model_type,omitempty"` // "standard", "thinking", "upgraded_thinking"
|
||||
}
|
||||
|
||||
// ChatResponse represents a chat response
|
||||
type ChatResponse struct {
|
||||
ID string `json:"id"`
|
||||
Message string `json:"message"`
|
||||
Role string `json:"role"`
|
||||
SessionID string `json:"session_id"`
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
Model string `json:"model"`
|
||||
TokenUsage TokenUsage `json:"token_usage"`
|
||||
ContextUsed []string `json:"context_used"`
|
||||
}
|
||||
|
||||
// TokenUsage represents token usage information
|
||||
type TokenUsage struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
}
|
||||
|
||||
// MistralMessage represents a message for Mistral API
|
||||
type MistralMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// MistralRequest represents a request to Mistral API
|
||||
type MistralRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []MistralMessage `json:"messages"`
|
||||
MaxTokens int `json:"max_tokens,omitempty"`
|
||||
Temperature float64 `json:"temperature,omitempty"`
|
||||
}
|
||||
|
||||
// MistralResponse represents a response from Mistral API
|
||||
type MistralResponse struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []struct {
|
||||
Index int `json:"index"`
|
||||
Message struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
} `json:"choices"`
|
||||
Usage struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
|
||||
var mistralConfig = MistralConfig{
|
||||
APIKey: os.Getenv("MISTRAL_API_KEY"),
|
||||
BaseURL: "https://api.mistral.ai/v1",
|
||||
Model: "mistral-small-latest", // Cheap and capable model
|
||||
MaxTokens: 4000,
|
||||
Temperature: 0.7,
|
||||
}
|
||||
|
||||
// GetMistralConfig returns current Mistral configuration
|
||||
func GetMistralConfig() MistralConfig {
|
||||
return mistralConfig
|
||||
}
|
||||
|
||||
// SendMessage handles chat message requests
|
||||
func SendMessage(c *gin.Context) {
|
||||
userID := c.GetUint("user_id")
|
||||
|
||||
var req ChatRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// Get or create session
|
||||
var session models.ChatSession
|
||||
if req.SessionID != nil {
|
||||
if err := models.DB.Where("id = ? AND user_id = ?", *req.SessionID, userID).First(&session).Error; err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "Session not found"})
|
||||
return
|
||||
}
|
||||
} else {
|
||||
// Create new session
|
||||
session = models.ChatSession{
|
||||
UserID: userID,
|
||||
Title: fmt.Sprintf("Chat %s", time.Now().Format("Jan 2, 3:04 PM")),
|
||||
IncludeBookmarks: true,
|
||||
IncludeTasks: true,
|
||||
IncludeFiles: true,
|
||||
IncludeNotes: true,
|
||||
}
|
||||
if req.Context != nil {
|
||||
session.IncludeBookmarks = req.Context["bookmarks"]
|
||||
session.IncludeTasks = req.Context["tasks"]
|
||||
session.IncludeFiles = req.Context["files"]
|
||||
session.IncludeNotes = req.Context["notes"]
|
||||
}
|
||||
if err := models.DB.Create(&session).Error; err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to create session"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Save user message
|
||||
userMessage := models.ChatMessage{
|
||||
UserID: userID,
|
||||
SessionID: strconv.Itoa(int(session.ID)),
|
||||
Content: req.Message,
|
||||
Role: "user",
|
||||
}
|
||||
if err := models.DB.Create(&userMessage).Error; err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to save message"})
|
||||
return
|
||||
}
|
||||
|
||||
// Get conversation history
|
||||
var messages []models.ChatMessage
|
||||
models.DB.Where("session_id = ?", session.ID).Order("created_at asc").Find(&messages)
|
||||
|
||||
// Build context from user data
|
||||
contextData, err := buildUserContext(userID, session)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to build context"})
|
||||
return
|
||||
}
|
||||
|
||||
// Build messages for AI provider (system + history)
|
||||
aiMessages := []services.Message{
|
||||
{
|
||||
Role: "system",
|
||||
Content: buildSystemPrompt(contextData),
|
||||
},
|
||||
}
|
||||
|
||||
// Add conversation history (limit to last 10 messages to manage token count)
|
||||
startIdx := 0
|
||||
if len(messages) > 11 { // system + 10 messages
|
||||
startIdx = len(messages) - 10
|
||||
}
|
||||
for i := startIdx; i < len(messages); i++ {
|
||||
aiMessages = append(aiMessages, services.Message{
|
||||
Role: messages[i].Role,
|
||||
Content: messages[i].Content,
|
||||
})
|
||||
}
|
||||
|
||||
// Determine AI provider
|
||||
aiProvider := services.ProviderMistral
|
||||
switch req.Provider {
|
||||
case "longcat":
|
||||
aiProvider = services.ProviderLongCat
|
||||
case "grok":
|
||||
aiProvider = services.ProviderGrok
|
||||
case "deepseek":
|
||||
aiProvider = services.ProviderDeepSeek
|
||||
case "ollama":
|
||||
aiProvider = services.ProviderOllama
|
||||
case "openrouter":
|
||||
aiProvider = services.ProviderOpenRouter
|
||||
}
|
||||
|
||||
aiService := services.NewAIService(aiProvider)
|
||||
aiReq := services.AIRequest{
|
||||
Messages: aiMessages,
|
||||
MaxTokens: 2000,
|
||||
Temperature: 0.7,
|
||||
ModelType: req.ModelType,
|
||||
}
|
||||
|
||||
// Call AI provider
|
||||
startTime := time.Now()
|
||||
var aiResp *services.AIResponse
|
||||
switch req.ModelType {
|
||||
case "thinking":
|
||||
aiResp, err = aiService.ChatCompletionWithThinking(aiReq)
|
||||
case "upgraded_thinking":
|
||||
aiResp, err = aiService.ChatCompletionWithUpgradedThinking(aiReq)
|
||||
default:
|
||||
aiResp, err = aiService.ChatCompletion(aiReq)
|
||||
}
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to call AI: " + err.Error()})
|
||||
return
|
||||
}
|
||||
processingMs := time.Since(startTime).Milliseconds()
|
||||
|
||||
if len(aiResp.Choices) == 0 {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "No response from AI"})
|
||||
return
|
||||
}
|
||||
|
||||
// Extract assistant response, handling thinking models where needed
|
||||
assistantContent := services.ParseThinkingResponse(aiResp, aiProvider, req.ModelType)
|
||||
|
||||
// Save assistant message
|
||||
assistantMessage := models.ChatMessage{
|
||||
UserID: userID,
|
||||
SessionID: strconv.Itoa(int(session.ID)),
|
||||
Content: assistantContent,
|
||||
Role: "assistant",
|
||||
ModelUsed: aiResp.Model,
|
||||
TokenCount: aiResp.Usage.TotalTokens,
|
||||
ProcessingMs: processingMs,
|
||||
ContextItems: getContextItemIDs(contextData),
|
||||
}
|
||||
if err := models.DB.Create(&assistantMessage).Error; err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to save response"})
|
||||
return
|
||||
}
|
||||
|
||||
// Update session
|
||||
session.MessageCount = len(messages) + 1
|
||||
now := time.Now()
|
||||
session.LastMessageAt = &now
|
||||
models.DB.Save(&session)
|
||||
|
||||
// Return response
|
||||
response := ChatResponse{
|
||||
ID: aiResp.ID,
|
||||
Message: assistantContent,
|
||||
Role: "assistant",
|
||||
SessionID: strconv.Itoa(int(session.ID)),
|
||||
Timestamp: time.Now(),
|
||||
Model: aiResp.Model,
|
||||
TokenUsage: TokenUsage{
|
||||
PromptTokens: aiResp.Usage.PromptTokens,
|
||||
CompletionTokens: aiResp.Usage.CompletionTokens,
|
||||
TotalTokens: aiResp.Usage.TotalTokens,
|
||||
},
|
||||
ContextUsed: getContextItemIDs(contextData),
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, response)
|
||||
}
|
||||
|
||||
// GetSessions retrieves user's chat sessions
|
||||
func GetSessions(c *gin.Context) {
|
||||
userID := c.GetUint("user_id")
|
||||
|
||||
var sessions []models.ChatSession
|
||||
models.DB.Where("user_id = ?", userID).Order("updated_at desc").Find(&sessions)
|
||||
|
||||
c.JSON(http.StatusOK, sessions)
|
||||
}
|
||||
|
||||
// GetSessionMessages retrieves messages for a specific session
|
||||
func GetSessionMessages(c *gin.Context) {
|
||||
userID := c.GetUint("user_id")
|
||||
sessionID := c.Param("id")
|
||||
|
||||
// Verify session belongs to user
|
||||
var session models.ChatSession
|
||||
if err := models.DB.Where("id = ? AND user_id = ?", sessionID, userID).First(&session).Error; err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "Session not found"})
|
||||
return
|
||||
}
|
||||
|
||||
var messages []models.ChatMessage
|
||||
models.DB.Where("session_id = ?", sessionID).Order("created_at asc").Find(&messages)
|
||||
|
||||
c.JSON(http.StatusOK, messages)
|
||||
}
|
||||
|
||||
// DeleteSession deletes a chat session and its messages
|
||||
func DeleteSession(c *gin.Context) {
|
||||
userID := c.GetUint("user_id")
|
||||
sessionID := c.Param("id")
|
||||
|
||||
// Verify session belongs to user
|
||||
var session models.ChatSession
|
||||
if err := models.DB.Where("id = ? AND user_id = ?", sessionID, userID).First(&session).Error; err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "Session not found"})
|
||||
return
|
||||
}
|
||||
|
||||
// Delete messages first
|
||||
models.DB.Where("session_id = ?", sessionID).Delete(&models.ChatMessage{})
|
||||
|
||||
// Delete session
|
||||
models.DB.Delete(&session)
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "Session deleted"})
|
||||
}
|
||||
|
||||
func callMistral(messages []MistralMessage) (*MistralResponse, error) {
|
||||
reqBody := MistralRequest{
|
||||
Model: mistralConfig.Model,
|
||||
Messages: messages,
|
||||
MaxTokens: mistralConfig.MaxTokens,
|
||||
Temperature: mistralConfig.Temperature,
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(reqBody)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req, err := http.NewRequest("POST", mistralConfig.BaseURL+"/chat/completions", strings.NewReader(string(jsonData)))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+mistralConfig.APIKey)
|
||||
|
||||
client := &http.Client{Timeout: 30 * time.Second}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("Mistral API returned status %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var mistralResp MistralResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&mistralResp); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &mistralResp, nil
|
||||
}
|
||||
Reference in New Issue
Block a user