|
|
import json |
|
|
from pathlib import Path |
|
|
from typing import List |
|
|
from functools import partial |
|
|
from .base_agent import BaseAgent |
|
|
from prompt.template import METHOD_CRITIQUE_PROMPT |
|
|
from utils.convert_format import markdown_to_json_method |
|
|
from utils.utils import parse_llm_output_to_json |
|
|
from utils.embedding import EmbeddingScorer |
|
|
|
|
|
import json |
|
|
|
|
|
|
|
|
|
|
|
class MethodScorer: |
|
|
|
|
|
def __init__(self, score_func, parent_weight=0.5, child_weight=0.5): |
|
|
self.parent_weight = parent_weight |
|
|
self.child_weight = child_weight |
|
|
self.score_func = score_func |
|
|
self.leaves = [] |
|
|
|
|
|
def process(self, data): |
|
|
self.leaves = [] |
|
|
for root_node in data: |
|
|
self._process_node(root_node, parent_scores=[]) |
|
|
for root_node in data: |
|
|
self._collect_leaves(root_node) |
|
|
return self.leaves |
|
|
|
|
|
def _process_node(self, node, parent_scores): |
|
|
if 'children' in node: |
|
|
children = node.get('children', []) |
|
|
if children: |
|
|
first_child = children[0] |
|
|
if 'method_class' in first_child: |
|
|
input_for_llm = [{"method": child["method_class"], "description": child.get("description", "")} for child in children] |
|
|
llm_result = self.score_func(input_for_llm) |
|
|
for idx, child in enumerate(children): |
|
|
if idx < len(llm_result): |
|
|
child['score'] = llm_result[idx]['score'] |
|
|
else: |
|
|
child['score'] = 0 |
|
|
current_score = node.get('score') |
|
|
new_parent = parent_scores.copy() |
|
|
if current_score is not None: |
|
|
new_parent.append(current_score) |
|
|
for child in children: |
|
|
self._process_node(child, new_parent) |
|
|
else: |
|
|
input_for_llm = [{"method": child["method"], "description": child.get("description", "")} for child in children] |
|
|
llm_result = self.score_func(input_for_llm) |
|
|
for idx, child in enumerate(children): |
|
|
if idx < len(llm_result): |
|
|
child_score = llm_result[idx]['score'] |
|
|
else: |
|
|
child_score = 0 |
|
|
child['score'] = child_score |
|
|
parent_avg = sum(parent_scores) / len(parent_scores) if parent_scores else 0 |
|
|
final_score = parent_avg * self.parent_weight + child_score * self.child_weight |
|
|
child['final_score'] = final_score |
|
|
|
|
|
def _collect_leaves(self, node): |
|
|
if 'children' in node: |
|
|
for child in node['children']: |
|
|
self._collect_leaves(child) |
|
|
else: |
|
|
if 'final_score' in node: |
|
|
self.leaves.append({ |
|
|
"method": node["method"], |
|
|
"description": node.get("description", ""), |
|
|
"score": node['final_score'] |
|
|
}) |
|
|
|
|
|
|
|
|
|
|
|
class MethodRanking(BaseAgent): |
|
|
def __init__(self, llm, rag=True): |
|
|
super().__init__(llm) |
|
|
self.rag = rag |
|
|
self.embedding_scorer = EmbeddingScorer() |
|
|
current_file = Path(__file__).resolve() |
|
|
json_path = current_file.parent.parent.parent / 'data/actor_data/docs/method_en_v1.json' |
|
|
md_path = current_file.parent.parent.parent / 'data/actor_data/docs/method_en_v1.md' |
|
|
|
|
|
with open(str(md_path), "r", encoding="utf-8") as f: |
|
|
self.markdown_text = f.read() |
|
|
self.method_tree = markdown_to_json_method(self.markdown_text) |
|
|
with open(json_path, "w+", encoding="utf-8") as f: |
|
|
json.dump(self.method_tree, f, ensure_ascii=False, indent=4) |
|
|
|
|
|
def llm_score_method(self, problem_description: str, methods: List[dict]): |
|
|
methods_str = '\n'.join([f"{i+1}. {method['method']} {method.get('description', '')}" for i, method in enumerate(methods)]) |
|
|
prompt = METHOD_CRITIQUE_PROMPT.format(problem_description=problem_description, methods=methods_str) |
|
|
answer = self.llm.generate(prompt) |
|
|
method_scores = parse_llm_output_to_json(answer).get('methods', []) |
|
|
method_scores = sorted(method_scores, key=lambda x: x['method_index']) |
|
|
for method in method_scores: |
|
|
method['score'] = sum(method['scores'].values()) / len(method['scores']) |
|
|
|
|
|
return method_scores |
|
|
|
|
|
def format_methods(self, methods: List[str]): |
|
|
return '\n'.join([f"**{method['method']}:** {method['description']}" for method in methods]) |
|
|
|
|
|
def top_methods(self, problem_description: str, top_k: int=6, method: str='embedding'): |
|
|
if self.rag: |
|
|
if method == 'embedding': |
|
|
score_func = partial(self.embedding_scorer.score_method, problem_description) |
|
|
else: |
|
|
score_func = partial(self.llm_score_method, problem_description) |
|
|
method_scores = MethodScorer(score_func).process(self.method_tree) |
|
|
method_scores.sort(key=lambda x: x['score'], reverse=True) |
|
|
return self.format_methods(method_scores[:top_k]) |
|
|
else: |
|
|
return self.markdown_text |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
from input.test_middle_result import problem_str |
|
|
from llm.llm import LLM |
|
|
llm = LLM('deepseek-chat') |
|
|
|
|
|
mr = MethodRanking(llm) |
|
|
|
|
|
|
