Update parse.py

parse.py

@@ -1,6 +1,13 @@
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage
 import os
+import pandas as pd
+
+from huggingface_hub import login
+
+login("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")
+
 
 # Load OpenRouter API Key
 openrouter_api_key = "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"
@@ -14,18 +21,18 @@ model = ChatOpenAI(
 # Create a chat prompt template
 template = (
     "You are tasked with extracting specific information from the following text content: {dom_content}. "
-    "Please follow these instructions carefully: \n\n"
-    "1. **Extract Information:** Only extract the information that directly matches the provided description: {parse_description}. "
-    "2. **No Extra Content:** Do not include any additional text, comments, or explanations in your response. "
-    "3. **Empty Response:** If no information matches the description, return an empty string ('')."
-    "4. **Direct Data Only:** Your output should contain only the data that is explicitly requested, with no other text."
-    "5. **Type:** The output should always be a table, and if there's more than one table, return every table separately. Use Markdown table format.\n"
-    "6. **Standardized Table Format:** Ensure each table is formatted as a Markdown table with clear headers and consistent column alignment.\n"
-    "7. **Accuracy:** The output should be as accurate as possible.\n"
-    "8. **Column Separators:** Use the pipe symbol (|) to clearly separate columns in the Markdown table.\n"
-    "9. **Header Row:** The first row of each table should be the header row, clearly labeling each column.\n"
-    "10. **Alignment Row:** The second row should contain hyphens (-) to indicate column alignment (e.g., --- for left alignment, :---: for center alignment, ---: for right alignment).\n"
-    "11. **Data Rows:** Subsequent rows should contain the data, with each cell aligned according to the alignment row.\n"
+    "Please follow these instructions carefully:\n\n"
+    "1. **Task:** Extract data from the provided text that matches the description: {parse_description}.\n"
+    "2. **Output Format:** Return the extracted data ONLY as one or more Markdown tables. Each table MUST be correctly formatted.\n"
+    "3. **Markdown Table Format:** Each table must adhere to the following Markdown format:\n"
+    "   - Start with a header row, clearly labeling each column, separated by pipes (|).\n"
+    "   - Follow the header row with an alignment row, using hyphens (-) to indicate column alignment (e.g., --- for left alignment).\n"
+    "   - Subsequent rows should contain the data, with cells aligned according to the alignment row.\n"
+    "   - Use pipes (|) to separate columns in each data row.\n"
+    "4. **No Explanations:** Do not include any introductory or explanatory text before or after the table(s).\n"
+    "5. **Empty Response:** If no information matches the description, return an empty string ('').\n"
+    "6. **Multiple Tables:** If the text contains multiple tables matching the description, return each table separately, following the Markdown format for each.\n"
+    "7. **Accuracy:** The extracted data must be accurate and reflect the information in the provided text.\n"
 )
 
 # Function to parse and extract information from the chunks
@@ -45,3 +52,35 @@ def parse(dom_chunks, parse_description):
 
     # Return the parsed results as a single string
     return "\n".join(parsed_results)
+
+def merge_tables_with_llm(tables, parse_description):
+    """Merges a list of Pandas DataFrames into a single Markdown table using LLM."""
+    from langchain_core.prompts import ChatPromptTemplate
+    from langchain_openai import ChatOpenAI
+    # Convert DataFrames to Markdown strings
+    table_strings = [table.to_markdown(index=False) for table in tables]
+
+    # Create a prompt for the LLM
+    merge_prompt = (
+        "You are tasked with merging the following Markdown tables into a single, comprehensive Markdown table.\n"
+        "The tables contain information related to: {parse_description}.\n"
+        "Please follow these instructions carefully:\n\n"
+        "1. **Task:** Merge the data from the following tables into a single table that matches the description: {parse_description}.\n"
+        "2. **Output Format:** Return the merged data ONLY as a single Markdown table. The table MUST be correctly formatted.\n"
+        "3. **Markdown Table Format:** The table must adhere to the following Markdown format:\n"
+        "   - Start with a header row, clearly labeling each column, separated by pipes (|).\n"
+        "   - Follow the header row with an alignment row, using hyphens (-) to indicate column alignment (e.g., --- for left alignment).\n"
+        "   - Subsequent rows should contain the data, with cells aligned according to the alignment row.\n"
+        "   - Use pipes (|) to separate columns in each data row.\n"
+        "4. **No Explanations:** Do not include any introductory or explanatory text before or after the table.\n"
+        "5. **Empty Response:** If no information matches the description, return an empty string ('') if no data can be merged.\n"
+        "6. **Duplicate Columns:** If there are duplicate columns, rename them to be unique.\n"
+        "7. **Missing Values:** If there are missing values, fill them with 'N/A'.\n\n"
+        "Here are the tables:\n\n" + "\n\n".join(table_strings) +
+        "\n\nReturn the merged table in Markdown format:"
+    )
+
+    # Invoke the LLM
+    message = HumanMessage(content=merge_prompt)
+    response = model.invoke([message])
+    return response.content