calebdre’s gists

calebdre / run-query.ts

Created June 16, 2024 01:52

	type QueryResult = {
	score: number,
	func: FunctionInfo
	}
	export async function runQuery(
	query: string,
	files: FileInfo[]
	): Promise<QueryResult[]> {
	// convert input query to embedding
	const queryEmbedding = await createEmbedding(query);

calebdre / cosine-similarity.ts

Created June 16, 2024 01:50

	// math -> https://mathjs.org/index.html
	export function cosineSimilarity(embedding1: number[], embedding2: number[]): number {
	const dotProduct = math.dot(embedding1, embedding2)
	const magnitude1 = math.norm(embedding1)
	const magnitude2 = math.norm(embedding2)
	return dotProduct / (magnitude1 * magnitude2)
	}

calebdre / create-embedding.ts

Created June 15, 2024 17:23

	export const createEmbedding = async (
	model: string,
	text: string
	) => {
	const resp = await ollama.embeddings({
	model: model,
	prompt: text
	})

	return resp.embedding

calebdre / data-collect.ts

Last active June 15, 2024 17:22

	const collectFileData = async (dir: string) => {
	// collect all the files in the directory
	const files = findTsFiles(dir)
	for (const file of files) {
	// extract functions from each file
	file.functions = extractFunctions(file.content)
	for (const func of file.functions) {
	// get the description for each function
	const description = await getFunctionDescription(file, func.functionName)
	if (description) {

calebdre / getFunctionDescription.ts

Created June 14, 2024 20:57

	export const getFunctionDescription = async (file: FileInfo, functionName: string, model: string = "llama3"): Promise<string \| null> => {
	const prompt = `given the following file ${file!.path}/${file!.filename}:
	\`\`\`typescript
	${file!.content}
	\`\`\`

	Provide a concise description for the function \`${functionName}\`. focusing on its specific role within the larger codebase.
	The description should briefly explain:
	- The purpose of the function
	- The function's inputs, outputs, and any notable side effects or dependencies

calebdre / extractFunctions.ts

Created June 14, 2024 20:55

	export function extractFunctions(fileContents: string): ExtractFunctionResponse[] {
	const functions: ExtractFunctionResponse[] = [];

	try {
	const ast = parser.parse(fileContents, {
	sourceType: 'module',
	plugins: ['typescript', 'jsx'],
	allowImportExportEverywhere: true,
	allowAwaitOutsideFunction: true,
	allowReturnOutsideFunction: true,

calebdre / findTsFiles.ts

Created June 14, 2024 20:54

	type FileInfo {
	content: string
	path: string
	filename: string
	}

	function findTsFiles(directory: string, exclude_dirs: string[]): void {
	const fileList: FileInfo[] = []

	function traverseDirectory(currentDir: string) {

calebdre / evaluateCodeSearch.ts

Last active June 14, 2024 01:11

	const evaluateCodeSearch = async (
	files: FileInfo[],
	numSamples: number = 15,
	k: number = 5
	) => {
	// create validation set
	const validationSet: ValidationQuery[] = await generateValidationQueries(files, numSamples)
	const resultSet: QueryResult[][] = []

	// run each validation query against the code search system

calebdre / calculateMap@k.ts

Created June 14, 2024 00:59

	export const calculateMAPAtK = (
	sampleQueries: SampleQuery[],
	retrievedResults: QueryResult[][],
	k: number
	) => {
	const apScores = sampleQueries.map((query, index) => {
	return calculateAveragePrecision(query, retrievedResults[index], k);
	});

	return apScores.reduce((sum, score) => sum + score, 0) / apScores.length;

calebdre / calculateAveragePrecision.ts

Last active June 14, 2024 00:56

	export const calculateAveragePrecision = (
	sampleQuery: SampleQuery,
	retrievedResults: QueryResult[],
	k: number
	) => {
	// get the function names from the validation set
	const relevantFunctions = sampleQuery.expectedFunctions.map((func) => func.functionName)

	// get the function names from the retrieved results
	const retrievedFunctionNames = retrievedResults.map((result) => result.funcName)

Caleb Lewis calebdre