langchain-ai/kork
1
0
Fork 0
mirror of https://github.com/langchain-ai/kork.git synced 2026-06-30 22:28:05 -04:00
Code Issues 6 Packages Projects Releases 2 Wiki Activity

Initial commit

Browse Source
This commit is contained in:
Eugene Yurtsev
2023-05-04 14:43:52 -04:00
commit 6e84264e0f
48 changed files with 10363 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/test.yml
+33
View File
@@ -0,0 +1,33 @@
name: test
on:
push:
branches: [main]
pull_request:
env:
POETRY_VERSION: "1.4.2"
jobs:
build:
runs-on: ubuntu-latest
strategy:
matrix:
python-version:
- "3.8"
- "3.9"
- "3.10"
- "3.11"
steps:
- uses: actions/checkout@v3
- name: Install poetry
run: pipx install poetry==$POETRY_VERSION
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v4
with:
python-version: ${{ matrix.python-version }}
cache: "poetry"
- name: Install dependencies
run: poetry install --with test
- name: Run unit tests
run: |
poetry run poe test
.gitignore
+129
View File
@@ -0,0 +1,129 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
pip-wheel-metadata/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
.python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
CONTRIBUTING.md
+39
View File
@@ -0,0 +1,39 @@
# Contributing to Kork
Thanks for your interest in contributing to Kork!
If you have ideas or features you would like to see implemented feel free to
open an issue and let me know.
PRs are welcome, but before starting to work on a substantial PR, please file
an issue to discuss the design and the code change.
## Setting up for development
The package uses [poetry](https://python-poetry.org/) together with
[poethepoet](https://github.com/nat-n/poethepoet).
### Install dependencies
```shell
poetry install --with dev,test
```
### List tasks
```shell
poe
```
### autoformat
```shell
poe autoformat
```
### test
```shell
poe test
```
LICENSE
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2023 Langchain AI
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
README.md
+78
View File
@@ -0,0 +1,78 @@
[![Unit Tests](https://github.com/langchain-ai/kork/actions/workflows/test.yml/badge.svg?branch=main&event=push)](https://github.com/langchain-ai/kork/actions/workflows/test.yml)
# Kork ![alt The Parrot](assets/parrot.png)
`Kork` is an *experimental* [Langchain chain](https://python.langchain.com/en/latest/modules/chains.html) that can write and execute code safely.
`Kork` **cannot** write *arbitrary* code. If that's what you need, save yourself some time and use [docker](https://www.docker.com/) and a real programming language.
`Kork` does not aim implement a fully functional programming language, but to make it
easier to get an LLM to code that is *safe* and *useful* for a *constrained* problem
domain.
A calculator powered by natural language is one example of a constrained problem domain.
The main feature of the `Kork` programming language and interpreter is that it does not have any features. To do useful things with `Kork` a user specifies *foreign functions* that can be called from the `Kork` interpreter by the chain. The chain uses an LLM to translate a user query into a `Kork` program, which is then executed by a `Kork` interpreter.
## Safety
Nothing is bulletproof when humans are involved.
* Don't do silly things like exposing `eval` as a foreign function.
* If a foreign function allocates memory, the LLM could ask for more memory than is available crashing the process.
* Think carefully about whether the LLM could behave maliciously before trusting its output.
## Quality
The short version: No benchmarks yet!
The long version: Quality of generated programs depends on many factors.
You can experiment with tweaking the prompt (e.g., explain the syntax of the language), changing the foreign function retriever (e.g., retrieve the most relevant foreign functions based on the user query), or providing examples in the form of (query, expected program).
The prompt, examples and syntax of the language can trip up the LLM to assume
it's programming in a specific language (e.g., typescript or python) and assume
that it can use language features that are not supported by `Kork` or to import
libraries that are not available.
## Limitations
- The `Kork` language and interpreter are limited to *function invocation*, *variable declaration* and
*basic arithmetic* (no function declaration, loops etc.)
- Only supporting `int`, `float`, `str`, `type(None)`, `bool` types. No support for `lists` or `object` types.
- Very limited type annotations.
## Future Work
Let us know if you're interested in contributing or have ideas for improvements!
- Allow enabling/disabling language features
- Allow changing underlying language syntax
- Add support for objects
- Add other language features (e.g., loops)
- Provide foreign function retriever implementations based on similarity to user query
## Compatibility
`Kork` is tested against python 3.8, 3.9, 3.10, 3.11.
## Installation
```sh
pip install kork
```
## 🙏 Appreciation
* [Lark](https://github.com/lark-parser/lark) -- For making it easy to define a grammar and parse it!
* [Bob Nystrom](https://github.com/munificent) -- For writing [Crafting Interpreters](https://www.craftinginterpreters.com/)!
## © Contributing
If you have any ideas or feature requests, please open an issue and share!
See [CONTRIBUTING.md](https://github.com/langchain-ai/kork/blob/main/CONTRIBUTING.md) for more information.
## 🎶 Why the name?
Fast to type and maybe sufficiently unique.
assets/parrot.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 18 KiB

docs/Makefile
+20
View File
@@ -0,0 +1,20 @@
# Minimal makefile for Sphinx documentation
#
# You can set these variables from the command line, and also
# from the environment for the first two.
SPHINXOPTS ?=
SPHINXBUILD ?= sphinx-build
SOURCEDIR = source
BUILDDIR = build
# Put it first so that "make" without argument is like "make help".
help:
@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
.PHONY: help Makefile
# Catch-all target: route all unknown targets to Sphinx using the new
# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS).
%: Makefile
@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
docs/make.bat
+35
View File
@@ -0,0 +1,35 @@
@ECHO OFF
pushd %~dp0
REM Command file for Sphinx documentation
if "%SPHINXBUILD%" == "" (
set SPHINXBUILD=sphinx-build
)
set SOURCEDIR=source
set BUILDDIR=build
%SPHINXBUILD% >NUL 2>NUL
if errorlevel 9009 (
echo.
echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
echo.installed, then set the SPHINXBUILD environment variable to point
echo.to the full path of the 'sphinx-build' executable. Alternatively you
echo.may add the Sphinx directory to PATH.
echo.
echo.If you don't have Sphinx installed, grab it from
echo.https://www.sphinx-doc.org/
exit /b 1
)
if "%1" == "" goto help
%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
goto end
:help
%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
:end
popd
docs/source/_static/parrot.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 18 KiB

docs/source/api.rst
+34
View File
@@ -0,0 +1,34 @@
.. _api:
.. currentmodule:: kork
API
----------
The main **Kork** API is shown here:
.. autosummary::
CodeChain
AstPrinter
AbstractRetriever
SimpleRetriever
Kork Interpreter
=================
.. autosummary::
InterpreterResult
run_interpreter
Environment
Index
=====
.. automodule:: generated.kork
docs/source/ast.ipynb
+348
View File
@@ -0,0 +1,348 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "62549efe",
"metadata": {
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"# Abstract Syntax Tree\n",
"\n",
"`Kork` parses the code into the following AST https://github.com/langchain-ai/kork/blob/main/kork/ast.py.\n",
"\n",
"Below we'll show a few of the utilities to parse and print the ast."
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "79fd8e41",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
},
"nbsphinx": "hidden",
"pycharm": {
"name": "#%%\n"
},
"tags": [
"remove_cell"
]
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"The autoreload extension is already loaded. To reload it, use:\n",
" %reload_ext autoreload\n"
]
}
],
"source": [
"%load_ext autoreload\n",
"%autoreload 2\n",
"\n",
"import sys\n",
"\n",
"sys.path.insert(0, \"../\")"
]
},
{
"cell_type": "markdown",
"id": "4fb92038-08ad-4cf6-91db-06fb25c281bd",
"metadata": {
"tags": []
},
"source": [
"## Parse"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "b5775536-ea67-4ed6-97e7-4d6f69f8d83a",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from kork.parser import parse"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "b6fec8a2-e4ce-48a8-b490-1e8be83886e4",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"Program(stmts=(Variable(name='x'),))"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"parse(\"x\")"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "4811e4e7-2a91-45f1-8d1f-3d9d75755698",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"Program(stmts=(Assign(name='x', value=Literal(value=1)),))"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"parse(\"x = 1\")"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "b247fe8c-f782-4442-afb5-f84976aab87b",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"Program(stmts=(ExternFunctionDef(name='foo', params=ParamList(params=[]), return_type='Any', implementation=None, doc_string=''), Assign(name='x', value=Literal(value=1)), Assign(name='y', value=Binary(left=Binary(left=Variable(name='x'), operator='*', right=Variable(name='x')), operator='+', right=FunctionCall(name='foo', args=[])))))"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"parse(\n",
" \"\"\"\n",
"extern fn foo() -> Any # Comment\n",
"\n",
"x = 1; // Comment\n",
"y = x * x + foo()\n",
"\"\"\"\n",
")"
]
},
{
"cell_type": "markdown",
"id": "fb6ac75e-6f3f-41b1-b0b8-c05db72a3549",
"metadata": {},
"source": [
"## AstPrinter\n",
"\n",
"The AST printer can be used to translate AST syntax into a string representation of the program.\n",
"\n",
"This functionality is used to generate the prompt for the LLM as we'll see later in the docs."
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "62630c9b-eb0b-49a3-9bb8-f7bdfd3846b3",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from kork import AstPrinter"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "9f22fd0e-5390-4102-9141-00c1f76ab618",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"program = parse(\n",
" \"\"\"\n",
" extern fn foo() -> Any # Comment\n",
"x=1; // Comment\n",
"y=x*x + foo()\n",
"\"\"\"\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "9cb394cd-4d41-4ad0-944e-359ca70c4b40",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"ParamList(params=[])"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"program.stmts[0].params"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "282a762d-5361-4d7d-b31d-b2f25912ff42",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"extern fn foo() -> Any\n",
"x = 1\n",
"y = x * x + foo()\n"
]
}
],
"source": [
"print(AstPrinter().visit(program))"
]
},
{
"cell_type": "markdown",
"id": "fe62fa93-90f6-4554-82b4-72ba4f873158",
"metadata": {},
"source": [
"## Nodes"
]
},
{
"cell_type": "markdown",
"id": "76a1c4bf-34d7-4536-a08a-34cd7bac4901",
"metadata": {},
"source": [
"The nodes in the ast are in the `ast` module."
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "55051ae1-3490-4840-aed4-b8b54d7976d0",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from kork.ast import FunctionCall, FunctionDef"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "37a7d6e7-5ea3-4c2d-8d3d-cd9e5c27a117",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"\u001b[0;31mInit signature:\u001b[0m \u001b[0mFunctionCall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'str'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'Sequence[Expr]'\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;31mDocstring:\u001b[0m Represent a function call.\n",
"\u001b[0;31mFile:\u001b[0m ~/src/kork/kork/ast.py\n",
"\u001b[0;31mType:\u001b[0m ABCMeta\n",
"\u001b[0;31mSubclasses:\u001b[0m "
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"?FunctionCall"
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "686ae81a-c866-4950-a852-f1288c7c6a9b",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"\u001b[0;31mInit signature:\u001b[0m\n",
"\u001b[0mFunctionDef\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
"\u001b[0;34m\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'str'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
"\u001b[0;34m\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'ParamList'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
"\u001b[0;34m\u001b[0m \u001b[0mbody\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'Sequence[Union[Stmt, Expr]]'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
"\u001b[0;34m\u001b[0m \u001b[0mreturn_type\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;34m'str'\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
"\u001b[0;34m\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;31mDocstring:\u001b[0m Represent a function definition with an implementation.\n",
"\u001b[0;31mFile:\u001b[0m ~/src/kork/kork/ast.py\n",
"\u001b[0;31mType:\u001b[0m ABCMeta\n",
"\u001b[0;31mSubclasses:\u001b[0m "
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"?FunctionDef"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.2"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
docs/source/calculator.ipynb
+335
View File
File diff suppressed because one or more lines are too long
docs/source/conf.py
+106
View File
@@ -0,0 +1,106 @@
# Configuration file for the Sphinx documentation builder.
#
# This file only contains a selection of the most common options. For a full
# list see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html
# -- Path setup --------------------------------------------------------------
# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#
# import os
# import sys
# sys.path.insert(0, os.path.abspath('.'))
# -- Project information -----------------------------------------------------
import pathlib
import sys
from typing import List
import toml
ROOT_FOLDER = str(pathlib.Path(__file__).parent.parent.parent)
# Add the project root to the path
sys.path.insert(0, ROOT_FOLDER)
with open("../../pyproject.toml") as f:
data = toml.load(f)
project = "Kork"
copyright = "2023, Langchain AI"
author = "Langchain AI"
version = data["tool"]["poetry"]["version"]
release = version
html_title = project + " " + version
# -- General configuration ---------------------------------------------------
# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
extensions = [
"sphinx.ext.autodoc",
"sphinx.ext.autodoc.typehints",
"sphinx.ext.autosummary",
"sphinx.ext.napoleon",
"sphinx.ext.viewcode",
"myst_nb",
"sphinx_copybutton",
"IPython.sphinxext.ipython_console_highlighting",
]
source_suffix = [".ipynb", ".html", ".md", ".rst"]
# Add any paths that contain templates here, relative to this directory.
templates_path = ["_templates"]
# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This pattern also affects html_static_path and html_extra_path.
exclude_patterns: List[str] = []
# -- Options for HTML output -------------------------------------------------
# The theme to use for HTML and HTML Help pages. See the documentation for
# a list of builtin themes.
#
html_theme = "sphinx_book_theme"
html_theme_options = {
"path_to_docs": "docs",
"repository_url": "https://github.com/langchain-ai/kork",
"home_page_in_toc": True,
"show_navbar_depth": 2,
"use_sidenotes": True,
"use_repository_button": True,
}
html_context = {
"display_github": True, # Integrate GitHub
"github_user": "langchain-ai", # Username
"github_repo": "kork", # Repo name
"github_version": "main", # Version
"conf_py_path": "/docs/", # Path in the checkout to the docs root
}
# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ["_static"]
# These paths are either relative to html_static_path
# or fully qualified paths (eg. https://...)
html_css_files = [
"css/custom.css",
]
html_logo = "_static/parrot.png"
nb_execution_mode = "off"
autosummary_generate = True
docs/source/examples.ipynb
+454
View File
@@ -0,0 +1,454 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "62549efe",
"metadata": {
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"# Examples\n",
"\n",
"`Kork` allows providing examples in the prompt. \n",
"\n",
"1. Examples help the LLM understand the syntax of the language\n",
"2. Examples can be dynamically generated based on the user query, to surface the most relevant programs.\n",
"\n",
"The format of examples is a sequence of 2-tuples of the form: (`str`, `Program`).\n",
"\n",
"Let's see how to specify them."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "79fd8e41",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
},
"nbsphinx": "hidden",
"pycharm": {
"name": "#%%\n"
},
"tags": [
"remove_cell"
]
},
"outputs": [],
"source": [
"%load_ext autoreload\n",
"%autoreload 2\n",
"\n",
"import sys\n",
"\n",
"sys.path.insert(0, \"../\")"
]
},
{
"cell_type": "markdown",
"id": "27975824-3db3-48ad-bbdd-6784827f7a99",
"metadata": {},
"source": [
"## The Simplest Way"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "5c5b3a8d-c464-499c-a5d1-275ab2ceb226",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from kork.parser import parse"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "1a6dfced-5e2c-4c34-8dcd-a5d207bcd92d",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"examples_as_strings = [\n",
" (\n",
" \"declare a variable called `y` and assign to it the value 8\",\n",
" \"var y = 8\",\n",
" )\n",
"]"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "23b683ed-f52a-4fce-b554-9fbbbaa4af4c",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"examples = [(query, parse(code)) for query, code in examples_as_strings]"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "03e16c63-f804-4076-a70a-2d324dc43e8a",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"[('declare a variable called `y` and assign to it the value 8',\n",
" Program(stmts=(VarDecl(name='y', value=Literal(value=8)),)))]"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"examples"
]
},
{
"cell_type": "markdown",
"id": "4fb92038-08ad-4cf6-91db-06fb25c281bd",
"metadata": {
"tags": []
},
"source": [
"## The Hardest Way\n",
"\n",
"Let's import AST elements and write an example query and a corresponding short program."
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "b5775536-ea67-4ed6-97e7-4d6f69f8d83a",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from kork.ast import Program, VarDecl, Literal"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "b6fec8a2-e4ce-48a8-b490-1e8be83886e4",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"examples = [\n",
" (\n",
" \"declare a variable called `y` and assign to it the value 8\",\n",
" Program(stmts=[VarDecl(\"y\", Literal(value=8))]),\n",
" )\n",
"]"
]
},
{
"cell_type": "markdown",
"id": "0b8b7689-2d41-4693-a389-7ccf5916aeb3",
"metadata": {
"tags": []
},
"source": [
"## Formatted Example\n",
"\n",
"The chain formats the example using an `AstPrinter` to format the examples as a sequence of the 2-tuple of (`str`, `str`).\n",
"\n",
"Here's what happens:"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "a1e2b3c0-e049-49c2-af20-d5c52e34db9c",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from kork.examples import format_examples\n",
"from kork import AstPrinter"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "83eea564-0102-4d89-b627-7781d4d76a53",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[('declare a variable called `y` and assign to it the value 8', '```SmirkingCat\\nvar y = 8\\n```')]\n"
]
}
],
"source": [
"print(format_examples(\"SmirkingCat\", examples, AstPrinter()))"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "2e8c20f7-b419-4955-b5e3-c8046ec29357",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[('Text: \"\"\"\\ndeclare a variable called `y` and assign to it the value 8\\n\"\"\"', '```SmirkingCat\\nvar y = 8\\n```')]\n"
]
}
],
"source": [
"print(\n",
" format_examples(\n",
" \"SmirkingCat\", examples, AstPrinter(), input_formatter=\"text_prefix\"\n",
" )\n",
")"
]
},
{
"cell_type": "markdown",
"id": "d0db700c-f182-4c24-a060-2e19f4b3102d",
"metadata": {},
"source": [
"## The Hardest Way Made Easier\n",
"\n",
"There are utility functions that make it easier to generate programs that invoke foreign functions -- which is\n",
"expected to be a common pattern."
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "0fe1ff0d-f703-47a6-8398-97c5a853710f",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from kork import c_, r_"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "fb84132e-a6ac-4b47-a891-2667420b24c3",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"\u001b[0;31mSignature:\u001b[0m \u001b[0mc_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mCallable\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mAny\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mkork\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mast\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mFunctionCall\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;31mDocstring:\u001b[0m Create a kork function call.\n",
"\u001b[0;31mFile:\u001b[0m ~/src/kork/kork/examples.py\n",
"\u001b[0;31mType:\u001b[0m function"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"?c_"
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "b5f678c5-4715-4e96-950a-9b82ed0485d4",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"\u001b[0;31mSignature:\u001b[0m \u001b[0mr_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mexpr\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mkork\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mast\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mExpr\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mkork\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mast\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mProgram\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;31mDocstring:\u001b[0m Assign last program expression to a result variable.\n",
"\u001b[0;31mFile:\u001b[0m ~/src/kork/kork/examples.py\n",
"\u001b[0;31mType:\u001b[0m function"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"?r_"
]
},
{
"cell_type": "code",
"execution_count": 14,
"id": "8aeef131-2c87-473c-9780-5d6e391993aa",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import math\n",
"import operator"
]
},
{
"cell_type": "markdown",
"id": "054f4260-7c49-4af9-8720-25c990379ee9",
"metadata": {
"tags": []
},
"source": [
"Here's a small program that invokes the foreign function `log10` with arg `2`"
]
},
{
"cell_type": "code",
"execution_count": 15,
"id": "3b411f83-4413-4ace-a786-5c9a13bca343",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"FunctionCall(name='log10', args=[Literal(value=2)])"
]
},
"execution_count": 15,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"c_(math.log10, 2)"
]
},
{
"cell_type": "markdown",
"id": "86e9bee5-558d-4344-a315-29770d9b1b3e",
"metadata": {
"tags": []
},
"source": [
"A program that invokes `log10` with arg of 2 and stores the result in a `result` variable in the global namespace."
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "07480868-5d28-48e1-8b58-1b72c0c65ee0",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"Program(stmts=[VarDecl(name='result', value=FunctionCall(name='log10', args=[Literal(value=2)]))])"
]
},
"execution_count": 16,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"r_(c_(math.log10, 2))"
]
},
{
"cell_type": "code",
"execution_count": 17,
"id": "d4113fd1-f17d-4177-a18d-e571088d329e",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"program = r_(c_(operator.add, 2, c_(math.log10, 3)))"
]
},
{
"cell_type": "code",
"execution_count": 20,
"id": "9facd1c6-b203-49d1-840e-12a3efaa5b11",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"var result = add(2, log10(3))\n"
]
}
],
"source": [
"print(AstPrinter().visit(program))"
]
},
{
"cell_type": "code",
"execution_count": 19,
"id": "77fcb89f-a8bb-4ff6-a63b-574f39087642",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"examples = [(\"take the log base 10 of 3 and add 2 to the result.\", program)]"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.2"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
docs/source/index.md
+100
View File
@@ -0,0 +1,100 @@
# Introduction
`Kork` is an *experimental* [Langchain chain](https://python.langchain.com/en/latest/modules/chains.html) that can write and execute code safely.
`Kork` **cannot** write *arbitrary* code. If that's what you need, save yourself some time and use [docker](https://www.docker.com/) and a real programming language.
`Kork` does not aim implement a fully functional programming language, but to make it
easier to get an LLM to code that is *safe* and *useful* for a *constrained* problem
domain.
A calculator powered by natural language is one example of a constrained problem domain.
The main feature of the `Kork` programming language and interpreter is that it does not have any features. To do useful things with `Kork` a user specifies *foreign functions* that can be called from the `Kork` interpreter by the chain. The chain uses an LLM to translate a user query into a `Kork` program, which is then executed by a `Kork` interpreter.
## Safety
Nothing is bulletproof when humans are involved.
* Don't do silly things like exposing `eval` as a foreign function.
* If a foreign function allocates memory, the LLM could ask for more memory than is available crashing the process.
* Think carefully about whether the LLM could behave maliciously before trusting its output.
## Quality
The short version: No benchmarks yet!
The long version: Quality of generated programs depends on many factors.
You can experiment with tweaking the prompt (e.g., explain the syntax of the language), changing the foreign function retriever (e.g., retrieve the most relevant foreign functions based on the user query), or providing examples in the form of (query, expected program).
The prompt, examples and syntax of the language can trip up the LLM to assume
it's programming in a specific language (e.g., typescript or python) and assume
that it can use language features that are not supported by `Kork` or to import
libraries that are not available.
## Limitations
- The `Kork` language and interpreter are limited to *function invocation*, *variable declaration* and
*basic arithmetic* (no function declaration, loops etc.)
- Only supporting `int`, `float`, `str`, `type(None)`, `bool` types. No support for `lists` or `object` types.
- Very limited type annotations.
## Future Work
Let us know if you're interested in contributing or have ideas for improvements!
- Allow enabling/disabling language features
- Allow changing underlying language syntax
- Add support for objects
- Add other language features (e.g., loops)
- Provide foreign function retriever implementations based on similarity to user query
## Compatibility
`Kork` is tested against python 3.8, 3.9, 3.10, 3.11.
## Installation
```sh
pip install kork
```
## 🙏 Appreciation
* [Lark](https://github.com/lark-parser/lark) -- For making it easy to define a grammar and parse it!
* [Bob Nystrom](https://github.com/munificent) -- For writing [Crafting Interpreters](https://www.craftinginterpreters.com/)!
## © Contributing
If you have any ideas or feature requests, please open an issue and share!
See [CONTRIBUTING.md](https://github.com/langchain-ai/kork/blob/main/CONTRIBUTING.md) for more information.
## 🎶 Why the name?
Fast to type and maybe sufficiently unique.
```{toctree}
:maxdepth: 2
:caption: Contents
introduction
language
ast
examples
prompt
retrievers
```
```{toctree}
:caption: Examples
calculator
```
```{toctree}
:maxdepth: 2
:caption: Advanced
api
```
docs/source/introduction.ipynb
+1107
View File
File diff suppressed because one or more lines are too long
docs/source/language.ipynb
+456
View File
@@ -0,0 +1,456 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "62549efe",
"metadata": {
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"# The Language 😼\n",
"\n",
"We'll call `Kork`'s programming language the **Smirking Cat** (😼) or `Kork` or simply 😼. Honestly, it doesn't really matter (as long as the LLM doesn't get confused!).\n",
"Set the language name in the prompt to whatever you'd like. Keep it short to avoid consuming too many tokens on things that don't matter. 🙃\n",
"\n",
"The grammar for the language is defined here: https://github.com/langchain-ai/kork/blob/main/kork/parser.py.\n",
"\n",
"This language isn't particularly good and is extremely limited (especially at the moment) -- it was clobbered together\n",
"in a few hours, is incorrect in some known ways (e.g., `const` isn't actually `const`) and most definitely incorrect in unknown ways as well. 🤫\n",
"\n",
"The goal is *NOT* to allow the LLM to write arbitrary python code. If you need arbitrary code execution use a real programming language.\n",
"\n",
"`Kork` is meant to accommodate scenarios in which one wants to produce small constrained programs that can help achieve some specific task phrased in natural language. \n",
"\n",
"It was designed with the following ideas in mind: \n",
"\n",
"* Keep the language minimal to limit what the LLM can code\n",
"* Encourage the LLM to rely heavily on external function invocation\n",
"* Keep the syntax similar to existing languages, so the LLM can learn how to code in it using information in the prompt\n",
"* Discourage the LLM from assuming a standard library is available\n",
"\n",
"## Alternative Approaches\n",
"\n",
"* Changing syntax: syntax based on S-expressions looks promising based on a few qualitative experiments.\n",
"* Implementing a light weight python / typescript interpreter in python that supports a minimal subset of features."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "79fd8e41",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
},
"nbsphinx": "hidden",
"pycharm": {
"name": "#%%\n"
},
"tags": [
"remove_cell"
]
},
"outputs": [],
"source": [
"%load_ext autoreload\n",
"%autoreload 2\n",
"\n",
"import sys\n",
"\n",
"sys.path.insert(0, \"../\")"
]
},
{
"cell_type": "markdown",
"id": "4fb92038-08ad-4cf6-91db-06fb25c281bd",
"metadata": {
"tags": []
},
"source": [
"## The interpreter\n",
"\n",
"Let's take a look at the **Smirking Cat** interpreter."
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "62630c9b-eb0b-49a3-9bb8-f7bdfd3846b3",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from kork import run_interpreter"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "f0714c89-d220-44f1-a5f7-822d3434502a",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"{'environment': Environment(parent=None, variables={'x': 10}), 'errors': []}"
]
},
"execution_count": 3,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"result = run_interpreter(\"var x = 1; x = x * 10\")\n",
"result"
]
},
{
"cell_type": "markdown",
"id": "f5a0586b-a6f3-4f03-8245-5e67b58f3ef3",
"metadata": {},
"source": [
"The environment reflects the state of the global symbols after the interpreter finished executing the code."
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "927549ac-c80b-4e71-bae2-27f5f75b30a0",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"{'x': 10}"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"result[\"environment\"].variables"
]
},
{
"cell_type": "markdown",
"id": "9e61c733-e1aa-4b01-ab8f-570274331770",
"metadata": {},
"source": [
"Syntax errors will lead the `errors` key being populated. The current error isn't super informative."
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "3f3e95a1-65a0-46fa-8538-8ef35eedbaf1",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"{'environment': Environment(parent=None, variables={}),\n",
" 'errors': [lark.exceptions.UnexpectedToken()]}"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"run_interpreter(\"1 = 2\")"
]
},
{
"cell_type": "markdown",
"id": "b47e4f89-9734-475b-ac40-248430236cb7",
"metadata": {},
"source": [
"RunTimeExceptions will also appear in the `errors` key."
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "02b5723d-985d-4187-b5b3-048e192af59d",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"{'environment': Environment(parent=None, variables={}),\n",
" 'errors': [kork.exceptions.KorkRunTimeException('Variable `x` not found')]}"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"run_interpreter(\"x + 1\")"
]
},
{
"cell_type": "markdown",
"id": "cc821fab-8b69-49f6-86fd-f1022215001b",
"metadata": {},
"source": [
"## Environment\n",
"\n",
"The interpreter can be run with a pre-populated environment"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "66fbfc14-4d4c-4933-8780-5b8d287b087a",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from kork import Environment"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "2f39a42e-1073-426b-ab02-994e9e710779",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"env = Environment()"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "a76df2c1-772a-4f3f-a028-3d5ca2cb36b9",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"10"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"env.set_symbol(\"x\", 10)"
]
},
{
"cell_type": "markdown",
"id": "105a1633-4640-4bd1-9c9b-28643edfc9c4",
"metadata": {},
"source": [
"Please note that `y` was created even though no `let` or `const` or `var` keyword was used."
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "b357aba2-1b50-48e9-a350-23f4ecd9dcbe",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"{'environment': Environment(parent=None, variables={'x': 10, 'y': 30}),\n",
" 'errors': []}"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"result = run_interpreter(\"y = x * 3\", environment=env)\n",
"result"
]
},
{
"cell_type": "markdown",
"id": "04a27351-3100-4c98-b8fa-9afabbb9da90",
"metadata": {},
"source": [
"## Foreign Functions\n",
"\n",
"A user can import existing python functions as foreign functions."
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "51ef8652-7e3a-4cae-8c28-39ed909c14b1",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from kork.foreign_funcs import to_extern_func_def"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "24577081-e906-4695-9df6-e06587a60a27",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"def foo(s: str) -> str:\n",
" \"\"\"Foo will reverse a string!\"\"\"\n",
" return s[::-1]"
]
},
{
"cell_type": "markdown",
"id": "4bccc8ec-33f1-400a-88df-19029b0ef207",
"metadata": {},
"source": [
"Below we'll convert the python function into an internal representation of an external function definition."
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "50707464-64cd-4d5e-ba8a-7cebdd86a756",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"extern_func_def = to_extern_func_def(foo)"
]
},
{
"cell_type": "code",
"execution_count": 14,
"id": "292e356a-fcb0-4694-ada6-6e561e4c22e4",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"ExternFunctionDef(name='foo', params=ParamList(params=[Param(name='s', type_='str')]), return_type='str', implementation=<function foo at 0x7f5e68873370>, doc_string='Foo will reverse a string!')"
]
},
"execution_count": 14,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"extern_func_def"
]
},
{
"cell_type": "markdown",
"id": "0c82a5d1-310d-4d49-b439-8290f7cd41e8",
"metadata": {},
"source": [
"Let's add this function to the environment"
]
},
{
"cell_type": "code",
"execution_count": 15,
"id": "6de0af02-ad56-443c-b3f8-f9b1be87f5e7",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"ExternFunctionDef(name='foo', params=ParamList(params=[Param(name='s', type_='str')]), return_type='str', implementation=<function foo at 0x7f5e68873370>, doc_string='Foo will reverse a string!')"
]
},
"execution_count": 15,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"env.set_symbol(\"foo\", extern_func_def)"
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "fef9f8e9-3254-4417-9170-3a995c7450e1",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"'woem'"
]
},
"execution_count": 16,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"result = run_interpreter('var z = foo(\"meow\")', environment=env)\n",
"\n",
"result[\"environment\"].variables[\"z\"]"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.2"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
docs/source/prompt.ipynb
+291
View File
@@ -0,0 +1,291 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "62549efe",
"metadata": {
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"# The Prompt\n",
"\n",
"Let's see how to customize the prompt."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "79fd8e41",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
},
"nbsphinx": "hidden",
"pycharm": {
"name": "#%%\n"
},
"tags": [
"remove_cell"
]
},
"outputs": [],
"source": [
"%load_ext autoreload\n",
"%autoreload 2\n",
"\n",
"import sys\n",
"\n",
"sys.path.insert(0, \"../\")"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "1d3a2212-38fa-4ab1-bec2-e1f4aeb8af47",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import math\n",
"import langchain\n",
"from kork import (\n",
" CodeChain,\n",
" ast,\n",
" AstPrinter,\n",
" c_,\n",
" r_,\n",
" run_interpreter,\n",
")\n",
"from langchain import PromptTemplate\n",
"from kork import SimpleContextRetriever"
]
},
{
"cell_type": "markdown",
"id": "4fb100da-29d9-442b-b89c-4b60bbc1967b",
"metadata": {},
"source": [
"The code below defines a place holder model for a chat model.\n",
"Feel free to replace it with a real language model."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "3c5e0604-b50c-470d-855c-9b5253206a81",
"metadata": {
"tags": [
"hide-input"
]
},
"outputs": [],
"source": [
"from typing import Any, List, Optional\n",
"\n",
"from langchain.chat_models.base import BaseChatModel\n",
"from langchain.schema import AIMessage, BaseMessage, ChatGeneration, ChatResult\n",
"from pydantic import Extra\n",
"\n",
"\n",
"class ToyChatModel(BaseChatModel):\n",
" response: str\n",
"\n",
" class Config:\n",
" \"\"\"Configuration for this pydantic object.\"\"\"\n",
"\n",
" extra = Extra.forbid\n",
" arbitrary_types_allowed = True\n",
"\n",
" def _generate(\n",
" self, messages: List[BaseMessage], stop: Optional[List[str]] = None\n",
" ) -> ChatResult:\n",
" message = AIMessage(content=self.response)\n",
" generation = ChatGeneration(message=message)\n",
" return ChatResult(generations=[generation])\n",
"\n",
" async def _agenerate(\n",
" self, messages: List[BaseMessage], stop: Optional[List[str]] = None\n",
" ) -> Any:\n",
" \"\"\"Async version of _generate.\"\"\"\n",
" message = AIMessage(content=self.response)\n",
" generation = ChatGeneration(message=message)\n",
" return ChatResult(generations=[generation])"
]
},
{
"cell_type": "markdown",
"id": "07c73183-ffd6-450c-abe0-685e89ca5054",
"metadata": {},
"source": [
"## Custom prompt\n",
"\n",
"Prompts take two optional input variables: the `language_name` and the `external_funcitons_block`. If the prompt uses such a variable, the variable will be auto-populated by the chain."
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "bd679388-6320-4933-aa1c-17ae8b719289",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"instruction_template = PromptTemplate(\n",
" template=\"\"\"\n",
"You are coding in a language called the `{language_name}`.\n",
"\n",
"\n",
"\n",
"{external_functions_block}\n",
"\n",
"Begin!\\n\n",
"\"\"\",\n",
" input_variables=[\"language_name\", \"external_functions_block\"],\n",
")"
]
},
{
"cell_type": "markdown",
"id": "ab4d3808-50ab-4ade-aab7-677e23624cab",
"metadata": {},
"source": [
"## Declare the chain"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "5e2df25f-169d-4961-9fd9-4b10bb9f07c9",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"chain = CodeChain.from_defaults(\n",
" llm=ToyChatModel(response=\"MEOW MEOW MEOW MEOW\"), # The LLM to use\n",
" examples=[\n",
" (\"2**5\", r_(c_(math.pow, 2, 5))),\n",
" (\"take the log base 2 of 2\", r_(c_(math.log2, 2))),\n",
" ], # Example programs\n",
" context=[math.pow, math.log2, math.log10],\n",
" instruction_template=instruction_template,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "7d92f7b6-d1f5-41b5-826b-b338edc3831e",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"environment, few_shot_prompt = chain.prepare_context(query=\"hello\")"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "e72e5f9e-fe2d-4f6b-81c1-52c89ecd032e",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"[ExternFunctionDef(name='pow', params=ParamList(params=[Param(name='x', type_='Any'), Param(name='y', type_='Any')]), return_type='Any', implementation=<built-in function pow>, doc_string='Return x**y (x to the power of y).'),\n",
" ExternFunctionDef(name='log2', params=ParamList(params=[Param(name='x', type_='Any')]), return_type='Any', implementation=<built-in function log2>, doc_string='Return the base 2 logarithm of x.'),\n",
" ExternFunctionDef(name='log10', params=ParamList(params=[Param(name='x', type_='Any')]), return_type='Any', implementation=<built-in function log10>, doc_string='Return the base 10 logarithm of x.')]"
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"environment.list_external_functions()"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "06b5f073-bb8c-40a5-a816-52075417dee2",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"You are coding in a language called the `😼`.\n",
"\n",
"\n",
"\n",
"You have access to the following external functions:\n",
"\n",
"```😼\n",
"\n",
"extern fn pow(x: Any, y: Any) -> Any // Return x**y (x to the power of y).\n",
"extern fn log2(x: Any) -> Any // Return the base 2 logarithm of x.\n",
"extern fn log10(x: Any) -> Any // Return the base 10 logarithm of x.\n",
"```\n",
"\n",
"\n",
"Begin!\n",
"\n",
"Input: \"\"\"\n",
"2**5\n",
"\"\"\"\n",
"Output: ```😼\n",
"var result = pow(2, 5)\n",
"```\n",
"\n",
"Input: \"\"\"\n",
"take the log base 2 of 2\n",
"\"\"\"\n",
"Output: ```😼\n",
"var result = log2(2)\n",
"```\n",
"\n",
"Input: [user input]\n",
"Output:\n"
]
}
],
"source": [
"print(few_shot_prompt.format_prompt(query=\"[user input]\").to_string())"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.2"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
docs/source/retrievers.ipynb
+392
View File
@@ -0,0 +1,392 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "62549efe",
"metadata": {
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"# Retrievers\n",
"\n",
"The `Kork` CodeChain accepts retriever interfaces for both examples and foreign functions.\n",
"\n",
"The built-in retrievers return the content that they were initialized with as is.\n",
"\n",
"You could implement dynamic retrievers that rely on a vector store and return the most relevant examples and foreign functions to achieve the task at hand."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "79fd8e41",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
},
"nbsphinx": "hidden",
"pycharm": {
"name": "#%%\n"
},
"tags": [
"remove_cell"
]
},
"outputs": [],
"source": [
"%load_ext autoreload\n",
"%autoreload 2\n",
"\n",
"import sys\n",
"\n",
"sys.path.insert(0, \"../\")"
]
},
{
"cell_type": "markdown",
"id": "4fb100da-29d9-442b-b89c-4b60bbc1967b",
"metadata": {},
"source": [
"The code below defines a place holder model for a chat model.\n",
"Feel free to replace it with a real language model."
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "3c5e0604-b50c-470d-855c-9b5253206a81",
"metadata": {
"tags": [
"hide-input"
]
},
"outputs": [],
"source": [
"from typing import Any, List, Optional\n",
"\n",
"from langchain.chat_models.base import BaseChatModel\n",
"from langchain.schema import AIMessage, BaseMessage, ChatGeneration, ChatResult\n",
"from pydantic import Extra\n",
"\n",
"\n",
"class ToyChatModel(BaseChatModel):\n",
" response: str\n",
"\n",
" class Config:\n",
" \"\"\"Configuration for this pydantic object.\"\"\"\n",
"\n",
" extra = Extra.forbid\n",
" arbitrary_types_allowed = True\n",
"\n",
" def _generate(\n",
" self, messages: List[BaseMessage], stop: Optional[List[str]] = None\n",
" ) -> ChatResult:\n",
" message = AIMessage(content=self.response)\n",
" generation = ChatGeneration(message=message)\n",
" return ChatResult(generations=[generation])\n",
"\n",
" async def _agenerate(\n",
" self, messages: List[BaseMessage], stop: Optional[List[str]] = None\n",
" ) -> Any:\n",
" \"\"\"Async version of _generate.\"\"\"\n",
" message = AIMessage(content=self.response)\n",
" generation = ChatGeneration(message=message)\n",
" return ChatResult(generations=[generation])"
]
},
{
"cell_type": "markdown",
"id": "77a7b08f-9d98-475c-b10e-711cbb9d1fbe",
"metadata": {},
"source": [
"## Context Retriever\n",
"\n",
"A context retriever (yes, it's poorly named) supports a method that takes the user `query` and returns a list of foreign functions.\n",
"\n",
"The built-in retriever returns the information it was initialized with, entirely ignoring the `query`.\n",
"\n",
"Sub-class from the abstract interface to provide your own implementation!"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "32df6b1b-edeb-4117-be67-c225db5c7a28",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from kork import SimpleContextRetriever\n",
"import math"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "a3db1e43-5445-40c7-b11c-6e27c5c3dfca",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"simple_context_retriever = SimpleContextRetriever.from_functions(\n",
" [math.pow, math.log2, math.log10]\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "bcd5a5c6-2b4e-4ef9-842f-992bd8a823a4",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"[ExternFunctionDef(name='pow', params=ParamList(params=[Param(name='x', type_='Any'), Param(name='y', type_='Any')]), return_type='Any', implementation=<built-in function pow>, doc_string='Return x**y (x to the power of y).'),\n",
" ExternFunctionDef(name='log2', params=ParamList(params=[Param(name='x', type_='Any')]), return_type='Any', implementation=<built-in function log2>, doc_string='Return the base 2 logarithm of x.'),\n",
" ExternFunctionDef(name='log10', params=ParamList(params=[Param(name='x', type_='Any')]), return_type='Any', implementation=<built-in function log10>, doc_string='Return the base 10 logarithm of x.')]"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"simple_context_retriever.retrieve(\"this is the query\")"
]
},
{
"cell_type": "markdown",
"id": "f9a05cc8-579f-4bf7-a8b9-61fe13b336ed",
"metadata": {},
"source": [
"## Example Retriever\n",
"\n",
"Supports the same `retrieve` method as a context retriever. \n",
"\n",
"The built-in example retriever returns the information it was initialized with, entirely ignoring the `query`.\n",
"\n",
"Sub-class from the abstract interface to provide your own implementation!"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "54d8822b-31f1-4e35-9399-2fb8e09ea91b",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from kork import c_, r_, AstPrinter, SimpleExampleRetriever"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "73a91b95-25a0-46c3-be2f-3810152971d6",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"language_name = \"RollingMeow\""
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "e0f8176e-6722-4306-8b6e-8d223f113b0e",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"retriever = SimpleExampleRetriever.from_programs(\n",
" language_name,\n",
" [\n",
" (\"2**5\", r_(c_(math.pow, 2, 5))),\n",
" (\"take the log base 2 of 2\", r_(c_(math.log2, 2))),\n",
" ],\n",
" AstPrinter(),\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "d99b3cd7-3ecc-4c8b-b99c-630b2b717b3a",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"text/plain": [
"[('2**5', '```RollingMeow\\nvar result = pow(2, 5)\\n```'),\n",
" ('take the log base 2 of 2', '```RollingMeow\\nvar result = log2(2)\\n```')]"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"retriever.retrieve(\"[ignore]\")"
]
},
{
"cell_type": "markdown",
"id": "ab4d3808-50ab-4ade-aab7-677e23624cab",
"metadata": {},
"source": [
"## Declare the chain"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "26874d63-5460-42c7-a346-5d70691ba02e",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"from kork import (\n",
" CodeChain,\n",
" ast,\n",
" AstPrinter,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "5e2df25f-169d-4961-9fd9-4b10bb9f07c9",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"chain = CodeChain.from_defaults(\n",
" llm=ToyChatModel(response=\"MEOW MEOW MEOW MEOW\"), # The LLM to use\n",
" ast_printer=AstPrinter(), # Knows how to print the AST\n",
" examples=retriever, # Example programs\n",
" context=simple_context_retriever,\n",
" language_name=language_name,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "7d92f7b6-d1f5-41b5-826b-b338edc3831e",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"_, few_shot_prompt = chain.prepare_context(query=\"hello\")"
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "06b5f073-bb8c-40a5-a816-52075417dee2",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"You are programming in a language called \"RollingMeow\".\n",
"\n",
"You are an expert programmer and must follow the instructions below exactly.\n",
"\n",
"Your goal is to translate a user query into a corresponding and valid RollingMeow\n",
"program.\n",
"\n",
"You have access to the following external functions:\n",
"\n",
"```RollingMeow\n",
"\n",
"extern fn pow(x: Any, y: Any) -> Any // Return x**y (x to the power of y).\n",
"extern fn log2(x: Any) -> Any // Return the base 2 logarithm of x.\n",
"extern fn log10(x: Any) -> Any // Return the base 10 logarithm of x.\n",
"```\n",
"\n",
"\n",
"Do not assume that any other functions except for the ones listed above exist.\n",
"\n",
"Wrap the program in ```RollingMeow and ``` to indicate the start and end of\n",
"the program.\n",
"\n",
"Store the solution to the query in a variable called \"result\".\n",
"\n",
"Here is a sample valid program:\n",
"\n",
"```RollingMeow\n",
"var x = 1 # Assign 1 to the variable x\n",
"var result = 1 + 2 # Calculate the sum of 1 + 2 and assign to result\n",
"var result = x # Assign the value of x to result\n",
"```\n",
"\n",
"Guidelines:\n",
"- Do not use operators, instead invoke appropriate external functions.\n",
"- Do not declare functions, do not use loops, do not use conditionals.\n",
"- Solve the problem only using variable declarations and function invocations.\n",
"\n",
"Begin!\n",
"Input: 2**5\n",
"Output: ```RollingMeow\n",
"var result = pow(2, 5)\n",
"```\n",
"\n",
"Input: take the log base 2 of 2\n",
"Output: ```RollingMeow\n",
"var result = log2(2)\n",
"```\n",
"\n",
"Input: [user input]\n",
"Output:\n"
]
}
],
"source": [
"print(few_shot_prompt.format_prompt(query=\"[user input]\").to_string())"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.2"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
kork/__init__.py
+31
View File
@@ -0,0 +1,31 @@
from kork import ast
from kork.ast_printer import AstPrinter
from kork.chain import CodeChain
from kork.environment import Environment
from kork.examples import (
AbstractExampleRetriever,
SimpleExampleRetriever,
c_,
format_examples,
r_,
)
from kork.exceptions import KorkException
from kork.interpreter import InterpreterResult, run_interpreter
from kork.retrieval import AbstractContextRetriever, SimpleContextRetriever
__all__ = [
"CodeChain",
"ast",
"AstPrinter",
"KorkException",
"InterpreterResult",
"run_interpreter",
"AbstractContextRetriever",
"SimpleContextRetriever",
"Environment",
"c_",
"r_",
"format_examples",
"AbstractExampleRetriever",
"SimpleExampleRetriever",
]
kork/ast.py
+174
View File
@@ -0,0 +1,174 @@
"""The AST for the language.
The AST is a bit messy right now in terms of what's a statement vs. an expression,
and will likely need to be cleaned up a bit in the near future.
"""
from __future__ import annotations
import abc
import dataclasses
from typing import Any, Callable, Optional, Sequence, TypeVar, Union
T = TypeVar("T")
class Visitor(abc.ABC):
"""Abstract visitor interface."""
def _to_snake_case(name: str) -> str:
"""Convert a name into snake_case."""
snake_case = ""
for i, char in enumerate(name):
if char.isupper() and i != 0:
snake_case += "_" + char.lower()
else:
snake_case += char.lower()
return snake_case
@dataclasses.dataclass(frozen=True)
class Expr(abc.ABC):
"""Abstract expression."""
def accept(self, visitor: Visitor) -> Any:
"""Accept implementation for a visitor."""
return getattr(visitor, f"visit_{_to_snake_case(self.__class__.__name__)}")(
self
)
@dataclasses.dataclass(frozen=True)
class Literal(Expr):
"""A literal expression."""
value: Union[float, int, bool, None, str]
@dataclasses.dataclass(frozen=True)
class Grouping(Expr):
"""A grouping expression."""
expr: Expr
@dataclasses.dataclass(frozen=True)
class Unary(Expr):
"""A unary expression."""
operator: str
right: Expr
@dataclasses.dataclass(frozen=True)
class Binary(Expr):
"""A binary expression."""
left: Expr
operator: str
right: Expr
@dataclasses.dataclass(frozen=True)
class Assign(Expr):
"""Assignment statement for a variable."""
name: str
value: Expr
@dataclasses.dataclass(frozen=True)
class Variable(Expr):
"""Variable reference."""
name: str
@dataclasses.dataclass(frozen=True)
class List_(Expr):
"""List literal."""
elements: Sequence[Expr]
@dataclasses.dataclass(frozen=True)
class Stmt(abc.ABC):
"""Abstract statement."""
def accept(self, visitor: Visitor) -> Any:
"""Accept implementation for a visitor."""
return getattr(visitor, f"visit_{_to_snake_case(self.__class__.__name__)}")(
self
)
@dataclasses.dataclass(frozen=True)
class Param(Stmt):
"""Represent a function parameter."""
name: str
type_: str
@dataclasses.dataclass(frozen=True)
class ParamList(Stmt):
"""Represent a list of function parameters."""
params: Sequence[Param]
@dataclasses.dataclass(frozen=True)
class FunctionDef(Stmt):
"""Represent a function definition with an implementation."""
name: str
params: ParamList
body: Sequence[Union[Stmt, Expr]]
# This is just an annotation so it's a string.
return_type: str
@dataclasses.dataclass(frozen=True)
class ExternFunctionDef(Stmt):
"""External function definition."""
name: str
params: ParamList
# This is just an annotation so it's a string.
return_type: str
implementation: Optional[Callable] = None
doc_string: str = ""
def add_implementation(self, implementation: Callable) -> "ExternFunctionDef":
"""Add an implementation to an external function definition."""
if self.implementation is not None:
raise ValueError("Cannot add implementation to an already implemented func")
return ExternFunctionDef(
name=self.name,
params=self.params,
return_type=self.return_type,
implementation=implementation,
)
@dataclasses.dataclass(frozen=True)
class FunctionCall(Expr):
"""Represent a function call."""
name: str # Name of function being invoked
args: Sequence[Expr] # Arguments to the function call
@dataclasses.dataclass(frozen=True)
class VarDecl(Stmt):
"""Represent a variable declaration."""
name: str # Name of variable
value: Expr # Expression to assign to the variable
@dataclasses.dataclass(frozen=True)
class Program(Stmt):
"""Represent a program."""
stmts: Sequence[Union[Stmt, Expr]]
kork/ast_printer.py
+116
View File
@@ -0,0 +1,116 @@
import abc
from typing import Union
from kork import ast
class AbstractAstPrinter(ast.Visitor, abc.ABC):
@abc.abstractmethod
def visit(self, element: Union[ast.Stmt, ast.Expr]) -> str:
"""Entry-point for printing the AST."""
class AstPrinter(AbstractAstPrinter):
"""Default AST Printer implementation."""
def visit(self, element: Union[ast.Stmt, ast.Expr]) -> str:
"""Entry-point for printing the AST."""
return element.accept(self)
def visit_program(self, program: ast.Program) -> str:
"""Print a program."""
return "\n".join(stmt.accept(self) for stmt in program.stmts)
def visit_extern_function_def(
self, extern_function_def: ast.ExternFunctionDef
) -> str:
"""Print an extern function definition."""
code = [
f"extern fn {extern_function_def.name}",
"(",
extern_function_def.params.accept(self),
")",
" -> ",
extern_function_def.return_type,
]
if extern_function_def.doc_string:
# TODO: This is hacky for prototyping, will need to have a better
# way to represent the doc string.
doc_string = extern_function_def.doc_string.strip().split("\n")[0]
code.append(f" // {doc_string}")
return "".join(code)
def visit_function_call(self, call: ast.FunctionCall) -> str:
"""Print a function call."""
return f"{call.name}({', '.join(arg.accept(self) for arg in call.args)})"
def visit_function_def(self, function_def: ast.FunctionDef) -> str:
"""Print a function definition."""
signature = [
f"fn {function_def.name}",
"(",
function_def.params.accept(self),
")",
" -> ",
function_def.return_type,
]
code = [" {"]
for stmt in function_def.body:
code.append(stmt.accept(self))
code.append("}")
return "".join(signature) + "\n".join(code)
def visit_variable(self, variable: ast.Variable) -> str:
"""Print a variable."""
return variable.name
def visit_var_decl(self, var_decl: ast.VarDecl) -> str:
"""Print a variable declaration."""
return f"var {var_decl.name} = {var_decl.value.accept(self)}"
def visit_assign(self, assign: ast.Assign) -> str:
"""Print an assignment."""
return f"{assign.name} = {assign.value.accept(self)}"
def visit_literal(self, literal: ast.Literal) -> str:
"""Print a literal."""
value = literal.value
if isinstance(value, type(None)):
return "null"
elif isinstance(value, bool):
return str(value).lower()
elif isinstance(value, str):
return f'"{value}"'
elif isinstance(value, (float, int)):
return str(value)
else:
raise AssertionError(f"Unknown literal type: {type(value)}")
def visit_param_list(self, param_list: ast.ParamList) -> str:
"""Print a parameter list."""
return ", ".join(param.accept(self) for param in param_list.params)
def visit_param(self, param: ast.Param) -> str:
"""Print a parameter."""
return f"{param.name}: {param.type_}"
def visit_list_(self, list_: ast.List_) -> str:
"""Print a list."""
return f"[{', '.join(element.accept(self) for element in list_.elements)}]"
def visit_unary(self, unary: ast.Unary) -> str:
"""Visit a unary expression."""
return f"{unary.operator}{unary.right.accept(self)}"
def visit_grouping(self, grouping: ast.Grouping) -> str:
return f"({grouping.expr.accept(self)})"
def visit_binary(self, binary: ast.Binary) -> str:
"""Print a binary expression."""
return (
f"{binary.left.accept(self)} "
f"{binary.operator} {binary.right.accept(self)}"
)
kork/chain.py
+337
View File
@@ -0,0 +1,337 @@
"""Implementation of a programming chain."""
from __future__ import annotations
from typing import (
Callable,
Dict,
List,
Optional,
Sequence,
Tuple,
TypedDict,
Union,
cast,
)
from langchain import LLMChain
from langchain.chains.base import Chain
from langchain.prompts import PromptTemplate
try:
from langchain.base_language import BaseLanguageModel
except ImportError:
from langchain.schema import BaseLanguageModel
from pydantic import Extra
from kork import ast
from kork.ast_printer import AbstractAstPrinter, AstPrinter
from kork.environment import Environment, create_environment
from kork.examples import (
AbstractExampleRetriever,
InputFormatter,
SimpleExampleRetriever,
format_examples,
format_text,
)
from kork.exceptions import LLMParseException
from kork.interpreter import InterpreterResult, run_interpreter
from kork.prompt_adapter import FewShotTemplate
from kork.retrieval import AbstractContextRetriever, FuncLike, SimpleContextRetriever
from kork.utils import unwrap_code
_INSTRUCTION_PROMPT = """\
You are programming in a language called "{language_name}".
You are an expert programmer and must follow the instructions below exactly.
Your goal is to translate a user query into a corresponding and valid {language_name}
program.
{external_functions_block}
Do not assume that any other functions except for the ones listed above exist.
Wrap the program in ```{language_name} and ``` to indicate the start and end of
the program.
Store the solution to the query in a variable called "result".
Here is a sample valid program:
```{language_name}
var x = 1 # Assign 1 to the variable x
var result = 1 + 2 # Calculate the sum of 1 + 2 and assign to result
var result = x # Assign the value of x to result
```
Guidelines:
- Do not use operators, instead invoke appropriate external functions.
- Do not declare functions, do not use loops, do not use conditionals.
- Solve the problem only using variable declarations and function invocations.
Begin!
"""
# PUBLIC API
DEFAULT_INSTRUCTION_PROMPT = PromptTemplate(
input_variables=[
"language_name",
"external_functions_block",
],
template=_INSTRUCTION_PROMPT,
)
class CodeResult(TypedDict):
"""Result of a code chain."""
errors: Sequence[Exception]
raw: str
code: str
environment: Optional[Environment] # Returned on success
# PUBLIC API
ExampleTuple = Tuple[str, ast.Program]
class CodeChain(Chain):
"""A coding chain."""
llm: BaseLanguageModel
"""The language model to use for the coding chain."""
interpreter: Callable[[str, Environment], InterpreterResult]
"""The interpreter to use for the coding chain."""
ast_printer: AbstractAstPrinter
"""The AST printer to use for the coding chain."""
context_retriever: Optional[AbstractContextRetriever] = None
"""Context to inject into the environment and prompt.
At the moment, this context is limited to external functions.
"""
example_retriever: Optional[AbstractExampleRetriever] = None
"""Examples that should be added to the prompt."""
instruction_template: PromptTemplate = DEFAULT_INSTRUCTION_PROMPT
"""Template for the instruction prompt."""
input_key: str = "query"
# Smirking kat language
language_name: str = "😼"
input_formatter: InputFormatter = "triple_quotes"
"""The smirking kat programming language; aka Kork; aka 😼"""
class Config:
"""Configuration for this pydantic object."""
extra = Extra.allow
arbitrary_types_allowed = True
@property
def input_keys(self) -> List[str]:
"""Will be whatever keys the prompt expects.
:meta private:
"""
return [self.input_key]
@property
def output_keys(self) -> List[str]:
"""Output keys to return.
:meta private:
"""
return ["raw", "errors", "code", "environment"]
def prepare_context(self, query: str) -> Tuple[Environment, FewShotTemplate]:
"""Get the pre-populated environment and the few shot template.
Args:
query: The query to prepare the context for.
Returns:
The prepopulated environment and a pre-formatted few shot template.
"""
if self.context_retriever:
external_function_definitions = self.context_retriever.retrieve(query)
else:
external_function_definitions = []
environment = create_environment(external_function_definitions)
if self.example_retriever:
formatted_examples = self.example_retriever.retrieve(query)
else:
formatted_examples = []
external_functions = "\n".join(
[
self.ast_printer.visit(func_declaration)
for func_declaration in environment.list_external_functions()
]
)
input_variables = self.instruction_template.input_variables
formatting_variables = {}
if "language_name" in input_variables:
formatting_variables["language_name"] = self.language_name
if external_functions:
external_functions_block = (
"You have access to the following external functions:\n\n"
f"```{self.language_name}\n\n"
f"{external_functions}\n"
"```\n"
)
else:
external_functions_block = "Do not assume any functions exist.\n"
if "external_functions_block" in input_variables:
formatting_variables["external_functions_block"] = external_functions_block
prompt = self.instruction_template.partial(
**formatting_variables,
)
return environment, FewShotTemplate(
instruction_template=prompt,
examples=formatted_examples,
input_variables=["query"],
)
def _call(self, inputs: Dict[str, str]) -> CodeResult: # type: ignore
"""Call the chain."""
# Note not using original query!!
# We remove the leading and trailing whitespace from the query
# to make things a bit more robust.
query = inputs[self.input_key].strip()
environment, few_shot_template = self.prepare_context(query)
chain = LLMChain(
prompt=few_shot_template,
llm=self.llm,
)
formatted_query = format_text(query, self.input_formatter)
llm_output = cast(str, chain.predict_and_parse(query=formatted_query))
code = unwrap_code(self.language_name, llm_output)
if not code:
return {
"errors": [
LLMParseException(
"Could not find code block. Please make sure the code "
"is wrapped in a code block."
)
],
"raw": llm_output,
"code": "",
"environment": None,
}
interpreter_result = self.interpreter(code, environment)
if interpreter_result["errors"]:
return {
"errors": interpreter_result["errors"],
"raw": llm_output,
"code": code,
"environment": None,
}
return {
"errors": [],
"raw": llm_output,
"code": code,
"environment": interpreter_result["environment"],
}
@property
def _chain_type(self) -> str:
"""Used for serialization."""
return "code_chain"
@classmethod
def from_defaults(
cls,
*,
llm: BaseLanguageModel,
interpreter: Callable[[str, Environment], InterpreterResult] = run_interpreter,
ast_printer: AbstractAstPrinter = AstPrinter(),
context: Union[AbstractContextRetriever, Sequence[FuncLike], None] = None,
examples: Union[AbstractExampleRetriever, Sequence[ExampleTuple], None] = None,
instruction_template: PromptTemplate = DEFAULT_INSTRUCTION_PROMPT,
input_key: str = "query",
# Smirking kat language
language_name: str = "😼",
input_formatter: InputFormatter = "markdown_text",
) -> CodeChain:
"""Create a code chain from pre-configured defaults.
Args:
llm: The language model to use for coding
interpreter: The interpreter that will be used to execute the program
ast_printer: An ast printer that can print the AST as a text string
context: Either a list of functions or a context retriever.
The list of functions can be a mixture of python
callables and Kork external functions.
All python functions will be converted into Kork
external functions, and everything passed into
the default context retriever.
examples: An example retriever or a list of examples.
If a list of examples, a simple example retriever
will be created.
instruction_template: Use to customize the instruction components of
the prompt.
language_name: The language name to use for the prompt.
input_formatter: A formatting that will be applied to the input part
of each example tuple, if passing in a list of examples.
input_key: The input key to use.
Returns:
A code chain
"""
if isinstance(examples, (AbstractExampleRetriever, type(None))):
_example_retriever = examples
elif isinstance(examples, Sequence):
formatted_examples = format_examples(
language_name,
examples,
ast_printer,
input_formatter=input_formatter,
)
_example_retriever = SimpleExampleRetriever(examples=formatted_examples)
else:
raise TypeError(
f"example_retriever must be of type `AbstractExampleRetriever` or "
f"`Sequence[ExampleTuple]`. Got {type(examples)} instead."
)
if isinstance(context, (AbstractContextRetriever, type(None))):
_context_retriever = context
elif isinstance(context, Sequence):
_context_retriever = SimpleContextRetriever.from_functions(context)
else:
raise TypeError(
f"context_retriever must be of type `AbstractContextRetriever` or "
f"`Sequence[FunctionDefinition]`. Got {type(context)} "
f"instead."
)
return cls(
llm=llm,
interpreter=interpreter,
ast_printer=ast_printer,
context_retriever=_context_retriever,
example_retriever=_example_retriever,
instruction_template=instruction_template,
input_key=input_key,
language_name=language_name,
input_formatter=input_formatter,
)
kork/display.py
+92
View File
@@ -0,0 +1,92 @@
"""Utils for displaying chain results in a notebook."""
import math
from html import escape
from typing import Any, Optional, Sequence, TypedDict, Union
from kork.chain import CodeResult
class HtmlResult(TypedDict):
"""A result that can be displayed in a notebook."""
query: str
errors: str
raw: str
code: str
result: str
expected: str
correct: str
def as_html_dict(
code_result: CodeResult,
query: Optional[str] = None,
expected_answer: Optional[Any] = None,
) -> HtmlResult:
"""Use to generate a dict that can be easily displayed in an IPython notebook."""
from IPython.display import Code
code = code_result["code"].strip()
if code:
_code = Code(data=code)._repr_html_().strip().replace("\n", "<br/>")
else:
_code = code
env = code_result["environment"]
result = env.variables.get("result", "") if env else ""
if isinstance(expected_answer, (float, int)) and isinstance(result, (float, int)):
correct = math.isclose(expected_answer, result)
else:
correct = bool(expected_answer == result)
query = query or code_result.get("query", "") # type: ignore
return {
"query": escape(str(query)),
"errors": escape(str(code_result["errors"])),
"raw": escape(str(code_result["raw"])).replace("\n", "<br/>"),
"code": _code,
"result": escape(str(result).strip()),
"expected": escape(str(expected_answer)),
"correct": "✅" if correct else "⛔",
}
def display_html_results(
html_results: Union[Sequence[HtmlResult], HtmlResult],
columns: Optional[Sequence[str]] = None,
) -> Any:
"""Display a sequence of HTML results as a table."""
import pandas as pd
from IPython import display
_results = html_results if isinstance(html_results, list) else [html_results]
_columns = (
columns
if columns
else ["query", "code", "result", "expected", "correct", "errors", "raw"]
)
df = pd.DataFrame(_results, columns=_columns)
if "query" in df.columns:
df["query"] = df["query"].str.wrap(30)
df["query"] = df["query"].str.replace("\n", "<br/>")
df = df.style.set_properties(**{"text-align": "left"})
return display.HTML(df.to_html(escape=False))
def display_code_result(
code_result: CodeResult,
query: Optional[str] = None,
expected_answer: Optional[Any] = None,
columns: Optional[Sequence[str]] = None,
) -> Any:
"""Display a code result as a table."""
return display_html_results(
as_html_dict(code_result, query=query, expected_answer=expected_answer),
columns=columns,
)
kork/environment.py
+78
View File
@@ -0,0 +1,78 @@
from __future__ import annotations
import copy
import dataclasses
from dataclasses import field
from typing import Dict, List, Optional, Sequence, Union
from kork import ast
from kork.exceptions import KorkRunTimeException
SymbolValue = Union[float, int, str, list, ast.FunctionDef, ast.ExternFunctionDef]
@dataclasses.dataclass
class Environment:
"""Environment for storing variables and function definitions."""
parent: Optional[Environment] = None
variables: Dict[str, SymbolValue] = field(default_factory=dict)
def get_symbol(self, name: str) -> SymbolValue:
"""Get a symbol from the environment.
Args:
name: symbol name to lookup
Returns:
SymbolValue for the given symbol if it exists, otherwise raises an error
"""
if name in self.variables:
return self.variables[name]
elif self.parent:
return self.parent.get_symbol(name)
else:
raise KorkRunTimeException(f"Variable `{name}` not found")
def set_symbol(self, name: str, value: SymbolValue) -> SymbolValue:
"""Set a symbol in the environment."""
# TODO: We need to determine whether want the variable to be
# declared prior to be being set.
# At the moment, this allows for `x = 1` without having to declare the
# variable.
self.variables[name] = value
return value
def clone(self) -> Environment:
"""Clone a root level environment, with deep copy on variables.
Cloning functionality is provided because the environment is
a mutable variable.
Biggest danger is mutating the variables dict which the caller
may retain a reference to.
TODO(Eugene): Refactor to make this less dangerous.
"""
if self.parent is not None:
raise AssertionError("Cannot clone an environment with a parent.")
return Environment(parent=None, variables=copy.deepcopy(self.variables))
def list_external_functions(self) -> List[ast.ExternFunctionDef]:
"""Get a list of the externally defined foreign functions."""
return [
func
for func in self.variables.values()
if isinstance(func, ast.ExternFunctionDef)
]
def create_environment(
extern_function_defs: Sequence[ast.ExternFunctionDef],
) -> Environment:
"""Create a new environment with pre-populated state."""
variables: Dict[str, SymbolValue] = {
func.name: func for func in extern_function_defs
}
environment = Environment(parent=None, variables=variables)
return environment
kork/examples.py
+160
View File
@@ -0,0 +1,160 @@
"""Interface to specify kork examples easily."""
import abc
from typing import Any, Callable, List, Literal, Sequence, Tuple, Union
from kork import ast
from kork.ast_printer import AbstractAstPrinter
from kork.foreign_funcs import to_kork_function_call
def _add_result_variable(expr: ast.Expr) -> ast.Program:
"""Assembles a solution from a program."""
stmt = ast.VarDecl(name="result", value=expr)
return ast.Program(stmts=[stmt])
# PUBLIC API
# Use to denote different types of formatters for the input.
InputFormatter = Union[
Literal["text_prefix"],
Literal["triple_quotes"],
Literal["markdown_text"],
None,
Callable[[str], str],
]
def format_text(text: str, input_formatter: InputFormatter = None) -> str:
"""An encoder for the input text.
Args:
text: the text to encode
input_formatter: the formatter to use for the input
* None: use for single sentences or single paragraphs, no formatting
* triple_quotes: surround input with \"\"\", use for long text
* text_prefix: same as triple_quote but with `TEXT: ` prefix
* Callable: user provided function
Returns:
The encoded text if it was encoded
"""
if input_formatter == "text_prefix":
return 'Text: """\n' + text + '\n"""'
elif input_formatter == "triple_quotes":
return '"""\n' + text + '\n"""'
elif input_formatter == "markdown_text":
return "```text\n" + text + "\n```"
elif input_formatter is None:
return text
else:
raise NotImplementedError(
f'No support for input encoding "{input_formatter}". '
' Use one of "long_text" or None.'
)
# Alias to create a kork function call
def c_(name: Callable, *args: Any) -> ast.FunctionCall:
"""Create a kork function call."""
return to_kork_function_call(name, *args)
# Alias to assign the expression to a variable called `result`
def r_(expr: ast.Expr) -> ast.Program:
"""Assign last program expression to a result variable."""
return _add_result_variable(expr)
def format_examples(
language_name: str,
examples: Sequence[Tuple[str, ast.Program]],
ast_printer: AbstractAstPrinter,
input_formatter: InputFormatter = None,
) -> List[Tuple[str, str]]:
"""Format examples."""
formatted_examples = []
for example_input, desired_output in examples:
formatted_input = format_text(example_input, input_formatter=input_formatter)
formatted_examples.append(
(
formatted_input,
f"```{language_name}\n{ast_printer.visit(desired_output)}\n```",
)
)
return formatted_examples
class AbstractExampleRetriever(abc.ABC):
"""Abstract interface for an example retriever.
An example interface must implement the `retrieve` method which
returns a list of relevant examples based on the given query.
"""
@abc.abstractmethod
def retrieve(self, query: str) -> List[Tuple[str, str]]:
"""Retrieve examples."""
class SimpleExampleRetriever(AbstractExampleRetriever):
"""Simple example retriever.
Simple example that returns the examples it was initialized with.
Supports initialization from a list of programs.
Example:
.. code-block:: python
from kork import SimpleExampleRetriever, AstPrinter, c_, r_
simple_example_retriever = SimpleExampleRetriever.from_programs(
language_name="kork",
examples=[
("add 1 2", r_(c_(add, 1, 2))),
("add 1 2 3", r_(c_(add, 1, 2, 3))),
],
ast_printer=AstPrinter(),
)
examples = simple_example_retriever.retrieve("add 1 2")
"""
def __init__(
self,
examples: Sequence[Tuple[str, str]],
) -> None:
"""Initialize the retriever with a list of examples."""
self._examples = examples
def retrieve(self, query: str) -> List[Tuple[str, str]]:
"""Retrieve examples that best match the given query."""
return list(self._examples)
@classmethod
def from_programs(
cls,
language_name: str,
examples: Sequence[Tuple[str, ast.Program]],
ast_printer: AbstractAstPrinter,
) -> "SimpleExampleRetriever":
"""Create a simple example retriever from programs.
Args:
language_name: The language name to use for the markdown code blocks.
examples: A sequence of tuples of the form (query, desired program).
ast_printer: The ast printer to use to format the desired output.
Returns:
A simple example retriever.
"""
return cls(
examples=format_examples(
language_name=language_name,
examples=examples,
ast_printer=ast_printer,
)
)
kork/exceptions.py
+21
View File
@@ -0,0 +1,21 @@
"""Definitions for custom Kork exceptions."""
class KorkException(Exception):
"""Generic Kork exception."""
class LLMParseException(KorkException):
"""Failed to parse LLM output."""
class KorkSyntaxException(KorkException):
"""Exceptions raised during syntax parsing."""
class KorkInterpreterException(KorkException):
"""Exceptions raised during interpretation."""
class KorkRunTimeException(KorkInterpreterException):
"""An exception that is raised during Kork interpreter runtime."""
kork/foreign_funcs.py
+172
View File
@@ -0,0 +1,172 @@
"""API to import foreign functions."""
import inspect
import sys
import types
import typing
from typing import Any, Callable, List, Mapping, Tuple, TypedDict
from kork import ast
PY_VERSION = (sys.version_info.major, sys.version_info.minor)
class FunctionInfo(TypedDict):
"""Information about a function."""
name: str
args: List[Tuple[str, Any, Any]]
return_type: Any
docstring: str
def _type_repr(obj: Any) -> str:
"""Return the repr() of an object, special-casing types (internal helper).
If obj is a type, we return a shorter version than the default
type.__repr__, based on the module and qualified name, which is
typically enough to uniquely identify a type. For everything
else, we fall back on repr(obj).
"""
if PY_VERSION >= (3, 9):
# Ignoring type for passing mypy testing with Python 3.8
if isinstance(obj, types.GenericAlias): # type: ignore[attr-defined]
return repr(obj)
if isinstance(obj, type):
if obj.__module__ == "builtins":
return obj.__qualname__
return f"{obj.__module__}.{obj.__qualname__}"
if obj is ...:
return "..."
if isinstance(obj, types.FunctionType):
return obj.__name__
return repr(obj)
def _type_repr_wrapper(obj: Any) -> str:
"""A hack to remove all references to `typing.` in type strings.
There's likely a much better solution, but this will do for prototyping.
Args:
obj (Any): The object to get the type string of.
Returns:
str: The type string of the object.
"""
return _type_repr(obj).replace("typing.", "")
def _get_stringified_type(type_hints: Mapping[str, Any], arg_name: str) -> str:
"""Get string version of an argument type from a mapping of type hints.
Args:
type_hints (Mapping[str, Any]): Mapping of argument names to types.
arg_name (str): Name of the argument or `return` for the return type.
Returns:
str: String representation of the type if it exists, otherwise Any.
"""
if arg_name in type_hints:
return _type_repr_wrapper(type_hints[arg_name])
else:
return "Any"
def _get_function_info(func: Callable) -> FunctionInfo:
"""Extract information about a function.
Args:
func (callable): The function to get information about.
Returns:
A dictionary with the following keys:
- name: The name of the function.
- args: A list of tuples containing the name and type of each argument.
- return_type: The type of the value returned by the function.
- docstring: The docstring of the function.
"""
# Get the name of the function
name = func.__name__
try:
# Get information about the function's arguments
type_hints = typing.get_type_hints(func)
except ValueError:
type_hints = {}
try:
signature = inspect.signature(func)
args = []
for arg_name, param in signature.parameters.items():
arg_type = _get_stringified_type(type_hints, arg_name)
arg_default = (
param.default if param.default != inspect.Parameter.empty else None
)
args.append((arg_name, arg_type, arg_default))
# Get the return type of the function
return_type = _get_stringified_type(type_hints, "return")
except ValueError: # TODO: Exception is too broad, limit to correct scope
return_type = "Any"
args = []
# Get the docstring of the function
docstring = func.__doc__
# Construct and return a dictionary with the function information
return {
"name": name,
"args": args,
"return_type": return_type,
"docstring": docstring.strip() if docstring else "",
}
def _convert_to_kork_expr(value: Any) -> ast.Expr:
"""Convert mixed-expression value to a kork expression.
This function converts a potentially mixed python, kork expression
into a pure kork expression.
Args:
value (Any): The value to convert.
Returns:
An `kork` expression.
"""
if isinstance(value, ast.Expr):
# If it's already a kork expression, return as is
return value
elif isinstance(value, (int, float, bool, type(None), str)):
return ast.Literal(value=value)
elif isinstance(value, (list, tuple)):
return ast.List_(elements=[_convert_to_kork_expr(v) for v in value])
else:
raise ValueError(f"Cannot convert {value} to an `kork` expression.")
# PUBLIC API
def to_extern_func_def(func: Callable) -> ast.ExternFunctionDef:
"""Convert a python function to a kork external function definition."""
func_info = _get_function_info(func)
params = [
ast.Param(name=arg_name, type_=arg_type)
for arg_name, arg_type, _ in func_info["args"]
]
return ast.ExternFunctionDef(
name=func_info["name"],
params=ast.ParamList(params=params),
return_type=func_info["return_type"],
implementation=func,
doc_string=func_info["docstring"],
)
def to_kork_function_call(func: Callable, *args: Any) -> ast.FunctionCall:
"""Convert a python function call to a kork function call."""
return ast.FunctionCall(
name=func.__name__,
args=[_convert_to_kork_expr(arg) for arg in args],
)
kork/interpreter.py
+167
View File
@@ -0,0 +1,167 @@
from typing import Any, Optional, Sequence, TypedDict, Union
from lark.exceptions import LarkError
from kork import ast
from kork.environment import Environment
from kork.exceptions import KorkRunTimeException
# TODO: Determine why mypy is not recognizing the import.
from kork.parser import parse # type: ignore[attr-defined]
class Interpreter(ast.Visitor):
"""Kork's default interpreter."""
def __init__(self, program: ast.Program, environment: Environment) -> None:
"""Initialize the interpreter."""
self.program = program
self.environment = environment
self.last_expression_value: Any = None
def intepret(self) -> None:
"""Interpret the program."""
self.program.accept(self)
def visit_program(self, program: ast.Program) -> None:
"""Visit a program."""
for stmt in program.stmts:
stmt.accept(self)
def visit_extern_function_def(
self, extern_function_def: ast.ExternFunctionDef
) -> None:
"""Visit an external function definition."""
self.environment.variables[extern_function_def.name] = extern_function_def
def visit_function_call(self, call: ast.FunctionCall) -> Any:
"""Visit a function call."""
function = self.environment.get_symbol(call.name)
args = [arg.accept(self) for arg in call.args]
if isinstance(function, ast.ExternFunctionDef):
if function.implementation is None:
raise KorkRunTimeException("External function has not been linked yet.")
try:
return function.implementation(*args)
except Exception as e:
raise KorkRunTimeException(
"Encountered an exception while invoking an external function "
f"`{call.name}` with arguments: `{args}`."
f"Details: {e}"
)
elif isinstance(function, ast.FunctionDef):
raise KorkRunTimeException("Function definition is unimplemented.")
else:
raise KorkRunTimeException(f"Cannot call {call.name}.")
def visit_function_def(self, function_def: ast.FunctionDef) -> None:
"""Visit a function definition."""
self.environment.variables[function_def.name] = function_def
def visit_variable(self, variable: ast.Variable) -> Any:
"""Visit a variable."""
return self.environment.get_symbol(variable.name)
def visit_var_decl(self, var_decl: ast.VarDecl) -> None:
"""Visit a variable declaration."""
self.environment.variables[var_decl.name] = var_decl.value.accept(self)
def visit_unary(self, unary: ast.Unary) -> Any:
"""Visit a unary expression."""
operator = unary.operator
right = unary.right.accept(self)
if operator == "-":
return -right
elif operator == "+":
return right
else:
raise ValueError(f"Unknown operator {operator}.")
def visit_binary(self, binary: ast.Binary) -> Any:
"""Visit a binary expression."""
left = binary.left.accept(self)
right = binary.right.accept(self)
operator = binary.operator
if operator == "+":
return left + right
elif operator == "-":
return left - right
elif operator == "*":
return left * right
elif operator == "/":
return left / right
elif operator == "**":
return left**right
elif operator == "%":
return left % right
elif (
operator == "^"
): # Please note that we do not follow python convention here
return left**right # Duplicating `**` operator
else:
raise ValueError(f"Unknown operator {operator}.")
def visit_grouping(self, grouping: ast.Grouping) -> Any:
"""Visit a grouping."""
return grouping.expr.accept(self)
def visit_assign(self, assign: ast.Assign) -> None:
"""Visit an assignment."""
self.environment.variables[assign.name] = assign.value.accept(self)
def visit_literal(self, literal: ast.Literal) -> Union[int, float, bool, str, None]:
"""Visit a number."""
return literal.value
def visit_list_(self, list_: ast.List_) -> list:
"""Visit a list."""
return [element.accept(self) for element in list_.elements]
# PUBLIC API
class InterpreterResult(TypedDict):
"""Use this to return the result of the interpreter."""
environment: Environment
errors: Sequence[Exception]
def run_interpreter(
code: str, environment: Optional[Environment] = None
) -> InterpreterResult:
"""Run the interpreter with the given code.
Args:
code: The code to run
environment: The environment to run the code in
Returns:
the final environment after running the code (this will likely change)
"""
# TODO(Eugene): May want to refactor to avoid cloning here without introducing
# mutability associated bugs.
environment = (
environment.clone() if environment else Environment(parent=None, variables={})
)
try:
program = parse(code)
interpreter = Interpreter(program, environment)
interpreter.intepret()
except (
LarkError,
KorkRunTimeException,
ValueError,
ZeroDivisionError,
TypeError,
) as e:
return InterpreterResult(
environment=environment,
errors=[e],
)
return InterpreterResult(
environment=interpreter.environment,
errors=[],
)
kork/parser.py
+229
View File
@@ -0,0 +1,229 @@
# type:ignore[no-untyped-def]
"""Kork's default AST parser.
Kork uses Lark to parse the AST. The grammar follows closely
the one used in Crafting Interpreters for the Lox Programming Language.
https://craftinginterpreters.com/appendix-i.html#expressions
The grammar and parser were clobbered together in a few hours of work.
If you deeply care about language design, please forgive any transgressions,
and feel free to help improve the grammar/parser/interpreter.
"""
from typing import cast
from lark import Lark, Transformer, v_args
from kork import ast
GRAMMAR = """
program: statement+
statement: function_decl
| extern_function_decl
| var_decl
| expr
function_decl: "fn" CNAME "(" [params] ")" "->" type block
extern_function_decl: "extern" "fn" CNAME "(" [params] ")" "->" type
var_decl: ("let" | "const" | "var") CNAME "=" expr [";"]
params: param ("," param)*
param: CNAME ":" type
type: CNAME
block: "{" [statement+] "}"
?expr: assignment
?assignment: CNAME "=" expr [";"] -> variable_assign
| term
!?term: term "+" factor -> binary
| term "-" factor -> binary
| factor
!?factor: factor "*" unary -> binary
| factor "/" unary -> binary
| factor "%" unary -> binary
| unary
!?unary: "-" unary -> unary_r
| "+" unary -> unary_r
| power
!?power: call "**" unary -> binary
| call "^" unary -> binary
| call
?call: CNAME "(" [args] ")" -> function_call
| primary
?primary: numbers
| list
| string
| "true" -> true
| "false" -> false
| "null" -> null
| "(" expr ")" -> group
| CNAME -> identifier
?numbers: SIGNED_NUMBER ("E"|"e") SIGNED_NUMBER -> scientific
| SIGNED_NUMBER -> number
args: expr ("," expr)*
string: ESCAPED_STRING
list: "[" [args] "]"
%import common.CNAME
%import common.SIGNED_NUMBER
%import common.ESCAPED_STRING
%import common.WS
%ignore WS
%ignore /#.*/
%ignore /\/\/.*/
"""
# mypy: ignore-errors
@v_args(inline=True)
class AstParser(Transformer):
"""An ast parser for the given language."""
def program(self, *items) -> ast.Program:
"""Create a program from the given statements."""
return ast.Program(stmts=items)
def function_decl(self, name, params, return_type, block) -> ast.FunctionDef:
"""Function declaration from the given name, params, return type and block."""
return ast.FunctionDef(
name=str(name),
params=params,
return_type=return_type,
body=block,
)
def extern_function_decl(self, name, params, return_type) -> ast.ExternFunctionDef:
"""Extern function declaration from the given name, params and return type."""
return ast.ExternFunctionDef(
name=str(name),
params=params or ast.ParamList(params=[]),
return_type=return_type,
)
def statement(self, item) -> ast.Stmt:
"""Create a statement from the given item."""
return item
def var_decl(self, name, value) -> ast.VarDecl:
"""Variable declaration from the given name and value."""
return ast.VarDecl(
name=str(name),
value=value,
)
def params(self, *items) -> ast.ParamList:
"""Create a param list from the given items."""
return ast.ParamList(params=items or [])
def unary_r(self, operator, right) -> ast.Unary:
"""Create a unary from the given operator and left."""
op = str(operator)
if op not in {"+", "-"}:
raise ValueError(f"Unknown operator: {op}.")
return ast.Unary(operator=str(operator), right=right)
def binary(self, left, operator, right) -> ast.Binary:
"""Create a binary from the given left, operator and right."""
op = str(operator)
if op not in {"+", "-", "*", "/", "**", "//", "%", "^"}:
raise ValueError(f"Unknown operator: {op}.")
return ast.Binary(left=left, operator=str(operator), right=right)
def param(self, name, type_) -> ast.Param:
"""Create a param from the given name and type."""
return ast.Param(
name=str(name),
type_=type_,
)
def type(self, name):
"""Create a type from the given name."""
return str(name)
def block(self, *statements):
"""Create a block from the given statements."""
return statements or []
def group(self, expr) -> ast.Grouping:
"""Create a group from the given expression."""
return ast.Grouping(expr=expr)
def function_call(self, name, args) -> ast.FunctionCall:
"""Create a function call from the given name and args."""
return ast.FunctionCall(name=str(name), args=args or [])
def identifier(self, name) -> ast.Variable:
"""Create a variable reference from the given name."""
return ast.Variable(
name=str(name),
)
def variable_assign(self, name, value) -> ast.Assign:
"""Create a variable assignment from the given name and value."""
return ast.Assign(
name=str(name),
value=value,
)
def expr(self, item) -> ast.Expr:
"""Create an expression from the given item."""
return item
def args(self, *items):
"""Create an argument list for a function from the given items."""
return items or []
def list(self, items) -> ast.List_:
"""Create a list literal from the given items."""
return ast.List_(
elements=items or [],
)
def true(self) -> ast.Literal:
"""Create a true literal."""
return ast.Literal(value=True)
def false(self) -> ast.Literal:
"""Create a false literal."""
return ast.Literal(value=False)
def null(self) -> ast.Literal:
"""Create a null literal."""
return ast.Literal(value=None)
def number(self, item) -> ast.Literal:
try:
return ast.Literal(value=int(item))
except ValueError:
return ast.Literal(value=float(item))
def scientific(self, coefficient, exponent) -> ast.Literal:
return ast.Literal(value=coefficient * 10**exponent)
def string(self, item) -> ast.Literal:
return ast.Literal(
value=str(item[1:-1]),
)
ast_parser = Lark(GRAMMAR, parser="lalr", transformer=AstParser(), start="program")
# PUBLIC API
def parse(source: str) -> ast.Program:
"""Parse the given source code into an AST."""
return cast(ast.Program, ast_parser.parse(source))
kork/prompt_adapter.py
+96
View File
@@ -0,0 +1,96 @@
"""A prompt adapter to allow working with both regular LLMs and Chat LLMs.
The prompt adapter supports breaking the prompt into:
1) Instruction Section
2) (Optional) Example Section
"""
from typing import Any, Callable, List, Sequence, Tuple
from langchain import BasePromptTemplate, PromptTemplate
from langchain.schema import BaseMessage, HumanMessage, PromptValue, SystemMessage
from pydantic import Extra
class FewShotPromptValue(PromptValue):
"""Integration with langchain prompt format."""
string: Callable[[], str]
messages: Callable[[], List[BaseMessage]]
class Config:
"""Configuration for this pydantic object."""
extra = Extra.forbid
arbitrary_types_allowed = True
def to_string(self) -> str:
"""Format the prompt to a string."""
return self.string()
def to_messages(self) -> List[BaseMessage]:
"""Get materialized messages."""
return self.messages()
class FewShotTemplate(BasePromptTemplate):
"""Code prompt template."""
instruction_template: PromptTemplate
examples: Sequence[Tuple[str, str]] = ()
@property
def _prompt_type(self) -> str:
"""Get the prompt type."""
return "FewShotTemplate"
def format_prompt(self, **kwargs: Any) -> PromptValue:
"""Format the prompt."""
if len(self.input_variables) != 1:
raise AssertionError(
f"Expected 1 input variable, got {len(self.input_variables)}"
)
query_key = self.input_variables[0]
query = kwargs[query_key]
def _lazy_string() -> str:
"""Lazy string."""
return self.to_string(query)
def _lazy_messages() -> List[BaseMessage]:
"""Lazy messages."""
return self.to_messages(query)
return FewShotPromptValue(
string=_lazy_string,
messages=_lazy_messages,
)
def format(self, **kwargs: Any) -> str:
"""Deprecated format method."""
raise NotImplementedError()
def to_string(self, query: str) -> str:
"""Format the prompt to a string."""
instruction_section = self.instruction_template.format()
examples_block = ""
for input_example, output_example in self.examples:
examples_block += f"Input: {input_example}\n\nOutput: {output_example}\n\n"
return instruction_section + examples_block + f"Input: {query}\n\nOutput:"
def to_messages(self, query: str) -> List[BaseMessage]:
"""Get materialized messages."""
instruction_segment = self.instruction_template.format().strip()
messages: List[BaseMessage] = [SystemMessage(content=instruction_segment)]
for input_example, output_example in self.examples:
messages.append(HumanMessage(content=input_example.strip()))
messages.append(SystemMessage(content=output_example.strip()))
messages.append(HumanMessage(content=query.strip()))
return messages
kork/py.typed
View File
kork/retrieval.py
+62
View File
@@ -0,0 +1,62 @@
"""Logic that attempts to surface the most relevant information for writing code."""
from __future__ import annotations
import abc
import dataclasses
from typing import Callable, Sequence, Union
from kork import ast
from kork.foreign_funcs import to_extern_func_def
@dataclasses.dataclass(frozen=True)
class AbstractContextRetriever(abc.ABC):
"""Abstract interface for retrieving programming context."""
@abc.abstractmethod
def retrieve(self, query: str) -> Sequence[ast.ExternFunctionDef]:
"""Retrieve the external function definitions."""
FuncLike = Union[ast.ExternFunctionDef, Callable]
@dataclasses.dataclass(frozen=True)
class SimpleContextRetriever(AbstractContextRetriever):
"""Retrieve information as was provided without any filtering or re-ranking."""
external_functions: Sequence[ast.ExternFunctionDef] = tuple()
def __post_init__(self) -> None:
"""Validate the external functions."""
for func in self.external_functions:
if not isinstance(func, ast.ExternFunctionDef):
raise ValueError(
"SimpleRetriever only accepts ExternFunctionDef objects in "
f"the external_functions list. Got: {func}"
)
def retrieve(self, query: str) -> Sequence[ast.ExternFunctionDef]:
"""Retrieve the external function definitions."""
return self.external_functions
@classmethod
def from_functions(cls, mixed_funcs: Sequence[FuncLike]) -> SimpleContextRetriever:
"""Create a simple retrieval from a sequence of functions.
Args:
mixed_funcs: A sequence of functions or external function definitions.
Returns:
A simple retriever.
"""
external_functions = []
for func in mixed_funcs:
if isinstance(func, ast.ExternFunctionDef):
external_functions.append(func)
else:
external_functions.append(to_extern_func_def(func))
return cls(external_functions=external_functions)
kork/utils.py
+29
View File
@@ -0,0 +1,29 @@
import re
from typing import Optional
def wrap_in_tag(tag_name: str, content: str) -> str:
"""Wrap the content in an HTML style tag."""
return f"<{tag_name}>{content}</{tag_name}>"
def unwrap_tag(tag_name: str, text: str) -> Optional[str]:
"""Extract content located inside a tag."""
pattern = f"<{tag_name}>(.*?)</{tag_name}>"
content = re.compile(pattern, re.DOTALL)
match = content.search(text)
if match:
return match.group(1)
else:
return None
def unwrap_code(language_name: str, text: str) -> Optional[str]:
"""Extract code located inside a code tag."""
pattern = f"```{language_name}\n(.*?)\n```"
content = re.compile(pattern, re.DOTALL)
match = content.search(text)
if match:
return match.group(1)
else:
return None
kork/version.py
+7
View File
@@ -0,0 +1,7 @@
"""Get the version of the package."""
from importlib import metadata
try:
__version__ = metadata.version("kork")
except metadata.PackageNotFoundError:
__version__ = "local"
poetry.lock
Generated
+3946
View File
File diff suppressed because it is too large Load Diff
pyproject.toml
+78
View File
@@ -0,0 +1,78 @@
[tool.poetry]
name = "kork"
version = "0.0.1"
description = "Programming with LLMs"
authors = ["LangChain"]
license = "MIT"
readme = "README.md"
repository = "https://github.com/langchain-ai/kork"
[tool.poetry.dependencies]
python = "^3.8.1"
openai = "^0.27"
langchain = ">=0.0.110"
lark = "^1.1.5"
sphinx-design = "^0.4.1"
[tool.poetry.group.dev.dependencies]
jupyterlab = "^3.6.1"
[tool.poetry.group.test.dependencies]
pytest = "^7.2.1"
black = { version="^23.1.0", extras=["jupyter"] }
poethepoet = "^0.18.1"
ruff = "^0.0.255"
pytest-cov = "^4.0.0"
[tool.poetry.group.docs.dependencies]
nbsphinx = "^0.8.9"
sphinx = "^4.5.0"
sphinx-autobuild = "^2021.3.14"
sphinx_book_theme = "^1.0.0"
myst-nb = "^0.17.1"
linkchecker = "^10.2.1"
toml = "^0.10.2"
sphinx-copybutton = "^0.5.1"
[tool.poetry.group.typing.dependencies]
mypy = "^0.991"
[tool.poetry.group.types.dependencies]
types-toml = "^0.10.8.5"
[tool.poe.tasks]
black = "black"
ruff = "ruff"
pytest.cmd = "py.test --durations=5 -W error::RuntimeWarning --cov --cov-config=.coveragerc --cov-report xml --cov-report term-missing:skip-covered"
mypy = "mypy . --pretty --show-error-codes"
fix = { shell = "poe black . && poe ruff --fix ." }
# Using a --preview feature. Commit your code prior to use this fix.
fix_strings = "black kork --preview"
test = { shell = "poe black . --check --diff && poe ruff . && poe pytest && poe mypy" }
# Use to auto-generate docs
apidoc = "sphinx-apidoc -o docs/source/generated kork"
[tool.ruff]
select = [
"E", # pycodestyle
"F", # pyflakes
"I", # isort
]
# Same as Black.
line-length = 88
[tool.mypy]
disallow_untyped_defs = "True"
ignore_missing_imports = "True"
[tool.coverage.run]
omit = [
"tests/*",
]
[build-system]
requires = ["poetry-core>=1.0.0"]
build-backend = "poetry.core.masonry.api"
tests/__init__.py
View File
tests/test_ast.py
+6
View File
@@ -0,0 +1,6 @@
from kork.ast import _to_snake_case
def test_snake_case() -> None:
assert _to_snake_case("Number") == "number"
assert _to_snake_case("NumberWoof") == "number_woof"
tests/test_chain.py
+114
View File
@@ -0,0 +1,114 @@
from kork import AstPrinter, CodeChain, SimpleContextRetriever, run_interpreter
from kork.examples import SimpleExampleRetriever
from kork.exceptions import KorkRunTimeException, LLMParseException
from kork.parser import parse # type: ignore
from .utils import ToyChatModel
def test_code_chain() -> None:
"""Test the code chain."""
example_retriever = SimpleExampleRetriever(examples=[("Add 1 and 2", "add_(1, 2)")])
llm = ToyChatModel(response="```kork\nvar x = 1\n```")
chain = CodeChain(
llm=llm,
retriever=SimpleContextRetriever(),
interpreter=run_interpreter,
ast_printer=AstPrinter(),
example_retriever=example_retriever,
language_name="kork",
)
response = chain(inputs={"query": "blah"})
# Why does the chain return a `query` key?
assert sorted(response) == ["code", "environment", "errors", "query", "raw"]
env = response.pop("environment")
assert response == {
"query": "blah",
"code": "var x = 1",
"errors": [],
"raw": "```kork\nvar x = 1\n```",
}
assert env.get_symbol("x") == 1
def test_bad_program() -> None:
"""Test the code chain."""
example_retriever = SimpleExampleRetriever(examples=[("Add 1 and 2", "add_(1, 2)")])
llm = ToyChatModel(response="```kork\nINVALID PROGRAM\n```")
chain = CodeChain(
llm=llm,
retriever=SimpleContextRetriever(),
interpreter=run_interpreter,
ast_printer=AstPrinter(),
example_retriever=example_retriever,
language_name="kork",
)
response = chain(inputs={"query": "blah"})
# Why does the chain return a `query` key?
assert sorted(response) == ["code", "environment", "errors", "query", "raw"]
assert response["raw"] == "```kork\nINVALID PROGRAM\n```"
assert response["code"] == "INVALID PROGRAM"
assert response["environment"] is None
assert response["errors"] is not None
assert isinstance(response["errors"][0], KorkRunTimeException)
def test_llm_output_missing_program() -> None:
"""Test the code chain with llm response that's missing code."""
llm = ToyChatModel(response="oops.")
example_retriever = SimpleExampleRetriever(examples=[("Add 1 and 2", "add_(1, 2)")])
chain = CodeChain(
llm=llm,
retriever=SimpleContextRetriever(),
interpreter=run_interpreter,
ast_printer=AstPrinter(),
example_retriever=example_retriever,
language_name="meow",
)
response = chain(inputs={"query": "blah"})
# Why does the chain return a `query` key?
assert sorted(response) == ["code", "environment", "errors", "query", "raw"]
assert response["raw"] == "oops."
assert response["code"] == ""
assert response["environment"] is None
errors = response["errors"]
assert len(errors) == 1
assert isinstance(response["errors"][0], LLMParseException)
def test_from_defaults_instantiation() -> None:
"""Test from default instantiation."""
llm = ToyChatModel(response="```kork\nvar x = 1;\n```")
chain = CodeChain.from_defaults(llm=llm, language_name="kork")
response = chain(inputs={"query": "blah"})
# Why does the chain return a `query` key?
assert sorted(response) == ["code", "environment", "errors", "query", "raw"]
assert response["environment"].get_symbol("x") == 1
# Instantiate with sequence of examples and foreign functions
def _do() -> int:
"""Do something"""
return 1
# Verify that we can instantiate
examples = [("Add 1 and 2", parse("add_(1, 2)"))]
chain2 = CodeChain.from_defaults(
llm=llm, examples=examples, context=[_do], language_name="kork"
)
assert isinstance(chain2.context_retriever, SimpleContextRetriever)
external_func_defs = chain2.context_retriever.retrieve("[anything]")
# Check foreign funcs
assert len(external_func_defs) == 1
assert external_func_defs[0].name == "_do"
# Check examples
assert isinstance(chain2.example_retriever, SimpleExampleRetriever)
assert chain2.example_retriever.retrieve("[anything]") == [
# Please note that an input formatted was applied by default to the example
("```text\nAdd 1 and 2\n```", "```kork\nadd_(1, 2)\n```")
]
tests/test_environment.py
+39
View File
@@ -0,0 +1,39 @@
import pytest
from kork import ast
from kork.environment import Environment
from kork.exceptions import KorkRunTimeException
def test_accessing_variables() -> None:
"""Test the instantiation of the environment."""
env = Environment()
with pytest.raises(KorkRunTimeException):
env.get_symbol("x")
env.set_symbol("x", 1)
assert env.get_symbol("x") == 1
def test_clone() -> None:
"""Test env cloning."""
env = Environment()
env.set_symbol("x", 1)
new_env = env.clone()
env.set_symbol("y", 2)
new_env.set_symbol("z", 3)
assert env.variables == {"x": 1, "y": 2}
assert new_env.variables == {"x": 1, "z": 3}
def test_list_external_functions() -> None:
"""Test list external functions."""
env = Environment()
assert env.list_external_functions() == []
func = ast.ExternFunctionDef(
name="foo", params=ast.ParamList([]), return_type="Any"
)
env.set_symbol("foo", func)
assert env.list_external_functions() == [func]
tests/test_examples.py
+30
View File
@@ -0,0 +1,30 @@
from kork.ast_printer import AstPrinter
from kork.examples import c_, format_examples, r_
def add_(x: int, y: int) -> int:
"""Add two numbers."""
return x + y
def test_format_examples() -> None:
"""Test format examples."""
examples = [
(
"Add 1 and 2",
r_(c_(add_, 1, 2)),
),
]
formatted_examples = format_examples(
language_name="meow",
examples=examples,
ast_printer=AstPrinter(),
)
assert formatted_examples == [
(
"Add 1 and 2",
"```meow\nvar result = add_(1, 2)\n```",
),
]
tests/test_foreign_functions.py
+56
View File
@@ -0,0 +1,56 @@
from typing import Any, Literal, Mapping, Sequence, Union
from kork.ast import ExternFunctionDef, Param, ParamList
from kork.foreign_funcs import to_extern_func_def
# Do not type the function below. We're testing initialization of the retriever.
def foo(): # type: ignore
pass
def bar(x: int) -> int:
"""Add one to x."""
return x + 1
def zoo(x: Mapping[str, Any], q: Sequence[Union[Literal["a"], Literal["b"]]]) -> Any:
"""Add one to x."""
pass
def test_to_extern_func_def() -> None:
"""Test conversion of a python function into an extern function definition."""
assert to_extern_func_def(foo) == ExternFunctionDef(
name="foo",
params=ParamList(params=[]),
# If the source function is un-typed, we assume
# that the return type is Any.
return_type="Any",
implementation=foo,
doc_string="",
)
assert to_extern_func_def(bar) == ExternFunctionDef(
name="bar",
params=ParamList(params=[Param(name="x", type_="int")]),
return_type="int",
implementation=bar,
doc_string="Add one to x.",
)
assert to_extern_func_def(zoo) == ExternFunctionDef(
name="zoo",
params=ParamList(
params=[
Param(name="x", type_="Mapping[str, Any]"),
Param(
name="q",
type_="Sequence[Union[Literal['a'], Literal['b']]]",
),
]
),
return_type="Any",
implementation=zoo,
doc_string="Add one to x.",
)
tests/test_interpreter.py
+203
View File
@@ -0,0 +1,203 @@
from typing import Any
import pytest
from kork import ast
from kork.ast_printer import AstPrinter
from kork.environment import Environment
from kork.exceptions import KorkRunTimeException
from kork.interpreter import run_interpreter
# TODO: Determine why mypy is not recognizing the import.
from kork.parser import parse # type: ignore[attr-defined]
def set_up_environment() -> Environment:
"""Prepopulate an environment with some symbols."""
environment = Environment(
parent=None,
variables={},
)
environment.set_symbol(
"add",
ast.ExternFunctionDef(
name="add",
params=ast.ParamList(
params=[
ast.Param(name="a", type_="number"),
ast.Param(name="b", type_="number"),
]
),
return_type="number",
).add_implementation(lambda a, b: a + b),
)
environment.set_symbol(
"time",
ast.ExternFunctionDef(
name="time",
params=ast.ParamList(params=[]),
return_type="number",
).add_implementation(lambda: 1.0),
)
return environment
@pytest.mark.parametrize(
"code, value",
[
# Int
("var x = 1", 1),
# None
("var x = null", None),
# Scientific
("var x = 1e2", 100),
("var x = 1E2", 100),
("var x = 1E-2", 0.01),
# String
('var x = "hello"', "hello"),
# Float
("var x = 5.0", 5.0),
("var x = 5.0 + 2", 7.0),
("var x = 5\nvar x = 2", 2.0),
("var x = 5;var x = 2", 2.0),
("var x = add(1,2)", 3.0),
("var x = time()", 1.0),
# Add
("var x = 1 + 1", 2),
# Subtract
("var x = 2 - 3", -1),
# Unary minus
("var x = -3", -3),
# Divide
("var x = 6 / 3", 2),
# Multiply
("var x = 2 * 3", 6),
# Operator precedence
("var x = 1 + 3 * 3", 10),
# Grouping
("var x = (1 + 1) * 3", 6),
# List
('var x = [1.0,2,"3"]', [1.0, 2, "3"]),
# Variable declaration and assignment
("var x = 1\nx=2", 2),
# Comments
("var x = 1 # comment", 1),
("var x = 1 // Also a comment", 1),
# Sum with identifiers
("var x = 2\nvar x = 1 + x", 3),
("var x = 2\nvar x = 1 / x", 0.5),
("var y = 1\nvar z = 2\nvar x = y + z", 3),
("var y = 1\nvar z = 2\nvar x = y * z", 2),
# Add binary operators
("var x = 1 + 1", 2),
("var x = 1 - 1", 0),
("var x = 1 * 3", 3),
("var x = 8 / 2", 4),
("var x = 9 % 3", 0),
("var x = 9 % 5", 4),
# unary rule
("var x = --3", 3),
# Take powers
("var x = 2 ** 4", 16),
("var x = 2 ^ 4", 16),
("var x = 2 ^ -1", 0.5),
("var x = 2 ^ -add(2, 2)", 1 / 16),
("var x = 2 ^ +add(2, 2)", 16),
("var y = 3; var x = 2 ^ y", 8),
("var y = 3; var z = 2; var x = z ^ y", 8),
("var y = 3; var z = 2; var x = z ** y", 8),
# Add binary operators involving primary and function call
("var z = 2\nvar x = 1 + add(z, z)", 5),
# Add binary operators involving function calls
("var z = 2\nvar x = add(z, z) * add(z, z)", 16),
# Repeat tests but with semi-colons
# Sum with identifiers
("var x = 2; x = 1 + x;", 3),
("var x = 2; x = 1 / x;", 0.5),
("var y = 1; z = 2; x = y + z;", 3),
("var y = 1; z = 2; x = y * z;", 2),
# Add binary operators involving primary and function call
("var z = 2; x = 1 + add(z, z);", 5),
# Add binary operators involving function calls
("var z = 2; x = add(z, z) * add(z, z);", 16),
],
)
def test_interpreter(code: str, value: Any) -> None:
"""Run program and verify value of `x` symbol is as expected."""
result = run_interpreter(code, environment=set_up_environment())
assert result["environment"].get_symbol("x") == value
def test_interpreter_extern_func_declaration() -> None:
"""Test that an extern function declaration is added to the environment."""
result = run_interpreter("extern fn meow(x : int, y:int) -> int")
env = result["environment"]
extern_fn_declaration = env.get_symbol("meow")
assert isinstance(extern_fn_declaration, ast.ExternFunctionDef)
# Verify that the function cannot be called.
result = run_interpreter("meow(1,2)", environment=env)
assert isinstance(result["errors"][0], KorkRunTimeException)
env.set_symbol("meow", extern_fn_declaration.add_implementation(lambda x, y: x + y))
result = run_interpreter("var x = meow(1,2)", environment=env)
assert result["environment"].get_symbol("x") == 3
def test_cannot_invoke_undefined_variable() -> None:
"""Test exception is raised when referencing an undefined variable."""
result = run_interpreter("var x = y")
assert isinstance(result["errors"][0], KorkRunTimeException)
def test_function_definition() -> None:
"""Test unimplemented features raise exception."""
result = run_interpreter("def fn meow(x : int, y:int) -> int {}")
assert isinstance(result["errors"][0], KorkRunTimeException)
def test_function_invocation() -> None:
"""Test unimplemented features raise exception."""
environment = Environment(parent=None, variables={})
environment.set_symbol(
"meow",
ast.FunctionDef(
name="meow",
body=[],
params=ast.ParamList(
params=[
ast.Param(name="a", type_="number"),
ast.Param(name="b", type_="number"),
]
),
return_type="number",
),
)
result = run_interpreter("var x = meow(1,2)", environment=environment)
run_interpreter("fn meow(x : int, y:int) -> int {}\nmeow(1,2)")
assert isinstance(result["errors"][0], KorkRunTimeException)
@pytest.mark.parametrize(
"code, expected",
[
("var x = 5", "var x = 5"),
("var x = 5\t\t\t\nadd(2,3 )", "var x = 5\nadd(2, 3)"),
# Verify whitespace insensitivity
("extern fn add(x : int, y:int)->int", "extern fn add(x: int, y: int) -> int"),
(
"extern fn add (x : int, y:int) -> int",
"extern fn add(x: int, y: int) -> int",
),
("fn add(x : int, y:int) -> int {}", "fn add(x: int, y: int) -> int {\n}"),
('var x = ["a", [], 1, 2.0, null]', 'var x = ["a", [], 1, 2.0, null]'),
("var x = []; var z = 1;", "var x = []\nvar z = 1"),
],
)
def test_ast_parsing_and_printing(code: str, expected: str) -> None:
"""Test parsing of code into an AST and then printing it back as code."""
program = parse(code)
assert AstPrinter().visit(program) == expected
tests/test_prompt_adapter.py
+58
View File
@@ -0,0 +1,58 @@
from langchain.prompts import PromptTemplate
from kork.prompt_adapter import FewShotPromptValue, FewShotTemplate
def test_few_shot_template() -> None:
"""Test few shot template."""
prompt_template = PromptTemplate(
template="meow\n\n",
input_variables=[],
)
few_shot_template = FewShotTemplate(
instruction_template=prompt_template,
examples=[("foo", "bar")],
input_variables=["query"],
)
prompt_value = few_shot_template.format_prompt(query="query")
assert isinstance(prompt_value, FewShotPromptValue)
# Test to_string()
assert (
prompt_value.to_string()
== "meow\n\nInput: foo\n\nOutput: bar\n\nInput: query\n\nOutput:"
)
# Test to_messages()
messages = prompt_value.to_messages()
assert len(messages) == 4
assert messages[0].content == "meow"
assert messages[1].content == "foo"
assert messages[2].content == "bar"
assert messages[3].content == "query"
def test_few_shot_template_with_no_examples() -> None:
"""Test few shot template without any examples"""
prompt_template = PromptTemplate(
template="meow\n\n",
input_variables=[],
)
few_shot_template = FewShotTemplate(
instruction_template=prompt_template,
examples=[],
input_variables=["query"],
)
prompt_value = few_shot_template.format_prompt(query="query")
assert isinstance(prompt_value, FewShotPromptValue)
# Test to_string()
assert prompt_value.to_string() == "meow\n\nInput: query\n\nOutput:"
# Test to_messages()
messages = prompt_value.to_messages()
assert len(messages) == 2
assert messages[0].content == "meow"
assert messages[1].content == "query"
tests/test_retrieval.py
+27
View File
@@ -0,0 +1,27 @@
from kork.ast import ExternFunctionDef, ParamList
from kork.retrieval import SimpleContextRetriever
def foo() -> None:
"""Do nothing."""
def bar(x: int) -> int:
"""Add one to x."""
return x + 1
def test_simple_retriever() -> None:
"""Test simple retriever"""
external_func = ExternFunctionDef(
name="meow",
params=ParamList(params=[]),
return_type="int",
implementation=None,
)
simple_retriever = SimpleContextRetriever.from_functions([foo, bar, external_func])
external_func_defs = simple_retriever.retrieve("does not matter")
assert len(external_func_defs) == 3
assert external_func_defs[0].name == "foo"
assert external_func_defs[1].name == "bar"
assert external_func_defs[2].name == "meow"
tests/test_utils.py
+27
View File
@@ -0,0 +1,27 @@
from kork.utils import unwrap_code, unwrap_tag, wrap_in_tag
def test_unwrap_tag() -> None:
"""Test unwrap_tag."""
# Test with an empty string
assert unwrap_tag("", "") is None
# Test with a string that doesn't contain the tag
assert unwrap_tag("table", "This is some text.") is None
# Test with a string that does contain the tag
assert unwrap_tag("table", "<table>hello</table>") == "hello"
# Test with a string that does contain the tag and some stuff before and after
assert unwrap_tag("table", "prefix<table>hello</table>suffix") == "hello"
def test_wrap_in_tag() -> None:
"""Test wrap_in_tag."""
assert wrap_in_tag("table", "hello") == "<table>hello</table>"
def test_unwrap_code() -> None:
"""Unwrap code."""
assert unwrap_code("python", "") is None
assert unwrap_code("python", "```python\nprint('hello')\n```") == "print('hello')"
tests/utils.py
+30
View File
@@ -0,0 +1,30 @@
from typing import Any, List, Optional
from langchain.chat_models.base import BaseChatModel
from langchain.schema import AIMessage, BaseMessage, ChatGeneration, ChatResult
from pydantic import Extra
class ToyChatModel(BaseChatModel):
response: str
class Config:
"""Configuration for this pydantic object."""
extra = Extra.forbid
arbitrary_types_allowed = True
def _generate(
self, messages: List[BaseMessage], stop: Optional[List[str]] = None
) -> ChatResult:
message = AIMessage(content=self.response)
generation = ChatGeneration(message=message)
return ChatResult(generations=[generation])
async def _agenerate(
self, messages: List[BaseMessage], stop: Optional[List[str]] = None
) -> Any:
"""Async version of _generate."""
message = AIMessage(content=self.response)
generation = ChatGeneration(message=message)
return ChatResult(generations=[generation])