chore: import upstream snapshot with attribution

examples/flows/standard/gen-docstring/README.md

@@ -0,0 +1,62 @@
+# Generate Python docstring
+This example can help you automatically generate Python code's docstring and return the modified code.
+
+Tools used in this flow：
+- `load_code` tool, it can load code from a file path.
+  - Load content from a local file.
+  - Loading content from a remote URL, currently loading HTML content, not just code.
+- `divide_code` tool, it can divide code into code blocks.
+  - To avoid files that are too long and exceed the token limit, it is necessary to split the file.
+  - Avoid using the same function (such as __init__(self)) to generate docstrings in the same one file, which may cause confusion when adding docstrings to the corresponding functions in the future.
+- `generate_docstring` tool, it can generate docstring for a code block, and merge docstring into origin code.
+
+## What you will learn
+
+In this flow, you will learn
+- How to compose an auto generate docstring flow.
+- How to use different LLM APIs to request LLM, including synchronous/asynchronous APIs, chat/completion APIs.
+- How to use asynchronous multiple coroutine approach to request LLM API.
+- How to construct a prompt.
+
+## Prerequisites
+
+### Install promptflow sdk and other dependencies:
+```bash
+pip install -r requirements.txt
+```
+
+### Create connection for LLM to use
+```bash
+# Override keys with --set to avoid yaml file changes
+pf connection create --file ../../../connections/azure_openai.yml --set api_key=<your_api_key> api_base=<your_api_base>
+```
+Note:  
+The [azure_openai.yml](../../../connections/azure_openai.yml) file is located in connections folder.  
+We are using connection named `open_ai_connection`in [flow.dag.yaml](flow.dag.yaml).
+
+## Execute with Promptflow
+### Execute with SDK
+`python main.py --source <your_file_path>`  
+**Note**: the file path should be a python file path, default is `./azure_open_ai.py`.
+
+A webpage will be generated, displaying diff:
+![result](result.png)
+
+
+### Execute with CLI
+```bash
+# run flow with default file path in flow.dag.yaml
+pf flow test --flow . 
+
+# run flow with file path
+pf flow test --flow . --inputs source="./azure_open_ai.py"
+```
+
+```bash
+# run flow with batch data
+pf run create --flow . --data ./data.jsonl --name auto_generate_docstring --column-mapping source='${data.source}'
+```
+Output the code after add the docstring.
+
+You can also skip providing `column-mapping` if provided data has same column name as the flow.
+Reference [here](https://aka.ms/pf/column-mapping) for default behavior when `column-mapping` not provided in CLI.

examples/flows/standard/gen-docstring/azure_open_ai.py

@@ -0,0 +1,290 @@
+import asyncio
+import logging
+import time
+import uuid
+from typing import List
+from openai.version import VERSION as OPENAI_VERSION
+import os
+from abc import ABC, abstractmethod
+import tiktoken
+from dotenv import load_dotenv
+from prompt import PromptLimitException
+
+
+class AOAI(ABC):
+    def __init__(self, **kwargs):
+        if OPENAI_VERSION.startswith("0."):
+            raise Exception(
+                "Please upgrade your OpenAI package to version >= 1.0.0 or "
+                "using the command: pip install --upgrade openai."
+            )
+        init_params = {}
+        api_type = os.environ.get("API_TYPE")
+        if os.getenv("OPENAI_API_VERSION") is not None:
+            init_params["api_version"] = os.environ.get("OPENAI_API_VERSION")
+        if os.getenv("OPENAI_ORG_ID") is not None:
+            init_params["organization"] = os.environ.get("OPENAI_ORG_ID")
+        if os.getenv("OPENAI_API_KEY") is None:
+            raise ValueError("OPENAI_API_KEY is not set in environment variables")
+        if os.getenv("OPENAI_API_BASE") is not None:
+            if api_type == "azure":
+                init_params["azure_endpoint"] = os.environ.get("OPENAI_API_BASE")
+            else:
+                init_params["base_url"] = os.environ.get("OPENAI_API_BASE")
+        init_params["api_key"] = os.environ.get("OPENAI_API_KEY")
+        # A few sanity checks
+        if api_type == "azure":
+            if init_params.get("azure_endpoint") is None:
+                raise ValueError(
+                    "OPENAI_API_BASE is not set in environment variables, this is required when api_type==azure"
+                )
+            if init_params.get("api_version") is None:
+                raise ValueError(
+                    "OPENAI_API_VERSION is not set in environment variables, this is required when api_type==azure"
+                )
+            if init_params["api_key"].startswith("sk-"):
+                raise ValueError(
+                    "OPENAI_API_KEY should not start with sk- when api_type==azure, "
+                    "are you using openai key by mistake?"
+                )
+            from openai import AzureOpenAI as Client
+            from openai import AsyncAzureOpenAI as AsyncClient
+        else:
+            from openai import OpenAI as Client
+            from openai import AsyncClient as AsyncClient
+        self.client = Client(**init_params)
+        self.async_client = AsyncClient(**init_params)
+        self.default_engine = None
+        self.engine = kwargs.pop('model', None) or os.environ.get("MODEL")
+        self.total_tokens = 4000
+        self.max_tokens = kwargs.pop('max_tokens', None) or os.environ.get("MAX_TOKENS") or 1200
+        if self.engine == "gpt-4-32k":
+            self.total_tokens = 31000
+        if self.engine == "gpt-4":
+            self.total_tokens = 7000
+        if self.engine == "gpt-3.5-turbo-16k":
+            self.total_tokens = 15000
+        if self.max_tokens > self.total_tokens:
+            raise ValueError(f"max_tokens must be less than total_tokens, "
+                             f"total_tokens is {self.total_tokens}, max_tokens is {self.max_tokens}")
+        self.tokens_limit = self.total_tokens - self.max_tokens
+
+    def count_tokens(self, text: str) -> int:
+        try:
+            encoding = tiktoken.encoding_for_model(self.engine)
+        except KeyError:
+            encoding = tiktoken.encoding_for_model(self.default_engine)
+        return len(encoding.encode(text))
+
+    def query(self, text, **kwargs):
+        stream = kwargs.pop("stream", False)
+        for i in range(3):
+            try:
+                if not stream:
+                    return self.query_with_no_stream(text, **kwargs)
+                else:
+                    return "".join(self.query_with_stream(text, **kwargs))
+            except Exception as e:
+                logging.error(f"Query failed, message={e}, "
+                              f"will retry request llm after {(i + 1) * (i + 1)} seconds.")
+                time.sleep((i + 1) * (i + 1))
+        raise Exception("Query failed, and retry 3 times, but still failed.")
+
+    async def async_query(self, text, **kwargs):
+        stream = kwargs.pop("stream", False)
+        for i in range(3):
+            try:
+                if not stream:
+                    res = await self.async_query_with_no_stream(text, **kwargs)
+                    return res
+                else:
+                    res = await self.async_query_with_stream(text, **kwargs)
+                    return "".join(res)
+            except Exception as e:
+                logging.error(f"llm response error, message={e}, "
+                              f"will retry request llm after {(i + 1) * (i + 1)} seconds.")
+                await asyncio.sleep((i + 1) * (i + 1))
+        raise Exception("llm response error, and retry 3 times, but still failed.")
+
+    @abstractmethod
+    def query_with_no_stream(self, text, **kwargs):
+        pass
+
+    @abstractmethod
+    def query_with_stream(self, text, **kwargs):
+        pass
+
+    @abstractmethod
+    async def async_query_with_no_stream(self, text, **kwargs):
+        pass
+
+    @abstractmethod
+    async def async_query_with_stream(self, text, **kwargs):
+        pass
+
+
+class ChatLLM(AOAI):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.default_engine = "gpt-3.5-turbo"
+        self.engine = self.engine or self.default_engine
+        self.system_prompt = "You are a Python engineer."
+        self.conversation = dict()
+
+    def query_with_no_stream(self, text, **kwargs):
+        conversation_id = kwargs.pop('conversation', None)
+        messages = self.create_prompt(text, conversation_id)
+        self.validate_tokens(messages)
+        temperature = kwargs.pop("temperature", 0.1)
+        response = self.client.chat.completions.create(
+            model=self.engine,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=self.max_tokens,
+            stream=False,
+            **kwargs,
+        )
+        response_role = response.choices[0].message.role
+        full_response = response.choices[0].message.content
+        self.add_to_conversation(text, "user", conversation_id=conversation_id)
+        self.add_to_conversation(full_response, response_role, conversation_id=conversation_id)
+        return full_response
+
+    def query_with_stream(self, text, **kwargs):
+        conversation_id = kwargs.pop('conversation', None)
+        messages = self.create_prompt(text, conversation_id)
+        self.validate_tokens(messages)
+        temperature = kwargs.pop("temperature", 0.1)
+        response = self.client.chat.completions.create(
+            model=self.engine,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=self.max_tokens,
+            stream=True,
+            **kwargs,
+        )
+
+        response_role = None
+        full_response = ""
+        for chunk in response:
+            delta = chunk.choices[0].delta
+            response_role = delta.role
+            if delta.content:
+                content = delta.content
+                full_response += content
+                yield content
+        self.add_to_conversation(text, "user", conversation_id=conversation_id)
+        self.add_to_conversation(full_response, response_role, conversation_id=conversation_id)
+
+    async def async_query_with_no_stream(self, text, **kwargs):
+        conversation_id = kwargs.pop('conversation', None)
+        messages = self.create_prompt(text, conversation_id)
+        self.validate_tokens(messages)
+        temperature = kwargs.pop("temperature", 0.1)
+        response = await self.async_client.chat.completions.create(
+            model=self.engine,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=self.max_tokens,
+            stream=False,
+            **kwargs,
+        )
+        response_role = response.choices[0].message.role
+        full_response = response.choices[0].message.content
+        self.add_to_conversation(text, "user", conversation_id=conversation_id)
+        self.add_to_conversation(full_response, response_role, conversation_id=conversation_id)
+        return full_response
+
+    async def async_query_with_stream(self, text, **kwargs):
+        conversation_id = kwargs.pop('conversation', None)
+        messages = self.create_prompt(text, conversation_id)
+        self.validate_tokens(messages)
+        temperature = kwargs.pop("temperature", 0.1)
+        response = await self.async_client.chat.completions.create(
+            model=self.engine,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=self.max_tokens,
+            stream=True,
+            **kwargs,
+        )
+
+        response_role = None
+        full_response = ""
+        for chunk in response:
+            delta = chunk.choices[0].delta
+            response_role = delta.role
+            if delta.content:
+                content = delta.content
+                full_response += content
+                yield content
+        self.add_to_conversation(text, "user", conversation_id=conversation_id)
+        self.add_to_conversation(full_response, response_role, conversation_id=conversation_id)
+
+    def get_unique_conversation_id(self):
+        return str(uuid.uuid4()).replace('-', '')
+
+    def add_to_conversation(self, message: str, role: str, conversation_id: str) -> None:
+        """
+        Add a message to the conversation
+        """
+        if type(conversation_id) is str:
+            self.conversation[conversation_id].append({"role": role, "content": message})
+
+    def del_conversation(self, conversation_id: str) -> None:
+        if conversation_id in self.conversation:
+            del self.conversation[conversation_id]
+
+    def init_conversation(self, conversation_id: str, system_prompt) -> None:
+        """
+        Init a new conversation
+        """
+        if type(conversation_id) is str:
+            self.conversation[conversation_id] = [{"role": "system", "content": system_prompt}]
+
+    def get_tokens_count(self, messages: List[dict]) -> int:
+        """
+        Get token count
+        """
+        num_tokens = 0
+        for message in messages:
+            # every message follows <im_start>{role/name}\n{content}<im_end>\n
+            num_tokens += 5
+            for key, value in message.items():
+                if value:
+                    num_tokens += self.count_tokens(value)
+                if key == "name":  # if there's a name, the role is omitted
+                    num_tokens += 5  # role is always required and always 1 token
+        num_tokens += 5  # every reply is primed with <im_start>assistant
+        return num_tokens
+
+    def validate_tokens(self, messages: List[dict]) -> None:
+        total_tokens = self.get_tokens_count(messages)
+        if total_tokens > self.tokens_limit:
+            message = f"token count {total_tokens} exceeds limit {self.tokens_limit}"
+            raise PromptLimitException(message)
+
+    def create_prompt(self, text: str, conversation_id: str = None):
+        unique_conversation_id = self.get_unique_conversation_id()
+        conversation_id = conversation_id or unique_conversation_id
+        if conversation_id not in self.conversation:
+            self.init_conversation(conversation_id=conversation_id, system_prompt=self.system_prompt)
+
+        _conversation = self.conversation[conversation_id] + [{"role": "user", "content": text}]
+
+        while self.get_tokens_count(_conversation) > self.tokens_limit and len(_conversation) > 2:
+            _conversation.pop(1)
+
+        if unique_conversation_id == conversation_id:
+            self.del_conversation(conversation_id=unique_conversation_id)
+
+        return _conversation
+
+
+if __name__ == "__main__":
+    load_dotenv()
+    llm = ChatLLM()
+    print(llm.query(text='how are you?'))
+    res = llm.query_with_stream(text='how are you?')
+    for item in res:
+        print(item)

examples/flows/standard/gen-docstring/combine_code.jinja2

@@ -0,0 +1 @@
+{{divided|join('')}}

examples/flows/standard/gen-docstring/combine_code_tool.py

@@ -0,0 +1,9 @@
+from promptflow.core import tool
+from divider import Divider
+from typing import List
+
+
+@tool
+def combine_code(divided: List[str]):
+    code = Divider.combine(divided)
+    return code

examples/flows/standard/gen-docstring/data.jsonl

@@ -0,0 +1,3 @@
+{"source": "./divider.py"}
+{"source":  "./azure_open_ai.py"}
+{"source":  "./generate_docstring_tool.py"}

examples/flows/standard/gen-docstring/diff.py

@@ -0,0 +1,18 @@
+import difflib
+import webbrowser
+
+
+def show_diff(left_content, right_content, name="file"):
+    d = difflib.HtmlDiff()
+    html = d.make_file(
+        left_content.splitlines(),
+        right_content.splitlines(),
+        "origin " + name,
+        "new " + name,
+        context=True,
+        numlines=20)
+    html = html.encode()
+    html_name = name + "_diff.html"
+    with open(html_name, "w+b") as fp:
+        fp.write(html)
+    webbrowser.open(html_name)

examples/flows/standard/gen-docstring/divide_code_tool.py

@@ -0,0 +1,9 @@
+from promptflow.core import tool
+from divider import Divider
+
+
+@tool
+def divide_code(file_content: str):
+    # Divide the code into several parts according to the global import/class/function.
+    divided = Divider.divide_file(file_content)
+    return divided

examples/flows/standard/gen-docstring/divider.py

@@ -0,0 +1,112 @@
+import logging
+import re
+from typing import List
+
+
+class Settings:
+    divide_file = {
+        "py": r"(?<!.)(class|def)",
+    }
+    divide_func = {
+        "py": r"((\n {,6})|^)(class|def)\s+(\S+(?=\())\s*(\([^)]*\))?\s*(->[^:]*:|:) *"
+    }
+
+
+class Divider:
+    language = 'py'
+
+    @classmethod
+    def divide_file(cls, text) -> List[str]:
+        matches = list(re.finditer(Settings.divide_file[Divider.language], text))
+        splitted_content = []
+        min_pos = matches[0].start() if len(matches) > 0 else len(text)
+        for i in range(len(matches)):
+            start = matches[i].start()
+            end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
+            splitted_content.append(text[start:end])
+        if min_pos != 0:
+            splitted_content.insert(0, text[0:min_pos])
+        return splitted_content
+
+    @classmethod
+    def divide_half(cls, text) -> List[str]:
+        """
+        Divide the content into two parts, but ensure that the function body is not split.
+        """
+        _, pos = Divider.get_functions_and_pos(text)
+        if len(pos) > 1:  # Divide the code into two parts and every part start with a function.
+            i = len(pos) // 2
+            return [text[0:pos[i][0]], text[pos[i][0]:]]
+        if len(pos) == 1:  # Divide the code into two parts, [function define + body, other body].
+            body = text[pos[0][1]:]
+            body_lines = body.split('\n')
+            body_ten_lines = '\n'.join(body_lines[0:10])
+            return [text[0:pos[0][1]] + body_ten_lines, body[len(body_ten_lines):]]
+        return [text]
+
+    @classmethod
+    def get_functions_and_pos(cls, text):
+        matches = re.finditer(Settings.divide_func[Divider.language], text)
+        functions = []
+        pos = []
+        for match in matches:
+            matched_text = match.group().replace('\n', '')
+            func = re.sub(r' +', ' ', matched_text).replace(' :', ':')
+            func = re.sub(r'[\s,]+\)', ')', func)
+            func = re.sub(r'\([\s,]+', '(', func)
+            functions.append(func.strip())
+            pos.append((match.start(), match.end()))
+        return functions, pos
+
+    @classmethod
+    def combine(cls, divided: List[str]):
+        return ''.join(divided)
+
+    @classmethod
+    def merge_doc2code(cls, docstring: str, origin_code: str) -> str:
+        funcs1, pos1 = Divider.get_functions_and_pos(docstring)
+        funcs2, pos2 = Divider.get_functions_and_pos(origin_code)
+        pattern = r'""".*?"""'
+        code = origin_code if len(funcs2) == 0 else origin_code[0:pos2[0][0]]
+        pos1.append((len(docstring), len(docstring)))  # avoid index out of range
+        pos2.append((len(origin_code), len(origin_code)))  # avoid index out of range
+        for i2 in range(len(funcs2)):  # add docstring for each function in origin_code
+            part_full_code = origin_code[pos2[i2][0]:pos2[i2 + 1][0]]
+            try:
+                i1 = funcs1.index(funcs2[i2])
+            except ValueError:
+                logging.warning(f"No docstring found for {funcs2[i2]}")
+                code += part_full_code
+                continue
+            new_doc = re.findall(pattern, docstring[pos1[i1][1]:pos1[i1 + 1][0]], re.DOTALL)
+            if new_doc:
+                func_line = origin_code[pos2[i2][0]:pos2[i2][1]].replace('\n', '')
+                empty_line_num = (len(func_line) - len(func_line.lstrip()) + 4)
+                func_body = origin_code[pos2[i2][1]:pos2[i2 + 1][0]]
+                code_doc = list(re.finditer(pattern, func_body, re.DOTALL))
+                format_new_doc = Divider.format_indentation(new_doc[0], empty_line_num)
+                is_replace_doc = len(code_doc) > 0 and (re.sub(r'\s+', '', func_body[0:code_doc[0].start()]) == '')
+                if is_replace_doc:
+                    code += part_full_code.replace(code_doc[0].group(), format_new_doc.strip(), 1)
+                else:
+                    code += origin_code[pos2[i2][0]:pos2[i2][1]] + '\n' + format_new_doc + '\n' + origin_code[
+                                                                                                  pos2[i2][1]:
+                                                                                                  pos2[i2 + 1][0]]
+            else:
+                code += part_full_code
+        return code
+
+    @classmethod
+    def format_indentation(cls, text, empty_line_num):
+        lines = text.splitlines()
+        last_line_space_num = len(lines[-1]) - len(lines[-1].lstrip())
+        need_add_space = max(empty_line_num - last_line_space_num, 0) * ' '
+        lines[0] = last_line_space_num * ' ' + lines[0].lstrip()  # Align the first row to the last row
+        indented_lines = [(need_add_space + line).rstrip() for line in lines]
+        indented_string = '\n'.join(indented_lines)
+        return indented_string
+
+    @classmethod
+    def has_class_or_func(cls, text):
+        funcs, _ = Divider.get_functions_and_pos(text)
+        return len(funcs) > 0

examples/flows/standard/gen-docstring/doc_format.jinja2

@@ -0,0 +1,48 @@
+This is the docstring style of sphinx:
+"""Description of the function.
+
+:param [ParamName]: [ParamDescription](, defaults to [DefaultParamVal].)
+:type [ParamName]: [ParamType](, optional)
+...
+:raises [ErrorType]: [ErrorDescription]
+...
+:return: [ReturnDescription]
+:rtype: [ReturnType]
+"""
+
+Note:
+For custom class types, please use the full path, for example:
+"~azure.ai.ml.entities._inputs_outputs.Input" is full path for "Input" because of "from azure.ai.ml.entities._inputs_outputs import Input, Output"
+"~import_node.Import" is full path for "Import" because of "import import_node.Import"
+
+Complete function docstring example:
+from azure.ai.ml.entities._inputs_outputs import Input, Output
+from azure.ai.ml.constants import JobType
+def output(input: Input, import_node: Import, startHnd=1, endHnd=None, uuids=None) -> Output:
+    """Create an Output object.
+
+    :param input: The input object.
+    :type input: ~azure.ai.ml.entities._inputs_outputs.Input
+    :param import_node: The Import object.
+    :type import_node: ~import_node.Import
+    :param startHnd: Start index, defaults to 1
+    :type startHnd: int, optional
+    :param endHnd: End index, defaults to None
+    :type endHnd: int, optional
+
+    :return: The Output object.
+    :rtype: ~azure.ai.ml.entities._inputs_outputs.Output
+    """
+    pass
+
+Here's some code for you:
+{{module}}
+{{code}}
+
+Please follow the sphinx style and refer above complete function docstring example, then output the docstring for the following class/functions.
+Please replace "{docstring}" with the actual docstring.
+{% for func in functions %}
+{{func}}
+{docstring}
+pass
+{% endfor %}

examples/flows/standard/gen-docstring/file.py

@@ -0,0 +1,83 @@
+import logging
+import os
+from urllib.parse import urlparse
+import requests
+
+
+class File:
+    def __init__(self, source: str):
+        self._source = source
+        self._is_url = source.startswith("http://") or source.startswith("https://")
+        if self._is_url:
+            parsed_url = urlparse(source)
+            path = parsed_url.path
+        else:
+            path = source
+        self._path = os.path.normpath(os.path.abspath(path))
+        self._dirname = os.path.dirname(self._path)
+        self._filename = os.path.basename(self._path).split(".")[0]
+        self._language = os.path.basename(self._path).split(".")[1]
+
+    def _read_content(self):
+        if self._is_url:
+            response = requests.get(self.source)
+            if response.status_code == 200:
+                content = response.text
+                return content
+            else:
+                print(f"Failed to retrieve content from URL: {self.source}")
+                return None
+        else:
+            try:
+                with open(self._path, "r") as file:
+                    content = file.read()
+                    return content
+            except FileNotFoundError:
+                print(f"File not found: {self.source}")
+                return None
+
+    @property
+    def content(self) -> str:
+        if not hasattr(self, "_text"):
+            self._content = self._read_content()
+        return self._content
+
+    @property
+    def language(self) -> str:
+        return self._language
+
+    @property
+    def filename(self) -> str:
+        return self._filename
+
+    @property
+    def dirname(self) -> str:
+        return self._dirname
+
+    @property
+    def source(self) -> str:
+        return self._source
+
+    def override_origin_file(self, content: str) -> None:
+        if not self._is_url:
+            with open(self._path, "w") as f:
+                f.write(content)
+        else:
+            logging.warning(
+                "Cannot override origin file from URL, create a new file instead."
+            )
+            self.create_new_file(content)
+
+    def create_new_file(self, content: str) -> None:
+        if self._is_url:
+            path = os.path.join(
+                "./",
+                self.filename + f"_doc.{self.language}",
+            )
+        else:
+            path = os.path.join(
+                self.dirname,
+                self.filename + f"_doc.{self.language}",
+            )
+        with open(path, "w") as f:
+            f.write(content)

examples/flows/standard/gen-docstring/flow.dag.yaml

@@ -0,0 +1,42 @@
+$schema: https://azuremlschemas.azureedge.net/promptflow/latest/Flow.schema.json
+inputs:
+  source:
+    type: string
+    default: ./azure_open_ai.py
+outputs:
+  code:
+    type: string
+    reference: ${combine_code.output}
+nodes:
+- name: load_code
+  type: python
+  source:
+    type: code
+    path: load_code_tool.py
+  inputs:
+    source: ${inputs.source}
+- name: divide_code
+  type: python
+  source:
+    type: code
+    path: divide_code_tool.py
+  inputs:
+    file_content: ${load_code.output}
+- name: generate_docstring
+  type: python
+  source:
+    type: code
+    path: generate_docstring_tool.py
+  inputs:
+    divided: ${divide_code.output}
+    connection: open_ai_connection
+    model: gpt-35-turbo
+- name: combine_code
+  type: prompt
+  source:
+    type: code
+    path: combine_code.jinja2
+  inputs:
+    divided: ${generate_docstring.output}
+environment:
+  python_requirements_txt: requirements.txt

examples/flows/standard/gen-docstring/generate_docstring_tool.py

@@ -0,0 +1,95 @@
+import ast
+import asyncio
+import logging
+import os
+import sys
+from typing import Union, List
+from promptflow.core import tool
+from azure_open_ai import ChatLLM
+from divider import Divider
+from prompt import docstring_prompt, PromptLimitException
+from promptflow.connections import AzureOpenAIConnection, OpenAIConnection
+
+
+def get_imports(content):
+    tree = ast.parse(content)
+    import_statements = []
+
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Import):
+            for n in node.names:
+                import_statements.append(f"import {n.name}")
+        elif isinstance(node, ast.ImportFrom):
+            module_name = node.module
+            for n in node.names:
+                import_statements.append(f"from {module_name} import {n.name}")
+
+    return import_statements
+
+
+async def async_generate_docstring(divided: List[str]):
+    llm = ChatLLM()
+    divided = list(reversed(divided))
+    all_divided = []
+
+    # If too many imports result in tokens exceeding the limit, please set an empty string.
+    modules = ''  # '\n'.join(get_imports(divided[-1]))
+    modules_tokens = llm.count_tokens(modules)
+    if modules_tokens > 300:
+        logging.warning(f'Too many imports, the number of tokens is {modules_tokens}')
+    if modules_tokens > 500:
+        logging.warning(f'Too many imports, the number of tokens is {modules_tokens}, will set an empty string.')
+        modules = ''
+
+    # Divide the code into two parts if the global class/function is too long.
+    while len(divided):
+        item = divided.pop()
+        try:
+            llm.validate_tokens(llm.create_prompt(docstring_prompt(code=item, module=modules)))
+        except PromptLimitException as e:
+            logging.warning(e.message + ', will divide the code into two parts.')
+            divided_tmp = Divider.divide_half(item)
+            if len(divided_tmp) > 1:
+                divided.extend(list(reversed(divided_tmp)))
+                continue
+        except Exception as e:
+            logging.warning(e)
+        all_divided.append(item)
+
+    tasks = []
+    last_code = ''
+    for item in all_divided:
+        if Divider.has_class_or_func(item):
+            tasks.append(llm.async_query(docstring_prompt(last_code=last_code, code=item, module=modules)))
+        else:  # If the code has not function or class, no need to generate docstring.
+            tasks.append(asyncio.sleep(0))
+        last_code = item
+    res_doc = await asyncio.gather(*tasks)
+    new_code = []
+    for i in range(len(all_divided)):
+        if type(res_doc[i]) is str:
+            new_code.append(Divider.merge_doc2code(res_doc[i], all_divided[i]))
+        else:
+            new_code.append(all_divided[i])
+
+    return new_code
+
+
+@tool
+def generate_docstring(divided: List[str],
+                       connection: Union[AzureOpenAIConnection, OpenAIConnection] = None,
+                       model: str = None):
+    if isinstance(connection, AzureOpenAIConnection):
+        os.environ["OPENAI_API_KEY"] = connection.api_key
+        os.environ["OPENAI_API_BASE"] = connection.api_base
+        os.environ["OPENAI_API_VERSION"] = connection.api_version
+        os.environ["API_TYPE"] = connection.api_type
+    elif isinstance(connection, OpenAIConnection):
+        os.environ["OPENAI_API_KEY"] = connection.api_key
+        os.environ["ORGANIZATION"] = connection.organization
+    if model:
+        os.environ["MODEL"] = model
+
+    if sys.platform.startswith("win"):
+        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+    return asyncio.run(async_generate_docstring(divided))

examples/flows/standard/gen-docstring/load_code_tool.py

@@ -0,0 +1,8 @@
+from promptflow.core import tool
+from file import File
+
+
+@tool
+def load_code(source: str):
+    file = File(source)
+    return file.content

examples/flows/standard/gen-docstring/main.py

@@ -0,0 +1,18 @@
+import argparse
+from file import File
+from diff import show_diff
+from load_code_tool import load_code
+from promptflow.client import PFClient
+from pathlib import Path
+
+
+if __name__ == "__main__":
+    current_folder = Path(__file__).absolute().parent
+    parser = argparse.ArgumentParser(description="The code path of code that need to generate docstring.")
+    parser.add_argument("--source", help="Path for the code file", default=str(current_folder / 'azure_open_ai.py'))
+    args = parser.parse_args()
+
+    pf = PFClient()
+    source = args.source
+    flow_result = pf.test(flow=str(current_folder), inputs={"source": source})
+    show_diff(load_code(source), flow_result['code'], File(source).filename)

examples/flows/standard/gen-docstring/prompt.py

@@ -0,0 +1,30 @@
+import sys
+from promptflow.tools.common import render_jinja_template
+from divider import Divider
+
+
+class PromptLimitException(Exception):
+    def __init__(self, message="", **kwargs):
+        super().__init__(message, **kwargs)
+        self._message = str(message)
+        self._kwargs = kwargs
+        self._inner_exception = kwargs.get("error")
+        self.exc_type, self.exc_value, self.exc_traceback = sys.exc_info()
+        self.exc_type = self.exc_type.__name__ if self.exc_type else type(self._inner_exception)
+        self.exc_msg = "{}, {}: {}".format(message, self.exc_type, self.exc_value)
+
+    @property
+    def message(self):
+        if self._message:
+            return self._message
+
+        return self.__class__.__name__
+
+
+def docstring_prompt(last_code: str = '', code: str = '', module: str = '') -> str:
+    functions, _ = Divider.get_functions_and_pos(code)
+    # Add the first few lines to the function, such as decorator, to make the docstring generated better by llm.
+    first_three_lines = '\n'.join(last_code.split('\n')[-3:])
+    with open('doc_format.jinja2') as file:
+        return render_jinja_template(prompt=file.read(), module=module.strip('\n'),
+                                     code=(first_three_lines + code).strip('\n'), functions=functions)

examples/flows/standard/gen-docstring/requirements.txt

@@ -0,0 +1,5 @@
+promptflow[azure]
+promptflow-tools
+python-dotenv
+jinja2
+tiktoken