robinroy03/github-bot

1from flask import Flask, request2import requests3import json4import os5 6from jwt_setup import get_installation_access_token7 8 9app = Flask(__name__)10 11def get_llm_output(query: str):12    LLM_output = requests.post('https://robinroy03-fury-engine.hf.space/api/google/generate', json={"query": query, "llm": "gemini-1.5-pro", "knn": 3, "stream": False})13    LLM_output = json.loads(LLM_output.text)14    LLM_output = LLM_output['response'] + '\n\n**References**\n' + LLM_output['references']15 16    return LLM_output17 18 19def get_root_comment_id(node_id: str) -> str:20    """21    Parameters22    ----------23    node_id: str24      The ID of the comment node.25 26    Returns27    -------28    root comment node_id: str29      The ID of the root comment node ID30 31    """32    token = get_installation_access_token(33        os.environ.get('CHATBOT_PUB_KEY'),34        os.environ.get('APP_ID'),35        os.environ.get('INSTALLATION_ID')36    )37    headers = {"Authorization": f"token {token}"}38    query = "query get_root_comment_ID ($id: ID!){node(id: $id) {... on DiscussionComment {replyTo {id}}}}"39    variables = {"id": node_id}40 41    json_return = requests.post('https://api.github.com/graphql', json={'query': query, 'variables': variables}, headers=headers)42    json_return = json.loads(json_return.text)43 44    return json_return['data']['node']['replyTo']['id']45 46 47def post_discussion_reply(discussion_id: str, replytoid: str, body: str):48    """49    Parameters50    ----------51    discussion_id: str52        ID of the discussion, available on the webhook payload53    replytoid: str54        ID of the root reply. This is not the ID of a threaded reply. Use get_root_comment_id() for the root ID.55    body: str56        Body of the reply57 58    """59    token = get_installation_access_token(60        os.environ.get('CHATBOT_PUB_KEY'),61        os.environ.get('APP_ID'),62        os.environ.get('INSTALLATION_ID')63    )64 65    headers = {"Authorization": f"token {token}"}66    query = "mutation AddDiscussionComment($discussionId: ID!, $body: String!, $replyToId: ID!) {addDiscussionComment(input: {discussionId: $discussionId, body: $body, replyToId: $replyToId}) {clientMutationId}}"67    variables = {"discussionId": discussion_id, "body": body, "replyToId": replytoid}68 69    requests.post('https://api.github.com/graphql', json={'query': query, 'variables': variables}, headers=headers)70 71 72def post_discussion_comment(discussion_id: str, body: str):73    token = get_installation_access_token(74        os.environ.get('CHATBOT_PUB_KEY'),75        os.environ.get('APP_ID'),76        os.environ.get('INSTALLATION_ID')77    )78 79    headers = {"Authorization": f"token {token}"}80    query = "mutation AddDiscussionComment($discussionId: ID!, $body: String!) {addDiscussionComment(input: {discussionId: $discussionId, body: $body}) {clientMutationId}}"81    variables = {"discussionId": discussion_id, "body": body}82 83    requests.post('https://api.github.com/graphql', json={'query': query, 'variables': variables}, headers=headers)84 85 86@app.route("/github", methods=['POST'])87def main():88    payload = request.get_json()89    URL = "https://api.github.com/graphql"90 91    if 'discussion' in payload and 'comment' not in payload:92        discussion_id = payload['discussion']['node_id']93        discussion_title = payload['discussion']['title']94        discussion_body = payload['discussion']['body']95 96        if ("[X] I want an AI-generated answer as a first response." in discussion_body) or ("[x] I want an AI-generated answer as a first response." in discussion_body):97            post_discussion_comment(discussion_id=discussion_id, body=get_llm_output(f"{discussion_title}\n{discussion_body}"))98        else:99            return {"status": False}    100 101        return {"status": True}102    103    elif 'discussion' in payload and 'comment' in payload:104        if payload['comment']['parent_id'] != None and payload['comment']['body'].startswith("!fury"):105            post_discussion_reply(106                payload['discussion']['node_id'],107                get_root_comment_id(payload['comment']['node_id']),108                get_llm_output(payload['comment']['body'][6:])    109            )110        elif payload['comment']['parent_id'] == None and payload['comment']['body'].startswith("!fury"):111            post_discussion_reply(112                payload['discussion']['node_id'],113                payload['comment']['node_id'],114                get_llm_output(payload['comment']['body'][6:])115            )116 117        return {"status": True}118 119    else:120        return {"status": False}121