Retrievers

The Kork CodeChain accepts retriever interfaces for both examples and foreign functions.

The built-in retrievers return the content that they were initialized with as is.

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.

The code below defines a place holder model for a chat model. Feel free to replace it with a real language model.

Show code cell source Hide code cell source

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])

Context Retriever

A context retriever (yes, it’s poorly named) supports a method that takes the user query and returns a list of foreign functions.

The built-in retriever returns the information it was initialized with, entirely ignoring the query.

Sub-class from the abstract interface to provide your own implementation!

from kork import SimpleContextRetriever
import math

simple_context_retriever = SimpleContextRetriever.from_functions(
    [math.pow, math.log2, math.log10]
)

simple_context_retriever.retrieve("this is the query")

[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).'),
 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.'),
 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.')]

Example Retriever

Supports the same retrieve method as a context retriever.

The built-in example retriever returns the information it was initialized with, entirely ignoring the query.

Sub-class from the abstract interface to provide your own implementation!

from kork import c_, r_, AstPrinter, SimpleExampleRetriever

language_name = "RollingMeow"

retriever = SimpleExampleRetriever.from_programs(
    language_name,
    [
        ("2**5", r_(c_(math.pow, 2, 5))),
        ("take the log base 2 of 2", r_(c_(math.log2, 2))),
    ],
    AstPrinter(),
)

retriever.retrieve("[ignore]")

[('2**5', '<code>var result = pow(2, 5)</code>'),
 ('take the log base 2 of 2', '<code>var result = log2(2)</code>')]

Declare the chain

from kork import (
    CodeChain,
    ast,
    AstPrinter,
)

chain = CodeChain.from_defaults(
    llm=ToyChatModel(response="MEOW MEOW MEOW MEOW"),  # The LLM to use
    ast_printer=AstPrinter(),  # Knows how to print the AST
    examples=retriever,  # Example programs
    context=simple_context_retriever,
    language_name=language_name,
)

_, few_shot_prompt = chain.prepare_context(query="hello")

print(few_shot_prompt.format_prompt(query="[user input]").to_string())

You are programming in a language called "RollingMeow".

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 RollingMeow
program.

You have access to the following external functions:

```RollingMeow
extern fn pow(x: Any, y: Any) -> Any // Return x**y (x to the power of y).
extern fn log2(x: Any) -> Any // Return the base 2 logarithm of x.
extern fn log10(x: Any) -> Any // Return the base 10 logarithm of x.
```


Do not assume that any other functions except for the ones listed above exist.

Wrap the program in <code> and </code> tags.

Store the solution to the query in a variable called "result".

Here is a sample valid program:

<code>
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
</code>

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!
Input: 2**5

Output: <code>var result = pow(2, 5)</code>

Input: take the log base 2 of 2

Output: <code>var result = log2(2)</code>

Input: [user input]

Output: