-
Notifications
You must be signed in to change notification settings - Fork 18
Expand file tree
/
Copy pathfsm_definition.py
More file actions
70 lines (55 loc) · 1.85 KB
/
fsm_definition.py
File metadata and controls
70 lines (55 loc) · 1.85 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
from chatfaq_sdk import ChatFAQSDK
from chatfaq_sdk.clients import llm_request
from chatfaq_sdk.fsm import FSMDefinition, State, Transition
from chatfaq_sdk.layers import Message
from pydantic import BaseModel
import json
import logging
logger = logging.getLogger(__name__)
logging.basicConfig(
level=logging.DEBUG,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
# Define your desired output structure
class UserInfo(BaseModel):
name: str
age: int
async def send_greeting(sdk: ChatFAQSDK, ctx: dict):
yield Message(
"I will extract the name and age of an user description, please input the description"
)
async def send_info(sdk: ChatFAQSDK, ctx: dict):
print('The last MML is:')
print(ctx["conv_mml"][-1])
logger.info("Extracting user info...")
response = await llm_request(
sdk,
"gemini-2.0-flash",
messages=[
{
"role": "system",
"content": "You are an assistant that extracts the user name and age from a description into a json object.",
},
],
conversation_id=ctx["conversation_id"],
bot_channel_name=ctx["bot_channel_name"],
use_conversation_context=True,
response_schema=UserInfo.model_json_schema(),
)
data = response.get("content", None)
if data:
yield Message(f"Here is the extracted information: {json.dumps(data)}.")
else:
yield Message("No information extracted")
greeting_state = State(name="Greeting", events=[send_greeting], initial=True)
send_info_state = State(
name="Send Info",
events=[send_info],
)
# After the initial state, always transition to the extraction state
_to_send_info = Transition(
dest=send_info_state,
)
fsm_definition = FSMDefinition(
states=[greeting_state, send_info_state], transitions=[_to_send_info]
)