feat: add type checker to api backend & align litellm and old backend (#647)

* move cache auto continue and retry to all api backend

* add type checker to json mode output

* fix CI

* feat: Add json_mode handling and streaming support in chat completion function

* lint

* fix a bug when returning a dict which value could contain int or bool

* remove litellm

---------

Co-authored-by: Xu Yang <xuyang1@microsoft.com>
Co-authored-by: Young <afe.young@gmail.com>
This commit is contained in:
Xu Yang
2025-02-28 15:13:43 +08:00
committed by GitHub
parent a7f0325447
commit 993ca37125
25 changed files with 634 additions and 606 deletions
@@ -6,7 +6,7 @@ import random
import re
from itertools import combinations
from pathlib import Path
from typing import Union
from typing import List, Union
from jinja2 import Environment, StrictUndefined
@@ -339,6 +339,7 @@ class CoSTEERRAGStrategyV2(RAGStrategy):
system_prompt=analyze_component_system_prompt,
user_prompt=analyze_component_user_prompt,
json_mode=True,
json_target_type=List[int],
),
)["component_no_list"]
return [all_component_nodes[index - 1] for index in sorted(list(set(component_no_list)))]
@@ -12,6 +12,7 @@ File structure
"""
import json
from typing import Dict
from rdagent.components.coder.CoSTEER import CoSTEER
from rdagent.components.coder.CoSTEER.evaluators import (
@@ -85,7 +86,10 @@ class EnsembleMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
for _ in range(5):
ensemble_code = json.loads(
APIBackend().build_messages_and_create_chat_completion(
user_prompt=user_prompt, system_prompt=system_prompt, json_mode=True
user_prompt=user_prompt,
system_prompt=system_prompt,
json_mode=True,
json_target_type=Dict[str, str],
)
)["code"]
if ensemble_code != workspace.file_dict.get("ensemble.py"):
@@ -1,4 +1,5 @@
import json
from typing import Dict
from rdagent.components.coder.CoSTEER import CoSTEER
from rdagent.components.coder.CoSTEER.evaluators import (
@@ -70,7 +71,10 @@ class FeatureMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
for _ in range(5):
feature_code = json.loads(
APIBackend().build_messages_and_create_chat_completion(
user_prompt=user_prompt, system_prompt=system_prompt, json_mode=True
user_prompt=user_prompt,
system_prompt=system_prompt,
json_mode=True,
json_target_type=Dict[str, str],
)
)["code"]
if feature_code != workspace.file_dict.get("feature.py"):
@@ -1,3 +1,5 @@
from typing import Dict
from rdagent.components.coder.CoSTEER import CoSTEER
from rdagent.components.coder.CoSTEER.evaluators import (
CoSTEERMultiEvaluator,
@@ -83,6 +85,7 @@ class ModelMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
user_prompt=user_prompt,
system_prompt=system_prompt,
json_mode=BatchEditOut.json_mode,
json_target_type=Dict[str, str],
)
)
@@ -24,6 +24,7 @@ File structure
import json
import re
from typing import Dict
from rdagent.app.data_science.conf import DS_RD_SETTING
from rdagent.components.coder.CoSTEER import CoSTEER
@@ -108,20 +109,30 @@ class DataLoaderMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
spec_session = APIBackend().build_chat_session(session_system_prompt=system_prompt)
data_loader_spec = json.loads(
spec_session.build_chat_completion(user_prompt=data_loader_prompt, json_mode=True)
spec_session.build_chat_completion(
user_prompt=data_loader_prompt, json_mode=True, json_target_type=Dict[str, str]
)
)["spec"]
feature_spec = json.loads(
spec_session.build_chat_completion(
user_prompt=feature_prompt, json_mode=True, json_target_type=Dict[str, str]
)
)["spec"]
model_spec = json.loads(
spec_session.build_chat_completion(
user_prompt=model_prompt, json_mode=True, json_target_type=Dict[str, str]
)
)["spec"]
ensemble_spec = json.loads(
spec_session.build_chat_completion(
user_prompt=ensemble_prompt, json_mode=True, json_target_type=Dict[str, str]
)
)["spec"]
workflow_spec = json.loads(
spec_session.build_chat_completion(
user_prompt=workflow_prompt, json_mode=True, json_target_type=Dict[str, str]
)
)["spec"]
feature_spec = json.loads(spec_session.build_chat_completion(user_prompt=feature_prompt, json_mode=True))[
"spec"
]
model_spec = json.loads(spec_session.build_chat_completion(user_prompt=model_prompt, json_mode=True))[
"spec"
]
ensemble_spec = json.loads(spec_session.build_chat_completion(user_prompt=ensemble_prompt, json_mode=True))[
"spec"
]
workflow_spec = json.loads(spec_session.build_chat_completion(user_prompt=workflow_prompt, json_mode=True))[
"spec"
]
else:
data_loader_spec = workspace.file_dict["spec/data_loader.md"]
feature_spec = workspace.file_dict["spec/feature.md"]
@@ -146,7 +157,10 @@ class DataLoaderMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
for _ in range(5):
data_loader_code = json.loads(
APIBackend().build_messages_and_create_chat_completion(
user_prompt=user_prompt, system_prompt=system_prompt, json_mode=True
user_prompt=user_prompt,
system_prompt=system_prompt,
json_mode=True,
json_target_type=Dict[str, str],
)
)["code"]
if data_loader_code != workspace.file_dict.get("load_data.py"):
@@ -1,4 +1,5 @@
import json
from typing import Dict
from rdagent.components.coder.CoSTEER import CoSTEER
from rdagent.components.coder.CoSTEER.evaluators import (
@@ -73,7 +74,10 @@ class WorkflowMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
for _ in range(5):
workflow_code = json.loads(
APIBackend().build_messages_and_create_chat_completion(
user_prompt=user_prompt, system_prompt=system_prompt, json_mode=True
user_prompt=user_prompt,
system_prompt=system_prompt,
json_mode=True,
json_target_type=Dict[str, str],
)
)["code"]
if workflow_code != workspace.file_dict.get("main.py"):
@@ -2,7 +2,7 @@ import io
import json
from abc import abstractmethod
from pathlib import Path
from typing import Tuple
from typing import Dict, Tuple
import pandas as pd
from jinja2 import Environment, StrictUndefined
@@ -212,7 +212,10 @@ class FactorOutputFormatEvaluator(FactorEvaluator):
try:
api = APIBackend() if attempts == 0 else APIBackend(use_chat_cache=False)
resp = api.build_messages_and_create_chat_completion(
user_prompt=gen_df_info_str, system_prompt=system_prompt, json_mode=True
user_prompt=gen_df_info_str,
system_prompt=system_prompt,
json_mode=True,
json_target_type=Dict[str, str | bool | int],
)
resp_dict = json.loads(resp)
resp_dict["output_format_decision"] = str(resp_dict["output_format_decision"]).lower() in ["true", "1"]
@@ -556,6 +559,7 @@ class FactorFinalDecisionEvaluator(FactorEvaluator):
system_prompt=system_prompt,
json_mode=True,
seed=attempts, # in case of useless retrying when cache enabled.
json_target_type=Dict[str, str | bool | int],
),
)
final_decision = final_evaluation_dict["final_decision"]
@@ -2,6 +2,7 @@ from __future__ import annotations
import json
from pathlib import Path
from typing import Dict
from jinja2 import Environment, StrictUndefined
@@ -168,7 +169,10 @@ class FactorMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
APIBackend(
use_chat_cache=FACTOR_COSTEER_SETTINGS.coder_use_cache
).build_messages_and_create_chat_completion(
user_prompt=user_prompt, system_prompt=system_prompt, json_mode=True
user_prompt=user_prompt,
system_prompt=system_prompt,
json_mode=True,
json_target_type=Dict[str, str],
)
)["code"]
return code
@@ -1,6 +1,6 @@
import json
from pathlib import Path
from typing import Tuple
from typing import Dict, Tuple
import numpy as np
from jinja2 import Environment, StrictUndefined
@@ -177,6 +177,7 @@ class ModelFinalEvaluator(CoSTEEREvaluator):
user_prompt=user_prompt,
system_prompt=system_prompt,
json_mode=True,
json_target_type=Dict[str, str | bool | int],
),
)
if isinstance(final_evaluation_dict["final_decision"], str) and final_evaluation_dict[
@@ -1,5 +1,6 @@
import json
from pathlib import Path
from typing import Dict
from jinja2 import Environment, StrictUndefined
@@ -96,6 +97,7 @@ class ModelMultiProcessEvolvingStrategy(MultiProcessEvolvingStrategy):
user_prompt=user_prompt,
system_prompt=system_prompt,
json_mode=True,
json_target_type=Dict[str, str],
),
)["code"]
return code