Skip to content

Commit

Permalink
Update to use turbo model
Browse files Browse the repository at this point in the history
  • Loading branch information
@handotdev
handotdev committed Apr 4, 2023
1 parent 2e1ca1e commit 3de9b09
Show file tree
Hide file tree
Showing 3 changed files with 251 additions and 132 deletions.
207 changes: 152 additions & 55 deletions server/brain/codex/docs.ts
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,35 @@
import { getLanguageContextById } from 'brain/languages';
import { addComments } from 'brain/helpers';
import { Synopsis, Property, ParamExplained, FunctionSynopsis, TypedefSynopsis } from 'parsing/types';
import { EXPLAIN_PARAM, EXPLAIN_SIMPLE, SUMMARIZE_FUNCTION,
EXPLAIN_CONTEXT, EXPLAIN_PROPERTY, SUMMARIZE_TYPE, GET_RETURN, SUMMARIZE_FUNCTION_SIMPLE, SUMMARIZE_CLASS, CodexCall, SUMMARIZE_CLASS_SIMPLE } from 'brain/codex/prompt';
import {
Synopsis,
Property,
ParamExplained,
FunctionSynopsis,
TypedefSynopsis,
} from 'parsing/types';
import {
EXPLAIN_PARAM,
EXPLAIN_SIMPLE,
SUMMARIZE_FUNCTION,
EXPLAIN_CONTEXT,
EXPLAIN_PROPERTY,
SUMMARIZE_TYPE,
GET_RETURN,
SUMMARIZE_FUNCTION_SIMPLE,
SUMMARIZE_CLASS,
OpenAPICall,
SUMMARIZE_CLASS_SIMPLE,
} from 'brain/codex/prompt';
import { writeFunctionInFormat } from 'formatting/functions';
import { formatTypedef } from 'formatting/typedef';
import { DocFormat } from 'constants/enums';
import { chooseDocstringPrompt, formatReturnExplained, getSummaryFromMultipleResponses, makeCodexCall, OpenAIResponse } from 'brain/codex/helpers';
import {
chooseDocstringPrompt,
formatReturnExplained,
getSummaryFromMultipleResponses,
makeCodexCall,
OpenAIResponse,
} from 'brain/codex/helpers';
import { AxiosResponse } from 'axios';
import { Custom } from 'routes/writer/helpers';
import { getMultipleTranslations } from 'services/translate';
Expand All @@ -16,60 +39,98 @@ const MAX_CHARACTER_COUNT = 12000;
export type DocstringPrompt = {
docstring: string;
promptId: string;
}
};

const responseToDocstringPrompt = (responses: AxiosResponse<OpenAIResponse>[], callTypes: CodexCall[]): DocstringPrompt[] => {
const responseToDocstringPrompt = (
responses: AxiosResponse<OpenAIResponse>[],
callTypes: OpenAPICall[]
): DocstringPrompt[] => {
const docstringPrompts: DocstringPrompt[] = [];
for(let i=0; i<responses.length; i++) {
docstringPrompts.push({ docstring: responses[i]?.data?.choices[0].text, promptId: callTypes[i].id})
for (let i = 0; i < responses.length; i++) {
docstringPrompts.push({
docstring: responses[i]?.data?.choices[0].text,
promptId: callTypes[i].id,
});
}
return docstringPrompts
}
return docstringPrompts;
};

const getFunctionDocstringPrompt = async (synopsis: FunctionSynopsis, code: string, languageCommented: string, docFormat: DocFormat, custom: Custom): Promise<DocstringPrompt> => {
const paramPromises = synopsis.params?.map((param) => {
return makeCodexCall(EXPLAIN_PARAM, code, languageCommented, { parameter: param.name });
}) ?? [];
const getFunctionDocstringPrompt = async (
synopsis: FunctionSynopsis,
code: string,
languageCommented: string,
docFormat: DocFormat,
custom: Custom
): Promise<DocstringPrompt> => {
const paramPromises =
synopsis.params?.map((param) => {
return makeCodexCall(EXPLAIN_PARAM, code, languageCommented, { parameter: param.name });
}) ?? [];

const returnPromise = synopsis.returns ? makeCodexCall(GET_RETURN, code, languageCommented) : null;
const returnPromise = synopsis.returns
? makeCodexCall(GET_RETURN, code, languageCommented)
: null;
const summaryPromise = makeCodexCall(SUMMARIZE_FUNCTION, code, languageCommented);
const summarySimplePromise = makeCodexCall(SUMMARIZE_FUNCTION_SIMPLE, code, languageCommented);
const docPromises = [returnPromise, summaryPromise, summarySimplePromise, ...paramPromises];
const docResponses = await Promise.all(docPromises);

let paramsExplained: ParamExplained[] = paramPromises.length > 0 ? docResponses.slice(3).map((param, i) => {
const explained = param.data.choices[0].text.trim();
return {
...synopsis.params[i],
explanation: explained
}
}) : [];
let paramsExplained: ParamExplained[] =
paramPromises.length > 0
? docResponses.slice(3).map((param, i) => {
const explained = param.data.choices[0].text.trim();
return {
...synopsis.params[i],
explanation: explained,
};
})
: [];

const [returnRes, summaryRes, summarySimpleRes] = docResponses;
const docstringPrompts = responseToDocstringPrompt([summaryRes, summarySimpleRes], [SUMMARIZE_FUNCTION, SUMMARIZE_FUNCTION_SIMPLE]);
const docstringPrompts = responseToDocstringPrompt(
[summaryRes, summarySimpleRes],
[SUMMARIZE_FUNCTION, SUMMARIZE_FUNCTION_SIMPLE]
);
const docstringPrompt = chooseDocstringPrompt(docstringPrompts);
let returnExplained = formatReturnExplained(returnRes?.data?.choices[0].text);

// Adding premium feature for additional language translation
if (custom.language) {
const paramsExplinations = paramsExplained.map((paramExplained) => paramExplained.explanation);
const translatedResults = await getMultipleTranslations([docstringPrompt.docstring, returnExplained, ...paramsExplinations], custom.language);
const [translatedDocstring, translatedReturnExplained, ...translatedParamsExplained] = translatedResults;
const translatedResults = await getMultipleTranslations(
[docstringPrompt.docstring, returnExplained, ...paramsExplinations],
custom.language
);
const [translatedDocstring, translatedReturnExplained, ...translatedParamsExplained] =
translatedResults;
docstringPrompt.docstring = translatedDocstring;
returnExplained = translatedReturnExplained;
paramsExplained = paramsExplained.map((paramExplained, i) => {
return {
...paramExplained,
explanation: translatedParamsExplained[i]
}
explanation: translatedParamsExplained[i],
};
});
}

const docstring = writeFunctionInFormat(docFormat, docstringPrompt.docstring, paramsExplained, returnExplained, synopsis.returnsType, custom);
const docstring = writeFunctionInFormat(
docFormat,
docstringPrompt.docstring,
paramsExplained,
returnExplained,
synopsis.returnsType,
custom
);
return { docstring, promptId: docstringPrompt.promptId };
}
};

const getTypedefDocstring = async (synopsis: TypedefSynopsis, code: string, languageCommented: string, docFormat: DocFormat, custom: Custom): Promise<string> => {
const getTypedefDocstring = async (
synopsis: TypedefSynopsis,
code: string,
languageCommented: string,
docFormat: DocFormat,
custom: Custom
): Promise<string> => {
const propertiesPromises = synopsis.properties?.map((property) => {
return makeCodexCall(EXPLAIN_PROPERTY, code, languageCommented, { property: property.name });
});
Expand All @@ -78,88 +139,124 @@ const getTypedefDocstring = async (synopsis: TypedefSynopsis, code: string, lang
const docPromises = [summaryPromise, ...propertiesPromises];
const docResponses = await Promise.all(docPromises);

let propertiesWithExplanation: Property[] = docResponses.slice(1)?.map((param, i) => {
let propertiesWithExplanation: Property[] = docResponses.slice(1)?.map((param, i) => {
return {
...synopsis.properties[i],
explanation: param.data.choices[0].text.trim()
}
explanation: param.data.choices[0].text.trim(),
};
});
const [summaryRes] = docResponses;
let summary = getSummaryFromMultipleResponses(summaryRes.data.choices[0].text);

// Adding languages feature
if (custom.language) {
const propertiesExplanations = propertiesWithExplanation.map((propertyExplained) => propertyExplained.explanation);
const translatedResults = await getMultipleTranslations([summary, ...propertiesExplanations], custom.language);
const propertiesExplanations = propertiesWithExplanation.map(
(propertyExplained) => propertyExplained.explanation
);
const translatedResults = await getMultipleTranslations(
[summary, ...propertiesExplanations],
custom.language
);
const [translatedSummary, ...translatedPropertiesExplained] = translatedResults;
summary = translatedSummary;
propertiesWithExplanation = propertiesWithExplanation.map((propertyExplained, i) => {
return {
...propertyExplained,
explanation: translatedPropertiesExplained[i]
}
explanation: translatedPropertiesExplained[i],
};
});
}

return formatTypedef(docFormat, summary, propertiesWithExplanation);
}
};

const getTypedefDocstringPrompt = async (synopsis: TypedefSynopsis, code: string, languageCommented: string, docFormat: DocFormat, custom: Custom): Promise<DocstringPrompt> => {
const getTypedefDocstringPrompt = async (
synopsis: TypedefSynopsis,
code: string,
languageCommented: string,
docFormat: DocFormat,
custom: Custom
): Promise<DocstringPrompt> => {
const docstring = await getTypedefDocstring(synopsis, code, languageCommented, docFormat, custom);
return { docstring, promptId: SUMMARIZE_TYPE.id};
}
return { docstring, promptId: SUMMARIZE_TYPE.id };
};

const getClassDocstring = async (code: string, languageCommented: string, custom: Custom): Promise<DocstringPrompt> => {
const callOptions: CodexCall[] = [SUMMARIZE_CLASS, SUMMARIZE_CLASS_SIMPLE];
const getClassDocstring = async (
code: string,
languageCommented: string,
custom: Custom
): Promise<DocstringPrompt> => {
const callOptions: OpenAPICall[] = [SUMMARIZE_CLASS, SUMMARIZE_CLASS_SIMPLE];
const promises = callOptions.map((option) => {
return makeCodexCall(option, code, languageCommented);
});
const responses = await Promise.all(promises);
const docstringPrompts: DocstringPrompt[] = callOptions.map((option, i) => {
return {
docstring: responses[i]?.data?.choices[0].text,
promptId: option.id
}
promptId: option.id,
};
});

const selectedPrompt = chooseDocstringPrompt(docstringPrompts);

if (custom.language) {
const translatedResults = await getMultipleTranslations([selectedPrompt.docstring], custom.language);
const translatedResults = await getMultipleTranslations(
[selectedPrompt.docstring],
custom.language
);
const [translatedDocstring] = translatedResults;
selectedPrompt.docstring = translatedDocstring;
}

return selectedPrompt;
};

const getSimpleSummary = async (code: string, languageCommented: string, context: string, custom: Custom): Promise<DocstringPrompt> => {
const getSimpleSummary = async (
code: string,
languageCommented: string,
context: string,
custom: Custom
): Promise<DocstringPrompt> => {
let summaryContextPromise;
if (context.length + code.length <= MAX_CHARACTER_COUNT) {
summaryContextPromise = makeCodexCall(EXPLAIN_CONTEXT, code, languageCommented, { context });
}
const summarySimplePromise = makeCodexCall(EXPLAIN_SIMPLE, code, languageCommented);

const [summaryContext, summarySimple] = await Promise.all([summaryContextPromise, summarySimplePromise]);
const [summaryContext, summarySimple] = await Promise.all([
summaryContextPromise,
summarySimplePromise,
]);
const docstringOptions: DocstringPrompt[] = [
{ docstring: summaryContext?.data.choices[0]?.text, promptId: EXPLAIN_CONTEXT.id},
{ docstring: summarySimple?.data.choices[0]?.text, promptId: EXPLAIN_SIMPLE.id }
{ docstring: summaryContext?.data.choices[0]?.text, promptId: EXPLAIN_CONTEXT.id },
{ docstring: summarySimple?.data.choices[0]?.text, promptId: EXPLAIN_SIMPLE.id },
];
const selectedPrompt = chooseDocstringPrompt(docstringOptions);

if (custom.language) {
const translatedResults = await getMultipleTranslations([selectedPrompt.docstring], custom.language);
const translatedResults = await getMultipleTranslations(
[selectedPrompt.docstring],
custom.language
);
const [translatedDocstring] = translatedResults;
selectedPrompt.docstring = translatedDocstring;
}

return selectedPrompt;
}
};

export const getDocstringPrompt = async (code: string, synopsis: Synopsis, languageId: string | null, docFormat: DocFormat, context: string, custom: Custom = {}): Promise<DocstringPrompt> => {
export const getDocstringPrompt = async (
code: string,
synopsis: Synopsis,
languageId: string | null,
docFormat: DocFormat,
context: string,
custom: Custom = {}
): Promise<DocstringPrompt> => {
const languageContext = getLanguageContextById(languageId);
const languageCommented = addComments(languageContext.name, languageId);
switch(synopsis.kind) {
switch (synopsis.kind) {
case 'function':
return await getFunctionDocstringPrompt(synopsis, code, languageCommented, docFormat, custom);
case 'class':
Expand All @@ -169,4 +266,4 @@ export const getDocstringPrompt = async (code: string, synopsis: Synopsis, langu
default:
return await getSimpleSummary(code, languageCommented, context, custom);
}
}
};

0 comments on commit 3de9b09

Please sign in to comment.