├── requirements.txt ├── pyproject.toml ├── .gitignore ├── util └── models.py ├── fsm_llm ├── __init__.py ├── utils.py ├── state_models.py ├── llm_handler.py └── fsm.py ├── setup.py ├── examples ├── switch_agent.py ├── support_agent.py ├── tutor_agent.py ├── tutor_agent_modified.py └── medical_agent.py ├── content ├── calculus_example.json └── calculus_content.json ├── README.md └── LICENSE /requirements.txt: -------------------------------------------------------------------------------- 1 | openai>=0.27.0 2 | pydantic>=1.10.0 3 | jinja2>=3.0.0 4 | python-dotenv>=0.15.0 5 | asyncio 6 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools", "wheel"] 3 | build-backend = "setuptools.build_meta" 4 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Ignore 2 | __pycache__/ 3 | *.py[cod] 4 | *.swp 5 | .DS_Store 6 | .idea 7 | .env 8 | dist 9 | 10 | # Ignore egg-info directories 11 | *.egg-info/ 12 | 13 | # Ignore build artifacts 14 | built/ 15 | bdist.win-amd64/ 16 | lib/ 17 | -------------------------------------------------------------------------------- /util/models.py: -------------------------------------------------------------------------------- 1 | import openai 2 | from dotenv import load_dotenv 3 | import os 4 | 5 | load_dotenv() 6 | 7 | openai.api_key = os.getenv("OPENAI_API_KEY") 8 | openai.organization = os.getenv("OPENAI_ORGANIZATION") 9 | 10 | client = openai.OpenAI() 11 | models_page = client.models.list() 12 | 13 | for model in models_page: 14 | print(model.id) 15 | -------------------------------------------------------------------------------- /fsm_llm/__init__.py: -------------------------------------------------------------------------------- 1 | # core/__init__.py 2 | from .fsm import LLMStateMachine 3 | from .state_models import FSMState, FSMRun, FSMError, ImmediateStateChange, DefaultResponse 4 | from .llm_handler import LLMUtilities 5 | from .utils import wrap_into_json_response 6 | 7 | __all__ = [ 8 | "LLMStateMachine", 9 | "FSMState", 10 | "FSMRun", 11 | "FSMError", 12 | "ImmediateStateChange", 13 | "DefaultResponse", 14 | "LLMUtilities", 15 | "wrap_into_json_response", 16 | ] 17 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | with open("README.md", "r", encoding="utf-8") as f: 4 | long_description = f.read() 5 | 6 | setup( 7 | name="fsm-llm", 8 | version="0.1.3", 9 | description="A Python framework for creating FSM-based LLM agents", 10 | long_description=long_description, 11 | long_description_content_type="text/markdown", 12 | author="Jeffrey Zhou", 13 | author_email="jeffreyzhou3@outlook.com", 14 | url="https://github.com/jsz-05/LLM-State-Machine", 15 | packages=find_packages(), 16 | include_package_data=True, 17 | install_requires=[ 18 | "openai>=0.27.0", 19 | "pydantic>=1.10.0", 20 | "jinja2>=3.0.0", 21 | "python-dotenv>=0.15.0", 22 | ], 23 | classifiers=[ 24 | "Programming Language :: Python :: 3", 25 | "License :: OSI Approved :: Apache Software License", 26 | "Operating System :: OS Independent", 27 | ], 28 | python_requires=">=3.8", 29 | ) 30 | -------------------------------------------------------------------------------- /fsm_llm/utils.py: -------------------------------------------------------------------------------- 1 | import json 2 | from typing import Callable, Dict, Any, Literal, Optional, Type, Union 3 | from pydantic import BaseModel, create_model 4 | 5 | from fsm_llm.state_models import FSMState, DefaultResponse 6 | 7 | 8 | def _generate_response_schema( 9 | current_state_model: Union[Type[BaseModel], None], 10 | transitions: dict[str, str], 11 | default_state: str, 12 | ) -> Type[BaseModel]: 13 | """Create a response model based on the current state model and transitions, this will be used for structured_response openai param.""" 14 | 15 | # Extract the transition keys as a tuple for the Literal type 16 | transition_keys = tuple([default_state] + list(transitions.keys())) 17 | 18 | next_state_key_type = Literal.__getitem__(transition_keys) 19 | 20 | if not current_state_model: 21 | current_state_model = DefaultResponse 22 | 23 | # Dynamically create the model with response and next_state_key fields 24 | return create_model( 25 | "EnclosedResponse", 26 | response=(current_state_model, ...), 27 | next_state_key=(next_state_key_type, ...), 28 | ) 29 | 30 | 31 | def _add_transitions(prompt_template: str, fsm_state: FSMState) -> str: 32 | """Add transitions to the system prompt.""" 33 | prompt_template += f"\n\nYou are currently in {fsm_state.key} and based on user input, you can transition to the following states (with conditions defined):" 34 | for key, value in fsm_state.transitions.items(): 35 | prompt_template += f"\n- {key}: {value}" 36 | 37 | prompt_template += "\n\nIn response add the state you want to transition to.. (or leave blank to stay in the current state)" 38 | return prompt_template 39 | 40 | 41 | def wrap_into_json_response(data: BaseModel, next_state: str) -> BaseModel: 42 | dict_res = {"response": data.model_dump(), "next_state_key": next_state} 43 | 44 | return json.dumps(dict_res) 45 | -------------------------------------------------------------------------------- /fsm_llm/state_models.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Callable, Optional, Type 2 | import pydantic 3 | from pydantic import BaseModel 4 | 5 | class FSMState(pydantic.BaseModel): 6 | """Defines a state in the FSM for managing conversation flow. 7 | 8 | Params: 9 | - key (str): Unique identifier for the state. 10 | - func (Callable): Function defining the state action. 11 | - prompt_template (str): System prompt for the model. 12 | - temperature (float): Model's response randomness. 13 | - transitions (dict[str, str]): Maps user inputs to next states. 14 | - response_model (Optional[Type[BaseModel]]): Model for parsing the AI's response. 15 | - preprocess_input (Optional[Callable]): Preprocess user input before state function. 16 | - preprocess_chat (Optional[Callable]): Preprocess chat history before state function. 17 | - preprocess_prompt_template (Optional[Callable]): Preprocess the system prompt. 18 | """ 19 | key: str 20 | func: Callable 21 | prompt_template: str 22 | temperature: float 23 | transitions: dict[str, str] 24 | response_model: Optional[Type[BaseModel]] 25 | preprocess_input: Optional[Callable] 26 | preprocess_chat: Optional[Callable] 27 | preprocess_prompt_template: Optional[Callable] 28 | 29 | class DefaultResponse(BaseModel): 30 | """Default response model for AI output. 31 | 32 | Params: 33 | - content (str): Content of the AI response. 34 | """ 35 | content: str 36 | 37 | class FSMRun(pydantic.BaseModel): 38 | """Outcome of a single FSM step. 39 | 40 | Params: 41 | - state (str): Current state key. 42 | - chat_history (list[dict]): History of conversation. 43 | - context_data (dict[str, Any]): Relevant contextual data. 44 | - response_raw (dict): Raw AI model response. 45 | - response (Any): Processed response. 46 | """ 47 | state: str 48 | chat_history: list[dict] 49 | context_data: dict[str, Any] 50 | response_raw: dict 51 | response: Any 52 | 53 | class FSMError(Exception): 54 | """Custom exception for FSM-related errors.""" 55 | pass 56 | 57 | class VerifiedResponse(BaseModel): 58 | """Model for verifying transition responses. 59 | 60 | Params: 61 | - message (str): Verification message. 62 | - is_valid (bool): Whether the verification passed. 63 | """ 64 | message: str 65 | is_valid: bool 66 | 67 | class ImmediateStateChange(BaseModel): 68 | """Triggers immediate state transition. 69 | 70 | Params: 71 | - next_state (str): The state to transition to. 72 | - input (str): Input for the new state (default: "Hey"). 73 | - keep_original_response (bool): Preserve and prepend the original response. 74 | - keep_original_seperator (str): Separator between original and new response. 75 | """ 76 | next_state: str 77 | input: str = "Default" 78 | keep_original_response: bool = False 79 | keep_original_seperator: str = " " 80 | -------------------------------------------------------------------------------- /examples/switch_agent.py: -------------------------------------------------------------------------------- 1 | import os 2 | from dotenv import load_dotenv 3 | import openai 4 | from fsm_llm import LLMStateMachine 5 | from fsm_llm.state_models import FSMRun 6 | 7 | # Global variable to track on-off state 8 | SWITCH_STATE = "OFF" 9 | 10 | # Load environment variables 11 | load_dotenv() 12 | 13 | # Initialize OpenAI API key 14 | openai.api_key = os.getenv("OPENAI_API_KEY") 15 | openai.organization = os.getenv("OPENAI_ORGANIZATION") 16 | 17 | # Create the FSM 18 | fsm = LLMStateMachine(initial_state="START", end_state="END") 19 | 20 | 21 | # Define the START state 22 | @fsm.define_state( 23 | state_key="START", 24 | prompt_template="You are an on-off switcher. Ask the user if they want to turn switch the on or off.", 25 | transitions={"STATE_ON": "If user wants to turn on the switch", "END": "If user wants to end the conversation"}, 26 | ) 27 | async def start_state(fsm: LLMStateMachine, response: str, will_transition: bool): 28 | global SWITCH_STATE 29 | if will_transition and fsm.get_next_state() == "STATE_ON": 30 | SWITCH_STATE = "ON" 31 | print("SWITCH TURNED ON") 32 | elif will_transition and fsm.get_next_state() == "END": 33 | return "Goodbye!" 34 | return response 35 | 36 | 37 | # Define the STATE_ON state 38 | @fsm.define_state( 39 | state_key="STATE_ON", 40 | prompt_template="The switch is now on. Ask the user if they want to turn off the switch or end the conversation.", 41 | transitions={"START": "If user wants to turn off the switch", "END": "If user wants to end the conversation"}, 42 | ) 43 | async def state_on(fsm: LLMStateMachine, response: str, will_transition: bool): 44 | global SWITCH_STATE 45 | if will_transition and fsm.get_next_state() == "START": 46 | SWITCH_STATE = "OFF" 47 | print("SWITCH TURNED OFF") 48 | elif will_transition and fsm.get_next_state() == "END": 49 | return "Goodbye!" 50 | return response 51 | 52 | 53 | # Define the END state 54 | @fsm.define_state( 55 | state_key="END", 56 | prompt_template="Goodbye!", 57 | ) 58 | async def end_state(fsm: LLMStateMachine, response: str, will_transition: bool): 59 | return "Goodbye!" 60 | 61 | 62 | async def main(): 63 | """Example of a simple on-off switch FSM using LLMStateMachine""" 64 | # Create the OpenAI client 65 | openai_client = openai.AsyncOpenAI() 66 | 67 | print("Agent: Hi. I am an on-off switch manager.") 68 | while not fsm.is_completed(): # Run until FSM reaches the END state 69 | user_input = input("Your input: ") 70 | if user_input.lower() in ["quit", "exit"]: 71 | fsm.set_next_state("END") 72 | break 73 | run_state: FSMRun = await fsm.run_state_machine(openai_client, user_input=user_input) 74 | print(f"Agent: {run_state.response}") 75 | print("CURRENT SWITCH STATE:", SWITCH_STATE) 76 | 77 | print("Agent: Goodbye.") 78 | 79 | 80 | if __name__ == "__main__": 81 | import asyncio 82 | asyncio.run(main()) 83 | -------------------------------------------------------------------------------- /fsm_llm/llm_handler.py: -------------------------------------------------------------------------------- 1 | from typing import Callable, Dict, Any, Optional, Type 2 | from pydantic import BaseModel 3 | import openai 4 | import jinja2 5 | from .state_models import FSMError, FSMState 6 | from .utils import _add_transitions 7 | 8 | class LLMUtilities: 9 | """Handles all LLM-related operations including prompt processing and API calls.""" 10 | 11 | @staticmethod 12 | async def get_completion( 13 | async_openai_instance: openai.AsyncOpenAI, 14 | chat_history: list, 15 | response_model: Type[BaseModel], 16 | llm_model: str, 17 | current_state: Optional[FSMState] = None, 18 | ) -> dict: 19 | """Get completion from LLM with optional state-specific processing""" 20 | if current_state: 21 | # Process state-specific prompt and chat history 22 | processed_prompt = LLMUtilities.process_prompt_template( 23 | current_state.prompt_template, 24 | getattr(current_state, 'user_defined_context', {}), 25 | current_state.preprocess_prompt_template 26 | ) 27 | processed_prompt = _add_transitions(processed_prompt, current_state) 28 | 29 | system_message = {"role": "system", "content": processed_prompt} 30 | chat_history = [system_message] + chat_history 31 | 32 | if current_state.preprocess_chat: 33 | chat_history = current_state.preprocess_chat(chat_history) 34 | 35 | # Execute LLM call 36 | completion = await async_openai_instance.beta.chat.completions.parse( 37 | model=llm_model, 38 | messages=chat_history, 39 | response_format=response_model, 40 | ) 41 | 42 | message = completion.choices[0].message 43 | if not message.parsed: 44 | raise FSMError(f"Error in parsing the completion: {message.refusal}") 45 | 46 | return message.parsed.model_dump() 47 | 48 | @staticmethod 49 | def process_prompt_template( 50 | prompt_template: str, 51 | context: Dict[str, Any], 52 | preprocess_prompt_template: Optional[Callable] = None, 53 | ) -> str: 54 | """Process the system prompt with Jinja2 templates and optional pre-processing""" 55 | # Pre-process system prompt with Jinja2 56 | template = jinja2.Template(prompt_template) 57 | processed_prompt = template.render(context) 58 | 59 | if preprocess_prompt_template: 60 | processed_prompt = ( 61 | preprocess_prompt_template(processed_prompt) 62 | or processed_prompt 63 | ) 64 | 65 | return processed_prompt 66 | 67 | @staticmethod 68 | def process_chat_history( 69 | chat_history: list, 70 | preprocess_chat: Optional[Callable] = None, 71 | fsm_instance = None, 72 | ) -> list: 73 | """Process chat history with optional pre-processing function""" 74 | if preprocess_chat: 75 | chat_history = preprocess_chat(chat_history, fsm_instance) 76 | return chat_history -------------------------------------------------------------------------------- /content/calculus_example.json: -------------------------------------------------------------------------------- 1 | [ 2 | { 3 | "id": "1", 4 | "content_id": "2", 5 | "content_name": "On Different Degrees of Smallness", 6 | "example": "Let us think of x as a quantity that can grow by a small amount so as to become x+dx, where dx is the small increment added by growth. The square of this is x^2+2x·dx+(dx)^2. The second term is not negligible because it is a first-order quantity; while the third term is of the second order of smallness, being a bit of, a bit of x^2. Thus if we took dx to mean numerically, say, 1/60 of x, then the second term would be 2/60 of x^2, whereas the third term would be 1/3600 of x^2. This last term is clearly less important than the second. But if we go further and take dx to mean only 1/1000 of x, then the second term will be 2/1000 of x^2, while the third term will be only 1/1,000,000 of x^2.\n\nGeometrically this may be depicted as follows: Draw a square (Figure 1) the side of which we will take to represent x. Now suppose the square to grow by having a bit dx added to its size each way. The enlarged square is made up of the original square x^2, the two rectangles at the top and on the right, each of which is of area x·dx (or together 2x·dx), and the little square at the top right-hand corner which is (dx)^2. In Figure 2 we have taken dx as quite a big fraction of x–about 1/5. But suppose we had taken it only 1/100–about the thickness of an inked line drawn with a fine pen. Then the little corner square will have an area of only 1/10,000 of x^2, and be practically invisible. Clearly (dx)^2 is negligible if only we consider the increment dx to be itself small enough." 7 | }, 8 | { 9 | "id": "2", 10 | "content_id": "4", 11 | "content_name": "Simplest Cases", 12 | "example": "Try differentiating y=x³ in the same way. We let y grow to y+dy, while x grows to x+dx. Then we have y+dy=(x+dx)³. Doing the cubing we obtain y+dy=x³+3x²·dx+3x(dx)²+(dx)³. Now we know that we may neglect small quantities of the second and third orders; since, when dy and dx are both made indefinitely small, (dx)² and (dx)³ will become indefinitely smaller by comparison. So, regarding them as negligible, we have left: y+dy=x³+3x²·dx. But y=x³; and, subtracting this, we have: dy and dydx=3x²·dx,=3x²." 13 | }, 14 | { 15 | "id": "3", 16 | "content_id": "4", 17 | "content_name": "Simplest Cases", 18 | "example": "Try differentiating y=x⁴. Starting as before by letting both y and x grow a bit, we have: y+dy=(x+dx)⁴. Working out the raising to the fourth power, we get y+dy=x⁴+4x³dx+6x²(dx)²+4x(dx)³+(dx)⁴. Then striking out the terms containing all the higher powers of dx, as being negligible by comparison, we have y+dy=x⁴+4x³dx. Subtracting the original y=x⁴, we have left dy and dydx=4x³dx,=4x³. Now all these cases are quite easy. Let us collect the results to see if we can infer any general rule. Put them in two columns, the values of y in one and the corresponding values found for dydx in the other: thus y dydx x² 2x x³ 3x² x⁴ 4x³ Just look at these results: the operation of differentiating appears to have had the effect of diminishing the power of x by 1 (for example in the last case reducing x⁴ to x³), and at the same time multiplying by a number (the same number in fact which originally appeared as the power). Now, when you have once seen this, you might easily conjecture how the others will run. You would expect that differentiating x⁵ would give 5x⁴, or differentiating x⁶ would give 6x⁵. If you hesitate, try one of these, and see whether the conjecture comes right." 19 | }, 20 | { 21 | "id": "4", 22 | "content_id": "4", 23 | "content_name": "Simplest Cases", 24 | "example": "" 25 | } 26 | 27 | ] 28 | -------------------------------------------------------------------------------- /examples/support_agent.py: -------------------------------------------------------------------------------- 1 | import os 2 | from dotenv import load_dotenv 3 | from pydantic import BaseModel 4 | import openai 5 | from fsm_llm.fsm import LLMStateMachine 6 | from fsm_llm.state_models import FSMRun, DefaultResponse 7 | 8 | # Load environment variables 9 | load_dotenv() 10 | 11 | # Initialize OpenAI API key 12 | openai.api_key = os.getenv("OPENAI_API_KEY") 13 | openai.organization = os.getenv("OPENAI_ORGANIZATION") 14 | 15 | # Create the FSM 16 | fsm = LLMStateMachine(initial_state="START", end_state="END") 17 | 18 | # Define the response model for user identification 19 | class UserIdentificationResponse(BaseModel): 20 | user_name: str 21 | phone_number: str 22 | 23 | class ConfirmationResponse(BaseModel): 24 | confirmation: str # Expect "yes" or "no" 25 | 26 | # Define the START state 27 | @fsm.define_state( 28 | state_key="START", 29 | prompt_template=( 30 | "You are a customer support bot. Your first task is to ask the user for their " 31 | "name and phone number. Please ensure the user provides both details before proceeding." 32 | ), 33 | response_model=UserIdentificationResponse, 34 | transitions={"CONFIRM": "Once the user provides their name and phone number"}, 35 | ) 36 | async def start_state( 37 | fsm: LLMStateMachine, response: UserIdentificationResponse, will_transition: bool 38 | ): 39 | if will_transition and fsm.get_next_state() == "CONFIRM": 40 | # Store user details in context 41 | fsm.set_context_data( 42 | "verified_user", 43 | {"user_name": response.user_name, "phone_number": response.phone_number}, 44 | ) 45 | return ( 46 | f"Thank you! You provided the following details:\n" 47 | f"Name: {response.user_name}\nPhone Number: {response.phone_number}\n" 48 | f"Is this information correct? (yes/no)" 49 | ) 50 | return "Please provide your name and phone number." 51 | 52 | 53 | 54 | # Define the CONFIRM state 55 | @fsm.define_state( 56 | state_key="CONFIRM", 57 | prompt_template="Please confirm the information you provided. Reply with 'yes' or 'no'.", 58 | response_model=ConfirmationResponse, 59 | transitions={ 60 | "IDENTIFIED": "If the user confirms the details are correct", 61 | "START": "If the user indicates the details are incorrect", 62 | }, 63 | ) 64 | async def confirm_state( 65 | fsm: LLMStateMachine, response: ConfirmationResponse, will_transition: bool 66 | ): 67 | if response.confirmation.lower() == "yes": 68 | fsm.set_next_state("IDENTIFIED") 69 | return "Thank you for confirming your details. How can I help you?" 70 | elif response.confirmation.lower() == "no": 71 | fsm.set_next_state("START") 72 | return "Let's try again. Please provide your name and phone number." 73 | else: 74 | return "Invalid response. Please reply with 'yes' or 'no'." 75 | 76 | 77 | 78 | # Define the IDENTIFIED state 79 | @fsm.define_state( 80 | state_key="IDENTIFIED", 81 | prompt_template=( 82 | "Thank you for identifying yourself. Is there anything else you need help with?" 83 | ), 84 | response_model=DefaultResponse, 85 | transitions={"END": "When the user indicates the conversation is over"}, 86 | ) 87 | async def identified_state(fsm: LLMStateMachine, response: DefaultResponse, will_transition: bool): 88 | if will_transition and fsm.get_next_state() == "END": 89 | return "Thank you! Have a great day!" 90 | return response.content or "You have been identified successfully. How can I assist you further?" 91 | 92 | 93 | 94 | # Define the END state 95 | @fsm.define_state( 96 | state_key="END", 97 | prompt_template="Thank you! Goodbye.", 98 | response_model=DefaultResponse, 99 | ) 100 | async def end_state(fsm: LLMStateMachine, response: DefaultResponse, will_transition: bool): 101 | return "Goodbye! If you need further assistance, feel free to reach out again." 102 | 103 | async def main(): 104 | openai_client = openai.AsyncOpenAI() 105 | 106 | print("Agent: Hello! I am your customer service assistant. Say something to get started.") 107 | while not fsm.is_completed(): # Run until FSM reaches the END state 108 | user_input = input("Your input: ") 109 | if user_input.lower() in ["quit", "exit"]: 110 | fsm.set_next_state("END") 111 | break 112 | run_state: FSMRun = await fsm.run_state_machine(openai_client, user_input=user_input) 113 | print(f"Agent: {run_state.response}") 114 | 115 | print("Agent: Conversation ended. Thank you!") 116 | 117 | 118 | 119 | if __name__ == "__main__": 120 | import asyncio 121 | asyncio.run(main()) 122 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # FSM-based LLM Conversational Agents 2 | [![PyPI version](https://img.shields.io/pypi/v/fsm-llm.svg?style=flat)](https://pypi.org/project/fsm-llm/) 3 | [![PyPI Downloads](https://static.pepy.tech/badge/fsm-llm)](https://pepy.tech/projects/fsm-llm) 4 | 5 | 6 | 7 | This project provides a package framework for creating conversational agents using a Finite State Machine (FSM) powered by Large Language Models (LLMs). It integrates with OpenAI's API and provides an easy way to define states, transitions, and interactions. 8 | 9 | This is currently an experimental setup, and also part of a research project I am doing for university. For now it is meant for developers and experimenters mainly. Requires an OpenAI API key (currently tested on gpt-4o-mini and gpt-4o). 10 | 11 | I intend to continue working on this in the (hopefully) near future! If there are bugs please create an issue and I will try my best to provide critical updates. 12 | 13 | Package last updated: Dec 2024 14 | 15 | ## Features 16 | 17 | - Define states and transitions for your agent using a simple decorator. 18 | - Handle dynamic conversation flow with flexible state management. 19 | - Integrates with GPT models to generate responses based on state context. 20 | 21 | 22 | ## Installation 23 | 24 | 1. Install the package directly from PyPI: 25 | ``` 26 | pip install fsm-llm 27 | ``` 28 | 29 | 30 | 2. Set up environment variables: 31 | Create a `.env` file and add your OpenAI API key: 32 | ``` 33 | OPENAI_API_KEY=your-api-key 34 | OPENAI_ORGANIZATION=your-organization-id 35 | ``` 36 | 37 | ## Usage Example (On/Off Switch) 38 | 39 | ### 0. What is an FSM? 40 | A Finite State Machine (FSM) is a computational model used to design systems that can exist in one of a finite number of states at any given time. The system transitions from one state to another based on specific conditions or inputs. 41 | 42 | ![on and off fsm](https://github.com/user-attachments/assets/a804de3d-47c4-4b02-a461-4f95340eab9f) 43 | 44 | In the example ```examples/switch_agent.py```, the FSM transitions between two states (OFF and ON) based on the input (user pressing the switch). If the user presses the switch while the light is already in the desired state (e.g., turning the light off when it’s already off), the FSM remains in the same state but can trigger additional responses. 45 | 46 | 47 | 48 | 49 | ### 1. **Creating the FSM (Finite State Machine)** 50 | The `LLMStateMachine` class is the core of the framework. It handles state transitions based on user input. 51 | ```python 52 | from fsm_llm import LLMStateMachine 53 | 54 | # Create the FSM 55 | fsm = LLMStateMachine(initial_state="START", end_state="END") 56 | ``` 57 | 58 | 59 | ### 2. **Defining States with the `@fsm.define_state` Decorator** 60 | 61 | The main feature of this framework is the ability to define states using the `@fsm.define_state` decorator. Each state has a unique key, a prompt that will be used for the LLM, and possible transitions to other states based on user input. 62 | 63 | #### `START` State 64 | 65 | ```python 66 | @fsm.define_state( 67 | state_key="START", 68 | prompt_template="You are an on-off switcher. Ask the user if they want to turn the switch on or off.", 69 | transitions={"STATE_ON": "If user wants to turn on the switch", "END": "If user wants to end the conversation"}, 70 | ) 71 | async def start_state(fsm: LLMStateMachine, response: str, will_transition: bool): 72 | global SWITCH_STATE 73 | if will_transition and fsm.get_next_state() == "STATE_ON": 74 | SWITCH_STATE = "ON" 75 | print("SWITCH TURNED ON") 76 | elif will_transition and fsm.get_next_state() == "END": 77 | return "Goodbye!" 78 | return response 79 | ``` 80 | 81 | - **`state_key="START"`**: The name of the state in the FSM. 82 | - **`prompt_template`**: The message that will be used to prompt the LLM. In this case, the agent asks the user if they want to turn the switch on or off. 83 | - **`transitions`**: This defines what happens next based on user input. For example: 84 | - If the user wants to turn the switch on, the FSM transitions to the `STATE_ON` state. 85 | - If the user wants to end the conversation, the FSM will transition to the `END` state. 86 | 87 | Inside the function `start_state`, we check whether the FSM will transition to the `STATE_ON` or `END` state. If the transition happens, we update the `SWITCH_STATE` to `"ON"`. 88 | 89 | #### `STATE_ON` State 90 | 91 | ```python 92 | @fsm.define_state( 93 | state_key="STATE_ON", 94 | prompt_template="The switch is now on. Ask the user if they want to turn off the switch or end the conversation.", 95 | transitions={"START": "If user wants to turn off the switch", "END": "If user wants to end the conversation"}, 96 | ) 97 | async def state_on(fsm: LLMStateMachine, response: str, will_transition: bool): 98 | global SWITCH_STATE 99 | if will_transition and fsm.get_next_state() == "START": 100 | SWITCH_STATE = "OFF" 101 | print("SWITCH TURNED OFF") 102 | elif will_transition and fsm.get_next_state() == "END": 103 | return "Goodbye!" 104 | return response 105 | ``` 106 | 107 | The logic inside `state_on` checks the transition. If the FSM is transitioning back to the `START` state, it sets the `SWITCH_STATE` to `"OFF"`. 108 | 109 | #### `END` State 110 | 111 | ```python 112 | @fsm.define_state( 113 | state_key="END", 114 | prompt_template="Goodbye!", 115 | ) 116 | async def end_state(fsm: LLMStateMachine, response: str, will_transition: bool): 117 | return "Goodbye!" 118 | ``` 119 | 120 | 121 | ### 3. **Running the Agent** 122 | 123 | To run the FSM-based agent, we use an asynchronous loop to interact with the user and process their input. 124 | ```python 125 | from fsm_llm.state_models import FSMRun 126 | ``` 127 | 128 | ```python 129 | async def main(): 130 | """Example of a simple on-off switch FSM using LLMStateMachine""" 131 | # Create the OpenAI client 132 | openai_client = openai.AsyncOpenAI() 133 | 134 | print("Agent: Hi. I am an on-off switch manager.") 135 | while not fsm.is_completed(): # Run until FSM reaches the END state 136 | user_input = input("Your input: ") 137 | if user_input.lower() in ["quit", "exit"]: 138 | fsm.set_next_state("END") 139 | break 140 | run_state: FSMRun = await fsm.run_state_machine(openai_client, user_input=user_input) 141 | print(f"Agent: {run_state.response}") 142 | print("CURRENT SWITCH STATE:", SWITCH_STATE) 143 | 144 | print("Agent: Goodbye.") 145 | ``` 146 | 147 | - **`while not fsm.is_completed()`**: The loop continues running until the FSM reaches the `END` state. 148 | - **`user_input`**: The user provides input, which the FSM processes. 149 | - **`fsm.run_state_machine`**: This method processes the current state and transitions based on the user's input. The OpenAI client is used to get the response. 150 | - **`SWITCH_STATE`**: After each interaction, the current state of the switch (on or off) is printed. 151 | 152 | 153 | 154 | 155 | ## Examples 156 | 157 | - **Light Switch Agent**: A simple agent that asks the user whether they want to turn a switch on or off. ```switch_agent.py``` 158 | - **Customer Support Agent**: A bot that collects user details and assists with customer queries. ```support_agent.py``` 159 | - **Medical Triage Agent**: A complex agent that helps assess if a medical situation is an emergency and collects patient data. ```medical_agent.py``` 160 | 161 | 163 | -------------------------------------------------------------------------------- /examples/tutor_agent.py: -------------------------------------------------------------------------------- 1 | import os 2 | from dotenv import load_dotenv 3 | import openai 4 | from fsm_llm import LLMStateMachine 5 | from fsm_llm.state_models import FSMRun 6 | import json 7 | import jinja2 8 | 9 | # Load environment variables 10 | load_dotenv() 11 | 12 | # Initialize OpenAI API key 13 | openai.api_key = os.getenv("OPENAI_API_KEY") 14 | openai.organization = os.getenv("OPENAI_ORGANIZATION") 15 | 16 | # Initialize the FSM 17 | fsm = LLMStateMachine(initial_state="show_content", end_state="END") 18 | 19 | # Global variables to track the learning state 20 | LEARNING_STATE = { 21 | "current_content_id": 1, # Tracks the ID of the content the user is currently on 22 | } 23 | 24 | # Keep track of last user input globally 25 | LAST_USER_INPUT = "" 26 | 27 | # Actions 28 | USER_ACTIONS = ["ua_next", "ua_ask_clarifying_content", "ua_ask_clarifying_example"] 29 | SYSTEM_ACTIONS = ["sa_show_content", "sa_show_example", "sa_show_quiz"] 30 | 31 | CONTENT_FILE = "content/calculus_content.json" 32 | # Function to load content dynamically from a file 33 | def load_content(content_id, file_path=CONTENT_FILE): 34 | try: 35 | with open(file_path, "r", encoding="utf-8") as file: 36 | content_data = json.load(file) # Assumes content is stored as JSON 37 | 38 | # Loop through the list to find the matching content ID 39 | for item in content_data: 40 | if item.get("id") == str(content_id): 41 | return item.get("content", "Content field not found.") # Return the 'content' field 42 | 43 | return "Content not found." # Default if ID not found 44 | except json.JSONDecodeError: 45 | return "Error: Failed to decode JSON." 46 | except FileNotFoundError: 47 | return f"Error: File '{file_path}' not found." 48 | except Exception as e: 49 | return f"Error loading content: {e}" 50 | 51 | 52 | 53 | EXAMPLE_FILE = "content/calculus_example.json" 54 | # Function to load content dynamically from a file 55 | def load_example(content_id, file_path=EXAMPLE_FILE): 56 | try: 57 | with open(file_path, "r", encoding="utf-8") as file: 58 | example_data = json.load(file) # Assumes content is stored as JSON 59 | 60 | # Loop through the list to find the matching content ID 61 | for item in example_data: 62 | if item.get("content_id") == str(content_id): 63 | return item.get("example", "example field not found.") # Return the 'content' field 64 | 65 | return "Example not found." # Default if ID not found 66 | except json.JSONDecodeError: 67 | return "Error: Failed to decode JSON." 68 | except FileNotFoundError: 69 | return f"Error: File '{file_path}' not found." 70 | except Exception as e: 71 | return f"Error loading content: {e}" 72 | 73 | SHOW_CONTENT_TEMPLATE = """ 74 | You are a friendly and helpful calculus tutor. 75 | The user said: "{{ user_input }}" 76 | 77 | Current Topic ID: {{ topic_id }} 78 | Content for this topic: 79 | {{ topic_content }} 80 | 81 | Explain this content in a helpful way. If the user wants more content, you can move to show_content. 82 | If they want an example, move to show_example. 83 | If they want a quiz, move to quiz. 84 | If they want to end, move to END. 85 | 86 | Include the content above in your explanation to the user. 87 | """ 88 | 89 | SHOW_EXAMPLE_TEMPLATE = """ 90 | You are a friendly and helpful calculus tutor. 91 | The user said: "{{ user_input }}" 92 | 93 | Current Topic ID: {{ topic_id }} 94 | Previously shown content: 95 | {{ topic_content }} 96 | 97 | Example for this topic: 98 | {{ topic_example }} 99 | 100 | Explain the example and how it relates to the content. If the user wants more content, move to show_content. 101 | If they want another example, move to show_example. 102 | If they want a quiz, move to quiz. 103 | If they want to end, move to END. 104 | """ 105 | 106 | QUIZ_TEMPLATE = """ 107 | You are a friendly and helpful calculus tutor. 108 | The user said: "{{ user_input }}" 109 | 110 | Current Topic ID: {{ topic_id }} 111 | Previously shown content: 112 | {{ topic_content }} 113 | 114 | Please create a short quiz related to the above content. Include a few questions and maybe some hints. 115 | If the user wants more content after this, move to show_content. 116 | If they want an example, move to show_example. 117 | If they want another quiz, move to quiz. 118 | If they want to end, move to END. 119 | """ 120 | 121 | END_TEMPLATE = "The learning session has concluded. Goodbye!" 122 | 123 | def preprocess_prompt_template(processed_prompt: str) -> str: 124 | """Dynamically fill in user input and content/example before sending to LLM.""" 125 | topic_id = LEARNING_STATE["current_content_id"] 126 | topic_content = load_content(topic_id, CONTENT_FILE) 127 | topic_example = load_example(topic_id, EXAMPLE_FILE) 128 | user_input = LAST_USER_INPUT 129 | 130 | # Use Jinja2 to render the template dynamically 131 | template = jinja2.Template(processed_prompt) 132 | rendered = template.render( 133 | user_input=user_input, 134 | topic_id=topic_id, 135 | topic_content=topic_content, 136 | topic_example=topic_example 137 | ) 138 | return rendered 139 | 140 | # Define the `show_content` state 141 | @fsm.define_state( 142 | state_key="show_content", 143 | prompt_template=SHOW_CONTENT_TEMPLATE, 144 | transitions={ 145 | "show_content": "If the user wants to move to the next section.", 146 | "show_example": "If the user asks for an example.", 147 | "quiz": "If the user asks for a quiz.", 148 | "END": "If the user wants to end the session." 149 | }, 150 | preprocess_prompt_template=preprocess_prompt_template 151 | ) 152 | async def show_content_state(fsm: LLMStateMachine, response: str, will_transition: bool): 153 | # If we are going to show_content again, increment the content_id 154 | if will_transition and fsm.get_next_state() == "show_content": 155 | LEARNING_STATE['current_content_id'] += 1 156 | # Return the LLM's response directly, which should now contain the content 157 | return response 158 | 159 | # Define the `show_example` state 160 | @fsm.define_state( 161 | state_key="show_example", 162 | prompt_template=SHOW_EXAMPLE_TEMPLATE, 163 | transitions={ 164 | "show_content": "If the user asks for more content.", 165 | "show_example": "If the user asks for another example.", 166 | "quiz": "If the user asks for a quiz.", 167 | "END": "If the user wants to end the session." 168 | }, 169 | preprocess_prompt_template=preprocess_prompt_template 170 | ) 171 | async def show_example_state(fsm: LLMStateMachine, response: str, will_transition: bool): 172 | if will_transition and fsm.get_next_state() == "show_content": 173 | LEARNING_STATE['current_content_id'] += 1 174 | return response 175 | 176 | 177 | # Define the `quiz` state 178 | @fsm.define_state( 179 | state_key="quiz", 180 | prompt_template=QUIZ_TEMPLATE, 181 | transitions={ 182 | "show_content": "If the user asks for more content.", 183 | "show_example": "If the user asks for another example.", 184 | "quiz": "If the user wants another quiz.", 185 | "END": "If the user wants to end the session." 186 | }, 187 | preprocess_prompt_template=preprocess_prompt_template 188 | ) 189 | async def quiz_state(fsm: LLMStateMachine, response: str, will_transition: bool): 190 | # If transitioning to show_content, increment 191 | if will_transition and fsm.get_next_state() == "show_content": 192 | LEARNING_STATE['current_content_id'] += 1 193 | return response 194 | 195 | # Define the END state 196 | @fsm.define_state( 197 | state_key="END", 198 | prompt_template=END_TEMPLATE 199 | ) 200 | async def end_state(fsm: LLMStateMachine, response: str): 201 | return "Thank you for learning! Goodbye!" 202 | 203 | 204 | # Simulated interaction loop 205 | async def main(): 206 | """Simulates a learning session with the tutor agent.""" 207 | import random 208 | 209 | # Create the OpenAI client 210 | openai_client = openai.AsyncOpenAI() 211 | 212 | print("Tutor: Welcome to the learning session!") 213 | while not fsm.is_completed(): 214 | user_input = input("User: ") 215 | if user_input.lower() in ["quit", "exit"]: 216 | fsm.set_next_state("END") 217 | break 218 | 219 | run_state: FSMRun = await fsm.run_state_machine(openai_client, user_input=user_input) 220 | print(f"Tutor: {run_state.response}") 221 | print("CURRENT CONTENT ID:", LEARNING_STATE["current_content_id"]) 222 | print("Agent: Goodbye.") 223 | 224 | 225 | if __name__ == "__main__": 226 | import asyncio 227 | asyncio.run(main()) 228 | -------------------------------------------------------------------------------- /examples/tutor_agent_modified.py: -------------------------------------------------------------------------------- 1 | import os 2 | from dotenv import load_dotenv 3 | import openai 4 | from fsm_llm import LLMStateMachine 5 | from fsm_llm.state_models import FSMRun 6 | import json 7 | import jinja2 8 | 9 | # Load environment variables 10 | load_dotenv() 11 | 12 | # Initialize OpenAI API key 13 | openai.api_key = os.getenv("OPENAI_API_KEY") 14 | openai.organization = os.getenv("OPENAI_ORGANIZATION") 15 | 16 | # Initialize the FSM 17 | fsm = LLMStateMachine(initial_state="show_content", end_state="END") 18 | 19 | # Global variables to track the learning state 20 | LEARNING_STATE = { 21 | "current_content_id": 1, # Tracks the ID of the content the user is currently on 22 | } 23 | 24 | # Keep track of last user input globally 25 | LAST_USER_INPUT = "" 26 | 27 | # Actions 28 | USER_ACTIONS = ["ua_next", "ua_ask_clarifying_content", "ua_ask_clarifying_example"] 29 | SYSTEM_ACTIONS = ["sa_show_content", "sa_show_example", "sa_show_quiz"] 30 | 31 | CONTENT_FILE = "content/calculus_content.json" 32 | EXAMPLE_FILE = "content/calculus_example.json" 33 | 34 | def load_content(content_id, file_path=CONTENT_FILE): 35 | try: 36 | with open(file_path, "r", encoding="utf-8") as file: 37 | content_data = json.load(file) 38 | for item in content_data: 39 | if item.get("id") == str(content_id): 40 | return item.get("content", "Content field not found.") 41 | return "Content not found." 42 | except json.JSONDecodeError: 43 | return "Error: Failed to decode JSON." 44 | except FileNotFoundError: 45 | return f"Error: File '{file_path}' not found." 46 | except Exception as e: 47 | return f"Error loading content: {e}" 48 | 49 | def load_example(content_id, file_path=EXAMPLE_FILE): 50 | try: 51 | with open(file_path, "r", encoding="utf-8") as file: 52 | example_data = json.load(file) 53 | for item in example_data: 54 | if item.get("content_id") == str(content_id): 55 | return item.get("example", "Example field not found.") 56 | return "Example not found." 57 | except json.JSONDecodeError: 58 | return "Error: Failed to decode JSON." 59 | except FileNotFoundError: 60 | return f"Error: File '{file_path}' not found." 61 | except Exception as e: 62 | return f"Error loading example: {e}" 63 | 64 | SHOW_CONTENT_TEMPLATE = """ 65 | You are a friendly and helpful calculus tutor. 66 | The user said: "{{ user_input }}" 67 | 68 | Current Topic ID: {{ topic_id }} 69 | Content for this topic: 70 | {{ topic_content }} 71 | 72 | Explain this content in a helpful way. If the user wants more content, you can move to show_content. 73 | If they want an example, move to show_example. 74 | If they want a quiz, move to quiz. 75 | If they want to end, move to END. 76 | 77 | Include the content above in your explanation to the user. 78 | """ 79 | 80 | SHOW_EXAMPLE_TEMPLATE = """ 81 | You are a friendly and helpful calculus tutor. 82 | The user said: "{{ user_input }}" 83 | 84 | Current Topic ID: {{ topic_id }} 85 | Previously shown content: 86 | {{ topic_content }} 87 | 88 | Example for this topic: 89 | {{ topic_example }} 90 | 91 | Explain the example and how it relates to the content. If the user wants more content, move to show_content. 92 | If they want another example, move to show_example. 93 | If they want a quiz, move to quiz. 94 | If they want to end, move to END. 95 | """ 96 | 97 | QUIZ_TEMPLATE = """ 98 | You are a friendly and helpful calculus tutor. 99 | The user said: "{{ user_input }}" 100 | 101 | Current Topic ID: {{ topic_id }} 102 | Previously shown content: 103 | {{ topic_content }} 104 | 105 | Please create a short quiz related to the above content. Include a few questions and maybe some hints. 106 | If the user wants more content after this, move to show_content. 107 | If they want an example, move to show_example. 108 | If they want another quiz, move to quiz. 109 | If they want to end, move to END. 110 | """ 111 | 112 | END_TEMPLATE = "The learning session has concluded. Goodbye!" 113 | 114 | def preprocess_prompt_template(processed_prompt: str) -> str: 115 | """Dynamically fill in user input and content/example before sending to LLM.""" 116 | topic_id = LEARNING_STATE["current_content_id"] 117 | topic_content = load_content(topic_id, CONTENT_FILE) 118 | topic_example = load_example(topic_id, EXAMPLE_FILE) 119 | user_input = LAST_USER_INPUT 120 | 121 | # Use Jinja2 to render the template dynamically 122 | template = jinja2.Template(processed_prompt) 123 | rendered = template.render( 124 | user_input=user_input, 125 | topic_id=topic_id, 126 | topic_content=topic_content, 127 | topic_example=topic_example 128 | ) 129 | return rendered 130 | 131 | # Define the `show_content` state 132 | @fsm.define_state( 133 | state_key="show_content", 134 | prompt_template=SHOW_CONTENT_TEMPLATE, 135 | transitions={ 136 | "show_content": "If the user wants to move to the next section.", 137 | "show_example": "If the user asks for an example.", 138 | "quiz": "If the user asks for a quiz.", 139 | "END": "If the user wants to end the session." 140 | }, 141 | preprocess_prompt_template=preprocess_prompt_template 142 | ) 143 | async def show_content_state(fsm: LLMStateMachine, response: str, will_transition: bool): 144 | # If we are going to show_content again, increment the content_id 145 | if will_transition and fsm.get_next_state() == "show_content": 146 | LEARNING_STATE['current_content_id'] += 1 147 | # Return the LLM's response directly, which should now contain the content 148 | return response 149 | 150 | # Define the `show_example` state 151 | @fsm.define_state( 152 | state_key="show_example", 153 | prompt_template=SHOW_EXAMPLE_TEMPLATE, 154 | transitions={ 155 | "show_content": "If the user asks for more content.", 156 | "show_example": "If the user asks for another example.", 157 | "quiz": "If the user asks for a quiz.", 158 | "END": "If the user wants to end the session." 159 | }, 160 | preprocess_prompt_template=preprocess_prompt_template 161 | ) 162 | async def show_example_state(fsm: LLMStateMachine, response: str, will_transition: bool): 163 | if will_transition and fsm.get_next_state() == "show_content": 164 | LEARNING_STATE['current_content_id'] += 1 165 | return response 166 | 167 | # Define the `quiz` state 168 | @fsm.define_state( 169 | state_key="quiz", 170 | prompt_template=QUIZ_TEMPLATE, 171 | transitions={ 172 | "show_content": "If the user asks for more content.", 173 | "show_example": "If the user asks for another example.", 174 | "quiz": "If the user wants another quiz.", 175 | "END": "If the user wants to end the session." 176 | }, 177 | preprocess_prompt_template=preprocess_prompt_template 178 | ) 179 | async def quiz_state(fsm: LLMStateMachine, response: str, will_transition: bool): 180 | # If transitioning to show_content, increment 181 | if will_transition and fsm.get_next_state() == "show_content": 182 | LEARNING_STATE['current_content_id'] += 1 183 | return response 184 | 185 | # Define the END state 186 | @fsm.define_state( 187 | state_key="END", 188 | prompt_template=END_TEMPLATE 189 | ) 190 | async def end_state(fsm: LLMStateMachine, response: str): 191 | return "Thank you for learning! Goodbye!" 192 | 193 | # Simulated interaction loop 194 | async def main(): 195 | """Simulates a learning session with the tutor agent.""" 196 | import random 197 | openai_client = openai.AsyncOpenAI() 198 | 199 | print("Tutor: Welcome to the learning session!") 200 | while not fsm.is_completed(): 201 | user_input = input("User: ") 202 | if user_input.lower() in ["quit", "exit"]: 203 | fsm.set_next_state("END") 204 | break 205 | 206 | global LAST_USER_INPUT 207 | LAST_USER_INPUT = user_input 208 | 209 | # Simulate user and system actions 210 | user_action = random.choice(USER_ACTIONS) if user_input.strip() == "" else user_input 211 | if user_action in USER_ACTIONS: 212 | print(f"[User Action Triggered: {user_action}]") 213 | if user_action == "ua_next": 214 | fsm.set_next_state("show_content") 215 | elif user_action == "ua_ask_clarifying_content": 216 | fsm.set_next_state("show_content") 217 | elif user_action == "ua_ask_clarifying_example": 218 | fsm.set_next_state("show_example") 219 | else: 220 | system_action = random.choice(SYSTEM_ACTIONS) 221 | print(f"[System Action Triggered: {system_action}]") 222 | if system_action == "sa_show_content": 223 | fsm.set_next_state("show_content") 224 | elif system_action == "sa_show_example": 225 | fsm.set_next_state("show_example") 226 | elif system_action == "sa_show_quiz": 227 | fsm.set_next_state("quiz") 228 | 229 | run_state: FSMRun = await fsm.run_state_machine(openai_client, user_input=user_input) 230 | print(f"Tutor: {run_state.response}") 231 | print("CURRENT CONTENT ID:", LEARNING_STATE["current_content_id"]) 232 | 233 | if __name__ == "__main__": 234 | import asyncio 235 | asyncio.run(main()) 236 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 5 | 1. Definitions. 6 | "License" shall mean the terms and conditions for use, reproduction, 7 | and distribution as defined by Sections 1 through 9 of this document. 8 | "Licensor" shall mean the copyright owner or entity authorized by 9 | the copyright owner that is granting the License. 10 | "Legal Entity" shall mean the union of the acting entity and all 11 | other entities that control, are controlled by, or are under common 12 | control with that entity. For the purposes of this definition, 13 | "control" means (i) the power, direct or indirect, to cause the 14 | direction or management of such entity, whether by contract or 15 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 16 | outstanding shares, or (iii) beneficial ownership of such entity. 17 | "You" (or "Your") shall mean an individual or Legal Entity 18 | exercising permissions granted by this License. 19 | "Source" form shall mean the preferred form for making modifications, 20 | including but not limited to software source code, documentation 21 | source, and configuration files. 22 | "Object" form shall mean any form resulting from mechanical 23 | transformation or translation of a Source form, including but 24 | not limited to compiled object code, generated documentation, 25 | and conversions to other media types. 26 | "Work" shall mean the work of authorship, whether in Source or 27 | Object form, made available under the License, as indicated by a 28 | copyright notice that is included in or attached to the work 29 | (an example is provided in the Appendix below). 30 | "Derivative Works" shall mean any work, whether in Source or Object 31 | form, that is based on (or derived from) the Work and for which the 32 | editorial revisions, annotations, elaborations, or other modifications 33 | represent, as a whole, an original work of authorship. For the purposes 34 | of this License, Derivative Works shall not include works that remain 35 | separable from, or merely link (or bind by name) to the interfaces of, 36 | the Work and Derivative Works thereof. 37 | "Contribution" shall mean any work of authorship, including 38 | the original version of the Work and any modifications or additions 39 | to that Work or Derivative Works thereof, that is intentionally 40 | submitted to Licensor for inclusion in the Work by the copyright owner 41 | or by an individual or Legal Entity authorized to submit on behalf of 42 | the copyright owner. For the purposes of this definition, "submitted" 43 | means any form of electronic, verbal, or written communication sent 44 | to the Licensor or its representatives, including but not limited to 45 | communication on electronic mailing lists, source code control systems, 46 | and issue tracking systems that are managed by, or on behalf of, the 47 | Licensor for the purpose of discussing and improving the Work, but 48 | excluding communication that is conspicuously marked or otherwise 49 | designated in writing by the copyright owner as "Not a Contribution." 50 | "Contributor" shall mean Licensor and any individual or Legal Entity 51 | on behalf of whom a Contribution has been received by Licensor and 52 | subsequently incorporated within the Work. 53 | 2. Grant of Copyright License. Subject to the terms and conditions of 54 | this License, each Contributor hereby grants to You a perpetual, 55 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 56 | copyright license to reproduce, prepare Derivative Works of, 57 | publicly display, publicly perform, sublicense, and distribute the 58 | Work and such Derivative Works in Source or Object form. 59 | 3. Grant of Patent License. Subject to the terms and conditions of 60 | this License, each Contributor hereby grants to You a perpetual, 61 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 62 | (except as stated in this section) patent license to make, have made, 63 | use, offer to sell, sell, import, and otherwise transfer the Work, 64 | where such license applies only to those patent claims licensable 65 | by such Contributor that are necessarily infringed by their 66 | Contribution(s) alone or by combination of their Contribution(s) 67 | with the Work to which such Contribution(s) was submitted. If You 68 | institute patent litigation against any entity (including a 69 | cross-claim or counterclaim in a lawsuit) alleging that the Work 70 | or a Contribution incorporated within the Work constitutes direct 71 | or contributory patent infringement, then any patent licenses 72 | granted to You under this License for that Work shall terminate 73 | as of the date such litigation is filed. 74 | 4. Redistribution. You may reproduce and distribute copies of the 75 | Work or Derivative Works thereof in any medium, with or without 76 | modifications, and in Source or Object form, provided that You 77 | meet the following conditions: 78 | (a) You must give any other recipients of the Work or 79 | Derivative Works a copy of this License; and 80 | (b) You must cause any modified files to carry prominent notices 81 | stating that You changed the files; and 82 | (c) You must retain, in the Source form of any Derivative Works 83 | that You distribute, all copyright, patent, trademark, and 84 | attribution notices from the Source form of the Work, 85 | excluding those notices that do not pertain to any part of 86 | the Derivative Works; and 87 | (d) If the Work includes a "NOTICE" text file as part of its 88 | distribution, then any Derivative Works that You distribute must 89 | include a readable copy of the attribution notices contained 90 | within such NOTICE file, excluding those notices that do not 91 | pertain to any part of the Derivative Works, in at least one 92 | of the following places: within a NOTICE text file distributed 93 | as part of the Derivative Works; within the Source form or 94 | documentation, if provided along with the Derivative Works; or, 95 | within a display generated by the Derivative Works, if and 96 | wherever such third-party notices normally appear. The contents 97 | of the NOTICE file are for informational purposes only and 98 | do not modify the License. You may add Your own attribution 99 | notices within Derivative Works that You distribute, alongside 100 | or as an addendum to the NOTICE text from the Work, provided 101 | that such additional attribution notices cannot be construed 102 | as modifying the License. 103 | You may add Your own copyright statement to Your modifications and 104 | may provide additional or different license terms and conditions 105 | for use, reproduction, or distribution of Your modifications, or 106 | for any such Derivative Works as a whole, provided Your use, 107 | reproduction, and distribution of the Work otherwise complies with 108 | the conditions stated in this License. 109 | 5. Submission of Contributions. Unless You explicitly state otherwise, 110 | any Contribution intentionally submitted for inclusion in the Work 111 | by You to the Licensor shall be under the terms and conditions of 112 | this License, without any additional terms or conditions. 113 | Notwithstanding the above, nothing herein shall supersede or modify 114 | the terms of any separate license agreement you may have executed 115 | with Licensor regarding such Contributions. 116 | 6. Trademarks. This License does not grant permission to use the trade 117 | names, trademarks, service marks, or product names of the Licensor, 118 | except as required for reasonable and customary use in describing the 119 | origin of the Work and reproducing the content of the NOTICE file. 120 | 7. Disclaimer of Warranty. Unless required by applicable law or 121 | agreed to in writing, Licensor provides the Work (and each 122 | Contributor provides its Contributions) on an "AS IS" BASIS, 123 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 124 | implied, including, without limitation, any warranties or conditions 125 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 126 | PARTICULAR PURPOSE. You are solely responsible for determining the 127 | appropriateness of using or redistributing the Work and assume any 128 | risks associated with Your exercise of permissions under this License. 129 | 8. Limitation of Liability. In no event and under no legal theory, 130 | whether in tort (including negligence), contract, or otherwise, 131 | unless required by applicable law (such as deliberate and grossly 132 | negligent acts) or agreed to in writing, shall any Contributor be 133 | liable to You for damages, including any direct, indirect, special, 134 | incidental, or consequential damages of any character arising as a 135 | result of this License or out of the use or inability to use the 136 | Work (including but not limited to damages for loss of goodwill, 137 | work stoppage, computer failure or malfunction, or any and all 138 | other commercial damages or losses), even if such Contributor 139 | has been advised of the possibility of such damages. 140 | 9. Accepting Warranty or Additional Liability. While redistributing 141 | the Work or Derivative Works thereof, You may choose to offer, 142 | and charge a fee for, acceptance of support, warranty, indemnity, 143 | or other liability obligations and/or rights consistent with this 144 | License. However, in accepting such obligations, You may act only 145 | on Your own behalf and on Your sole responsibility, not on behalf 146 | of any other Contributor, and only if You agree to indemnify, 147 | defend, and hold each Contributor harmless for any liability 148 | incurred by, or claims asserted against, such Contributor by reason 149 | of your accepting any such warranty or additional liability. 150 | END OF TERMS AND CONDITIONS 151 | APPENDIX: How to apply the Apache License to your work. 152 | To apply the Apache License to your work, attach the following 153 | boilerplate notice, with the fields enclosed by brackets "[]" 154 | replaced with your own identifying information. (Don't include 155 | the brackets!) The text should be enclosed in the appropriate 156 | comment syntax for the file format. We also recommend that a 157 | file or class name and description of purpose be included on the 158 | same "printed page" as the copyright notice for easier 159 | identification within third-party archives. 160 | Copyright 2024 Jeffrey Zhou 161 | Licensed under the Apache License, Version 2.0 (the "License"); 162 | you may not use this file except in compliance with the License. 163 | You may obtain a copy of the License at 164 | http://www.apache.org/licenses/LICENSE-2.0 165 | Unless required by applicable law or agreed to in writing, software 166 | distributed under the License is distributed on an "AS IS" BASIS, 167 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 168 | See the License for the specific language governing permissions and 169 | limitations under the License. 170 | -------------------------------------------------------------------------------- /fsm_llm/fsm.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from functools import wraps 3 | from typing import Any, Callable, Dict, Optional, Type 4 | import openai 5 | from pydantic import BaseModel, ValidationError 6 | 7 | from fsm_llm.utils import _add_transitions, _generate_response_schema 8 | from fsm_llm.llm_handler import LLMUtilities # Updated import 9 | 10 | from fsm_llm.state_models import ( 11 | FSMRun, 12 | FSMState, 13 | DefaultResponse, 14 | FSMError, 15 | ImmediateStateChange, 16 | ) 17 | 18 | logger = logging.getLogger("llmstatemachine") 19 | 20 | class LLMStateMachine: 21 | """ 22 | Finite State Machine for LLM-driven Agents. 23 | This class enables the creation of conversational agents by defining states, transitions using structured responses from LLMs. 24 | 25 | Parameters: 26 | - _state: The current active state of the FSM, representing the ongoing context of the conversation. 27 | - _initial_state: The starting state of the FSM, used for initialization and resetting. 28 | - _next_state: The next state to transition to, determined dynamically during execution. 29 | - _end_state: The terminal state of the FSM, where processing ends (default is "END"). 30 | - _state_registry: A dictionary that stores metadata about all defined states and their transitions. 31 | - _session_history: A list that records the current session's concise user and assistant interactions. 32 | - _full_session_history: A comprehensive log of all interactions and state transitions in the session. 33 | - _running_session_history: A temporary, live record of the conversation for intermediate processing. 34 | - user_defined_context: A dictionary for storing custom session-specific or user-specific context data. 35 | - _get_completion: A callable for fetching responses from the LLM, defaulting to `default_get_completion`. 36 | """ 37 | 38 | def __init__(self, initial_state: str, end_state: str = "END"): 39 | self._state = initial_state 40 | self._initial_state = initial_state 41 | self._next_state = None 42 | self._end_state = end_state 43 | self._state_registry = {} 44 | self._session_history = [] 45 | self._full_session_history = [] 46 | self._running_session_history = [] 47 | self.user_defined_context = {} 48 | self._llm_utils = LLMUtilities() 49 | 50 | def define_state( 51 | self, 52 | state_key: str, 53 | prompt_template: str, 54 | preprocess_prompt_template: Optional[Callable] = None, 55 | temperature: float = 0.5, 56 | transitions: Dict[str, str] = None, 57 | response_model: Optional[BaseModel] = None, 58 | preprocess_input: Optional[Callable] = None, 59 | preprocess_chat: Optional[Callable] = None, 60 | ): 61 | """ 62 | Decorator to define and register a state [@fsm.define_state(...)] in the FSM (Finite State Machine). 63 | 64 | This function simplifies the process of associating metadata (such as prompts and transitions) 65 | with a Python function that defines the behavior of the state. 66 | 67 | Parameters: 68 | - state_key (str): A unique identifier for the state. 69 | - prompt_template (str): Instructions provided to the LLM when this state is active. 70 | - preprocess_prompt_template (Optional[Callable]): A func to preprocess the sys prompt for the LLM. 71 | - temperature (float): Determines the randomness of LLM responses, defaults at 0.5. 72 | - transitions (Dict[str, str], optional): Maps possible next states to their conditions 73 | (e.g., {"NEXT_STATE": "If user agrees"}). Defaults to None. 74 | - response_model (Optional[BaseModel]): A Pydantic model for parsing and validating the LLM's response. 75 | - preprocess_input (Optional[Callable]): A func to preprocess user input before sending to the LLM. 76 | - preprocess_chat (Optional[Callable]): A func to preprocess the chat history for the LLM. 77 | 78 | Returns: 79 | - callable The original function wrapped and registered with the FSM. 80 | """ 81 | 82 | # Empty dictionary for transitions if none is provided 83 | if transitions is None: 84 | transitions = {} 85 | 86 | def decorator(func: Callable): 87 | @wraps(func) 88 | async def wrapper(*args, **kwargs): 89 | return await func(*args, **kwargs) 90 | 91 | # Register the state in the FSM's registry with the provided metadata 92 | self._state_registry[state_key] = FSMState( 93 | key=state_key, 94 | func=wrapper, 95 | prompt_template=prompt_template, 96 | preprocess_prompt_template=preprocess_prompt_template, 97 | temperature=temperature, 98 | transitions=transitions, 99 | response_model=response_model, 100 | preprocess_input=preprocess_input, 101 | preprocess_chat=preprocess_chat, 102 | ) 103 | return wrapper 104 | return decorator 105 | 106 | async def run_state_machine( 107 | self, 108 | async_openai_instance: openai.AsyncOpenAI, 109 | user_input: str, 110 | model: str = "gpt-4o", 111 | *args, 112 | **kwargs, 113 | ) -> FSMRun: 114 | """ 115 | Executes a single step of the Finite State Machine (FSM) using the provided user input. 116 | Processes the current state, generates a response via the LLM, and transitions the FSM. 117 | 118 | Parameters: 119 | - async_openai_instance (openai.AsyncOpenAI): OpenAI client for API calls. 120 | - user_input (str): User input for the FSM. 121 | - model (str): The LLM model to use for generating responses (default: "gpt-4o"). 122 | 123 | Returns: 124 | - FSMRun: A structured representation of the FSM's state, chat history, and response. 125 | """ 126 | 127 | current_state: FSMState = self._state_registry.get(self._state) 128 | if not current_state: 129 | raise FSMError(f"State '{self._state}' not found in the state registry.") 130 | 131 | if current_state.preprocess_input: 132 | user_input = current_state.preprocess_input(user_input, self) or user_input 133 | 134 | # Prepare chat history to use as context 135 | chat_history_copy = self._session_history.copy() 136 | full_session_history_copy = self._full_session_history.copy() 137 | chat_history_copy.append({"role": "user", "content": user_input}) 138 | full_session_history_copy.append({"role": "user", "content": user_input}) 139 | 140 | # Generate response using our LLMUtilities 141 | response_schema = _generate_response_schema( 142 | current_state.response_model, 143 | current_state.transitions, 144 | current_state.key 145 | ) 146 | 147 | # Generate a response using the LLM 148 | response_data = await self._llm_utils.get_completion( 149 | async_openai_instance, 150 | chat_history_copy, 151 | response_schema, 152 | model, 153 | current_state 154 | ) 155 | 156 | # Extract response and next state 157 | next_state_key = response_data.get("next_state_key", current_state.key) 158 | raw_response = response_data.get("response") 159 | 160 | # Validate response 161 | if current_state.response_model: 162 | try: 163 | parsed_response = current_state.response_model(**raw_response) 164 | except ValidationError as error: 165 | raise FSMError(f"Error parsing response: {error}") 166 | else: 167 | parsed_response = raw_response.get("content", raw_response) 168 | 169 | # Validate and update next state 170 | if next_state_key not in self._state_registry: 171 | next_state_key = current_state.key 172 | self._next_state = next_state_key 173 | 174 | # Execute state logic 175 | function_context = { 176 | "fsm": self, 177 | "response": parsed_response, 178 | "will_transition": self._state != self._next_state, 179 | **kwargs, 180 | } 181 | final_response = await current_state.func(**function_context) 182 | 183 | # Handle immediate state changes if needed 184 | if isinstance(final_response, ImmediateStateChange): 185 | self._state = final_response.next_state 186 | return await self.run_state_machine( 187 | async_openai_instance, 188 | final_response.input, 189 | model, 190 | *args, 191 | **kwargs 192 | ) 193 | 194 | # Finalize response (or assigns fallback) 195 | final_response_str = final_response or parsed_response 196 | 197 | # Update histories 198 | chat_history_copy.append({"role": "assistant", "content": final_response_str}) 199 | full_session_history_copy.append({"role": "assistant", "content": final_response_str}) 200 | self._session_history = chat_history_copy 201 | self._full_session_history = full_session_history_copy 202 | 203 | # Transition state 204 | previous_state = self._state 205 | self._state = self._next_state 206 | self._next_state = None 207 | 208 | return FSMRun( 209 | state=self._state, 210 | chat_history=chat_history_copy, 211 | context_data=self.user_defined_context, 212 | response_raw=response_data, 213 | response=final_response_str, 214 | ) 215 | 216 | def reset(self): 217 | """Resets the FSM to its initial state.""" 218 | self._state = self._initial_state 219 | self._next_state = None 220 | self._session_history = [] 221 | self.user_defined_context = {} 222 | 223 | def get_curr_state(self): 224 | """Returns the current state of the FSM.""" 225 | return self._state 226 | 227 | def get_next_state(self): 228 | """Returns the next state of the FSM.""" 229 | return self._next_state 230 | 231 | def set_next_state(self, next_state: str): 232 | """Sets the next state of the FSM.""" 233 | self._next_state = next_state 234 | 235 | def get_running_session_history(self): 236 | """Returns the current running chat history.""" 237 | return self._running_session_history 238 | 239 | def set_running_session_history(self, chat_history: list): 240 | """Sets the current running chat history.""" 241 | self._running_session_history = chat_history 242 | 243 | def get_full_session_history(self): 244 | """Returns the full chat history of the FSM.""" 245 | return self._full_session_history 246 | 247 | def set_context_data(self, key: str, value: Any): 248 | """Sets a key-value pair into the user-defined context.""" 249 | self.user_defined_context[key] = value 250 | 251 | def set_context_data_dict(self, data: Dict[str, Any]): 252 | """Sets multiple key-value pairs into the user-defined context.""" 253 | self.user_defined_context.update(data) 254 | 255 | def get_context_data(self, key: str, default: Any = None): 256 | """Gets a value from the user-defined context, with a default value.""" 257 | return self.user_defined_context.get(key, default) 258 | 259 | def get_full_context_data(self): 260 | """Returns the full user-defined context.""" 261 | return self.user_defined_context 262 | 263 | def is_completed(self): 264 | """Checks if the FSM has reached its final state.""" 265 | return self._state == self._end_state 266 | 267 | def is_urgent_shift(self): 268 | """Checks if the FSM is in an urgent shift state.""" 269 | return self._is_urgent_shift 270 | 271 | -------------------------------------------------------------------------------- /examples/medical_agent.py: -------------------------------------------------------------------------------- 1 | import os 2 | from datetime import datetime 3 | from typing import List, Optional 4 | from enum import Enum 5 | from pydantic import BaseModel, Field 6 | import openai 7 | from dotenv import load_dotenv 8 | from fsm_llm.fsm import LLMStateMachine 9 | from fsm_llm.state_models import FSMRun, DefaultResponse, ImmediateStateChange 10 | 11 | # Load environment variables 12 | load_dotenv() 13 | 14 | # ANSI color codes 15 | GREY = "\033[90m" 16 | LIGHT_BLUE = "\033[94m" 17 | ORANGE = "\033[38;5;208m" 18 | RESET = "\033[0m" 19 | 20 | class Severity(str, Enum): 21 | LOW = "low" 22 | MODERATE = "moderate" 23 | HIGH = "high" 24 | CRITICAL = "critical" 25 | 26 | class VitalSigns(BaseModel): 27 | blood_pressure: Optional[str] 28 | heart_rate: Optional[int] 29 | temperature: Optional[float] 30 | oxygen_saturation: Optional[int] 31 | respiratory_rate: Optional[int] 32 | 33 | class Symptom(BaseModel): 34 | name: str 35 | severity: Severity 36 | duration: str 37 | description: str 38 | 39 | class PatientInfo(BaseModel): 40 | name: str 41 | age: int 42 | gender: str 43 | existing_conditions: List[str] = Field(default_factory=list) 44 | current_medications: List[str] = Field(default_factory=list) 45 | allergies: List[str] = Field(default_factory=list) 46 | 47 | class EmergencyAssessment(BaseModel): 48 | is_emergency: bool 49 | reasoning: str 50 | recommended_action: str 51 | 52 | class DrugInteraction(BaseModel): 53 | severity: Severity 54 | description: str 55 | recommendation: str 56 | 57 | class SymptomAssessment(BaseModel): 58 | symptoms: List[Symptom] 59 | potential_causes: List[str] 60 | risk_factors: List[str] 61 | additional_questions: List[str] 62 | 63 | class TreatmentPlan(BaseModel): 64 | primary_recommendations: List[str] 65 | lifestyle_modifications: List[str] 66 | follow_up_timeline: str 67 | warning_signs: List[str] 68 | emergency_conditions: List[str] 69 | 70 | # Initialize FSM 71 | fsm = LLMStateMachine(initial_state="INITIAL_TRIAGE", end_state="END") 72 | 73 | @fsm.define_state( 74 | state_key="INITIAL_TRIAGE", 75 | prompt_template=""" 76 | You are an advanced medical triage system. First, assess if this is an immediate emergency requiring urgent care. 77 | Look for red flags such as: 78 | - Chest pain, difficulty breathing, severe bleeding 79 | - Stroke symptoms (FAST: Face drooping, Arm weakness, Speech difficulty, Time to call emergency) 80 | - Severe allergic reactions 81 | - Loss of consciousness 82 | - Severe head injuries 83 | 84 | Based on the user's initial complaint, determine if immediate emergency response is needed. 85 | """, 86 | response_model=EmergencyAssessment, 87 | transitions={ 88 | "EMERGENCY": "If immediate medical attention is required", 89 | "GATHER_PATIENT_INFO": "If situation is not immediately life-threatening", 90 | } 91 | ) 92 | async def initial_triage( 93 | fsm: LLMStateMachine, 94 | response: EmergencyAssessment, 95 | will_transition: bool 96 | ) -> str: 97 | fsm.set_context_data("emergency_assessment", response.model_dump()) 98 | 99 | if response.is_emergency: 100 | return ImmediateStateChange( 101 | next_state="EMERGENCY", 102 | input="Emergency situation detected", 103 | keep_original_response=True 104 | ) 105 | 106 | return ( 107 | "Assessment: " + response.reasoning + "\n" 108 | "Recommended Action: " + response.recommended_action + "\n" 109 | "Since this isn't an immediate emergency, I'll need to gather some information about you. " 110 | "Please provide your name, age, gender, any existing medical conditions, " 111 | "current medications, and allergies." 112 | ) 113 | 114 | @fsm.define_state( 115 | state_key="EMERGENCY", 116 | prompt_template=""" 117 | EMERGENCY PROTOCOL ACTIVATED 118 | Provide clear, urgent instructions while emergency services are contacted. 119 | Review the situation and provide immediate first-aid guidance if appropriate. 120 | """, 121 | response_model=DefaultResponse, 122 | transitions={"END": "After emergency instructions are provided"} 123 | ) 124 | async def emergency_handler( 125 | fsm: LLMStateMachine, 126 | response: DefaultResponse, 127 | will_transition: bool 128 | ) -> str: 129 | emergency_data = fsm.get_context_data("emergency_assessment") 130 | return ( 131 | "🚨 EMERGENCY SITUATION DETECTED 🚨\n" 132 | + response.content + "\n\n" 133 | "Please call emergency services immediately (911 in the US).\n" 134 | "Stay on the line while help is dispatched." 135 | ) 136 | 137 | @fsm.define_state( 138 | state_key="GATHER_PATIENT_INFO", 139 | prompt_template=""" 140 | Parse the user's information into structured patient data. 141 | Prompt for any missing critical information. 142 | Look for any red flags in the patient's history or medications. 143 | """, 144 | response_model=PatientInfo, 145 | transitions={ 146 | "SYMPTOM_ASSESSMENT": "When patient info is complete", 147 | "EMERGENCY": "If any red flags are detected in patient history", 148 | } 149 | ) 150 | async def gather_patient_info( 151 | fsm: LLMStateMachine, 152 | response: PatientInfo, 153 | will_transition: bool 154 | ) -> str: 155 | fsm.set_context_data("patient_info", response.model_dump()) 156 | 157 | return ( 158 | "Thank you. I've recorded your information. " 159 | "Now, please describe the symptoms you're experiencing, " 160 | "including when they started and how severe they are." 161 | ) 162 | 163 | @fsm.define_state( 164 | state_key="SYMPTOM_ASSESSMENT", 165 | prompt_template=""" 166 | Analyze the reported symptoms considering: 167 | - Patient's age, gender, and medical history 168 | - Symptom severity and duration 169 | - Potential interactions with existing conditions 170 | - Risk factors and warning signs 171 | 172 | Generate a structured assessment and identify any patterns or concerning combinations. 173 | """, 174 | response_model=SymptomAssessment, 175 | transitions={ 176 | "DRUG_INTERACTION_CHECK": "If symptoms are well understood and non-emergency", 177 | "EMERGENCY": "If symptoms suggest a serious condition", 178 | } 179 | ) 180 | async def assess_symptoms( 181 | fsm: LLMStateMachine, 182 | response: SymptomAssessment, 183 | will_transition: bool 184 | ) -> str: 185 | fsm.set_context_data("symptom_assessment", response.model_dump()) 186 | 187 | # Check for high-severity symptoms 188 | if any(s.severity == Severity.CRITICAL for s in response.symptoms): 189 | return ImmediateStateChange( 190 | next_state="EMERGENCY", 191 | input="Critical symptoms detected" 192 | ) 193 | 194 | assessment = "Symptom Assessment:\n" 195 | for symptom in response.symptoms: 196 | assessment += f"- {symptom.name} ({symptom.severity}): {symptom.description}\n" 197 | 198 | assessment += "\nPotential Causes:\n" 199 | for cause in response.potential_causes: 200 | assessment += f"- {cause}\n" 201 | 202 | if response.additional_questions: 203 | assessment += "\nI need some additional information:\n" 204 | for question in response.additional_questions: 205 | assessment += f"- {question}\n" 206 | 207 | return assessment 208 | 209 | @fsm.define_state( 210 | state_key="DRUG_INTERACTION_CHECK", 211 | prompt_template=""" 212 | Analyze the patient's current medications for potential interactions. 213 | Consider both existing conditions and reported symptoms. 214 | Flag any concerning combinations or contraindications. 215 | """, 216 | response_model=DrugInteraction, 217 | transitions={ 218 | "GENERATE_PLAN": "If no severe interactions are found", 219 | "EMERGENCY": "If dangerous drug interactions are detected", 220 | } 221 | ) 222 | async def check_drug_interactions( 223 | fsm: LLMStateMachine, 224 | response: DrugInteraction, 225 | will_transition: bool 226 | ) -> str: 227 | fsm.set_context_data("drug_interactions", response.model_dump()) 228 | 229 | if response.severity == Severity.CRITICAL: 230 | return ImmediateStateChange( 231 | next_state="EMERGENCY", 232 | input=f"Critical drug interaction detected: {response.description}" 233 | ) 234 | 235 | interaction_msg = ( 236 | f"Medication Analysis:\n" 237 | f"Severity: {response.severity}\n" 238 | f"Details: {response.description}\n" 239 | f"Recommendation: {response.recommendation}\n\n" 240 | ) 241 | 242 | return interaction_msg + "Now, let's generate your treatment plan." 243 | 244 | @fsm.define_state( 245 | state_key="GENERATE_PLAN", 246 | prompt_template=""" 247 | Based on the complete assessment, generate a comprehensive care plan. 248 | Consider: 249 | - Patient's specific circumstances and limitations 250 | - Interaction with existing conditions and medications 251 | - Clear follow-up timeline and monitoring plans 252 | - Specific warning signs to watch for 253 | """, 254 | response_model=TreatmentPlan, 255 | transitions={ 256 | "END": "After plan is generated and explained", 257 | "EMERGENCY": "If complications arise during plan generation", 258 | } 259 | ) 260 | async def generate_treatment_plan( 261 | fsm: LLMStateMachine, 262 | response: TreatmentPlan, 263 | will_transition: bool 264 | ) -> str: 265 | fsm.set_context_data("treatment_plan", response.model_dump()) 266 | 267 | plan = "Based on our assessment, here's your care plan:\n\n" 268 | 269 | plan += "Recommendations:\n" 270 | for rec in response.primary_recommendations: 271 | plan += f"- {rec}\n" 272 | 273 | plan += "\nLifestyle Modifications:\n" 274 | for mod in response.lifestyle_modifications: 275 | plan += f"- {mod}\n" 276 | 277 | plan += f"\nFollow-up Timeline: {response.follow_up_timeline}\n" 278 | 279 | plan += "\nWarning Signs (Seek immediate care if you experience):\n" 280 | for warning in response.warning_signs: 281 | plan += f"- {warning}\n" 282 | 283 | plan += "\nEmergency Conditions:\n" 284 | for condition in response.emergency_conditions: 285 | plan += f"- {condition}\n" 286 | 287 | # Generate audit log 288 | audit_log = { 289 | "timestamp": datetime.now().isoformat(), 290 | "patient_info": fsm.get_context_data("patient_info"), 291 | "emergency_assessment": fsm.get_context_data("emergency_assessment"), 292 | "drug_interactions": fsm.get_context_data("drug_interactions"), 293 | "symptom_assessment": fsm.get_context_data("symptom_assessment"), 294 | "treatment_plan": response.model_dump(), 295 | } 296 | fsm.set_context_data("audit_log", audit_log) 297 | 298 | return plan 299 | 300 | @fsm.define_state( 301 | state_key="END", 302 | prompt_template="Provide final instructions and documentation.", 303 | response_model=DefaultResponse, 304 | ) 305 | async def end_state( 306 | fsm: LLMStateMachine, 307 | response: DefaultResponse, 308 | will_transition: bool 309 | ) -> str: 310 | audit_log = fsm.get_context_data("audit_log") 311 | # Here you could save the audit log to a secure medical record system 312 | 313 | return ( 314 | "Your assessment is complete. Please follow the provided care plan " 315 | "and don't hesitate to seek emergency care if warning signs develop. " 316 | "A record of this assessment has been saved for future reference." 317 | ) 318 | 319 | async def main(): 320 | openai_client = openai.AsyncOpenAI() 321 | print(GREY + "Medical Triage System Initialized" + RESET) 322 | print(GREY + "Please describe your medical concern:" + RESET) 323 | 324 | while not fsm.is_completed(): 325 | # Instead of fsm.current_state, we print the state from the FSMRun 326 | run_state: FSMRun = await fsm.run_state_machine( 327 | openai_client, 328 | user_input=input(ORANGE + "You: " + RESET), 329 | model="gpt-4o-mini" 330 | ) 331 | # Assuming run_state has an attribute `state` representing the current state key. 332 | print(LIGHT_BLUE + f"Current state: {run_state.state}" + RESET) 333 | print(GREY + f"System: {run_state.response}" + RESET) 334 | 335 | if __name__ == "__main__": 336 | import asyncio 337 | asyncio.run(main()) 338 | -------------------------------------------------------------------------------- /content/calculus_content.json: -------------------------------------------------------------------------------- 1 | [ 2 | { 3 | "id": "1", 4 | "name": "Preliminaries", 5 | "content": "The preliminary terror, which chokes off most fifth-form boys from even attempting to learn how to calculate, can be abolished once for all by simply stating what is the meaning-in common-sense terms-of the two principal symbols that are used in calculating. These dreadful symbols are:\n\n(1) d which merely means “a little bit of.”\n\nThus dx means a little bit of x; or du means a little bit of u. Ordinary mathematicians think it more polite to say “an element of,” instead of “a little bit of.” Just as you please. But you will find that these little bits (or elements) may be considered to be indefinitely small.\n\n(2) ∫ which is merely a long S, and may be called (if you like) “the sum of.”\n\nThus ∫dx means the sum of all the little bits of x; or ∫dt means the sum of all the little bits of t. Ordinary mathematicians call this symbol “the integral of.” Now any fool can see that if x is considered as made up of a lot of little bits, each of which is called dx, if you add them all up together you get the sum of all the dx's, (which is the same thing as the whole of x). The word “integral” simply means “the whole.” If you think of the duration of time for one hour, you may (if you like) think of it as cut up into 3600 little bits called seconds. The whole of the 3600 little bits added up together make one hour.\n\nWhen you see an expression that begins with this terrifying symbol, you will henceforth know that it is put there merely to give you instructions that you are now to perform the operation (if you can) of totalling up all the little bits that are indicated by the symbols that follow.\n\nThat's all." 6 | }, 7 | { 8 | "id": "2", 9 | "name": "On Different Degrees of Smallness", 10 | "content": "We shall find that in our processes of calculation we have to deal with small quantities of various degrees of smallness.\n\nWe shall have also to learn under what circumstances we may consider small quantities to be so minute that we may omit them from consideration. Everything depends upon relative minuteness.\n\nBefore we fix any rules let us think of some familiar cases. There are 60 minutes in the hour, 24 hours in the day, 7 days in the week. There are therefore 1440 minutes in the day and 10080 minutes in the week.\n\nObviously 1 minute is a very small quantity of time compared with a whole week. Indeed, our forefathers considered it small as compared with an hour, and called it “one minùte,” meaning a minute fraction-namely one sixtieth-of an hour. When they came to require still smaller subdivisions of time, they divided each minute into 60 still smaller parts, which, in Queen Elizabeth's days, they called “second minùtes” (i.e.: small quantities of the second order of minuteness). Nowadays we call these small quantities of the second order of smallness “seconds.” But few people know why they are so called.\n\nNow if one minute is so small as compared with a whole day, how much smaller by comparison is one second!\n\nAgain, think of a farthing as compared with a sovereign: it is barely worth more than 1/1000 part. A farthing more or less is of precious little importance compared with a sovereign: it may certainly be regarded as a small quantity. But compare a farthing with £1000: relatively to this greater sum, the farthing is of no more importance than 1/1000 of a farthing would be to a sovereign. Even a golden sovereign is relatively a negligible quantity in the wealth of a millionaire.\n\nNow if we fix upon any numerical fraction as constituting the proportion which for any purpose we call relatively small, we can easily state other fractions of a higher degree of smallness. Thus if, for the purpose of time, 1/60 be called a small fraction, then 1/60 of 1/60 (being a small fraction of a small fraction) may be regarded as a small quantity of the second order of smallness.\n\nThe mathematicians talk about the second order of “magnitude” (i.e.: greatness) when they really mean second order of smallness. This is very confusing to beginners.\n\nOr, if for any purpose we were to take 1 per cent. (i.e.: 1/100) as a small fraction, then 1 per cent. of 1 per cent. (i.e.: 1/10,000) would be a small fraction of the second order of smallness; and 1/1,000,000 would be a small fraction of the third order of smallness, being 1 per cent. of 1 per cent. of 1 per cent.\n\nLastly, suppose that for some very precise purpose we should regard 1/1,000,000 as “small.” Thus, if a first-rate chronometer is not to lose or gain more than half a minute in a year, it must keep time with an accuracy of 1 part in 1,051,200. Now if, for such a purpose, we regard 1/1,000,000 (or one millionth) as a small quantity, then 1/1,000,000 of 1/1,000,000, that is 1/1,000,000,000,000 (or one billionth) will be a small quantity of the second order of smallness, and may be utterly disregarded, by comparison.\n\nThen we see that the smaller a small quantity itself is, the more negligible does the corresponding small quantity of the second order become. Hence we know that in all cases we are justified in neglecting the small quantities of the second-or third (or higher)-orders, if only we take the small quantity of the first order small enough in itself.\n\nBut, it must be remembered, that small quantities if they occur in our expressions as factors multiplied by some other factor, may become important if the other factor is itself large. Even a farthing becomes important if only it is multiplied by a few hundred.\n\nNow in the calculus we write dx for a little bit of x. These things such as dx, and du, and dy, are called “differentials,” the differential of x, or of u, or of y, as the case may be. [You read them as dee-eks, or dee-you, or dee-wy.] If dx be a small bit of x, and relatively small of itself, it does not follow that such quantities as x·dx, or x^2dx, or axdx are negligible. But dx×dx would be negligible, being a small quantity of the second order." 11 | }, 12 | { 13 | "id": "3", 14 | "name": "Calculus and Variables", 15 | "content": "All through the calculus we are dealing with quantities that are growing, and with rates of growth. We classify all quantities into two classes: constants and variables. Those which we regard as of fixed value, and call constants, we generally denote algebraically by letters from the beginning of the alphabet, such as a, b, or c; while those which we consider as capable of growing, or (as mathematicians say) of 'varying,' we denote by letters from the end of the alphabet, such as x, y, z, u, v, w, or sometimes t.\n\nMoreover, we are usually dealing with more than one variable at once, and thinking of the way in which one variable depends on the other: for instance, we think of the way in which the height reached by a projectile depends on the time of attaining that height. Or we are asked to consider a rectangle of given area, and to enquire how any increase in the length of it will compel a corresponding decrease in the breadth of it. Or we think of the way in which any variation in the slope of a ladder will cause the height that it reaches, to vary.\n\nSuppose we have got two such variables that depend one on the other. An alteration in one will bring about an alteration in the other, because of this dependence. Let us call one of the variables x, and the other that depends on it y.\n\nSuppose we make x to vary, that is to say, we either alter it or imagine it to be altered, by adding to it a bit which we call dx. We are thus causing x to become x+dx. Then, because x has been altered, y will have altered also, and will have become y+dy. Here the bit dy may be in some cases positive, in others negative; and it won't (except by a miracle) be the same size as dx.\n\nTake two examples.\n\n(1) Let x and y be respectively the base and the height of a right-angled triangle (Figure 4), of which the slope of the other side is fixed at 30°. If we suppose this triangle to expand and yet keep its angles the same as at first, then, when the base grows so as to become x+dx, the height becomes y+dy. Here, increasing x results in an increase of y. The little triangle, the height of which is dy, and the base of which is dx, is similar to the original triangle; and it is obvious that the value of the ratio dydx is the same as that of the ratio yx. As the angle is 30° it will be seen that here dydx=1/1.73.\n\n(2) Let x represent, in Figure 5, the horizontal distance, from a wall, of the bottom end of a ladder, AB, of fixed length; and let y be the height it reaches up the wall. Now y clearly depends on x. It is easy to see that, if we pull the bottom end A a bit further from the wall, the top end B will come down a little lower. Let us state this in scientific language. If we increase x to x+dx, then y will become y−dy; that is, when x receives a positive increment, the increment which results to y is negative.\n\nYes, but how much? Suppose the ladder was so long that when the bottom end A was 19 inches from the wall the top end B reached just 15 feet from the ground. Now, if you were to pull the bottom end out 1 inch more, how much would the top end come down? Put it all into inches: x=19 inches, y=180 inches. Now the increment of x which we call dx, is 1 inch: or x+dx=20 inches.\n\nHow much will y be diminished? The new height will be y−dy. If we work out the height by Euclid I. 47, then we shall be able to find how much dy will be. The length of the ladder is √(180^2 + 19^2) = 181 inches. Clearly then, the new height, which is y−dy, will be such that (y−dy)^2 = (181^2 − 20^2) = 32761−400 = 32361, √32361 = 179.89 inches. Now y is 180, so that dy is 180−179.89 = 0.11 inch.\n\nSo we see that making dx an increase of 1 inch has resulted in making dy a decrease of 0.11 inch.\n\nAnd the ratio of dy to dx may be stated thus: dydx = −0.11/1.\n\nIt is also easy to see that (except in one particular position) dy will be of a different size from dx.\n\nNow right through the differential calculus we are hunting, hunting, hunting for a curious thing, a mere ratio, namely, the proportion which dy bears to dx when both of them are indefinitely small.\n\nIt should be noted here that we can only find this ratio dydx when y and x are related to each other in some way, so that whenever x varies y does vary also. For instance, in the first example just taken, if the base x of the triangle be made longer, the height y of the triangle becomes greater also, and in the second example, if the distance x of the foot of the ladder from the wall be made to increase, the height y reached by the ladder decreases in a corresponding manner, slowly at first, but more and more rapidly as x becomes greater. In these cases the relation between x and y is perfectly definite, it can be expressed mathematically, being yx=tan30° and x^2 + y^2 = l^2 (where l is the length of the ladder) respectively, and dydx has the meaning we found in each case.\n\nIf, while x is, as before, the distance of the foot of the ladder from the wall, y is, instead of the height reached, the horizontal length of the wall, or the number of bricks in it, or the number of years since it was built, any change in x would naturally cause no change whatever in y; in this case dydx has no meaning whatever, and it is not possible to find an expression for it. Whenever we use differentials dx, dy, dz, etc., the existence of some kind of relation between x, y, z, etc., is implied, and this relation is called a 'function' in x, y, z, etc.; the two expressions given above, for instance, namely yx=tan30° and x^2 + y^2 = l^2, are functions of x and y. Such expressions contain implicitly (that is, contain without distinctly showing it) the means of expressing either x in terms of y or y in terms of x, and for this reason they are called implicit functions in x and y; they can be respectively put into the forms y = x tan30° or x = y tan30° = √(l^2 − x^2) or x = √(l^2 − y^2).\n\nThese last expressions state explicitly (that is, distinctly) the value of x in terms of y, or of y in terms of x, and they are for this reason called explicit functions of x or y. For example x^2 + 3 = 2y − 7 is an implicit function in x and y; it may be written y = (x^2 + 10)/2 (explicit function of x) or x = √(2y − 10) (explicit function of y). We see that an explicit function in x, y, z, etc., is simply something the value of which changes when x, y, z, etc., are changing, either one at the time or several together. Because of this, the value of the explicit function is called the dependent variable, as it depends on the value of the other variable quantities in the function; these other variables are called the independent variables because their value is not determined from the value assumed by the function. For example, if u = x^2 sin(θ), x and θ are the independent variables, and u is the dependent variable.\n\nSometimes the exact relation between several quantities x, y, z either is not known or it is not convenient to state it; it is only known, or convenient to state, that there is some sort of relation between these variables, so that one cannot alter either x or y or z singly without affecting the other quantities; the existence of a function in x, y, z is then indicated by the notation F(x, y, z) (implicit function) or by x = F(y, z), y = F(x, z) or z = F(x, y) (explicit function). Sometimes the letter f or ϕ is used instead of F, so that y=F(x), y=f(x) and y=ϕ(x) all mean the same thing, namely, that the value of y depends on the value of x in some way which is not stated.\n\nWe call the ratio dydx 'the differential coefficient of y with respect to x.' It is a solemn scientific name for this very simple thing. But we are not going to be frightened by solemn names, when the things themselves are so easy. Instead of being frightened we will simply pronounce a brief curse on the stupidity of giving long crack-jaw names; and, having relieved our minds, will go on to the simple thing itself, namely the ratio dydx.\n\nIn ordinary algebra which you learned at school, you were always hunting after some unknown quantity which you called x or y; or sometimes there were two unknown quantities to be hunted for simultaneously. You have now to learn to go hunting in a new way; the fox being now neither x nor y. Instead of this you have to hunt for this curious cub called dydx. The process of finding the value of dydx is called 'differentiating.' But, remember, what is wanted is the value of this ratio when both dy and dx are themselves indefinitely small. The true value of the differential coefficient is that to which it approximates in the limiting case when each of them is considered as infinitesimally minute.\n\nLet us now learn how to go in quest of dydx." 16 | }, 17 | { 18 | "id": "4", 19 | "name": "Simplest Cases", 20 | "content": "Let us begin with the simple expression y=x². Now remember that the fundamental notion about the calculus is the idea of growing. Mathematicians call it varying. Now as y and x² are equal to one another, it is clear that if x grows, x² will also grow. And if x² grows, then y will also grow. What we have got to find out is the proportion between the growing of y and the growing of x. In other words, our task is to find out the ratio between dy and dx, or, in brief, to find the value of dydx. Let x grow a little bit bigger and become x+dx; similarly, y will grow a bit bigger and will become y+dy. Then, clearly, it will still be true that the enlarged y will be equal to the square of the enlarged x. Writing this down, we have: y+dy = (x+dx)² = x² + 2x·dx + (dx)². What does (dx)² mean? Remember that dx meant a bit—a little bit—of x. Then (dx)² will mean a little bit of a little bit of x; that is, as explained above, it is a small quantity of the second order of smallness. It may therefore be discarded as quite inconsiderable in comparison with the other terms. Leaving it out, we then have: y+dy. Now y=x²; so let us subtract this from the equation and we have left dy. Dividing across by dx, we find dydx = 2x. Now this is what we set out to find. The ratio of the growing of y to the growing of x is, in the case before us, found to be 2x. This ratio dydx is the result of differentiating y with respect to x. Differentiating means finding the differential coefficient. Suppose we had some other function of x, as, for example, u=7x²+3. Then if we were told to differentiate this with respect to x, we should have to find dudx, or, what is the same thing, d(7x²+3)/dx. On the other hand, we may have a case in which time was the independent variable, such as y=b+12at². Then, if we were told to differentiate it, that means we must find its differential coefficient with respect to t, so that then our business would be to try to find dydt, that is, to find d(b+12at²)/dt. Suppose x=100 and y=10000. Let x grow till it becomes 101 (that is, let dx=1). Then the enlarged y will be 101×101=10201. But if we agree that we may ignore small quantities of the second order, 1 may be rejected as compared with 10000; so we may round off the enlarged y to 10200. y has grown from 10000 to 10200; the bit added on is dy, which is therefore 200. dydx=200/1=200. According to the algebra-working of the previous paragraph, we find dydx=2x. And so it is; for x=100 and 2x=200. But, you will say, we neglected a whole unit. Well, try again, making dx a still smaller bit. Try dx=0.1. Then x+dx=100.1, and (x+dx)²=100.1×100.1=10020.01. Now the last figure 1 is only one-millionth part of the 10000, and is utterly negligible; so we may take 10020 without the little decimal at the end. And this makes dy=20; and dydx=200/0.1=200, which is still the same as 2x." 21 | }, 22 | { 23 | "id": "5", 24 | "name": "Simplest Cases", 25 | "content": "Following out logically our observation, we should conclude that if we want to deal with any higher power,-call it n-we could tackle it in the same way. Let y=xⁿ, then, we should expect to find that dydx=nxⁿ⁻¹. For example, let n=8, then y=x⁸; and differentiating it would give dydx=8x⁷. And, indeed, the rule that differentiating xⁿ gives as the result nxⁿ⁻¹ is true for all cases where n is a whole number and positive. [Expanding (x+dx)ⁿ by the binomial theorem will at once show this.] But the question whether it is true for cases where n has negative or fractional values requires further consideration. Case of a negative power. Let y=x⁻². Then proceed as before: y+dy=(x+dx)⁻²=x⁻²(1+dx/x)⁻². Expanding this by the binomial theorem, we get =x⁻²[1-2dx/x+2(2+1)1×2(dx/x)²−etc.]=x⁻²−2x⁻³·dx+3x⁻⁴(dx)²−4x⁻⁵(dx)³+etc. So, neglecting the small quantities of higher orders of smallness, we have: y+dy=x⁻²−2x⁻³·dx. Subtracting the original y=x⁻², we find dydx=-2x⁻³, which is still in accordance with the rule inferred above. Case of a fractional power. Let y=x¹². Then, as before, y+dy=(x+dx)¹²=x¹²(1+dx/x)¹²=x√+12dx/x√−18(dx)²/x√+terms with higher powers of dx. Subtracting the original y=x¹², and neglecting higher powers we have left: dy=12dx/x√=12x⁻¹²·dx, and dydx=12x⁻¹². Agreeing with the general rule. Summary. Let us see how far we have got. We have arrived at the following rule: To differentiate xⁿ, multiply by the power and reduce the power by one, so giving us nxⁿ⁻¹ as the result." 26 | } 27 | ] 28 | --------------------------------------------------------------------------------