130 lines
3.6 KiB
Go
130 lines
3.6 KiB
Go
package ollama
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
|
|
"gitea.miguelmuniz.com/rogueking/junk2jive-server/internal/response"
|
|
)
|
|
|
|
type Ollama struct {
|
|
OllamaKey string
|
|
Message string
|
|
Image string
|
|
}
|
|
|
|
// ChatMessage represents a message in the chat completion request
|
|
type ChatMessage struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
}
|
|
|
|
// FileInfo represents a file in the chat completion request
|
|
type FileInfo struct {
|
|
Type string `json:"type"`
|
|
ID string `json:"id"`
|
|
}
|
|
|
|
// ChatCompletionRequest represents the request body for the chat completion API
|
|
type ChatCompletionRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []ChatMessage `json:"messages"`
|
|
Files []FileInfo `json:"files,omitempty"`
|
|
}
|
|
|
|
func NewOllama() *Ollama {
|
|
return &Ollama{
|
|
OllamaKey: os.Getenv("OLLAMA_API_KEY"),
|
|
}
|
|
}
|
|
|
|
func (o *Ollama) SendRequest(userMessage string, fileID string) (string, error) {
|
|
// Prepare request body
|
|
reqBody := ChatCompletionRequest{
|
|
Model: "gemma3:12b",
|
|
Messages: []ChatMessage{
|
|
{Role: "user", Content: userMessage},
|
|
},
|
|
}
|
|
|
|
// Add file if fileID is provided
|
|
if fileID != "" {
|
|
reqBody.Files = []FileInfo{
|
|
{Type: "file", ID: fileID},
|
|
}
|
|
}
|
|
|
|
// Marshal the request body to JSON
|
|
jsonBody, err := json.Marshal(reqBody)
|
|
if err != nil {
|
|
response.RespondWithError(nil, nil, http.StatusInternalServerError, "Error marshalling request", err)
|
|
return "", fmt.Errorf("error marshalling request: %w", err)
|
|
}
|
|
|
|
// Create the HTTP request
|
|
req, err := http.NewRequest("POST", "http://localhost:3000/api/chat/completions", bytes.NewBuffer(jsonBody))
|
|
if err != nil {
|
|
response.RespondWithError(nil, nil, http.StatusInternalServerError, "Error creating request", err)
|
|
return "", fmt.Errorf("error creating request: %w", err)
|
|
}
|
|
|
|
// Add headers
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Authorization", "Bearer "+o.OllamaKey)
|
|
|
|
// Send the request
|
|
client := &http.Client{}
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
response.RespondWithError(nil, nil, http.StatusInternalServerError, "Error sending request", err)
|
|
return "", fmt.Errorf("error sending request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
// Read the response
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
response.RespondWithError(nil, nil, http.StatusInternalServerError, "Error reading response", err)
|
|
return "", fmt.Errorf("error reading response: %w", err)
|
|
}
|
|
|
|
// Check for non-200 status code
|
|
if resp.StatusCode != http.StatusOK {
|
|
response.RespondWithError(nil, nil, http.StatusInternalServerError, "API request failed", err)
|
|
return "", fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
|
|
}
|
|
|
|
return string(body), nil
|
|
}
|
|
|
|
func OllamaRequest(w http.ResponseWriter, r *http.Request) {
|
|
// Parse request body
|
|
var requestData struct {
|
|
Message string `json:"message"`
|
|
FileID string `json:"fileId,omitempty"`
|
|
}
|
|
|
|
err := json.NewDecoder(r.Body).Decode(&requestData)
|
|
if err != nil {
|
|
response.RespondWithError(w, r, http.StatusBadRequest, "Invalid request body", err)
|
|
return
|
|
}
|
|
|
|
// Create ollama client using API key from environment
|
|
ollama := NewOllama()
|
|
|
|
// Send request to Ollama API
|
|
apiResponse, err := ollama.SendRequest(requestData.Message, requestData.FileID)
|
|
if err != nil {
|
|
response.RespondWithError(w, r, http.StatusInternalServerError, "Error sending request to Ollama API", err)
|
|
http.Error(w, fmt.Sprintf("Error from Ollama API: %v", err), http.StatusInternalServerError)
|
|
return
|
|
}
|
|
|
|
// Return response
|
|
response.RespondWithJSON(w, http.StatusOK, map[string]string{"response": apiResponse})
|
|
} |