langfun 0.0.2.dev20240409__tar.gz → 0.0.2.dev20240411__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (107) hide show
  1. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/PKG-INFO +1 -1
  2. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/eval/base_test.py +4 -4
  3. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/eval/matching_test.py +1 -1
  4. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/eval/scoring_test.py +2 -1
  5. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/langfunc_test.py +1 -1
  6. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/language_model.py +8 -1
  7. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/language_model_test.py +2 -2
  8. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/llms/cache/in_memory_test.py +22 -22
  9. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/llms/google_genai_test.py +8 -3
  10. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/llms/openai.py +2 -1
  11. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/llms/openai_test.py +0 -1
  12. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun.egg-info/PKG-INFO +1 -1
  13. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/LICENSE +0 -0
  14. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/README.md +0 -0
  15. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/__init__.py +0 -0
  16. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/__init__.py +0 -0
  17. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/coding/__init__.py +0 -0
  18. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/coding/python/__init__.py +0 -0
  19. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/coding/python/correction.py +0 -0
  20. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/coding/python/correction_test.py +0 -0
  21. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/coding/python/errors.py +0 -0
  22. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/coding/python/errors_test.py +0 -0
  23. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/coding/python/execution.py +0 -0
  24. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/coding/python/execution_test.py +0 -0
  25. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/coding/python/generation.py +0 -0
  26. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/coding/python/generation_test.py +0 -0
  27. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/coding/python/parsing.py +0 -0
  28. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/coding/python/parsing_test.py +0 -0
  29. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/coding/python/permissions.py +0 -0
  30. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/coding/python/permissions_test.py +0 -0
  31. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/component.py +0 -0
  32. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/component_test.py +0 -0
  33. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/concurrent.py +0 -0
  34. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/concurrent_test.py +0 -0
  35. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/console.py +0 -0
  36. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/console_test.py +0 -0
  37. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/eval/__init__.py +0 -0
  38. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/eval/base.py +0 -0
  39. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/eval/matching.py +0 -0
  40. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/eval/scoring.py +0 -0
  41. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/langfunc.py +0 -0
  42. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/llms/__init__.py +0 -0
  43. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/llms/cache/__init__.py +0 -0
  44. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/llms/cache/base.py +0 -0
  45. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/llms/cache/in_memory.py +0 -0
  46. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/llms/fake.py +0 -0
  47. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/llms/fake_test.py +0 -0
  48. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/llms/google_genai.py +0 -0
  49. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/llms/llama_cpp.py +0 -0
  50. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/llms/llama_cpp_test.py +0 -0
  51. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/memories/__init__.py +0 -0
  52. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/memories/conversation_history.py +0 -0
  53. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/memories/conversation_history_test.py +0 -0
  54. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/memory.py +0 -0
  55. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/message.py +0 -0
  56. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/message_test.py +0 -0
  57. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/modalities/__init__.py +0 -0
  58. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/modalities/image.py +0 -0
  59. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/modalities/image_test.py +0 -0
  60. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/modalities/mime.py +0 -0
  61. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/modalities/mime_test.py +0 -0
  62. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/modalities/video.py +0 -0
  63. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/modalities/video_test.py +0 -0
  64. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/modality.py +0 -0
  65. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/modality_test.py +0 -0
  66. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/natural_language.py +0 -0
  67. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/natural_language_test.py +0 -0
  68. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/sampling.py +0 -0
  69. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/sampling_test.py +0 -0
  70. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/__init__.py +0 -0
  71. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/completion.py +0 -0
  72. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/completion_test.py +0 -0
  73. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/description.py +0 -0
  74. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/description_test.py +0 -0
  75. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/mapping.py +0 -0
  76. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/mapping_test.py +0 -0
  77. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/parsing.py +0 -0
  78. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/parsing_test.py +0 -0
  79. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/prompting.py +0 -0
  80. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/prompting_test.py +0 -0
  81. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/schema.py +0 -0
  82. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/schema_generation.py +0 -0
  83. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/schema_generation_test.py +0 -0
  84. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/schema_test.py +0 -0
  85. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/scoring.py +0 -0
  86. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/structured/scoring_test.py +0 -0
  87. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/subscription.py +0 -0
  88. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/subscription_test.py +0 -0
  89. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/template.py +0 -0
  90. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/template_test.py +0 -0
  91. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/templates/__init__.py +0 -0
  92. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/templates/completion.py +0 -0
  93. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/templates/completion_test.py +0 -0
  94. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/templates/conversation.py +0 -0
  95. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/templates/conversation_test.py +0 -0
  96. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/templates/demonstration.py +0 -0
  97. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/templates/demonstration_test.py +0 -0
  98. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/templates/selfplay.py +0 -0
  99. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/templates/selfplay_test.py +0 -0
  100. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/text_formatting.py +0 -0
  101. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun/core/text_formatting_test.py +0 -0
  102. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun.egg-info/SOURCES.txt +0 -0
  103. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun.egg-info/dependency_links.txt +0 -0
  104. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun.egg-info/requires.txt +0 -0
  105. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/langfun.egg-info/top_level.txt +0 -0
  106. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/setup.cfg +0 -0
  107. {langfun-0.0.2.dev20240409 → langfun-0.0.2.dev20240411}/setup.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: langfun
3
- Version: 0.0.2.dev20240409
3
+ Version: 0.0.2.dev20240411
4
4
  Summary: Langfun: Language as Functions.
5
5
  Home-page: https://github.com/google/langfun
6
6
  Author: Langfun Authors
@@ -101,7 +101,7 @@ class EvaluationTest(unittest.TestCase):
101
101
  self.assertEqual(s.dir, os.path.join(s.root_dir, s.id))
102
102
  self.assertEqual(s.hash, s.clone().hash)
103
103
  # Test persistent hash.
104
- self.assertEqual(s.hash, 'abc7c29a')
104
+ self.assertEqual(s.hash, '4a2f9694')
105
105
  self.assertEqual(
106
106
  s.hash, s.clone(override={'max_workers': 2, 'lm.timeout': 20}).hash
107
107
  )
@@ -209,7 +209,7 @@ class EvaluationTest(unittest.TestCase):
209
209
  s.result,
210
210
  dict(
211
211
  experiment_setup=dict(
212
- id='Evaluation@17915dc6',
212
+ id='Evaluation@28030a68',
213
213
  dir=s.dir,
214
214
  model='StaticSequence',
215
215
  prompt_template='{{example.question}}',
@@ -321,7 +321,7 @@ class EvaluationTest(unittest.TestCase):
321
321
  s.children[0].dir, os.path.join(s.root_dir, s.children[0].id)
322
322
  )
323
323
  # Test persistent hash.
324
- self.assertEqual(s.hash, 'ca7f722b')
324
+ self.assertEqual(s.hash, 'e2a499d0')
325
325
 
326
326
  summary = s.run(verbose=True)
327
327
  self.assertEqual(len(summary.evaluations), 2)
@@ -448,7 +448,7 @@ class SuiteTest(unittest.TestCase):
448
448
  lm=lm
449
449
  )
450
450
  # Test for persistent hash.
451
- self.assertEqual(s.hash, '7285e52b')
451
+ self.assertEqual(s.hash, '8eb9e73f')
452
452
  s.run()
453
453
  expected = {
454
454
  s.children[0].id: dict(
@@ -103,7 +103,7 @@ class MatchingTest(unittest.TestCase):
103
103
  s.result,
104
104
  dict(
105
105
  experiment_setup=dict(
106
- id='MyTask@3d87f97f',
106
+ id='MyTask@9e3c5aa3',
107
107
  dir=s.dir,
108
108
  model='StaticSequence',
109
109
  prompt_template='{{example.question}}',
@@ -75,13 +75,14 @@ class ScoringTest(unittest.TestCase):
75
75
  ])
76
76
 
77
77
  s = eval_set(lm=lm)
78
+ self.maxDiff = None
78
79
  self.assertEqual(s.avg_score, 0.0)
79
80
  s.run()
80
81
  self.assertEqual(
81
82
  s.result,
82
83
  dict(
83
84
  experiment_setup=dict(
84
- id='ConstraintFollowing@9e51bb9e',
85
+ id='ConstraintFollowing@76285393',
85
86
  dir=s.dir,
86
87
  model='StaticSequence',
87
88
  prompt_template='{{example}}',
@@ -93,7 +93,7 @@ class LangFuncCallTest(unittest.TestCase):
93
93
  repr(l),
94
94
  "LangFunc(template_str='Hello', clean=True,"
95
95
  ' lm=ExcitedEchoer(sampling_options=LMSamplingOptions(temperature=0.0,'
96
- ' max_tokens=1024, n=1, top_k=40, top_p=None, stop=None,'
96
+ ' max_tokens=None, n=1, top_k=40, top_p=None, stop=None,'
97
97
  ' random_seed=None, logprobs=False, top_logprobs=None), cache=None,'
98
98
  ' max_concurrency=None, timeout=120.0, max_attempts=5,'
99
99
  ' retry_interval=(5, 60), exponential_backoff=True, debug=False))',
@@ -69,7 +69,14 @@ class LMSamplingOptions(component.Component):
69
69
  'OpenAI models have temperature range from 0.0 to 2.0.'
70
70
  )
71
71
  ] = 0.0
72
- max_tokens: Annotated[int, 'Per example max tokens to generate.'] = 1024
72
+ max_tokens: Annotated[
73
+ int | None,
74
+ (
75
+ 'Per example max tokens to generate. '
76
+ 'If None, use the model default.'
77
+ )
78
+ ] = None
79
+
73
80
  n: Annotated[int | None, 'Max number of samples to return.'] = 1
74
81
  top_k: Annotated[
75
82
  int | None,
@@ -73,13 +73,13 @@ class LMSamplingOptionsTest(unittest.TestCase):
73
73
  def test_cache_key(self):
74
74
  options = lm_lib.LMSamplingOptions()
75
75
  key1 = options.cache_key()
76
- self.assertEqual(key1, (0.0, 1024, 1, 40, None, None))
76
+ self.assertEqual(key1, (0.0, None, 1, 40, None, None))
77
77
  with options.override(temperature=1.0, max_tokens=256):
78
78
  key2 = options.cache_key()
79
79
  self.assertEqual(key2, (1.0, 256, 1, 40, None, None))
80
80
 
81
81
  # Make sure key1 does not change upon override.
82
- self.assertEqual(key1, (0.0, 1024, 1, 40, None, None))
82
+ self.assertEqual(key1, (0.0, None, 1, 40, None, None))
83
83
 
84
84
 
85
85
  class LanguageModelTest(unittest.TestCase):
@@ -44,19 +44,19 @@ class InMemoryLMCacheTest(unittest.TestCase):
44
44
  self.assertEqual(
45
45
  list(cache.keys()),
46
46
  [
47
- ('a', (0.0, 1024, 1, 40, None, None), 0),
48
- ('a', (0.0, 1024, 1, 40, None, None), 1),
49
- ('b', (0.0, 1024, 1, 40, None, None), 0),
50
- ('c', (0.0, 1024, 1, 40, None, None), 0),
47
+ ('a', (0.0, None, 1, 40, None, None), 0),
48
+ ('a', (0.0, None, 1, 40, None, None), 1),
49
+ ('b', (0.0, None, 1, 40, None, None), 0),
50
+ ('c', (0.0, None, 1, 40, None, None), 0),
51
51
  ],
52
52
  )
53
53
  self.assertEqual(
54
54
  list(cache.keys('StaticSequence')),
55
55
  [
56
- ('a', (0.0, 1024, 1, 40, None, None), 0),
57
- ('a', (0.0, 1024, 1, 40, None, None), 1),
58
- ('b', (0.0, 1024, 1, 40, None, None), 0),
59
- ('c', (0.0, 1024, 1, 40, None, None), 0),
56
+ ('a', (0.0, None, 1, 40, None, None), 0),
57
+ ('a', (0.0, None, 1, 40, None, None), 1),
58
+ ('b', (0.0, None, 1, 40, None, None), 0),
59
+ ('c', (0.0, None, 1, 40, None, None), 0),
60
60
  ],
61
61
  )
62
62
 
@@ -90,19 +90,19 @@ class InMemoryLMCacheTest(unittest.TestCase):
90
90
  list(cache.items()),
91
91
  [
92
92
  (
93
- ('a', (0.0, 1024, 1, 40, None, None), 0),
93
+ ('a', (0.0, None, 1, 40, None, None), 0),
94
94
  cache_entry('1'),
95
95
  ),
96
96
  (
97
- ('a', (0.0, 1024, 1, 40, None, None), 1),
97
+ ('a', (0.0, None, 1, 40, None, None), 1),
98
98
  cache_entry('2', 1),
99
99
  ),
100
100
  (
101
- ('b', (0.0, 1024, 1, 40, None, None), 0),
101
+ ('b', (0.0, None, 1, 40, None, None), 0),
102
102
  cache_entry('3'),
103
103
  ),
104
104
  (
105
- ('c', (0.0, 1024, 1, 40, None, None), 0),
105
+ ('c', (0.0, None, 1, 40, None, None), 0),
106
106
  cache_entry('4'),
107
107
  ),
108
108
  ],
@@ -111,19 +111,19 @@ class InMemoryLMCacheTest(unittest.TestCase):
111
111
  list(cache.items('StaticSequence')),
112
112
  [
113
113
  (
114
- ('a', (0.0, 1024, 1, 40, None, None), 0),
114
+ ('a', (0.0, None, 1, 40, None, None), 0),
115
115
  cache_entry('1'),
116
116
  ),
117
117
  (
118
- ('a', (0.0, 1024, 1, 40, None, None), 1),
118
+ ('a', (0.0, None, 1, 40, None, None), 1),
119
119
  cache_entry('2', 1),
120
120
  ),
121
121
  (
122
- ('b', (0.0, 1024, 1, 40, None, None), 0),
122
+ ('b', (0.0, None, 1, 40, None, None), 0),
123
123
  cache_entry('3'),
124
124
  ),
125
125
  (
126
- ('c', (0.0, 1024, 1, 40, None, None), 0),
126
+ ('c', (0.0, None, 1, 40, None, None), 0),
127
127
  cache_entry('4'),
128
128
  ),
129
129
  ],
@@ -161,8 +161,8 @@ class InMemoryLMCacheTest(unittest.TestCase):
161
161
  self.assertEqual(
162
162
  list(cache.keys()),
163
163
  [
164
- ('a', (0.0, 1024, 1, 40, None, None), 0),
165
- ('a', (1.0, 1024, 1, 40, None, None), 0),
164
+ ('a', (0.0, None, 1, 40, None, None), 0),
165
+ ('a', (1.0, None, 1, 40, None, None), 0),
166
166
  ],
167
167
  )
168
168
 
@@ -180,15 +180,15 @@ class InMemoryLMCacheTest(unittest.TestCase):
180
180
  self.assertEqual(
181
181
  list(cache.keys('StaticSequence')),
182
182
  [
183
- ('a', (0.0, 1024, 1, 40, None, None), 0),
184
- ('b', (0.0, 1024, 1, 40, None, None), 0),
183
+ ('a', (0.0, None, 1, 40, None, None), 0),
184
+ ('b', (0.0, None, 1, 40, None, None), 0),
185
185
  ],
186
186
  )
187
187
  self.assertEqual(
188
188
  list(cache.keys('Echo')),
189
189
  [
190
- ('a', (0.0, 1024, 1, 40, None, None), 0),
191
- ('b', (0.0, 1024, 1, 40, None, None), 0),
190
+ ('a', (0.0, None, 1, 40, None, None), 0),
191
+ ('b', (0.0, None, 1, 40, None, None), 0),
192
192
  ],
193
193
  )
194
194
  self.assertEqual(len(cache), 4)
@@ -152,10 +152,15 @@ class GenAITest(unittest.TestCase):
152
152
  )
153
153
 
154
154
  def test_model_hub(self):
155
+ orig_get_model = genai.get_model
156
+ genai.get_model = mock_get_model
157
+
155
158
  model = google_genai._GOOGLE_GENAI_MODEL_HUB.get('gemini-pro')
156
159
  self.assertIsNotNone(model)
157
160
  self.assertIs(google_genai._GOOGLE_GENAI_MODEL_HUB.get('gemini-pro'), model)
158
161
 
162
+ genai.get_model = orig_get_model
163
+
159
164
  def test_api_key_check(self):
160
165
  with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
161
166
  _ = google_genai.GeminiPro()._api_initialized
@@ -167,7 +172,7 @@ class GenAITest(unittest.TestCase):
167
172
 
168
173
  def test_call(self):
169
174
  with mock.patch(
170
- 'google.generativeai.generative_models.GenerativeModel.generate_content'
175
+ 'google.generativeai.GenerativeModel.generate_content',
171
176
  ) as mock_generate:
172
177
  orig_get_model = genai.get_model
173
178
  genai.get_model = mock_get_model
@@ -176,7 +181,7 @@ class GenAITest(unittest.TestCase):
176
181
  lm = google_genai.GeminiPro(api_key='test_key')
177
182
  self.maxDiff = None
178
183
  self.assertEqual(
179
- lm('hello', temperature=2.0, top_k=20).text,
184
+ lm('hello', temperature=2.0, top_k=20, max_tokens=1024).text,
180
185
  (
181
186
  'This is a response to hello with n=1, temperature=2.0, '
182
187
  'top_p=None, top_k=20, max_tokens=1024, stop=None.'
@@ -197,7 +202,7 @@ class GenAITest(unittest.TestCase):
197
202
  (
198
203
  "hello to models/text-bison-001 with {'temperature': 2.0, "
199
204
  "'top_k': 20, 'top_p': None, 'candidate_count': 1, "
200
- "'max_output_tokens': 1024, 'stop_sequences': None}"
205
+ "'max_output_tokens': None, 'stop_sequences': None}"
201
206
  ),
202
207
  )
203
208
  genai.get_model = orig_get_model
@@ -164,7 +164,6 @@ class OpenAI(lf.LanguageModel):
164
164
  args = dict(
165
165
  n=options.n,
166
166
  temperature=options.temperature,
167
- max_tokens=options.max_tokens,
168
167
  stream=False,
169
168
  timeout=self.timeout,
170
169
  logprobs=options.logprobs,
@@ -173,6 +172,8 @@ class OpenAI(lf.LanguageModel):
173
172
  # Completion and ChatCompletion uses different parameter name for model.
174
173
  args['model' if self.is_chat_model else 'engine'] = self.model
175
174
 
175
+ if options.max_tokens is not None:
176
+ args['max_tokens'] = options.max_tokens
176
177
  if options.top_p is not None:
177
178
  args['top_p'] = options.top_p
178
179
  if options.stop:
@@ -121,7 +121,6 @@ class OpenaiTest(unittest.TestCase):
121
121
  top_logprobs=None,
122
122
  n=1,
123
123
  temperature=1.0,
124
- max_tokens=1024,
125
124
  stream=False,
126
125
  timeout=120.0,
127
126
  stop=['\n'],
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: langfun
3
- Version: 0.0.2.dev20240409
3
+ Version: 0.0.2.dev20240411
4
4
  Summary: Langfun: Language as Functions.
5
5
  Home-page: https://github.com/google/langfun
6
6
  Author: Langfun Authors