truefoundry 0.2.10__py3-none-any.whl → 0.3.0__py3-none-any.whl

truefoundry/langchain/truefoundry_chat.py

@@ -0,0 +1,130 @@
+from typing import Any, Dict, List, Optional
+
+from langchain.chat_models.base import SimpleChatModel
+from langchain.pydantic_v1 import Extra, Field, root_validator
+from langchain.schema.messages import (
+    AIMessage,
+    BaseMessage,
+    ChatMessage,
+    HumanMessage,
+    SystemMessage,
+)
+
+from truefoundry.langchain.utils import (
+    requests_retry_session,
+    validate_tfy_environment,
+)
+from truefoundry.logger import logger
+
+
+class TrueFoundryChat(SimpleChatModel):
+    """`TrueFoundry LLM Gateway` chat models API.
+
+    To use, you must have the environment variable ``TFY_API_KEY`` set with your API key and ``TFY_HOST`` set with your host or pass it as a named parameter to the constructor.
+    """
+
+    model: str = Field(description="The model to use for chat.")
+    """The model to use for chat."""
+    tfy_llm_gateway_url: Optional[str] = Field(default=None)
+    """TrueFoundry LLM Gateway endpoint URL. Automatically inferred from env var `TFY_LLM_GATEWAY_URL` if not provided."""
+    tfy_api_key: Optional[str] = Field(default=None)
+    """TrueFoundry API Key. Automatically inferred from env var `TFY_API_KEY` if not provided."""
+    model_parameters: Optional[dict] = Field(default_factory=dict)
+    """Model parameters"""
+    request_timeout: int = Field(default=30)
+    """The timeout for the request in seconds."""
+    max_retries: int = Field(default=5)
+    """The number of retries for HTTP requests."""
+    retry_backoff_factor: float = Field(default=0.3)
+    """The backoff factor for exponential backoff during retries."""
+    system_prompt: str = Field(default="You are a AI assistant")
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+        allow_population_by_field_name = True
+
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        values = validate_tfy_environment(values)
+        if not values["tfy_api_key"]:
+            raise ValueError(
+                "Did not find `tfy_api_key`, please add an environment variable"
+                " `TFY_API_KEY` which contains it, or pass"
+                "  `tfy_api_key` as a named parameter."
+            )
+        if not values["tfy_llm_gateway_url"]:
+            raise ValueError(
+                "Did not find `tfy_llm_gateway_url`, please add an environment variable"
+                " `TFY_LLM_GATEWAY_URL` which contains it, or pass"
+                "  `tfy_llm_gateway_url` as a named parameter."
+            )
+        return values
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of chat model."""
+        return "truefoundry-chat"
+
+    def _call(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> str:
+        if len(messages) == 0:
+            raise ValueError("No messages provided to chat.")
+
+        if not isinstance(messages[0], SystemMessage):
+            messages.insert(0, SystemMessage(content=self.system_prompt))
+
+        message_dicts = [
+            TrueFoundryChat._convert_message_to_dict(message) for message in messages
+        ]
+
+        payload = {**self.model_parameters} if self.model_parameters else {}
+
+        if stop:
+            payload["stop_sequences"] = stop
+
+        payload["messages"] = message_dicts
+        payload["model"] = self.model
+
+        session = requests_retry_session(
+            retries=self.max_retries, backoff_factor=self.retry_backoff_factor
+        )
+
+        url = f"{self.tfy_llm_gateway_url}/openai/chat/completions"
+        logger.debug(f"Chat using - model: {self.model} at endpoint: {url}")
+        response = session.post(
+            url=url,
+            json=payload,
+            headers={
+                "Authorization": f"Bearer {self.tfy_api_key}",
+            },
+            timeout=self.request_timeout,
+        )
+        response.raise_for_status()
+        output = response.json()
+        return output["choices"][0]["message"]["content"]
+
+    @staticmethod
+    def _convert_message_to_dict(message: BaseMessage) -> dict:
+        if isinstance(message, ChatMessage):
+            message_dict = {"role": message.role, "content": message.content}
+        elif isinstance(message, HumanMessage):
+            message_dict = {"role": "user", "content": message.content}
+        elif isinstance(message, AIMessage):
+            message_dict = {"role": "assistant", "content": message.content}
+        elif isinstance(message, SystemMessage):
+            message_dict = {"role": "system", "content": message.content}
+        else:
+            raise ValueError(f"Got unknown message type: {message}")
+        if message.additional_kwargs:
+            logger.debug(
+                "Additional message arguments are unsupported by TrueFoundry LLM Gateway "
+                " and will be ignored: %s",
+                message.additional_kwargs,
+            )
+        return message_dict

truefoundry/langchain/truefoundry_embeddings.py

@@ -0,0 +1,171 @@
+import concurrent.futures
+import math
+from typing import Dict, List, Optional
+
+import tqdm
+from langchain.embeddings.base import Embeddings
+from langchain.pydantic_v1 import BaseModel, Extra, Field, root_validator
+
+from truefoundry.langchain.utils import (
+    requests_retry_session,
+    validate_tfy_environment,
+)
+from truefoundry.logger import logger
+
+EMBEDDER_BATCH_SIZE = 32
+PARALLEL_WORKERS = 4
+
+
+class TrueFoundryEmbeddings(BaseModel, Embeddings):
+    """`TrueFoundry LLM Gateway` embedding models API.
+
+    To use, you must have the environment variable ``TFY_API_KEY`` set with your API key and ``TFY_HOST`` set with your host or pass it
+    as a named parameter to the constructor.
+    """
+
+    model: str = Field(description="The model to use for embedding.")
+    """The model to use for embedding."""
+    tfy_llm_gateway_url: Optional[str] = Field(default=None)
+    """TrueFoundry LLM Gateway endpoint URL. Automatically inferred from env var `TFY_LLM_GATEWAY_URL` if not provided."""
+    tfy_api_key: Optional[str] = Field(default=None)
+    """TrueFoundry API Key. Automatically inferred from env var `TFY_API_KEY` if not provided."""
+    model_parameters: Optional[dict] = Field(default_factory=dict)
+    """Model parameters"""
+    request_timeout: int = Field(default=30)
+    """The timeout for the request in seconds."""
+    max_retries: int = Field(default=5)
+    """The number of retries for HTTP requests."""
+    retry_backoff_factor: float = Field(default=0.3)
+    """The backoff factor for exponential backoff during retries."""
+    batch_size: int = Field(default=EMBEDDER_BATCH_SIZE)
+    """The batch size to use for embedding."""
+    parallel_workers: int = Field(default=PARALLEL_WORKERS)
+    """The number of parallel workers to use for embedding."""
+
+    __private_attributes__ = {"_executor"}
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+        allow_population_by_field_name = True
+
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        values = validate_tfy_environment(values)
+        if not values["tfy_api_key"]:
+            raise ValueError(
+                "Did not find `tfy_api_key`, please add an environment variable"
+                " `TFY_API_KEY` which contains it, or pass"
+                "  `tfy_api_key` as a named parameter."
+            )
+        if not values["tfy_llm_gateway_url"]:
+            raise ValueError(
+                "Did not find `tfy_llm_gateway_url`, please add an environment variable"
+                " `TFY_LLM_GATEWAY_URL` which contains it, or pass"
+                "  `tfy_llm_gateway_url` as a named parameter."
+            )
+        return values
+
+    def _init_private_attributes(self):
+        self._executor = concurrent.futures.ThreadPoolExecutor(
+            max_workers=self.parallel_workers
+        )
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of embedding model."""
+        return "truefoundry-embeddings"
+
+    def __del__(self):
+        """
+        Destructor method to clean up the executor when the object is deleted.
+
+        Returns:
+            None
+        """
+        self._executor.shutdown()
+
+    def _remote_embed(self, texts, query_mode=False):
+        """
+        Perform remote embedding using a HTTP POST request to a designated endpoint.
+
+        Args:
+            texts (List[str]): A list of text strings to be embedded.
+            query_mode (bool): A flag to indicate if running in query mode or in embed mode (indexing).
+        Returns:
+            List[List[float]]: A list of embedded representations of the input texts.
+        """
+        session = requests_retry_session(
+            retries=self.max_retries, backoff_factor=self.retry_backoff_factor
+        )
+
+        payload = {
+            "input": texts,
+            "model": self.model,
+        }
+
+        url = f"{self.tfy_llm_gateway_url}/openai/embeddings"
+        logger.debug(
+            f"Embedding using - model: {self.model} at endpoint: {url}, for {len(texts)} texts"
+        )
+        response = session.post(
+            url=url,
+            json=payload,
+            headers={
+                "Authorization": f"Bearer {self.tfy_api_key}",
+            },
+            timeout=self.request_timeout,
+        )
+        response.raise_for_status()
+        output = response.json()
+        return [data["embedding"] for data in output["data"]]
+
+    def _embed(self, texts: List[str], query_mode: bool):
+        """
+        Perform embedding on a list of texts using remote embedding in chunks.
+
+        Args:
+            texts (List[str]): A list of text strings to be embedded.
+            query_mode (bool): A flag to indicate if running in query mode or in embed mode (indexing).
+        Returns:
+            List[List[float]]: A list of embedded representations of the input texts.
+        """
+        embeddings = []
+
+        def _feeder():
+            for i in range(0, len(texts), self.batch_size):
+                chunk = texts[i : i + self.batch_size]
+                yield chunk
+
+        embeddings = list(
+            tqdm.tqdm(
+                self._executor.map(self._remote_embed, _feeder()),
+                total=int(math.ceil(len(texts) / self.batch_size)),
+            )
+        )
+        return [item for batch in embeddings for item in batch]
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """
+        Embed a list of text documents.
+
+        Args:
+            texts (List[str]): A list of text documents to be embedded.
+
+        Returns:
+            List[List[float]]: A list of embedded representations of the input documents.
+        """
+        return self._embed(texts, query_mode=False)
+
+    def embed_query(self, text: str) -> List[float]:
+        """
+        Embed a query text.
+
+        Args:
+            text (str): The query text to be embedded.
+
+        Returns:
+            List[float]: The embedded representation of the input query text.
+        """
+        return self._embed([text], query_mode=True)[0]

truefoundry/langchain/truefoundry_llm.py

@@ -0,0 +1,106 @@
+from typing import Any, Dict, List, Optional
+
+from langchain.llms.base import LLM
+from langchain.pydantic_v1 import Extra, Field, root_validator
+
+from truefoundry.langchain.utils import (
+    requests_retry_session,
+    validate_tfy_environment,
+)
+from truefoundry.logger import logger
+
+
+class TrueFoundryLLM(LLM):
+    """`TrueFoundry LLM Gateway` completion models API.
+
+    To use, you must have the environment variable ``TFY_API_KEY`` set with your API key and ``TFY_HOST`` set with your host or pass it as a named parameter to the constructor.
+    """
+
+    model: str = Field(description="The model to use for completion.")
+    """The model to use for completion."""
+    tfy_llm_gateway_url: Optional[str] = Field(default=None)
+    """TrueFoundry LLM Gateway endpoint URL. Automatically inferred from env var `TFY_LLM_GATEWAY_URL` if not provided."""
+    tfy_api_key: Optional[str] = Field(default=None)
+    """TrueFoundry API Key. Automatically inferred from env var `TFY_API_KEY` if not provided."""
+    model_parameters: Optional[dict] = Field(default_factory=dict)
+    """Model parameters"""
+    request_timeout: int = Field(default=30)
+    """The timeout for the request in seconds."""
+    max_retries: int = Field(default=5)
+    """The number of retries for HTTP requests."""
+    retry_backoff_factor: float = Field(default=0.3)
+    """The backoff factor for exponential backoff during retries."""
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+        allow_population_by_field_name = True
+
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        values = validate_tfy_environment(values)
+        if not values["tfy_api_key"]:
+            raise ValueError(
+                "Did not find `tfy_api_key`, please add an environment variable"
+                " `TFY_API_KEY` which contains it, or pass"
+                "  `tfy_api_key` as a named parameter."
+            )
+        if not values["tfy_llm_gateway_url"]:
+            raise ValueError(
+                "Did not find `tfy_llm_gateway_url`, please add an environment variable"
+                " `TFY_LLM_GATEWAY_URL` which contains it, or pass"
+                "  `tfy_llm_gateway_url` as a named parameter."
+            )
+        return values
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of llm model."""
+        return "truefoundry-llm"
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Call out to the deployed model
+
+        Args:
+            prompt: The prompt to pass into the model.
+            stop: Optional list of stop words to use when generating.
+
+        Returns:
+            The string generated by the model.
+
+        Example:
+            .. code-block:: python
+
+                response = model("I have a joke for you...")
+        """
+
+        payload = {**self.model_parameters} if self.model_parameters else {}
+        if stop:
+            payload["stop_sequences"] = stop
+
+        payload["prompt"] = prompt
+        payload["model"] = self.model
+
+        session = requests_retry_session(
+            retries=self.max_retries, backoff_factor=self.retry_backoff_factor
+        )
+
+        url = f"{self.tfy_llm_gateway_url}/openai/completions"
+        logger.debug(f"Completion using - model: {self.model} at endpoint: {url}")
+        response = session.post(
+            url=url,
+            json=payload,
+            headers={
+                "Authorization": f"Bearer {self.tfy_api_key}",
+            },
+            timeout=self.request_timeout,
+        )
+        response.raise_for_status()
+        output = response.json()
+        return output["choices"][0]["text"]

truefoundry/langchain/utils.py

@@ -0,0 +1,85 @@
+import os
+from typing import Dict, List, Optional
+from urllib.parse import urljoin
+
+import requests
+from langchain.pydantic_v1 import BaseModel
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+from truefoundry.deploy.lib.auth.servicefoundry_session import ServiceFoundrySession
+
+
+class ModelParameters(BaseModel):
+    temperature: Optional[float]
+    maximum_length: Optional[int]
+    top_p: Optional[float]
+    top_k: Optional[int]
+    repetition_penalty: Optional[float]
+    frequency_penalty: Optional[float]
+    presence_penalty: Optional[float]
+    stop_sequences: Optional[List[str]]
+
+
+def validate_tfy_environment(values: Dict):
+    gateway_url = values["tfy_llm_gateway_url"] or os.getenv("TFY_LLM_GATEWAY_URL")
+    api_key = values["tfy_api_key"] or os.getenv("TFY_API_KEY")
+
+    if gateway_url and api_key:
+        values["tfy_llm_gateway_url"] = gateway_url
+        values["tfy_api_key"] = api_key
+        return values
+
+    sfy_session = ServiceFoundrySession()
+    if not sfy_session:
+        raise Exception(
+            "Unauthenticated: Please login using truefoundry login --host <https://example-domain.com>"
+        )
+
+    if not gateway_url:
+        gateway_url = urljoin(sfy_session.base_url, "/api/llm")
+
+    if not api_key:
+        api_key = sfy_session.access_token
+
+    values["tfy_llm_gateway_url"] = gateway_url
+    values["tfy_api_key"] = api_key
+    return values
+
+
+def requests_retry_session(
+    retries=5,
+    backoff_factor=0.3,
+    status_forcelist=(500, 502, 503, 504),
+    method_whitelist=frozenset({"GET", "POST"}),
+    session=None,
+):
+    """
+    Returns a `requests` session with retry capabilities for certain HTTP status codes.
+
+    Args:
+        retries (int): The number of retries for HTTP requests.
+        backoff_factor (float): The backoff factor for exponential backoff during retries.
+        status_forcelist (tuple): A tuple of HTTP status codes that should trigger a retry.
+        method_whitelist (frozenset): The set of HTTP methods that should be retried.
+        session (requests.Session, optional): An optional existing requests session to use.
+
+    Returns:
+        requests.Session: A session with retry capabilities.
+    """
+    # Implementation taken from https://www.peterbe.com/plog/best-practice-with-retries-with-requests
+    session = session or requests.Session()
+    retry = Retry(
+        total=retries,
+        read=retries,
+        connect=retries,
+        status=retries,
+        backoff_factor=backoff_factor,
+        allowed_methods=method_whitelist,
+        status_forcelist=status_forcelist,
+        respect_retry_after_header=True,
+    )
+    adapter = HTTPAdapter(max_retries=retry)
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
+    return session

truefoundry/logger.py

@@ -0,0 +1,17 @@
+# https://docs.python.org/3/howto/logging.html#library-config
+import logging
+
+from rich.logging import RichHandler
+
+logger = logging.getLogger("truefoundry")
+
+
+def add_cli_handler(level: int = logging.INFO, show_path=False):
+    # TODO (chiragjn): Probably don't use rich handler, it adds hard breaks
+    # See: https://github.com/Textualize/rich/discussions/344
+    # Maybe try: https://github.com/pycontribs/enrich/blob/main/src/enrich/logging.py
+    # or simpler override of logging.Handler with console.print(rich.Text("..."))
+    handler = RichHandler(level=level, show_path=show_path)
+    handler.setLevel(level)
+    logger.addHandler(handler)
+    logger.setLevel(logging.DEBUG)

truefoundry/pydantic_v1.py

@@ -0,0 +1,5 @@
+try:
+    from pydantic.v1 import *  # noqa: F403
+    from pydantic.v1 import utils  # noqa: F401
+except ImportError:
+    from pydantic import *  # noqa: F403

truefoundry/python_deploy_codegen.py

@@ -0,0 +1,132 @@
+import ast
+import io
+import re
+from typing import List
+
+from rich.console import Console
+from rich.pretty import pprint
+
+from truefoundry.deploy import Application
+
+
+def generate_code(
+    symbols_to_import: List[str],
+    application_type: str,
+    spec_repr: str,
+    workspace_fqn: str,
+):
+    symbols = ",".join(symbols_to_import)
+    code = f"""\
+import logging
+from truefoundry.deploy import (
+    {symbols},
+)
+logging.basicConfig(level=logging.INFO)
+
+{application_type} = {spec_repr}
+
+{application_type}.deploy(workspace_fqn="{workspace_fqn}")\
+"""
+    return code
+
+
+def extract_class_names(code):
+    tree = ast.parse(code)
+
+    # Function to extract keywords from the AST
+    def extract_class_names_from_ast_tree(node):
+        keywords = set()
+        for child_node in ast.iter_child_nodes(node):
+            if isinstance(child_node, ast.Call):
+                keywords.add(child_node.func.id)
+            keywords.update(extract_class_names_from_ast_tree(child_node))
+        return keywords
+
+    # Get keywords from the main body of the code
+    main_keywords = extract_class_names_from_ast_tree(tree)
+    return list(main_keywords)
+
+
+def replace_enums_with_values(raw_str):
+    # required to replace enums of format <AppProtocol.HTTP: 'http'> with 'http'
+    pattern = r'<([a-zA-Z0-9_]+).[a-zA-Z0-9_]+: [\'"](.+)[\'"]>'
+    replacement = r"'\2'"
+
+    result = re.sub(pattern, replacement, raw_str)
+    return result
+
+
+def remove_none_type_fields(code):
+    lines = code.split("\n")
+    new_lines = [
+        line
+        for line in lines
+        if not (line.endswith("=None") or line.endswith("=None,"))
+    ]
+    formatted_code = "\n".join(new_lines)
+    return formatted_code
+
+
+def remove_type_field(code):
+    lines = code.split("\n")
+    new_lines = [re.sub(r'^[ \t]*type=[\'"][^"]*[\'"],', "", line) for line in lines]
+    return "\n".join(new_lines)
+
+
+def add_deploy_line(code, workspace_fqn, application_type):
+    deploy_line = f"{application_type}.deploy('workspace_fqn={workspace_fqn}')"
+    return code + "\n" + deploy_line
+
+
+def get_python_repr(obj):
+    stream = io.StringIO()
+    console = Console(file=stream, no_color=True, highlighter=None)
+    pprint(obj, expand_all=True, console=console, indent_guides=False)
+    return stream.getvalue()
+
+
+COMMENT_FOR_LOCAL_SOURCE = """# Set build_source=LocalSource(local_build=False), in order to deploy code from your local.
+# With local_build=False flag, docker image will be built on cloud instead of local
+# Else it will try to use docker installed on your local machine to build the image"""
+
+
+def add_local_source_comment(code):
+    lines = code.split("\n")
+    new_lines = []
+    for line in lines:
+        if line.lstrip(" ").startswith("build_source=GitSource"):
+            new_lines.append(COMMENT_FOR_LOCAL_SOURCE)
+        new_lines.append(line)
+    return "\n".join(new_lines)
+
+
+def convert_deployment_config_to_python(workspace_fqn: str, deployment_config: dict):
+    """
+    Convert a deployment config to a python file that can be used to deploy to a workspace
+    """
+    application = Application.parse_obj(deployment_config)
+    application_type = application.__root__.type
+
+    spec_repr = get_python_repr(application.__root__)
+    spec_repr = replace_enums_with_values(spec_repr)
+    spec_repr = remove_none_type_fields(spec_repr)
+    spec_repr = remove_type_field(spec_repr)
+
+    # extract class names to import
+    symbols_to_import = extract_class_names(spec_repr)
+
+    # check if GitSource exists in array of symbols to import
+    if "GitSource" in symbols_to_import:
+        symbols_to_import.append("LocalSource")
+
+    generated_code = generate_code(
+        symbols_to_import=symbols_to_import,
+        application_type=application_type,
+        spec_repr=spec_repr,
+        workspace_fqn=workspace_fqn,
+    )
+
+    if "GitSource" in symbols_to_import:
+        generated_code = add_local_source_comment(generated_code)
+
+    return generated_code

truefoundry/version.py

@@ -0,0 +1,6 @@
+try:
+    import importlib.metadata
+
+    __version__ = importlib.metadata.version("truefoundry")
+except Exception:
+    __version__ = "NA"

truefoundry/workflow/__init__.py

@@ -0,0 +1,19 @@
+try:
+    from flytekit import task
+except ImportError:
+    print("To use workflows, please run 'pip install truefoundry[workflow]'.")
+
+from flytekit import conditional
+from flytekit.types.directory import FlyteDirectory
+
+from truefoundry.deploy.v2.lib.patched_models import (
+    ContainerTaskConfig,
+    PythonTaskConfig,
+    TaskDockerFileBuild,
+    TaskPythonBuild,
+)
+from truefoundry.workflow.container_task import ContainerTask
+from truefoundry.workflow.map_task import map_task
+from truefoundry.workflow.python_task import PythonFunctionTask
+from truefoundry.workflow.task import task
+from truefoundry.workflow.workflow import ExecutionConfig, workflow

truefoundry/workflow/container_task.py

@@ -0,0 +1,12 @@
+from flytekit import ContainerTask
+
+from truefoundry.workflow import ContainerTaskConfig
+
+
+class ContainerTask(ContainerTask):
+    def __init__(self, name: str, task_config: ContainerTaskConfig):
+        super().__init__(name=name, image="", command=[])
+        self.tfy_task_config = task_config
+
+    def get_custom(self, settings) -> dict:
+        return {"truefoundry": self.tfy_task_config.dict()}