Skip to content

A
agentchat
Commit 409192b5 authored by Wade's avatar Wade
Browse files
Options

add ollama models

parent e1fb0284
Expand all Hide whitespace changes
Inline Side-by-side
Showing with 960 additions and 591 deletions
agent_chat.log
View file @ 409192b5
This diff is collapsed.
docs/swagger.yaml
View file @ 409192b5
......@@ -149,6 +149,33 @@ paths:
data: '{"id": 1}'
code: 200
msg: "GraphRAG index data stored successfully"
/models:
post:
summary: List available models for a provider
description: Returns a list of available models for the specified provider (e.g., "deepseek" or "googleai").
tags:
- Chat
requestBody:
required: true
content:
application/json:
schema:
type: string
description: The provider to list models for
example: "deepseek"
responses:
'200':
description: Successful response with the list of models
content:
application/json:
schema:
$ref: '#/components/schemas/Response'
examples:
success:
value:
data: '{"deepseek-chat":{"label":"DeepSeek Chat","supports":{"multiturn":true,"systemRole":true,"media":false,"tools":true},"versions":[]},"deepseek-coder":{"label":"DeepSeek Coder","supports":{"multiturn":true,"systemRole":true,"media":false,"tools":true},"versions":[]},"deepseek-reasoner":{"label":"DeepSeek Reasoner","supports":{"multiturn":true,"systemRole":true,"media":false,"tools":true},"versions":[]}}'
code: 200
msg: "Models listed successfully"
/chat:
post:
summary: Send a chat message
......@@ -228,8 +255,29 @@ components:
msg:
type: string
description: A message describing the result
example: "Milvus index data stored successfully"
example: "Models listed successfully"
required:
- data
- code
- msg
# components:
# schemas:
# Response:
# type: object
# properties:
# data:
# type: string
# description: The response data, typically a JSON string or message
# example: '{"id": 1}'
# code:
# type: integer
# description: The response code (200 for success, 400 for invalid input, etc.)
# example: 200
# msg:
# type: string
# description: A message describing the result
# example: "Milvus index data stored successfully"
# required:
# - data
# - code
# - msg
flow.go 0 → 100644
View file @ 409192b5
This diff is collapsed.
main.go
View file @ 409192b5
......@@ -2,22 +2,15 @@ package main
import (
"context"
"encoding/json"
"flag"
"fmt"
"net/http"
"os"
"strings"
"time"
"github.com/firebase/genkit/go/ai"
"github.com/firebase/genkit/go/genkit"
"github.com/rs/zerolog"
"github.com/wade-liwei/agentchat/plugins/deepseek"
"github.com/wade-liwei/agentchat/plugins/graphrag"
"github.com/wade-liwei/agentchat/plugins/knowledge" // Import knowledge package
"github.com/wade-liwei/agentchat/plugins/graphrag" // Import knowledge package
"github.com/wade-liwei/agentchat/plugins/milvus"
"github.com/wade-liwei/agentchat/plugins/question"
"github.com/firebase/genkit/go/plugins/evaluators"
"github.com/firebase/genkit/go/plugins/googlegenai"
......@@ -27,7 +20,6 @@ import (
_ "github.com/wade-liwei/agentchat/docs" // 导入生成的 Swagger 文档
"github.com/rs/zerolog/log"
"github.com/wade-liwei/agentchat/util"
)
type ChatInput struct {
......@@ -145,8 +137,6 @@ func main() {
log.Fatal().Msg(err.Error())
}
embedder := googlegenai.GoogleAIEmbedder(g, "embedding-001")
if embedder == nil {
log.Fatal().Msg(err.Error())
......@@ -166,253 +156,16 @@ func main() {
log.Fatal().Msgf("DefineIndexerAndRetriever failed: %v", err)
}
// 定义文档索引流
genkit.DefineFlow(g, "index/document", func(ctx context.Context, input *DocumentInput) (Response, error) {
if input.Metadata == nil {
input.Metadata = make(map[string]any)
}
input.Metadata[util.UserIdKey] = input.UserID
input.Metadata[util.UserNameKey] = input.Username
doc := ai.DocumentFromText(input.Content, input.Metadata)
err := indexer.Index(ctx, &ai.IndexerRequest{
Documents: []*ai.Document{doc},
})
if err != nil {
return Response{
Code: 500,
Msg: err.Error(),
}, nil
}
return Response{
Code: 200,
Msg: "Document indexed successfully",
}, nil
})
graphIndexer, graphRetriever, err := graphrag.DefineIndexerAndRetriever(ctx, g)
if err != nil {
log.Fatal().Msgf("graphrag.DefineIndexerAndRetriever failed: %v", err)
}
genkit.DefineFlow(g, "index/graph", func(ctx context.Context, input *GraphInput) (Response, error) {
opt := graphrag.IndexReqOption{
UserId: input.UserID,
UserName: input.Username,
}
resDocName := ""
if v, ok := input.Metadata[graphrag.DocNameKey]; !ok {
// Generate random docName.
docName, err := graphrag.GenerateRandomDocName(8)
if err != nil {
return Response{
Code: 500,
Msg: fmt.Sprintf("generate random docName for document %w", err),
}, nil
}
input.Metadata[graphrag.DocNameKey] = docName
resDocName = docName
} else {
if str, isString := v.(string); isString {
resDocName = str
}
}
doc := ai.DocumentFromText(input.Content, input.Metadata)
err := graphIndexer.Index(ctx, &ai.IndexerRequest{
Documents: []*ai.Document{doc},
Options: &opt,
})
if err != nil {
return Response{
Code: 500,
Msg: fmt.Sprintf("index document: %w", err),
}, nil
}
return Response{
Code: 200,
Msg: fmt.Sprintf("Document indexed successfully, docname %s", resDocName),
}, nil
})
qa, err := question.InitQAStore(*pgConnString)
if err != nil {
log.Fatal().Msgf("InitQAStore failed: %v", err)
}
// Initialize KnowledgeClient with test parameters
kc := knowledge.NewKnowledgeClient(knowledge.ClientConfig{
Endpoint: "lkeap.tencentcloudapi.com",
Region: "ap-guangzhou",
})
if err := kc.Init(ctx); err != nil {
log.Fatal().Msgf("Failed to initialize KnowledgeClient: %v", err)
}
log.Info().Msg("KnowledgeClient initialized successfully")
// Define a simple flow that generates jokes about a given topic
genkit.DefineFlow(g, "chat", func(ctx context.Context, input *ChatInput) (Response, error) {
ctxAsJson, _ := json.Marshal(ctx)
log.Info().Msgf("input----ctxAsJson----%s", string(ctxAsJson))
inputAsJson, err := json.Marshal(input)
if err != nil {
return Response{
Code: 500,
Msg: fmt.Sprintf("json.Marshal: %w", err),
}, nil
}
log.Info().Msgf("input--------%s", string(inputAsJson))
idx, lastQa, lastok, err := qa.WriteAndGetLatestQA(context.Background(), question.QA{
FromID: &input.FromID,
From: &input.From,
Question: &input.Content,
To: &input.To,
ToID: &input.ToID,
Milvus: &input.Milvus,
Graph: &input.Graph,
})
if err != nil {
return Response{
Code: 500,
Msg: fmt.Sprintf("WriteAndGetLatestQA: %w", err),
}, nil
}
qaAsJson, err := json.Marshal(lastQa)
if err != nil {
return Response{
Code: 500,
Msg: fmt.Sprintf("json.Marshal(lastQa): %w", err),
}, nil
}
log.Info().Msgf("qaAsJson--------%s", string(qaAsJson))
promptInput := &simpleQaPromptInput{
Query: input.Content,
}
if lastok && lastQa.Summary != nil {
promptInput.Summary = *lastQa.Summary
}
metaData := make(map[string]any)
metaData[util.UserIdKey] = input.ToID
metaData[util.UserNameKey] = input.To
dRequest := ai.DocumentFromText(input.Content, metaData)
response, err := ai.Retrieve(ctx, retriever, ai.WithDocs(dRequest))
if err != nil {
log.Error().Msgf("milvus Retrieve err.Error() %s", err.Error())
} else {
var sb strings.Builder
for _, d := range response.Documents {
sb.WriteString(d.Content[0].Text)
sb.WriteByte('\n')
}
promptInput.Context = sb.String()
log.Info().Msgf("promptInput.Context: %s", promptInput.Context)
}
begin := time.Now()
graphResponse, err := ai.Retrieve(ctx, graphRetriever, ai.WithDocs(dRequest))
if err != nil {
log.Error().Msgf("graph Retrieve err.Error() %s", err.Error())
} else {
var sb strings.Builder
for _, d := range graphResponse.Documents {
sb.WriteString(d.Content[0].Text)
sb.WriteByte('\n')
}
promptInput.Graph = sb.String()
log.Info().Msgf("promptInput.Graph : %s", promptInput.Graph)
}
fmt.Println("graph time", time.Since(begin).Seconds())
simpleQaPrompt,err := defineSimpleQaPrompt(g,input.Model)
if err != nil {
return Response{
Code: 500,
Msg: fmt.Sprintf("index document: %w", err),
}, nil
}
resp, err := simpleQaPrompt.Execute(ctx, ai.WithInput(promptInput))
if err != nil {
return Response{
Code: 500,
Msg: fmt.Sprintf("index document: %w", err),
}, nil
}
if lastok {
if promptInput.Summary == "" {
promptInput.Summary = resp.Text()
}
log.Info().
Str("from", input.From).
Str("from_id", input.FromID).
Str("to", input.To).
Str("to_id", input.ToID).
Str("promptInput.Query", promptInput.Query).
Str("resp.Text()", resp.Text()).
Str("promptInput.Summary", promptInput.Summary).Msg("QueryRewriteWithSummary")
res, err := kc.QueryRewriteWithSummary(context.Background(), promptInput.Query, resp.Text(), promptInput.Summary)
if err != nil {
log.Error().Msg(err.Error())
} else {
qa.UpdateQAFields(context.Background(), idx, res.RewrittenQuery, resp.Text())
/*
{"RewrittenQuery":"Conversation summary: The available knowledge base does not contain information about the capital of the UK.","RawResponse":{"Response":{"Content":"Conversation summary: The available knowledge base does not contain information about the capital of the UK.","Usage":{"InputTokens":74,"OutputTokens":19,"TotalTokens":93},"RequestId":"15f1ce0c-a83f-4d95-af22-33a3bd829e8d"}}}
*/
}
} else {
qa.UpdateQAFields(context.Background(), idx, "", resp.Text())
}
log.Info().
Str("from", input.From).
Str("from_id", input.FromID).
Str("to", input.To).
Str("to_id", input.ToID).
Str("question", promptInput.Query).
Str("context", promptInput.Context).
Str("graph", promptInput.Graph).
Str("last summary", promptInput.Summary).
Str("answer", resp.Text()).
Msg("Question and answer pair recorded")
return Response{
Data: resp.Text(),
Code: 200,
}, nil
})
DefineDocumentFlow(g, indexer)
DefineGraphFlow(g, graphIndexer)
DefineModelsFlow(g)
DefineChatFlow(g, retriever, graphRetriever, *pgConnString)
// 配置限速器：每秒 10 次请求，突发容量 20，最大并发 5
rl := NewRateLimiter(10, 20, 5)
......@@ -445,49 +198,3 @@ type Response struct {
Code int `json:"code"`
Msg string `json:"msg"`
}
// defineSimpleQaPrompt 加载或定义指定名称的 Prompt
func defineSimpleQaPrompt(g *genkit.Genkit, promptName string) (*ai.Prompt, error) {
// 步骤 1：尝试查找现有的 Prompt
log.Info().
Str("method", "defineSimpleQaPrompt").
Str("prompt_name", promptName).
Msg("Attempting to lookup prompt")
prompt := genkit.LookupPrompt(g, promptName)
if prompt != nil {
log.Info().
Str("method", "defineSimpleQaPrompt").
Str("prompt_name", promptName).
Msg("Prompt found and loaded")
return prompt, nil
}
log.Info().
Str("method", "defineSimpleQaPrompt").
Str("prompt_name", promptName).
Msg("Prompt not found, defining new prompt")
// 步骤 2：如果未找到，定义新的 Prompt
simpleQaPrompt, err := genkit.DefinePrompt(g, promptName,
// ai.WithModelName("deepseek/deepseek-chat"),
ai.WithModelName(promptName),
ai.WithPrompt(simpleQaPromptTemplate),
ai.WithInputType(simpleQaPromptInput{}),
ai.WithOutputFormat(ai.OutputFormatText),
)
if err != nil {
log.Error().
Str("method", "defineSimpleQaPrompt").
Str("prompt_name", promptName).
Err(err).
Msg("Failed to define prompt")
return nil, err
}
log.Info().
Str("method", "defineSimpleQaPrompt").
Str("prompt_name", promptName).
Msg("Prompt defined successfully")
return simpleQaPrompt, nil
}
plugins/deepseek/deepseek.go
View file @ 409192b5
This diff is collapsed.
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment