benjaminshafii · August 10, 2023 00:40 · Aug 10, 2023
diff --git a/basel.js b/basel.js
@@ -0,0 +1,236 @@
+const Anthropic = require('@anthropic-ai/sdk');
+const path = require('path');
+const YAML = require('yaml');
+const fs = require('fs');
+
+// Initialize Anthropic SDK
+const anthropic = new Anthropic({
+  apiKey: process.env.ANTHROPIC_API_KEY,
+});
+
+// Main function
+async function main() {
+  // Get command line arguments
+  const args = process.argv.slice(2);
+  const roots = [];
+  const includes = [];
+  const excludes = [];
+  let i = 0;
+  while (i < args.length) {
+    if (args[i] === '--root') {
+      roots.push(args[i + 1]);
+      i += 2;
+    } else if (args[i] === '--include') {
+      includes.push(args[i + 1]);
+      i += 2;
+    } else if (args[i] === '--exclude') {
+      excludes.push(args[i + 1]);
+      i += 2;
+    } else {
+      i++;
+    }
+  }
+
+  // Get files
+  const files = getFiles({
+    roots: roots.length > 0 ? roots : ['app'],
+    includes: includes.length > 0 ? includes : ['route'],
+    excludes: excludes.length > 0 ? excludes : ['dist'],
+  });
+  console.log('List of files to generate OpenAPI spec for:')
+  console.log(files)
+
+  // Initialize OpenAPI spec
+  // let spec = initializeSpec();
+
+  // Generate spec for each endpoint
+  let yamlDraft = ''
+  for (let endpoint of files) {
+    const endpointSchema = await generateEndpointSchema(endpoint);
+    console.log(endpointSchema)
+    yamlDraft += endpointSchema
+
+  }
+  fs.writeFileSync('draft_spec.yaml', yamlDraft);
+
+  // Validate and write spec to file
+  const validSpec = await validateSpec(yamlDraft);
+  writeSpecToFile(validSpec);
+}
+
+
+
+// Generate endpoint schema
+async function generateEndpointSchema(endpoint) {
+  const completion = await anthropic.completions.create({
+    model: 'claude-2',
+    prompt: generatePrompt(endpoint),
+    max_tokens_to_sample: 600,
+  });
+  // get the yaml from the completion
+
+  const regex = /<yaml>(.*?)<\/yaml>/s;
+  const match = completion.completion.match(regex);
+  return match ? match[1] : null;
+
+}
+
+// Generate prompt for endpoint schema
+function generatePrompt(endpoint) {
+  return `
+Human: Write hello between <yaml></yaml>
+Assistant: <yaml>hello</yaml>
+Human: Please provide an OpenAPI 3.0 schema for this API endpoint between yaml:
+<yaml>
+${endpoint}:
+${fs.readFileSync(endpoint, 'utf8')}
+</yaml>
+Assistant:`;
+}
+// Validate spec
+async function validateSpec(spec) {
+  const specYaml = YAML.stringify(spec, {
+    aliasDuplicateObjects: true, indent: 2,
+    lineWidth: -1
+  });
+
+  const validYaml = await anthropic.completions.create({
+    model: 'claude-2',
+    max_tokens_to_sample: 100000,
+    prompt: generateValidationPrompt(specYaml)
+  });
+
+  const regex = /<yaml>(.*?)<\/yaml>/s;
+  const match = validYaml.completion.match(regex);
+  return match ? match[1] : null;
+}
+
+// Generate validation prompt
+function generateValidationPrompt(specYaml) {
+  return `
+Human: Output valid OpenAPI yaml within <yaml></yaml> always:
+Assistant: 
+<yaml>
+openapi: 3.0.0
+info:
+  title: Example API
+  description: An example to demonstrate OpenAPI 3.0
+  version: 1.0.0
+servers:
+  - url: https://api.example.com
+paths:
+  /users:
+    get:
+      summary: Gets a list of users
+      responses:
+        200:
+          description: Success
+          content:
+            application/json:
+              schema:
+                type: array
+                items:
+                  $ref: '#/components/schemas/User'
+  /users/{userId}:
+    get:
+      summary: Gets a user by ID
+      parameters:
+        - name: userId
+          in: path
+          required: true
+          schema:
+            type: integer
+      responses:
+        200:
+          description: Success
+          content:
+            application/json:
+              schema: 
+                $ref: '#/components/schemas/User'
+        404:
+          description: User not found
+    put:
+      summary: Updates a user
+      parameters:
+        - name: userId
+          in: path
+          required: true
+          schema:
+            type: integer
+      requestBody:
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/UserUpdate'
+      responses:
+        200:
+          description: Success
+        404:
+          description: User not found
+ 
+components:
+  schemas:
+    User:
+      type: object
+      properties:
+        id:
+          type: integer
+        name: 
+          type: string
+    UserUpdate:
+      type: object
+      properties:
+        name:
+          type: string
+</yaml>
+Human:  ${specYaml}
+Output valid OpenAPI yaml within <yaml></yaml> fix errors and output a single yaml file:
+Assistant:
+`;
+}
+
+// Write spec to file
+function writeSpecToFile(validSpec) {
+  if (validSpec) {
+    fs.writeFileSync('spec.yaml', validSpec);
+  }
+}
+
+// Get files
+function getFiles(options) {
+  const {
+    roots = ['./'],
+    includes = [],
+    excludes = ['/node_modules/'],
+    extensions = ['js', 'ts', 'jsx', 'tsx']
+  } = options;
+
+  let files = [];
+  for (const root of roots) {
+    files = files.concat(walk(root));
+  }
+
+  return files.filter(file => {
+    if (includes.length > 0) {
+      return includes.some(include => file.includes(include));
+    }
+    return !excludes.some(exclude => file.includes(exclude))
+  }).filter(file => extensions.includes(path.extname(file).slice(1)));
+}
+
+// Walk directory
+function walk(dir) {
+  let files = [];
+  for (const file of fs.readdirSync(dir)) {
+    const filepath = path.join(dir, file);
+    if (fs.statSync(filepath).isDirectory()) {
+      files = files.concat(walk(filepath));
+    } else {
+      files.push(filepath);
+    }
+  }
+  return files;
+}
+
+// Run main function
+main();
No results found