Shared llm package refactoring (#336)

Co-authored-by: = Enea_Gore <[email protected]>

.github/workflows/build.yml

@@ -56,6 +56,13 @@ jobs:
           docker build -t athena .
           cd ..
 
+      - name: Build llm_core image
+        id: set-image-core_llm
+        run: |
+          cd modules/llm_core
+          docker build -t llm_core .
+          cd ..
+
       - name: Docker Login
         id: docker-login
         run: |

docker-compose.yml

@@ -8,6 +8,13 @@ services:
     image: athena
     command: echo "Athena build succeeded, exiting (this is normal)"
 
+  llm_core:
+      build: modules/llm_core
+      depends_on:
+        - athena
+      image: llm_core
+      command: echo "llm_core build succeeded, exiting (this is normal)"
+
   assessment_module_manager:
     build: ./assessment_module_manager
     depends_on:
@@ -30,6 +37,7 @@ services:
     build: modules/programming/module_programming_llm
     depends_on:
       - athena
+      - llm_core
     ports:
       - "5002:5002"
 
@@ -38,6 +46,7 @@ services:
     build: modules/text/module_text_llm
     depends_on:
       - athena
+      - llm_core
     ports:
       - "5003:5003"
 
@@ -70,5 +79,6 @@ services:
     build: modules/modeling/module_modeling_llm
     depends_on:
       - athena
+      - llm_core
     ports:
       - "5008:5008"

modules/llm_core/.env.example

@@ -0,0 +1,39 @@
+# Comment out the variables that you define somewhere else
+# Environment variables are overwritten by .env file
+
+PRODUCTION=0
+SECRET=12345abcdef
+DATABASE_URL=sqlite:///../data/data.sqlite
+
+
+################################################################
+# LLM Credentials                                              #
+################################################################
+
+# Default model to use
+# See below for options, available models are also logged on startup
+LLM_DEFAULT_MODEL="azure_openai_gpt-35-turbo"
+
+# Enable LLM-as-a-judge approach 0 = disabled, 1 = enabled
+LLM_ENABLE_LLM_AS_A_JUDGE=1
+# Evaluation model to use for the LLM-as-a-judge approach [Only important if you want to use it in the /evaluate endpoint]
+# See below for options, available models are also logged on startup
+LLM_EVALUATION_MODEL="azure_openai_gpt-4o"
+
+# Standard OpenAI (Non-Azure) [leave blank if not used]
+# Model names prefixed with `openai_` followed by the model name, e.g. `openai_text-davinci-003`
+# A list of models can be found in `module_text_llm/helpers/models/openai.py` (openai_models)
+OPENAI_API_KEY="xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
+
+# Azure OpenAI [leave blank if not used]
+# Model names prefixed with `azure_openai_` followed by the deployment id, e.g. `azure_openai_gpt-35`
+AZURE_OPENAI_API_KEY="xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
+AZURE_OPENAI_ENDPOINT="https://ase-eu01.openai.azure.com/" # change base if needed
+OPENAI_API_VERSION="2024-06-01" # change base if needed
+
+# LangSmith (can be used for tracing LLMs) [leave blank if not used]
+# See https://docs.smith.langchain.com
+# LANGCHAIN_TRACING_V2=true
+# LANGCHAIN_ENDPOINT="https://api.smith.langchain.com"
+# LANGCHAIN_API_KEY="XXX"
+# LANGCHAIN_PROJECT="XXX"

modules/llm_core/Dockerfile

@@ -0,0 +1,24 @@
+# syntax=docker/dockerfile:1
+
+# This is the Dockerfile for the shared llm package.
+# Its output is used as a dependency in the module_* Dockerfiles.
+
+FROM python:3.11 as llm_core
+
+WORKDIR /code
+
+# Poetry
+RUN pip install --no-cache-dir poetry==1.5.0
+
+# Dependencies
+COPY pyproject.toml poetry.lock ./
+COPY --from=athena /code /athena
+
+RUN poetry config virtualenvs.create false \
+    && poetry install --no-interaction --no-ansi
+
+# Project files
+COPY . ./
+
+# Build the package
+RUN poetry build -f wheel

modules/llm_core/llm_core/__init__.py

@@ -0,0 +1,4 @@
+import dotenv
+
+# Load environment variables from .env file (for local development)
+dotenv.load_dotenv(override=True)

modules/modeling/module_modeling_llm/module_modeling_llm/models/__init__.py → modules/llm_core/llm_core/models/__init__.py

@@ -2,8 +2,7 @@
 from typing import Type, Union, List, Optional
 from langchain.base_language import BaseLanguageModel
 
-from module_modeling_llm.models.model_config import ModelConfig
-
+from llm_core.models.model_config import ModelConfig
 
 
 DefaultModelConfig: Type[ModelConfig]
@@ -15,7 +14,7 @@
 
 types: List[Type[ModelConfig]] = []
 try:
-    import module_modeling_llm.models.openai as openai_config
+    import llm_core.models.openai as openai_config
     types.append(openai_config.OpenAIModelConfig)
     if default_model_name in openai_config.available_models:
         DefaultModelConfig = openai_config.OpenAIModelConfig
@@ -36,4 +35,4 @@
     ModelConfigType = type0
 else:
     type1 = types[1]
-    ModelConfigType = Union[type0, type1] # type: ignore
+    ModelConfigType = Union[type0, type1] # type: ignore

modules/text/module_text_llm/module_text_llm/helpers/llm_utils.py → modules/llm_core/llm_core/utils/llm_utils.py

@@ -1,20 +1,15 @@
-from typing import Optional, Type, TypeVar, List
-from pydantic import BaseModel, ValidationError
+from typing import Type, TypeVar, List
+from pydantic import BaseModel
 import tiktoken
-
-from langchain.chains import LLMChain
 from langchain.chat_models import ChatOpenAI
 from langchain.base_language import BaseLanguageModel
 from langchain.prompts import (
     ChatPromptTemplate,
     SystemMessagePromptTemplate,
     HumanMessagePromptTemplate,
 )
-from langchain.chains.openai_functions import create_structured_output_chain
 from langchain.output_parsers import PydanticOutputParser
-from langchain.schema import OutputParserException
-
-from athena import emit_meta, get_experiment_environment
+from athena import emit_meta
 
 T = TypeVar("T", bound=BaseModel)
 
@@ -112,51 +107,4 @@ def get_chat_prompt_with_formatting_instructions(
     system_message_prompt.prompt.partial_variables = {"format_instructions": output_parser.get_format_instructions()}
     system_message_prompt.prompt.input_variables.remove("format_instructions")
     human_message_prompt = HumanMessagePromptTemplate.from_template(human_message + "\n\nJSON response following the provided schema:")
-    return ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
-
-
-async def predict_and_parse(
-        model: BaseLanguageModel, 
-        chat_prompt: ChatPromptTemplate, 
-        prompt_input: dict, 
-        pydantic_object: Type[T], 
-        tags: Optional[List[str]]
-    ) -> Optional[T]:
-    """Predicts an LLM completion using the model and parses the output using the provided Pydantic model
-
-    Args:
-        model (BaseLanguageModel): The model to predict with
-        chat_prompt (ChatPromptTemplate): Prompt to use
-        prompt_input (dict): Input parameters to use for the prompt
-        pydantic_object (Type[T]): Pydantic model to parse the output
-        tags (Optional[List[str]]: List of tags to tag the prediction with
-
-    Returns:
-        Optional[T]: Parsed output, or None if it could not be parsed
-    """
-    experiment = get_experiment_environment()
-
-    tags = tags or []
-    if experiment.experiment_id is not None:
-        tags.append(f"experiment-{experiment.experiment_id}")
-    if experiment.module_configuration_id is not None:
-        tags.append(f"module-configuration-{experiment.module_configuration_id}")
-    if experiment.run_id is not None:
-        tags.append(f"run-{experiment.run_id}")
-
-    if supports_function_calling(model):
-        chain = create_structured_output_chain(pydantic_object, llm=model, prompt=chat_prompt, tags=tags)
-
-        try:
-            return await chain.arun(**prompt_input)
-        except (OutputParserException, ValidationError):
-            # In the future, we should probably have some recovery mechanism here (i.e. fix the output with another prompt)
-            return None
-
-    output_parser = PydanticOutputParser(pydantic_object=pydantic_object)
-    chain = LLMChain(llm=model, prompt=chat_prompt, output_parser=output_parser, tags=tags)
-    try:
-        return await chain.arun(**prompt_input)
-    except (OutputParserException, ValidationError):
-        # In the future, we should probably have some recovery mechanism here (i.e. fix the output with another prompt)
-        return None
+    return ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])

modules/modeling/module_modeling_llm/module_modeling_llm/utils/predict_and_parse.py → modules/llm_core/llm_core/utils/predict_and_parse.py

@@ -8,12 +8,24 @@
 T = TypeVar("T", bound=BaseModel)
 
 async def predict_and_parse(
-        model: BaseLanguageModel,
-        chat_prompt: ChatPromptTemplate,
-        prompt_input: dict,
-        pydantic_object: Type[T],
+        model: BaseLanguageModel, 
+        chat_prompt: ChatPromptTemplate, 
+        prompt_input: dict, 
+        pydantic_object: Type[T], 
         tags: Optional[List[str]]
-) -> Optional[T]:
+    ) -> Optional[T]:
+    """Predicts an LLM completion using the model and parses the output using the provided Pydantic model
+
+    Args:
+        model (BaseLanguageModel): The model to predict with
+        chat_prompt (ChatPromptTemplate): Prompt to use
+        prompt_input (dict): Input parameters to use for the prompt
+        pydantic_object (Type[T]): Pydantic model to parse the output
+        tags (Optional[List[str]]: List of tags to tag the prediction with
+
+    Returns:
+        Optional[T]: Parsed output, or None if it could not be parsed
+    """
     experiment = get_experiment_environment()
 
     tags = tags or []