Update main.py

main.py

@@ -70,7 +70,7 @@ class APIGuardMiddleware(BaseHTTPMiddleware):
         # Get current client url and client IP address
         client_url = request.url.path
         client_ip_addr = request.client.host
-
+        
         # IP-based rate limitation
         async with r.pipeline(transaction=True) as pipeline:
             try:
@@ -80,6 +80,11 @@ class APIGuardMiddleware(BaseHTTPMiddleware):
                 res = None
     
             if res == None:
+                # lim = 25 if client_url.endswith('text2quiz-three.vercel.app') else 5
+                # lim = 10 if client_url.endswith('localhost:8100') else 5
+
+                app_logger.info(client_url)
+                
                 ok = await pipeline.set(client_ip_addr, 25).execute()
                 await r.expire(client_ip_addr, 60)
             elif res > 0:
@@ -600,6 +605,28 @@ async def convert2html(request):
 
     return JSONResponse(rets)
 
+async def llm_generate_text(request):
+    async with httpx.AsyncClient() as client:
+        o = await request.json()
+        _r = await client.post(
+                    url="https://api.awanllm.com/v1/chat/completions",
+                    headers={'Content-Type': 'application/json', 'Authorization': f'Bearer {LLM_API_KEY}'},
+                    data=json.dumps({
+                        "model": "Meta-Llama-3-8B-Instruct",
+                        "messages": [
+                            {"role": "user", "content": f"Generate a study note about {o['prompt']}. Use {o['lang']}."}
+                        ],
+                        "max_tokens": 4096,
+                        "presence_penalty":0.3,
+                        "temperature":0.5
+                    }),
+                    timeout=None)
+
+        try:
+            repl = _r.json()['choices'][0]['message']['content']
+            return PlainTextResponse(repl)
+        else:
+            raise HTTPException(status_code=422)
 
 async def get_flashcards(request):
     # [title, content, keywords]
@@ -612,7 +639,7 @@ async def get_flashcards(request):
 
     __tldr = await tldr(__content, __lang)
     print(__tldr)
-    __definitions = await get_definitions_from_words(__keywords, __tldr)
+    __definitions = await get_definitions_from_words(__keywords, __tldr, __lang)
     print(__definitions)
 
     return JSONResponse({"tldr": __tldr, "defs": __definitions, "imgs": await fetch_img_for_words(__keywords)})
@@ -685,7 +712,8 @@ app = Starlette(debug=True,routes=[
     Route('/convert2md', convert2md, methods=['POST']),
     Route('/mltest', __mltest, methods=['GET']),
     Route('/validateSimilarity', validate_similarity, methods=['POST']),
-    Route('/', root, methods=['GET'])
+    Route('/llmGenerateText', llm_generate_text, methods=['POST']),
+    Route('/', root, methods=['GET']),
 ],
 middleware=middleware)