Skip to content

Commit

Permalink
feat: Add publisher agent, merge code with streaming template (#324)
Browse files Browse the repository at this point in the history
---------
Co-authored-by: Marcus Schiesser <[email protected]>
  • Loading branch information
leehuwuj authored Oct 2, 2024
1 parent 3e8057a commit 04ddebc
Show file tree
Hide file tree
Showing 57 changed files with 1,408 additions and 884 deletions.
5 changes: 5 additions & 0 deletions .changeset/flat-singers-share.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"create-llama": patch
---

Add publisher agent to multi-agents for generating documents (PDF and HTML)
5 changes: 5 additions & 0 deletions .changeset/gorgeous-penguins-shout.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"create-llama": patch
---

Allow tool selection for multi-agents (Python and TS)
1 change: 1 addition & 0 deletions e2e/python/resolve_dependencies.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ if (
const toolOptions = [
"wikipedia.WikipediaToolSpec",
"google.GoogleSearchToolSpec",
"document_generator",
];

const dataSources = [
Expand Down
4 changes: 3 additions & 1 deletion e2e/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,9 @@ export async function runCreateLlama({
if (appType) {
commandArgs.push(appType);
}
if (!useLlamaParse) {
if (useLlamaParse) {
commandArgs.push("--use-llama-parse");
} else {
commandArgs.push("--no-llama-parse");
}

Expand Down
39 changes: 20 additions & 19 deletions helpers/env-variables.ts
Original file line number Diff line number Diff line change
Expand Up @@ -426,34 +426,35 @@ const getToolEnvs = (tools?: Tool[]): EnvVar[] => {
const getSystemPromptEnv = (
tools?: Tool[],
dataSources?: TemplateDataSource[],
framework?: TemplateFramework,
template?: TemplateType,
): EnvVar[] => {
const defaultSystemPrompt =
"You are a helpful assistant who helps users with their questions.";

const systemPromptEnv: EnvVar[] = [];
// build tool system prompt by merging all tool system prompts
let toolSystemPrompt = "";
tools?.forEach((tool) => {
const toolSystemPromptEnv = tool.envVars?.find(
(env) => env.name === TOOL_SYSTEM_PROMPT_ENV_VAR,
);
if (toolSystemPromptEnv) {
toolSystemPrompt += toolSystemPromptEnv.value + "\n";
}
});
// multiagent template doesn't need system prompt
if (template !== "multiagent") {
let toolSystemPrompt = "";
tools?.forEach((tool) => {
const toolSystemPromptEnv = tool.envVars?.find(
(env) => env.name === TOOL_SYSTEM_PROMPT_ENV_VAR,
);
if (toolSystemPromptEnv) {
toolSystemPrompt += toolSystemPromptEnv.value + "\n";
}
});

const systemPrompt = toolSystemPrompt
? `\"${toolSystemPrompt}\"`
: defaultSystemPrompt;
const systemPrompt = toolSystemPrompt
? `\"${toolSystemPrompt}\"`
: defaultSystemPrompt;

const systemPromptEnv = [
{
systemPromptEnv.push({
name: "SYSTEM_PROMPT",
description: "The system prompt for the AI model.",
value: systemPrompt,
},
];

});
}
if (tools?.length == 0 && (dataSources?.length ?? 0 > 0)) {
const citationPrompt = `'You have provided information from a knowledge base that has been passed to you in nodes of information.
Each node has useful metadata such as node ID, file name, page, etc.
Expand Down Expand Up @@ -559,7 +560,7 @@ export const createBackendEnvFile = async (
...getToolEnvs(opts.tools),
...getTemplateEnvs(opts.template),
...getObservabilityEnvs(opts.observability),
...getSystemPromptEnv(opts.tools, opts.dataSources, opts.framework),
...getSystemPromptEnv(opts.tools, opts.dataSources, opts.template),
];
// Render and write env file
const content = renderEnvVar(envVars);
Expand Down
40 changes: 32 additions & 8 deletions helpers/python.ts
Original file line number Diff line number Diff line change
Expand Up @@ -364,7 +364,12 @@ export const installPythonTemplate = async ({
| "modelConfig"
>) => {
console.log("\nInitializing Python project with template:", template, "\n");
const templatePath = path.join(templatesDir, "types", template, framework);
let templatePath;
if (template === "extractor") {
templatePath = path.join(templatesDir, "types", "extractor", framework);
} else {
templatePath = path.join(templatesDir, "types", "streaming", framework);
}
await copy("**", root, {
parents: true,
cwd: templatePath,
Expand Down Expand Up @@ -401,23 +406,42 @@ export const installPythonTemplate = async ({
cwd: path.join(compPath, "services", "python"),
});
}

if (template === "streaming") {
// For the streaming template only:
// Copy engine code
if (template === "streaming" || template === "multiagent") {
// Select and copy engine code based on data sources and tools
let engine;
if (dataSources.length > 0 && (!tools || tools.length === 0)) {
console.log("\nNo tools selected - use optimized context chat engine\n");
engine = "chat";
} else {
// Multiagent always uses agent engine
if (template === "multiagent") {
engine = "agent";
} else {
// For streaming, use chat engine by default
// Unless tools are selected, in which case use agent engine
if (dataSources.length > 0 && (!tools || tools.length === 0)) {
console.log(
"\nNo tools selected - use optimized context chat engine\n",
);
engine = "chat";
} else {
engine = "agent";
}
}

// Copy engine code
await copy("**", enginePath, {
parents: true,
cwd: path.join(compPath, "engines", "python", engine),
});
}

if (template === "multiagent") {
// Copy multi-agent code
await copy("**", path.join(root), {
parents: true,
cwd: path.join(compPath, "multiagent", "python"),
rename: assetRelocator,
});
}

console.log("Adding additional dependencies");

const addOnDependencies = getAdditionalDependencies(
Expand Down
23 changes: 23 additions & 0 deletions helpers/tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,29 @@ For better results, you can specify the region parameter to get results from a s
},
],
},
{
display: "Document generator",
name: "document_generator",
supportedFrameworks: ["fastapi", "nextjs", "express"],
dependencies: [
{
name: "xhtml2pdf",
version: "^0.2.14",
},
{
name: "markdown",
version: "^3.7",
},
],
type: ToolType.LOCAL,
envVars: [
{
name: TOOL_SYSTEM_PROMPT_ENV_VAR,
description: "System prompt for document generator tool.",
value: `If user request for a report or a post, use document generator tool to create a file and reply with the link to the file.`,
},
],
},
{
display: "Code Interpreter",
name: "interpreter",
Expand Down
5 changes: 4 additions & 1 deletion helpers/typescript.ts
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,10 @@ export const installTSTemplate = async ({
// Select and copy engine code based on data sources and tools
let engine;
tools = tools ?? [];
if (dataSources.length > 0 && tools.length === 0) {
// multiagent template always uses agent engine
if (template === "multiagent") {
engine = "agent";
} else if (dataSources.length > 0 && tools.length === 0) {
console.log("\nNo tools selected - use optimized context chat engine\n");
engine = "chat";
} else {
Expand Down
16 changes: 8 additions & 8 deletions questions.ts
Original file line number Diff line number Diff line change
Expand Up @@ -141,12 +141,10 @@ export const getDataSourceChoices = (
});
}
if (selectedDataSource === undefined || selectedDataSource.length === 0) {
if (template !== "multiagent") {
choices.push({
title: "No datasource",
value: "none",
});
}
choices.push({
title: "No datasource",
value: "none",
});
choices.push({
title:
process.platform !== "linux"
Expand Down Expand Up @@ -734,8 +732,10 @@ export const askQuestions = async (
}
}

if (!program.tools && program.template === "streaming") {
// TODO: allow to select tools also for multi-agent framework
if (
!program.tools &&
(program.template === "streaming" || program.template === "multiagent")
) {
if (ciInfo.isCI) {
program.tools = getPrefOrDefault("tools");
} else {
Expand Down
2 changes: 1 addition & 1 deletion templates/components/engines/python/agent/engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
from llama_index.core.tools.query_engine import QueryEngineTool


def get_chat_engine(filters=None, params=None, event_handlers=None):
def get_chat_engine(filters=None, params=None, event_handlers=None, **kwargs):
system_prompt = os.getenv("SYSTEM_PROMPT")
top_k = int(os.getenv("TOP_K", 0))
tools = []
Expand Down
27 changes: 20 additions & 7 deletions templates/components/engines/python/agent/tools/__init__.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
import importlib
import os

import yaml
import importlib
from llama_index.core.tools.tool_spec.base import BaseToolSpec
from llama_index.core.tools.function_tool import FunctionTool
from llama_index.core.tools.tool_spec.base import BaseToolSpec


class ToolType:
Expand Down Expand Up @@ -40,14 +41,26 @@ def load_tools(tool_type: str, tool_name: str, config: dict) -> list[FunctionToo
raise ValueError(f"Failed to load tool {tool_name}: {e}")

@staticmethod
def from_env() -> list[FunctionTool]:
tools = []
def from_env(
map_result: bool = False,
) -> list[FunctionTool] | dict[str, FunctionTool]:
"""
Load tools from the configured file.
Params:
- use_map: if True, return map of tool name and the tool itself
"""
if map_result:
tools = {}
else:
tools = []
if os.path.exists("config/tools.yaml"):
with open("config/tools.yaml", "r") as f:
tool_configs = yaml.safe_load(f)
for tool_type, config_entries in tool_configs.items():
for tool_name, config in config_entries.items():
tools.extend(
ToolFactory.load_tools(tool_type, tool_name, config)
)
tool = ToolFactory.load_tools(tool_type, tool_name, config)
if map_result:
tools[tool_name] = tool
else:
tools.extend(tool)
return tools
Loading

0 comments on commit 04ddebc

Please sign in to comment.