from parallel import Parallel

client = Parallel()

# Tip: If using Pydantic models, you can generate the schema automatically:
# class CompanyEnrichment(BaseModel):
#     ceo_name: str = Field(description="Name of the CEO")
#     founding_year: str = Field(description="Year the company was founded")
# output_schema = {"type": "json", "json_schema": CompanyEnrichment.model_json_schema()}

schema = client.beta.findall.enrich(
    findall_id="findall_40e0ab8c10754be0b7a16477abb38a2f",
    processor="core",
    output_schema={
        "type": "json",
        "json_schema": {
            "type": "object",
            "properties": {
                "ceo_name": {
                    "type": "string",
                    "description": "Name of the CEO"
                },
                "founding_year": {
                    "type": "string",
                    "description": "Year the company was founded"
                }
            },
            "required": ["ceo_name", "founding_year"],
            "additionalProperties": False
        }
    }
)

print(f"Enrichment added, schema: {schema.model_dump_json(indent=2)}")

{
  "objective": "Find all AI companies that raised Series A funding in 2024",
  "entity_type": "companies",
  "match_conditions": [
    {
      "name": "developing_ai_products_check",
      "description": "Company must be developing artificial intelligence (AI) products"
    }
  ],
  "enrichments": [
    {
      "processor": "core",
      "output_schema": {
        "json_schema": {
          "type": "object",
          "properties": {
            "ceo_name": {
              "type": "string",
              "description": "Name of the current CEO of the company. If the CEO is not publicly known, provide the name of the highest-ranking executive or founder. If no information is available, return null."
            }
          }
        },
        "type": "json"
      }
    }
  ],
  "generator": "core",
  "match_limit": 50
}