context_test.go 600 B

12345678910111213141516171819202122232425262728
  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. ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) // TODO maybe shorter?
  11. defer cancel()
  12. // Set up the test data
  13. req := api.GenerateRequest{
  14. Model: "llama2",
  15. Prompt: "Write me a story with a ton of emojis?",
  16. Stream: &stream,
  17. Options: map[string]interface{}{
  18. "temperature": 0,
  19. "seed": 123,
  20. "num_ctx": 128,
  21. },
  22. }
  23. GenerateTestHelper(ctx, t, req, []string{"once", "upon", "lived"})
  24. }