refactor: update for new flow

unit-mesh · Sep 20, 2023 · 31d6327 · 31d6327
1 parent 6755dfe
commit 31d6327
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 5 deletions.
diff --git a/rag-script/.gitignore b/rag-script/.gitignore
@@ -1,3 +1,4 @@
 temp
 dependencies/
-0_codes.json
+0_codes.json
+.env
diff --git a/rag-script/build.gradle.kts b/rag-script/build.gradle.kts
@@ -24,6 +24,8 @@ dependencies {
     implementation(libs.huggingface.tokenizers)
     implementation(libs.dataframe)
 
+    implementation("io.github.cdimascio:dotenv-kotlin:6.4.1")
+
     implementation("io.ktor:ktor-client-logging:2.3.4")
     implementation("io.ktor:ktor-client-core:2.3.4")
     implementation("io.ktor:ktor-client-cio:2.3.4")

diff --git a/rag-script/src/main/kotlin/cc/unitmesh/apply/Connection.kt b/rag-script/src/main/kotlin/cc/unitmesh/apply/Connection.kt
@@ -37,8 +37,8 @@ class Connection(val connectionType: ConnectionType, val apiKey: String, val api
         return completion(function())
     }
 
-    fun streamCompletion(msg: String): Flowable<String> {
-        return provider.streamCompletion(listOf(LlmMsg.ChatMessage(LlmMsg.ChatRole.User, msg)))
+    fun streamCompletion(function: () -> String): Flowable<String> {
+        return provider.streamCompletion(listOf(LlmMsg.ChatMessage(LlmMsg.ChatRole.User, function())))
     }
 }
 

diff --git a/rag-script/src/main/kotlin/cc/unitmesh/apply/Workflow.kt b/rag-script/src/main/kotlin/cc/unitmesh/apply/Workflow.kt
@@ -6,12 +6,22 @@ import cc.unitmesh.nlp.embedding.EmbeddingProvider
 import cc.unitmesh.rag.document.Document
 import cc.unitmesh.rag.store.EmbeddingMatch
 import cc.unitmesh.store.ElasticsearchStore
+import io.github.cdimascio.dotenv.Dotenv
+
+
+
 
 /**
  * Apply is a DSL for invoking a function in a template.
  */
 @ApplyDsl
 class Workflow(val name: String) {
+    var env = try {
+        Dotenv.load()
+    } catch (e: Exception) {
+        null
+    }
+
     var connection: Connection = Connection(ConnectionType.OpenAI, "", "")
     var vectorStore: VectorStore = VectorStore(StoreType.Elasticsearch)
     var embedding: EmbeddingEngine = EmbeddingEngine(EngineType.SentenceTransformers)
@@ -117,3 +127,8 @@ fun scripting(name: String, init: Workflow.() -> Unit): Workflow {
     workflow.init()
     return workflow
 }
+fun scripting(init: Workflow.() -> Unit): Workflow {
+    val workflow = Workflow("scripting")
+    workflow.init()
+    return workflow
+}
diff --git a/rag-script/src/test/kotlin/cc/unitmesh/apply/WorkflowTest.kt b/rag-script/src/test/kotlin/cc/unitmesh/apply/WorkflowTest.kt
@@ -13,8 +13,11 @@ class WorkflowTest {
     @Test
     @Ignore
     fun index_and_query() {
-        scripting("code") {
-            connection = Connection(ConnectionType.OpenAI)
+        scripting {
+            val apiKey = env?.get("OPENAI_API_KEY") ?: ""
+            val apiHost = env?.get("OPENAI_API_HOST") ?: ""
+
+            connection = Connection(ConnectionType.OpenAI, apiKey, apiHost)
             embedding = EmbeddingEngine(EngineType.SentenceTransformers)
             vectorStore = VectorStore(StoreType.Elasticsearch)