Skip to content

Commit f2b698f

Browse files
hustshawnhustshawn
authored
Merge pull request #1 from lindarr915/azure-openai
Fix on Claude 3.5 tool use
2 parents a40c87f + c2ed0df commit f2b698f

File tree

4 files changed

+172
-7
lines changed

4 files changed

+172
-7
lines changed

‎.gitignore

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -159,4 +159,5 @@ cython_debug/
159159
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
160160
.idea/
161161

162-
Config
162+
Config
163+
.vscode/launch.json

‎src/api/models/bedrock.py

Lines changed: 35 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,12 @@ class BedrockModel(BaseChatModel):
173173
"tool_call": True,
174174
"stream_tool_call": False,
175175
},
176+
"mistral.mistral-large-2407-v1:0": {
177+
"system": True,
178+
"multimodal": False,
179+
"tool_call": True,
180+
"stream_tool_call": False,
181+
},
176182
"cohere.command-r-v1:0": {
177183
"system": True,
178184
"multimodal": False,
@@ -336,8 +342,33 @@ def _parse_messages(self, chat_request: ChatRequest) -> list[dict]:
336342
elif isinstance(message, AssistantMessage):
337343
if message.content:
338344
# Text message
339-
messages.append(
340-
{"role": message.role, "content": [{"text": message.content}]}
345+
if message.tool_calls:
346+
tool_input = json.loads(message.tool_calls[0].function.arguments)
347+
messages.append(
348+
{
349+
"role": message.role,
350+
"content": [
351+
{
352+
# Tool use message
353+
"toolUse": {
354+
"toolUseId": message.tool_calls[0].id,
355+
"name": message.tool_calls[0].function.name,
356+
"input": tool_input,
357+
},
358+
}
359+
],
360+
}
361+
)
362+
else:
363+
messages.append(
364+
{
365+
"role": message.role,
366+
"content": [
367+
{
368+
"text": message.content,
369+
}
370+
],
371+
}
341372
)
342373
else:
343374
# Tool use message
@@ -351,7 +382,7 @@ def _parse_messages(self, chat_request: ChatRequest) -> list[dict]:
351382
"toolUseId": message.tool_calls[0].id,
352383
"name": message.tool_calls[0].function.name,
353384
"input": tool_input,
354-
}
385+
},
355386
}
356387
],
357388
}
@@ -462,7 +493,7 @@ def _create_response(
462493
)
463494
message.tool_calls = tool_calls
464495
message.content = None
465-
else:
496+
elif content:
466497
message.content = content[0]["text"]
467498

468499
response = ChatResponse(

‎src/requirements.txt

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,5 +5,5 @@ mangum==0.17.0
55
tiktoken==0.6.0
66
requests==2.32.3
77
numpy==1.26.4
8-
boto3==1.34.132
9-
botocore==1.34.132
8+
boto3==1.35.17
9+
botocore==1.35.17

‎test/test.ipynb

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": 15,
6+
"metadata": {},
7+
"outputs": [],
8+
"source": [
9+
"from langchain.agents import AgentExecutor, create_tool_calling_agent\n",
10+
"from langchain_core.messages import HumanMessage\n",
11+
"from langchain_core.prompts import *\n",
12+
"from langchain.tools import tool\n",
13+
"from langchain_openai import AzureOpenAIEmbeddings, AzureChatOpenAI\n",
14+
"from tenacity import retry, stop_after_attempt, wait_exponential\n",
15+
"\n",
16+
"@tool\n",
17+
"def trident_cve_tool(cve_id: str):\n",
18+
" \"\"\"\n",
19+
" This is a CVE tool to call trident CVE information.\n",
20+
" It should be used to fetch general information of a specific cve (format like CVE-XXX-XXX).\n",
21+
" The general information from Trident API includes:\n",
22+
" cve id, name, description, update date, publish date and some other public available information.\n",
23+
" \"\"\"\n",
24+
" return str(\"Trident CVE data\")\n",
25+
"\n",
26+
"prompt = ChatPromptTemplate.from_messages(\n",
27+
" [\n",
28+
" SystemMessagePromptTemplate(\n",
29+
" prompt=PromptTemplate(\n",
30+
" input_variables=[\"cve_id\"],\n",
31+
" template=\"You are a helpful assistant to explain a certain CVE {cve_id} in detail\",\n",
32+
" )\n",
33+
" ),\n",
34+
" MessagesPlaceholder(variable_name=\"chat_history\", optional=True),\n",
35+
" HumanMessagePromptTemplate(\n",
36+
" prompt=PromptTemplate(\n",
37+
" input_variables=[\"question\"],\n",
38+
" template=\"Please answer this question: {question}\",\n",
39+
" )\n",
40+
" ),\n",
41+
" MessagesPlaceholder(variable_name=\"agent_scratchpad\"),\n",
42+
" ]\n",
43+
")\n",
44+
"tools = [trident_cve_tool]\n",
45+
"\n",
46+
"@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))\n",
47+
"def create_llm():\n",
48+
" return AzureChatOpenAI(\n",
49+
" temperature=0.0,\n",
50+
" azure_endpoint=\"http://localhost:8000\",\n",
51+
" openai_api_key=\"bedrock\",\n",
52+
" openai_api_version=\"2024-06-01\",\n",
53+
" deployment_name=\"anthropic.claude-3-5-sonnet-20240620-v1:0\", \n",
54+
" # deployment_name=\"gpt-4o_2024-05-13\",\n",
55+
" openai_api_type=\"azure\",\n",
56+
" )\n",
57+
"\n",
58+
"\n"
59+
]
60+
},
61+
{
62+
"cell_type": "code",
63+
"execution_count": 16,
64+
"metadata": {},
65+
"outputs": [
66+
{
67+
"name": "stdout",
68+
"output_type": "stream",
69+
"text": [
70+
"\n",
71+
"\n",
72+
"\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
73+
"\u001b[32;1m\u001b[1;3m\n",
74+
"Invoking: `trident_cve_tool` with `{'cve_id': 'CVE-2020-17087'}`\n",
75+
"\n",
76+
"\n",
77+
"\u001b[0m\u001b[36;1m\u001b[1;3mTrident CVE data\u001b[0m\u001b[32;1m\u001b[1;3m\n",
78+
"\n",
79+
"CVE-2020-17087 is a vulnerability in the Linux kernel's VFIO device driver that could allow a privileged user to escalate privileges or cause a denial of service on the system.\n",
80+
"\n",
81+
"The VFIO (Virtual Function I/O) driver allows safe and efficient device access from userspace on Linux hosts. This vulnerability is a use-after-free issue that could occur when removing a VFIO device. An unprivileged user with access to a VFIO device could potentially exploit this to gain elevated privileges on the system.\n",
82+
"\n",
83+
"Some key details about CVE-2020-17087:\n",
84+
"\n",
85+
"- Published: 2020-08-12\n",
86+
"- CVSS Score: 7.8 (High severity)\n",
87+
"- Affects: Linux kernel versions 5.4 through 5.8.5\n",
88+
"- Type: Use-after-free vulnerability in vfio_unbound_driver()\n",
89+
"- Impact: Privilege escalation, denial of service\n",
90+
"\n",
91+
"The vulnerability was discovered by Cedric Hombracht and reported to Linux kernel security teams. It was patched in Linux kernel 5.8.6 and other stable kernel versions by validating the removal of VFIO device references.\n",
92+
"\n",
93+
"System administrators running affected Linux kernel versions should apply the patched updates from their Linux distributions to mitigate potential exploitation of this vulnerability. Disabling VFIO drivers if not needed is also recommended as a defense-in-depth measure.\u001b[0m\n",
94+
"\n",
95+
"\u001b[1m> Finished chain.\u001b[0m\n",
96+
"{'question': 'What is CVE-2020-17087?', 'cve_id': 'CVE-2020-17087', 'output': \"\\n\\nCVE-2020-17087 is a vulnerability in the Linux kernel's VFIO device driver that could allow a privileged user to escalate privileges or cause a denial of service on the system.\\n\\nThe VFIO (Virtual Function I/O) driver allows safe and efficient device access from userspace on Linux hosts. This vulnerability is a use-after-free issue that could occur when removing a VFIO device. An unprivileged user with access to a VFIO device could potentially exploit this to gain elevated privileges on the system.\\n\\nSome key details about CVE-2020-17087:\\n\\n- Published: 2020-08-12\\n- CVSS Score: 7.8 (High severity)\\n- Affects: Linux kernel versions 5.4 through 5.8.5\\n- Type: Use-after-free vulnerability in vfio_unbound_driver()\\n- Impact: Privilege escalation, denial of service\\n\\nThe vulnerability was discovered by Cedric Hombracht and reported to Linux kernel security teams. It was patched in Linux kernel 5.8.6 and other stable kernel versions by validating the removal of VFIO device references.\\n\\nSystem administrators running affected Linux kernel versions should apply the patched updates from their Linux distributions to mitigate potential exploitation of this vulnerability. Disabling VFIO drivers if not needed is also recommended as a defense-in-depth measure.\"}\n"
97+
]
98+
}
99+
],
100+
"source": [
101+
"llm = create_llm()\n",
102+
"llm.temperature = 0.0\n",
103+
"agent = create_tool_calling_agent(llm, tools=tools, prompt=prompt)\n",
104+
"agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)\n",
105+
"content = agent_executor.invoke(\n",
106+
" {\"question\": \"What is CVE-2020-17087?\", \"cve_id\": \"CVE-2020-17087\"}\n",
107+
")\n",
108+
"print(content)\n"
109+
]
110+
}
111+
],
112+
"metadata": {
113+
"kernelspec": {
114+
"display_name": "base",
115+
"language": "python",
116+
"name": "python3"
117+
},
118+
"language_info": {
119+
"codemirror_mode": {
120+
"name": "ipython",
121+
"version": 3
122+
},
123+
"file_extension": ".py",
124+
"mimetype": "text/x-python",
125+
"name": "python",
126+
"nbconvert_exporter": "python",
127+
"pygments_lexer": "ipython3",
128+
"version": "3.11.5"
129+
}
130+
},
131+
"nbformat": 4,
132+
"nbformat_minor": 2
133+
}

0 commit comments

Comments
 (0)