langfun 0.0.2.dev20240429__py3-none-any.whl → 0.1.2.dev202501150804__py3-none-any.whl

langfun/core/llms/groq_test.py

@@ -11,89 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for Groq models."""
-
 import os
-from typing import Any
 import unittest
-from unittest import mock
-from langfun.core import modalities as lf_modalities
+import langfun.core as lf
 from langfun.core.llms import groq
-import pyglove as pg
-import requests
-
-
-def mock_requests_post(url: str, json: dict[str, Any], **kwargs):
-  del url, kwargs
-
-  response = requests.Response()
-  response.status_code = 200
-  response._content = pg.to_json_str({
-      'choices': [{
-          'message': {
-              'content': [{
-                  'type': 'text',
-                  'text': (
-                      f'hello with temperature={json.get("temperature")}, '
-                      f'top_p={json.get("top_p")}, '
-                      f'max_tokens={json.get("max_tokens")}, '
-                      f'stop={json.get("stop")}.'
-                  ),
-              }],
-          }
-      }],
-      'usage': {
-          'prompt_tokens': 2,
-          'completion_tokens': 1,
-          'total_tokens': 3,
-      },
-  }).encode()
-  return response
-
-
-def mock_mm_requests_post(url: str, json: dict[str, Any], **kwargs):
-  del url, kwargs
-  v = json['messages'][0]['content'][0]
-  image = lf_modalities.Image.from_uri(v['image_url'])
-
-  response = requests.Response()
-  response.status_code = 200
-  response._content = pg.to_json_str({
-      'choices': [
-          {
-              'message': {
-                  'content': [{
-                      'type': 'text',
-                      'text': image.uri,
-                  }],
-              }
-          }
-      ],
-      'usage': {
-          'prompt_tokens': 2,
-          'completion_tokens': 1,
-          'total_tokens': 3,
-      },
-  }).encode()
-  return response
-
-
-def mock_requests_post_error(status_code, error_type, error_message):
-  def _mock_requests(url: str, json: dict[str, Any], **kwargs):
-    del url, json, kwargs
-    response = requests.Response()
-    response.status_code = status_code
-    response._content = pg.to_json_str(
-        {
-            'error': {
-                'type': error_type,
-                'message': error_message,
-            }
-        }
-    ).encode()
-    return response
-
-  return _mock_requests
 
 
 class AuthropicTest(unittest.TestCase):
@@ -101,69 +22,42 @@ class AuthropicTest(unittest.TestCase):
   def test_basics(self):
     self.assertEqual(groq.GroqMistral_8x7B().model_id, 'mixtral-8x7b-32768')
     self.assertEqual(groq.GroqMistral_8x7B().max_concurrency, 16)
+    self.assertEqual(groq.GroqMistral_8x7B().estimate_cost(100, 100), 4.8e-5)
+
+  def test_request_args(self):
+    args = groq.GroqMistral_8x7B()._request_args(
+        lf.LMSamplingOptions(
+            temperature=1.0, stop=['\n'], n=1, random_seed=123,
+            logprobs=True, top_logprobs=True
+        )
+    )
+    self.assertNotIn('logprobs', args)
+    self.assertNotIn('top_logprobs', args)
 
   def test_api_key(self):
     lm = groq.GroqMistral_8x7B()
     with self.assertRaisesRegex(ValueError, 'Please specify `api_key`'):
-      lm('hi')
-
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_requests_post
-
-      lm = groq.GroqMistral_8x7B(api_key='fake key')
-      self.assertRegex(lm('hi').text, 'hello.*')
-
-      os.environ['GROQ_API_KEY'] = 'abc'
-      lm = groq.GroqMistral_8x7B()
-      self.assertRegex(lm('hi').text, 'hello.*')
-      del os.environ['GROQ_API_KEY']
-
-  def test_call(self):
-    with mock.patch('requests.Session.post') as mock_request:
-      mock_request.side_effect = mock_requests_post
-      lm = groq.GroqLlama3_70B(api_key='fake_key')
-      response = lm(
-          'hello',
-          temperature=0.0,
-          max_tokens=1024,
-          top_k=0.1,
-          top_p=0.2,
-          stop=['\n'],
-      )
-      self.assertEqual(
-          response.text,
-          (
-              'hello with temperature=0.0, top_p=0.2, '
-              "max_tokens=1024, stop=['\\n']."
-          ),
-      )
-      self.assertIsNotNone(response.usage)
-      self.assertIsNotNone(response.usage.prompt_tokens, 2)
-      self.assertIsNotNone(response.usage.completion_tokens, 1)
-      self.assertIsNotNone(response.usage.total_tokens, 3)
+      _ = lm.headers
 
-  def test_mm_call(self):
-    with mock.patch('requests.Session.post') as mock_mm_request:
-      mock_mm_request.side_effect = mock_mm_requests_post
-      lm = groq.GroqLlama3_70B(multimodal=True, api_key='fake_key')
-      response = lm(lf_modalities.Image.from_uri('https://fake/image.jpg'))
-      self.assertEqual(response.text, 'https://fake/image.jpg')
+    lm = groq.GroqMistral_8x7B(api_key='fake key')
+    self.assertEqual(
+        lm.headers,
+        {
+            'Content-Type': 'application/json',
+            'Authorization': 'Bearer fake key',
+        }
+    )
 
-  def test_call_errors(self):
-    for status_code, error_type, error_message in [
-        (429, 'rate_limit', 'Rate limit exceeded.'),
-        (503, 'service_unavailable', 'Service unavailable.'),
-        (500, 'bad_request', 'Bad request.'),
-    ]:
-      with mock.patch('requests.Session.post') as mock_mm_request:
-        mock_mm_request.side_effect = mock_requests_post_error(
-            status_code, error_type, error_message
-        )
-        lm = groq.GroqLlama3_70B(api_key='fake_key')
-        with self.assertRaisesRegex(
-            Exception, f'{status_code}:.*{error_type}'
-        ):
-          lm('hello', lm=lm, max_attempts=1)
+    os.environ['GROQ_API_KEY'] = 'abc'
+    lm = groq.GroqMistral_8x7B()
+    self.assertEqual(
+        lm.headers,
+        {
+            'Content-Type': 'application/json',
+            'Authorization': 'Bearer abc',
+        }
+    )
+    del os.environ['GROQ_API_KEY']
 
 
 if __name__ == '__main__':

langfun/core/llms/llama_cpp.py

@@ -14,61 +14,34 @@
 """Language models from llama.cpp."""
 
 from typing import Annotated
+from langfun.core.llms import openai_compatible
+import pyglove as pg
 
-import langfun.core as lf
-import requests
 
-
-@lf.use_init_args(["url"])
-class LlamaCppRemote(lf.LanguageModel):
+@pg.use_init_args(['url', 'model'])
+@pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
+class LlamaCppRemote(openai_compatible.OpenAICompatible):
   """The remote LLaMA C++ model.
 
   The Remote LLaMA C++ models can be launched via
   https://github.com/ggerganov/llama.cpp/tree/master/examples/server
   """
-
   url: Annotated[
       str,
-      "The name of the model to use.",
-  ] = ""
+      'The URL of the LLaMA C++ server.',
+  ]
 
-  name: Annotated[
+  model: Annotated[
       str,
-      "The abbreviation for the LLaMA CPP-based model name.",
-  ] = ""
+      'The name of the model to use.',
+  ] = ''
+
+  @property
+  def api_endpoint(self) -> str:
+    return self.url + '/completion'
 
   @property
   def model_id(self) -> str:
     """Returns a string to identify the model."""
-    return f"LLaMAC++({self.name})"
-
-  def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]:
-    def _complete_fn(cur_prompts):
-      results = []
-      for prompt in cur_prompts:
-        result = lf.LMSamplingResult()
-        for _ in range(self.sampling_options.n or 1):
-          data = {
-              "prompt": prompt.text,
-              "n_predict": self.sampling_options.max_tokens,
-              "top_k": self.sampling_options.top_k or 50,
-              "top_p": self.sampling_options.top_p or 0.95,
-          }
-          if self.sampling_options.temperature is not None:
-            data["temperature"] = self.sampling_options.temperature
-
-          response = requests.post(
-              f"{self.url}/completion",
-              json=data,
-              headers={"Content-Type": "application/json"},
-              timeout=self.timeout,
-          )
-          decoded_response = response.json()
-          response = decoded_response["content"]
-          result.samples.append(lf.LMSample(response, score=0.0))
-        results.append(result)
-      return results
+    return f'LLaMAC++({self.model or ""})'
 
-    return self._parallel_execute_with_currency_control(
-        _complete_fn, [prompts]
-    )[0]

langfun/core/llms/llama_cpp_test.py

@@ -11,44 +11,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for llama cpp models."""
-
-import typing
 import unittest
-from unittest import mock
-
-import langfun.core as lf
 from langfun.core.llms import llama_cpp
 
 
-def mock_requests_post(url: str, json: typing.Dict[str, typing.Any], **kwargs):
-  del kwargs
-
-  class TEMP:
-
-    def json(self):
-      return {"content": json["prompt"] + "\n" + url}
-
-  return TEMP()
-
-
 class LlamaCppRemoteTest(unittest.TestCase):
   """Tests for the LlamaCppRemote model."""
 
-  def test_call_completion(self):
-    with mock.patch("requests.post") as mock_request:
-      mock_request.side_effect = mock_requests_post
-      lm = llama_cpp.LlamaCppRemote(url="http://127.0.0.1:8080")
-      response = lm("hello", sampling_options=lf.LMSamplingOptions(n=1))
-      self.assertEqual(
-          response.text,
-          "hello\nhttp://127.0.0.1:8080/completion",
-      )
-
-  def test_name(self):
-    lm = llama_cpp.LlamaCppRemote()
+  def test_basics(self):
+    lm = llama_cpp.LlamaCppRemote("http://127.0.0.1:8080")
+    self.assertEqual(lm.api_endpoint, "http://127.0.0.1:8080/completion")
     self.assertEqual(lm.model_id, "LLaMAC++()")
-    lm = llama_cpp.LlamaCppRemote(url="xxx", name="x")
+    lm = llama_cpp.LlamaCppRemote("xxx", model="x")
     self.assertEqual(lm.model_id, "LLaMAC++(x)")