create node.js pkg

Author: bit-web24

lib/index.js

@@ -0,0 +1,19 @@
+const genai = require('../index.node');
+
+class GeminiFlash {
+    constructor(api_key) {
+        this.api_key = api_key;
+    }
+
+    async generate_content(prompt) {
+        try {
+            const response = await genai.generate_content(this.api_key, prompt);
+            return response;
+        } catch (error) {
+            console.error("Error generating content:", error);
+            throw error;
+        }
+    }
+}
+
+module.exports = GeminiFlash;

package.json

@@ -2,7 +2,7 @@
   "name": "gemini_flash",
   "version": "0.1.0",
   "description": "Client for Google's Gemini-1.5-Flash Generative AI Model",
-  "main": "index.node",
+  "main": "lib/index.js",
   "scripts": {
     "test": "cargo test",
     "cargo-build": "cargo build --message-format=json > cargo.log",

src/lib.rs

@@ -1,22 +1,19 @@
 use curl::easy::{Easy, List};
 use neon::prelude::*;
+use pulldown_cmark::{Event, Options, Parser, Tag, TagEnd};
 use serde_json::to_string;
-use pulldown_cmark::{Event, Parser, Tag, TagEnd, Options};
 
 mod request;
 mod response;
 use request::*;
-// use response::Response;
 
 #[neon::main]
 fn main(mut cx: ModuleContext) -> NeonResult<()> {
     cx.export_function("generate_content", generate_content)?;
-    cx.export_function("markdown_to_text", markdown_to_text)?;
     Ok(())
 }
 
-// @args api_key: String, prompt: &str
-pub fn generate_content(mut cx: FunctionContext) -> JsResult<JsString> {
+pub fn generate_content(mut cx: FunctionContext) -> JsResult<JsObject> {
     let api_key = cx.argument::<JsString>(0)?.value(&mut cx);
     let prompt = cx.argument::<JsString>(1)?.value(&mut cx);
 
@@ -56,13 +53,13 @@ pub fn generate_content(mut cx: FunctionContext) -> JsResult<JsString> {
     }
 
     let json_data = String::from_utf8(response_data).unwrap();
-    // let response: Response = serde_json::from_str(json_data.as_str()).unwrap();
-    // response
-    Ok(cx.string(json_data))
+    let response: response::Response = serde_json::from_str(json_data.as_str()).unwrap();
+    let js_res_obj = to_jsobj(cx, response);
+
+    Ok(js_res_obj.unwrap())
 }
 
-pub fn markdown_to_text(mut cx: FunctionContext) -> JsResult<JsString> {
-    let markdown = cx.argument::<JsString>(0)?.value(&mut cx);
+pub fn markdown_to_text(markdown: &str) -> String {
     let options = Options::empty();
     let parser = Parser::new_ext(&markdown, options);
     let mut plain_text = String::new();
@@ -71,18 +68,40 @@ pub fn markdown_to_text(mut cx: FunctionContext) -> JsResult<JsString> {
         match event {
             Event::Text(text) | Event::Code(text) => plain_text.push_str(&text),
             Event::Start(tag) => match tag {
-                Tag::Paragraph | Tag::Heading{..} => plain_text.push_str("\n\n"),
+                Tag::Paragraph | Tag::Heading { .. } => plain_text.push_str("\n\n"),
                 Tag::List(_) => plain_text.push('\n'),
                 Tag::Item => plain_text.push_str("\n• "),
-                _ => {},
+                _ => {}
             },
             Event::End(tag) => match tag {
-                TagEnd::Paragraph | TagEnd::Heading{..} => plain_text.push('\n'),
-                _ => {},
+                TagEnd::Paragraph | TagEnd::Heading { .. } => plain_text.push('\n'),
+                _ => {}
             },
             _ => {}
         }
     }
 
-    Ok(cx.string(plain_text))
+    plain_text
+}
+
+fn to_jsobj(mut cx: FunctionContext, res: response::Response) -> JsResult<JsObject> {
+    let obj = cx.empty_object();
+
+    let text = cx.string(markdown_to_text(&res.candidates[0].content.parts[0].text));
+    obj.set(&mut cx, "text", text)?;
+
+    let finish_reason = cx.string(&res.candidates[0].finishReason);
+    obj.set(&mut cx, "finish_reason", finish_reason)?;
+
+    let usage_metadata = cx.string(to_string(&res.usageMetadata).unwrap());
+    obj.set(&mut cx, "usage_metadata", usage_metadata)?;
+
+    let role = cx.string(&res.candidates[0].content.role);
+    obj.set(&mut cx, "role", role)?;
+
+    let prompt_feedback = to_string(&res.candidates[0].safetyRatings).unwrap();
+    let prompt_feedback_str = cx.string(prompt_feedback);
+    obj.set(&mut cx, "prompt_feedback", prompt_feedback_str)?;
+
+    Ok(obj)
 }

src/response.rs

@@ -14,7 +14,6 @@ pub struct Candidate {
     pub finishReason: String,
     pub index: u32,
     pub safetyRatings: Vec<SafetyRating>,
-    pub citationMetadata: CitationMetadata,
 }
 
 #[allow(non_snake_case)]
@@ -37,12 +36,6 @@ pub struct SafetyRating {
     pub probability: String,
 }
 
-#[allow(non_snake_case)]
-#[derive(Serialize, Deserialize, Debug)]
-pub struct CitationMetadata {
-    pub citationSources: Vec<CitationSource>,
-}
-
 #[allow(non_snake_case)]
 #[derive(Serialize, Deserialize, Debug)]
 pub struct CitationSource {
@@ -59,17 +52,3 @@ pub struct UsageMetadata {
     pub candidatesTokenCount: u32,
     pub totalTokenCount: u32,
 }
-
-// impl Response {
-//     pub fn text(&self) -> String {
-//         utils::markdown_to_text(self.candidates[0].content.parts[0].text.as_str())
-//     }
-
-//     pub fn prompt_feedback(&self) -> &Vec<SafetyRating> {
-//         self.candidates[0].safetyRatings.as_ref()
-//     }
-
-//     pub fn candidates(&self) -> &Vec<Candidate> {
-//         self.candidates.as_ref()
-//     }
-// }