okr/utils.py
2025-01-03 11:28:40 +01:00

72 lines
2.2 KiB
Python

import json
import streamlit as st
import re
# Function to construct the prompt
def construct_prompt(prompt_template: str, user_input: str) -> str:
return prompt_template.format(user_input=user_input)
def parse_json_content(cleaned_content: str):
"""
Parses the cleaned content to extract valid JSON data.
Args:
cleaned_content (str): The raw content containing JSON data.
Returns:
dict or list: The parsed JSON object.
"""
import re
# Step 1: Strip unwanted characters and clean the content
cleaned_content = cleaned_content.strip()
# Step 2: Use regex to extract only the valid JSON block (e.g., starts with [ or {)
json_match = re.search(r"(\{.*\}|\[.*\])", cleaned_content, re.DOTALL)
if not json_match:
raise ValueError("No valid JSON found in the content.")
# Step 3: Extract and parse the valid JSON
valid_json = json_match.group(0) # Extract matched JSON block
try:
extracted_data = json.loads(valid_json)
except json.JSONDecodeError as e:
raise ValueError(f"Failed to decode JSON. Error: {e}\nContent:\n{valid_json}")
return extracted_data
# Function to extract and parse JSON response
def extract_llm_response(response):
"""
Extracts and parses the JSON response from the API.
Args:
response (dict): The API response containing a hint and proposals.
Returns:
tuple: A tuple containing the objective (str), key results (list), and hint (str).
"""
raw_message_content = response["choices"][0]["message"]["content"]
# Clean and parse the JSON content
cleaned_content = raw_message_content.replace("`", "").split("json")[-1]
parsed_data = parse_json_content(cleaned_content=cleaned_content)
hint = parsed_data.get("hint", "")
proposals = parsed_data.get("proposal", [])
if proposals:
# Extract the first proposal's objective and key results
first_proposal = proposals[0] # Get the first proposal (assuming it's a list)
objective = first_proposal.get("objective", "")
key_results = first_proposal.get("key_results", [])
else:
objective = ""
key_results = []
return objective, key_results, hint