genesis-flow 1.0.0__py3-none-any.whl

mlflow/langchain/__init__.py

@@ -0,0 +1,24 @@
+from mlflow.langchain.autolog import autolog
+from mlflow.langchain.constants import FLAVOR_NAME
+from mlflow.version import IS_TRACING_SDK_ONLY
+
+__all__ = ["autolog", "FLAVOR_NAME"]
+
+# Import model logging APIs only if mlflow skinny or full package is installed,
+# i.e., skip if only mlflow-tracing package is installed.
+if not IS_TRACING_SDK_ONLY:
+    from mlflow.langchain.model import (
+        _LangChainModelWrapper,
+        _load_pyfunc,
+        load_model,
+        log_model,
+        save_model,
+    )
+
+    __all__ += [
+        "_LangChainModelWrapper",
+        "_load_pyfunc",
+        "load_model",
+        "log_model",
+        "save_model",
+    ]

mlflow/langchain/api_request_parallel_processor.py

@@ -0,0 +1,330 @@
+# Based ons: https://github.com/openai/openai-cookbook/blob/6df6ceff470eeba26a56de131254e775292eac22/examples/api_request_parallel_processor.py
+# Several changes were made to make it work with MLflow.
+# Currently, only chat completion is supported.
+
+"""
+API REQUEST PARALLEL PROCESSOR
+
+Using the LangChain API to process lots of text quickly takes some care.
+If you trickle in a million API requests one by one, they'll take days to complete.
+This script parallelizes requests using LangChain API.
+
+Features:
+- Streams requests from file, to avoid running out of memory for giant jobs
+- Makes requests concurrently, to maximize throughput
+- Logs errors, to diagnose problems with requests
+"""
+
+from __future__ import annotations
+
+import logging
+import queue
+import threading
+import time
+import traceback
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+
+import langchain.chains
+from langchain.callbacks.base import BaseCallbackHandler
+
+import mlflow
+from mlflow.exceptions import MlflowException
+from mlflow.langchain.utils.chat import (
+    transform_request_json_for_chat_if_necessary,
+    try_transform_response_iter_to_chat_format,
+    try_transform_response_to_chat_format,
+)
+from mlflow.langchain.utils.serialization import convert_to_serializable
+from mlflow.pyfunc.context import Context, get_prediction_context
+from mlflow.tracing.utils import maybe_set_prediction_context
+
+_logger = logging.getLogger(__name__)
+
+
+@dataclass
+class StatusTracker:
+    """
+    Stores metadata about the script's progress. Only one instance is created.
+    """
+
+    num_tasks_started: int = 0
+    num_tasks_in_progress: int = 0  # script ends when this reaches 0
+    num_tasks_succeeded: int = 0
+    num_tasks_failed: int = 0
+    num_api_errors: int = 0  # excluding rate limit errors, counted above
+    lock: threading.Lock = threading.Lock()
+
+    def start_task(self):
+        with self.lock:
+            self.num_tasks_started += 1
+            self.num_tasks_in_progress += 1
+
+    def complete_task(self, *, success: bool):
+        with self.lock:
+            self.num_tasks_in_progress -= 1
+            if success:
+                self.num_tasks_succeeded += 1
+            else:
+                self.num_tasks_failed += 1
+
+    def increment_num_api_errors(self):
+        with self.lock:
+            self.num_api_errors += 1
+
+
+@dataclass
+class APIRequest:
+    """
+    Stores an API request's inputs, outputs, and other metadata. Contains a method to make an API
+    call.
+
+    Args:
+        index: The request's index in the tasks list
+        lc_model: The LangChain model to call
+        request_json: The request's input data
+        results: The list to append the request's output data to, it's a list of tuples
+            (index, response)
+        errors: A dictionary to store any errors that occur
+        convert_chat_responses: Whether to convert the model's responses to chat format
+        did_perform_chat_conversion: Whether the input data was converted to chat format
+            based on the model's type and input data.
+        stream: Whether the request is a stream request
+        prediction_context: The prediction context to use for the request
+    """
+
+    index: int
+    lc_model: langchain.chains.base.Chain
+    request_json: dict[str, Any]
+    results: list[tuple[int, str]]
+    errors: dict[int, str]
+    convert_chat_responses: bool
+    did_perform_chat_conversion: bool
+    stream: bool
+    params: dict[str, Any]
+    prediction_context: Optional[Context] = None
+
+    def _predict_single_input(self, single_input, callback_handlers, **kwargs):
+        config = kwargs.pop("config", {})
+        config["callbacks"] = config.get("callbacks", []) + (callback_handlers or [])
+        if self.stream:
+            return self.lc_model.stream(single_input, config=config, **kwargs)
+        if hasattr(self.lc_model, "invoke"):
+            return self.lc_model.invoke(single_input, config=config, **kwargs)
+        else:
+            # for backwards compatibility, __call__ is deprecated and will be removed in 0.3.0
+            # kwargs shouldn't have config field if invoking with __call__
+            return self.lc_model(single_input, callbacks=callback_handlers, **kwargs)
+
+    def _try_convert_response(self, response):
+        if self.stream:
+            return try_transform_response_iter_to_chat_format(response)
+        else:
+            return try_transform_response_to_chat_format(response)
+
+    def single_call_api(self, callback_handlers: Optional[list[BaseCallbackHandler]]):
+        from langchain.schema import BaseRetriever
+
+        from mlflow.langchain.utils.logging import langgraph_types, lc_runnables_types
+
+        if isinstance(self.lc_model, BaseRetriever):
+            # Retrievers are invoked differently than Chains
+            response = self.lc_model.get_relevant_documents(
+                **self.request_json, callbacks=callback_handlers, **self.params
+            )
+        elif isinstance(self.lc_model, lc_runnables_types() + langgraph_types()):
+            if isinstance(self.request_json, dict):
+                # This is a temporary fix for the case when spark_udf converts
+                # input into pandas dataframe with column name, while the model
+                # does not accept dictionaries as input, it leads to errors like
+                # Expected Scalar value for String field 'query_text'
+                try:
+                    response = self._predict_single_input(
+                        self.request_json, callback_handlers, **self.params
+                    )
+                except TypeError as e:
+                    _logger.debug(
+                        f"Failed to invoke {self.lc_model.__class__.__name__} "
+                        f"with {self.request_json}. Error: {e!r}. Trying to "
+                        "invoke with the first value of the dictionary."
+                    )
+                    self.request_json = next(iter(self.request_json.values()))
+                    (
+                        prepared_request_json,
+                        did_perform_chat_conversion,
+                    ) = transform_request_json_for_chat_if_necessary(
+                        self.request_json, self.lc_model
+                    )
+                    self.did_perform_chat_conversion = did_perform_chat_conversion
+
+                    response = self._predict_single_input(
+                        prepared_request_json, callback_handlers, **self.params
+                    )
+            else:
+                response = self._predict_single_input(
+                    self.request_json, callback_handlers, **self.params
+                )
+
+            if self.did_perform_chat_conversion or self.convert_chat_responses:
+                response = self._try_convert_response(response)
+        else:
+            # return_only_outputs is invalid for stream call
+            if isinstance(self.lc_model, langchain.chains.base.Chain) and not self.stream:
+                kwargs = {"return_only_outputs": True}
+            else:
+                kwargs = {}
+            kwargs.update(**self.params)
+            response = self._predict_single_input(self.request_json, callback_handlers, **kwargs)
+
+            if self.did_perform_chat_conversion or self.convert_chat_responses:
+                response = self._try_convert_response(response)
+            elif isinstance(response, dict) and len(response) == 1:
+                # to maintain existing code, single output chains will still return
+                # only the result
+                response = response.popitem()[1]
+
+        return convert_to_serializable(response)
+
+    def call_api(
+        self, status_tracker: StatusTracker, callback_handlers: Optional[list[BaseCallbackHandler]]
+    ):
+        """
+        Calls the LangChain API and stores results.
+        """
+        _logger.debug(f"Request #{self.index} started with payload: {self.request_json}")
+
+        try:
+            with maybe_set_prediction_context(self.prediction_context):
+                response = self.single_call_api(callback_handlers)
+            _logger.debug(f"Request #{self.index} succeeded with response: {response}")
+            self.results.append((self.index, response))
+            status_tracker.complete_task(success=True)
+        except Exception as e:
+            self.errors[self.index] = (
+                f"error: {e!r} {traceback.format_exc()}\n request payload: {self.request_json}"
+            )
+            status_tracker.increment_num_api_errors()
+            status_tracker.complete_task(success=False)
+
+
+def process_api_requests(
+    lc_model,
+    requests: Optional[list[Union[Any, dict[str, Any]]]] = None,
+    max_workers: int = 10,
+    callback_handlers: Optional[list[BaseCallbackHandler]] = None,
+    convert_chat_responses: bool = False,
+    params: Optional[dict[str, Any]] = None,
+    context: Optional[Context] = None,
+):
+    """
+    Processes API requests in parallel.
+    """
+
+    # initialize trackers
+    retry_queue = queue.Queue()
+    status_tracker = StatusTracker()  # single instance to track a collection of variables
+    next_request = None  # variable to hold the next request to call
+    context = context or get_prediction_context()
+
+    results = []
+    errors = {}
+
+    # Note: we should call `transform_request_json_for_chat_if_necessary`
+    # for the whole batch data, because the conversion should obey the rule
+    # that if any record in the batch can't be converted, then all the record
+    # in this batch can't be converted.
+    (
+        converted_chat_requests,
+        did_perform_chat_conversion,
+    ) = transform_request_json_for_chat_if_necessary(requests, lc_model)
+
+    requests_iter = enumerate(converted_chat_requests)
+    with ThreadPoolExecutor(
+        max_workers=max_workers, thread_name_prefix="MlflowLangChainApi"
+    ) as executor:
+        while True:
+            # get next request (if one is not already waiting for capacity)
+            if not retry_queue.empty():
+                next_request = retry_queue.get_nowait()
+                _logger.warning(f"Retrying request {next_request.index}: {next_request}")
+            elif req := next(requests_iter, None):
+                # get new request
+                index, converted_chat_request_json = req
+                next_request = APIRequest(
+                    index=index,
+                    lc_model=lc_model,
+                    request_json=converted_chat_request_json,
+                    results=results,
+                    errors=errors,
+                    convert_chat_responses=convert_chat_responses,
+                    did_perform_chat_conversion=did_perform_chat_conversion,
+                    stream=False,
+                    prediction_context=context,
+                    params=params,
+                )
+                status_tracker.start_task()
+            else:
+                next_request = None
+
+            # if enough capacity available, call API
+            if next_request:
+                # call API
+                executor.submit(
+                    next_request.call_api,
+                    status_tracker=status_tracker,
+                    callback_handlers=callback_handlers,
+                )
+
+            # if all tasks are finished, break
+            # check next_request to avoid terminating the process
+            # before extra requests need to be processed
+            if status_tracker.num_tasks_in_progress == 0 and next_request is None:
+                break
+
+            time.sleep(0.001)  # avoid busy waiting
+
+        # after finishing, log final status
+        if status_tracker.num_tasks_failed > 0:
+            raise mlflow.MlflowException(
+                f"{status_tracker.num_tasks_failed} tasks failed. Errors: {errors}"
+            )
+
+        return [res for _, res in sorted(results)]
+
+
+def process_stream_request(
+    lc_model,
+    request_json: Union[Any, dict[str, Any]],
+    callback_handlers: Optional[list[BaseCallbackHandler]] = None,
+    convert_chat_responses: bool = False,
+    params: Optional[dict[str, Any]] = None,
+):
+    """
+    Process single stream request.
+    """
+    if not hasattr(lc_model, "stream"):
+        raise MlflowException(
+            f"Model {lc_model.__class__.__name__} does not support streaming prediction output. "
+            "No `stream` method found."
+        )
+
+    (
+        converted_chat_requests,
+        did_perform_chat_conversion,
+    ) = transform_request_json_for_chat_if_necessary(request_json, lc_model)
+
+    api_request = APIRequest(
+        index=0,
+        lc_model=lc_model,
+        request_json=converted_chat_requests,
+        results=None,
+        errors=None,
+        convert_chat_responses=convert_chat_responses,
+        did_perform_chat_conversion=did_perform_chat_conversion,
+        stream=True,
+        prediction_context=get_prediction_context(),
+        params=params,
+    )
+    with maybe_set_prediction_context(api_request.prediction_context):
+        return api_request.single_call_api(callback_handlers)

mlflow/langchain/autolog.py

@@ -0,0 +1,147 @@
+import logging
+
+from mlflow.langchain.constant import FLAVOR_NAME
+from mlflow.utils.annotations import experimental
+from mlflow.utils.autologging_utils import autologging_integration
+from mlflow.utils.autologging_utils.config import AutoLoggingConfig
+from mlflow.utils.autologging_utils.safety import safe_patch
+
+logger = logging.getLogger(__name__)
+
+
+@experimental(version="2.10.0")
+@autologging_integration(FLAVOR_NAME)
+def autolog(
+    disable=False,
+    exclusive=False,
+    disable_for_unsupported_versions=False,
+    silent=False,
+    log_traces=True,
+):
+    """
+    Enables (or disables) and configures autologging from Langchain to MLflow.
+
+    Args:
+        disable: If ``True``, disables the Langchain autologging integration. If ``False``,
+            enables the Langchain autologging integration.
+        exclusive: If ``True``, autologged content is not logged to user-created fluent runs.
+            If ``False``, autologged content is logged to the active fluent run,
+            which may be user-created.
+        disable_for_unsupported_versions: If ``True``, disable autologging for versions of
+            langchain that have not been tested against this version of the MLflow
+            client or are incompatible.
+        silent: If ``True``, suppress all event logs and warnings from MLflow during Langchain
+            autologging. If ``False``, show all events and warnings during Langchain
+            autologging.
+        log_traces: If ``True``, traces are logged for Langchain models by using
+            MlflowLangchainTracer as a callback during inference. If ``False``, no traces are
+            collected during inference. Default to ``True``.
+    """
+    try:
+        from langchain_core.callbacks import BaseCallbackManager
+
+        safe_patch(
+            FLAVOR_NAME,
+            BaseCallbackManager,
+            "__init__",
+            _patched_callback_manager_init,
+        )
+    except Exception as e:
+        logger.warning(f"Failed to enable tracing for LangChain. Error: {e}")
+
+    # Special handlings for edge cases.
+    try:
+        from langchain_core.callbacks import BaseCallbackManager
+        from langchain_core.runnables import RunnableSequence
+
+        safe_patch(
+            FLAVOR_NAME,
+            RunnableSequence,
+            "batch",
+            _patched_runnable_sequence_batch,
+        )
+
+        safe_patch(
+            FLAVOR_NAME,
+            BaseCallbackManager,
+            "merge",
+            _patched_callback_manager_merge,
+        )
+    except Exception:
+        logger.debug("Failed to patch RunnableSequence or BaseCallbackManager.", exc_info=True)
+
+
+def _patched_callback_manager_init(original, self, *args, **kwargs):
+    from mlflow.langchain.langchain_tracer import MlflowLangchainTracer
+
+    original(self, *args, **kwargs)
+
+    if not AutoLoggingConfig.init(FLAVOR_NAME).log_traces:
+        return
+
+    for handler in self.inheritable_handlers:
+        if isinstance(handler, MlflowLangchainTracer):
+            return
+
+    _handler = MlflowLangchainTracer()
+    self.add_handler(_handler, inherit=True)
+
+
+def _patched_callback_manager_merge(original, self, *args, **kwargs):
+    """
+    Patch BaseCallbackManager.merge to avoid a duplicated callback issue.
+
+    In the above patched __init__, we check `inheritable_handlers` to see if the MLflow tracer
+    is already propagated. This works when the `inheritable_handlers` is specified as constructor
+    arguments. However, in the `merge` method, LangChain does not use constructor but set
+    callbacks via the setter method. This causes duplicated callbacks injection.
+    https://github.com/langchain-ai/langchain/blob/d9a069c414a321e7a3f3638a32ecf8a37ec2d188/libs/core/langchain_core/callbacks/base.py#L962-L982
+    """
+    from mlflow.langchain.langchain_tracer import MlflowLangchainTracer
+
+    # Get the MLflow callback inherited from parent
+    inherited = self.inheritable_handlers + args[0].inheritable_handlers
+    inherited_mlflow_cb = next(
+        (cb for cb in inherited if isinstance(cb, MlflowLangchainTracer)), None
+    )
+
+    if not inherited_mlflow_cb:
+        return original(self, *args, **kwargs)
+
+    merged = original(self, *args, **kwargs)
+    # If a new MLflow callback is generated inside __init__, remove it
+    duplicate_mlflow_cbs = [
+        cb
+        for cb in merged.inheritable_handlers
+        if isinstance(cb, MlflowLangchainTracer) and cb != inherited_mlflow_cb
+    ]
+    for cb in duplicate_mlflow_cbs:
+        merged.remove_handler(cb)
+
+    return merged
+
+
+def _patched_runnable_sequence_batch(original, self, *args, **kwargs):
+    """
+    Patch to terminate span context attachment during batch execution.
+
+    RunnableSequence's batch() methods are implemented in a peculiar way
+    that iterates on steps->items sequentially within the same thread. For example, if a
+    sequence has 2 steps and the batch size is 3, the execution flow will be:
+      - Step 1 for item 1
+      - Step 1 for item 2
+      - Step 1 for item 3
+      - Step 2 for item 1
+      - Step 2 for item 2
+      - Step 2 for item 3
+    Due to this behavior, we cannot attach the span to the context for this particular
+    API, otherwise spans for different inputs will be mixed up.
+    """
+    from mlflow.langchain.langchain_tracer import _should_attach_span_to_context
+
+    original_state = _should_attach_span_to_context.get()
+    _should_attach_span_to_context.set(False)
+    try:
+        return original(self, *args, **kwargs)
+    finally:
+        _should_attach_span_to_context.set(original_state)