Pass tool call and app function response back in metadata (#193)

This commit is contained in:
Adil Hafeez 2024-10-18 13:25:39 -07:00 committed by GitHub
parent 62a000036e
commit dd1c7be706
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 169 additions and 112 deletions

View file

@ -13,62 +13,38 @@ logger = get_model_server_logger()
class Message(BaseModel):
role: str
content: str
content: str = ""
tool_calls: List[Dict[str, Any]] = []
tool_call_id: str = ""
class ChatMessage(BaseModel):
messages: list[Message]
tools: List[Dict[str, Any]]
# TODO: make it default none
metadata: Dict[str, str] = {}
def process_state(arch_state, history: list[Message]):
logger.info("state: {}".format(arch_state))
state_json = json.loads(arch_state)
state_map = {}
if state_json:
for tools_state in state_json:
for tool_state in tools_state:
state_map[tool_state["key"]] = tool_state
logger.info(f"state_map: {json.dumps(state_map)}")
sha_history = []
def process_messages(history: list[Message]):
updated_history = []
for hist in history:
updated_history.append({"role": hist.role, "content": hist.content})
if hist.role == "user":
sha_history.append(hist.content)
sha256_hash = hashlib.sha256()
joined_key_str = ("#.#").join(sha_history)
sha256_hash.update(joined_key_str.encode())
sha_key = sha256_hash.hexdigest()
logger.info(f"sha_key: {sha_key}")
if sha_key in state_map:
tool_call_state = state_map[sha_key]
if "tool_call" in tool_call_state:
tool_call_str = json.dumps(tool_call_state["tool_call"])
updated_history.append(
{
"role": "assistant",
"content": f"<tool_call>\n{tool_call_str}\n</tool_call>",
}
)
if "tool_response" in tool_call_state:
tool_resp = tool_call_state["tool_response"]
# TODO: try with role = user as well
updated_history.append(
{
"role": "user",
"content": f"<tool_response>\n{tool_resp}\n</tool_response>",
}
)
# we dont want to match this state with any other messages
del state_map[sha_key]
if hist.tool_calls:
if len(hist.tool_calls) > 1:
raise ValueError("Only one tool call is supported")
tool_call_str = json.dumps(hist.tool_calls[0]["function"])
updated_history.append(
{
"role": "assistant",
"content": f"<tool_call>\n{tool_call_str}\n</tool_call>",
}
)
elif hist.role == "tool":
updated_history.append(
{
"role": "user",
"content": f"<tool_response>\n{hist.content}\n</tool_response>",
}
)
else:
updated_history.append({"role": hist.role, "content": hist.content})
return updated_history
@ -79,10 +55,7 @@ async def chat_completion(req: ChatMessage, res: Response):
messages = [{"role": "system", "content": tools_encoded}]
metadata = req.metadata
arch_state = metadata.get("x-arch-state", "[]")
updated_history = process_state(arch_state, req.messages)
updated_history = process_messages(req.messages)
for message in updated_history:
messages.append({"role": message["role"], "content": message["content"]})