From 278c7e6b315f6cf39818b525148495dcb0c2e388 Mon Sep 17 00:00:00 2001 From: Ido Perlmuter Date: Sun, 11 Dec 2022 19:18:22 +0200 Subject: [PATCH] Add support for OpenAI API aiac will now use OpenAI's official (beta) API by default. This means we can use the GPT-3 language mode, which returns code but not a Markdown of explanations like ChatGPT does. The API uses a simple API key that can be provided via the `--api-key` flag, or the `OPENAI_API_KEY` environment variable. The previous ChatGPT usage is still supported, but via the `--chat-gpt` flag. When used, `--session-token` should also be provided (or the `CHATGPT_SESSION_TOKEN` environment variable). The application now moves its logic behind a command named "get" (alias "generate"), so it's not a part of the prompt. Instead, the library will always prepend the word "generate" to the prompt (if we simply used "get", the model will not necessarily generate code). --- README.md | 85 ++++++++----- go.mod | 4 +- go.sum | 23 ++-- libaiac/libaiac.go | 309 +++++++++++++++++++++++++++++++-------------- main.go | 74 ++++++++--- 5 files changed, 330 insertions(+), 165 deletions(-) diff --git a/README.md b/README.md index c1cc6e7..3f07987 100644 --- a/README.md +++ b/README.md @@ -5,8 +5,9 @@ * [Description](#description) -* [Caveats](#caveats) * [Quick Start](#quick-start) + * [Usage via OpenAI API](#usage-via-openai-api) + * [Usage via ChatGPT](#usage-via-chatgpt) * [Example Prompts](#example-prompts) * [Example Output](#example-output) * [Acknowledgements](#acknowledgements) @@ -17,35 +18,53 @@ ## Description `aiac` is a command line tool to generate IaC (Infrastructure as Code) templates -via [OpenAI](https://openai.com/)'s [ChatGPT](https://chat.openai.com/). The CLI allows you to ask ChatGPT to generate templates -for different scenarios (e.g. "generate terraform for AWS EC2"). It then makes -the request on your behalf, accepts the response, and extracts the example code -from it. This extracted code is either printed to standard output, or saved to a -file. Optionally, the CLI will append the complete message (i.e. the full -Markdown response including ChatGPT's explanations) to a separate markdown file. +via [OpenAI](https://openai.com/)'s API or via ChatGPT. The CLI allows you to ask the model to generate templates +for different scenarios (e.g. "generate terraform for AWS EC2"). It will make the +request, and store the resulting code to a file, or simply print it to standard +output. -## Caveats - -- ChatGPT's API is not public, and is likely to change frequently, which - may break this program. Please inform us via the [issues page](https://github.com/gofireflyio/aiac/issues) if this happens. -- ChatGPT may rate limit your requests, and is prone to answer slowly or not at - all when under heavy load. -- ChatGPT's responses to the same prompt may differ between executions. If you - are unhappy with the results, try again or modify your prompt. -- You will currently have to manually copy a session token from an actual browser - session. Directions in the [Quick Start](#quick-start) section. +When using ChatGPT, the server returns a Markdown file with code and explanations. +The CLI will extract the code in this case, and optionally store the entire +Markdown of explanations to a separate file. ## Quick Start First, install `aiac`: - go get github.com/gofireflyio/aiac + go generate github.com/gofireflyio/aiac Alternatively, clone the repository and build from source: git clone https://github.com/gofireflyio/aiac.git go build +### Usage via OpenAI API + +You will need to provide `aiac` with an API key. Create your API key [here](https://beta.openai.com/account/api-keys). +You can either provide the API key via the `--api-key` command line flag, or via +the `OPENAI_API_KEY` environment variable. + +By default, `aiac` simply prints the extracted code to standard output + + aiac --api-key=API_KEY generate terraform for AWS EC2 + +To store the resulting code to a file: + + aiac --api-key=API_KEY \ + --output-file="aws_ec2.tf" \ + get terraform for AWS EC2 + +### Usage via ChatGPT + +There are several caveats to using `aiac` in ChatGPT mode: + +- ChatGPT's API is not public, and is likely to change frequently, which + may break this program. Please inform us via the [issues page](https://github.com/gofireflyio/aiac/issues) if this happens. +- ChatGPT may rate limit your requests, and is prone to answer slowly or not at + all when under heavy load. +- You will currently have to manually copy a session token from an actual browser + session in order to authenticate (instructions follow). + You will need to provide `aiac` with a session token. Since ChatGPT doesn't currently support programmatic usage, you will need to do this via your browser (this is hopefully temporary until we can implement OpenAI authentication). @@ -55,38 +74,34 @@ To get a token, follow these steps: 2. Open the Web Developer Tools (usually Ctrl+Shift+I). 3. Go to the "Storage" tab, and move to the list of "Cookies". 4. Find the cookie called "__Secure-next-auth.session-token". -5. Copy its value. This is the session token. +5. Copy its value. This is the session token. You can store it in the + `CHATGPT_SESSION_TOKEN` environment variable, or provide it via the + `--session-token` command line flag. ![](/authentication.jpg) -By default, `aiac` simply prints the extracted code to standard output - - aiac --session-token=TOKEN generate terraform for AWS EC2 +Then run: -To store the resulting code to a file, and to append the explanations to a -markdown file, run: - - aiac --session-token=TOKEN \ - --output-file="aws_ec2.tf" \ + aiac --chat-gpt \ + --session-token=TOKEN \ + --output-file="ec2.tf" \ --readme-file="README.md" \ - generate terraform for AWS EC2 + get terraform for AWS EC2 ## Example Prompts The following prompts are known to work: -- generate Terraform for AWS EC2 -- generate Dockerfile for NodeJS -- generate GitHub action for deploying Terraform -- generate Python code for Pulumi that deploys Azure VPC +- get Terraform for AWS EC2 +- get Dockerfile for NodeJS with comments +- get GitHub action for deploying Terraform +- get Python code for Pulumi that deploys Azure VPC ## Example Output Command line prompt: -```sh -aiac --session-token=TOKEN generate Dockerfile for NodeJS -``` + aiac get dockerfile for nodejs with comments Output: diff --git a/go.mod b/go.mod index fa3fcb3..4247cc8 100644 --- a/go.mod +++ b/go.mod @@ -3,8 +3,9 @@ module github.com/gofireflyio/aiac go 1.19 require ( + github.com/adrg/xdg v0.4.0 github.com/alecthomas/kong v0.7.1 - github.com/google/uuid v1.3.0 + github.com/google/uuid v1.1.2 github.com/ido50/requests v1.5.0 ) @@ -12,4 +13,5 @@ require ( go.uber.org/atomic v1.10.0 // indirect go.uber.org/multierr v1.8.0 // indirect go.uber.org/zap v1.24.0 // indirect + golang.org/x/sys v0.0.0-20211025201205-69cdffdb9359 // indirect ) diff --git a/go.sum b/go.sum index 657272f..179831d 100644 --- a/go.sum +++ b/go.sum @@ -36,13 +36,15 @@ cloud.google.com/go/storage v1.8.0/go.mod h1:Wv1Oy7z6Yz3DshWRJFhqM/UCfaWIRTdp0RX cloud.google.com/go/storage v1.10.0/go.mod h1:FLPqc6j+Ki4BU591ie1oL6qBQGu2Bl/tZ9ullr3+Kg0= cloud.google.com/go/storage v1.14.0/go.mod h1:GrKmX003DSIwi9o29oFT7YDnHYwZoctc3fOKtUw0Xmo= dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU= -github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ= github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo= +github.com/adrg/xdg v0.4.0 h1:RzRqFcjH4nE5C6oTAxhBtoE2IRyjBSa62SCbyPidvls= +github.com/adrg/xdg v0.4.0/go.mod h1:N6ag73EX4wyxeaoeHctc1mas01KZgsj5tYiAIwqJE/E= github.com/alecthomas/assert/v2 v2.1.0 h1:tbredtNcQnoSd3QBhQWI7QZ3XHOVkw1Moklp2ojoH/0= github.com/alecthomas/kong v0.7.1 h1:azoTh0IOfwlAX3qN9sHWTxACE2oV8Bg2gAwBsMwDQY4= github.com/alecthomas/kong v0.7.1/go.mod h1:n1iCIO2xS46oE8ZfYCNDqdR0b0wZNrXAIAqro/2132U= github.com/alecthomas/repr v0.1.0 h1:ENn2e1+J3k09gyj2shc0dHr/yjaWSHRlrJ4DPMevDqE= +github.com/benbjohnson/clock v1.1.0 h1:Q92kusRqC1XV2MjkWETPvjJVqKetz1OzxZB7mHJLju8= github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI= github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI= @@ -113,9 +115,8 @@ github.com/google/pprof v0.0.0-20201023163331-3e6fc7fc9c4c/go.mod h1:kpwsk12EmLe github.com/google/pprof v0.0.0-20201203190320-1bf35d6f28c2/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE= github.com/google/pprof v0.0.0-20201218002935-b9804c9f04c2/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE= github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI= +github.com/google/uuid v1.1.2 h1:EVhdT+1Kseyi1/pUmXKaFxYsDNy9RQYkMWRH68J/W7Y= github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= -github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I= -github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg= github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk= github.com/googleapis/google-cloud-go-testing v0.0.0-20200911160855-bcd43fbb19e8/go.mod h1:dvDLG8qkwmyD9a/MJJN3XJcT3xFxOKAvTZGvuZmac9g= @@ -151,7 +152,6 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4= github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA= -github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk= github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= @@ -164,18 +164,15 @@ go.opencensus.io v0.22.2/go.mod h1:yxeiOL68Rb0Xd1ddK5vPZ/oVn4vY4Ynel7k9FzqtOIw= go.opencensus.io v0.22.3/go.mod h1:yxeiOL68Rb0Xd1ddK5vPZ/oVn4vY4Ynel7k9FzqtOIw= go.opencensus.io v0.22.4/go.mod h1:yxeiOL68Rb0Xd1ddK5vPZ/oVn4vY4Ynel7k9FzqtOIw= go.opencensus.io v0.22.5/go.mod h1:5pWMHQbX5EPX2/62yrJeAkowc+lfs/XD7Uxpq3pI6kk= -go.uber.org/atomic v1.6.0 h1:Ezj3JGmsOnG1MoRWQkPBsKLe9DwWD9QeXzTRzzldNVk= go.uber.org/atomic v1.6.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ= go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc= go.uber.org/atomic v1.10.0 h1:9qC72Qh0+3MqyJbAn8YU5xVq1frD8bn3JtD2oXtafVQ= go.uber.org/atomic v1.10.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= -go.uber.org/multierr v1.5.0 h1:KCa4XfM8CWFCpxXRGok+Q0SS/0XBhMDbHHGABQLvD2A= +go.uber.org/goleak v1.1.11 h1:wy28qYRKZgnJTxGxvye5/wgWr1EKjmUDGYox5mGlRlI= go.uber.org/multierr v1.5.0/go.mod h1:FeouvMocqHpRaaGuG9EjoKcStLC43Zu/fmqdUMPcKYU= go.uber.org/multierr v1.8.0 h1:dg6GjLku4EH+249NNmoIciG9N/jURbDG+pFlTkhzIC8= go.uber.org/multierr v1.8.0/go.mod h1:7EAYxJLBy9rStEaz58O2t4Uvip6FSURkq8/ppBp95ak= -go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee h1:0mgffUl7nfd+FpvXMVz4IDEaUSmT1ysygQC7qYo7sG4= go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee/go.mod h1:vJERXedbb3MVM5f9Ejo0C68/HhF8uaILCdgjnY+goOA= -go.uber.org/zap v1.14.1 h1:nYDKopTbvAPq/NrUVZwT15y2lpROBiLLyoRTbXOYWOo= go.uber.org/zap v1.14.1/go.mod h1:Mb2vm2krFEG5DV0W9qcHBYFtp/Wku1cvYaqPsS/WYfc= go.uber.org/zap v1.24.0 h1:FiJd5l1UOLj0wCgbSE0rwwXHzEdAZS6hiiSnxJN/D60= go.uber.org/zap v1.24.0/go.mod h1:2kMP+WWQ8aoFoedH3T2sq6iJ2yDWpHbP0f6MQbS9Gkg= @@ -208,7 +205,6 @@ golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHl golang.org/x/lint v0.0.0-20191125180803-fdd1cda4f05f/go.mod h1:5qLYkcX4OjUUV8bRuDixDT3tpyyb+LUpUlRWLxfhWrs= golang.org/x/lint v0.0.0-20200130185559-910be7a94367/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY= golang.org/x/lint v0.0.0-20200302205851-738671d3881b/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY= -golang.org/x/lint v0.0.0-20201208152925-83fdc39ff7b5 h1:2M3HP5CCK1Si9FQhwnzYhXdG6DXeebvUHFpre8QvbyI= golang.org/x/lint v0.0.0-20201208152925-83fdc39ff7b5/go.mod h1:3xt1FjdF8hUf6vQPIChWIBhFzV8gjjsPE/fR3IyQdNY= golang.org/x/mobile v0.0.0-20190312151609-d3739f865fa6/go.mod h1:z+o9i4GpDbdi3rU15maQ/Ox0txvL9dWGYEHz965HBQE= golang.org/x/mobile v0.0.0-20190719004257-d2bd2a29d028/go.mod h1:E/iHnbuqvinMTCcRqshq8CkpyQDoeVncDDYHnLhea+o= @@ -219,7 +215,6 @@ golang.org/x/mod v0.1.1-0.20191107180719-034126e5016b/go.mod h1:QqPTAvyqsEbceGzB golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.4.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= -golang.org/x/mod v0.4.1 h1:Kvvh58BN8Y9/lBi7hTekvtMpm07eUZ0ck5pRHpsMWrY= golang.org/x/mod v0.4.1/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= @@ -304,8 +299,9 @@ golang.org/x/sys v0.0.0-20210104204734-6f8348627aad/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.0.0-20210119212857-b64e53b001e4/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210225134936-a50acf3fe073/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210423185535-09eb48e85fd7/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= -golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1 h1:SrN+KX8Art/Sf4HNj6Zcz06G7VEz+7w9tdXTPOZ7+l4= golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20211025201205-69cdffdb9359 h1:2B5p2L5IfGiD7+b9BOoRMC6DgObAVZV+Fsp050NqXik= +golang.org/x/sys v0.0.0-20211025201205-69cdffdb9359/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= @@ -365,12 +361,10 @@ golang.org/x/tools v0.0.0-20201201161351-ac6f37ff4c2a/go.mod h1:emZCQorbCU4vsT4f golang.org/x/tools v0.0.0-20201208233053-a543418bbed2/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.0.0-20210105154028-b0ab187a4818/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.0.0-20210108195828-e2f9c7f1fc8e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= -golang.org/x/tools v0.1.0 h1:po9/4sTYwZU9lPhi1tOrb4hCv3qrhiQ77LZfGa2OjwY= golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= -golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE= golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= google.golang.org/api v0.4.0/go.mod h1:8k5glujaEP+g9n7WNsDg8QP6cUVNI86fCNMcbazEtwE= google.golang.org/api v0.7.0/go.mod h1:WtwebWUNSVBH/HAw79HIFXZNqEvBhG+Ra+ax0hx3E3M= @@ -464,16 +458,15 @@ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= -gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= honnef.co/go/tools v0.0.0-20190106161140-3f1c8253044a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= honnef.co/go/tools v0.0.0-20190418001031-e561f6794a2a/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= honnef.co/go/tools v0.0.1-2019.2.3/go.mod h1:a3bituU0lyd329TUQxRnasdCoJDkEUEAqEt0JzvZhAg= honnef.co/go/tools v0.0.1-2020.1.3/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9vFzvIQ3k= -honnef.co/go/tools v0.0.1-2020.1.4 h1:UoveltGrhghAA7ePc+e+QYDHXrBps2PqFZiHkGR/xK8= honnef.co/go/tools v0.0.1-2020.1.4/go.mod h1:X/FiERA/W4tHapMX5mGpAtMSVEeEUOyHaw9vFzvIQ3k= rsc.io/binaryregexp v0.2.0/go.mod h1:qTv7/COck+e2FymRvadv62gMdZztPaShugOCi3I+8D8= rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0= diff --git a/libaiac/libaiac.go b/libaiac/libaiac.go index deebef7..a757894 100644 --- a/libaiac/libaiac.go +++ b/libaiac/libaiac.go @@ -1,17 +1,20 @@ package libaiac import ( - "bufio" "bytes" + "bufio" "context" "encoding/json" "errors" "fmt" + "github.com/adrg/xdg" "github.com/google/uuid" "io" "net/http" "os" + "path/filepath" "strings" + "time" "github.com/ido50/requests" ) @@ -27,24 +30,48 @@ var ErrNoCode = errors.New("no code generated") type Client struct { *requests.HTTPClient - sessionToken string - accessToken string - conversationID string - parentID string + token string + chatGPT bool } -func NewClient(sessionToken string) *Client { - sessionToken = strings.TrimPrefix(sessionToken, SessionTokenCookie+"=") +func NewClient(chatGPT bool, token string) *Client { + cli := &Client{ + token: token, + chatGPT: chatGPT, + } - return &Client{ - sessionToken: sessionToken, - HTTPClient: requests.NewClient("https://"+ChatGPTHost). + if !chatGPT { + cli.HTTPClient = requests.NewClient("https://api.openai.com/v1"). Accept("application/json"). - Header("User-Agent", DefaultUserAgent). - Header("X-Openai-Assistant-App-Id", ""). - Header("Accept-Language", "en-US,en;q=0.9"). - Header("Referer", "https://"+ChatGPTHost+"/chat"), + Header("Authorization", fmt.Sprintf("Bearer %s", token)). + ErrorHandler(func( + httpStatus int, + contentType string, + body io.Reader, + ) error { + var res struct { + Error struct { + Message string `json:"message"` + Type string `json:"type"` + } `json:"error"` + } + err := json.NewDecoder(body).Decode(&res) + if err != nil { + return fmt.Errorf( + "OpenAI returned response %s", + http.StatusText(httpStatus), + ) + } + + return fmt.Errorf("[%s] %s", res.Error.Type, res.Error.Message) + }) + } else { + cli.HTTPClient = requests.NewClient(fmt.Sprintf("https://%s", ChatGPTHost)). + Header("User-Agent", "Mozilla/5.0 (Windows NT 10.0; rv:107.0) Gecko/20100101 Firefox/107.0"). + Header("Accept-Language", "en-US,en;q=0.9") } + + return cli } func (client *Client) Ask( @@ -53,35 +80,72 @@ func (client *Client) Ask( outputPath string, readmePath string, ) (err error) { - requestID, err := uuid.NewRandom() + var code, readme string + + if client.chatGPT { + code, readme, err = client.askViaChatGPT(ctx, prompt) + } else { + code, err = client.askViaAPI(ctx, prompt) + } + if err != nil { - return fmt.Errorf("failed generating UUID: %w", err) + return err } - messages := make(chan []byte) + var codeFd io.Writer - if client.accessToken == "" { - // get an access token - var session struct { - AccessToken string `json:"accessToken"` + if outputPath == "-" { + codeFd = os.Stdout + } else { + f, err := os.Create(outputPath) + if err != nil { + return fmt.Errorf( + "failed creating output file %s: %w", + outputPath, err, + ) } - err = client.NewRequest("GET", "/api/auth/session"). - Cookie(&http.Cookie{ - Name: SessionTokenCookie, - Value: client.sessionToken, - Path: "/", - Domain: ChatGPTHost, - }). - Into(&session). - RunContext(ctx) + defer f.Close() + + codeFd = f + } + + fmt.Fprint(codeFd, code) + + if readmePath != "" { + f, err := os.Create(readmePath) if err != nil { - return fmt.Errorf("failed getting session details: %w", err) + return fmt.Errorf( + "failed creating readme file %s: %w", + readmePath, err, + ) } - client.accessToken = session.AccessToken + defer f.Close() + + fmt.Fprintf(f, readme) + } + + return nil +} + +func (client *Client) askViaChatGPT(ctx context.Context, prompt string) ( + code string, + readme string, + err error, +) { + requestID, err := uuid.NewRandom() + if err != nil { + return code, readme, fmt.Errorf("failed generating UUID: %w", err) } + accessToken, err := client.loadAccessToken(ctx) + if err != nil { + return code, readme, fmt.Errorf("failed loading access token: %w", err) + } + + cacheAccessToken(accessToken) + // start a conversation body := map[string]interface{}{ "action": "next", @@ -99,24 +163,22 @@ func (client *Client) Ask( "conversation_id": nil, } - if client.conversationID != "" { - body["conversation_id"] = client.conversationID - } - if client.parentID == "" { - parentID, err := uuid.NewRandom() - if err != nil { - return fmt.Errorf("failed generating initial parent ID: %w", err) - } + parentID, err := uuid.NewRandom() + if err != nil { + return code, readme, fmt.Errorf( + "failed generating initial parent ID: %w", + err, + ) + } - client.parentID = parentID.String() - } + body["parent_message_id"] = parentID.String() - body["parent_message_id"] = client.parentID + messages := make(chan []byte) err = client.NewRequest("POST", "/backend-api/conversation"). Accept("text/event-stream"). JSONBody(body). - Header("Authorization", fmt.Sprintf("Bearer %s", client.accessToken)). + Header("Authorization", fmt.Sprintf("Bearer %s", accessToken)). Cookie(&http.Cookie{ Name: CallbackURLCookie, Value: "https://" + ChatGPTHost + "/", @@ -125,13 +187,13 @@ func (client *Client) Ask( }). Cookie(&http.Cookie{ Name: SessionTokenCookie, - Value: client.sessionToken, + Value: client.token, Path: "/", Domain: ChatGPTHost, }). Subscribe(ctx, messages) if err != nil { - return fmt.Errorf("failed starting a conversation: %w", err) + return code, readme, fmt.Errorf("failed starting a conversation: %w", err) } type ChatGPTMessage struct { @@ -158,12 +220,12 @@ func (client *Client) Ask( var msg ChatGPTMessage err = json.Unmarshal(bmsg, &msg) if err != nil { - return fmt.Errorf("failed parsing ChatGPT response: %w", err) + return code, readme, fmt.Errorf( + "failed parsing ChatGPT response: %w", + err, + ) } - client.conversationID = msg.ConversationID - client.parentID = msg.Message.ID - if len(msg.Message.Content.Parts) < 1 { continue } @@ -171,55 +233,14 @@ func (client *Client) Ask( finalMessage = msg.Message.Content.Parts[0] } - var output io.Writer - - if outputPath == "-" { - output = os.Stdout - } else { - f, err := os.Create(outputPath) - if err != nil { - return fmt.Errorf( - "failed creating output file %s: %w", - outputPath, err, - ) - } - - defer f.Close() - - output = f - } - - var readme io.Writer - - if readmePath != "" { - f, err := os.OpenFile(readmePath, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0644) - if err != nil { - return fmt.Errorf( - "failed creating/opening readme file %s: %w", - readmePath, err, - ) - } - - defer f.Close() - - readme = f - } - - if readme != nil { - fmt.Fprintf(readme, "# %s\n", prompt) - } - scanner := bufio.NewScanner(strings.NewReader(finalMessage)) var writeOutput, alreadyHadCode bool + var b strings.Builder for scanner.Scan() { line := scanner.Text() - if readme != nil { - fmt.Fprintln(readme, line) - } - if line == "```" { if !alreadyHadCode { writeOutput = !writeOutput @@ -228,16 +249,110 @@ func (client *Client) Ask( } } } else if writeOutput { - fmt.Fprintln(output, line) + fmt.Fprintln(&b, line) } } - if readme != nil { - fmt.Fprintln(readme, "") + return b.String(), finalMessage, nil +} + +func (client *Client) askViaAPI(ctx context.Context, prompt string) ( + code string, + err error, +) { + var answer struct { + Choices []struct { + Text string `json:"text"` + Index int64 `json:"index"` + FinishReason string `json:"finish_reason"` + } `json:"choices"` + } + + var status int + err = client.NewRequest("POST", "/completions"). + JSONBody(map[string]interface{}{ + "model": "text-davinci-003", + "prompt": prompt, + "max_tokens": 4097 - len(prompt), + }). + Into(&answer). + StatusInto(&status). + RunContext(ctx) + if err != nil { + return code, fmt.Errorf("failed sending prompt: %w", err) + } + + if len(answer.Choices) == 0 { + return code, fmt.Errorf("no results returned from API") + } + + if answer.Choices[0].FinishReason != "stop" { + return code, fmt.Errorf( + "result was truncated by API due to %s", + answer.Choices[0].FinishReason, + ) + } + + return strings.TrimSpace(answer.Choices[0].Text), nil +} + +type CacheFile struct { + AccessToken string `json:"accessToken"` + Expiry int64 `json:"expiry"` +} + +func (client *Client) loadAccessToken(ctx context.Context) (token string, err error) { + // try to load access token from cache file + f, err := os.Open(filepath.Join(xdg.ConfigHome, "aiac.token")) + if err == nil { + defer f.Close() + + var tokenData CacheFile + err = json.NewDecoder(f).Decode(&tokenData) + if err == nil && time.Now().Unix() < tokenData.Expiry { + // cached token has not expired yet + return tokenData.AccessToken, nil + } + } + + // get an access token from ChatGPT + var session struct { + AccessToken string `json:"accessToken"` + } + + err = client.NewRequest("GET", "/api/auth/session"). + Cookie(&http.Cookie{ + Name: SessionTokenCookie, + Value: client.token, + Path: "/", + Domain: ChatGPTHost, + }). + Into(&session). + RunContext(ctx) + if err != nil { + return token, fmt.Errorf( + "failed getting session details: %w", + err, + ) + } + + return session.AccessToken, nil +} + +func cacheAccessToken(token string) error { + f, err := os.Create(filepath.Join(xdg.ConfigHome, "aiac.token")) + if err != nil { + return fmt.Errorf("failed creating token file: %w", err) } - if !alreadyHadCode { - return ErrNoCode + defer f.Close() + + err = json.NewEncoder(f).Encode(CacheFile{ + AccessToken: token, + Expiry: time.Now().Add(3600 * time.Second).Unix(), + }) + if err != nil { + return fmt.Errorf("failed encoding token data: %w", err) } return nil diff --git a/main.go b/main.go index 0f6c4f7..b487225 100644 --- a/main.go +++ b/main.go @@ -2,6 +2,7 @@ package main import ( "context" + "errors" "fmt" "os" "strings" @@ -10,34 +11,73 @@ import ( "github.com/gofireflyio/aiac/libaiac" ) -var cli struct { - SessionToken string `help:"ChatGPT session token" required:""` - OutputFile string `help:"Output file to push resulting code to, defaults to stdout" default:"-" type:"path"` - ReadmeFile string `help:"Markdown file to update with explanations" type:"path"` - Ask []string `arg:"" help:"What to ask ChatGPT to do"` +type flags struct { + APIKey string `help:"OpenAI API key (env: OPENAI_API_KEY)" optional:""` + SessionToken string `help:"Session token for ChatGPT (env: CHATGPT_SESSION_TOKEN)" optional:""` + ChatGPT bool `help:"Use ChatGPT instead of the OpenAI API (requires --session-token)" default:false` + OutputFile string `help:"Output file to push resulting code to, defaults to stdout" default:"-" type:"path"` + ReadmeFile string `help:"Markdown file to push explanations to" optional:"" type:"path"` + Get struct { + What []string `arg:"" help:"What to ask ChatGPT to generate"` + } `cmd:"" help:"Generate IaC code" aliases:"generate"` } func main() { - kong.Parse(&cli) + var cli flags + cmd := kong.Parse(&cli) - client := libaiac.NewClient(cli.SessionToken) + if cmd.Command() != "get " { + fmt.Fprintln(os.Stderr, "Unknown command") + os.Exit(1) + } + + var token string + + if !cli.ChatGPT { + token = cli.APIKey + if token == "" { + var ok bool + token, ok = os.LookupEnv("OPENAI_API_KEY") + + if !ok { + fmt.Fprintf(os.Stderr, "You must provide an OpenAI API key\n") + os.Exit(1) + } + } + } else { + token = cli.SessionToken + if token == "" { + var ok bool + token, ok = os.LookupEnv("CHATGPT_SESSION_TOKEN") + + if !ok { + fmt.Fprintf(os.Stderr, "You must provide a ChatGPT session token\n") + os.Exit(1) + } + } + } + + client := libaiac.NewClient(cli.ChatGPT, token) err := client.Ask( context.TODO(), - strings.Join(cli.Ask, " "), + // NOTE: we are prepending the word "generate" to the prompt, this + // ensures the language model actually generates code. The word "get", + // on the other hand, doesn't necessarily result in code being generated. + fmt.Sprintf("generate %s", strings.Join(cli.Get.What, " ")), cli.OutputFile, cli.ReadmeFile, ) if err != nil { - if err == libaiac.ErrNoCode { - fmt.Fprintln( - os.Stderr, - "It doesn't look like ChatGPT generated any code, please make "+ - "sure that you're prompt properly guides ChatGPT to do so.", - ) - } else { - fmt.Fprintf(os.Stderr, "Request failed: %s\n", err) - } + if errors.Is(err, libaiac.ErrNoCode) { + fmt.Fprintln( + os.Stderr, + "It doesn't look like ChatGPT generated any code, please make "+ + "sure that you're prompt properly guides ChatGPT to do so.", + ) + } else { + fmt.Fprintf(os.Stderr, "Request failed: %s\n", err) + } os.Exit(1) } }