diff --git a/main.py b/main.py index f18f58f..a5a02f6 100644 --- a/main.py +++ b/main.py @@ -1,7 +1,7 @@ from fastapi import FastAPI from config import __CONFIG__ -# import mysql_connector +import mysql_connector from fastgpt_uploader import upload2fastgpt from semanticscholar import search_paper @@ -12,7 +12,8 @@ def query(query:str): list = search_paper(query) for i in list: res.append({ - 'id':'','datasetId':__CONFIG__['fastgpt_setId'], + 'id':i['paperId'], + 'datasetId':__CONFIG__['fastgpt_setId'], 'collectionId':__CONFIG__['fastgpt_colId'], 'sourceName':'Semantic Scholar', 'sourceId?':'', @@ -25,16 +26,23 @@ def query(query:str): } ] }) - print(i['title']) return res @app.get("/fastdoi") async def get_reference(questions): res = query(questions) - upload2fastgpt(res) - return res + new = [] + print(res) + for i in res: + if not mysql_connector.is_loaded(i['id']): + new.append(i) + if(upload2fastgpt(new)): + for i in new: + mysql_connector.new_load(i['id']) + print('new: '+i['q']) + return new if __name__ == '__main__': import uvicorn uvicorn.run(app, host="127.0.0.1", port=8964) - # mysql_connector.end_mysql() \ No newline at end of file + mysql_connector.end_mysql() \ No newline at end of file diff --git a/mysql_connector.py b/mysql_connector.py index 63bd0d0..b1c999b 100644 --- a/mysql_connector.py +++ b/mysql_connector.py @@ -23,3 +23,13 @@ def end_mysql(): cnx.close() except: print("No Mysql Opened.") + +def is_loaded(paperId): + query = ("SELECT COUNT(*) FROM uploaded_papers WHERE paperId = %s") + cursor.execute(query, (paperId,)) + return cursor.fetchone()[0] > 0 + +def new_load(paperId): + query = ("INSERT INTO uploaded_papers (paperId) VALUES (%s)") + cursor.execute(query, (paperId,)) + cnx.commit()