methods/icl_socratic/socratic_questioning_textOnly/code/socratic_tree.py

from socratic_node import SocraticNode

import os
import json
import re
import pdb
import torch
import time
from openai import OpenAI
import requests

client = OpenAI(
        api_key=os.environ.get("CHATANYWHERE_API_KEY", ""),
        base_url='https://api.chatanywhere.tech/v1', 
    )

# client = OpenAI(api_key=os.environ.get("DEEPSEEK_API_KEY", ""), base_url="https://api.deepseek.com")

# [ragqa] 可插拔 LLM：设为 callable(system_define:str, prompt:str)->str 即用它替代默认的 gpt-4o-mini。
_LLM_OVERRIDE = None


def set_llm(fn):
    """注入自定义 LLM（如本地模型）。传 None 恢复默认 gpt-4o-mini。"""
    global _LLM_OVERRIDE
    _LLM_OVERRIDE = fn


# [ragqa] 可插拔检索器：设为 callable(query:str)->str（top-1 passage 文本）即用它替代默认 HTTP(8895)。
# 进程内函数版默认注入空检索桩，从而完全不发 HTTP；纯查询分解不需要真实检索。
_RETRIEVER_OVERRIDE = None


def set_retriever(fn):
    """注入自定义检索器。传 None 恢复默认 HTTP ColBERT。"""
    global _RETRIEVER_OVERRIDE
    _RETRIEVER_OVERRIDE = fn


def get_retrieval(query_item):
        if _RETRIEVER_OVERRIDE is not None:
            return _RETRIEVER_OVERRIDE(query_item)
        url = 'http://localhost:8895/api/search?query='+query_item+'&k=1'
        response = requests.get(url=url)
        res_dic = response.json()
        corpus_list_topk = res_dic['topk']
        top1_passage = corpus_list_topk[0]['text']
        
        return top1_passage

class SocraticTree:
    def __init__(self, backbone, api, prompt_map, dataset, num_question, max_turn, max_depth, save_dir):
        self.backbone = backbone
        if backbone == 'gpt':
            pass
           
        elif backbone == 'falcon':
            # host model by text-generation: 
            from text_generation import Client
            self.pipeline = Client("http://127.0.0.1:8080")
            
        self.prompt_map = prompt_map
        self.dataset = dataset
        self.num_question = num_question
        self.max_turn = max_turn
        self.max_depth = max_depth
        
        self.root_node = None
        
        self.save_dir = save_dir    
        self.log_path = None
    
        
    def select_prompt(self, isMultipleChoice, type):
        try:
            # question_type = 'multipleChoice' if isMultipleChoice else 'normalQA'
            question_type = 'multipleChoice_cot' if isMultipleChoice else 'normalQA'
            prompt = self.prompt_map[type][question_type][self.dataset]['prompt']
            tip = self.prompt_map[type][question_type][self.dataset]['tip']
            system_define = self.prompt_map[type][question_type][self.dataset]['system_define']
        except Exception as e:
            print(e)
            pdb.set_trace()
        return prompt, tip, system_define
    
# while True:
#             try:
#                 global client
#                 response = client.chat.completions.create(
#                     model='deepseek-chat',

#                     messages=[
#                         {"role": "system", "content": system_define},
#                         {"role": "user", "content": prompt},
#                     ]
#                 )
#                 break
#             except Exception as e:
#                 # current time in human-readable format
#                 obj = time.localtime()
#                 t = time.asctime(obj)
#                 print(t, 'Request failed. Retrying...')
#                 pdb.set_trace()
#                 # time.sleep(10)
    def request_gpt(self, system_define, prompt):

        # [ragqa] 若注入了自定义 LLM，则用它（进程内本地模型等），否则走默认 gpt-4o-mini。
        if _LLM_OVERRIDE is not None:
            return _LLM_OVERRIDE(system_define, prompt).strip()

        response = client.chat.completions.create(
                    model='gpt-4o-mini',
                    messages=[
                        {"role": "system", "content": system_define},
                        {"role": "user", "content": prompt},
                    ]
                )
                
        return response.choices[0].message.content.strip()
    
    
    
    def request_falcon(self, system_define, prompt):
        input_str = system_define + '\n\n\n' + prompt
        answer = self.pipeline.generate(input_str, max_new_tokens=300).generated_text 
        return answer.strip()
    
    def request(self, system_define, prompt):
        if self.backbone == 'gpt':
            return self.request_gpt(system_define, prompt)
        elif self.backbone == 'falcon':
            return self.request_falcon(system_define, prompt)
    
    def qa2hint(self, question, answer):
        # request chatGPT
        system_define = """Imagine you are an editor. You are given a question-and-answer pair. You need to merge the question and answer into a statement sentence."""
        prompt = "Question: " + question + "\nAnswer: " + answer
        return self.request(system_define, prompt)
    
    def get_clean_answer(self, raw_answer):
        answer = raw_answer
        # remove " in the raw_answer string
        while '"' in answer:
            answer = answer.replace('"', '')
        if ': ' in answer:
            answer = answer.split(': ')[1]
        if 'the answer is ' in answer:
            idx = answer.index('the answer is ')
            answer = answer[idx+14]
        if 'Option ' in answer:
            answer = answer.split('Option ')[1][0]
        if len(answer) > 0 and answer[-1] == '.':
            answer = answer[:-1]
        if 'A. ' in answer:
            answer = 'A'
        elif 'B. ' in answer:
            answer = 'B'
        elif 'C. ' in answer:
            answer = 'C'
        elif 'D. ' in answer:
            answer = 'D'
        elif '. ' in raw_answer:
            answer = raw_answer.split('. ')[1]
        if len(answer) == 2 and answer[1] == '.':
            answer = answer[0]
        return answer
    
        
    def answer_question(self, node):
        passage = get_retrieval(node.question)
        
        system_define = "You are an expert in open domain Q&A and your task is to assess how relevant the retrieved passages are to the question. According to the degree of help of the retrieved corpus to answer the question, it can be divided into three correlation degrees: high, middle and low."
        tip = "Only allowed to answer high or middle or low. If you are not sure, choose the one you think is more likely."
        
        prompt = ''
        prompt += 'Question: '
        prompt += node.question + '\n'
        prompt += 'Passage: '
        prompt += passage + '\n'
        
        prompt += tip + '\n'
        prompt += 'Answer: '
        relevant = self.request(system_define, prompt)
        
        if 'high' in relevant or 'High' in relevant:
            confidence = 'high'
        elif 'middle' in relevant or 'Middle' in relevant:
            confidence = 'middle'
        elif 'low' in relevant or 'Low' in relevant:
            confidence = 'low'
        else:
            confidence = 'middle'
        
        return confidence
  
            
    
    def raise_question(self, node):
        system_define = "Imagine that you are a thoughtful and logical problem solver. You will get a question. However, this question is too complex or lacking information to answer. You need to break down the original question into several simpler subquestions to help the information retrieval system retrieve the relevant information. Important: Do not use pronouns or indefinite pronoun phrases in generated questions. The questions asked must be self-contained questions. Each question can contain only one parameter. Don't just ask yes/no questions."
        
        prompt = ''
        prompt += "For example, Question: Could the Great Wall of China connect the Dodgers to the White Sox? Note: The raised question has to be a self-contained question. Do not use pronouns or indefinite pronoun phrases in the generated questions. Copy context from the original question if needed. Deep Questions: 1. What is the most commonly cited figure for the total length of the Great Wall? 2. What is the straight-line distance between Chicago, Illinois, and Los Angeles, California?" + '\n'
        prompt += 'Question: '
        prompt += node.question + '\n'
        prompt += 'Deep Questions: '
        
        raw_answer = self.request(system_define, prompt)
        
        # post-process answer
        ls_questions = raw_answer.split('\n')
        ls_questions = [re.sub(r'^\d+\.\s+', '', q) for q in ls_questions]
        ls_questions = ls_questions[:self.num_question] if len(ls_questions) > self.num_question else ls_questions
        return ls_questions
    
    
        
    def self_questioning(self, node):
        confidence = self.answer_question(node)      
        if (confidence == 'high' or node.isLeaf) and (node.hasHint() or node.depth != 1):
            # self.log(node, confidence)
            return False, confidence
        else:
            ls_deepQuestion = self.raise_question(node)
            # self.log(node, confidence, ls_deepQuestion=ls_deepQuestion)
            return True, ls_deepQuestion
        
    
    def run(self, node, collected_questions=None):
        if collected_questions is None:
            collected_questions = []

        # self questioning
        continue_deeper, output = self.self_questioning(node)
        if continue_deeper:  # output is subquestion list
            for deep_question in output:
                # Collect the deep question
                collected_questions.append(deep_question)

                deep_node = SocraticNode(
                    deep_question,
                    node.turn,
                    node.depth + 1,
                    self.max_turn,
                    self.max_depth,
                    context=node.context
                )
                deep_answer = self.run(deep_node, collected_questions)
                node.add_hint(1)
            node.update_turn_num()
            return self.run(node, collected_questions)

        else:  # output is answer
            return collected_questions

        
        
    def start(self, id, question, options, context=None):
        self.log_path = self.save_dir + '/log/' + str(id) + '.txt'
        if os.path.exists(self.log_path):
            os.remove(self.log_path)
        self.root_node = SocraticNode(question, 1, 1, self.max_turn, self.max_depth, context=context, options=options)
        collected_questions = self.run(self.root_node)
        return collected_questions
    
        
    def log(self, node, raw_answer, ls_deepQuestion=None):
        turn = node.turn
        depth = node.depth
        isLeaf = node.isLeaf
        context = node.context
        question = node.question
        options = None
        if node.isMultipleChoice():
            options = node.add_optionID_toText()
        hints = node.get_textHints()
        raw_answer = raw_answer
        answer = str(node.answer)
        ls_deepQuestion = ls_deepQuestion       
        
        with open(self.log_path, 'a') as f:
            f.write('=========================================================================================================\n')
            f.write('Turn: ' + str(turn) + ', Depth: ' + str(depth) + ', isLeaf: ' + str(isLeaf) + '\n')
            if context is not None:
                f.write('Context:\t' + context + '\n')
            f.write('Question:\t' + question + '\n')
            if options is not None:
                f.write('Options:\t' + options + '\n')
            if len(hints) > 0:
                f.write('Hints:\t' + hints + '\n')
            f.write('Raw Answer:\t' + raw_answer + '\n')
            f.write('Answer:\t' + answer + '\n')
            if ls_deepQuestion is not None:
                f.write('\nDeep Questions:\n')
                for i, deepQuestion in enumerate(ls_deepQuestion):
                    f.write(str(i+1) + '. ' + deepQuestion + '\n')
            f.write('\n')