Files changed (1) hide show
  1. handler.py +60 -5
handler.py CHANGED
@@ -104,18 +104,73 @@ class EndpointHandler:
104
  # Split response into text and function call
105
  parts = response.split("Calling function:", 1)
106
  text_response = parts[0].strip()
107
- function_call = "Calling function:" + parts[1].strip()
108
 
109
  logger.info(f"Function call: {function_call}")
110
  logger.info(f"Text response: {text_response}")
111
- # Return both text and tool message
112
- return {"result": [{"text": text_response}]},
113
- #{"generations": [{"text": function_call, "type": "tool"}]}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
114
  else:
115
  response = output_text
116
 
117
  logger.info(f"Generated response: {json.dumps(response)}")
118
- return {"result": [{"text": response}]}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
119
 
120
  except Exception as e:
121
  logger.error(f"Error during generation: {str(e)}", exc_info=True)
 
104
  # Split response into text and function call
105
  parts = response.split("Calling function:", 1)
106
  text_response = parts[0].strip()
107
+ function_call = parts[1].strip()
108
 
109
  logger.info(f"Function call: {function_call}")
110
  logger.info(f"Text response: {text_response}")
111
+
112
+ # Return in Anthropic format with tool calls
113
+ return {
114
+ "result": {
115
+ "metadata": {
116
+ "finishReason": "end_turn",
117
+ "contentFilterMetadata": None
118
+ },
119
+ "output": {
120
+ "messageType": "ASSISTANT",
121
+ "metadata": {
122
+ "messageType": "ASSISTANT"
123
+ },
124
+ "toolCalls": [{
125
+ "type": "function",
126
+ "function": {
127
+ "name": function_call.split("(")[0].strip(),
128
+ "arguments": function_call.split("(")[1].rstrip(")")
129
+ }
130
+ }],
131
+ "content": text_response
132
+ }
133
+ },
134
+ "metadata": {
135
+ "id": f"msg_{int(time.time()*1000)}",
136
+ "model": "granite-3.1-8b",
137
+ "usage": {
138
+ "promptTokens": len(inputs["input_ids"][0]),
139
+ "generationTokens": len(output_tokens[0]) - len(inputs["input_ids"][0]),
140
+ "totalTokens": len(output_tokens[0])
141
+ }
142
+ }
143
+ }
144
  else:
145
  response = output_text
146
 
147
  logger.info(f"Generated response: {json.dumps(response)}")
148
+ # Return in Anthropic format without tool calls
149
+ return {
150
+ "result": {
151
+ "metadata": {
152
+ "finishReason": "end_turn",
153
+ "contentFilterMetadata": None
154
+ },
155
+ "output": {
156
+ "messageType": "ASSISTANT",
157
+ "metadata": {
158
+ "messageType": "ASSISTANT"
159
+ },
160
+ "toolCalls": [],
161
+ "content": response
162
+ }
163
+ },
164
+ "metadata": {
165
+ "id": f"msg_{int(time.time()*1000)}",
166
+ "model": "granite-3.1-8b",
167
+ "usage": {
168
+ "promptTokens": len(inputs["input_ids"][0]),
169
+ "generationTokens": len(output_tokens[0]) - len(inputs["input_ids"][0]),
170
+ "totalTokens": len(output_tokens[0])
171
+ }
172
+ }
173
+ }
174
 
175
  except Exception as e:
176
  logger.error(f"Error during generation: {str(e)}", exc_info=True)