langfun 0.1.2.dev202511030805__py3-none-any.whl → 0.1.2.dev202511050805__py3-none-any.whl

langfun/core/llms/vertexai.py

@@ -44,13 +44,32 @@ except ImportError:
 
 @pg.use_init_args(['api_endpoint'])
 class VertexAI(rest.REST):
-  """Base class for VertexAI models.
+  """Base class for models served on Vertex AI.
 
-  This class handles the authentication of vertex AI models. Subclasses
-  should implement `request` and `result` methods, as well as the `api_endpoint`
-  property. Or let users to provide them as __init__ arguments.
+  This class handles authentication for Vertex AI models. Subclasses,
+  such as `VertexAIGemini`, `VertexAIAnthropic`, and `VertexAILlama`,
+  provide specific implementations for different model families hosted
+  on Vertex AI.
 
-  Please check out VertexAIGemini in `gemini.py` as an example.
+  **Quick Start:**
+
+  If you are using Langfun from a Google Cloud environment (e.g., GCE, GKE)
+  that has service account credentials, authentication is handled automatically.
+  Otherwise, you might need to set up credentials:
+
+  ```bash
+  gcloud auth application-default login
+  ```
+
+  Then you can use a Vertex AI model:
+
+  ```python
+  import langfun as lf
+
+  lm = lf.llms.VertexAIGemini25Flash(project='my-project', location='global')
+  r = lm('Who are you?')
+  print(r)
+  ```
   """
 
   model: pg.typing.Annotated[
@@ -158,7 +177,21 @@ class VertexAI(rest.REST):
 @pg.use_init_args(['model'])
 @pg.members([('api_endpoint', pg.typing.Str().freeze(''))])
 class VertexAIGemini(VertexAI, gemini.Gemini):
-  """Gemini models served by Vertex AI.."""
+  """Gemini models served on Vertex AI.
+
+  **Quick Start:**
+
+  ```python
+  import langfun as lf
+
+  # Call Gemini 1.5 Flash on Vertex AI.
+  # If project and location are not specified, they will be read from
+  # environment variables 'VERTEXAI_PROJECT' and 'VERTEXAI_LOCATION'.
+  lm = lf.llms.VertexAIGemini25Flash(project='my-project', location='global')
+  r = lm('Who are you?')
+  print(r)
+  ```
+  """
 
   # Set default location to us-central1.
   location = 'us-central1'

langfun/core/logging.py

@@ -310,7 +310,7 @@ def warning(
     console: bool = False,
     **kwargs
 ) -> LogEntry:
-  """Logs an info message to the session."""
+  """Logs a warning message to the session."""
   return log('warning', message, indent=indent, console=console, **kwargs)
 
 

langfun/core/mcp/client.py

@@ -23,33 +23,53 @@ import pyglove as pg
 
 
 class McpClient(pg.Object):
-  """Base class for MCP client.
+  """Interface for Model Context Protocol (MCP) client.
 
-  Usage:
+  An MCP client serves as a bridge to an MCP server, enabling users to interact
+  with tools hosted on the server. It provides methods for listing available
+  tools and creating sessions for tool interaction.
+
+  There are three types of MCP clients:
+
+  * **Stdio-based client**: Ideal for interacting with tools exposed as
+    command-line executables through stdin/stdout.
+    Created by `lf.mcp.McpClient.from_command`.
+  * **HTTP-based client**: Designed for tools accessible via HTTP,
+    supporting Server-Sent Events (SSE) for streaming.
+    Created by `lf.mcp.McpClient.from_url`.
+  * **In-memory client**: Useful for testing or embedding MCP servers
+    within the same process.
+    Created by `lf.mcp.McpClient.from_fastmcp`.
+
+  **Example Usage:**
 
   ```python
+  import langfun as lf
 
-  def tool_use():
-    client = lf.mcp.McpClient.from_command('<MCP_CMD>', ['<ARG1>', 'ARG2'])
-    tools = client.list_tools()
-    tool_cls = tools['<TOOL_NAME>']
+  # Example 1: Stdio-based client
+  client = lf.mcp.McpClient.from_command('<MCP_CMD>', ['<ARG1>', 'ARG2'])
+  tools = client.list_tools()
+  tool_cls = tools['<TOOL_NAME>']
 
-    # Print the python definition of the tool.
-    print(tool_cls.python_definition())
+  # Print the Python definition of the tool.
+  print(tool_cls.python_definition())
 
-    with client.session() as session:
-      return tool_cls(x=1, y=2)(session)
+  with client.session() as session:
+    result = tool_cls(x=1, y=2)(session)
+    print(result)
 
-  async def tool_use_async_version():
+  # Example 2: HTTP-based client (async)
+  async def main():
     client = lf.mcp.McpClient.from_url('http://localhost:8000/mcp')
     tools = client.list_tools()
     tool_cls = tools['<TOOL_NAME>']
 
-    # Print the python definition of the tool.
+    # Print the Python definition of the tool.
     print(tool_cls.python_definition())
 
     async with client.session() as session:
-      return await tool_cls(x=1, y=2).acall(session)
+      result = await tool_cls(x=1, y=2).acall(session)
+      print(result)
   ```
   """
 
@@ -60,7 +80,15 @@ class McpClient(pg.Object):
   def list_tools(
       self, refresh: bool = False
   ) -> dict[str, Type[mcp_tool.McpTool]]:
-    """Lists all MCP tools."""
+    """Lists all available tools on the MCP server.
+
+    Args:
+      refresh: If True, forces a refresh of the tool list from the server.
+        Otherwise, a cached list may be returned.
+
+    Returns:
+      A dictionary mapping tool names to their corresponding `McpTool` classes.
+    """
     if self._tools is None or refresh:
       with self.session() as session:
         self._tools = session.list_tools()
@@ -68,11 +96,23 @@ class McpClient(pg.Object):
 
   @abc.abstractmethod
   def session(self) -> mcp_session.McpSession:
-    """Creates a MCP session."""
+    """Creates a new session for interacting with MCP tools.
+
+    Returns:
+      An `McpSession` object.
+    """
 
   @classmethod
   def from_command(cls, command: str, args: list[str]) -> 'McpClient':
-    """Creates a MCP client from a tool."""
+    """Creates an MCP client from a command-line executable.
+
+    Args:
+      command: The command to execute.
+      args: A list of arguments to pass to the command.
+
+    Returns:
+      A `McpClient` instance that communicates via stdin/stdout.
+    """
     return _StdioMcpClient(command=command, args=args)
 
   @classmethod
@@ -81,12 +121,27 @@ class McpClient(pg.Object):
       url: str,
       headers: dict[str, str] | None = None
   ) -> 'McpClient':
-    """Creates a MCP client from a URL."""
+    """Creates an MCP client from an HTTP URL.
+
+    Args:
+      url: The URL of the MCP server.
+      headers: An optional dictionary of HTTP headers to include in requests.
+
+    Returns:
+      A `McpClient` instance that communicates via HTTP.
+    """
     return _HttpMcpClient(url=url, headers=headers or {})
 
   @classmethod
   def from_fastmcp(cls, fastmcp: fastmcp_lib.FastMCP) -> 'McpClient':
-    """Creates a MCP client from a MCP server."""
+    """Creates an MCP client from an in-memory FastMCP instance.
+
+    Args:
+      fastmcp: An instance of `fastmcp_lib.FastMCP`.
+
+    Returns:
+      A `McpClient` instance that communicates with the in-memory server.
+    """
     return _InMemoryFastMcpClient(fastmcp=fastmcp)
 
 
@@ -97,18 +152,18 @@ class _StdioMcpClient(McpClient):
   args: Annotated[list[str], 'Arguments to pass to the command.']
 
   def session(self) -> mcp_session.McpSession:
-    """Creates a MCP session."""
+    """Creates an McpSession from command."""
     return mcp_session.McpSession.from_command(self.command, self.args)
 
 
 class _HttpMcpClient(McpClient):
-  """Server-Sent Events (SSE)/Streamable HTTP-based MCP client."""
+  """HTTP-based MCP client."""
 
   url: Annotated[str, 'URL to connect to.']
   headers: Annotated[dict[str, str], 'Headers to send with the request.'] = {}
 
   def session(self) -> mcp_session.McpSession:
-    """Creates a MCP session."""
+    """Creates an McpSession from URL."""
     return mcp_session.McpSession.from_url(self.url, self.headers)
 
 
@@ -118,5 +173,5 @@ class _InMemoryFastMcpClient(McpClient):
   fastmcp: Annotated[fastmcp_lib.FastMCP, 'MCP server to connect to.']
 
   def session(self) -> mcp_session.McpSession:
-    """Creates a MCP session."""
+    """Creates an McpSession from an in-memory FastMCP instance."""
     return mcp_session.McpSession.from_fastmcp(self.fastmcp)

langfun/core/mcp/session.py

@@ -26,10 +26,37 @@ from mcp.shared import memory
 
 
 class McpSession:
-  """Langfun's MCP session.
+  """Represents a session for interacting with an MCP server.
 
-  Compared to the standard mcp.ClientSession, Langfun's MCP session could be
-  used both synchronously and asynchronously.
+  `McpSession` provides the context for making calls to tools hosted on an
+  MCP server. It wraps the standard `mcp.ClientSession` to offer both
+  synchronous and asynchronous usage patterns.
+
+  Sessions are created using `lf.mcp.McpClient.session()` and should be used
+  as context managers (either sync or async) to ensure proper initialization
+  and teardown of the connection to the server.
+
+  **Example Sync Usage:**
+
+  ```python
+  import langfun as lf
+
+  client = lf.mcp.McpClient.from_command(...)
+  with client.session() as session:
+    tools = session.list_tools()
+    result = tools['my_tool'](x=1)(session)
+  ```
+
+  **Example Async Usage:**
+
+  ```python
+  import langfun as lf
+
+  client = lf.mcp.McpClient.from_url(...)
+  async with client.session() as session:
+    tools = await session.alist_tools()
+    result = await tools['my_tool'](x=1).acall(session)
+  ```
   """
 
   def __init__(self, stream) -> None:
@@ -74,11 +101,19 @@ class McpSession:
     self._session = None
 
   def list_tools(self) -> dict[str, Type[mcp_tool.McpTool]]:
-    """Lists all MCP tools synchronously."""
+    """Lists all available tools on the MCP server synchronously.
+
+    Returns:
+      A dictionary mapping tool names to their corresponding `McpTool` classes.
+    """
     return async_support.invoke_sync(self.alist_tools)
 
   async def alist_tools(self) -> dict[str, Type[mcp_tool.McpTool]]:
-    """Lists all MCP tools asynchronously."""
+    """Lists all available tools on the MCP server asynchronously.
+
+    Returns:
+      A dictionary mapping tool names to their corresponding `McpTool` classes.
+    """
     assert self._session is not None, 'MCP session is not entered.'
     return {
         t.name: mcp_tool.McpTool.make_class(t)
@@ -91,7 +126,16 @@ class McpSession:
       *,
       returns_message: bool = False
   ) -> Any:
-    """Calls a MCP tool synchronously."""
+    """Calls an MCP tool synchronously.
+
+    Args:
+      tool: The `McpTool` instance to call.
+      returns_message: If True, the tool call will return an `mcp.Message`
+        object; otherwise, it returns the tool's direct result.
+
+    Returns:
+      The result of the tool call.
+    """
     return tool(self, returns_message=returns_message)
 
   async def acall_tool(
@@ -100,7 +144,16 @@ class McpSession:
       *,
       returns_message: bool = False
   ) -> Any:
-    """Calls a MCP tool asynchronously."""
+    """Calls an MCP tool asynchronously.
+
+    Args:
+      tool: The `McpTool` instance to call.
+      returns_message: If True, the tool call will return an `mcp.Message`
+        object; otherwise, it returns the tool's direct result.
+
+    Returns:
+      The result of the tool call.
+    """
     return await tool.acall(self, returns_message=returns_message)
 
   @classmethod
@@ -109,7 +162,15 @@ class McpSession:
       command: str,
       args: list[str] | None = None
   ) -> 'McpSession':
-    """Creates a MCP session from a command."""
+    """Creates an MCP session from a command-line executable.
+
+    Args:
+      command: The command to execute.
+      args: An optional list of arguments to pass to the command.
+
+    Returns:
+      An `McpSession` instance.
+    """
     return cls(
         mcp.stdio_client(
             mcp.StdioServerParameters(command=command, args=args or [])
@@ -122,7 +183,18 @@ class McpSession:
       url: str,
       headers: dict[str, str] | None = None
   ) -> 'McpSession':
-    """Creates a MCP session from a URL."""
+    """Creates an MCP session from an HTTP URL.
+
+    The transport protocol (e.g., 'mcp' or 'sse') is inferred from the
+    last part of the URL path.
+
+    Args:
+      url: The URL of the MCP server.
+      headers: An optional dictionary of HTTP headers to include in requests.
+
+    Returns:
+      An `McpSession` instance.
+    """
     transport = url.removesuffix('/').split('/')[-1].lower()
     if transport == 'mcp':
       return cls(streamable_http.streamablehttp_client(url, headers or {}))
@@ -136,12 +208,20 @@ class McpSession:
       cls,
       fastmcp: fastmcp_lib.FastMCP
   ):
+    """Creates an MCP session from an in-memory FastMCP instance.
+
+    Args:
+      fastmcp: An instance of `fastmcp_lib.FastMCP`.
+
+    Returns:
+      An `McpSession` instance.
+    """
     return cls(_client_streams_from_fastmcp(fastmcp))
 
 
 @contextlib.asynccontextmanager
 async def _client_streams_from_fastmcp(fastmcp: fastmcp_lib.FastMCP):
-  """Creates client streams from a MCP server."""
+  """Creates client streams from an in-memory FastMCP instance."""
   server = fastmcp._mcp_server  # pylint: disable=protected-access
   async with memory.create_client_server_memory_streams(
       ) as (client_streams, server_streams):

langfun/core/mcp/tool.py

@@ -31,7 +31,31 @@ class _McpToolMeta(pg.symbolic.ObjectMeta):
 
 
 class McpTool(pg.Object, metaclass=_McpToolMeta):
-  """Base class for MCP tools."""
+  """Represents a tool available on an MCP server.
+
+  `McpTool` is the base class for all tool proxies generated from an MCP
+  server's tool definitions. Users do not typically subclass `McpTool` directly.
+  Instead, tool classes are obtained by calling `lf.mcp.McpClient.list_tools()`
+  or `lf.mcp.McpSession.list_tools()`.
+
+  Once a tool class is obtained, it can be instantiated with input parameters
+  and called via an `McpSession` to execute the tool on the server.
+
+  **Example Usage:**
+
+  ```python
+  import langfun as lf
+
+  client = lf.mcp.McpClient.from_command(...)
+  with client.session() as session:
+    # List tools and get the 'math' tool class.
+    math_tool_cls = session.list_tools()['math']
+
+    # Instantiate the tool with parameters and call it.
+    result = math_tool_cls(x=1, y=2, op='+')(session)
+    print(result)
+  ```
+  """
 
   TOOL_NAME: Annotated[
       ClassVar[str],
@@ -40,7 +64,17 @@ class McpTool(pg.Object, metaclass=_McpToolMeta):
 
   @classmethod
   def python_definition(cls, markdown: bool = True) -> str:
-    """Returns the Python definition of the tool."""
+    """Returns the Python definition of this tool's input schema.
+
+    This is useful for generating prompts that instruct a language model
+    on how to use the tool.
+
+    Args:
+      markdown: If True, formats the output as a Markdown code block.
+
+    Returns:
+      A string containing the Python definition of the tool's input schema.
+    """
     return lf_schema.Schema.from_value(cls).schema_str(
         protocol='python', markdown=markdown
     )
@@ -49,16 +83,17 @@ class McpTool(pg.Object, metaclass=_McpToolMeta):
   def result_to_message(
       cls, result: mcp.types.CallToolResult
   ) -> lf_message.ToolMessage:
-    """Converts a tool call result to a message.
+    """Converts an `mcp.types.CallToolResult` to an `lf.ToolMessage`.
 
-    This method allows users to convert an existing mcp.CallToolResult to a
-    Langfun ToolMessage.
+    This method translates results from the MCP protocol, including text,
+    image, and audio content, into a Langfun `ToolMessage`, making it easy
+    to integrate tool results into Langfun workflows.
 
     Args:
-      result: The MCP tool call result.
+      result: The `mcp.types.CallToolResult` object to convert.
 
     Returns:
-      A ToolMessage object.
+      An `lf.ToolMessage` instance representing the tool result.
     """
     chunks = []
     for item in result.content:
@@ -81,16 +116,18 @@ class McpTool(pg.Object, metaclass=_McpToolMeta):
       session,
       *,
       returns_message: bool = False) -> Any:
-    """Calls a MCP tool synchronously.
+    """Calls the MCP tool synchronously within a given session.
 
     Args:
-      session: A MCP session.
-      returns_message: If True, always returns a ToolMessage. Otherwise,
-        return structured content (result) if available, otherwise return
-        ToolMessage if there is multi-modal content, otherwise return text.
+      session: An `McpSession` object.
+      returns_message: If True, the raw `lf.ToolMessage` is returned.
+        If False(default), the method attempts to return a more specific result:
+          - The `result` field of the `ToolMessage` if it's populated.
+          - The `ToolMessage` itself if it contains multi-modal content.
+          - The `text` field of the `ToolMessage` otherwise.
 
     Returns:
-      The call result, or the result from structured content, or content.
+      The result of the tool call, processed according to `returns_message`.
     """
     return async_support.invoke_sync(
         self.acall,
@@ -104,16 +141,18 @@ class McpTool(pg.Object, metaclass=_McpToolMeta):
       *,
       returns_message: bool = False
   ) -> Any:
-    """Calls a MCP tool asynchronously.
+    """Calls the MCP tool asynchronously within a given session.
 
     Args:
-      session: McpSession or mcp.ClientSession.
-      returns_message: If True, always returns a ToolMessage. Otherwise,
-        return structured content (result) if available, otherwise return
-        ToolMessage if there is multi-modal content, otherwise return text.
+      session: An `McpSession` object or an `mcp.ClientSession`.
+      returns_message: If True, the raw `lf.ToolMessage` is returned.
+        If False(default), the method attempts to return a more specific result:
+          - The `result` field of the `ToolMessage` if it's populated.
+          - The `ToolMessage` itself if it contains multi-modal content.
+          - The `text` field of the `ToolMessage` otherwise.
 
     Returns:
-      The call result, or the result from structured content, or content.
+      The result of the tool call, processed according to `returns_message`.
     """
     if not isinstance(session, mcp.ClientSession):
       session = getattr(session, '_session', None)
@@ -131,7 +170,12 @@ class McpTool(pg.Object, metaclass=_McpToolMeta):
     return message.text
 
   def input_parameters(self) -> dict[str, Any]:
-    """Returns the input parameters of the tool."""
+    """Returns the input parameters for the tool call.
+
+    Returns:
+      A dictionary containing the input parameters, formatted for an
+      MCP `call_tool` request.
+    """
     # Optional fields are represented as fields with default values. Therefore,
     # we need to remove the default values from the JSON representation of the
     # tool.
@@ -147,7 +191,15 @@ class McpTool(pg.Object, metaclass=_McpToolMeta):
 
   @classmethod
   def make_class(cls, tool_definition: mcp.Tool) -> type['McpTool']:
-    """Makes a MCP tool class from tool definition."""
+    """Creates an `McpTool` subclass from an MCP tool definition.
+
+    Args:
+      tool_definition: An `mcp.Tool` object containing the tool's metadata.
+
+    Returns:
+      A dynamically generated class that inherits from `McpTool` and
+      represents the defined tool.
+    """
 
     class _McpTool(cls):
       auto_schema = False
@@ -170,11 +222,19 @@ class _McpToolInputMeta(pg.symbolic.ObjectMeta):
 
 
 class McpToolInput(pg.Object, metaclass=_McpToolInputMeta):
-  """Base class for MCP tool inputs."""
+  """Base class for generated MCP tool input schemas."""
 
   @classmethod
   def make_class(cls, name: str, schema: pg.Schema):
-    """Converts a schema to an input class."""
+    """Creates an `McpToolInput` subclass from a schema.
+
+    Args:
+      name: The name of the input class to generate.
+      schema: A `pg.Schema` object defining the input fields.
+
+    Returns:
+      A dynamically generated class that inherits from `McpToolInput`.
+    """
 
     class _McpToolInput(cls):
       pass

langfun/core/memory.py

@@ -43,6 +43,7 @@ class Memory(NaturalLanguageFormattable, Component):
       value: Any,
       **kwargs
   ) -> None:
+    """Remembers a value."""
     self._remember(value, **kwargs)
 
   def reset(self, **kwargs) -> None: