utils.py

import datetime
import json
import logging
import os
import re
from abc import ABC, abstractmethod

import dirtyjson
import hjson
import numpy as np
import openai
from fuzzywuzzy import process
from sklearn.metrics.pairwise import cosine_similarity

api_configs = {
    "SambaNova": {
        "api_key": os.environ.get("SAMBANOVA_API_KEY"),
        "url_base": "https://api.sambanova.ai/v1"
    },
    "Together": {
        "api_key": os.environ.get("TOGETHER_API_KEY"),
        "url_base": "https://api.together.xyz/v1"
    }
    # You can add more API configurations here for other providers
}

class Agent(ABC):
    def __init__(self, prompt_template, llm_client):
        self.prompt_template = prompt_template
        self.llm_client = llm_client

    @abstractmethod
    def generate_response(self, prompt):
        pass


def setup_logging():
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    log_filename = f"logs/swiftsage_log_{timestamp}.txt"
    
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        filename=log_filename,
        filemode='w'
    )
    
    # Also print to console
    console = logging.StreamHandler()
    console.setLevel(logging.INFO)
    formatter = logging.Formatter('%(name)-12s: %(levelname)-8s %(message)s')
    console.setFormatter(formatter)
    logging.getLogger('').addHandler(console)

    return logging.getLogger('SwiftSage')

 

def extract_and_parse_markup(text):
    keys = ["reasoning_steps", "final_answer", "feedback", "score", "critical_feedback", "revised_plan", "solved", "plan", "code"]
    result = {}
    if "<final_answer>" in text and "</final_answer>" not in text:
        text = text + "</final_answer>"

    for key in keys:
        # Create a pattern for each key
        pattern = f'<{key}>(.*?)</{key}>'
        
        # Search for the pattern in the text
        match = re.search(pattern, text, re.DOTALL)
        
        if match:
            # Extract the content, strip whitespace, and add to the result
            content = match.group(1).strip()
            result[key] = content

    if "code" in result.keys():
        result["code"] = result["code"].replace("```python", "").replace("```", "").strip()

    return result


class PromptTemplate:
    def __init__(self, template_dir):
        self.template_dir = template_dir
        self.templates = {}
        self.load_templates()

    def load_templates(self):
        for filename in ['swift_template.md', 'sage_template.md', 'reward_template.md']:
            with open(os.path.join(self.template_dir, filename), 'r') as f:
                key = filename.split('_')[0]
                self.templates[key] = f.read()

    def format(self, key, **kwargs):
        template = self.templates.get(key, "")
        for k, v in kwargs.items():
            template = template.replace("<" + k + ">", str(v)) 
        return template


class LLMClient:
    def __init__(self, model_id, api_config, temperature=0.3, top_p=1.0, max_tokens=3000, logger=None):
        self.client = openai.OpenAI(
            api_key=api_config['api_key'],
            base_url=api_config['url_base']
        )
        self.model_id = model_id
        self.temperature = temperature
        self.top_p = top_p
        self.max_tokens = max_tokens
        self.logger = logger

    def generate_response(self, messages):
        self.logger.info(f"Sending request to {self.model_id}")
        self.logger.info(f"Messages: {messages}")
        response = self.client.chat.completions.create(
            model=self.model_id,
            messages=messages,
            temperature=self.temperature,
            top_p=self.top_p,
            max_tokens=self.max_tokens
        )
        content = response.choices[0].message.content
        self.logger.info(f"Response from {self.model_id}:\n{content}")
        return content





if __name__ == "__main__":
    test_text = "test"
     
    print(extract_and_parse_markup(test_text))



"""

def extract_and_parse_json(text):

    keys_and_types = [
        ("reasoning_steps", list),
        ("final_answer", str),
        ("feedback", str),
        ("score", str),
        ("score", int),
        ("feedback", str),
        ("solved", str),
        ("critical_feedback", str),
        ("revised_plan", list),
    ]

    # Try to parse the JSON first
    try:
        # find the first and last curly braces and parse the json
        first_brace = text.find("{")
        last_brace = text.rfind("}")  
        if last_brace == -1:
            text = text + "\"}"
        if first_brace != -1 and last_brace != -1 and first_brace < last_brace:
            data = json.loads(text[first_brace:last_brace+1])
        return data
    except Exception as e:
        data = {}
        try: 
            data = dirtyjson.loads(text) 
        except Exception as e:
            pass
        # If JSON parsing fails, use regex to extract key-value pairs
        
        for key, _ in keys_and_types:
            # pattern = rf'"{key}"\s*:\s*([\[{{].*?[\]}}]|".*?")'
            pattern = rf'"{key}"\s*:\s*([\[{{].*?[\]}}]|".*?"|[-+]?\d+)'
            match = re.search(pattern, text, re.DOTALL)
            if match:
                try:
                    value = json.loads(match.group(1))
                except Exception as e:
                    value = match.group(1).strip('"')
                data[key] = value

    result = {}
    for key, expected_type in keys_and_types:
        if key in result.keys() and result[key] is not None:
            continue
        # Use fuzzy matching to find the closest key
        try:
            closest_key, score = process.extractOne(key, data.keys())
        except Exception as e:
            continue
        if score > 80:  # You can adjust this threshold
            value = data[closest_key]
            
            # Type checking and conversion
            if expected_type == list and isinstance(value, str):
                value = [item.strip() for item in value.strip('[]').split(',')]
            elif expected_type == str and isinstance(value, list):
                value = ', '.join(value)
            elif expected_type == int and value is not None:
                try:
                    value = int(value)
                except ValueError:
                    value = None
            
            result[key] = value
        else:
            result[key] = None 
    
    for key in list(result.keys()):
        if result[key] is None:
            del result[key]
    return result

def extract_and_parse_json_v1(text):
    def find_json_objects(s):
        # Find all substrings that look like JSON objects
        json_like_strs = re.findall(r'\{(?:[^{}]|\{[^{}]*\})*\}', s)
        return json_like_strs

    def try_parse_json(s):
        try:
            return json.loads(s)
        except json.JSONDecodeError:
            try:
                s = s.replace("\n", "")
                return hjson.loads(s)
            except json.JSONDecodeError:
                return None
            return None 

    # First, try to find JSON within code blocks
    code_block_pattern = r'```(?:json)?\s*([\s\S]*?)\s*```'
    code_blocks = re.findall(code_block_pattern, text, re.IGNORECASE)
    
    all_json_candidates = []
    
    # Add JSON candidates from code blocks
    for block in code_blocks:
        all_json_candidates.extend(find_json_objects(block))
    
    # Add JSON candidates from the entire text
    all_json_candidates.extend(find_json_objects(text))
    
    # Sort candidates by length, descending
    all_json_candidates.sort(key=len, reverse=True)
    
    # Try to parse each candidate
    for candidate in all_json_candidates:
        parsed_json = try_parse_json(candidate)
        if parsed_json is not None:
            return parsed_json
    
    raise ValueError("No valid JSON object found in the text")

 
 

"""