45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
import ollama
|
|
import re
|
|
# def clean_json_response(raw_string):
|
|
# """
|
|
# Removes Markdown code blocks and extra whitespace from AI responses.
|
|
# """
|
|
# # Use regex to find content between ```json and ``` or just ```
|
|
# match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', raw_string)
|
|
# if match:
|
|
# return match.group(1).strip()
|
|
# return raw_string.strip()
|
|
|
|
import json
|
|
import re
|
|
|
|
def robust_json_parser(raw_output):
|
|
# 1. Strip Markdown blocks
|
|
clean = re.sub(r'```(?:json)?|```', '', raw_output).strip()
|
|
|
|
# 2. Fix trailing commas before closing braces/brackets
|
|
clean = re.sub(r',\s*([\]}])', r'\1', clean)
|
|
|
|
try:
|
|
return json.loads(clean)
|
|
except json.JSONDecodeError:
|
|
# 3. Last resort: try to find the first '{' and last '}'
|
|
start_idx = clean.find('{')
|
|
end_idx = clean.rfind('}')
|
|
if start_idx != -1 and end_idx != -1:
|
|
try:
|
|
return json.loads(clean[start_idx:end_idx+1])
|
|
except:
|
|
pass
|
|
raise
|
|
|
|
def get_model_reponse(prompt):
|
|
response=ollama.chat(
|
|
model='alibayram/smollm3:latest',
|
|
messages=[{'role': 'user', 'content': prompt}],
|
|
stream=False # Set to True for real-time streaming
|
|
)
|
|
# print(response['message']['content'])
|
|
response=robust_json_parser(response['message']['content'])
|
|
return response
|