Updated cargo.toml

.gitignore

@@ -96,3 +96,5 @@ valgrind-python.supp
 *.pyd
 lcov.info
 netlify_build/
+
+build.sh

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "nova-python"
-version = "0.1.2"
+version = "0.1.5"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -9,6 +9,7 @@ name = "nova_python"
 crate-type = ["cdylib"]
 
 [dependencies]
-pyo3 = "0.19.0"
+pyo3 = {version = "0.19.2", features = ["extension-module", "generate-import-lib"]}
 reqwest = "0.11.18"
 tokio = { version = "1.29.1", features = ["rt-multi-thread", "time"] }
+serde_json = "1.0.104" 

README.md

@@ -36,7 +36,7 @@ Now, to make a request, use the `make_request` function. For example:
 from nova_python import Endpoints, Models, NovaClient
 client = NovaClient("YOUR_API_KEY")
 
-client.make_request(
+reponse = client.make_request(
     endpoint=Endpoints.CHAT_COMPLETION,
     model=Models.GPT3,
     data=[
@@ -44,6 +44,8 @@ client.make_request(
         {"role": "user", "content": "Hello!"}
     ]
 )
+
+...
 ```
 
 or
@@ -52,15 +54,36 @@ or
 from nova_python import Endpoints, Models, NovaClient
 client = NovaClient("YOUR_API_KEY")
 
-client.make_request(
+reponse = client.make_request(
     endpoint=Endpoints.MODERATION,
     model=Models.MODERATION_STABLE,
     data=[{"input": "I'm going to kill them."}]
 )
+
+...
+```
+If everything goes to plan, you'll receive a string containing JSON-Data, which you can then use in your project.  
+
+Note, that when using chat completion, as special ChatResponse-Instance get's returned.  
+You can access the reponse's json.data, by casting it to a string using the `str` method, like this:
+```python
+...
+
+str(reponse)
+
+...
+```
+
+but more importantly, you can use it's `get_message_content` function, to directly get access to the chat reponse. Used like this:
+```
+...
+
+content = reponse.get_message_content()
+
+...
 ```
 
 
-If everything goes to plan, you'll receive a string containing JSON-Data, which you can then use in your project.  
 *Happy prompting!*
 <br><br>
 ## FAQ ##

src/lib.rs

@@ -54,7 +54,7 @@ impl NovaClient {
     }
 
     // Used to make a request to the Nova API
-    fn make_request(&self, endpoint: Endpoints, model: Models, data: Vec<Py<PyDict>>, seconds_until_timeout: Option<String>) -> PyResult<String> {
+    fn make_request(&self, endpoint: Endpoints, model: Models, data: Vec<Py<PyDict>>, seconds_until_timeout: Option<usize>) -> PyResult<PyObject> {
         if !model_is_compatible(&endpoint, &model) {
             return Err(NovaClient::get_endpoint_not_compatible_error());
         }
@@ -64,13 +64,13 @@ impl NovaClient {
 
         let rt = tokio::runtime::Runtime::new().unwrap();
         let seconds_until_timeout = match seconds_until_timeout {
-            Some(seconds_until_timeout) => seconds_until_timeout.parse::<u64>().unwrap(),
+            Some(seconds_until_timeout) => seconds_until_timeout,
             None => 30
         };
 
-        let response: Result<String, reqwest::Error> = rt.block_on(async {
+        let unmatched_response: Result<String, reqwest::Error> = rt.block_on(async {
             let client = reqwest::Client::builder()
-                .timeout(time::Duration::from_secs(seconds_until_timeout))
+                .timeout(time::Duration::from_secs(seconds_until_timeout as u64))
                 .user_agent("Mozilla/5.0")
                 .build()
                 .unwrap();
@@ -86,9 +86,25 @@ impl NovaClient {
             Ok(text)
         });
 
-        match response {
-            Ok(response) => Ok(response),
-            Err(response) => Err(pyo3::exceptions::PyRuntimeError::new_err(response.to_string()))
+        let reponse = match unmatched_response {
+            Ok(unmatched_response) => Ok(unmatched_response),
+            Err(unmatched_response) => Err(pyo3::exceptions::PyRuntimeError::new_err(unmatched_response.to_string()))
+        }.unwrap();
+
+        if endpoint == Endpoints::ChatCompletion {
+            let final_reponse = Python::with_gil(|py| {
+                let reponse = ChatResponse::new(reponse);
+                reponse.into_py(py)
+            });
+
+            return Ok(final_reponse);
+            
+        } else {
+            let final_reponse = Python::with_gil(|py| {
+                reponse.into_py(py)
+            });
+
+            return Ok(final_reponse);
         }
     }
 
@@ -181,9 +197,36 @@ impl NovaClient {
     fn get_invalid_model_error() -> PyErr {
         pyo3::exceptions::PyValueError::new_err("Invalid model")
     }
+}
 
-    fn get_request_failed_error() -> PyErr {
-        pyo3::exceptions::PyRuntimeError::new_err("Request failed for unknown reasons.")
+#[pyclass(module = "nova_python", frozen)]
+struct ChatResponse {
+    json: String,
+}
+
+#[pymethods]
+impl ChatResponse {
+    #[new]
+    fn new(json: String) -> Self {
+        ChatResponse {
+            json
+        }
+    }
+
+    fn get_message_content(&self) -> PyResult<String> {
+        let json = &self.json;
+        let json: serde_json::Value = serde_json::from_str(json).unwrap();
+
+        let content = json["choices"][0]["message"]["content"].as_str().unwrap();
+        Ok(content.trim().to_string())
+    }
+
+    fn __str__(&self) -> PyResult<String> {
+        Ok(self.json.clone())
+    }
+
+    fn __repr__(&self) -> PyResult<String> {
+        Ok(format!("ChatResponse(json={})", self.json))
     }
 }