Testing a Gemini codegen run
1# /// script
2# dependencies = [
3# "google-genai"
4# ]
5# ///
6import os
7import json
8from pydantic import BaseModel
9from google import genai
10
11class OCamlFile(BaseModel):
12 filename: str
13 contents: str
14
15# Create output directory if it doesn't exist
16output_dir = "output"
17if not os.path.exists(output_dir):
18 os.makedirs(output_dir)
19
20# Read the RFC text files
21with open("rfc8620.txt", "r") as f:
22 rfc8620_content = f.read()
23
24with open("rfc8621.txt", "r") as f:
25 rfc8621_content = f.read()
26
27# Read API key from file
28with open("api-key", "r") as f:
29 api_key = f.read().strip()
30
31client = genai.Client(api_key=api_key)
32
33prompt = """I wish to generate a set of OCaml module signatures and types (no implementations) that will type check, for an implementation of the JMAP protocol (RFC8620) and the associated email extensions (RFC8621). The code you generate should have ocamldoc that references the relevant sections of the RFC it is implementing, using <https://www.rfc-editor.org/rfc/rfc8620.html#section-1.2> as a template for the hyperlinks (replace the fragment with the appropriate section identifier).
34
35The architecture of the modules should be one portable set that implement core JMAP (RFC8620) as an OCaml module called `Jmap` (with module aliases to the submodules that implement that). Then generate another set of modules that implement the email-specific extensions (RFC8621) including flag handling for (e.g.) Apple Mail under a module called `Jmap_email`. These should all be portable OCaml type signatures (the mli files), and then generate another module that implements the interface for a Unix implementation that uses the Unix module to perform real connections. You do not need to implement TLS support for this first iteration of the code interfaces.
36
37You should also generate a module index file called jmap.mli that explains how all the generated modules fit together, along with a sketch of some example OCaml code that uses it to connect to a JMAP server and list recent unread emails from a particular sender.
38
39When selecting dependencies, ONLY use Yojson, Uri and Unix in your type signatures aside from the OCaml standard library. The standard Hashtbl is fine for any k/v datastructures and do not use Maps or other functor applications for this. DO NOT generate any AST attributes, and do not use any PPX derivers or other syntax extensions. Just generate clean, conventional OCaml type signatures.
40
41Also generate the dune scaffolding in your output suitable to build the project. This is:
421) a dune-project file that contains "(lang dune 3.17)"
432) a dune file with:
44
45(library
46 (name jmap)
47 (public_name jmap)
48 (libraries yojson uri)
49 (modules_without_implementation jmap <the other modules you generated>))
50
51
52Output each file you generate in JSON format.
53
54File = {'filename': str, 'contents': str}
55Return: list[File]
56"""
57
58response = client.models.generate_content(
59 model="gemini-2.5-pro-exp-03-25",
60 contents=[rfc8620_content, rfc8621_content, prompt],
61 config={
62 'response_mime_type': 'application/json',
63 'response_schema': list[OCamlFile],
64 },
65)
66
67try:
68 # Use the Pydantic-parsed response
69 files: list[OCamlFile] = response.parsed
70
71 # Iterate through each file and write it to the output directory
72 for file_entry in files:
73 filename = file_entry.filename
74 contents = file_entry.contents
75
76 if filename and contents:
77 output_path = os.path.join(output_dir, filename)
78 with open(output_path, 'w') as f:
79 f.write(contents)
80 print(f"Saved {filename} to {output_path}")
81
82 # Also save the full JSON response for reference
83 with open(os.path.join(output_dir, "full_response.json"), 'w') as f:
84 json.dump([{"filename": file.filename, "contents": file.contents} for file in files], f, indent=2)
85 print(f"Saved full response to {os.path.join(output_dir, 'full_response.json')}")
86
87except json.JSONDecodeError:
88 print("Failed to parse response as JSON. Raw response:")
89 print(response.text)
90
91 # Save the raw response for debugging
92 with open(os.path.join(output_dir, "raw_response.txt"), 'w') as f:
93 f.write(response.text)
94 print(f"Saved raw response to {os.path.join(output_dir, 'raw_response.txt')}")