Phase 3: PDF/DOCX extraction, chunking, LLM client with mock interface
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Jānis Kacēns
@@ -0,0 +1,119 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
openai "github.com/sashabaranov/go-openai"
|
||||
)
|
||||
|
||||
// ParsedQuestion is a question extracted from a document chunk by the LLM.
|
||||
type ParsedQuestion struct {
|
||||
Question string
|
||||
Answers []ParsedAnswer
|
||||
}
|
||||
|
||||
// ParsedAnswer is one answer choice for a ParsedQuestion.
|
||||
type ParsedAnswer struct {
|
||||
Text string
|
||||
Correct bool
|
||||
}
|
||||
|
||||
// ChatClient is the interface for creating chat completions.
|
||||
// The concrete *openai.Client satisfies this interface.
|
||||
type ChatClient interface {
|
||||
CreateChatCompletion(ctx context.Context, req openai.ChatCompletionRequest) (openai.ChatCompletionResponse, error)
|
||||
}
|
||||
|
||||
// Client wraps a ChatClient with question-extraction logic.
|
||||
type Client struct {
|
||||
cc ChatClient
|
||||
model string
|
||||
}
|
||||
|
||||
// New creates a Client backed by the real OpenAI API.
|
||||
func New(apiKey, model string) *Client {
|
||||
if model == "" {
|
||||
model = "gpt-4o-mini"
|
||||
}
|
||||
return &Client{cc: openai.NewClient(apiKey), model: model}
|
||||
}
|
||||
|
||||
// NewWithClient creates a Client with an injected ChatClient (useful for tests).
|
||||
func NewWithClient(cc ChatClient, model string) *Client {
|
||||
return &Client{cc: cc, model: model}
|
||||
}
|
||||
|
||||
const systemPrompt = `You extract multiple-choice questions from study material. Return every question found. Exactly one answer per question must be marked correct. If the source doesn't clearly mark a correct answer, omit that question entirely. Do not invent questions not present in the text.
|
||||
|
||||
Respond with JSON matching this schema exactly:
|
||||
{"questions":[{"question":"<text>","answers":[{"text":"<text>","correct":false},{"text":"<text>","correct":true}]}]}`
|
||||
|
||||
type llmResponse struct {
|
||||
Questions []struct {
|
||||
Question string `json:"question"`
|
||||
Answers []struct {
|
||||
Text string `json:"text"`
|
||||
Correct bool `json:"correct"`
|
||||
} `json:"answers"`
|
||||
} `json:"questions"`
|
||||
}
|
||||
|
||||
// ExtractQuestions sends chunk to the LLM and returns validated, deduplicated questions.
|
||||
// Questions that do not have exactly one correct answer are silently dropped.
|
||||
func (c *Client) ExtractQuestions(ctx context.Context, chunk string) ([]ParsedQuestion, error) {
|
||||
resp, err := c.cc.CreateChatCompletion(ctx, openai.ChatCompletionRequest{
|
||||
Model: c.model,
|
||||
Messages: []openai.ChatCompletionMessage{
|
||||
{Role: openai.ChatMessageRoleSystem, Content: systemPrompt},
|
||||
{Role: openai.ChatMessageRoleUser, Content: chunk},
|
||||
},
|
||||
ResponseFormat: &openai.ChatCompletionResponseFormat{
|
||||
Type: openai.ChatCompletionResponseFormatTypeJSONObject,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("openai: %w", err)
|
||||
}
|
||||
if len(resp.Choices) == 0 {
|
||||
return nil, fmt.Errorf("openai: empty response")
|
||||
}
|
||||
|
||||
var raw llmResponse
|
||||
if err := json.Unmarshal([]byte(resp.Choices[0].Message.Content), &raw); err != nil {
|
||||
return nil, fmt.Errorf("parse llm response: %w", err)
|
||||
}
|
||||
|
||||
seen := make(map[string]bool)
|
||||
var out []ParsedQuestion
|
||||
for _, q := range raw.Questions {
|
||||
var nCorrect int
|
||||
for _, a := range q.Answers {
|
||||
if a.Correct {
|
||||
nCorrect++
|
||||
}
|
||||
}
|
||||
if nCorrect != 1 {
|
||||
continue
|
||||
}
|
||||
key := textHash(q.Question)
|
||||
if seen[key] {
|
||||
continue
|
||||
}
|
||||
seen[key] = true
|
||||
|
||||
pq := ParsedQuestion{Question: q.Question}
|
||||
for _, a := range q.Answers {
|
||||
pq.Answers = append(pq.Answers, ParsedAnswer{Text: a.Text, Correct: a.Correct})
|
||||
}
|
||||
out = append(out, pq)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func textHash(s string) string {
|
||||
h := sha256.Sum256([]byte(s))
|
||||
return fmt.Sprintf("%x", h[:8])
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
package llm_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
openai "github.com/sashabaranov/go-openai"
|
||||
|
||||
"qbank/internal/llm"
|
||||
)
|
||||
|
||||
// mockChat implements llm.ChatClient for testing.
|
||||
type mockChat struct{ body string }
|
||||
|
||||
func (m *mockChat) CreateChatCompletion(_ context.Context, _ openai.ChatCompletionRequest) (openai.ChatCompletionResponse, error) {
|
||||
return openai.ChatCompletionResponse{
|
||||
Choices: []openai.ChatCompletionChoice{
|
||||
{Message: openai.ChatCompletionMessage{Content: m.body}},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func mockClient(t *testing.T, questions []map[string]any) *llm.Client {
|
||||
t.Helper()
|
||||
body, err := json.Marshal(map[string]any{"questions": questions})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return llm.NewWithClient(&mockChat{body: string(body)}, "test-model")
|
||||
}
|
||||
|
||||
func TestExtractQuestions_HappyPath(t *testing.T) {
|
||||
qs, err := mockClient(t, []map[string]any{
|
||||
{
|
||||
"question": "What is 2+2?",
|
||||
"answers": []map[string]any{
|
||||
{"text": "3", "correct": false},
|
||||
{"text": "4", "correct": true},
|
||||
{"text": "5", "correct": false},
|
||||
},
|
||||
},
|
||||
}).ExtractQuestions(context.Background(), "text")
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("ExtractQuestions: %v", err)
|
||||
}
|
||||
if len(qs) != 1 {
|
||||
t.Fatalf("want 1 question, got %d", len(qs))
|
||||
}
|
||||
if qs[0].Question != "What is 2+2?" {
|
||||
t.Errorf("wrong question text: %q", qs[0].Question)
|
||||
}
|
||||
if len(qs[0].Answers) != 3 {
|
||||
t.Errorf("want 3 answers, got %d", len(qs[0].Answers))
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractQuestions_DropsInvalid(t *testing.T) {
|
||||
qs, err := mockClient(t, []map[string]any{
|
||||
{
|
||||
"question": "Two correct — should drop",
|
||||
"answers": []map[string]any{
|
||||
{"text": "A", "correct": true},
|
||||
{"text": "B", "correct": true},
|
||||
},
|
||||
},
|
||||
{
|
||||
"question": "Zero correct — should drop",
|
||||
"answers": []map[string]any{
|
||||
{"text": "A", "correct": false},
|
||||
{"text": "B", "correct": false},
|
||||
},
|
||||
},
|
||||
{
|
||||
"question": "Valid question",
|
||||
"answers": []map[string]any{
|
||||
{"text": "Wrong", "correct": false},
|
||||
{"text": "Right", "correct": true},
|
||||
},
|
||||
},
|
||||
}).ExtractQuestions(context.Background(), "text")
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("ExtractQuestions: %v", err)
|
||||
}
|
||||
if len(qs) != 1 {
|
||||
t.Fatalf("want 1 question after dropping invalid, got %d", len(qs))
|
||||
}
|
||||
if qs[0].Question != "Valid question" {
|
||||
t.Errorf("wrong question kept: %q", qs[0].Question)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractQuestions_Dedup(t *testing.T) {
|
||||
qs, err := mockClient(t, []map[string]any{
|
||||
{
|
||||
"question": "Duplicate?",
|
||||
"answers": []map[string]any{
|
||||
{"text": "Yes", "correct": true},
|
||||
{"text": "No", "correct": false},
|
||||
},
|
||||
},
|
||||
{
|
||||
"question": "Duplicate?",
|
||||
"answers": []map[string]any{
|
||||
{"text": "Yes", "correct": true},
|
||||
{"text": "No", "correct": false},
|
||||
},
|
||||
},
|
||||
}).ExtractQuestions(context.Background(), "text")
|
||||
|
||||
if err != nil {
|
||||
t.Fatalf("ExtractQuestions: %v", err)
|
||||
}
|
||||
if len(qs) != 1 {
|
||||
t.Errorf("want 1 unique question after dedup, got %d", len(qs))
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractQuestions_EmptyResponse(t *testing.T) {
|
||||
qs, err := mockClient(t, []map[string]any{}).ExtractQuestions(context.Background(), "text")
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if len(qs) != 0 {
|
||||
t.Errorf("want 0 questions for empty response, got %d", len(qs))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user