Job

Submodules

llmsmith.job.base module

class llmsmith.job.base.Job

Bases: Generic[T, U], ABC

A subclass of Job can run a list of tasks and store their inputs and outputs. An instance of llmsmith.task.base.JobMemory is used internally for storing the task inputs/outputs.

abstract async run(user_input: T)

Abstract method to run the Job with the given initial user input.

Parameters:: user_input (T) – The input provided by the user.

task_input(key: str) → TaskInput | None

Return the input for a task.

Parameters:: key (str) – The key of the task.
Returns:: The input of the task or None if not found.
Return type:: Union[TaskInput, None]

task_output(key: str) → TaskOutput | None

Return the output of a task.

Parameters:: key (str) – The key of the task.
Returns:: The output of the task or None if not found.
Return type:: Union[TaskOutput, None]

class llmsmith.job.base.JobMemory

Bases: object

JobMemory is a key-value store which stores the input and output values for the tasks (llmsmith.task.base.Task) present in the job.

add_task_input(key: str, task_input: TaskInput)

Stores the input value passed to a task against the specified key.

Parameters:

key (str) – key against which input value is stored
task_input (llmsmith.task.models.TaskInput) – input of the task

add_task_output(key: str, task_output: TaskOutput)

Stores the output value returned by a task against the specified key.

Parameters:

key (str) – key against which output value is stored
task_output (llmsmith.task.models.TaskOutput) – output of the task

Raises:

KeyError – if there is no input value stored against the task

get_task_input(key: str) → TaskInput | None

Returns the task input from the memory.

Parameters:: key (str) – key against which input value is stored
Returns:: input value passed to the task. Returns None if the key is not available in the memory
Return type:: llmsmith.task.models.TaskInput

get_task_output(key: str) → TaskOutput | None

Returns the task output from the memory.

Parameters:: key (str) – key against which output value is stored
Returns:: output value returned by the task. Returns None if the key is not available in the memory
Return type:: llmsmith.task.models.TaskOutput

llmsmith.job.job module

class llmsmith.job.job.ConcurrentJob

Bases: Job

An implementation of llmsmith.job.base.Job which executes the given tasks concurrently. Every task added to an instance of ConcurrentJob will share the same input, which is the initial user input passed to the job.

Consider the below tasks for example:

Write a crime thriller movie idea based on the information given by the user.
Write a horror movie idea based on the information given by the user.

Since both tasks are based on the same input provided by the user, we can do it concurrently.

Following is the code for the above flow using llmsmith.

import openai

from llmsmith.job.job import ConcurrentJob
from llmsmith.task.textgen.options.openai import OpenAITextGenOptions
from llmsmith.task.textgen.openai import OpenAITextGenTask

llm = openai.AsyncOpenAI(api_key=OPEN_AI_API_KEY)

user_input = "Write a movie idea based on the below information:\n Protagonist: introvert college student\n Country: Japan\n Location: university\n"

crime_movie_task = OpenAITextGenTask(
    name="openai-crime-movie-idea",
    llm=llm,
    llm_options=OpenAITextGenOptions(
        model="gpt-3.5-turbo",
        temperature=0.3,
        system_prompt="You are a movie script writer working on a crime thriller movie script"
    ),
)

horror_movie_task = OpenAITextGenTask(
    name="openai-horror-movie-idea",
    llm=llm,
    llm_options=OpenAITextGenOptions(
        model="gpt-3.5-turbo",
        temperature=0.3,
        system_prompt="You are a movie script writer working on a horror movie script"
    ),
)

job = ConcurrentJob()

# Add task for writing crime thriller movie idea
job.add_task(crime_movie_task)

# Add task for writing horror movie idea
job.add_task(horror_movie_task)

# Run the job. The 2 steps will be executed concurrently
await job.run(user_input)

# Print the output of the both tasks
print(job.task_output("openai-crime-movie-idea"))
print(job.task_output("openai-horror-movie-idea"))

add_task(task: Task) → Self

Add a task to the job.

Parameters:: task (llmsmith.task.base.Task) – task to be added to the job
Returns:: Self
Return type:: llmsmith.job.job.ConcurrentJob

async run(user_input: T)

Run the tasks concurrently.

Parameters:: user_input (T) – The initial input for the job

class llmsmith.job.job.SequentialJob

Bases: Job

An implementation of llmsmith.job.base.Job which executes the given tasks sequentially. When adding a task, it is possible to pass the input/output values of previous tasks via an input template with placeholders. The placeholders can be in the following formats:

For replacing with the input value of a previous task: {{task-name.input}}
For replacing with the output value of a previous task: {{task-name.output}}
For replacing with the initial user input which is passed to the job while running it: {{root}}

A simple RAG implementation can be used as an example here for showcasing the above points.

Consider the below flow:

An user query is passed as input to a retriever (chromaDB)
Retriever output is passed to an LLM (OpenAI) to rephrase the query
Rephrased query is used as input to an LLM (OpenAI) to get the answer

Following is the code for the above flow using llmsmith.

import chromadb
import openai

from chromadb.utils import embedding_functions

from llmsmith.job.job import SequentialJob

from llmsmith.task.retrieval.vector.chromadb import ChromaDBRetriever
from llmsmith.task.textgen.options.openai import OpenAITextGenOptions
from llmsmith.task.textgen.openai import OpenAITextGenTask

chroma_client = chromadb.HttpClient(host=CHROMA_DB_HOST, port=8000)
llm = openai.AsyncOpenAI(api_key=OPEN_AI_API_KEY)
collection = chroma_client.get_collection(
    name="university_faq", embedding_function=embedding_functions.ONNXMiniLM_L6_V2()
)

user_input = "I'm interested in ML and AI. Tell me about the courses offered here which will match my interests."

retrieval_task = ChromaDBRetriever(
    name="chromadb-retriever",
    collection=collection,
)

rephrase_task = OpenAITextGenTask(
    name="openai-rephraser",
    llm=llm,
    llm_options=OpenAITextGenOptions(model="gpt-3.5-turbo", temperature=0),
)

generate_answer_task = OpenAITextGenTask(
    name="openai-answer-generator",
    llm=llm,
    llm_options=OpenAITextGenOptions(model="gpt-3.5-turbo", temperature=0),
)

job = SequentialJob()

# First step - retrieve the relevant documents from Chroma DB
job.add_task(retrieval_task)

# Second step - Rephrase the question based on the documents retrieved from Chroma DB.
# Note the placeholders {{root}} and {{chromadb-retriever.output}}.
# {{chromadb-retriever.output}} will be replaced by the output value of Chroma DB retriever
# {{root}} will be replaced with the initial user input
job.add_task(
    rephrase_task,
    input_template="Rephrase the question based on the context: \n\n QUESTION:\n{{root}}\n\nCONTEXT:\n{{chromadb-retriever.output}}",
)

# Third step - Answer the question based on the rephrased question and the relevant context documents retrieved from Chroma DB.
# Note the placeholders {{openai-rephraser.output}} and {{chromadb-retriever.output}}.
# {{openai-rephraser.output}} will be replaced by the output value of query rephraser task
# {{chromadb-retriever.output}} will be replaced by the output value of Chroma DB retriever
job.add_task(
    generate_answer_task,
    input_template="Answer the question based on the context: \n\n QUESTION:\n{{openai-rephraser.output}}\n\nCONTEXT:\n{{chromadb-retriever.output}}",
)

# Run the job. The 3 steps will be executed sequentially
await job.run(user_input)

# Print the output of the final task in the job
print(job.task_output("openai-answer-generator"))

add_task(task: Task, input_template: str | None = '{{root}}') → Self

Add a task to the job. An optional input template can also be passed which can be used to pass the input/output values of previous tasks via placeholders.

Parameters:

task (llmsmith.task.base.Task) – task to be added to the job
input_template (str, optional) – string template with placeholders referring to inputs/outputs of previous tasks. Defaults to {{root}} which refers to initial user input.

Returns:

Self

Return type:

llmsmith.job.job.SequentialJob

async run(user_input: T)

Run the tasks sequentially.

Parameters:: user_input (T) – The initial input for the job

Job

Submodules

llmsmith.job.base module

llmsmith.job.job module

Module contents