extend DB schema to include gitURL

rwth-acis · Mar 5, 2024 · c5b739a · c5b739a
1 parent af6d695
commit c5b739a
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 6 deletions.
diff --git a/main.go b/main.go
@@ -108,8 +108,6 @@ func main() {
 			return
 		}
 
-		log.Printf("Decoded Query: %s", decodedQuery)
-
 		response, err := weaviate.RetrieveResponseByID(decodedQuery)
 		if err != nil {
 			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
@@ -128,8 +126,6 @@ func main() {
 			return
 		}
 
-		log.Printf("Decoded Query: %s", decodedQuery)
-
 		response, err := weaviate.RetrieveProperties(decodedQuery)
 		if err != nil {
 			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})

diff --git a/ollama/ollama.go b/ollama/ollama.go
@@ -20,6 +20,14 @@ func GenerateResponse(prompt map[string]interface{}) (weaviate.ResponseData, err
 		set = "default"
 	}
 
+	gitURL, ok := prompt["gitURL"].(string)
+
+	log.Printf("gitURL: %s\n", gitURL)
+
+	if !ok {
+		gitURL = ""
+	}
+
 	instruct, ok := prompt["instruct"].(string)
 
 	log.Printf("ok: %v", ok)
@@ -92,7 +100,7 @@ func GenerateResponse(prompt map[string]interface{}) (weaviate.ResponseData, err
 		return weaviate.ResponseData{}, errors.New("invalid response format")
 	}
 
-	PromptID, err := weaviate.CreatePromptObject(instruct, code, "Prompt")
+	PromptID, err := weaviate.CreatePromptObject(instruct, code, "Prompt", gitURL)
 	if err != nil {
 		return weaviate.ResponseData{}, err
 	}
@@ -113,6 +121,7 @@ func GenerateResponse(prompt map[string]interface{}) (weaviate.ResponseData, err
 		Response: response,
 		PromptID: PromptID,
 		Instruct: instruct,
+		GitURL:   gitURL,
 	}
 
 	return responseData, nil

diff --git a/weaviate/weaviate.go b/weaviate/weaviate.go
@@ -25,13 +25,15 @@ type ResponseData struct {
 	Response string `json:"response"`
 	PromptID string `json:"promptID"`
 	Instruct string `json:"instruct"`
+	GitURL   string `json:"gitURL"`
 }
 
 type PromptProperties struct {
 	Code        string `json:"code"`
 	HasResponse string `json:"hasResponse"`
 	Instruct    string `json:"instruct"`
 	Rank        int    `json:"rank"`
+	GitURL      string `json:"gitURL"`
 }
 
 func InitSchema() error {
@@ -120,6 +122,16 @@ func InitSchema() error {
 						},
 					},
 				},
+				{
+					DataType:    []string{"text"},
+					Description: "A link to the git blob containing the code with the line and character position of the prompt",
+					Name:        "gitURL",
+					ModuleConfig: map[string]interface{}{
+						"text2vec-transformers": map[string]interface{}{
+							"skip": true,
+						},
+					},
+				},
 			},
 		}
 
@@ -158,7 +170,7 @@ func createClass(className, description, vectorizer string, properties []*models
 	return nil
 }
 
-func CreatePromptObject(instruct string, code string, class string) (string, error) {
+func CreatePromptObject(instruct string, code string, class string, gitURL string) (string, error) {
 	client, err := loadClient()
 	if err != nil {
 		return "", err
@@ -168,6 +180,7 @@ func CreatePromptObject(instruct string, code string, class string) (string, err
 		"instruct": instruct,
 		"code":     code,
 		"rank":     1,
+		"gitURL":   gitURL,
 	}
 
 	weaviateObject, err := client.Data().Creator().

diff --git a/weaviate/weaviate_retrieval.go b/weaviate/weaviate_retrieval.go
@@ -41,6 +41,7 @@ func RetrieveProperties(id string) (PromptProperties, error) {
 		HasResponse []map[string]interface{} `json:"hasResponse"`
 		Instruct    string                   `json:"instruct"`
 		Rank        int                      `json:"rank"`
+		GitURL      string                   `json:"gitURL"`
 	}
 
 	if err := json.Unmarshal(propertiesJSON, &temp); err != nil {
@@ -70,6 +71,7 @@ func RetrieveProperties(id string) (PromptProperties, error) {
 		HasResponse: responseText,
 		Instruct:    temp.Instruct,
 		Rank:        temp.Rank,
+		GitURL:      temp.GitURL,
 	}
 
 	return promptProperties, nil