Enhance analysis command to handle current directory and update model…

… mappings for anthropic
PrashamTrivedi · Nov 12, 2024 · 81d81e7 · 81d81e7
1 parent 89831c9
commit 81d81e7
Show file tree

Hide file tree

Showing 4 changed files with 121 additions and 11 deletions.
diff --git a/_tests_/analyse.test.mts b/_tests_/analyse.test.mts
@@ -8,6 +8,16 @@ import fs from 'fs'
 import ora from 'ora'
 import * as inquirer from '@inquirer/prompts'
 import * as setExpertise from '../commands/setExpertise.mjs'
+import chalk from "chalk"
+
+// Remove this mock as it won't work correctly
+// const mockProcessCwd = vi.fn(() => '/test/current-project')
+// vi.mock('process', () => ({
+//   default: {
+//     cwd: mockProcessCwd
+//   },
+//   cwd: mockProcessCwd
+// }))
 
 vi.mock('../directoryProcessor.mjs')
 vi.mock('../modelUtils.mjs')
@@ -20,6 +30,10 @@ vi.mock('../commands/setExpertise.mjs')
 describe('analyse command', () => {
   beforeEach(() => {
     vi.resetAllMocks()
+    // Mock process.cwd
+    const mockCwd = vi.spyOn(process, 'cwd')
+    mockCwd.mockReturnValue('/test/current-project')
+
     vi.mocked(ora).mockReturnValue({
       start: vi.fn().mockReturnThis(),
       stop: vi.fn().mockReturnThis(),
@@ -33,6 +47,7 @@ describe('analyse command', () => {
 
   afterEach(() => {
     vi.resetAllMocks()
+    vi.restoreAllMocks() // Add this to cleanup process.cwd mock
   })
 
   it('should analyze a non-monorepo project correctly', async () => {
@@ -282,4 +297,61 @@ describe('analyse command', () => {
     })
     expect(setExpertise.handler).toHaveBeenCalled()
   })
+
+  it('should handle current directory analysis correctly', async () => {
+    const mockArgv = {
+      path: '.',
+      verbose: false,
+      openai: true,
+      streaming: false,
+      ignore: [],
+    }
+
+    const mockDirectoryStructure = {
+      name: 'current-project',
+      children: [
+        {name: 'src', children: []},
+        {name: 'package.json', content: '{}'},
+      ],
+    }
+
+    const mockDirectoryInference = {
+      isMonorepo: false,
+      workflow: 'nodejs',
+      dependenciesFile: 'package.json',
+      programmingLanguage: 'javascript',
+    }
+
+    vi.mocked(directoryProcessor.getDirStructure).mockResolvedValue(mockDirectoryStructure as any)
+    const mockModelUtils = {
+      initializeModels: vi.fn().mockResolvedValue(undefined),
+      getLlmInterface: vi.fn().mockReturnValue({
+        getName: vi.fn().mockReturnValue('Mocked model'),
+        inferProjectDirectory: vi.fn().mockResolvedValue(JSON.stringify(mockDirectoryInference)),
+        inferCode: vi.fn().mockResolvedValue('Mocked code inference'),
+        inferInterestingCode: vi.fn().mockResolvedValue('Mocked interesting code'),
+        inferDependency: vi.fn().mockResolvedValue('Mocked dependency inference'),
+      }),
+    }
+    vi.mocked(ModelUtils.getInstance).mockReturnValue(mockModelUtils as any)
+    vi.mocked(utils.readConfig).mockReturnValue({ANALYSIS_DIR: '/test'} as any)
+    vi.mocked(fs.readFileSync).mockReturnValue('{}')
+
+    const consoleSpy = vi.spyOn(console, 'log')
+
+    await handler(mockArgv as any)
+
+    expect(consoleSpy).toHaveBeenCalledWith('Analyzing current directory: current-project')
+    expect(consoleSpy).toHaveBeenCalledWith(`Analysing ${chalk.redBright('current-project')}'s file structure to getting started.`)
+    expect(directoryProcessor.getDirStructure).toHaveBeenCalledWith('.', [], false)
+    expect(utils.writeAnalysis).toHaveBeenCalledWith(
+      '/test/.SourceSailor/current-project',
+      'directoryStructure',
+      expect.any(Object),
+      true,
+      false
+    )
+
+    consoleSpy.mockRestore()
+  })
 })
diff --git a/_tests_/anthropic.test.mts b/_tests_/anthropic.test.mts
@@ -52,7 +52,24 @@ describe('AnthropicInterface', () => {
 
     describe('listModels', () => {
         it('should return a list of available models', async () => {
-            const expectedModels = ['claude-3-haiku', 'claude-3-sonnet', 'claude-3-opus', 'claude-3.5-sonnet', 'haiku-3', 'sonnet-3', 'opus-3', 'sonnet-3.5']
+            const expectedModels = [
+                'claude-3-haiku',
+                'claude-3.5-haiku',
+                'claude-3.5-haiku-latest',
+                'claude-3-sonnet',
+                'claude-3-opus',
+                'claude-3.5-sonnet',
+                'claude-3.5-sonnet-legacy',
+                'claude-3.5-sonnet-latest',
+                'haiku-3',
+                'haiku-3.5',
+                'haiku-3.5-latest',
+                'sonnet-3',
+                'opus-3',
+                'sonnet-3.5',
+                'sonnet-3.5-legacy',
+                'sonnet-3.5-latest'
+            ]
             const result = await anthropicInterface.listModels(false)
             expect(result).toEqual(expectedModels)
         })

diff --git a/anthropic.mts b/anthropic.mts
@@ -10,12 +10,16 @@ type modelType = keyof typeof modelMapping
 
 const modelMapping = {
     "claude-3-haiku": 'claude-3-haiku-20240307',
+    "claude-3.5-haiku": 'claude-3-5-haiku-20241022',
+    "claude-3.5-haiku-latest": 'claude-3-5-haiku-latest',
     'claude-3-sonnet': 'claude-3-sonnet-20240229',
     'claude-3-opus': 'claude-3-opus-20240229',
     'claude-3.5-sonnet': 'claude-3-5-sonnet-20241022',
     'claude-3.5-sonnet-legacy': 'claude-3-5-sonnet-20240620',
     'claude-3.5-sonnet-latest': 'claude-3-5-sonnet-latest',
     'haiku-3': 'claude-3-haiku-20240307',
+    'haiku-3.5': 'claude-3-5-haiku-20241022',
+    'haiku-3.5-latest': 'claude-3-5-haiku-latest',
     'sonnet-3': 'claude-3-sonnet-20240229',
     'opus-3': 'claude-3-opus-20240229',
     'sonnet-3.5': 'claude-3-5-sonnet-20241022',

diff --git a/commands/analyse.mts b/commands/analyse.mts
@@ -6,8 +6,7 @@ import fs from 'fs'
 import {addAnalysisInGitIgnore, readConfig, writeAnalysis, writeError} from "../utils.mjs"
 import ora from 'ora'
 import {Arguments} from 'yargs'
-import {ChatCompletionChunk} from "openai/resources/index.mjs"
-import {Stream} from "openai/streaming.mjs"
+
 import chalk from "chalk"
 import {confirm} from '@inquirer/prompts'
 import {handler as setExpertiseHandler} from './setExpertise.mjs'
@@ -77,21 +76,29 @@ export async function handler(argv: Arguments) {
     const ignore = argv.ignore as string[] || argv.i as string[] || []
     const modelName = argv.model as string || argv.m as string
 
-
     const config = readConfig()
 
     if (isVerbose) {
         console.log(`Analyse the given directory structure to understand the project structure and dependencies: ${argv.path}`)
-
         console.log(`Using model: ${modelName || config.DEFAULT_OPENAI_MODEL}`)
     }
 
     const rootDir = config.ANALYSIS_DIR
     const userExpertise = JSON.stringify(config.userExpertise)
 
-    const projectName = argv.path as string
+    // Handle current directory case
+    const inputPath = argv.path as string
+    const projectName = inputPath === '.' ? process.cwd().split('/').pop() ?? "" : inputPath
+
+    // Log current directory name if analyzing current directory
+    if (inputPath === '.') {
+        console.log(`Analyzing current directory: ${projectName}`)
+    }
 
     const isProjectRoot = rootDir === 'p'
+    if (isVerbose) {
+        console.log({rootDir, isProjectRoot, projectName})
+    }
 
     const modelUtils = ModelUtils.getInstance()
     await modelUtils.initializeModels()
@@ -107,14 +114,24 @@ export async function handler(argv: Arguments) {
     }
     console.log(`Analysing ${chalk.redBright(projectName)}'s file structure to getting started.`)
     // const defaultSpinner = ora().start()
-    const path = argv.path as string
+    const sourceCodePath = inputPath
     const isRoot = true
-    const sourceCodePath = argv.path as string
     const dirToWriteAnalysis = isProjectRoot ? `${sourceCodePath}/.SourceSailor` : `${rootDir}/.SourceSailor/${projectName}`
+if (isVerbose) {
+        console.log({sourceCodePath, isProjectRoot, isRoot, dirToWriteAnalysis})
+    }
 
-
-    const {directoryInferrence, directoryStructureWithContent} = await analyseDirectoryStructure(path, isVerbose, isRoot,
-        dirToWriteAnalysis, isProjectRoot, ignore, llmInterface, userExpertise, selectedModelName)
+    const {directoryInferrence, directoryStructureWithContent} = await analyseDirectoryStructure(
+        sourceCodePath,  // Use sourceCodePath instead of undefined path
+        isVerbose,
+        isRoot,
+        dirToWriteAnalysis,
+        isProjectRoot,
+        ignore,
+        llmInterface,
+        userExpertise,
+        selectedModelName
+    )
 
     if (isVerbose) {
         console.log({project: argv.path, directoryInferrence})