@@ -67,16 +67,26 @@ pub(crate) async fn sagemaker_compatibility(
6767 default_return_full_text : Extension < bool > ,
6868 infer : Extension < Infer > ,
6969 compute_type : Extension < ComputeType > ,
70+ context : Extension < Option < opentelemetry:: Context > > ,
7071 info : Extension < Info > ,
7172 Json ( req) : Json < SagemakerRequest > ,
7273) -> Result < Response , ( StatusCode , Json < ErrorResponse > ) > {
7374 match req {
7475 SagemakerRequest :: Generate ( req) => {
75- compat_generate ( default_return_full_text, infer, compute_type, Json ( req) ) . await
76+ compat_generate (
77+ default_return_full_text,
78+ infer,
79+ compute_type,
80+ context,
81+ Json ( req) ,
82+ )
83+ . await
84+ }
85+ SagemakerRequest :: Chat ( req) => {
86+ chat_completions ( infer, compute_type, info, context, Json ( req) ) . await
7687 }
77- SagemakerRequest :: Chat ( req) => chat_completions ( infer, compute_type, info, Json ( req) ) . await ,
7888 SagemakerRequest :: Completion ( req) => {
79- completions ( infer, compute_type, info, Json ( req) ) . await
89+ completions ( infer, compute_type, info, context , Json ( req) ) . await
8090 }
8191 }
8292}
0 commit comments