alex/LocalAGI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

rag: add KB to conversation

Browse Source
This commit is contained in:
Ettore Di Giacinto
2024-04-09 22:34:22 +02:00
parent 36abf837a9
commit 78ba7871e9
4 changed files with 87 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
llm/rag.go
View File

@@ -7,10 +7,7 @@ import (
"github.com/sashabaranov/go-openai"
)
func StoreStringEmbeddingInVectorDB(apiHost string, openaiClient *openai.Client, s string) error {
// Example usage
client := NewStoreClient(apiHost)
func StoreStringEmbeddingInVectorDB(client *StoreClient, openaiClient *openai.Client, s string) error {
resp, err := openaiClient.CreateEmbeddings(context.TODO(),
openai.EmbeddingRequestStrings{
Input: []string{s},
@@ -39,8 +36,7 @@ func StoreStringEmbeddingInVectorDB(apiHost string, openaiClient *openai.Client,
return nil
}
func FindSimilarStrings(apiHost string, openaiClient *openai.Client, s string, similarEntries int) ([]string, error) {
client := NewStoreClient(apiHost)
func FindSimilarStrings(client *StoreClient, openaiClient *openai.Client, s string, similarEntries int) ([]string, error) {
resp, err := openaiClient.CreateEmbeddings(context.TODO(),
openai.EmbeddingRequestStrings{

21
llm/store.go
View File

@@ -10,8 +10,9 @@ import (
// Define a struct to hold your store API client
type StoreClient struct {
BaseURL string
Client *http.Client
BaseURL string
APIToken string
Client *http.Client
}
// Define request and response struct formats based on the API documentation
@@ -45,10 +46,11 @@ type FindResponse struct {
}
// Constructor for StoreClient
func NewStoreClient(baseUrl string) *StoreClient {
func NewStoreClient(baseUrl, apiToken string) *StoreClient {
return &StoreClient{
BaseURL: baseUrl,
Client: &http.Client{},
BaseURL: baseUrl,
APIToken: apiToken,
Client: &http.Client{},
}
}
@@ -105,6 +107,10 @@ func (c *StoreClient) doRequest(path string, data interface{}) error {
if err != nil {
return err
}
// Set Bearer token
if c.APIToken != "" {
req.Header.Set("Authorization", "Bearer "+c.APIToken)
}
req.Header.Set("Content-Type", "application/json")
resp, err := c.Client.Do(req)
@@ -132,7 +138,10 @@ func (c *StoreClient) doRequestWithResponse(path string, data interface{}) ([]by
return nil, err
}
req.Header.Set("Content-Type", "application/json")
// Set Bearer token
if c.APIToken != "" {
req.Header.Set("Authorization", "Bearer "+c.APIToken)
}
resp, err := c.Client.Do(req)
if err != nil {
return nil, err