Реализована операции Milvus для управления документами и встраиванием, включая функции вставки, запроса и удаления. Внедрите архитектуру RAG с LLM и сервисами встраивания. Добавьте обработку текста для фрагментации и конкатенации. Создайте автономный скрипт для настройки и управления Milvus. Разработайте комплексные тесты API для обработки документов и взаимодействия с LLM, включая имитации для сервисов. Расширьте возможности конфигурации пользователя с помощью дополнительных настроек YAML.
This commit is contained in:
Dmitriy Fofanov
313
tests/api_test.go
Normal file
313
tests/api_test.go
Normal file
@@ -0,0 +1,313 @@
|
||||
package api_test
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"easy_rag/internal/models"
|
||||
"encoding/json"
|
||||
"github.com/google/uuid"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"easy_rag/api"
|
||||
"easy_rag/internal/pkg/rag"
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
)
|
||||
|
||||
// Example: Test UploadHandler
|
||||
func TestUploadHandler(t *testing.T) {
|
||||
e := echo.New()
|
||||
|
||||
// Create a mock for the LLM, Embeddings, and Database
|
||||
mockLLM := new(MockLLMService)
|
||||
mockEmbeddings := new(MockEmbeddingsService)
|
||||
mockDB := new(MockDatabase)
|
||||
|
||||
// Setup the Rag object
|
||||
r := &rag.Rag{
|
||||
LLM: mockLLM,
|
||||
Embeddings: mockEmbeddings,
|
||||
Database: mockDB,
|
||||
}
|
||||
|
||||
// We expect calls to these mocks in the background goroutine, for each document.
|
||||
|
||||
// The request body
|
||||
requestBody := api.RequestUpload{
|
||||
Docs: []api.UploadDoc{
|
||||
{
|
||||
Content: "Test document content",
|
||||
Link: "http://example.com/doc",
|
||||
Filename: "doc1.txt",
|
||||
Category: "TestCategory",
|
||||
Metadata: map[string]string{"Author": "Me"},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// Convert requestBody to JSON
|
||||
reqBodyBytes, _ := json.Marshal(requestBody)
|
||||
|
||||
// Create a new request
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/upload", bytes.NewReader(reqBodyBytes))
|
||||
req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
|
||||
|
||||
// Create a ResponseRecorder
|
||||
rec := httptest.NewRecorder()
|
||||
|
||||
// New echo context
|
||||
c := e.NewContext(req, rec)
|
||||
// Set the rag object in context
|
||||
c.Set("Rag", r)
|
||||
|
||||
// Because the UploadHandler spawns a goroutine, we only test the immediate HTTP response.
|
||||
// We can still set expectations for the calls that happen in the goroutine to ensure they're invoked.
|
||||
// For example, we expect the summary to be generated, so:
|
||||
|
||||
testSummary := "Test summary from LLM"
|
||||
mockLLM.On("Generate", mock.Anything).Return(testSummary, nil).Maybe() // .Maybe() because the concurrency might not complete by the time we assert
|
||||
|
||||
// The embedding vector returned from the embeddings service
|
||||
testVector := [][]float32{{0.1, 0.2, 0.3, 0.4}}
|
||||
|
||||
// We'll mock calls to Vectorize() for summary and each chunk
|
||||
mockEmbeddings.On("Vectorize", mock.AnythingOfType("string")).Return(testVector, nil).Maybe()
|
||||
|
||||
// The database SaveDocument / SaveEmbeddings calls
|
||||
mockDB.On("SaveDocument", mock.AnythingOfType("models.Document")).Return(nil).Maybe()
|
||||
mockDB.On("SaveEmbeddings", mock.AnythingOfType("[]models.Embedding")).Return(nil).Maybe()
|
||||
|
||||
// Invoke the handler
|
||||
err := api.UploadHandler(c)
|
||||
|
||||
// Check no immediate errors
|
||||
assert.NoError(t, err)
|
||||
|
||||
// Check the response
|
||||
assert.Equal(t, http.StatusAccepted, rec.Code)
|
||||
|
||||
var resp map[string]interface{}
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
|
||||
// We expect certain fields in the JSON response
|
||||
assert.Equal(t, "v1", resp["version"])
|
||||
assert.NotEmpty(t, resp["task_id"])
|
||||
assert.Equal(t, "Processing started", resp["status"])
|
||||
|
||||
// Typically, you might want to wait or do more advanced concurrency checks if you want to test
|
||||
// the logic in the goroutine, but that goes beyond a simple unit test.
|
||||
// The background process can be tested more thoroughly in integration or end-to-end tests.
|
||||
|
||||
// Optionally assert that our mocks were called
|
||||
mockLLM.AssertExpectations(t)
|
||||
mockEmbeddings.AssertExpectations(t)
|
||||
mockDB.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// Example: Test ListAllDocsHandler
|
||||
func TestListAllDocsHandler(t *testing.T) {
|
||||
e := echo.New()
|
||||
|
||||
mockLLM := new(MockLLMService)
|
||||
mockEmbeddings := new(MockEmbeddingsService)
|
||||
mockDB := new(MockDatabase)
|
||||
|
||||
r := &rag.Rag{
|
||||
LLM: mockLLM,
|
||||
Embeddings: mockEmbeddings,
|
||||
Database: mockDB,
|
||||
}
|
||||
|
||||
// Mock data
|
||||
doc1 := models.Document{
|
||||
ID: uuid.NewString(),
|
||||
Filename: "doc1.txt",
|
||||
Summary: "summary doc1",
|
||||
}
|
||||
doc2 := models.Document{
|
||||
ID: uuid.NewString(),
|
||||
Filename: "doc2.txt",
|
||||
Summary: "summary doc2",
|
||||
}
|
||||
docs := []models.Document{doc1, doc2}
|
||||
|
||||
// Expect the database to return the docs
|
||||
mockDB.On("ListDocuments").Return(docs, nil)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/docs", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
c.Set("Rag", r)
|
||||
|
||||
err := api.ListAllDocsHandler(c)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var resp map[string]interface{}
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
assert.Equal(t, "v1", resp["version"])
|
||||
|
||||
// The "docs" field should match the ones we returned
|
||||
docsInterface, ok := resp["docs"].([]interface{})
|
||||
assert.True(t, ok)
|
||||
assert.Len(t, docsInterface, 2)
|
||||
|
||||
// Verify mocks
|
||||
mockDB.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// Example: Test GetDocHandler
|
||||
func TestGetDocHandler(t *testing.T) {
|
||||
e := echo.New()
|
||||
|
||||
mockLLM := new(MockLLMService)
|
||||
mockEmbeddings := new(MockEmbeddingsService)
|
||||
mockDB := new(MockDatabase)
|
||||
|
||||
r := &rag.Rag{
|
||||
LLM: mockLLM,
|
||||
Embeddings: mockEmbeddings,
|
||||
Database: mockDB,
|
||||
}
|
||||
|
||||
// Mock a single doc
|
||||
docID := "123"
|
||||
testDoc := models.Document{
|
||||
ID: docID,
|
||||
Filename: "doc3.txt",
|
||||
Summary: "summary doc3",
|
||||
}
|
||||
|
||||
mockDB.On("GetDocument", docID).Return(testDoc, nil)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/v1/doc/123", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
// set path param
|
||||
c.SetParamNames("id")
|
||||
c.SetParamValues(docID)
|
||||
c.Set("Rag", r)
|
||||
|
||||
err := api.GetDocHandler(c)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var resp map[string]interface{}
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
assert.Equal(t, "v1", resp["version"])
|
||||
|
||||
docInterface := resp["doc"].(map[string]interface{})
|
||||
assert.Equal(t, "doc3.txt", docInterface["filename"])
|
||||
|
||||
// Verify mocks
|
||||
mockDB.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// Example: Test AskDocHandler
|
||||
func TestAskDocHandler(t *testing.T) {
|
||||
e := echo.New()
|
||||
|
||||
mockLLM := new(MockLLMService)
|
||||
mockEmbeddings := new(MockEmbeddingsService)
|
||||
mockDB := new(MockDatabase)
|
||||
|
||||
r := &rag.Rag{
|
||||
LLM: mockLLM,
|
||||
Embeddings: mockEmbeddings,
|
||||
Database: mockDB,
|
||||
}
|
||||
|
||||
// 1) We expect to Vectorize the question
|
||||
question := "What is the summary of doc?"
|
||||
questionVector := [][]float32{{0.5, 0.2, 0.1}}
|
||||
mockEmbeddings.On("Vectorize", question).Return(questionVector, nil)
|
||||
|
||||
// 2) We expect a DB search
|
||||
emb := []models.Embedding{
|
||||
{
|
||||
ID: "emb1",
|
||||
DocumentID: "doc123",
|
||||
TextChunk: "Relevant content chunk",
|
||||
Score: 0.99,
|
||||
},
|
||||
}
|
||||
mockDB.On("Search", questionVector).Return(emb, nil)
|
||||
|
||||
// 3) We expect the LLM to generate an answer from the chunk
|
||||
generatedAnswer := "Here is an answer from the chunk"
|
||||
// The prompt we pass is something like: "Given the following information: chunk ... Answer the question: question"
|
||||
mockLLM.On("Generate", mock.AnythingOfType("string")).Return(generatedAnswer, nil)
|
||||
|
||||
// Prepare request
|
||||
reqBody := api.RequestQuestion{
|
||||
Question: question,
|
||||
}
|
||||
reqBytes, _ := json.Marshal(reqBody)
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/ask", bytes.NewReader(reqBytes))
|
||||
req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
c.Set("Rag", r)
|
||||
|
||||
// Execute
|
||||
err := api.AskDocHandler(c)
|
||||
|
||||
// Verify
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var resp map[string]interface{}
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
assert.Equal(t, "v1", resp["version"])
|
||||
assert.Equal(t, generatedAnswer, resp["answer"])
|
||||
|
||||
// The docs field should have the docID "doc123"
|
||||
docsInterface := resp["docs"].([]interface{})
|
||||
assert.Len(t, docsInterface, 1)
|
||||
assert.Equal(t, "doc123", docsInterface[0])
|
||||
|
||||
// Verify mocks
|
||||
mockLLM.AssertExpectations(t)
|
||||
mockEmbeddings.AssertExpectations(t)
|
||||
mockDB.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// Example: Test DeleteDocHandler
|
||||
func TestDeleteDocHandler(t *testing.T) {
|
||||
e := echo.New()
|
||||
mockLLM := new(MockLLMService)
|
||||
mockEmbeddings := new(MockEmbeddingsService)
|
||||
mockDB := new(MockDatabase)
|
||||
|
||||
r := &rag.Rag{
|
||||
LLM: mockLLM,
|
||||
Embeddings: mockEmbeddings,
|
||||
Database: mockDB,
|
||||
}
|
||||
|
||||
docID := "abc"
|
||||
mockDB.On("DeleteDocument", docID).Return(nil)
|
||||
|
||||
req := httptest.NewRequest(http.MethodDelete, "/api/v1/doc/abc", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
c.SetParamNames("id")
|
||||
c.SetParamValues(docID)
|
||||
c.Set("Rag", r)
|
||||
|
||||
err := api.DeleteDocHandler(c)
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var resp map[string]interface{}
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &resp)
|
||||
assert.Equal(t, "v1", resp["version"])
|
||||
|
||||
// docs should be nil according to DeleteDocHandler
|
||||
assert.Nil(t, resp["docs"])
|
||||
|
||||
// Verify mocks
|
||||
mockDB.AssertExpectations(t)
|
||||
}
|
||||
89
tests/mock_test.go
Normal file
89
tests/mock_test.go
Normal file
@@ -0,0 +1,89 @@
|
||||
package api_test
|
||||
|
||||
import (
|
||||
"easy_rag/internal/models"
|
||||
"github.com/stretchr/testify/mock"
|
||||
)
|
||||
|
||||
// --------------------
|
||||
// Mock LLM
|
||||
// --------------------
|
||||
type MockLLMService struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *MockLLMService) Generate(prompt string) (string, error) {
|
||||
args := m.Called(prompt)
|
||||
return args.String(0), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockLLMService) GetModel() string {
|
||||
args := m.Called()
|
||||
return args.String(0)
|
||||
}
|
||||
|
||||
// --------------------
|
||||
// Mock Embeddings
|
||||
// --------------------
|
||||
type MockEmbeddingsService struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *MockEmbeddingsService) Vectorize(text string) ([][]float32, error) {
|
||||
args := m.Called(text)
|
||||
return args.Get(0).([][]float32), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockEmbeddingsService) GetModel() string {
|
||||
args := m.Called()
|
||||
return args.String(0)
|
||||
}
|
||||
|
||||
// --------------------
|
||||
// Mock Database
|
||||
// --------------------
|
||||
type MockDatabase struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
// GetDocumentInfo(id string) (models.DocumentInfo, error)
|
||||
func (m *MockDatabase) GetDocumentInfo(id string) (models.Document, error) {
|
||||
args := m.Called(id)
|
||||
return args.Get(0).(models.Document), args.Error(1)
|
||||
}
|
||||
|
||||
// SaveDocument(document Document) error
|
||||
func (m *MockDatabase) SaveDocument(doc models.Document) error {
|
||||
args := m.Called(doc)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
// SaveEmbeddings([]Embedding) error
|
||||
func (m *MockDatabase) SaveEmbeddings(emb []models.Embedding) error {
|
||||
args := m.Called(emb)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
// ListDocuments() ([]Document, error)
|
||||
func (m *MockDatabase) ListDocuments() ([]models.Document, error) {
|
||||
args := m.Called()
|
||||
return args.Get(0).([]models.Document), args.Error(1)
|
||||
}
|
||||
|
||||
// GetDocument(id string) (Document, error)
|
||||
func (m *MockDatabase) GetDocument(id string) (models.Document, error) {
|
||||
args := m.Called(id)
|
||||
return args.Get(0).(models.Document), args.Error(1)
|
||||
}
|
||||
|
||||
// DeleteDocument(id string) error
|
||||
func (m *MockDatabase) DeleteDocument(id string) error {
|
||||
args := m.Called(id)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
// Search(vector []float32) ([]models.Embedding, error)
|
||||
func (m *MockDatabase) Search(vector [][]float32) ([]models.Embedding, error) {
|
||||
args := m.Called(vector)
|
||||
return args.Get(0).([]models.Embedding), args.Error(1)
|
||||
}
|
||||
125
tests/openrouter_test.go
Normal file
125
tests/openrouter_test.go
Normal file
@@ -0,0 +1,125 @@
|
||||
package api_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
openroute2 "easy_rag/internal/llm/openroute"
|
||||
"fmt"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestFetchChatCompletions(t *testing.T) {
|
||||
client := openroute2.NewOpenRouterClient("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")
|
||||
|
||||
request := openroute2.Request{
|
||||
Model: "qwen/qwen2.5-vl-72b-instruct:free",
|
||||
Messages: []openroute2.MessageRequest{
|
||||
{openroute2.RoleUser, "Привет!", "", ""},
|
||||
},
|
||||
}
|
||||
|
||||
output, err := client.FetchChatCompletions(request)
|
||||
if err != nil {
|
||||
t.Errorf("error %v", err)
|
||||
}
|
||||
|
||||
t.Logf("output: %v", output.Choices[0].Message.Content)
|
||||
}
|
||||
|
||||
func TestFetchChatCompletionsStreaming(t *testing.T) {
|
||||
client := openroute2.NewOpenRouterClient("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")
|
||||
|
||||
request := openroute2.Request{
|
||||
Model: "qwen/qwen2.5-vl-72b-instruct:free",
|
||||
Messages: []openroute2.MessageRequest{
|
||||
{openroute2.RoleUser, "Привет!", "", ""},
|
||||
},
|
||||
Stream: true,
|
||||
}
|
||||
|
||||
outputChan := make(chan openroute2.Response)
|
||||
processingChan := make(chan interface{})
|
||||
errChan := make(chan error)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
go client.FetchChatCompletionsStream(request, outputChan, processingChan, errChan, ctx)
|
||||
|
||||
for {
|
||||
select {
|
||||
case output := <-outputChan:
|
||||
if len(output.Choices) > 0 {
|
||||
t.Logf("%s", output.Choices[0].Delta.Content)
|
||||
}
|
||||
case <-processingChan:
|
||||
t.Logf("Обработка\n")
|
||||
case err := <-errChan:
|
||||
if err != nil {
|
||||
t.Errorf("Ошибка: %v", err)
|
||||
return
|
||||
}
|
||||
return
|
||||
case <-ctx.Done():
|
||||
fmt.Println("Контекст отменен:", ctx.Err())
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchChatCompletionsAgentStreaming(t *testing.T) {
|
||||
client := openroute2.NewOpenRouterClient("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")
|
||||
agent := openroute2.NewRouterAgent(client, "qwen/qwen2.5-vl-72b-instruct:freet", openroute2.RouterAgentConfig{
|
||||
Temperature: 0.7,
|
||||
MaxTokens: 100,
|
||||
})
|
||||
|
||||
outputChan := make(chan openroute2.Response)
|
||||
processingChan := make(chan interface{})
|
||||
errChan := make(chan error)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
chat := []openroute2.MessageRequest{
|
||||
{Role: openroute2.RoleSystem, Content: "Вы полезный помощник."},
|
||||
{Role: openroute2.RoleUser, Content: "Привет!"},
|
||||
}
|
||||
|
||||
go agent.ChatStream(chat, outputChan, processingChan, errChan, ctx)
|
||||
|
||||
for {
|
||||
select {
|
||||
case output := <-outputChan:
|
||||
if len(output.Choices) > 0 {
|
||||
t.Logf("%s", output.Choices[0].Delta.Content)
|
||||
}
|
||||
case <-processingChan:
|
||||
t.Logf("Обработка\n")
|
||||
case err := <-errChan:
|
||||
if err != nil {
|
||||
t.Errorf("Ошибка: %v", err)
|
||||
return
|
||||
}
|
||||
return
|
||||
case <-ctx.Done():
|
||||
fmt.Println("Контекст отменен:", ctx.Err())
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchChatCompletionsAgentSimpleChat(t *testing.T) {
|
||||
client := openroute2.NewOpenRouterClient("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")
|
||||
agent := openroute2.NewRouterAgentChat(client, "qwen/qwen2.5-vl-72b-instruct:free", openroute2.RouterAgentConfig{
|
||||
Temperature: 0.0,
|
||||
MaxTokens: 100,
|
||||
}, "Вы полезный помощник, отвечайте короткими словами.")
|
||||
|
||||
agent.Chat("Запомни это: \"wojtess\"")
|
||||
agent.Chat("Что я просил вас запомнить?")
|
||||
|
||||
for _, msg := range agent.Messages {
|
||||
content, ok := msg.Content.(string)
|
||||
if ok {
|
||||
t.Logf("%s: %s", msg.Role, content)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user