Skip to content

Commit a66c7e2

Browse files
committed
ray bot
Signed-off-by: Max Pumperla <[email protected]>
1 parent 15ee3c8 commit a66c7e2

File tree

1 file changed

+54
-26
lines changed

1 file changed

+54
-26
lines changed

main.py

Lines changed: 54 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,6 @@
1717
files_to_diff_dict
1818
)
1919

20-
ray.init(object_store_memory=78643200)
21-
2220

2321
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
2422
logger = logging.getLogger("Docu Mentor")
@@ -55,35 +53,68 @@
5553
openai.api_key = os.environ.get("ANYSCALE_API_KEY")
5654

5755

58-
SYSTEM_CONTENT = """
59-
You are a helpful assistant.
56+
SYSTEM_CONTENT = """You are a helpful assistant.
6057
Improve the following <content>. Criticise syntax, grammar, punctuation, style, etc.
6158
Recommend common technical writing knowledge, such as used in Vale
6259
and the Google developer documentation style guide.
6360
If the content is good, don't comment on it.
64-
Do not comment on file names, just the actual text.
61+
You can use GitHub-flavored markdown syntax in your answer.
62+
"""
63+
64+
PROMPT = """Improve this content.
65+
Don't comment on file names or other meta data, just the actual text.
6566
The <content> will be in JSON format and contains file name keys and text values.
66-
You can use GitHub-flavored markdown syntax.
6767
Make sure to give very concise feedback per file.
6868
"""
6969

7070
def mentor(
7171
content,
7272
model="meta-llama/Llama-2-70b-chat-hf",
7373
system_content=SYSTEM_CONTENT,
74-
extra_instructions="Improve this content."
74+
prompt=PROMPT
7575
):
76-
"""The content can be any string in principle, but the system prompt is
77-
crafted for dictionary data of the form {'file_name': 'file_content'}.
78-
"""
79-
return openai.ChatCompletion.create(
76+
result = openai.ChatCompletion.create(
8077
model=model,
8178
messages=[
8279
{"role": "system", "content": system_content},
83-
{"role": "user", "content": f"This is the content: {content}. {extra_instructions}"},
80+
{"role": "user", "content": f"This is the content: {content}. {prompt}"},
8481
],
8582
temperature=0.7,
8683
)
84+
usage = result.get("usage")
85+
prompt_tokens = usage.get("prompt_tokens")
86+
completion_tokens = usage.get("completion_tokens")
87+
content = result["choices"][0]["message"]["content"]
88+
89+
return content, model, prompt_tokens, completion_tokens
90+
91+
try:
92+
ray.init()
93+
except:
94+
logger.info("Ray init failed.")
95+
96+
97+
@ray.remote
98+
def mentor_task(content, model, system_content, prompt):
99+
return mentor(content, model, system_content, prompt)
100+
101+
def ray_mentor(
102+
content: dict,
103+
model="meta-llama/Llama-2-70b-chat-hf",
104+
system_content=SYSTEM_CONTENT,
105+
prompt="Improve this content."
106+
):
107+
futures = [
108+
mentor_task.remote(v, model, system_content, prompt)
109+
for v in content.values()
110+
]
111+
suggestions = ray.get(futures)
112+
content = {k: v[0] for k, v in zip(content.keys(), suggestions)}
113+
prompt_tokens = sum(v[2] for v in suggestions)
114+
completion_tokens = sum(v[3] for v in suggestions)
115+
116+
return content, model, prompt_tokens, completion_tokens
117+
87118

88119

89120
app = FastAPI()
@@ -163,35 +194,32 @@ async def handle_webhook(request: Request):
163194
diff_response = await client.get(url, headers=headers)
164195
diff = diff_response.text
165196

166-
files_with_diff = files_to_diff_dict(diff)
197+
files = files_to_diff_dict(diff)
167198

168199
# Filter the dictionary
169200
if files_to_keep:
170-
files_with_diff = {
171-
k: files_with_diff[k]
172-
for k in files_with_diff
201+
files = {
202+
k: files[k]
203+
for k in files
173204
if any(sub in k for sub in files_to_keep)
174205
}
175-
176-
logger.info(files_with_diff.keys())
206+
logger.info(files.keys())
177207

178208
# Get suggestions from Docu Mentor
179-
chat_completion = mentor(files_with_diff)
209+
content, model, prompt_tokens, completion_tokens = ray_mentor(files) if ray.is_initialized() else mentor(files)
180210

181-
logger.info(chat_completion)
182-
model = chat_completion.get("model")
183-
usage = chat_completion.get("usage")
184-
prompt_tokens = usage.get("prompt_tokens")
185-
completion_tokens = usage.get("completion_tokens")
186-
content = chat_completion["choices"][0]["message"]["content"]
211+
print_content = ""
212+
for k, v in content.items():
213+
print_content += f"{k}:\n\t\{v}\n\n"
214+
logger.info(print_content)
187215

188216
# Let's comment on the PR
189217
await client.post(
190218
f"{comment['issue_url']}/comments",
191219
json={
192220
"body": f":rocket: Docu Mentor finished analysing your PR! :rocket:\n\n"
193221
+ "Take a look at your results:\n"
194-
+ f"{content}\n\n"
222+
+ f"{print_content}\n\n"
195223
+ "This bot is proudly powered by [Anyscale Endpoints](https://app.endpoints.anyscale.com/).\n"
196224
+ f"It used the model {model}, used {prompt_tokens} prompt tokens, "
197225
+ f"and {completion_tokens} completion tokens in total."

0 commit comments

Comments
 (0)