From 06b20a1772860075d1590356e449bde73b92f784 Mon Sep 17 00:00:00 2001
From: thucpn <thucsh2@gmail.com>
Date: Tue, 26 Dec 2023 15:59:51 +0700
Subject: [PATCH] fix: code review and bugs
---
.vscode/settings.json | 7 +-
packages/create-llama/questions.ts | 157 ++++++++----------
.../vectordbs/python/none/context.py | 12 +-
.../components/vectordbs/python/none/index.py | 2 +-
.../src/controllers/chat.controller.ts | 2 +-
.../simple/fastapi/app/api/routers/chat.py | 3 +-
.../types/simple/fastapi/app/engine/index.py | 34 +---
.../src/controllers/chat.controller.ts | 2 +-
.../streaming/fastapi/README-template.md | 4 +-
.../streaming/fastapi/app/engine/index.py | 34 +---
.../streaming/nextjs/app/api/chat/route.ts | 2 +-
11 files changed, 102 insertions(+), 157 deletions(-)
diff --git a/.vscode/settings.json b/.vscode/settings.json
index 9f6017380..329335c14 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -5,5 +5,8 @@
"[xml]": {
"editor.defaultFormatter": "redhat.vscode-xml"
},
- "jest.rootPath": "./packages/core"
-}
\ No newline at end of file
+ "jest.rootPath": "./packages/core",
+ "[python]": {
+ "editor.defaultFormatter": "ms-python.black-formatter"
+ }
+}
diff --git a/packages/create-llama/questions.ts b/packages/create-llama/questions.ts
index 42fd144c0..5b4a9ca20 100644
--- a/packages/create-llama/questions.ts
+++ b/packages/create-llama/questions.ts
@@ -189,95 +189,84 @@ export const askQuestions = async (
}
}
- if (
- program.framework === "express" ||
- program.framework === "nextjs" ||
- program.framework === "fastapi"
- ) {
- if (!program.model) {
- if (ciInfo.isCI) {
- program.model = getPrefOrDefault("model");
- } else {
- const { model } = await prompts(
- {
- type: "select",
- name: "model",
- message: "Which model would you like to use?",
- choices: [
- { title: "gpt-3.5-turbo", value: "gpt-3.5-turbo" },
- { title: "gpt-4", value: "gpt-4" },
- { title: "gpt-4-1106-preview", value: "gpt-4-1106-preview" },
- {
- title: "gpt-4-vision-preview",
- value: "gpt-4-vision-preview",
- },
- ],
- initial: 0,
- },
- handlers,
- );
- program.model = model;
- preferences.model = model;
- }
+ if (!program.model) {
+ if (ciInfo.isCI) {
+ program.model = getPrefOrDefault("model");
+ } else {
+ const { model } = await prompts(
+ {
+ type: "select",
+ name: "model",
+ message: "Which model would you like to use?",
+ choices: [
+ { title: "gpt-3.5-turbo", value: "gpt-3.5-turbo" },
+ { title: "gpt-4", value: "gpt-4" },
+ { title: "gpt-4-1106-preview", value: "gpt-4-1106-preview" },
+ {
+ title: "gpt-4-vision-preview",
+ value: "gpt-4-vision-preview",
+ },
+ ],
+ initial: 0,
+ },
+ handlers,
+ );
+ program.model = model;
+ preferences.model = model;
+ }
+ }
+
+ if (!program.engine) {
+ if (ciInfo.isCI) {
+ program.engine = getPrefOrDefault("engine");
+ } else {
+ const { engine } = await prompts(
+ {
+ type: "select",
+ name: "engine",
+ message: "Which data source would you like to use?",
+ choices: [
+ {
+ title: "No data, just a simple chat",
+ value: "simple",
+ },
+ { title: "Use an example PDF", value: "context" },
+ ],
+ initial: 1,
+ },
+ handlers,
+ );
+ program.engine = engine;
+ preferences.engine = engine;
}
}
if (
- program.framework === "express" ||
- program.framework === "nextjs" ||
- program.framework === "fastapi"
+ program.engine !== "simple" &&
+ !program.vectorDb &&
+ program.framework !== "fastapi"
) {
- if (!program.engine) {
- if (ciInfo.isCI) {
- program.engine = getPrefOrDefault("engine");
- } else {
- const { engine } = await prompts(
- {
- type: "select",
- name: "engine",
- message: "Which data source would you like to use?",
- choices: [
- {
- title: "No data, just a simple chat",
- value: "simple",
- },
- { title: "Use an example PDF", value: "context" },
- ],
- initial: 1,
- },
- handlers,
- );
- program.engine = engine;
- preferences.engine = engine;
- }
- }
- if (
- program.engine !== "simple" &&
- !program.vectorDb &&
- program.framework !== "fastapi"
- ) {
- if (ciInfo.isCI) {
- program.vectorDb = getPrefOrDefault("vectorDb");
- } else {
- const { vectorDb } = await prompts(
- {
- type: "select",
- name: "vectorDb",
- message: "Would you like to use a vector database?",
- choices: [
- {
- title: "No, just store the data in the file system",
- value: "none",
- },
- { title: "MongoDB", value: "mongo" },
- ],
- initial: 0,
- },
- handlers,
- );
- program.vectorDb = vectorDb;
- preferences.vectorDb = vectorDb;
- }
+ if (ciInfo.isCI) {
+ program.vectorDb = getPrefOrDefault("vectorDb");
+ } else {
+ const { vectorDb } = await prompts(
+ {
+ type: "select",
+ name: "vectorDb",
+ message: "Would you like to use a vector database?",
+ choices: [
+ {
+ title: "No, just store the data in the file system",
+ value: "none",
+ },
+ { title: "MongoDB", value: "mongo" },
+ ],
+ initial: 0,
+ },
+ handlers,
+ );
+ program.vectorDb = vectorDb;
+ preferences.vectorDb = vectorDb;
}
}
diff --git a/packages/create-llama/templates/components/vectordbs/python/none/context.py b/packages/create-llama/templates/components/vectordbs/python/none/context.py
index 4756d813d..271ac8726 100644
--- a/packages/create-llama/templates/components/vectordbs/python/none/context.py
+++ b/packages/create-llama/templates/components/vectordbs/python/none/context.py
@@ -1,9 +1,17 @@
-from llama_index import ServiceContext
+import os
-from app.context import create_base_context
+from llama_index import ServiceContext
+from llama_index.llms import OpenAI
from app.engine.constants import CHUNK_SIZE, CHUNK_OVERLAP
+def create_base_context():
+ model = os.getenv("MODEL", "gpt-3.5-turbo")
+ return ServiceContext.from_defaults(
+ llm=OpenAI(model=model),
+ )
+
+
def create_service_context():
base = create_base_context()
return ServiceContext.from_defaults(
diff --git a/packages/create-llama/templates/components/vectordbs/python/none/index.py b/packages/create-llama/templates/components/vectordbs/python/none/index.py
index 8f7d36030..0170d6e83 100644
--- a/packages/create-llama/templates/components/vectordbs/python/none/index.py
+++ b/packages/create-llama/templates/components/vectordbs/python/none/index.py
@@ -14,7 +14,7 @@ def get_chat_engine():
# check if storage already exists
if not os.path.exists(STORAGE_DIR):
raise Exception(
- "StorageContext is empty - call 'npm run generate' to generate the storage first"
+ "StorageContext is empty - call 'python app/engine/generate.py' to generate the storage first"
)
logger = logging.getLogger("uvicorn")
# load the existing index
diff --git a/packages/create-llama/templates/types/simple/express/src/controllers/chat.controller.ts b/packages/create-llama/templates/types/simple/express/src/controllers/chat.controller.ts
index 63a1b1ddd..46dfa5684 100644
--- a/packages/create-llama/templates/types/simple/express/src/controllers/chat.controller.ts
+++ b/packages/create-llama/templates/types/simple/express/src/controllers/chat.controller.ts
@@ -33,7 +33,7 @@ export const chat = async (req: Request, res: Response) => {
}
const llm = new OpenAI({
- model: process.env.MODEL,
+ model: process.env.MODEL || "gpt-3.5-turbo",
});
const lastMessageContent = getLastMessageContent(
diff --git a/packages/create-llama/templates/types/simple/fastapi/app/api/routers/chat.py b/packages/create-llama/templates/types/simple/fastapi/app/api/routers/chat.py
index b90820550..9b1b1a1a8 100644
--- a/packages/create-llama/templates/types/simple/fastapi/app/api/routers/chat.py
+++ b/packages/create-llama/templates/types/simple/fastapi/app/api/routers/chat.py
@@ -2,7 +2,8 @@ from typing import List
from fastapi import APIRouter, Depends, HTTPException, status
from llama_index import VectorStoreIndex
-from llama_index.llms.base import MessageRole, ChatMessage
+from llama_index.llms.base import ChatMessage
+from llama_index.llms.types import MessageRole
from pydantic import BaseModel
from app.engine.index import get_chat_engine
diff --git a/packages/create-llama/templates/types/simple/fastapi/app/engine/index.py b/packages/create-llama/templates/types/simple/fastapi/app/engine/index.py
index 8a5bfa5c5..47a9b083c 100644
--- a/packages/create-llama/templates/types/simple/fastapi/app/engine/index.py
+++ b/packages/create-llama/templates/types/simple/fastapi/app/engine/index.py
@@ -1,17 +1,9 @@
import os
-import logging
-from llama_index import (
- SimpleDirectoryReader,
- StorageContext,
- VectorStoreIndex,
- load_index_from_storage,
- ServiceContext,
-)
+from llama_index import ServiceContext
+from llama_index.chat_engine import SimpleChatEngine
from llama_index.llms import OpenAI
-STORAGE_DIR = "./storage" # directory to cache the generated index
-DATA_DIR = "./data" # directory containing the documents to index
def create_base_context():
model = os.getenv("MODEL", "gpt-3.5-turbo")
@@ -19,26 +11,6 @@ def create_base_context():
llm=OpenAI(model=model),
)
-def get_index():
- service_context = create_base_context()
- logger = logging.getLogger("uvicorn")
- # check if storage already exists
- if not os.path.exists(STORAGE_DIR):
- logger.info("Creating new index")
- # load the documents and create the index
- documents = SimpleDirectoryReader(DATA_DIR).load_data()
- index = VectorStoreIndex.from_documents(documents,service_context=service_context)
- # store it for later
- index.storage_context.persist(STORAGE_DIR)
- logger.info(f"Finished creating new index. Stored in {STORAGE_DIR}")
- else:
- # load the existing index
- logger.info(f"Loading index from {STORAGE_DIR}...")
- storage_context = StorageContext.from_defaults(persist_dir=STORAGE_DIR)
- index = load_index_from_storage(storage_context,service_context=service_context)
- logger.info(f"Finished loading index from {STORAGE_DIR}")
- return index
def get_chat_engine():
- index = get_index()
- return index.as_chat_engine()
+ return SimpleChatEngine.from_defaults(service_context=create_base_context())
diff --git a/packages/create-llama/templates/types/streaming/express/src/controllers/chat.controller.ts b/packages/create-llama/templates/types/streaming/express/src/controllers/chat.controller.ts
index f1f50b48b..4bd1c8da6 100644
--- a/packages/create-llama/templates/types/streaming/express/src/controllers/chat.controller.ts
+++ b/packages/create-llama/templates/types/streaming/express/src/controllers/chat.controller.ts
@@ -35,7 +35,7 @@ export const chat = async (req: Request, res: Response) => {
}
const llm = new OpenAI({
- model: process.env.MODEL,
+ model: process.env.MODEL || "gpt-3.5-turbo",
});
const chatEngine = await createChatEngine(llm);
diff --git a/packages/create-llama/templates/types/streaming/fastapi/README-template.md b/packages/create-llama/templates/types/streaming/fastapi/README-template.md
index a91b9f15a..7f659b689 100644
--- a/packages/create-llama/templates/types/streaming/fastapi/README-template.md
+++ b/packages/create-llama/templates/types/streaming/fastapi/README-template.md
@@ -17,10 +17,10 @@ Example `.env` file:
OPENAI_API_KEY=<openai_api_key>
```
-Second, generate the embeddings of the documents in the `./data` directory:
+Second, generate the embeddings of the documents in the `./data` directory (if this folder exists - otherwise, skip this step):
```
-python app/engine/generate.py
+python app/engine/generate.py
```
Third, run the development server:
diff --git a/packages/create-llama/templates/types/streaming/fastapi/app/engine/index.py b/packages/create-llama/templates/types/streaming/fastapi/app/engine/index.py
index 8a5bfa5c5..47a9b083c 100644
--- a/packages/create-llama/templates/types/streaming/fastapi/app/engine/index.py
+++ b/packages/create-llama/templates/types/streaming/fastapi/app/engine/index.py
@@ -1,17 +1,9 @@
import os
-import logging
-from llama_index import (
- SimpleDirectoryReader,
- StorageContext,
- VectorStoreIndex,
- load_index_from_storage,
- ServiceContext,
-)
+from llama_index import ServiceContext
+from llama_index.chat_engine import SimpleChatEngine
from llama_index.llms import OpenAI
-STORAGE_DIR = "./storage" # directory to cache the generated index
-DATA_DIR = "./data" # directory containing the documents to index
def create_base_context():
model = os.getenv("MODEL", "gpt-3.5-turbo")
@@ -19,26 +11,6 @@ def create_base_context():
llm=OpenAI(model=model),
)
-def get_index():
- service_context = create_base_context()
- logger = logging.getLogger("uvicorn")
- # check if storage already exists
- if not os.path.exists(STORAGE_DIR):
- logger.info("Creating new index")
- # load the documents and create the index
- documents = SimpleDirectoryReader(DATA_DIR).load_data()
- index = VectorStoreIndex.from_documents(documents,service_context=service_context)
- # store it for later
- index.storage_context.persist(STORAGE_DIR)
- logger.info(f"Finished creating new index. Stored in {STORAGE_DIR}")
- else:
- # load the existing index
- logger.info(f"Loading index from {STORAGE_DIR}...")
- storage_context = StorageContext.from_defaults(persist_dir=STORAGE_DIR)
- index = load_index_from_storage(storage_context,service_context=service_context)
- logger.info(f"Finished loading index from {STORAGE_DIR}")
- return index
def get_chat_engine():
- index = get_index()
- return index.as_chat_engine()
+ return SimpleChatEngine.from_defaults(service_context=create_base_context())
diff --git a/packages/create-llama/templates/types/streaming/nextjs/app/api/chat/route.ts b/packages/create-llama/templates/types/streaming/nextjs/app/api/chat/route.ts
index b83855bc1..ff00a3894 100644
--- a/packages/create-llama/templates/types/streaming/nextjs/app/api/chat/route.ts
+++ b/packages/create-llama/templates/types/streaming/nextjs/app/api/chat/route.ts
@@ -42,7 +42,7 @@ export async function POST(request: NextRequest) {
}
const llm = new OpenAI({
- model: process.env.MODEL,
+ model: process.env.MODEL || "gpt-3.5-turbo",
maxTokens: 512,
});
--
GitLab