PDF Extractor with Native Multi Step Tool Use

1
import os
2
3
import cohere
4
import pandas as pd
5
import json
6
from unstructured.partition.pdf import partition_pdf

1
# uncomment to install dependencies
2
# !pip install cohere unstructured

1
# versions
2
print('cohere version:', cohere.__version__)

cohere version: 5.5.1

1
COHERE_API_KEY = os.environ.get("CO_API_KEY")
2
COHERE_MODEL = 'command-r-plus'
3
co = cohere.Client(api_key=COHERE_API_KEY)

1
def convert_to_json(text: str) -> dict:
2
    """
3
    Given text files, convert to json object and saves to csv.
4
5
    Args:
6
        text (str): The text to extract information from.
7
8
    Returns:
9
        dict: A dictionary containing the result of the conversion process.
10
    """
11
12
    MANDATORY_FIELDS = [
13
        "total_amount",
14
        "invoice_number",
15
    ]
16
17
    message = """# Instruction
18
    Given the text, convert to json object with the following keys:
19
    total_amount, invoice_number
20
21
    # Output format json:
22
    {{
23
        "total_amount": "<extracted amount="" invoice="" total="">",
24
        "invoice_number": "<extracted invoice="" number="">",
25
    }}
26
27
    Do not output code blocks.
28
29
    # Extracted PDF
30
    {text}
31
    """
32
33
    result = co.chat(
34
        message=message.format(text=text), model=COHERE_MODEL, preamble=None
35
    ).text
36
37
    try:
38
        result = json.loads(result)
39
        # check if all keys are present
40
        if not all(i in result.keys() for i in MANDATORY_FIELDS):
41
            return {"result": f"ERROR: Keys are missing. Please check your result {result}"}
42
43
        df = pd.DataFrame(result, index=[0])
44
        df.to_csv("output.csv", index=False)
45
        return {"result": "SUCCESS. All steps have been completed."}
46
47
    except Exception as e:
48
        return {"result": f"ERROR: Could not load the result as json. Please check the result: {result} and ERROR: {e}"}

1
def cohere_agent(
2
    message: str,
3
    preamble: str,
4
    verbose: bool = False,
5
) -> str:
6
    """
7
    Function to handle multi-step tool use api.
8
9
    Args:
10
        message (str): The message to send to the Cohere AI model.
11
        preamble (str): The preamble or context for the conversation.
12
        verbose (bool, optional): Whether to print verbose output. Defaults to False.
13
14
    Returns:
15
        str: The final response from the call.
16
    """
17
18
    functions_map = {
19
        "convert_to_json": convert_to_json,
20
    }
21
22
    tools = [
23
        {
24
            "name": "convert_to_json",
25
            "description": "Given a text, convert it to json object.",
26
            "parameter_definitions": {
27
                "text": {
28
                    "description": "text to be converted into json",
29
                    "type": "str",
30
                    "required": True,
31
                },
32
            },
33
        }
34
    ]
35
36
    counter = 1
37
38
    response = co.chat(
39
        model=COHERE_MODEL,
40
        message=message,
41
        preamble=preamble,
42
        tools=tools,
43
    )
44
45
    if verbose:
46
        print(f"\nrunning step 0")
47
        print(response.text)
48
49
    while response.tool_calls:
50
        tool_results = []
51
52
        if verbose:
53
            print(f"\nrunning step {counter}")
54
        for tool_call in response.tool_calls:
55
            print("tool_call.parameters:", tool_call.parameters)
56
            if tool_call.parameters:
57
                output = functions_map[tool_call.name](**tool_call.parameters)
58
            else:
59
                output = functions_map[tool_call.name]()
60
61
            outputs = [output]
62
            tool_results.append({"call": tool_call, "outputs": outputs})
63
64
            if verbose:
65
                print(
66
                    f"= running tool {tool_call.name}, with parameters: {tool_call.parameters}"
67
                )
68
                print(f"== tool results: {outputs}")
69
70
        response = co.chat(
71
            model=COHERE_MODEL,
72
            message="",
73
            chat_history=response.chat_history,
74
            preamble=preamble,
75
            tools=tools,
76
            tool_results=tool_results,
77
        )
78
79
        if verbose:
80
            print(response.text)
81
            counter += 1
82
83
    return response.text

1
def extract_pdf(path):
2
    """
3
    Function to extract text from a PDF file.
4
    """
5
    elements = partition_pdf(path)
6
    return "\n".join([str(el) for el in elements])
7
8
9
def pdf_extractor(pdf_path):
10
    """
11
    Main function that extracts pdf and calls the cohere agent.
12
    """
13
    pdf_text = extract_pdf(pdf_path)
14
15
    prompt = f"""
16
    # Instruction
17
    You are expert at extracting invoices from PDF. The text of the PDF file is given below.
18
19
    You must follow the steps below:
20
    1. Summarize the text and extract only the most information: total amount billed and invoice number.
21
    2. Using the summary above, call convert_to_json tool, which uses the summary from step 1.
22
    If you run into issues. Identifiy the issue and retry.
23
    You are not done unless you see SUCCESS in the tool output.
24
25
    # File Name:
26
    {pdf_path}
27
28
    # Extracted Text:
29
    {pdf_text}
30
    """
31
    output = cohere_agent(prompt, None, verbose=True)
32
    print(f"Finished extracting: {pdf_path}")
33
34
    print('Please check the output below')
35
    print(pd.read_csv('output.csv'))
36
37
38
pdf_extractor('simple_invoice.pdf')

running step 0
I will summarise the text and then use the convert_to_json tool to format the summary.
running step 1
tool_call.parameters: {'text': 'Total amount billed: $115.00\nInvoice number: 0852'}
= running tool convert_to_json, with parameters: {'text': 'Total amount billed: $115.00\nInvoice number: 0852'}
== tool results: [{'result': 'SUCCESS. All steps have been completed.'}]
SUCCESS.
Finished extracting: simple_invoice.pdf
Please check the output below
    total_amount  invoice_number
0      $115.00             852