From e78f864afdc3a5c4e6869d2bc8109fb059cc7247 Mon Sep 17 00:00:00 2001
From: Phodal Huang <h@phodal.com>
Date: Fri, 25 Aug 2023 09:26:19 +0800
Subject: [PATCH] feat(counit): init basic apis

---
 .../unitmesh/devti/counit/PatternMatching.kt  |  5 +++++
 .../cc/unitmesh/devti/counit/PreProcessor.kt  | 11 +++++++++++
 .../unitmesh/devti/counit/client/CoUnitApi.kt | 19 +++++++++++++++++++
 .../devti/counit/model/ExplainQuery.kt        | 18 ++++++++++++++++++
 .../devti/counit/model/PromptResult.kt        |  5 +++++
 .../cc/unitmesh/devti/counit/model/Tool.kt    |  8 ++++++++
 6 files changed, 66 insertions(+)
 create mode 100644 src/main/kotlin/cc/unitmesh/devti/counit/PatternMatching.kt
 create mode 100644 src/main/kotlin/cc/unitmesh/devti/counit/PreProcessor.kt
 create mode 100644 src/main/kotlin/cc/unitmesh/devti/counit/client/CoUnitApi.kt
 create mode 100644 src/main/kotlin/cc/unitmesh/devti/counit/model/ExplainQuery.kt
 create mode 100644 src/main/kotlin/cc/unitmesh/devti/counit/model/PromptResult.kt
 create mode 100644 src/main/kotlin/cc/unitmesh/devti/counit/model/Tool.kt

diff --git a/src/main/kotlin/cc/unitmesh/devti/counit/PatternMatching.kt b/src/main/kotlin/cc/unitmesh/devti/counit/PatternMatching.kt
new file mode 100644
index 0000000000..16f887c38e
--- /dev/null
+++ b/src/main/kotlin/cc/unitmesh/devti/counit/PatternMatching.kt
@@ -0,0 +1,5 @@
+package cc.unitmesh.devti.counit
+
+class PatternMatching {
+
+}
\ No newline at end of file
diff --git a/src/main/kotlin/cc/unitmesh/devti/counit/PreProcessor.kt b/src/main/kotlin/cc/unitmesh/devti/counit/PreProcessor.kt
new file mode 100644
index 0000000000..04e9191512
--- /dev/null
+++ b/src/main/kotlin/cc/unitmesh/devti/counit/PreProcessor.kt
@@ -0,0 +1,11 @@
+package cc.unitmesh.devti.counit
+
+import cc.unitmesh.devti.settings.configurable.coUnitSettings
+import com.intellij.openapi.project.Project
+
+class PreProcessor(val project: Project) {
+    fun isCoUnit(input: String): Boolean {
+        return project.coUnitSettings.enableCoUnit && input.startsWith("/counit")
+    }
+}
+
diff --git a/src/main/kotlin/cc/unitmesh/devti/counit/client/CoUnitApi.kt b/src/main/kotlin/cc/unitmesh/devti/counit/client/CoUnitApi.kt
new file mode 100644
index 0000000000..a7d3700f60
--- /dev/null
+++ b/src/main/kotlin/cc/unitmesh/devti/counit/client/CoUnitApi.kt
@@ -0,0 +1,19 @@
+package cc.unitmesh.devti.counit.client
+
+import cc.unitmesh.devti.counit.model.PromptResult
+import cc.unitmesh.devti.counit.model.Tool
+import retrofit2.Call
+import retrofit2.http.GET
+import retrofit2.http.POST
+import retrofit2.http.Query
+
+interface CoUnitApi {
+    @GET("/prompt/explain")
+    fun explainQuery(@Query("q") q: String): Call<PromptResult>
+
+    @POST("/prompt/functions/matching")
+    fun toolPrompter(@Query("q") q: String): Call<PromptResult>
+
+    @GET("/prompt/functions/list")
+    fun functions(): Call<List<Tool>>
+}
diff --git a/src/main/kotlin/cc/unitmesh/devti/counit/model/ExplainQuery.kt b/src/main/kotlin/cc/unitmesh/devti/counit/model/ExplainQuery.kt
new file mode 100644
index 0000000000..156abe9032
--- /dev/null
+++ b/src/main/kotlin/cc/unitmesh/devti/counit/model/ExplainQuery.kt
@@ -0,0 +1,18 @@
+package cc.unitmesh.devti.counit.model
+
+import kotlinx.serialization.SerialName
+import kotlinx.serialization.Serializable
+
+@Serializable
+data class ExplainQuery(
+    val domain: String,
+    val query: String,
+    @SerialName("hypothetical_document")
+    val hypotheticalDocument: String
+)
+
+@Serializable
+data class QAExample(
+    val question: String,
+    val answer: ExplainQuery
+)
\ No newline at end of file
diff --git a/src/main/kotlin/cc/unitmesh/devti/counit/model/PromptResult.kt b/src/main/kotlin/cc/unitmesh/devti/counit/model/PromptResult.kt
new file mode 100644
index 0000000000..a68d1cf197
--- /dev/null
+++ b/src/main/kotlin/cc/unitmesh/devti/counit/model/PromptResult.kt
@@ -0,0 +1,5 @@
+package cc.unitmesh.devti.counit.model
+
+data class PromptResult(
+    val prompt: String
+)
\ No newline at end of file
diff --git a/src/main/kotlin/cc/unitmesh/devti/counit/model/Tool.kt b/src/main/kotlin/cc/unitmesh/devti/counit/model/Tool.kt
new file mode 100644
index 0000000000..5a1063ad4d
--- /dev/null
+++ b/src/main/kotlin/cc/unitmesh/devti/counit/model/Tool.kt
@@ -0,0 +1,8 @@
+package cc.unitmesh.devti.counit.model
+
+data class Tool(
+    val name: String,
+    val description: String,
+    val schema: String,
+    val examples: String
+)
\ No newline at end of file