context_test.go 620 B

123456789101112131415161718192021222324252627282930
  1. //go:build integration
  2. package integration
  3. import (
  4. "context"
  5. "testing"
  6. "time"
  7. "github.com/ollama/ollama/api"
  8. )
  9. func TestContextExhaustion(t *testing.T) {
  10. // Longer needed for small footprint GPUs
  11. ctx, cancel := context.WithTimeout(context.Background(), 6*time.Minute)
  12. defer cancel()
  13. // Set up the test data
  14. req := api.GenerateRequest{
  15. Model: "llama2",
  16. Prompt: "Write me a story with a ton of emojis?",
  17. Stream: &stream,
  18. Options: map[string]interface{}{
  19. "temperature": 0,
  20. "seed": 123,
  21. "num_ctx": 128,
  22. },
  23. }
  24. GenerateTestHelper(ctx, t, req, []string{"once", "upon", "lived"})
  25. }