1
import os
2
import requests
3
from vellum import Vellum
4
from vellum.workflows.errors import WorkflowErrorCode
5
from vellum.workflows.exceptions import NodeException
6
from vellum.workflows.nodes import BaseNode, RetryNode
7
8
api_key = os.getenv("VELLUM_API_KEY")
9
if api_key is None:
10
    raise ValueError("VELLUM_API_KEY environment variable is not set")
11
client = Vellum(api_key=api_key)
12
13
@RetryNode.wrap(max_attempts=20, retry_on_error_code=WorkflowErrorCode.USER_DEFINED_ERROR)
14
class RetrieveData(BaseNode):
15
    document_id = Inputs.document_id
16
17
    class Outputs(BaseNode.Outputs):
18
        document_content: str
19
20
    def run(self) -> BaseNode.Outputs:
21
        # Retrieve the document status from Vellum
22
        try:
23
            response = client.documents.retrieve(id=self.document_id)
24
            indexing_state = response.document_to_document_indexes[0].indexing_state
25
            processing_state = response.processing_state
26
            text_file_url = response.document_to_document_indexes[0].extracted_text_file_url
27
28
            if indexing_state == "INDEXED":
29
                return self.Outputs(document_content=requests.get(text_file_url).text)
30
            elif indexing_state == "FAILED" or processing_state == "FAILED":
31
                raise NodeException("Indexing or processing failed")
32
33
            raise NodeException("Document processing not complete yet.", retry_on_error_code=WorkflowErrorCode.USER_DEFINED_ERROR)
34
        except Exception as e:
35
            raise NodeException(f"An error occurred: {str(e)}")

1
from pydantic import BaseModel, Field
2
from vellum import ChatMessagePromptBlock, JinjaPromptBlock, PromptParameters
3
from vellum.workflows.nodes import InlinePromptNode
4
5
# Define Pydantic models to match the JSON structure
6
class CategoryMetrics(BaseModel):
7
    participants: int = Field(..., description="Number of participants in the category.")
8
    ballots_completed: int = Field(..., alias="Ballots Completed", description="Number of ballots completed.")
9
    ballots_incomplete_terminated: int = Field(
10
        ..., alias="Ballots Incomplete/Terminated", description="Number of ballots incomplete or terminated."
11
    )
12
    accuracy: str = Field(..., description="Accuracy percentage with sample size.")
13
    time_to_complete: str = Field(..., alias="Time to Complete", description="Time taken to complete ballots.")
14
15
class ExtractedSchema(BaseModel):
16
    example_table: str = Field(..., alias="Example Table", description="Description of the table.")
17
    disability_category: DisabilityCategory = Field(..., alias="Disability Category")
18
19
class DataExtractionNode(InlinePromptNode):
20
    ml_model = "gpt-4o-mini"
21
    blocks = [
22
        ChatMessagePromptBlock(
23
            chat_role="SYSTEM",
24
            blocks=[
25
                JinjaPromptBlock(
26
                    template="""Analyze the following document content and extract all identifiable key-value pairs. Present the extracted information in JSON format, ensuring compliance with the following schema:
27
<document_content>
28
{{ document_content }}
29
</document_content>
30
Extracted Data (in JSON format):
31
""",
32
                ),
33
            ],
34
        ),
35
    ]
36
    prompt_inputs = {"document_content": RetrieveData.Outputs.document_content}
37
    parameters = PromptParameters(
38
        temperature=0,
39
        max_tokens=1000,
40
        top_p=1,
41
        custom_parameters={
42
            "json_schema": {"name": "data_extraction_schema", "schema": ExtractedSchema.model_json_schema()}
43
        },
44
    )