@@ -9,9 +9,11 @@ import (
99 "os"
1010 "strings"
1111
12+ "github.com/chand1012/memory"
1213 "github.com/chand1012/ottodocs/pkg/calc"
1314 "github.com/chand1012/ottodocs/pkg/config"
1415 "github.com/chand1012/ottodocs/pkg/constants"
16+ "github.com/chand1012/ottodocs/pkg/git"
1517 "github.com/chand1012/ottodocs/pkg/textfile"
1618 "github.com/chand1012/ottodocs/pkg/utils"
1719 l "github.com/charmbracelet/log"
@@ -96,6 +98,77 @@ Example: otto edit main.go --start 1 --end 10 --goal "Refactor the function"`,
9698 prompt = "GOAL: " + chatPrompt + "\n \n FILE: " + filePath + "\n \n " + contents
9799 }
98100
101+ client := openai .NewClient (c .APIKey )
102+
103+ if repoContext {
104+ repo , err := git .GetRepo ("." , "" , false )
105+ if err != nil {
106+ log .Errorf ("Error getting repo: %s" , err )
107+ os .Exit (1 )
108+ }
109+
110+ m , _ , err := memory .New (":memory:" )
111+ if err != nil {
112+ log .Errorf ("Error creating memory: %s" , err )
113+ os .Exit (1 )
114+ }
115+
116+ for _ , file := range repo .Files {
117+ err = m .Add (file .Path , file .Contents )
118+ if err != nil {
119+ log .Errorf ("Error indexing file: %s" , err )
120+ os .Exit (1 )
121+ }
122+ }
123+
124+ queryConstructorPrompt := []openai.ChatCompletionMessage {
125+ {
126+ Role : openai .ChatMessageRoleUser ,
127+ Content : "Convert this goal into a terms to use for a search query. They do not have to be organized, nor a complete sentence. Use no form of punctuation or quotations. Only return the query and nothing else: " + chatPrompt ,
128+ },
129+ }
130+
131+ utils .PrintColoredText ("Otto: " , c .OttoColor )
132+ fmt .Println ("Ok! Here is the query, taking your input into account." )
133+ utils .PrintColoredText ("Otto: " , c .OttoColor )
134+ stream , err := client .CreateChatCompletionStream (context .Background (), openai.ChatCompletionRequest {
135+ Model : "gpt-3.5-turbo" ,
136+ Messages : queryConstructorPrompt ,
137+ })
138+
139+ if err != nil {
140+ log .Errorf ("Error requesting from OpenAI: %s" , err )
141+ os .Exit (1 )
142+ }
143+
144+ query , err := utils .PrintChatCompletionStream (stream )
145+ if err != nil {
146+ log .Errorf ("Error printing chat completion stream: %s" , err )
147+ os .Exit (1 )
148+ }
149+
150+ utils .PrintColoredText ("Otto: " , c .OttoColor )
151+ fmt .Println ("Searching repo for files that match the query..." )
152+
153+ resp , err := m .Search (query )
154+ if err != nil {
155+ log .Errorf ("Error searching memory: %s" , err )
156+ os .Exit (1 )
157+ }
158+
159+ contextFiles = []string {}
160+ for _ , file := range resp {
161+ contextFiles = append (contextFiles , file .ID )
162+ }
163+ log .Debugf ("context files: %s" , contextFiles )
164+ }
165+ tokens , err := calc .PreciseTokens (prompt )
166+ if err != nil {
167+ log .Errorf ("Error calculating tokens: %s" , err )
168+ os .Exit (1 )
169+ }
170+
171+ maxTokens := calc .GetMaxTokens (c .Model )
99172 if len (contextFiles ) > 0 {
100173 var contextContent string
101174 for _ , contextFile := range contextFiles {
@@ -104,11 +177,20 @@ Example: otto edit main.go --start 1 --end 10 --goal "Refactor the function"`,
104177 log .Errorf ("Error loading context file: %s" , err )
105178 continue
106179 }
180+ contentTokens , err := calc .PreciseTokens (contextContent )
181+ if err != nil {
182+ log .Errorf ("Error loading context file: %s" , err )
183+ continue
184+ }
185+ if contentTokens + tokens > maxTokens {
186+ break
187+ }
107188 prompt += "\n \n CONTEXT: " + contextFile + "\n \n " + contextContent
189+ tokens += contentTokens
108190 }
109191 }
110192
111- client := openai . NewClient ( c . APIKey )
193+ log . Debugf ( "prompt: %s" , prompt )
112194
113195 messages = []openai.ChatCompletionMessage {
114196 {
@@ -220,6 +302,7 @@ func init() {
220302 editCmd .Flags ().BoolVarP (& force , "force" , "f" , false , "Force overwrite of existing files" )
221303 editCmd .Flags ().BoolVarP (& verbose , "verbose" , "v" , false , "Verbose output" )
222304 editCmd .Flags ().BoolVarP (& appendFile , "append" , "a" , false , "Append to the end of a file instead of overwriting it" )
305+ editCmd .Flags ().BoolVarP (& repoContext , "repo" , "r" , false , "Use the current repo as context" )
223306 editCmd .Flags ().IntVarP (& startLine , "start" , "s" , 1 , "Start line" )
224307 editCmd .Flags ().IntVarP (& endLine , "end" , "e" , 0 , "End line" )
225308 editCmd .Flags ().StringVarP (& chatPrompt , "goal" , "g" , "" , "Goal of the edit" )
0 commit comments