forked from philippgille/chromem-go
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request philippgille#33 from philippgille/improve-example
Improve example
- Loading branch information
Showing
2 changed files
with
108 additions
and
65 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,79 @@ | ||
package main | ||
|
||
import ( | ||
"context" | ||
"net/http" | ||
"strings" | ||
"text/template" | ||
|
||
"github.com/sashabaranov/go-openai" | ||
) | ||
|
||
const ( | ||
// We use a local LLM running in Ollama for asking the question: https://ollama.com/ | ||
ollamaBaseURL = "http://localhost:11434/v1" | ||
// We use a very small model that doesn't need much resources and is fast, but | ||
// doesn't have much knowledge: https://ollama.com/library/gemma | ||
// We found Gemma 2B to be superior to TinyLlama (1.1B), Stable LM 2 (1.6B) | ||
// and Phi-2 (2.7B) for the retrieval augmented generation (RAG) use case. | ||
llmModel = "gemma:2b" | ||
) | ||
|
||
// There are many different ways to provide the context to the LLM. | ||
// You can pass each context as user message, or the list as one user message, | ||
// or pass it in the system prompt. The system prompt itself also has a big impact | ||
// on how well the LLM handles the context, especially for LLMs with < 7B parameters. | ||
// The prompt engineering is up to you, it's out of scope for the vector database. | ||
var systemPromptTpl = template.Must(template.New("system_prompt").Parse(` | ||
You are a helpful assistant with access to a knowlege base, tasked with answering questions about the world and its history, people, places and other things. | ||
Answer the question in a very concise manner. Use an unbiased and journalistic tone. Do not repeat text. Don't make anything up. If you are not sure about something, just say that you don't know. | ||
{{- /* Stop here if no context is provided. The rest below is for handling contexts. */ -}} | ||
{{- if . -}} | ||
Answer the question solely based on the provided search results from the knowledge base. If the search results from the knowledge base are not relevant to the question at hand, just say that you don't know. Don't make anything up. | ||
Anything between the following 'context' HTML blocks is retrieved from the knowledge base, not part of the conversation with the user. The bullet points are ordered by relevance, so the first one is the most relevant. | ||
<context> | ||
{{- if . -}} | ||
{{- range $context := .}} | ||
- {{.}}{{end}} | ||
{{- end}} | ||
</context> | ||
{{- end -}} | ||
Don't mention the knowledge base, context or search results in your answer. | ||
`)) | ||
|
||
func askLLM(ctx context.Context, contexts []string, question string) string { | ||
// We can use the OpenAI client because Ollama is compatible with OpenAI's API. | ||
openAIClient := openai.NewClientWithConfig(openai.ClientConfig{ | ||
BaseURL: ollamaBaseURL, | ||
HTTPClient: http.DefaultClient, | ||
}) | ||
sb := &strings.Builder{} | ||
err := systemPromptTpl.Execute(sb, contexts) | ||
if err != nil { | ||
panic(err) | ||
} | ||
messages := []openai.ChatCompletionMessage{ | ||
{ | ||
Role: openai.ChatMessageRoleSystem, | ||
Content: sb.String(), | ||
}, { | ||
Role: openai.ChatMessageRoleUser, | ||
Content: "Question: " + question, | ||
}, | ||
} | ||
res, err := openAIClient.CreateChatCompletion(ctx, openai.ChatCompletionRequest{ | ||
Model: llmModel, | ||
Messages: messages, | ||
}) | ||
if err != nil { | ||
panic(err) | ||
} | ||
reply := res.Choices[0].Message.Content | ||
reply = strings.TrimSpace(reply) | ||
|
||
return reply | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters