-
-
Notifications
You must be signed in to change notification settings - Fork 24
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
v2.0.93; search and load old conversations
- Loading branch information
1 parent
ba7ece8
commit 7b80c43
Showing
10 changed files
with
230 additions
and
23 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
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
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
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
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,146 @@ | ||
""" | ||
LetMeDoIt AI Plugin - search chat records | ||
search and open old chat records | ||
[FUNCTION_CALL] | ||
""" | ||
|
||
from letmedoit import config | ||
from letmedoit.health_check import HealthCheck | ||
from pathlib import Path | ||
from chromadb.config import Settings | ||
import uuid, os, chromadb, re | ||
from letmedoit.utils.shared_utils import SharedUtil | ||
from prompt_toolkit import print_formatted_text, HTML | ||
|
||
chat_store = os.path.join(config.getFiles(), "chats") | ||
Path(chat_store).mkdir(parents=True, exist_ok=True) | ||
chroma_client = chromadb.PersistentClient(chat_store, Settings(anonymized_telemetry=False)) | ||
|
||
def get_or_create_collection(collection_name): | ||
collection = chroma_client.get_or_create_collection( | ||
name=collection_name, | ||
metadata={"hnsw:space": "cosine"}, | ||
embedding_function=HealthCheck.getEmbeddingFunction(), | ||
) | ||
return collection | ||
|
||
def add_vector(collection, text, metadata): | ||
id = str(uuid.uuid4()) | ||
collection.add( | ||
documents = [text], | ||
metadatas = [metadata], | ||
ids = [id] | ||
) | ||
|
||
def query_vectors(collection, query, n): | ||
return collection.query( | ||
query_texts=[query], | ||
n_results = n, | ||
) | ||
|
||
def save_chat_record(timestamp, order, record): | ||
role = record.get("role", "") | ||
content = record.get("content", "") | ||
if role and role in ("user", "assistant") and content: | ||
collection = get_or_create_collection("chats") | ||
metadata = { | ||
"timestamp": timestamp, | ||
"order": order, | ||
"role": role, | ||
} | ||
add_vector(collection, content, metadata) | ||
config.save_chat_record = save_chat_record | ||
|
||
def search_chats(function_args): | ||
query = function_args.get("query") # required | ||
config.print3(f"""Query: {query}""") | ||
collection = get_or_create_collection("chats") | ||
res = query_vectors(collection, query, config.chatRecordClosestMatches) | ||
config.stopSpinning() | ||
if res: | ||
exampleID = "" | ||
# display search results | ||
config.print2(config.divider) | ||
print(">>> retrieved chat records: ") | ||
for metadata, document in zip(res["metadatas"][0], res["documents"][0]): | ||
config.print(config.divider) | ||
config.print3(f"""Chat ID: {metadata["timestamp"]}""") | ||
if not exampleID: | ||
exampleID = metadata["timestamp"] | ||
config.print3(f"""Order: {metadata["order"]}""") | ||
config.print3(f"""Role: {metadata["role"]}""") | ||
config.print3(f"""Content: {document}""") | ||
config.print(config.divider) | ||
config.print2("Tips: You can load old chat records by quoting a chat ID or timestamp, e.g.") | ||
config.print(f">>> Load chat records with this ID: {exampleID}") | ||
config.print2(config.divider) | ||
return "" | ||
|
||
def load_chats(function_args): | ||
config.stopSpinning() | ||
timestamp = function_args.get("id") # required | ||
config.print3(f"Loading chat records: {timestamp} ...") | ||
|
||
try: | ||
folderPath = os.path.join(config.getFiles(), "chats", re.sub("^([0-9]+?\-[0-9]+?)\-.*?$", r"\1", timestamp)) | ||
chatFile = os.path.join(folderPath, f"{timestamp}.txt") | ||
with open(chatFile, "r", encoding="utf-8") as fileObj: | ||
messages = fileObj.read() | ||
currentMessages = eval(messages) | ||
if type(currentMessages) == list: | ||
config.currentMessages = currentMessages | ||
# display loaded messages | ||
print("") | ||
for index, i in enumerate(config.currentMessages): | ||
role = i.get("role", "") | ||
content = i.get("content", "") | ||
if role and role in ("user", "assistant") and content: | ||
if role == "user": | ||
print_formatted_text(HTML(f"<{config.terminalPromptIndicatorColor1}>>>> </{config.terminalPromptIndicatorColor1}><{config.terminalCommandEntryColor1}>{content}</{config.terminalCommandEntryColor1}>")) | ||
else: | ||
config.print(content) | ||
if role == 'assistant' and not index == len(config.currentMessages) - 2: | ||
print("") | ||
else: | ||
config.print3(f"Failed to load chat records '{timestamp}' due to invalid format!") | ||
except: | ||
config.print3(f"Failed to load chat records: {timestamp}\n") | ||
SharedUtil.showErrors() | ||
|
||
functionSignature1 = { | ||
"name": "search_chats", | ||
"description": """Search chat records""", | ||
"parameters": { | ||
"type": "object", | ||
"properties": { | ||
"query": { | ||
"type": "string", | ||
"description": "The search query in detail" | ||
}, | ||
}, | ||
"required": ["query"] | ||
} | ||
} | ||
|
||
functionSignature2 = { | ||
"name": "load_chats", | ||
"description": """Load or open old chat records if chat ID or timestamp is given""", | ||
"parameters": { | ||
"type": "object", | ||
"properties": { | ||
"id": { | ||
"type": "string", | ||
"description": "The chat ID or timestamp" | ||
}, | ||
}, | ||
"required": ["id"] | ||
} | ||
} | ||
|
||
config.inputSuggestions += ["Search chat records: ", "Load chat records with this ID: "] | ||
config.pluginsWithFunctionCall += ["search_chats", "load_chats"] | ||
config.chatGPTApiFunctionSignatures += [functionSignature1, functionSignature2] | ||
config.chatGPTApiAvailableFunctions["search_chats"] = search_chats | ||
config.chatGPTApiAvailableFunctions["load_chats"] = load_chats |
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
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
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
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
Oops, something went wrong.