feat: implement new functions and update API

Author: bregydoc

examples/pi.rs

@@ -1,7 +1,6 @@
 use auto_rust::auto_implement;
 
 #[auto_implement]
-/// Algoritmo basdo en trignometria para calcular el valor de pi
 fn calculate_pi_with_n_iterations(n: u64) -> f64 {
     todo!()
 }

examples/quine.rs

@@ -0,0 +1,14 @@
+use auto_rust::auto_implement;
+
+#[auto_implement]
+/// Only return the string literal of the quine
+/// Remember to use `to_string()` to convert it to a string
+fn calculate_rust_quine() -> String {
+    todo!()
+}
+
+fn main() {
+    let quine = calculate_rust_quine().to_string();
+
+    print!("{}", quine)
+}

examples/with_context.rs

@@ -0,0 +1,15 @@
+use auto_rust::auto_implement;
+
+#[auto_implement(
+    context = "This algorithm works by calculating the area of a circle using the trigonometric formula for the area of a triangle. Don't
+    use powi, use powf instead. Don't use the sqrt function, use the hypot function instead. Don't use the sin function, use the sin_cos function instead. Don't use the cos function, use the sin_cos function instead"
+)]
+fn calculate_pi_with_n_iterations(n: u64) -> f64 {
+    todo!()
+}
+
+fn main() {
+    let result = calculate_pi_with_n_iterations(100_000);
+    println!("pi: {}", result);
+    // assert_eq!(result, 3.0418396189294032);
+}

src/api.rs

@@ -8,7 +8,8 @@ pub fn open_ai_chat_completions(
     system_message: String,
     user_message: String,
 ) -> Result<ChatCompletionResponse, Box<dyn std::error::Error>> {
-    let open_ai_key = env::var("OPENAI_API_KEY").unwrap_or("".to_string());
+    let open_ai_key = env::var("OPENAI_API_KEY").unwrap();
+    let model_name = env::var("OPENAI_MODEL_NAME").unwrap_or("gpt-3.5-turbo".to_string());
 
     let mut headers = header::HeaderMap::new();
 
@@ -23,23 +24,20 @@ pub fn open_ai_chat_completions(
         .build()
         .unwrap();
 
+    let body = json!({
+        "model": model_name,
+        "messages": [
+            {"role": "system", "content": system_message},
+            {"role": "user", "content": user_message}
+        ]
+    });
+
     let res = client
         .post("https://api.openai.com/v1/chat/completions")
         .headers(headers)
-        .body(
-            json!({
-                "model": "gpt-3.5-turbo",
-                "messages": [
-                    {"role": "system", "content": system_message},
-                    {"role": "user", "content": user_message}
-                ]
-            })
-            .to_string(),
-        )
+        .body(body.to_string())
         .send()?
         .json::<ChatCompletionResponse>()?;
 
-    // println!("{:?}", res);
-
     Ok(res)
 }

src/generator.rs

@@ -2,7 +2,10 @@ use std::error::Error;
 
 use crate::api::open_ai_chat_completions;
 
-pub fn generate_body_function_from_head(head: String) -> Result<String, Box<dyn Error>> {
+pub fn generate_body_function_from_head(
+    head: String,
+    extra_context: Option<String>,
+) -> Result<String, Box<dyn Error>> {
     let system_message = "You are an AI code assistant trained on the GPT-4 architecture. Your task is to generate Rust function body implementations based only on the provided function signatures. When the user provides a function signature using the command '/complete', your response must be the plain text function body, without any explanations, formatting, or code blocks. Do not include the function signature, function name, or any other information in your response. Triple backticks (```) and function signatures are strictly prohibited in your response. Responding with any prohibited content will result in a penalty. 
     example 1:
     INPUT: /implement fn my_ip() -> String
@@ -16,13 +19,20 @@ pub fn generate_body_function_from_head(head: String) -> Result<String, Box<dyn
         ip_addr.to_string() 
     example 2:
     INPUT: /implement fn hello_world() -> String
-    OUTPUT: \"Hello World\".to_string()
+    OUTPUT: 
+        \"Hello World\".to_string()
     example 3:
     INPUT: /implement fn hello_world(name: String) -> String
-    OUTPUT: format!(\"Hello {}!\", name)
+    OUTPUT: 
+        format!(\"Hello {}!\", name)
 ".to_string();
 
-    let user_message = format!("/implement {}", head);
+    let user_message = extra_context
+        .map(|c| format!("Extra context: {}\n", c))
+        .unwrap_or("".to_string())
+        + &format!("/implement {}", head);
+
+    // println!("User message: {}", user_message);
 
     let res = open_ai_chat_completions(system_message, user_message).unwrap();
 

src/lib.rs

@@ -14,18 +14,20 @@ pub fn implement(_item: TokenStream) -> TokenStream {
     // TODO: Evaluate the use of dotenv in this crate
     dotenv().ok();
 
-    let implemented_fn = generate_body_function_from_head(_item.to_string()).unwrap();
+    let implemented_fn = generate_body_function_from_head(_item.to_string(), None).unwrap();
 
-    println!("{}", implemented_fn);
+    // println!("{}", implemented_fn);
 
     implemented_fn.parse().unwrap()
 }
 
 #[proc_macro_attribute]
-pub fn auto_implement(_args: TokenStream, input: TokenStream) -> TokenStream {
+pub fn auto_implement(args: TokenStream, input: TokenStream) -> TokenStream {
     let ast: ItemFn = syn::parse(input).expect("Failed to parse input as a function");
 
-    // Search for the information within the attributes.
+    let context = args.to_string();
+
+    // println!("Context: {}", context);
 
     let mut prompt_input = String::new();
 
@@ -46,11 +48,12 @@ pub fn auto_implement(_args: TokenStream, input: TokenStream) -> TokenStream {
 
     dotenv().ok();
 
-    let implemented_fn = generate_body_function_from_head(prompt_input).unwrap();
+    let implemented_fn = generate_body_function_from_head(prompt_input, Some(context)).unwrap();
 
     // println!("\n{}\n", implemented_fn);
 
     // #[allow(long_running_const_eval)]
+
     // loop {
     // let mut line = String::new();
     // let b1 = std::io::stdin().read_line(&mut line).unwrap();