Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commitf78c7df

Browse files
committed
start converting to async
1 parenta39496d commitf78c7df

File tree

2 files changed

+61
-44
lines changed

2 files changed

+61
-44
lines changed

‎pgml-apps/discord-bot/bot.py‎

Lines changed: 50 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,17 @@
11
# We start by importing necessary packages.
22
importdiscord
3-
frompsycopg_poolimportConnectionPool
3+
frompsycopg_poolimportAsyncConnectionPool
44
frompgmlimportDatabase
55
fromlangchain.document_loadersimportDirectoryLoader
66

7+
78
# Create the Bot class
89
classBot:
910
# Initialize the Bot with connection info and set up a database connection pool
1011
def__init__(self,conninfo:str):
1112
self.conninfo=conninfo
12-
self.pool=ConnectionPool(conninfo)
13-
self.pgml=Database(conninfo)
13+
self.pool=AsyncConnectionPool(conninfo)
14+
self.pgml=Database(conninfo)
1415

1516
# Initializes the Discord client with certain intents
1617
definit_discord_client(self):
@@ -19,54 +20,59 @@ def init_discord_client(self):
1920
returndiscord.Client(intents=intents)
2021

2122
# Ingests data from a directory, process it and register models, text splitters and generate chunks and embeddings
22-
defingest(self,path:str,collection_name:str):
23-
docs=self.load_documents(path)
24-
content=self.create_content(docs)
25-
collection=self.pgml.create_or_get_collection(collection_name)
26-
collection.upsert_documents(content)
27-
embedding_model_id=self.register_embedding_model(collection)
23+
asyncdefingest(self,path:str,collection_name:str):
24+
docs=awaitself.load_documents(path)
25+
content=awaitself.create_content(docs)
26+
collection=awaitself.pgml.create_or_get_collection(collection_name)
27+
awaitcollection.upsert_documents(content)
28+
embedding_model_id=awaitself.register_embedding_model(collection)
2829
splitter_id=self.register_text_splitter(collection)
29-
collection.generate_chunks(splitter_id=splitter_id)
30-
collection.generate_embeddings(model_id=embedding_model_id,splitter_id=splitter_id)
31-
32-
defcreate_or_get_collection(self,collection_name:str):
33-
returnself.pgml.create_or_get_collection(collection_name)
30+
awaitcollection.generate_chunks(splitter_id=splitter_id)
31+
awaitcollection.generate_embeddings(
32+
model_id=embedding_model_id,splitter_id=splitter_id
33+
)
34+
35+
asyncdefcreate_or_get_collection(self,collection_name:str):
36+
returnawaitself.pgml.create_or_get_collection(collection_name)
3437

3538
# Loads markdown documents from a directory
36-
defload_documents(self,path:str):
39+
asyncdefload_documents(self,path:str):
3740
print(f"Loading documents from{path}")
38-
loader=DirectoryLoader(path,glob='*.md')
41+
loader=DirectoryLoader(path,glob="*.md")
3942
docs=loader.load()
4043
print(f"Loaded{len(docs)} documents")
4144
returndocs
4245

4346
# Prepare content by iterating over each document
44-
defcreate_content(self,docs):
45-
return [{"text":doc.page_content,"source":doc.metadata['source']}fordocindocs]
46-
47+
asyncdefcreate_content(self,docs):
48+
return [
49+
{"text":doc.page_content,"source":doc.metadata["source"]}fordocindocs
50+
]
4751

4852
# Register an embedding model to the collection
49-
defregister_embedding_model(self,collection):
50-
embedding_model_id=collection.register_model(
53+
asyncdefregister_embedding_model(self,collection):
54+
embedding_model_id=awaitcollection.register_model(
5155
model_name="hkunlp/instructor-xl",
5256
model_params={"instruction":"Represent the document for retrieval: "},
5357
)
5458
returnembedding_model_id
5559

5660
# Register a text splitter to the collection
57-
defregister_text_splitter(self,collection,chunk_size:int=1500,chunk_overlap:int=40):
58-
splitter_id=collection.register_text_splitter(
61+
asyncdefregister_text_splitter(
62+
self,collection,chunk_size:int=1500,chunk_overlap:int=40
63+
):
64+
splitter_id=awaitcollection.register_text_splitter(
5965
splitter_name="RecursiveCharacterTextSplitter",
6066
splitter_params={"chunk_size":chunk_size,"chunk_overlap":chunk_overlap},
6167
)
6268
returnsplitter_id
6369

6470
# Run an SQL query and return the result
6571
asyncdefrun_query(self,statement:str,sql_params:tuple=None):
66-
conn=self.pool.getconn()
72+
conn=awaitself.pool.getconn()
6773
cur=conn.cursor()
6874
try:
69-
cur.execute(statement,sql_params)
75+
awaitcur.execute(statement,sql_params)
7076
returncur.fetchone()
7177
exceptExceptionase:
7278
print(e)
@@ -75,40 +81,45 @@ async def run_query(self, statement: str, sql_params: tuple = None):
7581

7682
# Query a collection with a string and return vector search results
7783
asyncdefquery_collection(self,collection_name:str,query:str):
78-
collection=self.pgml.create_or_get_collection(collection_name)
84+
collection=awaitself.pgml.create_or_get_collection(collection_name)
7985
returncollection.vector_search(
8086
query,
8187
top_k=3,
8288
model_id=2,
8389
splitter_id=2,
84-
query_parameters={"instruction":"Represent the question for retrieving supporting documents: "},
90+
query_parameters={
91+
"instruction":"Represent the question for retrieving supporting documents: "
92+
},
8593
)
8694

8795
# Start the Discord bot, listen to messages in 'bot-testing' channel and handle the messages
88-
defstart(self,collection_name:str,discord_token:str,channel_name:str):
96+
asyncdefstart(self,collection_name:str,discord_token:str,channel_name:str):
8997
self.discord_token=discord_token
9098
self.discord_client=self.init_discord_client()
9199

92100
@self.discord_client.event
93101
asyncdefon_ready():
94-
print(f'We have logged in as{self.discord_client.user}')
102+
print(f"We have logged in as{self.discord_client.user}")
95103

96104
@self.discord_client.event
97105
asyncdefon_message(message):
98-
print(f"Message from{message.author}:{message.content}")
99-
100-
ifmessage.author!=self.discord_client.userandmessage.channel.name==channel_name:
106+
print(f"Message from{message.author}:{message.content}")
107+
108+
if (
109+
message.author!=self.discord_client.user
110+
andmessage.channel.name==channel_name
111+
):
101112
awaitself.handle_message(collection_name,message)
102113

103114
self.discord_client.run(self.discord_token)
104115

105116
# Handle incoming messages, perform a search on the collection, and respond with a generated answer
106117
asyncdefhandle_message(self,collection_name,message):
107-
print(f'Message from{message.author}:{message.content}')
118+
print(f"Message from{message.author}:{message.content}")
108119
print("Searching the vector database")
109120
res=awaitself.query_collection(collection_name,message.content)
110121
print(f"Found{len(res)} results")
111-
context=self.build_context(res)
122+
context=awaitself.build_context(res)
112123
print("Running Completion query")
113124
completion=awaitself.run_transform_sql(context,message.content)
114125
print("Preparing response")
@@ -117,8 +128,8 @@ async def handle_message(self, collection_name, message):
117128
awaitmessage.channel.send(response)
118129

119130
# Build the context for the message from search results
120-
defbuild_context(self,res):
121-
return'\n'.join([f'{r["chunk"]}'forrinres])
131+
asyncdefbuild_context(self,res):
132+
return"\n".join([f'{r["chunk"]}'forrinres])
122133

123134
# Run a SQL function 'pgml.transform' to get a generated answer for the message
124135
asyncdefrun_transform_sql(self,context,message_content):
@@ -144,10 +155,10 @@ def prepare_prompt(self, context, message_content):
144155
145156
Context:
146157
{context}
147-
QUESTION<<{message_content}
158+
QUESTION<<{message_content}>>
148159
ANSWER<<"""
149160

150161
# Prepare the bot's response by removing the original prompt from the generated text
151162
defprepare_response(self,completion,context,message_content):
152-
generated_text=completion[0][0][0]['generated_text']
153-
returngenerated_text.replace(self.prepare_prompt(context,message_content),'')
163+
generated_text=completion[0][0][0]["generated_text"]
164+
returngenerated_text.replace(self.prepare_prompt(context,message_content),"")

‎pgml-apps/discord-bot/start.py‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,25 @@
11
# Importing libraries
22
frombotimportBot
33
fromdotenvimportload_dotenv
4+
importasyncio
5+
46
importos
7+
58
load_dotenv()
69

710

811
# get environment variables
912
pg_connection_string=os.getenv("PGML_CONNECTION_STR")
1013
collection_name=os.getenv("COLLECTION_NAME")
11-
markdown_folder_path=os.getenv("CONTENT_PATH")
1214
discord_token=os.getenv("DISCORD_TOKEN")
1315
channel_name=os.getenv("DISCORD_CHANNEL")
1416

15-
## initialize bot
16-
pgml_bot=Bot(conninfo=pg_connection_string)
1717

18-
## start discord bot
19-
pgml_bot.start(collection_name,discord_token,channel_name)
18+
asyncdefmain():
19+
pgml_bot=Bot(conninfo=pg_connection_string)
20+
awaitpgml_bot.start(collection_name,discord_token,channel_name)
21+
22+
23+
asyncio.run(main())
24+
25+
# ## start discord bot

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp