use token limit to shrink report content (#24)

Co-authored-by: xuyang1 <xuyang1@microsoft.com>

rdagent/core/conf.py

@@ -42,6 +42,10 @@ class FincoSettings(BaseSettings):
     chat_frequency_penalty: float = 0.0
     chat_presence_penalty: float = 0.0
 
+    chat_token_limit: int = (
+        100000  # 100000 is the maximum limit of gpt4, which might increase in the future version of gpt
+    )
+
     default_system_prompt: str = "You are an AI assistant who helps to answer user's questions about finance."
 
     # Embedding configs

rdagent/document_process/document_analysis.py

@@ -69,14 +69,17 @@ def classify_report_from_dict(
             res_dict[file_name] = {"class": 0}
             continue
 
-        if not any(substring in content for substring in substrings):
+        if not any(substring in content for substring in substrings) and False:
             res_dict[file_name] = {"class": 0}
         else:
-            gpt_4_max_token = 128000
-            if input_max_token < gpt_4_max_token:
-                content = enc.encode(content)
-                max_token_1 = max(0, min(len(content), input_max_token) - 1)
-                content = enc.decode(content[:max_token_1])
+            while (
+                APIBackend().build_messages_and_calculate_token(
+                    user_prompt=content,
+                    system_prompt=classify_prompt,
+                )
+                > Config().chat_token_limit
+            ):
+                content = content[: -(Config().chat_token_limit // 100)]
 
             vote_list = []
             for _ in range(vote_time):
@@ -450,6 +453,8 @@ def __kmeans_embeddings(embeddings: np.ndarray, k: int = 20) -> list[list[str]]:
 
 
 def __deduplicate_factor_dict(factor_dict: dict[str, dict[str, str]]) -> list[list[str]]:
+    if len(factor_dict) == 0:
+        return []
     factor_df = pd.DataFrame(factor_dict).T
     factor_df.index.names = ["factor_name"]
 

rdagent/document_process/prompts.yaml

@@ -75,7 +75,7 @@ extract_factor_formulation_user: |-
     ===========================Factor list in dataframe=============================
     {{ factor_dict }}
 
-classify_system: |-
+classify_system_chinese: |-
     你是一个研报分类助手。用户会输入一篇金融研报。请按照要求回答：
     因子指能够解释资产收益率或价格等的变量；而模型则指机器学习或深度学习模型，利用因子等变量来预测价格或收益率变化。
 
@@ -86,6 +86,20 @@ classify_system: |-
 
     请使用json进行回答。json key为：class
 
+classify_system: |-
+    Your job is classify whether the user input document is a quantitative investment research report. The user will input a document and you should classify it based on the following conditions:
+    1. The document is about finance other than other fields like biology, physics, chemistry, etc.
+    2. The document is a research report on stock selection (which needs to be strictly separated from time selection and base selection) in the field of metalworking quantification.
+    3. The document involves the composition of factors or models, or tests their performance.
+
+    If the document meets all the above conditions, please return 1; otherwise, please return 0.
+
+    Please respond with your decision in JSON format. Just respond the output json string without any interaction and explanation.
+    The JSON schema should include:
+    {
+        "class": 1
+    }
+
 factor_viability_system: |-
     User has designed several factors in quant investment. Please help the user to check the viability of these factors.
     These factors are used to build a daily frequency strategy in China A-share market.

rdagent/factor_implementation/share_modules/factor_implementation_config.py

@@ -29,10 +29,6 @@ class FactorImplementSettings(FincoSettings):
     v2_error_summary: bool = False
     v2_knowledge_sampler: float = 1.0
 
-    chat_token_limit: int = (
-        100000  # 100000 is the maximum limit of gpt4, which might increase in the future version of gpt
-    )
-
     implementation_factors_per_round: int = 100  # how many factors to choose for each round of evolving
     evo_multi_proc_n: int = 16  # how many processes to use for evolving (including eval & generation)