langfun 0.1.2.dev202510200805__py3-none-any.whl → 0.1.2.dev202511160804__py3-none-any.whl

langfun/core/llms/deepseek.py

@@ -93,8 +93,36 @@ _SUPPORTED_MODELS_BY_ID = {m.model_id: m for m in SUPPORTED_MODELS}
 # DeepSeek API uses an API format compatible with OpenAI.
 # Reference: https://api-docs.deepseek.com/
 @lf.use_init_args(['model'])
-class DeepSeek(openai_compatible.OpenAICompatible):
-  """DeepSeek model."""
+class DeepSeek(openai_compatible.OpenAIChatCompletionAPI):
+  """DeepSeek models.
+
+  **Quick Start:**
+
+  ```python
+  import langfun as lf
+
+  # Call DeepSeek-V3 using API key from environment variable
+  # 'DEEPSEEK_API_KEY'.
+  lm = lf.llms.DeepSeekV3()
+  r = lm('Who are you?')
+  print(r)
+  ```
+
+  **Setting up API key:**
+
+  The DeepSeek API key can be specified in following ways:
+
+  1. At model instantiation:
+
+     ```python
+     lm = lf.llms.DeepSeekV3(api_key='MY_API_KEY')
+     ```
+  2. via environment variable `DEEPSEEK_API_KEY`.
+
+  **References:**
+
+  *   https://api-docs.deepseek.com/
+  """
 
   model: pg.typing.Annotated[
       pg.typing.Enum(

langfun/core/llms/fake.py

@@ -20,7 +20,38 @@ import langfun.core as lf
 
 
 class Fake(lf.LanguageModel):
-  """The base class for all fake language models."""
+  """Base class for fake language models, used for testing.
+
+  Fake models simulate the behavior of real language models but return
+  pre-defined responses, making them useful for testing prompts,
+  data processing logic, and agent behavior without incurring API costs
+  or relying on external services.
+
+  Langfun provides several fake models:
+  * `lf.llms.Echo`: Echoes the prompt back as the response.
+  * `lf.llms.StaticResponse`: Returns a fixed, pre-defined response for
+    any prompt.
+  * `lf.llms.StaticMapping`: Returns responses based on a prompt-to-response
+    dictionary.
+  * `lf.llms.StaticSequence`: Returns responses from a pre-defined sequence
+    in order.
+
+  **Example:**
+
+  ```python
+  import langfun as lf
+
+  # Use Echo model for testing
+  lm = lf.llms.Echo()
+  response = lm('hello')
+  assert response.text == 'hello'
+
+  # Use StaticResponse model
+  lm = lf.llms.StaticResponse('world')
+  response = lm('hello')
+  assert response.text == 'world'
+  ```
+  """
 
   def _score(self, prompt: lf.Message| list[lf.Message],
              completions: list[lf.Message]):

langfun/core/llms/gemini.py

@@ -696,7 +696,15 @@ _SUPPORTED_MODELS_BY_ID = {m.model_id: m for m in SUPPORTED_MODELS}
 
 @pg.use_init_args(['model'])
 class Gemini(rest.REST):
-  """Language models provided by Google GenAI."""
+  """Base class for Gemini models served on Google GenAI and Vertex AI.
+
+  This class implements the Gemini API protocol, shared by
+  `lf.llms.GoogleGenAI` and `lf.llms.VertexAI`, providing common request
+  formatting and response parsing for Gemini models.
+
+  It is not intended to be used directly. Please use `lf.llms.GoogleGenAI` or
+  `lf.llms.VertexAI` instead.
+  """
 
   model: pg.typing.Annotated[
       pg.typing.Enum(
@@ -752,11 +760,8 @@ class Gemini(rest.REST):
         prompt.as_format('gemini', chunk_preprocessor=modality_conversion)
     )
     request['contents'] = contents
-    # Users could use `metadata_gemini_tools` to pass Gemini tools. For example,
-    # for enabling Search Grounding, users could pass:
-    # metadata_gemini_tools=[{'google_search': {}}]
-    if tools := prompt.metadata.get('gemini_tools'):
-      request['tools'] = tools
+    if sampling_options.extras:
+      request.update(sampling_options.extras)
     return request
 
   def _generation_config(
@@ -833,9 +838,9 @@ class Gemini(rest.REST):
     )
 
   def _error(self, status_code: int, content: str) -> lf.LMError:
-    if (
-        status_code == 400
-        and b'exceeds the maximum number of tokens' in content
+    if status_code == 400 and (
+        b'exceeds the maximum number of tokens' in content
+        or b'Reduce the input token count and try again.' in content
     ):
       return lf.ContextLimitError(f'{status_code}: {content}')
     return super()._error(status_code, content)

langfun/core/llms/google_genai.py

@@ -25,7 +25,35 @@ import pyglove as pg
 @lf.use_init_args(['model'])
 @pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
 class GenAI(gemini.Gemini):
-  """Language models provided by Google GenAI."""
+  """Google GenAI models.
+
+  **Quick Start:**
+
+  ```python
+  import langfun as lf
+
+  # Call Gemini 1.5 Flash using API key from environment variable
+  # 'GOOGLE_API_KEY'.
+  lm = lf.llms.Gemini15Flash()
+  r = lm('Who are you?')
+  print(r)
+  ```
+
+  **Setting up API key:**
+
+  The Google API key can be specified in following ways:
+
+  1. At model instantiation:
+
+     ```python
+     lm = lf.llms.Gemini15Flash(api_key='MY_API_KEY')
+     ```
+  2. via environment variable `GOOGLE_API_KEY`.
+
+  **References:**
+
+  *   https://ai.google.dev/docs
+  """
 
   model: pg.typing.Annotated[
       pg.typing.Enum(

langfun/core/llms/groq.py

@@ -259,10 +259,35 @@ _SUPPORTED_MODELS_BY_ID = {m.model_id: m for m in SUPPORTED_MODELS}
 
 
 @lf.use_init_args(['model'])
-class Groq(openai_compatible.OpenAICompatible):
-  """Groq LLMs through REST APIs (OpenAI compatible).
+class Groq(openai_compatible.OpenAIChatCompletionAPI):
+  """Groq models.
 
-  See https://platform.openai.com/docs/api-reference/chat
+  **Quick Start:**
+
+  ```python
+  import langfun as lf
+
+  # Call Llama 3.3 70B on Groq using API key from environment variable
+  # 'GROQ_API_KEY'.
+  lm = lf.llms.GroqLlama33_70B_Versatile()
+  r = lm('Who are you?')
+  print(r)
+  ```
+
+  **Setting up API key:**
+
+  The Groq API key can be specified in following ways:
+
+  1. At model instantiation:
+
+     ```python
+     lm = lf.llms.GroqLlama33_70B_Versatile(api_key='MY_API_KEY')
+     ```
+  2. via environment variable `GROQ_API_KEY`.
+
+  **References:**
+
+  *   https://console.groq.com/docs
   """
 
   model: pg.typing.Annotated[

langfun/core/llms/llama_cpp.py

@@ -20,11 +20,30 @@ import pyglove as pg
 
 @pg.use_init_args(['url', 'model'])
 @pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
-class LlamaCppRemote(openai_compatible.OpenAICompatible):
-  """The remote LLaMA C++ model.
+class LlamaCppRemote(openai_compatible.OpenAIChatCompletionAPI):
+  """LLaMA C++ models served via a remote server.
 
-  The Remote LLaMA C++ models can be launched via
-  https://github.com/ggerganov/llama.cpp/tree/master/examples/server
+  This class provides an interface to interact with language models
+  hosted on a LLaMA C++ server, which is compatible with the OpenAI
+  Chat Completions API format.
+
+  **Quick Start:**
+
+  Assuming a LLaMA C++ server is running at `http://localhost:8080`,
+  you can interact with it as follows:
+
+  ```python
+  import langfun as lf
+
+  # If model name is not specified, it will use server's default.
+  lm = lf.llms.LlamaCppRemote(url='http://localhost:8080')
+  r = lm('Who are you?')
+  print(r)
+  ```
+
+  **References:**
+
+  *   https://github.com/ggerganov/llama.cpp/tree/master/examples/server
   """
   url: Annotated[
       str,

langfun/core/llms/openai.py

@@ -1031,8 +1031,36 @@ _SUPPORTED_MODELS_BY_MODEL_ID = {m.model_id: m for m in SUPPORTED_MODELS}
 
 
 @lf.use_init_args(['model'])
-class OpenAI(openai_compatible.OpenAICompatible):
-  """OpenAI model."""
+class OpenAI(openai_compatible.OpenAIResponsesAPI):
+  """OpenAI models.
+
+  **Quick Start:**
+
+  ```python
+  import langfun as lf
+
+  # Call GPT-4o using API key from environment variable 'OPENAI_API_KEY'.
+  lm = lf.llms.Gpt4o()
+  r = lm('Who are you?')
+  print(r)
+  ```
+
+  **Setting up API key:**
+
+  The OpenAI API key can be specified in following ways:
+
+  1. At model instantiation:
+
+     ```python
+     lm = lf.llms.Gpt4o(api_key='MY_API_KEY')
+     ```
+  2. via environment variable `OPENAI_API_KEY`.
+
+  **References:**
+
+  *   https://platform.openai.com/docs/models
+  *   https://platform.openai.com/docs/api-reference
+  """
 
   model: pg.typing.Annotated[
       pg.typing.Enum(
@@ -1041,7 +1069,12 @@ class OpenAI(openai_compatible.OpenAICompatible):
       'The name of the model to use.',
   ]
 
-  api_endpoint: str = 'https://api.openai.com/v1/chat/completions'
+  # Disable message storage by default.
+  sampling_options = lf.LMSamplingOptions(
+      extras={'store': False}
+  )
+
+  api_endpoint: str = 'https://api.openai.com/v1/responses'
 
   api_key: Annotated[
       str | None,

langfun/core/llms/openai_compatible.py

@@ -23,8 +23,18 @@ import pyglove as pg
 
 
 @lf.use_init_args(['api_endpoint', 'model'])
-class OpenAICompatible(rest.REST):
-  """Base for OpenAI compatible models."""
+class OpenAIChatCompletionAPI(rest.REST):
+  """Base class for models compatible with OpenAI's Chat Completion API.
+
+  This class provides a common interface for language models that adhere to
+  the OpenAI Chat Completion API format, which is used by providers like
+  Groq, DeepSeek, and others. It standardizes request formatting and
+  response parsing for these models.
+
+  **References:**
+
+  *   https://platform.openai.com/docs/api-reference/chat
+  """
 
   model: Annotated[
       str, 'The name of the model to use.',
@@ -42,12 +52,14 @@ class OpenAICompatible(rest.REST):
     # Reference:
     # https://platform.openai.com/docs/api-reference/completions/create
     # NOTE(daiyip): options.top_k is not applicable.
-    args = dict(
-        n=options.n,
-        top_logprobs=options.top_logprobs,
-    )
+    args = {}
+
     if self.model:
       args['model'] = self.model
+    if options.n != 1:
+      args['n'] = options.n
+    if options.top_logprobs is not None:
+      args['top_logprobs'] = options.top_logprobs
     if options.logprobs:
       args['logprobs'] = options.logprobs
     if options.temperature is not None:
@@ -62,6 +74,8 @@ class OpenAICompatible(rest.REST):
       args['seed'] = options.random_seed
     if options.reasoning_effort is not None:
       args['reasoning_effort'] = options.reasoning_effort
+    if options.extras:
+      args.update(options.extras)
     return args
 
   def request(
@@ -72,27 +86,13 @@ class OpenAICompatible(rest.REST):
     """Returns the JSON input for a message."""
     request_args = self._request_args(sampling_options)
 
-    # Users could use `metadata_json_schema` to pass additional
-    # request arguments.
-    json_schema = prompt.metadata.get('json_schema')
-    if json_schema is not None:
-      if not isinstance(json_schema, dict):
-        raise ValueError(
-            f'`json_schema` must be a dict, got {json_schema!r}.'
-        )
-      if 'title' not in json_schema:
-        raise ValueError(
-            f'The root of `json_schema` must have a `title` field, '
-            f'got {json_schema!r}.'
-        )
+    # Handle structured output.
+    output_schema = self._structure_output_schema(prompt)
+    if output_schema is not None:
       request_args.update(
           response_format=dict(
               type='json_schema',
-              json_schema=dict(
-                  schema=json_schema,
-                  name=json_schema['title'],
-                  strict=True,
-              )
+              json_schema=output_schema,
           )
       )
       prompt.metadata.formatted_text = (
@@ -118,17 +118,43 @@ class OpenAICompatible(rest.REST):
       assert isinstance(system_message, lf.SystemMessage), type(system_message)
       messages.append(
           system_message.as_format(
-              'openai', chunk_preprocessor=modality_check
+              'openai_chat_completion_api', chunk_preprocessor=modality_check
           )
       )
     messages.append(
-        prompt.as_format('openai', chunk_preprocessor=modality_check)
+        prompt.as_format(
+            'openai_chat_completion_api',
+            chunk_preprocessor=modality_check
+        )
     )
     request = dict()
     request.update(request_args)
     request['messages'] = messages
     return request
 
+  def _structure_output_schema(
+      self, prompt: lf.Message
+  ) -> dict[str, Any] | None:
+    # Users could use `metadata_json_schema` to pass additional
+    # request arguments.
+    json_schema = prompt.metadata.get('json_schema')
+    if json_schema is not None:
+      if not isinstance(json_schema, dict):
+        raise ValueError(
+            f'`json_schema` must be a dict, got {json_schema!r}.'
+        )
+      if 'title' not in json_schema:
+        raise ValueError(
+            f'The root of `json_schema` must have a `title` field, '
+            f'got {json_schema!r}.'
+        )
+      return dict(
+          schema=json_schema,
+          name=json_schema['title'],
+          strict=True,
+      )
+    return None
+
   def _parse_choice(self, choice: dict[str, Any]) -> lf.LMSample:
     # Reference:
     # https://platform.openai.com/docs/api-reference/chat/object
@@ -144,7 +170,10 @@ class OpenAICompatible(rest.REST):
           for t in choice_logprobs['content']
       ]
     return lf.LMSample(
-        lf.Message.from_value(choice['message'], format='openai'),
+        lf.Message.from_value(
+            choice['message'],
+            format='openai_chat_completion_api'
+        ),
         score=0.0,
         logprobs=logprobs,
     )
@@ -169,3 +198,95 @@ class OpenAICompatible(rest.REST):
         or (status_code == 400 and b'string_above_max_length' in content)):
       return lf.ContextLimitError(f'{status_code}: {content}')
     return super()._error(status_code, content)
+
+
+class OpenAIResponsesAPI(OpenAIChatCompletionAPI):
+  """Base class for models compatible with OpenAI's Responses API.
+
+  This class provides a common interface for language models that adhere to
+  the new OpenAI Responses API format. It standardizes request formatting
+  and response parsing for these models, including handling instructions
+  (system messages) and structured outputs.
+
+  **References:**
+
+  *   https://platform.openai.com/docs/api-reference/responses
+  """
+
+  def _request_args(
+      self, options: lf.LMSamplingOptions) -> dict[str, Any]:
+    """Returns a dict as request arguments."""
+    if options.logprobs:
+      raise ValueError('logprobs is not supported on Responses API.')
+    if options.n != 1:
+      raise ValueError('n must be 1 for Responses API.')
+    return super()._request_args(options)
+
+  def request(
+      self,
+      prompt: lf.Message,
+      sampling_options: lf.LMSamplingOptions
+  ) -> dict[str, Any]:
+    """Returns the JSON input for a message."""
+    request_args = self._request_args(sampling_options)
+
+    # Handle structured output.
+    output_schema = self._structure_output_schema(prompt)
+    if output_schema is not None:
+      output_schema['type'] = 'json_schema'
+      request_args.update(text=dict(format=output_schema))
+      prompt.metadata.formatted_text = (
+          prompt.text
+          + '\n\n [RESPONSE FORMAT (not part of prompt)]\n'
+          + pg.to_json_str(request_args['text'], json_indent=2)
+      )
+
+    request = dict()
+    request.update(request_args)
+
+    # Users could use `metadata_system_message` to pass system message.
+    system_message = prompt.metadata.get('system_message')
+    if system_message:
+      assert isinstance(system_message, lf.SystemMessage), type(system_message)
+      request['instructions'] = system_message.text
+
+    # Prepare input.
+    def modality_check(chunk: str | lf.Modality) -> Any:
+      if (isinstance(chunk, lf_modalities.Mime)
+          and not self.supports_input(chunk.mime_type)):
+        raise ValueError(
+            f'Unsupported modality: {chunk!r}.'
+        )
+      return chunk
+
+    request['input'] = [
+        prompt.as_format(
+            'openai_responses_api',
+            chunk_preprocessor=modality_check
+        )
+    ]
+    return request
+
+  def _parse_output(self, output: dict[str, Any]) -> lf.LMSample:
+    for item in output:
+      if isinstance(item, dict) and item.get('type') == 'message':
+        return lf.LMSample(
+            lf.Message.from_value(item, format='openai_responses_api'),
+            score=0.0,
+        )
+    raise ValueError('No message found in output.')
+
+  def result(self, json: dict[str, Any]) -> lf.LMSamplingResult:
+    """Returns a LMSamplingResult from a JSON response."""
+    usage = json['usage']
+    return lf.LMSamplingResult(
+        samples=[self._parse_output(json['output'])],
+        usage=lf.LMSamplingUsage(
+            prompt_tokens=usage['input_tokens'],
+            completion_tokens=usage['output_tokens'],
+            total_tokens=usage['total_tokens'],
+            completion_tokens_details=usage.get(
+                'output_tokens_details', None
+            ),
+        ),
+    )