diff --git a/dist/rbabel-1.0.0-py3-none-any.whl b/dist/rbabel-1.0.0-py3-none-any.whl index 5630668..edd99d4 100644 Binary files a/dist/rbabel-1.0.0-py3-none-any.whl and b/dist/rbabel-1.0.0-py3-none-any.whl differ diff --git a/dist/rbabel-1.0.0.tar.gz b/dist/rbabel-1.0.0.tar.gz index 4d17259..77489c0 100644 Binary files a/dist/rbabel-1.0.0.tar.gz and b/dist/rbabel-1.0.0.tar.gz differ diff --git a/src/rbabel/__init__.py b/src/rbabel/__init__.py index d3b7485..dcc6aaa 100644 --- a/src/rbabel/__init__.py +++ b/src/rbabel/__init__.py @@ -1,9 +1,9 @@ -import logging +import logging logging.basicConfig( level=logging.INFO, - format='%(asctime)s - %(levelname)s - %(message)s', - datefmt='%Y-%m-%d %H:%M:%S' + format="%(asctime)s - %(levelname)s - %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", ) log = logging.getLogger(__name__) diff --git a/src/rbabel/__main__.py b/src/rbabel/__main__.py index 23d38fd..df6d703 100644 --- a/src/rbabel/__main__.py +++ b/src/rbabel/__main__.py @@ -1,13 +1,14 @@ -from aiohttp import web +from aiohttp import web + from rbabel.app import create_app from rbabel.args import parse_args + def main(): args = parse_args() app = create_app(args.url) - web.run_app(app,host=args.host,port=args.port) + web.run_app(app, host=args.host, port=args.port) -if __name__ == '__main__': +if __name__ == "__main__": nain() - diff --git a/src/rbabel/agent.py b/src/rbabel/agent.py index 8d7453a..f8ac734 100644 --- a/src/rbabel/agent.py +++ b/src/rbabel/agent.py @@ -19,10 +19,12 @@ some modifications at the end of the file. """ try: - import sys import os + import sys + sys.path.append(os.getcwd()) import env + API_KEY = env.API_KEY ASSISTANT_ID = env.ASSISTANT_ID except: @@ -193,14 +195,15 @@ class Agent: print("Exiting..") break + async def main(): """ - Example main function. The keys here are not real but look exactly like - the real ones for example purposes and that you're sure your key is in the + Example main function. The keys here are not real but look exactly like + the real ones for example purposes and that you're sure your key is in the right format. """ agent = Agent(api_key=API_KEY, assistant_id=ASSISTANT_ID) - + # Run interactive chat await agent.cli() diff --git a/src/rbabel/app.py b/src/rbabel/app.py index 4cfd044..bea75e6 100644 --- a/src/rbabel/app.py +++ b/src/rbabel/app.py @@ -1,9 +1,10 @@ -from aiohttp import web -from yura.client import AsyncClient as LLMClient -from rbabel.agent import Agent -from rbabel import log -import sys import os +import sys + +from aiohttp import web + +from rbabel.agent import Agent + sys.path.append(os.getcwd()) import env @@ -12,30 +13,36 @@ LLM_NAME = "rbabel" LLM_CONTEXT = "You are an English grammar corrector. You repsond with only the corrected English of by user given prompt and nothing more. Also replace numbers with the word variant." LLM_MODEL = "gemma2:latest" + class Application(web.Application): - - def __init__(self, llm_url=None, llm_name=LLM_NAME, llm_model=LLM_MODEL,llm_context=LLM_CONTEXT, *args, **kwargs): + def __init__( + self, + llm_url=None, + llm_name=LLM_NAME, + llm_model=LLM_MODEL, + llm_context=LLM_CONTEXT, + *args, + **kwargs, + ): self.agent = Agent(env.API_KEY, env.ASSISTANT_ID) super().__init__(*args, **kwargs) self.router.add_post("/", self.optimize_grammar_handler) - async def fix_grammar(self, content): async for message in self.agent.chat(content): if message: - print("<<<<",message,"<<<<") + print("<<<<", message, "<<<<") return message - async def optimize_grammar_handler(self,request): + async def optimize_grammar_handler(self, request): text = await request.json() - assert(type(text) ==str) + assert type(text) == str corrected = await self.fix_grammar(text) - print("was:",text) - print("became:",corrected) - return web.json_response(corrected,content_type="application/json") + print("was:", text) + print("became:", corrected) + return web.json_response(corrected, content_type="application/json") def create_app(llm_url): return Application(llm_url, LLM_NAME, LLM_MODEL, LLM_CONTEXT) - diff --git a/src/rbabel/args.py b/src/rbabel/args.py index 6277288..5ac1fc5 100644 --- a/src/rbabel/args.py +++ b/src/rbabel/args.py @@ -1,4 +1,5 @@ -import argparse +import argparse + def parse_args(): parser = argparse.ArgumentParser( @@ -6,30 +7,19 @@ def parse_args(): ) parser.add_argument( - '--url', + "--url", type=str, required=False, default="http://127.0.0.1:3011/", - help='URL of Katya LLM Server.' + help="URL of Katya LLM Server.", ) parser.add_argument( - '--host', - type=str, - default="127.0.0.1", - required=False, - help='Host to bind to.' + "--host", type=str, default="127.0.0.1", required=False, help="Host to bind to." ) - parser.add_argument( - '--port', - type=int, - default=3011, - required=False, - help='Port to bind to.' + "--port", type=int, default=3011, required=False, help="Port to bind to." ) return parser.parse_args() - - diff --git a/src/rbabel/cli.py b/src/rbabel/cli.py index 865c918..bbec672 100644 --- a/src/rbabel/cli.py +++ b/src/rbabel/cli.py @@ -1,37 +1,40 @@ import asyncio +import time + from aiohttp import ClientSession + from rbabel.args import parse_args -import time async def cli_client(url): - + while True: sentence = input("> ") async with ClientSession() as session: - async with session.post("http://localhost:8080",json=sentence) as response: + async with session.post("http://localhost:8080", json=sentence) as response: try: print(await response.json()) except Exception as ex: print(ex) print(await response.text()) + async def bench(url): index = 0 while True: index += 1 - sentence = "I bougt {} woden shoe".format(index) + sentence = f"I bougt {index} woden shoe" time_start = time.time() async with ClientSession() as session: - async with session.post(url,json=sentence) as response: + async with session.post(url, json=sentence) as response: try: print(await response.json()) except Exception as ex: print(ex) print(await response.text()) time_end = time.time() - print("Duration: {}".format(time_end-time_start)) - + print(f"Duration: {time_end - time_start}") + def cli_bench(): args = parse_args() @@ -42,9 +45,6 @@ def main(): args = parse_args() asyncio.run(cli_client(args.url)) -if __name__ == '__main__': + +if __name__ == "__main__": main() - - - -