diff --git a/main.py b/main.py index 873d0b4..051bf04 100644 --- a/main.py +++ b/main.py @@ -1,13 +1,11 @@ import os -from docarray import DocumentArray, Document -from jina import Client - from src import gpt, jina_cloud -from src.constants import TAG_TO_FILE_NAME, EXECUTOR_FOLDER -from src.prompt_examples import executor_example, docarray_example +from src.constants import TAG_TO_FILE_NAME, EXECUTOR_FOLDER, CLIENT_FILE_NAME +from src.jina_cloud import run_client_file +from src.prompt_examples import executor_example, docarray_example, client_example from src.prompt_tasks import general_guidelines, executor_file_task, requirements_file_task, \ - test_executor_file_task, docker_file_task + test_executor_file_task, docker_file_task, client_file_task from src.utils.io import recreate_folder from src.utils.string import find_between, clean_content @@ -37,27 +35,50 @@ metas: with open('executor/config.yml', 'w') as f: f.write(config_content) +def get_all_executor_files_with_content(): + folder_path = 'executor' + file_name_to_content = {} + for filename in os.listdir(folder_path): + file_path = os.path.join(folder_path, filename) -def main(executor_name, input_executor_description, input_test_description): + if os.path.isfile(file_path): + with open(file_path, 'r', encoding='utf-8') as file: + content = file.read() + file_name_to_content[filename] = content + + return file_name_to_content + +def main( + executor_name, + input_executor_description, + input_modality, + input_doc_field, + output_modality, + output_doc_field, + input_test_in, + input_test_out +): recreate_folder(EXECUTOR_FOLDER) system_definition = ( "You are a principal engineer working at Jina - an open source company." "Using the Jina framework, users can define executors." + executor_example + docarray_example + + client_example ) user_query = ( - input_executor_description - + general_guidelines - + executor_file_task() - + test_executor_file_task(executor_name) + general_guidelines() + + executor_file_task(executor_name, input_executor_description, input_modality, input_doc_field, + output_modality, output_doc_field) + + test_executor_file_task(executor_name, input_test_in, input_test_out) + requirements_file_task() + docker_file_task() - + input_test_description + + client_file_task() ) plain_text = gpt.get_response(system_definition, user_query) + extract_and_write(plain_text) write_config_yml(executor_name) @@ -66,26 +87,19 @@ def main(executor_name, input_executor_description, input_test_description): host = jina_cloud.deploy_flow(executor_name) - client = Client(host=host) + run_client_file(f'executor/{CLIENT_FILE_NAME}', host) - d = Document(uri='data/txt.png') - d.load_uri_to_blob() - response = client.post('/index', inputs=DocumentArray([d])) - response[0].summary() + return get_all_executor_files_with_content() if __name__ == '__main__': main( - executor_name='MyBelovedOcrExecutor', - input_executor_description=( - "Write an executor that takes image bytes as input (document.blob within a DocumentArray) " - # "and use BytesIO to convert it to PIL " \ - "and detects ocr " - "and returns the texts as output (as DocumentArray). " - ), - - input_test_description='The test downloads the image ' \ - 'https://miro.medium.com/v2/resize:fit:1024/0*4ty0Adbdg4dsVBo3.png ' \ - ' loads it as bytes, takes it as input to the executor and asserts that the output is "> Hello, world!_".', - + executor_name='MyCoolOcrExecutor', + input_executor_description="OCR detector", + input_modality='image', + input_doc_field='uri', + output_modality='text', + output_doc_field='text', + input_test_in='https://miro.medium.com/v2/resize:fit:1024/0*4ty0Adbdg4dsVBo3.png', + input_test_out='> Hello, world!_', ) diff --git a/server.py b/server.py new file mode 100644 index 0000000..f58c933 --- /dev/null +++ b/server.py @@ -0,0 +1,41 @@ +from fastapi import FastAPI +from pydantic import BaseModel, HttpUrl +from typing import Optional, Dict + +from main import main + +app = FastAPI() + +# Define the request model +class CreateRequest(BaseModel): + executor_name: str + input_executor_description: str + input_modality: str + input_doc_field: str + output_modality: str + output_doc_field: str + input_test_in: HttpUrl + input_test_out: str + +# Define the response model +class CreateResponse(BaseModel): + result: Dict[str, str] + success: bool + message: Optional[str] + +@app.post("/create", response_model=CreateResponse) +async def create_endpoint(request: CreateRequest): + try: + result = main( + executor_name=request.executor_name, + input_executor_description=request.input_executor_description, + input_modality=request.input_modality, + input_doc_field=request.input_doc_field, + output_modality=request.output_modality, + output_doc_field=request.output_doc_field, + input_test_in=request.input_test_in, + input_test_out=request.input_test_out, + ) + return CreateResponse(result=result, success=True, message=None) + except Exception as e: + return CreateResponse(result=None, success=False, message=str(e)) \ No newline at end of file diff --git a/src/constants.py b/src/constants.py index d977ff4..f8f8c2c 100644 --- a/src/constants.py +++ b/src/constants.py @@ -2,17 +2,21 @@ EXECUTOR_FILE_NAME = 'executor.py' TEST_EXECUTOR_FILE_NAME = 'test_executor.py' REQUIREMENTS_FILE_NAME = 'requirements.txt' DOCKER_FILE_NAME = 'Dockerfile' +CLIENT_FILE_NAME = 'client.py' EXECUTOR_FILE_TAG = 'executor' TEST_EXECUTOR_FILE_TAG = 'test_executor' REQUIREMENTS_FILE_TAG = 'requirements' DOCKER_FILE_TAG = 'dockerfile' +CLIENT_FILE_TAG = 'client' TAG_TO_FILE_NAME = { EXECUTOR_FILE_TAG: EXECUTOR_FILE_NAME, TEST_EXECUTOR_FILE_TAG: TEST_EXECUTOR_FILE_NAME, REQUIREMENTS_FILE_TAG: REQUIREMENTS_FILE_NAME, - DOCKER_FILE_TAG: DOCKER_FILE_NAME + DOCKER_FILE_TAG: DOCKER_FILE_NAME, + CLIENT_FILE_TAG: CLIENT_FILE_NAME } EXECUTOR_FOLDER = 'executor' +FLOW_URL_PLACEHOLDER = 'jcloud.jina.ai' \ No newline at end of file diff --git a/src/gpt.py b/src/gpt.py index 9486986..111b6bf 100644 --- a/src/gpt.py +++ b/src/gpt.py @@ -2,9 +2,13 @@ import os import openai +from src.utils.string import print_colored + openai.api_key = os.environ['OPENAI_API_KEY'] def get_response(system_definition, user_query): + print_colored('system_definition', system_definition, 'magenta') + print_colored('user_query', user_query, 'blue') response = openai.ChatCompletion.create( temperature=0, model="gpt-4", @@ -23,5 +27,5 @@ def get_response(system_definition, user_query): ] ) content = response['choices'][0]['message']['content'] - print(content) + print_colored('agent response', content, 'green') return content \ No newline at end of file diff --git a/src/jina_cloud.py b/src/jina_cloud.py index c429356..4eafc8e 100644 --- a/src/jina_cloud.py +++ b/src/jina_cloud.py @@ -1,9 +1,12 @@ import os +from multiprocessing.connection import Client import hubble from jcloud.flow import CloudFlow from jina import Flow +from src.constants import FLOW_URL_PLACEHOLDER + def push_executor(): cmd = 'jina hub push executor/. --verbose' @@ -46,4 +49,24 @@ executors: with flow: pass - return CloudFlow(path=full_flow_path).__enter__().endpoints['gateway'] \ No newline at end of file + return CloudFlow(path=full_flow_path).__enter__().endpoints['gateway'] + +def replace_client_line(file_content: str, replacement: str) -> str: + lines = file_content.split('\n') + for index, line in enumerate(lines): + if 'Client(' in line: + lines[index] = replacement + break + return '\n'.join(lines) + +def run_client_file(file_path, host): + with open(file_path, 'r') as file: + content = file.read() + + replaced_content = replace_client_line(content, f"client = Client(host='{host}')") + + + with open(file_path, 'w') as file: + file.write(replaced_content) + + import executor.client # runs the client script for validation diff --git a/src/prompt_examples.py b/src/prompt_examples.py index 81ee1f4..b596ab8 100644 --- a/src/prompt_examples.py +++ b/src/prompt_examples.py @@ -1,3 +1,5 @@ +from src.constants import FLOW_URL_PLACEHOLDER + executor_example = "Here is an example of how an executor can be defined. It always starts with a comment:" ''' @@ -14,7 +16,7 @@ class MyExecutor(Executor): d.text = 'hello world'" return docs ''' -"An executor gets a DocumentArray as input and returns a DocumentArray as output." +"An executor gets a DocumentArray as input and returns a DocumentArray as output. " docarray_example = ( "A DocumentArray is a python class that can be seen as a list of Documents. " @@ -60,7 +62,7 @@ message DocumentProto { from jina import DocumentArray, Document d1 = Document(text='hello') -d2 = Document(blob=b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x03L\x00\x00\x01\x18\x08\x06\x00\x00\x00o...') +d2 = Document(blob=b'\\x89PNG\\r\\n\\x1a\\n\\x00\\x00\\x00\\rIHDR\\x00\\x00\\x03L\\x00\\x00\\x01\\x18\\x08\\x06\\x00\\x00\\x00o...') d3 = Document(tensor=numpy.array([1, 2, 3]), chunks=[Document(uri=/local/path/to/file)] d4 = Document( uri='https://docs.docarray.org', @@ -69,9 +71,24 @@ d4 = Document( d5 = Document() d5.tensor = np.ones((2,4)) d6 = Document() -d6.blob = b'RIFF\x00\x00\x00\x00WAVEfmt \x10\x00...' +d6.blob = b'RIFF\\x00\\x00\\x00\\x00WAVEfmt \\x10\\x00...' docs = DocumentArray([ d1, d2, d3, d4 ]) +# the document has a helper function load_uri_to_blob: +# For instance, d4.load_uri_to_blob() downloads the file from d4.uri and stores it in d4.blob. +# If d4.uri was something like 'https://website.web/img.jpg', then d4.blob would be something like b'\xff\xd8\xff\xe0\x00\x10JFIF\x00\x01\x01... ''' -) \ No newline at end of file +) + +client_example = ( +"After the executor is deployed, it can be called via Jina Client. " +"Here is an example of a client file: " +f''' +from jina import Client, Document, DocumentArray +client = Client(host='{FLOW_URL_PLACEHOLDER}') +d = Document(uri='data/img.png') +d.load_uri_to_blob() +response = client.post('/process', inputs=DocumentArray([d])) +response[0].summary() +''') \ No newline at end of file diff --git a/src/prompt_tasks.py b/src/prompt_tasks.py index 1a69b19..0147854 100644 --- a/src/prompt_tasks.py +++ b/src/prompt_tasks.py @@ -1,36 +1,51 @@ from src.constants import EXECUTOR_FILE_NAME, REQUIREMENTS_FILE_NAME, TEST_EXECUTOR_FILE_NAME, DOCKER_FILE_NAME, \ - DOCKER_FILE_TAG + DOCKER_FILE_TAG, CLIENT_FILE_TAG, CLIENT_FILE_NAME -general_guidelines = ( - "The code you write is production ready. " - "Every file starts with comments describing what the code is doing before the first import. " - "Comments can only be written between tags. " - "Start from top-level and then fully implement all methods." -) + +def general_guidelines(): + return ( + "General guidelines: " + "The code you write is production ready. " + "Every file starts with comments describing what the code is doing before the first import. " + "Then all imports are listed. It is important to import all modules that could be needed in the executor code." + "Comments can only be written between tags. " + "Start from top-level and then fully implement all methods. " + "\n" + ) def _task(task, tag_name, file_name): - return task + f"The code will go into {file_name}. Wrap the code in the string $$$start_{tag_name}$$$...$$$end_{tag_name}$$$ " + return task + f"The code will go into {file_name}. Wrap the code in the string $$$start_{tag_name}$$$...$$$end_{tag_name}$$$ \n\n" - -def executor_file_task(): - return _task("Write the executor code. ", 'executor', EXECUTOR_FILE_NAME) +def executor_file_task(executor_name, input_executor_description, input_modality, input_doc_field, + output_modality, output_doc_field): + return _task( + f"Write the executor called '{executor_name}'. " + f"It matches the following description: '{input_executor_description}'. " + f"It gets a DocumentArray as input where each document has the input modality '{input_modality}' that is stored in document.{input_doc_field}. " + f"It returns a DocumentArray as output where each document has the output modality '{output_modality}' that is stored in document.{output_doc_field}. ", + 'executor', + EXECUTOR_FILE_NAME + ) def requirements_file_task(): return _task("Write the content of the requirements.txt file. " "Make sure to include pytest. " - "All versions are fixed. " , 'requirements', + "All versions are fixed. ", 'requirements', REQUIREMENTS_FILE_NAME) -def test_executor_file_task(executor_name): +def test_executor_file_task(executor_name, input_test_in, input_test_out): return _task( "Write a small unit test for the executor. " "Start the test with an extensive comment about the test case. " - "Use the following import to import the executor: " - f"from executor import {executor_name}", + + ( + "Test that the executor converts the input '" + input_test_in + "' to the output '" + input_test_out + "'. " + ) if input_test_in and input_test_out else "" + "Use the following import to import the executor: " + f"from executor import {executor_name} ", 'test_executor', TEST_EXECUTOR_FILE_NAME ) @@ -48,3 +63,9 @@ def docker_file_task(): "The Dockerfile runs the test during the build process. " , DOCKER_FILE_TAG, DOCKER_FILE_NAME) + +def client_file_task(): + return _task( + "Write the client file. " + , CLIENT_FILE_TAG, CLIENT_FILE_NAME + ) diff --git a/src/utils/string.py b/src/utils/string.py index b02a299..269bc1e 100644 --- a/src/utils/string.py +++ b/src/utils/string.py @@ -9,3 +9,26 @@ def find_between(input_string, start, end): def clean_content(content): return content.replace('```', '').strip() + +def print_colored(headline, text, color_code): + if color_code == 'black': + color_code = '30' + elif color_code == 'red': + color_code = '31' + elif color_code == 'green': + color_code = '32' + elif color_code == 'yellow': + color_code = '33' + elif color_code == 'blue': + color_code = '34' + elif color_code == 'magenta': + color_code = '35' + elif color_code == 'cyan': + color_code = '36' + elif color_code == 'white': + color_code = '37' + color_start = f"\033[{color_code}m" + reset = "\033[0m" + bold_start = "\033[1m" + print(f"{bold_start}{color_start}{headline}{reset}") + print(f"{color_start}{text}{reset}")