larger token and also country of origin

app/main.py

@@ -84,7 +84,7 @@ class GenerateRequest(BaseModel):
     model: str
     business_idea: Optional[str] = ""
 
-def generate_model_output(model: str, provider: str, api_key: str, prompt: str, max_tokens: int = 500) -> str:
+def generate_model_output(model: str, provider: str, api_key: str, prompt: str, max_tokens: int = 4000) -> str:
     """
     A helper function that wraps the Hugging Face Inference API call.
     """
@@ -92,13 +92,13 @@ def generate_model_output(model: str, provider: str, api_key: str, prompt: str,
         logging.info(f"Initializing InferenceClient with provider: {provider}")
         client = InferenceClient(provider=provider, api_key=api_key)
         
-        logging.info(f"Sending request to model: {model} with prompt length: {len(prompt)}")
+        logging.info(f"Sending request to model: {model} with prompt length: {len(prompt)} and max_tokens: {max_tokens}")
         completion = client.chat.completions.create(
             model=model,
             messages=[{"role": "user", "content": prompt}],
             max_tokens=max_tokens,
         )
-        logging.info(f"Successfully received response from model: {model}")
+        logging.info(f"Successfully received response from model: {model} with content length: {len(completion.choices[0].message.content)}")
         return completion.choices[0].message.content
     except Exception as e:
         error_message = f"Error generating output with model {model}: {str(e)}"
@@ -113,9 +113,9 @@ class HFInferenceLLM(LLM):
     model: str = None
     provider: str = "hf-inference"
     api_key: str = ""
-    max_tokens: int = 500
+    max_tokens: int = 4000  # Increased from 500 to 4000 by default
 
-    def __init__(self, model: str, provider: str = "hf-inference", api_key: str = "", max_tokens: int = 500):
+    def __init__(self, model: str, provider: str = "hf-inference", api_key: str = "", max_tokens: int = 4000):
         super().__init__()
         self.model = model
         self.provider = provider
@@ -174,7 +174,7 @@ def get_llm(model_name: str):
         if not openai_api_key:
             logging.error("OPENAI_API_KEY environment variable not set")
             raise ValueError("OPENAI_API_KEY environment variable is required for GPT models")
-        return ChatOpenAI(model_name="gpt-4o-mini-2024-07-18", temperature=0.7)
+        return ChatOpenAI(model_name="gpt-4o-mini-2024-07-18", temperature=0.7, max_tokens=4000)
     elif model_name == "Llama":
         logging.info("Initializing Llama model - note this is a gated model and requires special access")
         try:
@@ -183,7 +183,7 @@ def get_llm(model_name: str):
                 model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
                 provider="hf-inference",
                 api_key=hf_token,
-                max_tokens=500
+                max_tokens=4000  # Increased token limit
             )
         except Exception as e:
             logging.error(f"Failed to initialize Llama model: {str(e)}")
@@ -198,7 +198,7 @@ def get_llm(model_name: str):
                 model="Qwen/Qwen1.5-0.5B-Chat",  # Much smaller model (0.5B vs 7B)
                 provider="hf-inference",
                 api_key=hf_token,
-                max_tokens=500
+                max_tokens=4000
             )
         except Exception as e:
             error_msg = str(e)
@@ -217,7 +217,7 @@ def get_llm(model_name: str):
                 model="mistralai/Mistral-7B-Instruct-v0.3",
                 provider="hf-inference",
                 api_key=hf_token,
-                max_tokens=500
+                max_tokens=4000
             )
         except Exception as e:
             error_msg = str(e)
@@ -231,7 +231,7 @@ def get_llm(model_name: str):
                         model="mistralai/Mistral-7B-Instruct-v0.1",
                         provider="hf-inference",
                         api_key=hf_token,
-                        max_tokens=500
+                        max_tokens=4000
                     )
                 except Exception:
                     pass
@@ -245,7 +245,7 @@ def get_llm(model_name: str):
                 model="google/gemma-2b-it",
                 provider="hf-inference",
                 api_key=hf_token,
-                max_tokens=500
+                max_tokens=4000
             )
         except Exception as e:
             logging.error(f"Failed to initialize Gemma model: {str(e)}")
@@ -257,7 +257,7 @@ def get_llm(model_name: str):
                 model="microsoft/phi-3-mini-4k-instruct",  # Use mini-4k which is smaller
                 provider="hf-inference",
                 api_key=hf_token,
-                max_tokens=500
+                max_tokens=4000
             )
         except Exception as e:
             logging.error(f"Failed to initialize Phi-3 model: {str(e)}")
@@ -284,7 +284,7 @@ def get_llm(model_name: str):
                         model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
                         provider="hf-inference",
                         api_key=hf_token,
-                        max_tokens=500
+                        max_tokens=4000
                     )
                 except Exception as fallback_error:
                     logging.error(f"Fallback model also failed: {str(fallback_error)}")
@@ -379,38 +379,45 @@ async def generate_business_plan(data: GenerateRequest):
     if data.business_idea:
         business_context = f"\nBusiness Idea: {data.business_idea}\n"
     
+    # Add South Africa as the country context
+    country_context = "\nCountry: South Africa\n"
+    business_context += country_context
+    
     plan_template_str = (
         "Based on the following responses to business questions:{business_context}\n"
         "{q_and_a}\n\n"
-        "Generate a detailed business plan in Markdown format. Your response must include exactly the following sections and subheadings using proper Markdown syntax:\n\n"
+        "Generate a detailed business plan in Markdown format specifically for a South African business. "
+        "All financial figures should be in South African Rand (ZAR). Consider South African business regulations, "
+        "market conditions, economic factors, and cultural context throughout the plan.\n\n"
+        "Your response must include exactly the following sections and subheadings using proper Markdown syntax:\n\n"
         "## Executive Summary\n"
         "Provide a brief summary of the business, its purpose, and key objectives.\n\n"
         "## Market Analysis\n"
         "### Industry Overview\n"
-        "Describe the overall industry trends and market dynamics.\n\n"
+        "Describe the overall industry trends and market dynamics in South Africa.\n\n"
         "### Target Market\n"
-        "Define the primary target audience, including demographics and psychographics.\n\n"
+        "Define the primary target audience in South Africa, including demographics and psychographics.\n\n"
         "### Competitive Analysis\n"
-        "Analyze the competitive landscape and explain what differentiates the business from competitors.\n\n"
+        "Analyze the South African competitive landscape and explain what differentiates the business from competitors.\n\n"
         "## Operations Plan\n"
         "### Location\n"
-        "Specify the business location(s) and rationale for the choice.\n\n"
+        "Specify the business location(s) in South Africa and rationale for the choice.\n\n"
         "### Distribution\n"
-        "Explain the distribution channels and logistics plan for delivering the product or service.\n\n"
+        "Explain the distribution channels and logistics plan for delivering the product or service within South Africa.\n\n"
         "### Staffing\n"
-        "Outline the staffing requirements and key roles necessary for operations.\n\n"
+        "Outline the staffing requirements and key roles necessary for operations, including considerations for South African labor laws.\n\n"
         "## Financial Plan\n"
         "### Revenue Streams\n"
-        "Identify the primary and secondary sources of revenue.\n\n"
+        "Identify the primary and secondary sources of revenue, with all figures in South African Rand (ZAR).\n\n"
         "### Cost Structure\n"
-        "Detail the major cost components and how costs will be managed.\n\n"
+        "Detail the major cost components and how costs will be managed, with consideration for South African pricing and expenses.\n\n"
         "### Funding Requirements\n"
-        "Specify the funding needed to launch and sustain the business.\n\n"
+        "Specify the funding needed to launch and sustain the business in South Africa, including any local funding options.\n\n"
         "### Profitability\n"
-        "Outline the financial projections and timeline to profitability.\n\n"
+        "Outline the financial projections and timeline to profitability within the South African market.\n\n"
         "## Conclusion\n"
-        "Summarize the vision, key takeaways, and future direction of the business.\n\n"
-        "Follow this template exactly, and be concise."
+        "Summarize the vision, key takeaways, and future direction of the business in South Africa.\n\n"
+        "Make sure to complete ALL sections above. Do not cut off your response in the middle of a section. All financial figures should be in South African Rand (ZAR) with appropriate market rates."
     )
     
     plan_prompt = PromptTemplate(
@@ -422,6 +429,13 @@ async def generate_business_plan(data: GenerateRequest):
         # Get the LLM instance based on the selected model.
         logging.info(f"Initializing model: {data.model}")
         llm_selected = get_llm(data.model)
+        
+        # Ensure max_tokens is set high enough to avoid truncation
+        if hasattr(llm_selected, 'max_tokens'):
+            original_max_tokens = llm_selected.max_tokens
+            llm_selected.max_tokens = 6000  # Increase max tokens to ensure full completion
+            logging.info(f"Increased max_tokens from {original_max_tokens} to {llm_selected.max_tokens}")
+        
         plan_chain = LLMChain(llm=llm_selected, prompt=plan_prompt)
         
         logging.info(f"Generated prompt for business plan with model {data.model}")
@@ -430,8 +444,13 @@ async def generate_business_plan(data: GenerateRequest):
             logging.info(f"Generating business plan with model: {data.model}")
             full_plan = await asyncio.to_thread(plan_chain.run, {"q_and_a": q_and_a, "business_context": business_context})
             logging.info(f"Successfully generated business plan with model: {data.model}")
+            
+            # Check if plan seems truncated
+            if "## Conclusion" not in full_plan:
+                logging.warning("Business plan appears to be truncated (missing Conclusion section)")
+            
             return {
-                "summary": "Generated Business Plan",
+                "summary": "Generated Business Plan for South Africa",
                 "plan": full_plan,
             }
         except Exception as e: