use token limit to shrink report content (#24)

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

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.